diff --git a/backend/app/main.py b/backend/app/main.py index a3724a27e..5cbd5f81c 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -137,6 +137,7 @@ async def lifespan(app: FastAPI): import sys import os from contextlib import AsyncExitStack + from app.services.scheduler import start_scheduler from app.services.trigger_daemon import start_trigger_daemon from app.services.tool_seeder import seed_builtin_tools from app.services.template_seeder import seed_agent_templates @@ -304,7 +305,10 @@ def _bg_task_error(t): task_specs = [] if _role_enabled("all", "worker"): - task_specs.append(("trigger_daemon", start_trigger_daemon())) + task_specs.extend([ + ("trigger_daemon", start_trigger_daemon()), + ("agent_schedule_scheduler", start_scheduler()), + ]) if _role_enabled("all", "connector"): task_specs.extend([ ("feishu_ws", feishu_ws_manager.start_all()), diff --git a/backend/app/services/scheduler.py b/backend/app/services/scheduler.py index 8aeb8eaf2..8b9f00232 100644 --- a/backend/app/services/scheduler.py +++ b/backend/app/services/scheduler.py @@ -61,7 +61,7 @@ async def _tick(): agent = agent_result.scalar_one_or_none() if ( agent is None - or agent.status != "running" + or agent.status not in {"creating", "running", "idle"} or is_agent_expired(agent) ): logger.info( diff --git a/backend/tests/test_schedule_scheduler.py b/backend/tests/test_schedule_scheduler.py new file mode 100644 index 000000000..d9158477a --- /dev/null +++ b/backend/tests/test_schedule_scheduler.py @@ -0,0 +1,145 @@ +"""Regression coverage for automatic AgentSchedule consumption.""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch +import uuid + +import pytest + +from app.services.scheduler import _tick + + +class _Result: + def __init__(self, *, rows: list[object] | None = None, value: object | None = None) -> None: + self._rows = rows + self._value = value + + def scalars(self) -> "_Result": + return self + + def all(self) -> list[object]: + return list(self._rows or []) + + def scalar_one_or_none(self) -> object | None: + return self._value + + +class _Session: + def __init__(self, schedule: object, agent: object) -> None: + self._results = [_Result(rows=[schedule]), _Result(value=agent)] + self.commits = 0 + self.rollbacks = 0 + + async def execute(self, _statement: object) -> _Result: + if not self._results: + raise AssertionError("unexpected database query") + return self._results.pop(0) + + async def commit(self) -> None: + self.commits += 1 + + async def rollback(self) -> None: + self.rollbacks += 1 + + +class _SessionContext: + def __init__(self, session: _Session) -> None: + self._session = session + + async def __aenter__(self) -> _Session: + return self._session + + async def __aexit__(self, _exc_type, _exc, _traceback) -> bool: + return False + + +def _records(*, status: str = "idle") -> tuple[SimpleNamespace, SimpleNamespace]: + due_at = datetime.now(timezone.utc) - timedelta(minutes=1) + schedule = SimpleNamespace( + id=uuid.uuid4(), + agent_id=uuid.uuid4(), + name="daily-summary", + instruction="Prepare the daily summary", + cron_expr="0 9 * * *", + is_enabled=True, + last_run_at=None, + next_run_at=due_at, + run_count=0, + ) + agent = SimpleNamespace( + id=schedule.agent_id, + status=status, + is_expired=False, + expires_at=None, + ) + return schedule, agent + + +@pytest.mark.asyncio +@pytest.mark.parametrize("agent_status", ["creating", "running", "idle"]) +async def test_due_schedule_for_active_agent_is_enqueued_and_advanced( + agent_status: str, +) -> None: + schedule, agent = _records(status=agent_status) + session = _Session(schedule, agent) + handle = SimpleNamespace(run_id=uuid.uuid4()) + enqueue = AsyncMock(return_value=handle) + + with ( + patch("app.database.async_session", return_value=_SessionContext(session)), + patch("app.services.audit_logger.write_audit_log", new=AsyncMock()), + patch("app.services.heartbeat_runtime.enqueue_schedule_runtime", new=enqueue), + ): + await _tick() + + enqueue.assert_awaited_once() + assert enqueue.await_args.kwargs["agent"] is agent + assert enqueue.await_args.kwargs["schedule_id"] == schedule.id + assert session.commits == 1 + assert session.rollbacks == 0 + assert schedule.run_count == 1 + assert schedule.last_run_at is not None + assert schedule.next_run_at > schedule.last_run_at + + +@pytest.mark.asyncio +async def test_due_schedule_is_not_advanced_when_runtime_is_disabled() -> None: + schedule, agent = _records(status="idle") + original_next_run = schedule.next_run_at + session = _Session(schedule, agent) + + with ( + patch("app.database.async_session", return_value=_SessionContext(session)), + patch("app.services.audit_logger.write_audit_log", new=AsyncMock()), + patch( + "app.services.heartbeat_runtime.enqueue_schedule_runtime", + new=AsyncMock(return_value=None), + ) as enqueue, + ): + await _tick() + + enqueue.assert_awaited_once() + assert session.commits == 0 + assert session.rollbacks == 1 + assert schedule.run_count == 0 + assert schedule.last_run_at is None + assert schedule.next_run_at == original_next_run + + +@pytest.mark.asyncio +async def test_due_schedule_does_not_enqueue_for_stopped_agent() -> None: + schedule, agent = _records(status="stopped") + session = _Session(schedule, agent) + enqueue = AsyncMock() + + with ( + patch("app.database.async_session", return_value=_SessionContext(session)), + patch("app.services.audit_logger.write_audit_log", new=AsyncMock()), + patch("app.services.heartbeat_runtime.enqueue_schedule_runtime", new=enqueue), + ): + await _tick() + + enqueue.assert_not_awaited() diff --git a/backend/tests/test_schedule_scheduler_startup.py b/backend/tests/test_schedule_scheduler_startup.py new file mode 100644 index 000000000..f8414e12a --- /dev/null +++ b/backend/tests/test_schedule_scheduler_startup.py @@ -0,0 +1,78 @@ +"""Regression coverage for AgentSchedule scheduler startup wiring.""" + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock + +import pytest + +import app.main as main +from app.services import audit_logger, scheduler, trigger_daemon +from app.services.agent_runtime import worker_service + + +class _Task: + def __init__(self, coro, name: str) -> None: + self._name = name + coro.close() + + def add_done_callback(self, _callback) -> None: + return None + + def get_name(self) -> str: + return self._name + + def exception(self): + return None + + +async def _collect_background_task_names(monkeypatch, *, process_role: str) -> list[str]: + created: list[str] = [] + + def create_task(coro, *, name: str) -> _Task: + created.append(name) + return _Task(coro, name) + + @asynccontextmanager + async def runtime_context(**_kwargs): + yield + + monkeypatch.setattr(main.settings, "PROCESS_ROLE", process_role) + monkeypatch.setattr(main, "configure_logging", lambda: None) + monkeypatch.setattr(main, "intercept_standard_logging", lambda: None) + monkeypatch.setattr(main, "_log_bwrap_startup_status", lambda: None) + monkeypatch.setattr(asyncio, "create_task", create_task) + monkeypatch.setattr(main, "_start_ss_local", AsyncMock()) + monkeypatch.setattr(main, "close_redis", AsyncMock()) + monkeypatch.setattr(main.realtime_router, "start", AsyncMock()) + monkeypatch.setattr(main.realtime_router, "stop", AsyncMock()) + monkeypatch.setattr(audit_logger, "write_audit_log", AsyncMock()) + monkeypatch.setattr(trigger_daemon, "start_trigger_daemon", AsyncMock()) + monkeypatch.setattr(scheduler, "start_scheduler", AsyncMock()) + monkeypatch.setattr(worker_service, "running_runtime_worker_context", runtime_context) + + async with main.lifespan(main.app): + pass + + return created + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("process_role", "expected"), + [ + ("worker", True), + ("api", False), + ], +) +async def test_agent_schedule_scheduler_follows_worker_role( + monkeypatch, + process_role: str, + expected: bool, +) -> None: + names = await _collect_background_task_names(monkeypatch, process_role=process_role) + + assert ("trigger_daemon" in names) is expected + assert ("agent_schedule_scheduler" in names) is expected diff --git a/restart.sh b/restart.sh index c5e717d17..a4e7253b9 100755 --- a/restart.sh +++ b/restart.sh @@ -233,9 +233,6 @@ start_backend() { echo -e "${YELLOW}🔄 Running LangGraph checkpoint migrations...${NC}" .venv/bin/python -m app.scripts.setup_langgraph_checkpoints - # Auto-run data migrations (idempotent) - echo -e "${YELLOW}🔄 Running data migrations...${NC}" - .venv/bin/python -m app.scripts.migrate_schedules_to_triggers || true start_detached "$BACKEND_DIR" "$BACKEND_LOG" "$BACKEND_PID" \ env PYTHONUNBUFFERED=1 \ AGENT_RUNTIME_V2_ENABLED=true \