mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 10:14:32 +00:00
feat: default tools (#2485)
* feat: default tools * fix: tests * feat scheduler * fix: scheduler UI * Agent switch breadcrumbs * fix: minor fixes to scheduler * fix: tests
This commit is contained in:
107 files changed
+14212
-458
No files matched your search
@@ -246,6 +246,68 @@ class TestCompleteStreamMethod:
|
||||
mock_reserve.assert_called_once()
|
||||
mock_finalize.assert_called_once()
|
||||
|
||||
def test_tool_executor_conversation_id_set_after_reserve(
|
||||
self, mock_mongo_db, flask_app,
|
||||
):
|
||||
"""Regression: ``save_user_question`` may mint a fresh
|
||||
``conversation_id`` (first turn). The propagation MUST land on
|
||||
``agent.tool_executor.conversation_id`` BEFORE ``agent.gen`` runs,
|
||||
so tools needing a conversation home (``scheduler`` in an agentless
|
||||
chat) see it on the very first call.
|
||||
"""
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
fresh_conv_id = str(uuid.uuid4())
|
||||
seen_conv_id_on_gen: dict = {}
|
||||
|
||||
mock_agent = MagicMock()
|
||||
tool_executor = MagicMock()
|
||||
# Start with no conversation_id — the propagation must set it.
|
||||
tool_executor.conversation_id = None
|
||||
mock_agent.tool_executor = tool_executor
|
||||
|
||||
def _gen(**_kwargs):
|
||||
# Capture the executor's id at the exact moment gen runs;
|
||||
# this is what tools see when called from the agent loop.
|
||||
seen_conv_id_on_gen["value"] = (
|
||||
mock_agent.tool_executor.conversation_id
|
||||
)
|
||||
yield {"answer": "ok"}
|
||||
|
||||
mock_agent.gen.side_effect = _gen
|
||||
mock_agent.gen.return_value = None # use side_effect instead
|
||||
|
||||
with patch.object(
|
||||
resource.conversation_service, "save_user_question"
|
||||
) as mock_reserve, patch.object(
|
||||
resource.conversation_service, "finalize_message",
|
||||
return_value=True,
|
||||
):
|
||||
mock_reserve.return_value = {
|
||||
"conversation_id": fresh_conv_id,
|
||||
"message_id": str(uuid.uuid4()),
|
||||
"request_id": "req-prop",
|
||||
}
|
||||
|
||||
list(
|
||||
resource.complete_stream(
|
||||
question="schedule something",
|
||||
agent=mock_agent,
|
||||
conversation_id=None, # caller had no conv yet
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": "user-prop"},
|
||||
should_save_conversation=True,
|
||||
)
|
||||
)
|
||||
|
||||
# The fresh id reserved by save_user_question must reach the
|
||||
# tool_executor before agent.gen consumes it.
|
||||
assert seen_conv_id_on_gen["value"] == fresh_conv_id
|
||||
assert tool_executor.conversation_id == fresh_conv_id
|
||||
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
|
||||
@@ -429,7 +429,7 @@ class TestConfigureRetriever:
|
||||
assert sp.retriever_config["retriever_name"] == "hybrid_search"
|
||||
assert sp.retriever_config["chunks"] == 5
|
||||
|
||||
def test_request_overrides_agent(self):
|
||||
def test_agent_wins_over_request_on_agent_bound(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
@@ -438,9 +438,33 @@ class TestConfigureRetriever:
|
||||
)
|
||||
sp._agent_data = {"retriever": "hybrid_search", "chunks": 5}
|
||||
sp._configure_retriever()
|
||||
assert sp.retriever_config["retriever_name"] == "hybrid_search"
|
||||
assert sp.retriever_config["chunks"] == 5
|
||||
|
||||
def test_body_wins_on_agentless(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor(
|
||||
{"retriever": "duckdb", "chunks": 7}, {"sub": "u"},
|
||||
)
|
||||
sp._configure_retriever()
|
||||
assert sp.retriever_config["retriever_name"] == "duckdb"
|
||||
assert sp.retriever_config["chunks"] == 7
|
||||
|
||||
def test_agent_bound_drops_body_chunks_and_retriever(self):
|
||||
# Missing agent values fall back to system defaults, not body's.
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor(
|
||||
{"retriever": "duckdb", "chunks": 7}, {"sub": "u"},
|
||||
)
|
||||
sp._agent_data = {}
|
||||
sp._configure_retriever()
|
||||
assert sp.retriever_config["retriever_name"] == "classic"
|
||||
assert sp.retriever_config["chunks"] == 2
|
||||
|
||||
def test_invalid_agent_chunks_falls_back(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
@@ -569,10 +593,11 @@ class TestPreFetchTools:
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is None
|
||||
|
||||
def test_no_user_tools_returns_none(self, pg_conn):
|
||||
def test_no_template_skips_default_tool_prefetch(self, pg_conn):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
|
||||
sp = StreamProcessor({}, {"sub": "no-tools-user"})
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
@@ -580,3 +605,347 @@ class TestPreFetchTools:
|
||||
):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is None
|
||||
|
||||
def test_no_template_skips_only_default_rows_not_explicit(self, pg_conn):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
from application.storage.db.repositories.user_tools import (
|
||||
UserToolsRepository,
|
||||
)
|
||||
|
||||
UserToolsRepository(pg_conn).create(
|
||||
user_id="u-explicit-prefetch", name="read_webpage", status=True
|
||||
)
|
||||
sp = StreamProcessor({}, {"sub": "u-explicit-prefetch"})
|
||||
fetched = []
|
||||
|
||||
def _fake_fetch(tool_doc, required_actions):
|
||||
fetched.append(tool_doc)
|
||||
return {"ok": True}
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
True,
|
||||
), patch.object(sp, "_fetch_tool_data", _fake_fetch):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is not None
|
||||
assert "read_webpage" in got
|
||||
assert all(not d.get("default") for d in fetched)
|
||||
assert any(d.get("name") == "read_webpage" for d in fetched)
|
||||
|
||||
def test_default_tool_prefetched_when_template_references_it(
|
||||
self, pg_conn
|
||||
):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
|
||||
sp = StreamProcessor({}, {"sub": "u-tpl-default"})
|
||||
sp._required_tool_actions = {"read_webpage": {None}}
|
||||
fetched = []
|
||||
|
||||
def _fake_fetch(tool_doc, required_actions):
|
||||
fetched.append(tool_doc)
|
||||
return {"ok": True}
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
True,
|
||||
), patch.object(sp, "_fetch_tool_data", _fake_fetch):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is not None
|
||||
assert any(
|
||||
d.get("name") == "read_webpage" and d.get("default")
|
||||
for d in fetched
|
||||
)
|
||||
# Defaults are reachable by synthetic id only — not by name.
|
||||
assert default_tool_id("read_webpage") in got
|
||||
|
||||
def test_agent_bound_invocation_omits_default_tool_prefetch(self, pg_conn):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
|
||||
sp = StreamProcessor({"agent_id": "agent-xyz"}, {"sub": "u-ag"})
|
||||
sp._required_tool_actions = {"read_webpage": {None}}
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
True,
|
||||
):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is None
|
||||
|
||||
def test_template_name_key_favors_explicit_over_default(self, pg_conn):
|
||||
"""An explicit row and the synthesized default of the same name
|
||||
coexist: name key stays on the explicit, default reachable by
|
||||
synthetic id only."""
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
from application.storage.db.repositories.user_tools import (
|
||||
UserToolsRepository,
|
||||
)
|
||||
|
||||
user = "u-collision"
|
||||
explicit = UserToolsRepository(pg_conn).create(
|
||||
user_id=user, name="read_webpage", status=True,
|
||||
)
|
||||
explicit_id = str(explicit["id"])
|
||||
default_id = default_tool_id("read_webpage")
|
||||
|
||||
sp = StreamProcessor({}, {"sub": user})
|
||||
sp._required_tool_actions = {"read_webpage": {None}}
|
||||
|
||||
def _fake_fetch(tool_doc, required_actions):
|
||||
return {
|
||||
"is_default": bool(tool_doc.get("default")),
|
||||
"id": str(tool_doc.get("_id") or tool_doc.get("id")),
|
||||
}
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
|
||||
True,
|
||||
), patch.object(sp, "_fetch_tool_data", _fake_fetch):
|
||||
got = sp.pre_fetch_tools()
|
||||
assert got is not None
|
||||
assert got["read_webpage"]["is_default"] is False
|
||||
assert got["read_webpage"]["id"] == explicit_id
|
||||
assert got[explicit_id]["is_default"] is False
|
||||
assert got[default_id]["is_default"] is True
|
||||
|
||||
|
||||
class TestValidateAndSetModelAgentAuthority:
|
||||
"""Agent-bound chats: agent's ``default_model_id`` is authoritative."""
|
||||
|
||||
def test_agent_bound_ignores_body_model_id(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor({"model_id": "body-model"}, {"sub": "caller"})
|
||||
sp._agent_data = {"user": "owner"}
|
||||
sp.agent_config = {
|
||||
"default_model_id": "agent-model",
|
||||
"user_id": "owner",
|
||||
}
|
||||
captured: list = []
|
||||
|
||||
def _fake_validate(model_id, user_id=None):
|
||||
captured.append((model_id, user_id))
|
||||
return True
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.stream_processor.validate_model_id",
|
||||
side_effect=_fake_validate,
|
||||
), patch(
|
||||
"application.api.answer.services.stream_processor.get_default_model_id",
|
||||
return_value="global-default",
|
||||
):
|
||||
sp._validate_and_set_model()
|
||||
assert sp.model_id == "agent-model"
|
||||
# Resolved under the agent owner, not the caller.
|
||||
assert sp.model_user_id == "owner"
|
||||
assert ("agent-model", "owner") in captured
|
||||
|
||||
def test_agent_bound_no_default_falls_back_to_system(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor({"model_id": "body-model"}, {"sub": "u"})
|
||||
sp._agent_data = {"user": "u"}
|
||||
sp.agent_config = {"default_model_id": "", "user_id": "u"}
|
||||
with patch(
|
||||
"application.api.answer.services.stream_processor.validate_model_id",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"application.api.answer.services.stream_processor.get_default_model_id",
|
||||
return_value="global-default",
|
||||
):
|
||||
sp._validate_and_set_model()
|
||||
assert sp.model_id == "global-default"
|
||||
assert sp.model_user_id is None
|
||||
|
||||
def test_agentless_body_model_still_wins(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
sp = StreamProcessor({"model_id": "body-model"}, {"sub": "u"})
|
||||
sp._agent_data = None
|
||||
with patch(
|
||||
"application.api.answer.services.stream_processor.validate_model_id",
|
||||
return_value=True,
|
||||
):
|
||||
sp._validate_and_set_model()
|
||||
assert sp.model_id == "body-model"
|
||||
assert sp.model_user_id == "u"
|
||||
|
||||
|
||||
class TestGetDataFromApiKeySourceUnion:
|
||||
"""`_get_data_from_api_key`: primary ∪ extras, deduplicated, primary first."""
|
||||
|
||||
def _make_sp(self):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
return StreamProcessor({}, {"sub": "u"})
|
||||
|
||||
def test_union_primary_and_extras(self, pg_conn):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-merge-both"
|
||||
sources_repo = SourcesRepository(pg_conn)
|
||||
primary = sources_repo.create(name="primary", user_id=owner)
|
||||
extra1 = sources_repo.create(name="extra1", user_id=owner)
|
||||
extra2 = sources_repo.create(name="extra2", user_id=owner)
|
||||
|
||||
agent = AgentsRepository(pg_conn).create(
|
||||
owner, "agent-merge", "published",
|
||||
key="merge-key",
|
||||
source_id=str(primary["id"]),
|
||||
extra_source_ids=[str(extra1["id"]), str(extra2["id"])],
|
||||
retriever="hybrid",
|
||||
chunks=5,
|
||||
)
|
||||
assert agent is not None
|
||||
|
||||
sp = self._make_sp()
|
||||
with _patch_db(pg_conn):
|
||||
data = sp._get_data_from_api_key("merge-key")
|
||||
ids = [s["id"] for s in data["sources"]]
|
||||
assert ids == [
|
||||
str(primary["id"]),
|
||||
str(extra1["id"]),
|
||||
str(extra2["id"]),
|
||||
]
|
||||
assert data["source"] == str(primary["id"])
|
||||
|
||||
def test_only_primary(self, pg_conn):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-merge-primary-only"
|
||||
primary = SourcesRepository(pg_conn).create(
|
||||
name="primary", user_id=owner,
|
||||
)
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
owner, "primary-only", "published",
|
||||
key="primary-only-key",
|
||||
source_id=str(primary["id"]),
|
||||
extra_source_ids=[],
|
||||
)
|
||||
|
||||
sp = self._make_sp()
|
||||
with _patch_db(pg_conn):
|
||||
data = sp._get_data_from_api_key("primary-only-key")
|
||||
assert [s["id"] for s in data["sources"]] == [str(primary["id"])]
|
||||
assert data["source"] == str(primary["id"])
|
||||
|
||||
def test_only_extras(self, pg_conn):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-merge-extras-only"
|
||||
e1 = SourcesRepository(pg_conn).create(name="e1", user_id=owner)
|
||||
e2 = SourcesRepository(pg_conn).create(name="e2", user_id=owner)
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
owner, "extras-only", "published",
|
||||
key="extras-only-key",
|
||||
extra_source_ids=[str(e1["id"]), str(e2["id"])],
|
||||
)
|
||||
|
||||
sp = self._make_sp()
|
||||
with _patch_db(pg_conn):
|
||||
data = sp._get_data_from_api_key("extras-only-key")
|
||||
assert [s["id"] for s in data["sources"]] == [
|
||||
str(e1["id"]), str(e2["id"]),
|
||||
]
|
||||
assert data["source"] is None
|
||||
|
||||
def test_dedupe_primary_repeated_in_extras(self, pg_conn):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-merge-dedupe"
|
||||
primary = SourcesRepository(pg_conn).create(
|
||||
name="dup-primary", user_id=owner,
|
||||
)
|
||||
extra = SourcesRepository(pg_conn).create(
|
||||
name="dup-extra", user_id=owner,
|
||||
)
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
owner, "dedupe", "published",
|
||||
key="dedupe-key",
|
||||
source_id=str(primary["id"]),
|
||||
extra_source_ids=[str(primary["id"]), str(extra["id"])],
|
||||
)
|
||||
|
||||
sp = self._make_sp()
|
||||
with _patch_db(pg_conn):
|
||||
data = sp._get_data_from_api_key("dedupe-key")
|
||||
ids = [s["id"] for s in data["sources"]]
|
||||
assert ids == [str(primary["id"]), str(extra["id"])]
|
||||
|
||||
|
||||
class TestAgentBoundFieldsAuthoritative:
|
||||
"""End-to-end regression: agent's source/model/chunks/retriever win."""
|
||||
|
||||
def test_agent_values_win_over_body(self, pg_conn):
|
||||
from application.api.answer.services.stream_processor import (
|
||||
StreamProcessor,
|
||||
)
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
owner = "u-regr-agent-authority"
|
||||
primary = SourcesRepository(pg_conn).create(
|
||||
name="primary", user_id=owner,
|
||||
)
|
||||
extra = SourcesRepository(pg_conn).create(
|
||||
name="extra", user_id=owner,
|
||||
)
|
||||
AgentsRepository(pg_conn).create(
|
||||
owner, "authoritative", "published",
|
||||
key="auth-key",
|
||||
source_id=str(primary["id"]),
|
||||
extra_source_ids=[str(extra["id"])],
|
||||
default_model_id="model-A",
|
||||
retriever="hybrid",
|
||||
chunks=5,
|
||||
)
|
||||
|
||||
# Body sends different values for every field; all must be ignored.
|
||||
body = {
|
||||
"api_key": "auth-key",
|
||||
"model_id": "body-model-Z",
|
||||
"retriever": "duckdb",
|
||||
"chunks": 99,
|
||||
"active_docs": "body-source-id",
|
||||
}
|
||||
sp = StreamProcessor(body, {"sub": owner})
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.answer.services.stream_processor.validate_model_id",
|
||||
return_value=True,
|
||||
), patch(
|
||||
"application.api.answer.services.stream_processor.get_default_model_id",
|
||||
return_value="system-default",
|
||||
):
|
||||
sp._configure_agent()
|
||||
sp._validate_and_set_model()
|
||||
sp._configure_source()
|
||||
sp._configure_retriever()
|
||||
|
||||
assert sp.model_id == "model-A"
|
||||
assert sp.model_user_id == owner
|
||||
assert sp.agent_config["default_model_id"] == "model-A"
|
||||
assert sp.retriever_config["chunks"] == 5
|
||||
assert sp.retriever_config["retriever_name"] == "hybrid"
|
||||
assert sp.source == {
|
||||
"active_docs": [str(primary["id"]), str(extra["id"])],
|
||||
}
|
||||
@@ -0,0 +1,401 @@
|
||||
"""Tests for the scheduler dispatcher (engine-level, no Celery worker)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.scheduler_dispatcher import dispatch_due_runs
|
||||
from application.storage.db.repositories.schedule_runs import (
|
||||
ScheduleRunsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.schedules import SchedulesRepository
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
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])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patched_engine(pg_engine, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_dispatcher.get_engine",
|
||||
lambda: pg_engine,
|
||||
)
|
||||
yield pg_engine
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stub_enqueue(monkeypatch):
|
||||
"""Capture every execute_scheduled_run.apply_async."""
|
||||
enqueued: list[str] = []
|
||||
|
||||
class _Task:
|
||||
@staticmethod
|
||||
def apply_async(args=None, **kwargs):
|
||||
if args:
|
||||
enqueued.append(args[0])
|
||||
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.tasks.execute_scheduled_run", _Task
|
||||
)
|
||||
return enqueued
|
||||
|
||||
|
||||
def _create_schedule(engine, **kwargs):
|
||||
with engine.begin() as conn:
|
||||
return SchedulesRepository(conn).create(**kwargs)
|
||||
|
||||
|
||||
def _set_postgres_uri(monkeypatch, pg_engine):
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_dispatcher.settings",
|
||||
type("S", (), {
|
||||
"POSTGRES_URI": str(pg_engine.url),
|
||||
"SCHEDULE_MISFIRE_GRACE": 60,
|
||||
})(),
|
||||
)
|
||||
|
||||
|
||||
class TestDispatcherBasic:
|
||||
def test_due_recurring_enqueues_once_and_advances(
|
||||
self, pg_engine, patched_engine, stub_enqueue, monkeypatch,
|
||||
):
|
||||
_set_postgres_uri(monkeypatch, pg_engine)
|
||||
with pg_engine.begin() as conn:
|
||||
agent_id = _make_agent(conn)
|
||||
schedule = _create_schedule(
|
||||
pg_engine,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="* * * * *",
|
||||
next_run_at=_now() - timedelta(seconds=5),
|
||||
)
|
||||
counts = dispatch_due_runs()
|
||||
assert counts["enqueued"] == 1
|
||||
assert len(stub_enqueue) == 1
|
||||
with pg_engine.connect() as conn:
|
||||
row = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert row["next_run_at"] is not None
|
||||
|
||||
def test_once_dispatch_nulls_next_run_at_keeps_active(
|
||||
self, pg_engine, patched_engine, stub_enqueue, monkeypatch,
|
||||
):
|
||||
"""Once: dispatcher nulls next_run_at but leaves status='active' for the worker."""
|
||||
_set_postgres_uri(monkeypatch, pg_engine)
|
||||
with pg_engine.begin() as conn:
|
||||
agent_id = _make_agent(conn)
|
||||
schedule = _create_schedule(
|
||||
pg_engine,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="once",
|
||||
instruction="i", run_at=_now() + timedelta(seconds=1),
|
||||
next_run_at=_now() - timedelta(seconds=5),
|
||||
)
|
||||
counts = dispatch_due_runs()
|
||||
assert counts["enqueued"] == 1
|
||||
with pg_engine.connect() as conn:
|
||||
row = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert row["status"] == "active"
|
||||
assert row["next_run_at"] is None
|
||||
|
||||
|
||||
class TestDedupConstraint:
|
||||
def test_double_dispatch_only_one_run(
|
||||
self, pg_engine, patched_engine, stub_enqueue, monkeypatch,
|
||||
):
|
||||
_set_postgres_uri(monkeypatch, pg_engine)
|
||||
with pg_engine.begin() as conn:
|
||||
agent_id = _make_agent(conn)
|
||||
schedule = _create_schedule(
|
||||
pg_engine,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="*/5 * * * *",
|
||||
next_run_at=_now() - timedelta(seconds=2),
|
||||
)
|
||||
# Pre-claim simulates a racing dispatcher tick.
|
||||
with pg_engine.begin() as conn:
|
||||
row = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]),
|
||||
"u1",
|
||||
str(row["agent_id"]),
|
||||
row["next_run_at"],
|
||||
)
|
||||
counts = dispatch_due_runs()
|
||||
assert counts["enqueued"] == 0
|
||||
assert stub_enqueue == []
|
||||
|
||||
|
||||
class TestMisfireGrace:
|
||||
def test_stale_tick_recorded_skipped(
|
||||
self, pg_engine, patched_engine, stub_enqueue, monkeypatch,
|
||||
):
|
||||
_set_postgres_uri(monkeypatch, pg_engine)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_dispatcher.settings",
|
||||
type("S", (), {
|
||||
"POSTGRES_URI": str(pg_engine.url),
|
||||
"SCHEDULE_MISFIRE_GRACE": 30,
|
||||
})(),
|
||||
)
|
||||
with pg_engine.begin() as conn:
|
||||
agent_id = _make_agent(conn)
|
||||
schedule = _create_schedule(
|
||||
pg_engine,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="*/5 * * * *",
|
||||
next_run_at=_now() - timedelta(hours=2),
|
||||
)
|
||||
counts = dispatch_due_runs()
|
||||
assert counts["enqueued"] == 0
|
||||
assert counts["skipped"] >= 1
|
||||
with pg_engine.connect() as conn:
|
||||
runs = ScheduleRunsRepository(conn).list_runs(
|
||||
str(schedule["id"]), "u1",
|
||||
)
|
||||
assert any(r["error_type"] == "missed" for r in runs)
|
||||
|
||||
|
||||
class TestOverlap:
|
||||
def test_active_run_blocks_dispatch(
|
||||
self, pg_engine, patched_engine, stub_enqueue, monkeypatch,
|
||||
):
|
||||
_set_postgres_uri(monkeypatch, pg_engine)
|
||||
with pg_engine.begin() as conn:
|
||||
agent_id = _make_agent(conn)
|
||||
schedule = _create_schedule(
|
||||
pg_engine,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="*/5 * * * *",
|
||||
next_run_at=_now() - timedelta(seconds=2),
|
||||
)
|
||||
# Pre-create a running run with a different scheduled_for so overlap fires.
|
||||
with pg_engine.begin() as conn:
|
||||
row = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]),
|
||||
"u1",
|
||||
str(agent_id),
|
||||
_now() - timedelta(minutes=10),
|
||||
)
|
||||
ScheduleRunsRepository(conn).mark_running(row["id"], "t1")
|
||||
counts = dispatch_due_runs()
|
||||
assert counts["enqueued"] == 0
|
||||
|
||||
def test_once_overlap_clears_next_run_at(
|
||||
self, pg_engine, patched_engine, stub_enqueue, monkeypatch,
|
||||
):
|
||||
"""Once + overlap nulls next_run_at so the dispatcher stops re-picking."""
|
||||
_set_postgres_uri(monkeypatch, pg_engine)
|
||||
with pg_engine.begin() as conn:
|
||||
agent_id = _make_agent(conn)
|
||||
schedule = _create_schedule(
|
||||
pg_engine,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="once",
|
||||
instruction="i", run_at=_now() + timedelta(seconds=30),
|
||||
next_run_at=_now() - timedelta(seconds=5),
|
||||
)
|
||||
with pg_engine.begin() as conn:
|
||||
existing = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]),
|
||||
"u1",
|
||||
str(agent_id),
|
||||
_now() - timedelta(minutes=10),
|
||||
)
|
||||
ScheduleRunsRepository(conn).mark_running(existing["id"], "t-prev")
|
||||
dispatch_due_runs()
|
||||
with pg_engine.connect() as conn:
|
||||
row = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert row["status"] == "active"
|
||||
assert row["next_run_at"] is None
|
||||
|
||||
|
||||
class TestAgentlessSchedules:
|
||||
def test_dispatcher_claims_and_enqueues_agentless_once(
|
||||
self, pg_engine, patched_engine, stub_enqueue, monkeypatch,
|
||||
):
|
||||
"""``agent_id IS NULL`` rows are claimed like any other once-schedule."""
|
||||
_set_postgres_uri(monkeypatch, pg_engine)
|
||||
schedule = _create_schedule(
|
||||
pg_engine,
|
||||
user_id="u-agentless", agent_id=None, trigger_type="once",
|
||||
instruction="agentless ping",
|
||||
run_at=_now() + timedelta(seconds=30),
|
||||
next_run_at=_now() - timedelta(seconds=5),
|
||||
origin_conversation_id=None,
|
||||
created_via="chat",
|
||||
)
|
||||
counts = dispatch_due_runs()
|
||||
assert counts["enqueued"] == 1
|
||||
assert len(stub_enqueue) == 1
|
||||
with pg_engine.connect() as conn:
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
run_row = conn.execute(
|
||||
text(
|
||||
"SELECT * FROM schedule_runs "
|
||||
"WHERE schedule_id = CAST(:s AS uuid)"
|
||||
),
|
||||
{"s": str(schedule["id"])},
|
||||
).fetchone()
|
||||
# Once: dispatcher nulled next_run_at, schedule still active.
|
||||
assert sched["status"] == "active"
|
||||
assert sched["next_run_at"] is None
|
||||
# The pending run carries NULL agent_id (matches the parent).
|
||||
assert run_row._mapping["agent_id"] is None
|
||||
assert run_row._mapping["user_id"] == "u-agentless"
|
||||
|
||||
|
||||
class TestAgentlessRoundTrip:
|
||||
"""Agentless chat → tool → dispatcher → headless run → message appended."""
|
||||
|
||||
def test_agentless_dispatch_executes_and_appends_message(
|
||||
self, pg_engine, patched_engine, stub_enqueue, monkeypatch,
|
||||
):
|
||||
from unittest.mock import patch
|
||||
|
||||
from application.api.user.scheduler_worker import (
|
||||
execute_scheduled_run_body,
|
||||
)
|
||||
|
||||
_set_postgres_uri(monkeypatch, pg_engine)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_worker.get_engine",
|
||||
lambda: pg_engine,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_worker.settings",
|
||||
type("S", (), {
|
||||
"POSTGRES_URI": str(pg_engine.url),
|
||||
"SCHEDULE_AUTOPAUSE_FAILURES": 3,
|
||||
})(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_worker.publish_user_event",
|
||||
lambda *a, **k: "1-0",
|
||||
)
|
||||
|
||||
with pg_engine.begin() as conn:
|
||||
conv_id = conn.execute(
|
||||
text(
|
||||
"INSERT INTO conversations (user_id, name) "
|
||||
"VALUES ('u-e2e', 'agentless-chat') RETURNING id"
|
||||
)
|
||||
).fetchone()[0]
|
||||
schedule = _create_schedule(
|
||||
pg_engine,
|
||||
user_id="u-e2e", agent_id=None, trigger_type="once",
|
||||
instruction="ping later",
|
||||
run_at=_now() + timedelta(seconds=30),
|
||||
next_run_at=_now() - timedelta(seconds=5),
|
||||
origin_conversation_id=str(conv_id),
|
||||
created_via="chat",
|
||||
)
|
||||
|
||||
counts = dispatch_due_runs()
|
||||
assert counts["enqueued"] == 1
|
||||
run_id = stub_enqueue[0]
|
||||
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
return_value={
|
||||
"answer": "agentless e2e done",
|
||||
"tool_calls": [], "sources": [], "thought": "",
|
||||
"prompt_tokens": 1, "generated_tokens": 1,
|
||||
"denied": [], "error_type": None, "model_id": "fake",
|
||||
},
|
||||
):
|
||||
result = execute_scheduled_run_body(run_id, "celery-e2e")
|
||||
assert result["status"] == "success"
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
run = ScheduleRunsRepository(conn).get_internal(run_id)
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
messages = conn.execute(
|
||||
text(
|
||||
"SELECT * FROM conversation_messages "
|
||||
"WHERE conversation_id = CAST(:c AS uuid)"
|
||||
),
|
||||
{"c": str(conv_id)},
|
||||
).fetchall()
|
||||
assert run["status"] == "success"
|
||||
assert run["agent_id"] is None
|
||||
assert sched["status"] == "completed"
|
||||
assert sched["agent_id"] is None
|
||||
assert len(messages) == 1
|
||||
meta = messages[0]._mapping["message_metadata"]
|
||||
assert meta.get("scheduled") is True
|
||||
|
||||
|
||||
class TestOnceRoundTrip:
|
||||
"""End-to-end: chat-driven once-schedule executes and the schedule completes."""
|
||||
|
||||
def test_once_dispatch_executes_and_completes_schedule(
|
||||
self, pg_engine, patched_engine, stub_enqueue, monkeypatch,
|
||||
):
|
||||
from unittest.mock import patch
|
||||
|
||||
from application.api.user.scheduler_worker import (
|
||||
execute_scheduled_run_body,
|
||||
)
|
||||
|
||||
_set_postgres_uri(monkeypatch, pg_engine)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_worker.get_engine",
|
||||
lambda: pg_engine,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_worker.settings",
|
||||
type("S", (), {
|
||||
"POSTGRES_URI": str(pg_engine.url),
|
||||
"SCHEDULE_AUTOPAUSE_FAILURES": 3,
|
||||
})(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_worker.publish_user_event",
|
||||
lambda *a, **k: "1-0",
|
||||
)
|
||||
with pg_engine.begin() as conn:
|
||||
agent_id = _make_agent(conn)
|
||||
schedule = _create_schedule(
|
||||
pg_engine,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="once",
|
||||
instruction="follow up", run_at=_now() + timedelta(seconds=1),
|
||||
next_run_at=_now() - timedelta(seconds=5),
|
||||
)
|
||||
counts = dispatch_due_runs()
|
||||
assert counts["enqueued"] == 1
|
||||
assert len(stub_enqueue) == 1
|
||||
run_id = stub_enqueue[0]
|
||||
with pg_engine.connect() as conn:
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert sched["status"] == "active"
|
||||
assert sched["next_run_at"] is None
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
return_value={
|
||||
"answer": "done",
|
||||
"tool_calls": [], "sources": [], "thought": "",
|
||||
"prompt_tokens": 1, "generated_tokens": 1,
|
||||
"denied": [], "error_type": None, "model_id": "fake",
|
||||
},
|
||||
):
|
||||
result = execute_scheduled_run_body(run_id, "celery-c1")
|
||||
assert result["status"] == "success"
|
||||
with pg_engine.connect() as conn:
|
||||
run = ScheduleRunsRepository(conn).get_internal(run_id)
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert run["status"] == "success"
|
||||
assert run["output"] == "done"
|
||||
assert sched["status"] == "completed"
|
||||
assert sched["next_run_at"] is None
|
||||
@@ -0,0 +1,282 @@
|
||||
"""Tests for the scheduler reconciliation sweep + cleanup task."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
from application.storage.db.repositories.schedule_runs import (
|
||||
ScheduleRunsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.schedules import SchedulesRepository
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _make_pending_run(conn, *, user_id="u1"):
|
||||
agent_id = str(
|
||||
conn.execute(
|
||||
text(
|
||||
"INSERT INTO agents (user_id, name, status) "
|
||||
"VALUES (:u, 'a', 'draft') RETURNING id"
|
||||
),
|
||||
{"u": user_id},
|
||||
).fetchone()[0]
|
||||
)
|
||||
schedule = SchedulesRepository(conn).create(
|
||||
user_id=user_id, agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="*/5 * * * *",
|
||||
next_run_at=_now() + timedelta(minutes=5),
|
||||
)
|
||||
run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]), user_id, agent_id, _now(),
|
||||
)
|
||||
return schedule, run, agent_id
|
||||
|
||||
|
||||
def _make_once_pending_run(conn, *, user_id="u1"):
|
||||
"""Once-schedule + pending run variant of _make_pending_run."""
|
||||
agent_id = str(
|
||||
conn.execute(
|
||||
text(
|
||||
"INSERT INTO agents (user_id, name, status) "
|
||||
"VALUES (:u, 'a', 'draft') RETURNING id"
|
||||
),
|
||||
{"u": user_id},
|
||||
).fetchone()[0]
|
||||
)
|
||||
fire = _now() + timedelta(minutes=5)
|
||||
schedule = SchedulesRepository(conn).create(
|
||||
user_id=user_id, agent_id=agent_id, trigger_type="once",
|
||||
instruction="do once", run_at=fire,
|
||||
next_run_at=fire,
|
||||
)
|
||||
run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]), user_id, agent_id, fire,
|
||||
)
|
||||
return schedule, run, agent_id
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patched_engine(pg_engine, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.reconciliation.get_engine",
|
||||
lambda: pg_engine,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.reconciliation.settings",
|
||||
type("S", (), {
|
||||
"POSTGRES_URI": str(pg_engine.url),
|
||||
"SCHEDULE_RUN_TIMEOUT": 60,
|
||||
})(),
|
||||
)
|
||||
yield pg_engine
|
||||
|
||||
|
||||
class TestReconciler:
|
||||
def test_stuck_running_flipped_to_timeout(self, pg_engine, patched_engine):
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, run, _ = _make_pending_run(conn)
|
||||
ScheduleRunsRepository(conn).mark_running(run["id"], "t1")
|
||||
conn.execute(
|
||||
text(
|
||||
"UPDATE schedule_runs "
|
||||
"SET started_at = now() - interval '120 minutes' "
|
||||
"WHERE id = CAST(:i AS uuid)"
|
||||
),
|
||||
{"i": run["id"]},
|
||||
)
|
||||
|
||||
summary = run_reconciliation()
|
||||
assert summary["schedule_runs_failed"] >= 1
|
||||
with pg_engine.connect() as conn:
|
||||
row = ScheduleRunsRepository(conn).get_internal(run["id"])
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert row["status"] == "timeout"
|
||||
assert row["error_type"] == "timeout"
|
||||
assert sched["consecutive_failure_count"] == 1
|
||||
|
||||
def test_stuck_pending_flipped_to_failed(self, pg_engine, patched_engine):
|
||||
"""A pending run whose worker never started reconciles to 'failed'."""
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, run, _ = _make_pending_run(conn)
|
||||
conn.execute(
|
||||
text(
|
||||
"UPDATE schedule_runs "
|
||||
"SET created_at = now() - interval '120 minutes' "
|
||||
"WHERE id = CAST(:i AS uuid)"
|
||||
),
|
||||
{"i": run["id"]},
|
||||
)
|
||||
|
||||
summary = run_reconciliation()
|
||||
assert summary["schedule_runs_failed"] >= 1
|
||||
with pg_engine.connect() as conn:
|
||||
row = ScheduleRunsRepository(conn).get_internal(run["id"])
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert row["status"] == "failed"
|
||||
assert row["error_type"] == "internal"
|
||||
assert "worker_never_started" in (row["error"] or "")
|
||||
assert sched["consecutive_failure_count"] == 1
|
||||
|
||||
def test_once_schedule_with_stuck_running_run_marked_completed(
|
||||
self, pg_engine, patched_engine,
|
||||
):
|
||||
"""Once + stuck 'running' run -> parent flipped to 'completed'."""
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, run, _ = _make_once_pending_run(conn)
|
||||
ScheduleRunsRepository(conn).mark_running(run["id"], "t-once")
|
||||
conn.execute(
|
||||
text(
|
||||
"UPDATE schedule_runs "
|
||||
"SET started_at = now() - interval '120 minutes' "
|
||||
"WHERE id = CAST(:i AS uuid)"
|
||||
),
|
||||
{"i": run["id"]},
|
||||
)
|
||||
|
||||
run_reconciliation()
|
||||
with pg_engine.connect() as conn:
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
row = ScheduleRunsRepository(conn).get_internal(run["id"])
|
||||
assert row["status"] == "timeout"
|
||||
assert sched["status"] == "completed", (
|
||||
"stuck once-run must terminal-flip the parent schedule"
|
||||
)
|
||||
assert sched["next_run_at"] is None
|
||||
|
||||
def test_once_schedule_with_stuck_pending_run_marked_completed(
|
||||
self, pg_engine, patched_engine,
|
||||
):
|
||||
"""Once + stuck 'pending' run -> parent flipped to 'completed'."""
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, run, _ = _make_once_pending_run(conn)
|
||||
conn.execute(
|
||||
text(
|
||||
"UPDATE schedule_runs "
|
||||
"SET created_at = now() - interval '120 minutes' "
|
||||
"WHERE id = CAST(:i AS uuid)"
|
||||
),
|
||||
{"i": run["id"]},
|
||||
)
|
||||
|
||||
run_reconciliation()
|
||||
with pg_engine.connect() as conn:
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
row = ScheduleRunsRepository(conn).get_internal(run["id"])
|
||||
assert row["status"] == "failed"
|
||||
assert sched["status"] == "completed", (
|
||||
"stuck pending once-run must terminal-flip the parent schedule"
|
||||
)
|
||||
assert sched["next_run_at"] is None
|
||||
|
||||
def test_agentless_once_stuck_running_marked_completed(
|
||||
self, pg_engine, patched_engine,
|
||||
):
|
||||
"""Stuck-run terminal flip works for agentless once-schedules."""
|
||||
with pg_engine.begin() as conn:
|
||||
fire = _now() + timedelta(minutes=5)
|
||||
schedule = SchedulesRepository(conn).create(
|
||||
user_id="u-agentless", agent_id=None, trigger_type="once",
|
||||
instruction="agentless go", run_at=fire,
|
||||
next_run_at=fire,
|
||||
created_via="chat",
|
||||
)
|
||||
run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]), "u-agentless", None, fire,
|
||||
)
|
||||
ScheduleRunsRepository(conn).mark_running(run["id"], "t-stuck")
|
||||
conn.execute(
|
||||
text(
|
||||
"UPDATE schedule_runs "
|
||||
"SET started_at = now() - interval '120 minutes' "
|
||||
"WHERE id = CAST(:i AS uuid)"
|
||||
),
|
||||
{"i": run["id"]},
|
||||
)
|
||||
|
||||
run_reconciliation()
|
||||
with pg_engine.connect() as conn:
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
row = ScheduleRunsRepository(conn).get_internal(run["id"])
|
||||
assert row["status"] == "timeout"
|
||||
assert sched["status"] == "completed"
|
||||
assert sched["next_run_at"] is None
|
||||
|
||||
def test_recurring_schedule_with_stuck_run_stays_active(
|
||||
self, pg_engine, patched_engine,
|
||||
):
|
||||
"""Recurring keeps firing; only the run flips, not the parent."""
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, run, _ = _make_pending_run(conn)
|
||||
ScheduleRunsRepository(conn).mark_running(run["id"], "t-rec")
|
||||
conn.execute(
|
||||
text(
|
||||
"UPDATE schedule_runs "
|
||||
"SET started_at = now() - interval '120 minutes' "
|
||||
"WHERE id = CAST(:i AS uuid)"
|
||||
),
|
||||
{"i": run["id"]},
|
||||
)
|
||||
|
||||
run_reconciliation()
|
||||
with pg_engine.connect() as conn:
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert sched["status"] == "active"
|
||||
assert sched["consecutive_failure_count"] == 1
|
||||
|
||||
|
||||
class TestCleanup:
|
||||
def test_cleanup_schedule_runs_trims_old_rows(self, pg_engine, monkeypatch):
|
||||
from application.api.user.tasks import cleanup_schedule_runs as _task
|
||||
|
||||
monkeypatch.setattr(
|
||||
"application.storage.db.engine.get_engine",
|
||||
lambda: pg_engine,
|
||||
)
|
||||
|
||||
class S:
|
||||
POSTGRES_URI = str(pg_engine.url)
|
||||
SCHEDULE_RUN_OUTPUT_RETENTION_DAYS = 30
|
||||
monkeypatch.setattr("application.api.user.tasks.settings", S, raising=False)
|
||||
monkeypatch.setattr(
|
||||
"application.core.settings.settings.POSTGRES_URI",
|
||||
str(pg_engine.url),
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.settings.settings.SCHEDULE_RUN_OUTPUT_RETENTION_DAYS",
|
||||
30,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, _, _ = _make_pending_run(conn)
|
||||
for i in range(60):
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO schedule_runs (
|
||||
schedule_id, user_id, agent_id, status,
|
||||
scheduled_for, created_at
|
||||
) VALUES (
|
||||
CAST(:s AS uuid), 'u1',
|
||||
CAST(:a AS uuid), 'success',
|
||||
now() - interval '100 days' - (:i * interval '1 second'),
|
||||
now() - interval '100 days'
|
||||
)
|
||||
"""
|
||||
),
|
||||
{"s": str(schedule["id"]),
|
||||
"a": str(schedule["agent_id"]),
|
||||
"i": i},
|
||||
)
|
||||
|
||||
result = _task.run()
|
||||
assert isinstance(result.get("deleted"), int)
|
||||
assert result["deleted"] >= 1
|
||||
@@ -0,0 +1,428 @@
|
||||
"""Tests for execute_scheduled_run_body (mocked agent run)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.scheduler_worker import execute_scheduled_run_body
|
||||
from application.storage.db.repositories.schedule_runs import (
|
||||
ScheduleRunsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.schedules import SchedulesRepository
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _make_agent(conn, user_id: str = "u1") -> str:
|
||||
row = conn.execute(
|
||||
text(
|
||||
"INSERT INTO agents (user_id, name, status, default_model_id) "
|
||||
"VALUES (:u, 'a', 'draft', '') RETURNING id"
|
||||
),
|
||||
{"u": user_id},
|
||||
).fetchone()
|
||||
return str(row[0])
|
||||
|
||||
|
||||
def _make_pending_run(conn, *, user_id="u1"):
|
||||
agent_id = _make_agent(conn, user_id)
|
||||
schedule = SchedulesRepository(conn).create(
|
||||
user_id=user_id, agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="hello", cron="* * * * *",
|
||||
next_run_at=_now() + timedelta(minutes=5),
|
||||
)
|
||||
run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]),
|
||||
user_id,
|
||||
agent_id,
|
||||
_now(),
|
||||
)
|
||||
return schedule, run, agent_id
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patched_engine(pg_engine, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_worker.get_engine",
|
||||
lambda: pg_engine,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_worker.settings",
|
||||
type("S", (), {
|
||||
"POSTGRES_URI": str(pg_engine.url),
|
||||
"SCHEDULE_AUTOPAUSE_FAILURES": 2,
|
||||
})(),
|
||||
)
|
||||
yield pg_engine
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stub_events(monkeypatch):
|
||||
captured: list[tuple] = []
|
||||
|
||||
def _fake_publish(user_id, event_type, payload, *, scope=None):
|
||||
captured.append((event_type, payload, scope))
|
||||
return "1-0"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.scheduler_worker.publish_user_event",
|
||||
_fake_publish,
|
||||
)
|
||||
return captured
|
||||
|
||||
|
||||
class TestExecuteScheduledRunBody:
|
||||
def test_success_flow(self, pg_engine, patched_engine, stub_events):
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, run, _ = _make_pending_run(conn)
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
return_value={
|
||||
"answer": "all done",
|
||||
"tool_calls": [],
|
||||
"sources": [],
|
||||
"thought": "",
|
||||
"prompt_tokens": 10,
|
||||
"generated_tokens": 5,
|
||||
"denied": [],
|
||||
"error_type": None,
|
||||
"model_id": "fake-model",
|
||||
},
|
||||
):
|
||||
result = execute_scheduled_run_body(str(run["id"]), "celery-1")
|
||||
assert result["status"] == "success"
|
||||
with pg_engine.connect() as conn:
|
||||
row = ScheduleRunsRepository(conn).get_internal(str(run["id"]))
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert row["status"] == "success"
|
||||
assert row["output"] == "all done"
|
||||
assert row["prompt_tokens"] == 10
|
||||
assert sched["consecutive_failure_count"] == 0
|
||||
event_types = [e[0] for e in stub_events]
|
||||
assert "schedule.run.completed" in event_types
|
||||
|
||||
def test_agent_exception_marks_failed_and_bumps(
|
||||
self, pg_engine, patched_engine, stub_events,
|
||||
):
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, run, _ = _make_pending_run(conn)
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
side_effect=RuntimeError("boom"),
|
||||
):
|
||||
result = execute_scheduled_run_body(str(run["id"]), "celery-2")
|
||||
assert result["status"] == "failed"
|
||||
with pg_engine.connect() as conn:
|
||||
row = ScheduleRunsRepository(conn).get_internal(str(run["id"]))
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert row["status"] == "failed"
|
||||
assert row["error_type"] == "agent_error"
|
||||
assert sched["consecutive_failure_count"] == 1
|
||||
assert "schedule.run.failed" in {e[0] for e in stub_events}
|
||||
|
||||
def test_autopause_after_threshold(
|
||||
self, pg_engine, patched_engine, stub_events,
|
||||
):
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, run, agent_id = _make_pending_run(conn)
|
||||
SchedulesRepository(conn).bump_failure_count(str(schedule["id"]))
|
||||
another_run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]),
|
||||
"u1",
|
||||
agent_id,
|
||||
_now() + timedelta(seconds=1),
|
||||
)
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
side_effect=RuntimeError("boom"),
|
||||
):
|
||||
execute_scheduled_run_body(str(another_run["id"]), "celery-3")
|
||||
with pg_engine.connect() as conn:
|
||||
sched = SchedulesRepository(conn).get_internal(str(schedule["id"]))
|
||||
assert sched["status"] == "paused"
|
||||
assert "schedule.autopaused" in {e[0] for e in stub_events}
|
||||
|
||||
def test_denied_with_empty_output_marks_tool_not_allowed(
|
||||
self, pg_engine, patched_engine, stub_events,
|
||||
):
|
||||
with pg_engine.begin() as conn:
|
||||
schedule, run, _ = _make_pending_run(conn)
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
return_value={
|
||||
"answer": "",
|
||||
"tool_calls": [],
|
||||
"sources": [],
|
||||
"thought": "",
|
||||
"prompt_tokens": 1,
|
||||
"generated_tokens": 0,
|
||||
"denied": [{"tool_name": "telegram"}],
|
||||
"error_type": "tool_not_allowed",
|
||||
"model_id": "fake",
|
||||
},
|
||||
):
|
||||
execute_scheduled_run_body(str(run["id"]), "celery-4")
|
||||
with pg_engine.connect() as conn:
|
||||
row = ScheduleRunsRepository(conn).get_internal(str(run["id"]))
|
||||
assert row["status"] == "failed"
|
||||
assert row["error_type"] == "tool_not_allowed"
|
||||
|
||||
def test_one_time_loads_chat_history(
|
||||
self, pg_engine, patched_engine, stub_events,
|
||||
):
|
||||
with pg_engine.begin() as conn:
|
||||
agent_id = _make_agent(conn)
|
||||
schedule = SchedulesRepository(conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="once",
|
||||
instruction="follow up", run_at=_now() + timedelta(seconds=5),
|
||||
next_run_at=_now(),
|
||||
)
|
||||
conv_id = conn.execute(
|
||||
text(
|
||||
"INSERT INTO conversations (user_id, agent_id, name) "
|
||||
"VALUES ('u1', CAST(:a AS uuid), 'origin') RETURNING id"
|
||||
),
|
||||
{"a": agent_id},
|
||||
).fetchone()[0]
|
||||
SchedulesRepository(conn).update_internal(
|
||||
str(schedule["id"]),
|
||||
{"origin_conversation_id": str(conv_id)},
|
||||
)
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO conversation_messages
|
||||
(conversation_id, position, prompt, response, user_id)
|
||||
VALUES (CAST(:c AS uuid), 0, 'hello', 'hi', 'u1')
|
||||
"""
|
||||
),
|
||||
{"c": str(conv_id)},
|
||||
)
|
||||
run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]), "u1", agent_id, _now(),
|
||||
)
|
||||
captured: dict = {}
|
||||
def _fake_run(agent_config, query, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return {
|
||||
"answer": "follow-up answer",
|
||||
"tool_calls": [], "sources": [], "thought": "",
|
||||
"prompt_tokens": 1, "generated_tokens": 1,
|
||||
"denied": [], "error_type": None, "model_id": "fake",
|
||||
}
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless", _fake_run,
|
||||
):
|
||||
execute_scheduled_run_body(str(run["id"]), "celery-h")
|
||||
assert len(captured.get("chat_history", [])) == 1
|
||||
assert captured["chat_history"][0]["prompt"] == "hello"
|
||||
|
||||
def test_agentless_schedule_uses_system_defaults_and_appends(
|
||||
self, pg_engine, patched_engine, stub_events,
|
||||
):
|
||||
"""Agentless ``once`` schedule → ephemeral classic agent → message appended."""
|
||||
with pg_engine.begin() as conn:
|
||||
conv_id = conn.execute(
|
||||
text(
|
||||
"INSERT INTO conversations (user_id, name) "
|
||||
"VALUES ('u1', 'agentless-origin') RETURNING id"
|
||||
)
|
||||
).fetchone()[0]
|
||||
schedule = SchedulesRepository(conn).create(
|
||||
user_id="u1", agent_id=None, trigger_type="once",
|
||||
instruction="follow up agentless",
|
||||
run_at=_now() + timedelta(seconds=5),
|
||||
next_run_at=_now(),
|
||||
origin_conversation_id=str(conv_id),
|
||||
created_via="chat",
|
||||
)
|
||||
run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]), "u1", None, _now(),
|
||||
)
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run(agent_config, query, **kwargs):
|
||||
captured["agent_config"] = agent_config
|
||||
captured["kwargs"] = kwargs
|
||||
return {
|
||||
"answer": "agentless ran",
|
||||
"tool_calls": [], "sources": [], "thought": "",
|
||||
"prompt_tokens": 4, "generated_tokens": 6,
|
||||
"denied": [], "error_type": None, "model_id": "fake",
|
||||
}
|
||||
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
_fake_run,
|
||||
):
|
||||
result = execute_scheduled_run_body(str(run["id"]), "celery-agentless")
|
||||
assert result["status"] == "success"
|
||||
# Ephemeral classic config: no source, default retriever, no agent id.
|
||||
cfg = captured["agent_config"]
|
||||
assert cfg["id"] is None
|
||||
assert cfg["user_id"] == "u1"
|
||||
assert cfg["agent_type"] == "classic"
|
||||
assert cfg["retriever"] == "classic"
|
||||
assert cfg["prompt_id"] == "default"
|
||||
with pg_engine.connect() as conn:
|
||||
row = ScheduleRunsRepository(conn).get_internal(str(run["id"]))
|
||||
messages = conn.execute(
|
||||
text(
|
||||
"SELECT * FROM conversation_messages "
|
||||
"WHERE conversation_id = CAST(:c AS uuid)"
|
||||
),
|
||||
{"c": str(conv_id)},
|
||||
).fetchall()
|
||||
assert row["status"] == "success"
|
||||
assert row["output"] == "agentless ran"
|
||||
assert row["conversation_id"] is not None
|
||||
assert len(messages) == 1
|
||||
# The published event payload tolerates a NULL agent_id.
|
||||
appended_events = [e for e in stub_events if e[0] == "schedule.message.appended"]
|
||||
assert appended_events
|
||||
|
||||
def test_agentless_ephemeral_config_omits_tools_snapshot(
|
||||
self, pg_engine, patched_engine, stub_events,
|
||||
):
|
||||
"""Dead ``tools`` snapshot dropped — toolset is rebuilt at fire time."""
|
||||
with pg_engine.begin() as conn:
|
||||
conv_id = conn.execute(
|
||||
text(
|
||||
"INSERT INTO conversations (user_id, name) "
|
||||
"VALUES ('u1', 'no-tools-snap') RETURNING id"
|
||||
)
|
||||
).fetchone()[0]
|
||||
schedule = SchedulesRepository(conn).create(
|
||||
user_id="u1", agent_id=None, trigger_type="once",
|
||||
instruction="x", run_at=_now() + timedelta(seconds=5),
|
||||
next_run_at=_now(),
|
||||
origin_conversation_id=str(conv_id),
|
||||
created_via="chat",
|
||||
)
|
||||
run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]), "u1", None, _now(),
|
||||
)
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run(agent_config, query, **kwargs):
|
||||
captured["agent_config"] = agent_config
|
||||
return {
|
||||
"answer": "ok", "tool_calls": [], "sources": [], "thought": "",
|
||||
"prompt_tokens": 1, "generated_tokens": 1,
|
||||
"denied": [], "error_type": None, "model_id": "fake",
|
||||
}
|
||||
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
_fake_run,
|
||||
):
|
||||
execute_scheduled_run_body(str(run["id"]), "celery-no-snap")
|
||||
cfg = captured["agent_config"]
|
||||
# ``tools`` MUST NOT be in the ephemeral shape — the runtime
|
||||
# toolset is rebuilt by ``ToolExecutor`` (which honours headless
|
||||
# filtering for chat-only tools like ``scheduler``).
|
||||
assert "tools" not in cfg
|
||||
|
||||
def test_agentless_token_usage_row_has_null_agent_id(
|
||||
self, pg_engine, patched_engine, stub_events,
|
||||
):
|
||||
"""token_usage row for an agentless run carries ``agent_id IS NULL``."""
|
||||
with pg_engine.begin() as conn:
|
||||
conv_id = conn.execute(
|
||||
text(
|
||||
"INSERT INTO conversations (user_id, name) "
|
||||
"VALUES ('u1', 'agentless-tu') RETURNING id"
|
||||
)
|
||||
).fetchone()[0]
|
||||
schedule = SchedulesRepository(conn).create(
|
||||
user_id="u1", agent_id=None, trigger_type="once",
|
||||
instruction="tu", run_at=_now() + timedelta(seconds=5),
|
||||
next_run_at=_now(),
|
||||
origin_conversation_id=str(conv_id),
|
||||
created_via="chat",
|
||||
)
|
||||
run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]), "u1", None, _now(),
|
||||
)
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
return_value={
|
||||
"answer": "yes",
|
||||
"tool_calls": [], "sources": [], "thought": "",
|
||||
"prompt_tokens": 11, "generated_tokens": 7,
|
||||
"denied": [], "error_type": None, "model_id": "fake",
|
||||
},
|
||||
):
|
||||
execute_scheduled_run_body(str(run["id"]), "celery-tu")
|
||||
with pg_engine.connect() as conn:
|
||||
tu_row = conn.execute(
|
||||
text(
|
||||
"SELECT * FROM token_usage "
|
||||
"WHERE request_id = :r"
|
||||
),
|
||||
{"r": str(run["id"])},
|
||||
).fetchone()
|
||||
assert tu_row is not None
|
||||
assert tu_row._mapping["agent_id"] is None
|
||||
assert tu_row._mapping["source"] == "schedule"
|
||||
|
||||
def test_one_time_appends_message(
|
||||
self, pg_engine, patched_engine, stub_events,
|
||||
):
|
||||
with pg_engine.begin() as conn:
|
||||
agent_id = _make_agent(conn)
|
||||
schedule = SchedulesRepository(conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="once",
|
||||
instruction="hello", run_at=_now() + timedelta(seconds=5),
|
||||
next_run_at=_now(),
|
||||
)
|
||||
conv_id = conn.execute(
|
||||
text(
|
||||
"INSERT INTO conversations (user_id, agent_id, name) "
|
||||
"VALUES ('u1', CAST(:a AS uuid), 'origin') RETURNING id"
|
||||
),
|
||||
{"a": agent_id},
|
||||
).fetchone()[0]
|
||||
SchedulesRepository(conn).update_internal(
|
||||
str(schedule["id"]),
|
||||
{"origin_conversation_id": str(conv_id)},
|
||||
)
|
||||
run = ScheduleRunsRepository(conn).record_pending(
|
||||
str(schedule["id"]), "u1", agent_id, _now(),
|
||||
)
|
||||
with patch(
|
||||
"application.api.user.scheduler_worker.run_agent_headless",
|
||||
return_value={
|
||||
"answer": "scheduled answer",
|
||||
"tool_calls": [],
|
||||
"sources": [],
|
||||
"thought": "",
|
||||
"prompt_tokens": 2,
|
||||
"generated_tokens": 3,
|
||||
"denied": [],
|
||||
"error_type": None,
|
||||
"model_id": "fake",
|
||||
},
|
||||
):
|
||||
execute_scheduled_run_body(str(run["id"]), "celery-5")
|
||||
with pg_engine.connect() as conn:
|
||||
row = ScheduleRunsRepository(conn).get_internal(str(run["id"]))
|
||||
messages = conn.execute(
|
||||
text(
|
||||
"SELECT * FROM conversation_messages "
|
||||
"WHERE conversation_id = CAST(:c AS uuid)"
|
||||
),
|
||||
{"c": str(conv_id)},
|
||||
).fetchall()
|
||||
assert row["conversation_id"] is not None
|
||||
assert row["message_id"] is not None
|
||||
assert len(messages) == 1
|
||||
meta = messages[0]._mapping["message_metadata"]
|
||||
assert meta.get("scheduled") is True
|
||||
assert "schedule.message.appended" in {e[0] for e in stub_events}
|
||||
@@ -0,0 +1,508 @@
|
||||
"""Tests for the schedules REST API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.storage.db.repositories.schedules import SchedulesRepository
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app():
|
||||
return Flask(__name__)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.user.schedules.routes.db_session", _yield,
|
||||
), patch(
|
||||
"application.api.user.schedules.routes.db_readonly", _yield,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
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])
|
||||
|
||||
|
||||
class TestCreateRecurring:
|
||||
def test_unauthorized(self, app):
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/agents/x/schedules", method="POST", json={},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = None
|
||||
resp = AgentSchedules().post("x")
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_agent_not_found(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
"/api/agents/00000000-0000-0000-0000-000000000000/schedules",
|
||||
method="POST", json={"instruction": "x", "cron": "* * * * *"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = AgentSchedules().post(
|
||||
"00000000-0000-0000-0000-000000000000",
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_invalid_cron(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/agents/{agent_id}/schedules",
|
||||
method="POST",
|
||||
json={"instruction": "x", "cron": "not a cron"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = AgentSchedules().post(agent_id)
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_create_success(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/agents/{agent_id}/schedules",
|
||||
method="POST",
|
||||
json={
|
||||
"instruction": "weekly digest",
|
||||
"cron": "0 9 * * 1",
|
||||
"timezone": "Europe/Warsaw",
|
||||
"tool_allowlist": [],
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = AgentSchedules().post(agent_id)
|
||||
assert resp.status_code == 201
|
||||
body = resp.get_json()
|
||||
assert body["schedule"]["cron"] == "0 9 * * 1"
|
||||
assert body["schedule"]["timezone"] == "Europe/Warsaw"
|
||||
|
||||
|
||||
class TestCreateOnce:
|
||||
def test_creates_once_with_run_at(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
run_at = (_now() + timedelta(hours=2)).isoformat().replace(
|
||||
"+00:00", "Z",
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/agents/{agent_id}/schedules",
|
||||
method="POST",
|
||||
json={
|
||||
"instruction": "remind me",
|
||||
"trigger_type": "once",
|
||||
"run_at": run_at,
|
||||
"timezone": "UTC",
|
||||
"tool_allowlist": [],
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = AgentSchedules().post(agent_id)
|
||||
assert resp.status_code == 201
|
||||
body = resp.get_json()
|
||||
assert body["schedule"]["trigger_type"] == "once"
|
||||
assert body["schedule"]["run_at"] is not None
|
||||
|
||||
def test_once_requires_run_at(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/agents/{agent_id}/schedules",
|
||||
method="POST",
|
||||
json={
|
||||
"instruction": "remind me",
|
||||
"trigger_type": "once",
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = AgentSchedules().post(agent_id)
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_once_rejects_past_run_at(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
past = (_now() - timedelta(hours=1)).isoformat().replace(
|
||||
"+00:00", "Z",
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/agents/{agent_id}/schedules",
|
||||
method="POST",
|
||||
json={
|
||||
"instruction": "x",
|
||||
"trigger_type": "once",
|
||||
"run_at": past,
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = AgentSchedules().post(agent_id)
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_recurring_default_when_trigger_type_omitted(self, app, pg_conn):
|
||||
"""Backwards compat: a payload with cron but no trigger_type still works."""
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/agents/{agent_id}/schedules",
|
||||
method="POST",
|
||||
json={
|
||||
"instruction": "hourly",
|
||||
"cron": "0 * * * *",
|
||||
"timezone": "UTC",
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = AgentSchedules().post(agent_id)
|
||||
assert resp.status_code == 201
|
||||
assert resp.get_json()["schedule"]["trigger_type"] == "recurring"
|
||||
|
||||
|
||||
class TestListForAgent:
|
||||
def test_list(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
SchedulesRepository(pg_conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="* * * * *",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/agents/{agent_id}/schedules", method="GET",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = AgentSchedules().get(agent_id)
|
||||
assert resp.status_code == 200
|
||||
assert len(resp.get_json()["schedules"]) == 1
|
||||
|
||||
|
||||
class TestGetEditPatchDelete:
|
||||
def _make(self, conn, **kwargs):
|
||||
return SchedulesRepository(conn).create(**kwargs)
|
||||
|
||||
def test_get_owner_scoped(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleResource
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = self._make(
|
||||
pg_conn,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="* * * * *",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}", method="GET",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u2"}
|
||||
resp = ScheduleResource().get(str(s["id"]))
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_pause_then_resume(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleResource
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = self._make(
|
||||
pg_conn,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="* * * * *",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}",
|
||||
method="PATCH",
|
||||
json={"action": "pause"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = ScheduleResource().patch(str(s["id"]))
|
||||
assert resp.status_code == 200
|
||||
assert resp.get_json()["schedule"]["status"] == "paused"
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}",
|
||||
method="PATCH",
|
||||
json={"action": "resume"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = ScheduleResource().patch(str(s["id"]))
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()
|
||||
assert body["schedule"]["status"] == "active"
|
||||
assert body["schedule"]["next_run_at"] is not None
|
||||
|
||||
def test_delete_owner_scoped(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleResource
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = self._make(
|
||||
pg_conn,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="* * * * *",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}", method="DELETE",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u2"}
|
||||
resp = ScheduleResource().delete(str(s["id"]))
|
||||
assert resp.status_code == 404
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}", method="DELETE",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = ScheduleResource().delete(str(s["id"]))
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_put_invalid_cron(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleResource
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = self._make(
|
||||
pg_conn,
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="* * * * *",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}",
|
||||
method="PUT", json={"cron": "bad"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = ScheduleResource().put(str(s["id"]))
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
class TestRunNow:
|
||||
def test_runs_returns_202(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleRunNow
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = SchedulesRepository(pg_conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="* * * * *",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.tasks.execute_scheduled_run",
|
||||
type("T", (), {"apply_async": staticmethod(lambda **k: None)}),
|
||||
), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}/run", method="POST",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = ScheduleRunNow().post(str(s["id"]))
|
||||
assert resp.status_code == 202
|
||||
|
||||
def test_second_run_blocked_by_active(self, app, pg_conn):
|
||||
"""Run-Now serializes via FOR UPDATE + has_active_run; second 409s."""
|
||||
from application.api.user.schedules.routes import ScheduleRunNow
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = SchedulesRepository(pg_conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="0 9 * * 1",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.tasks.execute_scheduled_run",
|
||||
type("T", (), {"apply_async": staticmethod(lambda **k: None)}),
|
||||
), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}/run", method="POST",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
first = ScheduleRunNow().post(str(s["id"]))
|
||||
assert first.status_code == 202
|
||||
second = ScheduleRunNow().post(str(s["id"]))
|
||||
assert second.status_code == 409
|
||||
|
||||
|
||||
class TestMinInterval:
|
||||
def test_create_rejects_below_min_interval(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import AgentSchedules
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/agents/{agent_id}/schedules",
|
||||
method="POST",
|
||||
json={"instruction": "x", "cron": "* * * * *"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = AgentSchedules().post(agent_id)
|
||||
assert resp.status_code == 400
|
||||
assert "minimum interval" in resp.get_json()["message"]
|
||||
|
||||
def test_put_rejects_below_min_interval(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleResource
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = SchedulesRepository(pg_conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="0 9 * * 1",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}", method="PUT",
|
||||
json={"cron": "*/5 * * * *"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = ScheduleResource().put(str(s["id"]))
|
||||
assert resp.status_code == 400
|
||||
assert "minimum interval" in resp.get_json()["message"]
|
||||
|
||||
|
||||
class TestResumeOnceStale:
|
||||
def test_stale_run_at_returns_clear_409(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleResource
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = SchedulesRepository(pg_conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="once",
|
||||
instruction="i", run_at=_now() + timedelta(hours=1),
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
status="paused",
|
||||
)
|
||||
pg_conn.execute(
|
||||
text(
|
||||
"UPDATE schedules SET run_at = now() - interval '1 day' "
|
||||
"WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": str(s["id"])},
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}", method="PATCH",
|
||||
json={"action": "resume"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = ScheduleResource().patch(str(s["id"]))
|
||||
assert resp.status_code == 409
|
||||
assert "elapsed" in resp.get_json()["message"]
|
||||
|
||||
def test_resume_accepts_new_run_at(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleResource
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = SchedulesRepository(pg_conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="once",
|
||||
instruction="i", run_at=_now() + timedelta(hours=1),
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
status="paused",
|
||||
)
|
||||
pg_conn.execute(
|
||||
text(
|
||||
"UPDATE schedules SET run_at = now() - interval '1 day' "
|
||||
"WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": str(s["id"])},
|
||||
)
|
||||
new_run_at = (_now() + timedelta(hours=3)).isoformat()
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}", method="PATCH",
|
||||
json={"action": "resume", "run_at": new_run_at},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = ScheduleResource().patch(str(s["id"]))
|
||||
assert resp.status_code == 200
|
||||
body = resp.get_json()["schedule"]
|
||||
assert body["status"] == "active"
|
||||
|
||||
|
||||
class TestRunList:
|
||||
def test_list_owner_scoped(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleRunList
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = SchedulesRepository(pg_conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="* * * * *",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}/runs", method="GET",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u2"}
|
||||
resp = ScheduleRunList().get(str(s["id"]))
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestUnexpectedExceptionMasked:
|
||||
"""Unexpected exceptions log full trace + return generic 500 (no leak)."""
|
||||
|
||||
def test_unexpected_repo_error_returns_generic_500(self, app, pg_conn):
|
||||
from application.api.user.schedules.routes import ScheduleResource
|
||||
|
||||
agent_id = _make_agent(pg_conn)
|
||||
s = SchedulesRepository(pg_conn).create(
|
||||
user_id="u1", agent_id=agent_id, trigger_type="recurring",
|
||||
instruction="i", cron="* * * * *",
|
||||
next_run_at=_now() + timedelta(hours=1),
|
||||
)
|
||||
|
||||
# Mock the repo's ``get`` to raise an unexpected RuntimeError carrying
|
||||
# internal-sounding detail; the response must NOT echo that string.
|
||||
def _boom(*args, **kwargs):
|
||||
raise RuntimeError("internal detail: secret connection string")
|
||||
|
||||
with _patch_db(pg_conn), patch.object(
|
||||
SchedulesRepository, "get", side_effect=_boom,
|
||||
), app.test_request_context(
|
||||
f"/api/schedules/{s['id']}", method="GET",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u1"}
|
||||
resp = ScheduleResource().get(str(s["id"]))
|
||||
assert resp.status_code == 500
|
||||
body = resp.get_json()
|
||||
assert body["success"] is False
|
||||
assert body["message"] == "internal error"
|
||||
# Defensive: the internal-sounding detail must not leak through.
|
||||
assert "secret connection string" not in resp.get_data(as_text=True)
|
||||
@@ -222,7 +222,7 @@ class TestSetupPeriodicTasks:
|
||||
|
||||
setup_periodic_tasks(sender)
|
||||
|
||||
assert sender.add_periodic_task.call_count == 8
|
||||
assert sender.add_periodic_task.call_count == 11
|
||||
|
||||
calls = sender.add_periodic_task.call_args_list
|
||||
|
||||
@@ -246,6 +246,14 @@ class TestSetupPeriodicTasks:
|
||||
# message_events retention sweep (24h)
|
||||
assert calls[7][0][0] == timedelta(hours=24)
|
||||
assert calls[7][1].get("name") == "cleanup-message-events"
|
||||
# orphan memories sweep (24h)
|
||||
assert calls[8][0][0] == timedelta(hours=24)
|
||||
assert calls[8][1].get("name") == "cleanup-orphan-memories"
|
||||
# scheduler dispatcher
|
||||
assert calls[9][1].get("name") == "dispatch-scheduled-runs"
|
||||
# schedule runs cleanup (24h)
|
||||
assert calls[10][0][0] == timedelta(hours=24)
|
||||
assert calls[10][1].get("name") == "cleanup-schedule-runs"
|
||||
|
||||
|
||||
class TestMcpOauthTask:
|
||||
@@ -296,6 +304,7 @@ class TestDurableTaskRetryPolicy:
|
||||
"cleanup_pending_tool_state",
|
||||
"reconciliation_task",
|
||||
"version_check_task",
|
||||
"cleanup_orphan_memories",
|
||||
],
|
||||
)
|
||||
def test_short_periodic_tasks_have_no_retry_config(self, task_name):
|
||||
@@ -515,6 +524,72 @@ class TestCleanupMessageEventsTask:
|
||||
assert [r["sequence_no"] for r in rows] == [1]
|
||||
|
||||
|
||||
class TestCleanupOrphanMemoriesTask:
|
||||
"""Sweeps orphan memories from the FK-to-trigger orphan window."""
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_skips_when_postgres_uri_missing(self, monkeypatch):
|
||||
from application.api.user.tasks import cleanup_orphan_memories
|
||||
from application.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "POSTGRES_URI", None, raising=False)
|
||||
|
||||
result = cleanup_orphan_memories.run()
|
||||
assert result == {"deleted": 0, "skipped": "POSTGRES_URI not set"}
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_deletes_orphan_keeps_synthetic_and_live(
|
||||
self, pg_conn, monkeypatch
|
||||
):
|
||||
import uuid
|
||||
|
||||
from sqlalchemy import text as _text
|
||||
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tasks import cleanup_orphan_memories
|
||||
from application.core.settings import settings
|
||||
from application.storage.db.repositories.memories import (
|
||||
MemoriesRepository,
|
||||
)
|
||||
|
||||
repo = MemoriesRepository(pg_conn)
|
||||
synthetic_id = default_tool_id("memory")
|
||||
live_id = str(
|
||||
pg_conn.execute(
|
||||
_text(
|
||||
"INSERT INTO user_tools (user_id, name) "
|
||||
"VALUES ('u-task-mem', 'memory') RETURNING id"
|
||||
)
|
||||
).scalar()
|
||||
)
|
||||
orphan_id = str(uuid.uuid4())
|
||||
repo.upsert("u-task-mem", synthetic_id, "/syn.txt", "keep")
|
||||
repo.upsert("u-task-mem", live_id, "/live.txt", "keep")
|
||||
repo.upsert("u-task-mem", orphan_id, "/orphan.txt", "drop")
|
||||
|
||||
monkeypatch.setattr(
|
||||
settings, "POSTGRES_URI", "postgresql://stub", raising=False
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _fake_begin():
|
||||
yield pg_conn
|
||||
|
||||
fake_engine = MagicMock()
|
||||
fake_engine.begin = _fake_begin
|
||||
|
||||
with patch(
|
||||
"application.storage.db.engine.get_engine",
|
||||
return_value=fake_engine,
|
||||
):
|
||||
result = cleanup_orphan_memories.run()
|
||||
|
||||
assert result == {"deleted": 1}
|
||||
assert repo.get_by_path("u-task-mem", synthetic_id, "/syn.txt")
|
||||
assert repo.get_by_path("u-task-mem", live_id, "/live.txt")
|
||||
assert repo.get_by_path("u-task-mem", orphan_id, "/orphan.txt") is None
|
||||
|
||||
|
||||
class TestIngestIdempotency:
|
||||
"""Same short-circuit applies to the ingest task path."""
|
||||
|
||||
|
||||
@@ -1066,6 +1066,11 @@ def _seed_tool(pg_conn, user="u-tools", name="read_webpage", config=None):
|
||||
|
||||
class TestGetToolsHappy:
|
||||
def test_returns_user_tools(self, app, pg_conn):
|
||||
from application.agents.default_tools import (
|
||||
BUILTIN_AGENT_TOOLS,
|
||||
loaded_builtin_agent_tools,
|
||||
loaded_default_tools,
|
||||
)
|
||||
from application.api.user.tools.routes import GetTools
|
||||
|
||||
user = "u-get-tools"
|
||||
@@ -1079,7 +1084,26 @@ class TestGetToolsHappy:
|
||||
response = GetTools().get()
|
||||
assert response.status_code == 200
|
||||
assert response.json["success"] is True
|
||||
assert len(response.json["tools"]) == 1
|
||||
tools = response.json["tools"]
|
||||
# Response shape: 1 explicit + every default + builtins not already
|
||||
# surfaced as a default (dual-registered ``scheduler`` is dedup'd).
|
||||
defaults_count = len(loaded_default_tools())
|
||||
builtins_count = len(loaded_builtin_agent_tools())
|
||||
dual = sum(
|
||||
1 for name in loaded_default_tools() if name in BUILTIN_AGENT_TOOLS
|
||||
)
|
||||
assert len(tools) == 1 + defaults_count + (builtins_count - dual)
|
||||
explicit = [
|
||||
t for t in tools
|
||||
if not t.get("default") and not t.get("builtin")
|
||||
]
|
||||
defaults = [t for t in tools if t.get("default")]
|
||||
builtins = [t for t in tools if t.get("builtin")]
|
||||
assert len(explicit) == 1
|
||||
assert len(defaults) == defaults_count
|
||||
# Dual-registered tools (scheduler) appear once with both flags;
|
||||
# ``builtins`` here counts them via ``builtin=True``.
|
||||
assert len(builtins) == (builtins_count - dual) + dual
|
||||
|
||||
def test_db_error_returns_400(self, app):
|
||||
from application.api.user.tools.routes import GetTools
|
||||
@@ -1379,8 +1403,291 @@ class TestGetArtifactHappy:
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Default chat tools — synthetic-id branching in the tool endpoints
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestDefaultToolsRoutes:
|
||||
def test_get_tools_flags_defaults(self, app, pg_conn):
|
||||
from application.agents.default_tools import (
|
||||
default_tool_id,
|
||||
loaded_default_tools,
|
||||
)
|
||||
from application.api.user.tools.routes import GetTools
|
||||
|
||||
user = "u-def-get"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context("/api/get_tools"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = GetTools().get()
|
||||
assert response.status_code == 200
|
||||
defaults = [t for t in response.json["tools"] if t.get("default")]
|
||||
names = {t["name"] for t in defaults}
|
||||
assert names == set(loaded_default_tools())
|
||||
for tool in defaults:
|
||||
assert tool["id"] == default_tool_id(tool["name"])
|
||||
assert tool["status"] is True
|
||||
|
||||
def test_get_tools_surfaces_scheduler_with_both_flags(self, app, pg_conn):
|
||||
"""Dual-registered scheduler appears once with default+builtin flags."""
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import GetTools
|
||||
|
||||
user = "u-sched-dual"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context("/api/get_tools"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = GetTools().get()
|
||||
assert response.status_code == 200
|
||||
scheduler_id = default_tool_id("scheduler")
|
||||
scheduler_rows = [
|
||||
t for t in response.json["tools"] if t["id"] == scheduler_id
|
||||
]
|
||||
assert len(scheduler_rows) == 1 # dedup at the routes layer
|
||||
row = scheduler_rows[0]
|
||||
assert row["default"] is True
|
||||
assert row["builtin"] is True
|
||||
assert row["name"] == "scheduler"
|
||||
|
||||
def test_get_tools_status_reflects_opt_out(self, app, pg_conn):
|
||||
from application.api.user.tools.routes import GetTools
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
|
||||
user = "u-def-optout"
|
||||
UsersRepository(pg_conn).set_default_tool_enabled(
|
||||
user, "read_webpage", False
|
||||
)
|
||||
with _patch_tools_db(pg_conn), app.test_request_context("/api/get_tools"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = GetTools().get()
|
||||
by_name = {
|
||||
t["name"]: t
|
||||
for t in response.json["tools"]
|
||||
if t.get("default")
|
||||
}
|
||||
assert by_name["read_webpage"]["status"] is False
|
||||
assert by_name["memory"]["status"] is True
|
||||
|
||||
def test_update_tool_status_toggles_default_off(self, app, pg_conn):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateToolStatus
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
|
||||
user = "u-def-toggle"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool_status",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("memory"), "status": False},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateToolStatus().post()
|
||||
assert response.status_code == 200
|
||||
user_doc = UsersRepository(pg_conn).get(user)
|
||||
assert user_doc["tool_preferences"]["disabled_default_tools"] == ["memory"]
|
||||
|
||||
def test_update_tool_status_toggles_default_back_on(self, app, pg_conn):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateToolStatus
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
|
||||
user = "u-def-on"
|
||||
UsersRepository(pg_conn).set_default_tool_enabled(user, "memory", False)
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool_status",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("memory"), "status": True},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateToolStatus().post()
|
||||
assert response.status_code == 200
|
||||
user_doc = UsersRepository(pg_conn).get(user)
|
||||
assert user_doc["tool_preferences"]["disabled_default_tools"] == []
|
||||
|
||||
def test_update_tool_toggles_default_via_status(self, app, pg_conn):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateTool
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
|
||||
user = "u-def-updtool"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("read_webpage"), "status": False},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateTool().post()
|
||||
assert response.status_code == 200
|
||||
user_doc = UsersRepository(pg_conn).get(user)
|
||||
assert user_doc["tool_preferences"]["disabled_default_tools"] == [
|
||||
"read_webpage"
|
||||
]
|
||||
|
||||
def test_delete_tool_rejects_default(self, app, pg_conn):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import DeleteTool
|
||||
|
||||
user = "u-def-del"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/delete_tool",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("memory")},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = DeleteTool().post()
|
||||
assert response.status_code == 400
|
||||
assert response.json["success"] is False
|
||||
|
||||
def test_update_tool_default_without_status_is_rejected(
|
||||
self, app, pg_conn
|
||||
):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateTool
|
||||
|
||||
user = "u-def-noedit"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("memory"), "displayName": "Renamed"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateTool().post()
|
||||
assert response.status_code == 400
|
||||
assert response.json["success"] is False
|
||||
assert "not editable" in response.json["message"]
|
||||
|
||||
def test_update_tool_config_rejects_default(self, app, pg_conn):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateToolConfig
|
||||
|
||||
user = "u-def-cfg"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool_config",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("memory"), "config": {"x": 1}},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateToolConfig().post()
|
||||
assert response.status_code == 400
|
||||
assert response.json["success"] is False
|
||||
assert "config-free" in response.json["message"]
|
||||
|
||||
def test_update_tool_actions_rejects_default(self, app, pg_conn):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateToolActions
|
||||
|
||||
user = "u-def-act"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool_actions",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("memory"), "actions": []},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateToolActions().post()
|
||||
assert response.status_code == 400
|
||||
assert response.json["success"] is False
|
||||
assert "not editable" in response.json["message"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dual-registered tools (scheduler) — toggle MUST hit the default-tool path,
|
||||
# not the builtin "not editable" rejection. Regression for the iter-6 issue
|
||||
# where ``is_builtin_agent_tool_id`` was checked first, silently dropping the
|
||||
# write on a dual-registered uuid5.
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestDualRegisteredToggle:
|
||||
def test_update_tool_status_off_writes_disabled_default(self, app, pg_conn):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateToolStatus
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
|
||||
user = "u-sched-off"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool_status",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("scheduler"), "status": False},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateToolStatus().post()
|
||||
assert response.status_code == 200
|
||||
assert response.json["success"] is True
|
||||
user_doc = UsersRepository(pg_conn).get(user)
|
||||
assert "scheduler" in (
|
||||
user_doc["tool_preferences"]["disabled_default_tools"]
|
||||
)
|
||||
|
||||
def test_update_tool_status_on_removes_disabled_default(self, app, pg_conn):
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateToolStatus
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
|
||||
user = "u-sched-on"
|
||||
UsersRepository(pg_conn).set_default_tool_enabled(
|
||||
user, "scheduler", False,
|
||||
)
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool_status",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("scheduler"), "status": True},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateToolStatus().post()
|
||||
assert response.status_code == 200
|
||||
assert response.json["success"] is True
|
||||
user_doc = UsersRepository(pg_conn).get(user)
|
||||
assert "scheduler" not in (
|
||||
user_doc["tool_preferences"]["disabled_default_tools"]
|
||||
)
|
||||
|
||||
def test_update_tool_status_round_trip(self, app, pg_conn):
|
||||
"""Off → on returns to the empty-list baseline."""
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateToolStatus
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
|
||||
user = "u-sched-rt"
|
||||
scheduler_id = default_tool_id("scheduler")
|
||||
for status in (False, True):
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool_status",
|
||||
method="POST",
|
||||
json={"id": scheduler_id, "status": status},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateToolStatus().post()
|
||||
assert response.status_code == 200
|
||||
user_doc = UsersRepository(pg_conn).get(user)
|
||||
assert user_doc["tool_preferences"]["disabled_default_tools"] == []
|
||||
|
||||
def test_update_tool_with_status_writes_disabled_default(self, app, pg_conn):
|
||||
"""The /api/update_tool route also honours the default branch first."""
|
||||
from application.agents.default_tools import default_tool_id
|
||||
from application.api.user.tools.routes import UpdateTool
|
||||
from application.storage.db.repositories.users import UsersRepository
|
||||
|
||||
user = "u-sched-upd"
|
||||
with _patch_tools_db(pg_conn), app.test_request_context(
|
||||
"/api/update_tool",
|
||||
method="POST",
|
||||
json={"id": default_tool_id("scheduler"), "status": False},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateTool().post()
|
||||
assert response.status_code == 200
|
||||
assert response.json["success"] is True
|
||||
user_doc = UsersRepository(pg_conn).get(user)
|
||||
assert "scheduler" in (
|
||||
user_doc["tool_preferences"]["disabled_default_tools"]
|
||||
)
|
||||
|
||||
|
||||
Reference in new issue
Block a user