From 6a7fa91bbfe2b5ce5518d648f46dfdc43dced01d Mon Sep 17 00:00:00 2001 From: Y1fe1Zh0u Date: Wed, 22 Jul 2026 17:59:27 +0800 Subject: [PATCH] Restore automatic execution for Agent task schedules The Tasks UI still persists legacy AgentSchedule rows, while worker startup only launched the AgentTrigger daemon. Start the existing schedule consumer in the worker role, accept active Agent lifecycle states, and stop auto-migrating schedules during restart. Constraint: Tasks UI continues to write AgentSchedule in this targeted fix Rejected: Rely on one-time AgentSchedule migration | schedules created after service startup would remain unconsumed Confidence: high Scope-risk: narrow Directive: Keep timezone semantics out of this fix and address them in the scheduled-task refactor Tested: 18 focused scheduler/runtime tests; full backend suite with 1993 passing; Ruff on changed scheduler/tests; bash -n restart.sh Not-tested: 3010 deployment and real wall-clock timezone behavior --- backend/app/main.py | 6 +- backend/app/services/scheduler.py | 2 +- backend/tests/test_schedule_scheduler.py | 145 ++++++++++++++++++ .../tests/test_schedule_scheduler_startup.py | 78 ++++++++++ restart.sh | 3 - 5 files changed, 229 insertions(+), 5 deletions(-) create mode 100644 backend/tests/test_schedule_scheduler.py create mode 100644 backend/tests/test_schedule_scheduler_startup.py 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 \