mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-05 02:13:24 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
443 lines
16 KiB
Python
443 lines
16 KiB
Python
"""Tests for the SchedulerTool."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import uuid
|
|
from datetime import datetime, timedelta, timezone
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
|
|
# Pre-import to stabilise the ToolManager.load_tools walk's import order
|
|
# (avoids the mcp_tool ↔ docsgpt.api.user circular when ToolManager
|
|
# instantiation is the first reachable importer in a test process).
|
|
import docsgpt.api.user.tools.mcp # noqa: F401
|
|
|
|
from docsgpt.agents.tools.scheduler import SchedulerTool # noqa: E402
|
|
from docsgpt.core.settings import settings # noqa: E402
|
|
from docsgpt.storage.db.repositories.schedules import SchedulesRepository # noqa: E402
|
|
|
|
|
|
@pytest.fixture
|
|
def patch_sessions(pg_conn):
|
|
from contextlib import contextmanager
|
|
|
|
@contextmanager
|
|
def _ctx():
|
|
yield pg_conn
|
|
|
|
with patch(
|
|
"docsgpt.agents.tools.scheduler.db_session", _ctx,
|
|
), patch(
|
|
"docsgpt.agents.tools.scheduler.db_readonly", _ctx,
|
|
):
|
|
yield
|
|
|
|
|
|
def _make_agent(conn, user_id: str = "u1") -> str:
|
|
row = conn.execute(
|
|
text(
|
|
"INSERT INTO agents (user_id, name, status) "
|
|
"VALUES (:u, 'a', 'draft') RETURNING id"
|
|
),
|
|
{"u": user_id},
|
|
).fetchone()
|
|
return str(row[0])
|
|
|
|
|
|
def _make_tool(name="scheduler", *, user_id="u1", agent_id=None, conversation_id=None):
|
|
return SchedulerTool(
|
|
tool_config={
|
|
"agent_id": agent_id,
|
|
"conversation_id": conversation_id,
|
|
},
|
|
user_id=user_id,
|
|
)
|
|
|
|
|
|
class TestGuards:
|
|
def test_requires_user_id(self):
|
|
tool = SchedulerTool(tool_config={"agent_id": str(uuid.uuid4())})
|
|
assert "user_id" in tool.execute_action("schedule_task", instruction="x")
|
|
|
|
def test_rejects_invalid_agent_id(self):
|
|
tool = _make_tool(user_id="u1", agent_id="not-a-uuid")
|
|
assert "invalid agent_id" in tool.execute_action(
|
|
"schedule_task", instruction="x"
|
|
)
|
|
|
|
def test_requires_agent_or_conversation(self):
|
|
# Neither agent_id nor conversation_id → hard error (webhook caller
|
|
# outside any chat); scheduler can't operate without a conversation home.
|
|
tool = _make_tool(user_id="u1", agent_id=None, conversation_id=None)
|
|
out = tool.execute_action("schedule_task", instruction="x")
|
|
assert "conversation_id" in out or "conversation home" in out
|
|
|
|
|
|
class TestScheduleTask:
|
|
def test_creates_with_delay(self, pg_conn, patch_sessions):
|
|
agent_id = _make_agent(pg_conn)
|
|
tool = _make_tool(user_id="u1", agent_id=agent_id, conversation_id=None)
|
|
out = tool.execute_action(
|
|
"schedule_task", instruction="say hi", delay="2h",
|
|
)
|
|
parsed = json.loads(out)
|
|
assert "task_id" in parsed
|
|
assert "resolved_run_at" in parsed
|
|
row = SchedulesRepository(pg_conn).get(parsed["task_id"], "u1")
|
|
assert row is not None
|
|
assert row["trigger_type"] == "once"
|
|
assert row["created_via"] == "chat"
|
|
fire = datetime.fromisoformat(parsed["resolved_run_at"].replace("Z", "+00:00"))
|
|
delta = fire - datetime.now(timezone.utc)
|
|
assert timedelta(minutes=119) <= delta <= timedelta(minutes=121)
|
|
|
|
def test_creates_with_run_at(self, pg_conn, patch_sessions):
|
|
agent_id = _make_agent(pg_conn)
|
|
tool = _make_tool(user_id="u1", agent_id=agent_id)
|
|
fire = (datetime.now(timezone.utc) + timedelta(hours=3)).isoformat()
|
|
out = tool.execute_action(
|
|
"schedule_task", instruction="x", run_at=fire,
|
|
)
|
|
parsed = json.loads(out)
|
|
assert "task_id" in parsed
|
|
|
|
def test_rejects_both_delay_and_run_at(self, pg_conn, patch_sessions):
|
|
agent_id = _make_agent(pg_conn)
|
|
tool = _make_tool(user_id="u1", agent_id=agent_id)
|
|
out = tool.execute_action(
|
|
"schedule_task", instruction="x", delay="30m",
|
|
run_at="2030-01-01T00:00:00Z",
|
|
)
|
|
assert "only one" in out
|
|
|
|
def test_rejects_past_run_at(self, pg_conn, patch_sessions):
|
|
agent_id = _make_agent(pg_conn)
|
|
tool = _make_tool(user_id="u1", agent_id=agent_id)
|
|
past = (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat()
|
|
out = tool.execute_action("schedule_task", instruction="x", run_at=past)
|
|
assert "past" in out
|
|
|
|
def test_rejects_beyond_horizon(
|
|
self, pg_conn, patch_sessions, monkeypatch
|
|
):
|
|
monkeypatch.setattr(settings, "SCHEDULE_ONCE_MAX_HORIZON", 3600)
|
|
agent_id = _make_agent(pg_conn)
|
|
tool = _make_tool(user_id="u1", agent_id=agent_id)
|
|
far = (datetime.now(timezone.utc) + timedelta(hours=10)).isoformat()
|
|
out = tool.execute_action("schedule_task", instruction="x", run_at=far)
|
|
assert "horizon" in out
|
|
|
|
|
|
class TestQuota:
|
|
def test_quota_enforced(self, pg_conn, patch_sessions, monkeypatch):
|
|
monkeypatch.setattr(settings, "SCHEDULE_MAX_PER_USER", 2)
|
|
agent_id = _make_agent(pg_conn)
|
|
tool = _make_tool(user_id="u1", agent_id=agent_id)
|
|
for _ in range(2):
|
|
out = tool.execute_action(
|
|
"schedule_task", instruction="x", delay="1h",
|
|
)
|
|
assert "task_id" in out
|
|
out = tool.execute_action(
|
|
"schedule_task", instruction="x", delay="1h",
|
|
)
|
|
assert "maximum" in out
|
|
|
|
|
|
class TestListAndCancel:
|
|
def test_list_returns_pending(self, pg_conn, patch_sessions):
|
|
agent_id = _make_agent(pg_conn)
|
|
tool = _make_tool(user_id="u1", agent_id=agent_id)
|
|
for _ in range(3):
|
|
tool.execute_action(
|
|
"schedule_task", instruction="x", delay="1h",
|
|
)
|
|
listed = json.loads(tool.execute_action("list_scheduled_tasks"))
|
|
assert len(listed["tasks"]) == 3
|
|
assert all(t["status"] == "active" for t in listed["tasks"])
|
|
|
|
def test_cancel_flips_status(self, pg_conn, patch_sessions):
|
|
agent_id = _make_agent(pg_conn)
|
|
tool = _make_tool(user_id="u1", agent_id=agent_id)
|
|
created = json.loads(
|
|
tool.execute_action("schedule_task", instruction="x", delay="1h")
|
|
)
|
|
out = tool.execute_action(
|
|
"cancel_scheduled_task", task_id=created["task_id"]
|
|
)
|
|
assert "cancelled" in out
|
|
row = SchedulesRepository(pg_conn).get(created["task_id"], "u1")
|
|
assert row["status"] == "cancelled"
|
|
|
|
def test_cancel_unknown_id_rejected(self, pg_conn, patch_sessions):
|
|
agent_id = _make_agent(pg_conn)
|
|
tool = _make_tool(user_id="u1", agent_id=agent_id)
|
|
out = tool.execute_action(
|
|
"cancel_scheduled_task", task_id="not-a-uuid",
|
|
)
|
|
assert "valid id" in out
|
|
|
|
|
|
class TestActionsMetadata:
|
|
def test_actions_listed(self):
|
|
tool = SchedulerTool()
|
|
names = {a["name"] for a in tool.get_actions_metadata()}
|
|
assert names == {
|
|
"schedule_task", "list_scheduled_tasks", "cancel_scheduled_task",
|
|
}
|
|
|
|
|
|
class TestAgentlessInvocation:
|
|
def test_agentless_creates_schedule_with_null_agent_id(
|
|
self, pg_conn, patch_sessions,
|
|
):
|
|
"""Agentless chat → scheduler.schedule_task → row with NULL agent_id."""
|
|
conv_id = pg_conn.execute(
|
|
text(
|
|
"INSERT INTO conversations (user_id, name) "
|
|
"VALUES ('u1', 'origin') RETURNING id"
|
|
)
|
|
).fetchone()[0]
|
|
tool = _make_tool(
|
|
user_id="u1", agent_id=None, conversation_id=str(conv_id),
|
|
)
|
|
out = tool.execute_action(
|
|
"schedule_task", instruction="ping me later", delay="1h",
|
|
)
|
|
parsed = json.loads(out)
|
|
assert "task_id" in parsed
|
|
row = SchedulesRepository(pg_conn).get(parsed["task_id"], "u1")
|
|
assert row is not None
|
|
assert row["agent_id"] is None
|
|
assert row["trigger_type"] == "once"
|
|
assert row["created_via"] == "chat"
|
|
assert str(row["origin_conversation_id"]) == str(conv_id)
|
|
|
|
def test_agentless_list_scoped_to_conversation(
|
|
self, pg_conn, patch_sessions,
|
|
):
|
|
"""Agentless list_scheduled_tasks scopes to user + origin conversation."""
|
|
conv_a = pg_conn.execute(
|
|
text(
|
|
"INSERT INTO conversations (user_id, name) "
|
|
"VALUES ('u1', 'a') RETURNING id"
|
|
)
|
|
).fetchone()[0]
|
|
conv_b = pg_conn.execute(
|
|
text(
|
|
"INSERT INTO conversations (user_id, name) "
|
|
"VALUES ('u1', 'b') RETURNING id"
|
|
)
|
|
).fetchone()[0]
|
|
tool_a = _make_tool(
|
|
user_id="u1", agent_id=None, conversation_id=str(conv_a),
|
|
)
|
|
tool_b = _make_tool(
|
|
user_id="u1", agent_id=None, conversation_id=str(conv_b),
|
|
)
|
|
tool_a.execute_action(
|
|
"schedule_task", instruction="in-a", delay="1h",
|
|
)
|
|
tool_a.execute_action(
|
|
"schedule_task", instruction="in-a-2", delay="2h",
|
|
)
|
|
tool_b.execute_action(
|
|
"schedule_task", instruction="in-b", delay="3h",
|
|
)
|
|
listed_a = json.loads(tool_a.execute_action("list_scheduled_tasks"))
|
|
listed_b = json.loads(tool_b.execute_action("list_scheduled_tasks"))
|
|
assert len(listed_a["tasks"]) == 2
|
|
assert len(listed_b["tasks"]) == 1
|
|
assert all(t["status"] == "active" for t in listed_a["tasks"])
|
|
|
|
def test_agentless_cancel_blocked_for_other_conversation(
|
|
self, pg_conn, patch_sessions,
|
|
):
|
|
"""A user can't cancel tasks created in another agentless chat."""
|
|
conv_a = pg_conn.execute(
|
|
text(
|
|
"INSERT INTO conversations (user_id, name) "
|
|
"VALUES ('u1', 'a') RETURNING id"
|
|
)
|
|
).fetchone()[0]
|
|
conv_b = pg_conn.execute(
|
|
text(
|
|
"INSERT INTO conversations (user_id, name) "
|
|
"VALUES ('u1', 'b') RETURNING id"
|
|
)
|
|
).fetchone()[0]
|
|
tool_a = _make_tool(
|
|
user_id="u1", agent_id=None, conversation_id=str(conv_a),
|
|
)
|
|
tool_b = _make_tool(
|
|
user_id="u1", agent_id=None, conversation_id=str(conv_b),
|
|
)
|
|
created = json.loads(
|
|
tool_a.execute_action(
|
|
"schedule_task", instruction="x", delay="1h",
|
|
)
|
|
)
|
|
out = tool_b.execute_action(
|
|
"cancel_scheduled_task", task_id=created["task_id"],
|
|
)
|
|
assert "not found" in out
|
|
|
|
def test_agentless_cancel_succeeds_in_own_conversation(
|
|
self, pg_conn, patch_sessions,
|
|
):
|
|
conv = pg_conn.execute(
|
|
text(
|
|
"INSERT INTO conversations (user_id, name) "
|
|
"VALUES ('u1', 'a') RETURNING id"
|
|
)
|
|
).fetchone()[0]
|
|
tool = _make_tool(
|
|
user_id="u1", agent_id=None, conversation_id=str(conv),
|
|
)
|
|
created = json.loads(
|
|
tool.execute_action("schedule_task", instruction="x", delay="1h")
|
|
)
|
|
out = tool.execute_action(
|
|
"cancel_scheduled_task", task_id=created["task_id"],
|
|
)
|
|
assert "cancelled" in out
|
|
|
|
def test_agentless_snapshot_allowlist_lists_user_tools(
|
|
self, pg_conn, patch_sessions,
|
|
):
|
|
"""Agentless schedule captures the user's non-approval tools at fire-time."""
|
|
from docsgpt.agents.tools.scheduler import _safe_default_allowlist
|
|
from docsgpt.storage.db.repositories.user_tools import (
|
|
UserToolsRepository,
|
|
)
|
|
|
|
# Seed an explicit non-approval user tool.
|
|
user_tool = UserToolsRepository(pg_conn).create(
|
|
"u1", "read_webpage", config={}, actions=[
|
|
{"name": "fetch", "active": True, "require_approval": False},
|
|
], status=True,
|
|
)
|
|
conv = pg_conn.execute(
|
|
text(
|
|
"INSERT INTO conversations (user_id, name) "
|
|
"VALUES ('u1', 'a') RETURNING id"
|
|
)
|
|
).fetchone()[0]
|
|
tool = _make_tool(
|
|
user_id="u1", agent_id=None, conversation_id=str(conv),
|
|
)
|
|
out = tool.execute_action(
|
|
"schedule_task", instruction="x", delay="1h",
|
|
)
|
|
parsed = json.loads(out)
|
|
row = SchedulesRepository(pg_conn).get(parsed["task_id"], "u1")
|
|
# The explicit user_tools row is in the snapshot (approval=False).
|
|
assert str(user_tool["id"]) in (row["tool_allowlist"] or [])
|
|
# Direct allowlist call returns the same set.
|
|
ids = _safe_default_allowlist(None, "u1")
|
|
assert str(user_tool["id"]) in ids
|
|
|
|
|
|
class TestAllowlistSnapshotSemantics:
|
|
"""The schedule's ``tool_allowlist`` is a **pre-auth snapshot**, not a
|
|
visibility cap. The LLM sees the user's *current* tools at fire time
|
|
(via ``ToolExecutor._get_user_tools``); the snapshot only governs
|
|
whether an approval-gated tool can run unattended."""
|
|
|
|
def test_tool_added_after_creation_is_visible_at_fire_time(
|
|
self, pg_conn, patch_sessions,
|
|
):
|
|
"""Schedule captures the allowlist at creation; a tool added later is
|
|
visible at fire time (resolver re-queries) but isn't in the snapshot."""
|
|
from docsgpt.agents.tools.scheduler import _safe_default_allowlist
|
|
from docsgpt.storage.db.repositories.user_tools import (
|
|
UserToolsRepository,
|
|
)
|
|
|
|
pg_conn.execute(
|
|
text(
|
|
"INSERT INTO conversations (user_id, name) "
|
|
"VALUES ('u1', 'snap-add') RETURNING id"
|
|
)
|
|
).fetchone()
|
|
# Snapshot the allowlist BEFORE adding the new tool.
|
|
snapshot_before = _safe_default_allowlist(None, "u1")
|
|
|
|
# User adds an approval-gated tool AFTER schedule creation.
|
|
added = UserToolsRepository(pg_conn).create(
|
|
"u1", "telegram",
|
|
config={}, actions=[
|
|
{"name": "send", "active": True, "require_approval": True},
|
|
], status=True,
|
|
)
|
|
|
|
# The snapshot does NOT include the post-creation tool.
|
|
assert str(added["id"]) not in snapshot_before
|
|
# …but the LLM sees it at fire time (current resolver state).
|
|
snapshot_after = _safe_default_allowlist(None, "u1")
|
|
# An approval-gated tool is excluded from the snapshot regardless,
|
|
# but it IS in ``list_active_for_user`` (what the LLM's tool_executor
|
|
# uses) — make that explicit:
|
|
ids_now = {
|
|
str(r["id"]) for r in
|
|
UserToolsRepository(pg_conn).list_active_for_user("u1")
|
|
}
|
|
assert str(added["id"]) in ids_now
|
|
# And approval-gated still skipped from the safe allowlist.
|
|
assert str(added["id"]) not in snapshot_after
|
|
|
|
def test_tool_deleted_between_creation_and_fire_is_invisible(
|
|
self, pg_conn, patch_sessions,
|
|
):
|
|
"""A tool deleted between schedule creation and fire is gone for the
|
|
LLM at fire time (the resolver lists the current state)."""
|
|
from docsgpt.agents.tools.scheduler import _safe_default_allowlist
|
|
from docsgpt.storage.db.repositories.user_tools import (
|
|
UserToolsRepository,
|
|
)
|
|
|
|
repo = UserToolsRepository(pg_conn)
|
|
existing = repo.create(
|
|
"u1", "read_webpage",
|
|
config={}, actions=[
|
|
{"name": "fetch", "active": True, "require_approval": False},
|
|
], status=True,
|
|
)
|
|
# Snapshot at creation includes it (non-approval).
|
|
snapshot = _safe_default_allowlist(None, "u1")
|
|
assert str(existing["id"]) in snapshot
|
|
|
|
# User deletes it; fire-time resolver no longer surfaces it.
|
|
repo.delete(str(existing["id"]), "u1")
|
|
ids_now = {r["id"] for r in repo.list_active_for_user("u1")}
|
|
assert str(existing["id"]) not in ids_now
|
|
# And the freshly-recomputed allowlist drops it too.
|
|
snapshot_after = _safe_default_allowlist(None, "u1")
|
|
assert str(existing["id"]) not in snapshot_after
|
|
|
|
|
|
class TestInternalFlag:
|
|
def test_internal_true(self):
|
|
assert SchedulerTool.internal is True
|
|
|
|
def test_not_in_tool_manager_auto_load(self):
|
|
from docsgpt.agents.tools.tool_manager import ToolManager
|
|
|
|
tm = ToolManager(config={})
|
|
assert "scheduler" not in tm.tools
|
|
|
|
def test_load_tool_special_case_still_works(self):
|
|
from docsgpt.agents.tools.tool_manager import ToolManager
|
|
|
|
tm = ToolManager(config={})
|
|
tool = tm.load_tool(
|
|
"scheduler",
|
|
tool_config={"agent_id": str(uuid.uuid4())},
|
|
user_id="u1",
|
|
)
|
|
assert isinstance(tool, SchedulerTool)
|
|
assert tool.user_id == "u1"
|