diff --git a/backend/app/scheduler/service.py b/backend/app/scheduler/service.py index 3080fe008..9bcfc268a 100644 --- a/backend/app/scheduler/service.py +++ b/backend/app/scheduler/service.py @@ -76,6 +76,19 @@ class ScheduledTaskService: return "failed" return "enabled" + @staticmethod + def _task_status_for_launch(task: dict[str, Any], *, trigger: str) -> str: + # The task-level status to write once _launch_run has produced a live + # run. A `once` task stays "running" until handle_run_completion + # observes the real terminal outcome; declaring "completed" at launch + # would stick if the run fails or the process dies (startup + # reconciliation is cancel_stuck_once_tasks). + if task["schedule_type"] == "once": + return "running" + if trigger == "manual" and task.get("status") == "paused": + return "paused" + return "enabled" + @staticmethod def _task_status_for_skip(task: dict[str, Any]) -> str: if task["schedule_type"] == "once": @@ -131,6 +144,21 @@ class ScheduledTaskService: if trigger == "manual": return self._active_run_conflict_result(execution_thread_id) return await self._record_scheduled_skip(task, thread_id=execution_thread_id, now=now, trigger=trigger) + # Track whether _launch_run has produced a live run. A bookkeeping + # failure AFTER launch (the queued->running write, or the parent task + # update) must NOT be recorded as "failed": "failed" is outside the + # partial unique index uq_scheduled_task_run_active, so it would release + # the task's single active slot and the next dispatch would launch a + # duplicate run. Once launch succeeds we keep the row "running" and + # retain the launched run_id regardless of bookkeeping errors. + launched_run_id: str | None = None + launched_thread_id: str | None = None + # Flip immediately after _launch_run returns, before any further code + # that can raise (e.g. result["run_id"] on a malformed result). The + # retention branch keys off this flag, not `launched_run_id is not + # None`, so a launch that succeeded but whose result-unpacking raised + # still takes the retention path instead of the release-the-slot path. + launch_succeeded = False try: result = await self._launch_run( thread_id=execution_thread_id, @@ -143,26 +171,20 @@ class ScheduledTaskService: "scheduled_trigger": trigger, }, ) + launch_succeeded = True + launched_run_id = result["run_id"] + launched_thread_id = result["thread_id"] next_at = next_run_at( task["schedule_type"], task["schedule_spec"], task["timezone"], now=now, ) - if task["schedule_type"] == "once": - # Stay "running" until handle_run_completion sees the real - # terminal outcome; declaring "completed" at launch would stick - # if the run fails or the process dies (startup reconciliation - # is cancel_stuck_once_tasks). - task_status = "running" - elif trigger == "manual" and task.get("status") == "paused": - task_status = "paused" - else: - task_status = "enabled" + task_status = self._task_status_for_launch(task, trigger=trigger) await self._task_run_repo.update_status( task_run_id, status="running", - run_id=result["run_id"], + run_id=launched_run_id, started_at=now, # A fast-failing run can reach handle_run_completion before this # write resumes; never clobber its terminal status. @@ -173,8 +195,8 @@ class ScheduledTaskService: status=task_status, next_run_at=next_at, last_run_at=now, - last_run_id=result["run_id"], - last_thread_id=result["thread_id"], + last_run_id=launched_run_id, + last_thread_id=launched_thread_id, last_error=None, increment_run_count=True, # Same race as the run-row write above: a fast-failing run's @@ -184,20 +206,92 @@ class ScheduledTaskService: return { "outcome": "launched", "task_run_id": task_run_id, - "run_id": result["run_id"], - "thread_id": result["thread_id"], + "run_id": launched_run_id, + "thread_id": launched_thread_id, "error": None, } except Exception as exc: + if not launch_succeeded and self._is_overlap_conflict(exc) and trigger == "scheduled" and task.get("overlap_policy", "skip") == "skip": + # Pre-launch overlap conflict (e.g. same-thread multitask): no + # run was started, so recording a skip and releasing the slot is + # safe. Guarded by ``not launch_succeeded`` because a run that + # already launched must never be reclassified as a skip / failed. + return await self._finalize_skip( + task, + task_run_id=task_run_id, + thread_id=execution_thread_id, + now=now, + error=str(exc), + ) + next_at = next_run_at( task["schedule_type"], task["schedule_spec"], task["timezone"], now=now, ) - if self._is_overlap_conflict(exc) and trigger == "scheduled" and task.get("overlap_policy", "skip") == "skip": - return await self._finalize_skip(task, task_run_id=task_run_id, thread_id=execution_thread_id, now=now, error=str(exc)) + if launch_succeeded: + # _launch_run succeeded, so a run is live even though + # post-launch bookkeeping raised. Keep the task-run row + # "running" so it keeps holding the task's single active slot + # (preventing a duplicate launch on the next dispatch) and + # persist the run_id on the parent task for recovery / + # reconciliation / cancellation. These writes are best-effort: + # if the DB is still down the row stays "queued" -- still + # active, still holding the slot -- so we log and still report + # the run as launched so callers know a run is in flight. + task_status = self._task_status_for_launch(task, trigger=trigger) + try: + await self._task_run_repo.update_status( + task_run_id, + status="running", + run_id=launched_run_id, + started_at=now, + protect_terminal=True, + ) + except Exception: + logger.exception( + "Scheduled task-run %s: post-launch bookkeeping failed; run %s is still live (task %s)", + task_run_id, + launched_run_id, + task["id"], + ) + try: + await self._task_repo.update_after_launch( + task["id"], + status=task_status, + next_run_at=next_at, + last_run_at=now, + last_run_id=launched_run_id, + last_thread_id=launched_thread_id, + # The bookkeeping exception is an infrastructure-level + # transient, not a run-level failure: the run launched + # and is still in flight. Clear last_error like the + # success path so the task list does not show an error + # on a task whose run is actively running; the real + # terminal outcome is written by handle_run_completion. + # The transient itself is logged above. + last_error=None, + increment_run_count=True, + protect_terminal=True, + ) + except Exception: + logger.exception( + "Scheduled task %s: post-launch update failed; run %s is still live", + task["id"], + launched_run_id, + ) + return { + "outcome": "launched", + "task_run_id": task_run_id, + "run_id": launched_run_id, + "thread_id": launched_thread_id, + "error": str(exc), + } + + # _launch_run itself failed (or a step before it did): no live run + # was created, so it is safe to release the active slot. task_status = self._task_status_for_failure(task, trigger=trigger) await self._task_run_repo.update_status( task_run_id, diff --git a/backend/tests/test_scheduled_task_service.py b/backend/tests/test_scheduled_task_service.py index 07ee0b8a9..a071ad298 100644 --- a/backend/tests/test_scheduled_task_service.py +++ b/backend/tests/test_scheduled_task_service.py @@ -541,3 +541,255 @@ async def test_launch_bookkeeping_passes_protect_terminal(): await service.dispatch_task(task_repo.rows[0], now=datetime.now(UTC), trigger="scheduled") assert task_repo.updated[1]["protect_terminal"] is True + + +class _StatefulRunRepo: + """Stateful fake ``ScheduledTaskRunRepository`` for the #4452 tests. + + Mirrors just enough of the real repository to let a second dispatch + observe the active slot held by the first: + + * ``create()`` tracks each row by id, carrying its ``status`` and + ``run_id``; + * with ``fail_first_update=True`` the very FIRST ``update_status()`` + call raises, simulating a transient DB failure on the + ``queued -> running`` write that fires right after ``_launch_run`` + returns a live ``run_id``; every later ``update_status()`` applies; + * ``has_active_runs()`` reflects whether any tracked row for the task + is still in an active status (``queued``/``running``), exactly like + the partial unique index ``uq_scheduled_task_run_active``. + """ + + _ACTIVE = {"queued", "running"} + + def __init__(self, *, fail_first_update: bool = False) -> None: + self.created: list[dict] = [] + self.updates: list[tuple[str, dict]] = [] + self.rows: dict[str, dict] = {} + self._fail_first_update = fail_first_update + self._first_update_raised = False + + async def count_active_runs(self) -> int: + return sum(1 for row in self.rows.values() if row["status"] in self._ACTIVE) + + async def create(self, **kwargs) -> dict: + self.created.append(kwargs) + self.rows[kwargs["run_record_id"]] = { + "task_id": kwargs["task_id"], + "status": kwargs["status"], + "run_id": None, + } + return {"id": kwargs["run_record_id"]} + + async def update_status(self, run_record_id: str, **kwargs) -> None: + self.updates.append((run_record_id, kwargs)) + if self._fail_first_update and not self._first_update_raised: + # The launch-path queued->running write fails once, AFTER + # _launch_run has already returned a live run_id. + self._first_update_raised = True + raise RuntimeError("simulated transient DB error on queued->running write") + row = self.rows.get(run_record_id) + if row is None: + return + if "status" in kwargs: + row["status"] = kwargs["status"] + if kwargs.get("run_id") is not None: + row["run_id"] = kwargs["run_id"] + + async def has_active_runs(self, task_id: str) -> bool: + return any(row["task_id"] == task_id and row["status"] in self._ACTIVE for row in self.rows.values()) + + async def mark_stale_active_runs(self, *, error: str) -> int: + return 0 + + +@pytest.mark.asyncio +async def test_post_launch_bookkeeping_failure_does_not_release_active_slot(): + """Regression for issue #4452. + + A transient failure in the ``queued -> running`` bookkeeping write + (after ``_launch_run`` has already returned a live ``run_id``) must NOT + flip the task-run row to ``failed``: ``failed`` is outside the partial + unique index ``uq_scheduled_task_run_active``, so releasing the slot + would let the next dispatch launch a DUPLICATE run. The fix keeps the + row ``running`` with the launched ``run_id`` retained for recovery, + reconciliation, and cancellation. + """ + launched: list[dict] = [] + + async def fake_launch(**kwargs): + launched.append(kwargs) + return {"run_id": f"run-{len(launched)}", "thread_id": kwargs["thread_id"]} + + task_repo = DummyTaskRepo( + [ + { + "id": "task-4452", + "user_id": "user-1", + "thread_id": None, + "context_mode": "fresh_thread_per_run", + "assistant_id": "lead_agent", + "prompt": "do the thing", + "schedule_type": "cron", + "schedule_spec": {"cron": "*/5 * * * *"}, + "timezone": "UTC", + "status": "enabled", + "overlap_policy": "skip", + } + ] + ) + run_repo = _StatefulRunRepo(fail_first_update=True) + service = ScheduledTaskService( + task_repo=task_repo, + task_run_repo=run_repo, + launch_run=fake_launch, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_runs=3, + ) + + now = datetime.now(UTC) + task = dict(task_repo.rows[0]) + + first = await service.dispatch_task(task, now=now, trigger="scheduled") + # The run launched despite the post-launch bookkeeping error; the + # outcome and run_id reflect that a live run is in flight. + assert first["outcome"] == "launched" + assert first["run_id"] == "run-1" + assert first["error"] is not None # the bookkeeping error is surfaced, not hidden + + # Second dispatch must observe the active slot held by run-1 and NOT + # launch a duplicate. On main (bug) this would launch run-2 here. + second = await service.dispatch_task(task, now=now, trigger="scheduled") + assert len(launched) == 1, launched + assert second["outcome"] in {"skipped", "conflict"}, second + + # The launched run_id is retained on the task-run row (status "running", + # not "failed") so reconciliation / cancellation can still reach it. + first_row_id = run_repo.created[0]["run_record_id"] + assert run_repo.rows[first_row_id]["status"] == "running" + assert run_repo.rows[first_row_id]["run_id"] == "run-1" + + # The bookkeeping transient is NOT surfaced as the parent task's + # last_error: the run launched and is still in flight, so the task list + # must not show an error on an actively running task (matching the + # success path's clear-on-launch model). The real terminal outcome is + # written by handle_run_completion. + assert task_repo.updated[1]["last_error"] is None + + +@pytest.mark.asyncio +async def test_pre_launch_failure_still_releases_active_slot(): + """Complement to the #4452 fix: when ``_launch_run`` itself fails (no run + was ever started), the task-run row is marked ``failed`` and the active + slot is released as before -- the post-launch retention path does not + apply because there is no live run to protect. + """ + launched: list[dict] = [] + + async def fake_launch(**kwargs): + launched.append(kwargs) + raise RuntimeError("runtime refused to start the run") + + task_repo = DummyTaskRepo( + [ + { + "id": "task-4452-pre", + "user_id": "user-1", + "thread_id": None, + "context_mode": "fresh_thread_per_run", + "assistant_id": "lead_agent", + "prompt": "do the thing", + "schedule_type": "cron", + "schedule_spec": {"cron": "*/5 * * * *"}, + "timezone": "UTC", + "status": "enabled", + "overlap_policy": "skip", + } + ] + ) + run_repo = _StatefulRunRepo() + service = ScheduledTaskService( + task_repo=task_repo, + task_run_repo=run_repo, + launch_run=fake_launch, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_runs=3, + ) + + result = await service.dispatch_task(dict(task_repo.rows[0]), now=datetime.now(UTC), trigger="scheduled") + + assert result["outcome"] == "failed" + assert result["run_id"] is None + # launch was attempted (and raised), so exactly one launch attempt, and + # the row is terminal -> the slot is released for the next dispatch. + assert len(launched) == 1 + first_row_id = run_repo.created[0]["run_record_id"] + assert run_repo.rows[first_row_id]["status"] == "failed" + assert run_repo.rows[first_row_id]["run_id"] is None + + +@pytest.mark.asyncio +async def test_malformed_launch_result_still_retains_active_slot(): + """Defense-in-depth for the #4452 invariant. + + If ``_launch_run`` returns a malformed result (e.g. missing ``run_id``), + the unpacking line raises AFTER a live run was already created. The + dispatch must still take the retention path (keep the row active so the + slot stays held and no duplicate launches) rather than the pre-launch + generic-failure path, which would mark the row ``failed`` and release + the slot while a run is in flight. Keyed off ``launch_succeeded``, not + ``launched_run_id is not None``. + """ + launched: list[dict] = [] + + async def fake_launch(**kwargs): + launched.append(kwargs) + # Live run started, but the result payload is malformed. + return {"thread_id": kwargs["thread_id"]} + + task_repo = DummyTaskRepo( + [ + { + "id": "task-4452-malformed", + "user_id": "user-1", + "thread_id": None, + "context_mode": "fresh_thread_per_run", + "assistant_id": "lead_agent", + "prompt": "do the thing", + "schedule_type": "cron", + "schedule_spec": {"cron": "*/5 * * * *"}, + "timezone": "UTC", + "status": "enabled", + "overlap_policy": "skip", + } + ] + ) + run_repo = _StatefulRunRepo() + service = ScheduledTaskService( + task_repo=task_repo, + task_run_repo=run_repo, + launch_run=fake_launch, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_runs=3, + ) + + now = datetime.now(UTC) + task = dict(task_repo.rows[0]) + + first = await service.dispatch_task(task, now=now, trigger="scheduled") + # Launch succeeded, so the outcome is "launched" (a run is in flight) + # even though the result unpacking raised; run_id is unknown. + assert first["outcome"] == "launched" + assert first["run_id"] is None + + # Second dispatch must observe the active slot still held (row stays in + # an active status, NOT "failed") and NOT launch a duplicate. + second = await service.dispatch_task(task, now=now, trigger="scheduled") + assert len(launched) == 1, launched + assert second["outcome"] in {"skipped", "conflict"}, second + + first_row_id = run_repo.created[0]["run_record_id"] + assert run_repo.rows[first_row_id]["status"] == "running"