From 95667638bf509fa8b3a7b5a8119bc0989eabc0b9 Mon Sep 17 00:00:00 2001 From: hetaoBackend Date: Sat, 25 Jul 2026 22:38:28 +0800 Subject: [PATCH] fix(runtime): close pre-attachment run lifecycle gap (#2532) --- README.md | 2 + backend/AGENTS.md | 1 + backend/README.md | 2 + backend/app/gateway/services.py | 128 +++++++++------- .../harness/deerflow/runtime/runs/manager.py | 76 ++++++++-- backend/tests/test_gateway_services.py | 143 ++++++++++++++++-- .../tests/test_multi_worker_run_ownership.py | 47 +++--- backend/tests/test_run_manager.py | 34 +++++ 8 files changed, 330 insertions(+), 103 deletions(-) diff --git a/README.md b/README.md index 732977912..38bc800b7 100644 --- a/README.md +++ b/README.md @@ -277,6 +277,8 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser > [!IMPORTANT] > The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), and `run_ownership.heartbeat_enabled: true`. The bridge shares SSE delivery and `Last-Event-ID` replay across workers, while lease reconciliation marks runs from dead workers as errors, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE and `/wait` consumers also refresh durable status on heartbeats as a fallback if terminal publication fails. The rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. IM channel state and other process-local services still need their own multi-worker coordination. +Gateway run creation completes input, configuration, checkpoint, authenticated-context, and thread-metadata setup before reserving the thread with a durable pending run. Worker attachment is conditional on that pending run still being active, and in multi-worker mode the initial lease is the attachment deadline rather than a renewable heartbeat until a worker exists. Validation/cancellation failures therefore cannot strand a pending reservation or start an agent after cancellation. + See [CONTRIBUTING.md](CONTRIBUTING.md) for detailed Docker development guide. #### Option 2: Local Development diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 1bb3a785a..026f58370 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -16,6 +16,7 @@ DeerFlow is a LangGraph-based AI super agent system with a full-stack architectu - `make dev`, Docker dev, and production all run the agent runtime in Gateway via `RunManager` + `run_agent()` + `StreamBridge` (`packages/harness/deerflow/runtime/`). Nginx exposes that runtime at `/api/langgraph/*` and rewrites it to Gateway's native `/api/*` routers. - Gateway streams `write_file` and `str_replace` argument deltas in bounded batches when clients also subscribe to `values`; messages-only consumers retain the original per-chunk contract, while `values` preserves the complete tool call. - With `stream_subgraphs`, subgraph frames keep their namespace in the SSE event name (`values|`, LangGraph Platform style) instead of impersonating root frames — a delegated subagent inherits the parent checkpoint namespace, so publishing its `values` snapshot as bare `values` replaces the whole thread view in SDK clients (#4399). Root-only consumers (file-tool chunk batcher, subagent event persistence, LLM error-fallback detection) ignore namespaced frames. The web frontend does not request subgraph streaming; subtask progress rides root-namespace `task_*` custom events. +- Gateway run admission happens only after input/config/checkpoint validation, authenticated runtime-context assembly, and thread-metadata setup complete. After `create_or_reject()` persists the pending row, `RunManager.attach_task_if_active()` performs a lock-protected status check plus task creation; cancellation on either side of that boundary closes the admitted run. In heartbeat-enabled multi-worker mode, an unattached pending run keeps its initial lease as a launch deadline and is marked `error` when it expires instead of being renewed indefinitely (#2532). - Scheduled-task executions must reuse that same Gateway run lifecycle. The scheduler may decide *when* work runs, but it must dispatch through the existing run path rather than introducing a parallel execution stack. - Scheduled-task dispatch enforces "at most one active run per task when `overlap_policy=skip`" at the DB layer via the partial unique index `uq_scheduled_task_run_active` (`scheduled_task_runs.task_id WHERE status IN ('queued','running')`). `ScheduledTaskService.dispatch_task`'s `has_active_runs` check is a non-atomic fast path (its own session, separated from the `create()` insert by `await` points), so two concurrent dispatches — a manual `POST /scheduled-tasks/{id}/trigger` racing the poller, a double-click, or a client retry — can both pass it; the index is the atomic arbiter, and the losing `create` surfaces as `ActiveScheduledRunConflict` (translated from `IntegrityError` in the repository) and collapses to the same outcome as the fast path (manual → 409 conflict, scheduled → a `"skipped"` tombstone). The scheduled-skip tombstone is created directly as terminal `"skipped"` (not a transient `"queued"`) so it never occupies the active slot the pre-existing run still holds. Sibling of the `runs` table's `uq_runs_thread_active` (PR #4003), which keys on `thread_id` and so does not cover the default `fresh_thread_per_run` context where every dispatch gets a new thread. Index is status-only, not `overlap_policy`-conditional (the policy is fixed to `"skip"` in the MVP). diff --git a/backend/README.md b/backend/README.md index 6e9818673..e6e57a7fd 100644 --- a/backend/README.md +++ b/backend/README.md @@ -111,6 +111,8 @@ LLM-powered persistent context retention across conversations: FastAPI application providing REST endpoints for frontend integration: +Run launch preparation (input/config/checkpoint validation, authenticated context assembly, and thread metadata setup) completes before durable run admission. `RunManager` then attaches the background worker only if the pending run is still active. In heartbeat-enabled multi-worker mode, the initial lease is also the worker-attachment deadline and is not renewed until a task exists, preventing leaked pending rows from reserving a thread indefinitely. + | Route | Purpose | |-------|---------| | `GET /api/models` | List available LLM models | diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index e125d4eae..6b1516f4e 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -958,49 +958,10 @@ async def start_run( owner_context_token = set_current_user(SimpleNamespace(id=owner_user_id)) if owner_user_id else None try: - try: - async with goal_thread_lock(thread_id): - record = await run_mgr.create_or_reject( - thread_id, - body.assistant_id, - on_disconnect=disconnect, - metadata=body.metadata or {}, - # Persist a secret-redacted copy of the config: the run record is - # written to runs.kwargs_json and echoed by the run API, so a - # request-scoped secret (#3861) must not ride along. The live - # config built below keeps the secrets for the actual run. - kwargs={"input": body.input, "config": redact_config_secrets(body.config)}, - multitask_strategy=body.multitask_strategy, - model_name=model_name, - user_id=owner_user_id, - ) - except ConflictError as exc: - raise HTTPException(status_code=409, detail=str(exc)) from exc - except UnsupportedStrategyError as exc: - raise HTTPException(status_code=501, detail=str(exc)) from exc - - # Upsert thread metadata so the thread appears in /threads/search, - # even for threads that were never explicitly created via POST /threads - # (e.g. stateless runs). - try: - existing = await run_ctx.thread_store.get(thread_id) - if existing is None and owner_user_id: - unscoped_existing = await run_ctx.thread_store.get(thread_id, user_id=None) - if unscoped_existing is not None: - if unscoped_existing.get("user_id") != owner_user_id: - await run_ctx.thread_store.update_owner(thread_id, owner_user_id, user_id=None) - existing = await run_ctx.thread_store.get(thread_id) - if existing is None: - await run_ctx.thread_store.create( - thread_id, - assistant_id=body.assistant_id, - metadata=body.metadata, - ) - else: - await run_ctx.thread_store.update_status(thread_id, "running") - except Exception: - logger.warning("Failed to upsert thread_meta for %s (non-fatal)", sanitize_log_param(thread_id)) - + # Complete every fallible/awaited launch-preparation step before durable + # admission. Once create_or_reject() registers a pending run, the only + # remaining transition before worker execution is the manager-guarded + # task attachment below. agent_factory = resolve_agent_factory(body.assistant_id) is_internal_caller = getattr(getattr(request, "state", None), "auth_source", None) == AUTH_SOURCE_INTERNAL command = getattr(body, "command", None) @@ -1028,22 +989,73 @@ async def start_run( request_context=getattr(body, "context", None), ) - task = asyncio.create_task( - run_agent( - bridge, - run_mgr, - record, - ctx=run_ctx, - agent_factory=agent_factory, - graph_input=graph_input, - config=config, - stream_modes=stream_modes, - stream_subgraphs=body.stream_subgraphs, - interrupt_before=body.interrupt_before, - interrupt_after=body.interrupt_after, + # Upsert thread metadata so the thread appears in /threads/search, + # even for threads that were never explicitly created via POST /threads + # (e.g. stateless runs). This awaited setup stays before admission: a + # request cancelled here must not leave a pending run behind. + try: + existing = await run_ctx.thread_store.get(thread_id) + if existing is None and owner_user_id: + unscoped_existing = await run_ctx.thread_store.get(thread_id, user_id=None) + if unscoped_existing is not None: + if unscoped_existing.get("user_id") != owner_user_id: + await run_ctx.thread_store.update_owner(thread_id, owner_user_id, user_id=None) + existing = await run_ctx.thread_store.get(thread_id) + if existing is None: + await run_ctx.thread_store.create( + thread_id, + assistant_id=body.assistant_id, + metadata=body.metadata, + ) + else: + await run_ctx.thread_store.update_status(thread_id, "running") + except Exception: + logger.warning("Failed to upsert thread_meta for %s (non-fatal)", sanitize_log_param(thread_id)) + + try: + async with goal_thread_lock(thread_id): + record = await run_mgr.create_or_reject( + thread_id, + body.assistant_id, + on_disconnect=disconnect, + metadata=body.metadata or {}, + # Persist a secret-redacted copy of the config: the run record is + # written to runs.kwargs_json and echoed by the run API, so a + # request-scoped secret (#3861) must not ride along. The live + # config built below keeps the secrets for the actual run. + kwargs={"input": body.input, "config": redact_config_secrets(body.config)}, + multitask_strategy=body.multitask_strategy, + model_name=model_name, + user_id=owner_user_id, + ) + except ConflictError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc + except UnsupportedStrategyError as exc: + raise HTTPException(status_code=501, detail=str(exc)) from exc + + try: + await run_mgr.attach_task_if_active( + record.run_id, + lambda: run_agent( + bridge, + run_mgr, + record, + ctx=run_ctx, + agent_factory=agent_factory, + graph_input=graph_input, + config=config, + stream_modes=stream_modes, + stream_subgraphs=body.stream_subgraphs, + interrupt_before=body.interrupt_before, + interrupt_after=body.interrupt_after, + ), ) - ) - record.task = task + except asyncio.CancelledError: + # create_or_reject() has already crossed the durable admission + # boundary. Close the pending row if the request is cancelled + # before task attachment can complete. + await asyncio.shield(run_mgr.cancel(record.run_id)) + raise # Title sync is handled by worker.py's finally block which reads the # title from the checkpoint and calls thread_store.update_display_name diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 04a5d335c..208014e48 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -7,7 +7,7 @@ import logging import socket import sqlite3 import uuid -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Coroutine from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta from enum import StrEnum @@ -30,6 +30,7 @@ logger = logging.getLogger(__name__) ORPHAN_RECOVERY_STOP_REASON = "orphan_recovered" STARTUP_ORPHAN_RECOVERY_ERROR = "Gateway restarted before this run reached a durable final state." LEASE_ORPHAN_RECOVERY_ERROR = "Run lease expired — owning worker is unreachable." +RUN_LAUNCH_TIMEOUT_ERROR = "Run did not attach a worker before its launch lease expired." _RETRYABLE_SQLITE_MESSAGES = ( "database is locked", @@ -1093,13 +1094,45 @@ class RunManager: interrupted_records.append(r) # Outside the lock: persist interrupted status for locally-cancelled - # runs. Store-side claimed rows are already finalised. - for interrupted_record in interrupted_records: - await self._persist_status(interrupted_record, RunStatus.interrupted) + # runs. Store-side claimed rows are already finalised. Cancellation at + # this point happens after the replacement was admitted, so close that + # new run before propagating cancellation to the caller. + try: + for interrupted_record in interrupted_records: + await self._persist_status(interrupted_record, RunStatus.interrupted) + except asyncio.CancelledError: + try: + await asyncio.shield(self.cancel(record.run_id)) + except Exception: + logger.warning("Failed to close run %s after admission was cancelled", record.run_id, exc_info=True) + raise logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) return record + async def attach_task_if_active( + self, + run_id: str, + task_factory: Callable[[], Coroutine[Any, Any, None]], + ) -> asyncio.Task[None] | None: + """Attach and schedule a worker only while a pending run is active. + + Cancellation can race the gap between durable admission and worker + attachment. Creating the task while holding the manager lock makes + the status check and attachment one atomic lifecycle transition: a + run interrupted before attachment never starts executing. + """ + async with self._lock: + record = self._runs.get(run_id) + if record is None or record.status != RunStatus.pending or record.abort_event.is_set(): + return None + if record.task is not None: + raise RuntimeError(f"Run {run_id} already has an attached task") + + task = asyncio.create_task(task_factory()) + record.task = task + return task + async def reconcile_orphaned_inflight_runs( self, *, @@ -1294,24 +1327,35 @@ class RunManager: self._schedule_orphan_reconciliation() async def _renew_leases(self) -> None: - """Renew the lease on every locally-owned active run.""" + """Renew attached workers and expire pending runs that never launched.""" if self._store is None or self._run_ownership_config is None: return lease_seconds = self._run_ownership_config.lease_seconds new_expiry = (datetime.now(UTC) + timedelta(seconds=lease_seconds)).isoformat() + expired_unattached: list[RunRecord] = [] async with self._lock: - # Renew any pending/running run owned by this worker unless its - # background task has already completed. A pending run whose task - # has not been spawned yet (``task is None``) is still live from - # this worker's perspective — between ``create_run_atomic`` - # inserting the row and the worker layer spawning the agent task - # there is a brief window. If we drop those records here and the - # window stretches past ``lease_seconds`` (e.g. event-loop - # saturation, slow checkpoint hydrate on a fresh worker), peer - # reconciliation will reclaim the run as an orphan and mark it - # ``error`` even though this worker still intends to execute it. - active_runs = [(rid, record) for rid, record in self._runs.items() if record.status in (RunStatus.pending, RunStatus.running) and record.owner_worker_id == self._worker_id and (record.task is None or not record.task.done())] + active_runs: list[tuple[str, RunRecord]] = [] + for run_id, record in self._runs.items(): + if record.status not in (RunStatus.pending, RunStatus.running) or record.owner_worker_id != self._worker_id: + continue + if record.task is None and record.status == RunStatus.pending: + # Before attachment the initial lease is a launch deadline, + # not proof of a live worker. Renewing it would let a leaked + # pending admission reserve the thread forever. + if is_lease_expired(record.lease_expires_at, grace_seconds=0): + record.status = RunStatus.error + record.error = RUN_LAUNCH_TIMEOUT_ERROR + record.abort_event.set() + record.updated_at = _now_iso() + expired_unattached.append(record) + continue + if record.task is None or not record.task.done(): + active_runs.append((run_id, record)) + + for record in expired_unattached: + await self._persist_status(record, RunStatus.error, error=RUN_LAUNCH_TIMEOUT_ERROR) + logger.error("Run %s exceeded its worker-attachment deadline", record.run_id) for run_id, record in active_runs: try: diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index e237ce9e0..e7b1e6d13 100644 --- a/backend/tests/test_gateway_services.py +++ b/backend/tests/test_gateway_services.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from typing import Any import pytest @@ -1062,14 +1063,33 @@ def test_inject_authenticated_user_context_strips_internal_spoofed_attribution() assert "oauth_id" not in config["context"] -async def _capture_start_run_graph_input(body, *, auth_source=None): - from types import SimpleNamespace +async def _capture_start_run_graph_input(body: Any, *, auth_source: str | None = None) -> object: from unittest.mock import patch + from app.gateway.services import start_run + + request, _run_manager = _make_start_run_test_request(auth_source=auth_source) + captured: dict[str, object] = {} + + async def fake_run_agent(*_args: Any, **kwargs: Any) -> None: + captured["graph_input"] = kwargs["graph_input"] + + with ( + patch("app.gateway.services.resolve_agent_factory", return_value=object()), + patch("app.gateway.services.run_agent", side_effect=fake_run_agent), + ): + record = await start_run(body, "thread-command-test", request) + await record.task + + return captured["graph_input"] + + +def _make_start_run_test_request(*, auth_source: str | None = None) -> tuple[Any, Any]: + from types import SimpleNamespace + from langgraph.checkpoint.memory import InMemorySaver from langgraph.store.memory import InMemoryStore - from app.gateway.services import start_run from deerflow.persistence.thread_meta.memory import MemoryThreadMetaStore from deerflow.runtime import RunManager from deerflow.runtime.runs.store.memory import MemoryRunStore @@ -1089,19 +1109,124 @@ async def _capture_start_run_graph_input(body, *, auth_source=None): state=SimpleNamespace(auth_source=auth_source), app=SimpleNamespace(state=state), ) - captured: dict[str, object] = {} + return request, run_manager - async def fake_run_agent(*args, **kwargs): - captured["graph_input"] = kwargs["graph_input"] + +@pytest.mark.asyncio +async def test_start_run_validation_failure_does_not_leave_active_run(_stub_app_config: None) -> None: + """A rejected checkpoint must not reserve the thread for later runs.""" + from unittest.mock import patch + + from fastapi import HTTPException + + from app.gateway.routers.thread_runs import RunCreateRequest + from app.gateway.services import start_run + from deerflow.runtime.runs.schemas import RunStatus + + request, run_manager = _make_start_run_test_request() + + async def fake_run_agent(_bridge: Any, manager: Any, record: Any, **_kwargs: Any) -> None: + await manager.set_status(record.run_id, RunStatus.success) + + invalid = RunCreateRequest( + input={"messages": [{"role": "user", "content": "hello"}]}, + checkpoint_id="missing-checkpoint", + ) + valid = RunCreateRequest(input={"messages": [{"role": "user", "content": "retry"}]}) with ( patch("app.gateway.services.resolve_agent_factory", return_value=object()), patch("app.gateway.services.run_agent", side_effect=fake_run_agent), ): - record = await start_run(body, "thread-command-test", request) - await record.task + with pytest.raises(HTTPException) as exc_info: + await start_run(invalid, "thread-validation-failure", request) - return captured["graph_input"] + assert exc_info.value.status_code == 404 + assert not await run_manager.has_inflight("thread-validation-failure") + + retry = await start_run(valid, "thread-validation-failure", request) + assert retry.task is not None + await retry.task + + assert retry.status == RunStatus.success + + +@pytest.mark.asyncio +async def test_start_run_does_not_attach_worker_after_pre_attach_cancel( + _stub_app_config: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cancellation after admission must prevent the worker from starting.""" + import asyncio + from unittest.mock import patch + + from app.gateway.routers.thread_runs import RunCreateRequest + from app.gateway.services import start_run + from deerflow.runtime.runs.manager import CancelOutcome + from deerflow.runtime.runs.schemas import RunStatus + + request, run_manager = _make_start_run_test_request() + create_or_reject = run_manager.create_or_reject + worker_started = asyncio.Event() + + async def admit_then_cancel(*args: Any, **kwargs: Any) -> Any: + record = await create_or_reject(*args, **kwargs) + assert await run_manager.cancel(record.run_id) == CancelOutcome.cancelled + return record + + async def fake_run_agent(*_args: Any, **_kwargs: Any) -> None: + worker_started.set() + + monkeypatch.setattr(run_manager, "create_or_reject", admit_then_cancel) + body = RunCreateRequest(input={"messages": [{"role": "user", "content": "hello"}]}) + + with ( + patch("app.gateway.services.resolve_agent_factory", return_value=object()), + patch("app.gateway.services.run_agent", side_effect=fake_run_agent), + ): + record = await start_run(body, "thread-pre-attach-cancel", request) + + assert record.status == RunStatus.interrupted + assert record.task is None + await asyncio.sleep(0) + assert not worker_started.is_set() + + +@pytest.mark.asyncio +async def test_start_run_cancelled_during_thread_metadata_does_not_admit_run( + _stub_app_config: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cancellation during metadata setup must not strand a pending run.""" + import asyncio + from unittest.mock import patch + + from app.gateway.routers.thread_runs import RunCreateRequest + from app.gateway.services import start_run + + request, run_manager = _make_start_run_test_request() + thread_store = request.app.state.thread_store + original_get = thread_store.get + metadata_started = asyncio.Event() + release_metadata = asyncio.Event() + + async def blocking_get(*args: Any, **kwargs: Any) -> Any: + metadata_started.set() + await release_metadata.wait() + return await original_get(*args, **kwargs) + + monkeypatch.setattr(thread_store, "get", blocking_get) + body = RunCreateRequest(input={"messages": [{"role": "user", "content": "hello"}]}) + + with patch("app.gateway.services.resolve_agent_factory", return_value=object()): + start_task = asyncio.create_task(start_run(body, "thread-metadata-cancel", request)) + await metadata_started.wait() + start_task.cancel() + release_metadata.set() + with pytest.raises(asyncio.CancelledError): + await start_task + + assert not await run_manager.has_inflight("thread-metadata-cancel") def test_start_run_translates_resume_command_to_langgraph_command(_stub_app_config): diff --git a/backend/tests/test_multi_worker_run_ownership.py b/backend/tests/test_multi_worker_run_ownership.py index 8d0939dcf..1b8d7e7af 100644 --- a/backend/tests/test_multi_worker_run_ownership.py +++ b/backend/tests/test_multi_worker_run_ownership.py @@ -677,18 +677,8 @@ async def test_heartbeat_renews_active_run_leases(): @pytest.mark.anyio -async def test_heartbeat_renews_pending_run_before_task_is_spawned(): - """A run sitting in ``pending`` between ``create_run_atomic`` and task - spawn must still have its lease renewed. - - Pre-fix the renewal filter required ``record.task is not None``, so a - pending run with no task yet (the brief window after - ``create_run_atomic`` inserts the row before the worker layer spawns - the agent task) was silently skipped. If that window stretched past - ``lease_seconds`` — e.g. event-loop saturation, slow checkpoint - hydrate — peer reconciliation reclaimed the run as an orphan and - marked it ``error`` even though this worker still intended to run it. - """ +async def test_heartbeat_does_not_renew_unattached_pending_before_launch_deadline(): + """The initial lease is a launch deadline until a worker is attached.""" config = _lease_config(lease_seconds=30, heartbeat_enabled=True) store = MemoryRunStore() manager = _make_manager(store=store, run_ownership_config=config) @@ -701,18 +691,35 @@ async def test_heartbeat_renews_pending_run_before_task_is_spawned(): original_lease = record.lease_expires_at assert original_lease is not None - # Force a measurable gap so the renewed lease strictly post-dates the - # original — without this the two timestamps land in the same - # microsecond on fast hosts and the strict comparison fails trivially. - await asyncio.sleep(0.001) - store.update_lease = AsyncMock(wraps=store.update_lease) await manager._renew_leases() - store.update_lease.assert_awaited_once() - assert record.lease_expires_at is not None - assert record.lease_expires_at > original_lease + store.update_lease.assert_not_awaited() + assert record.lease_expires_at == original_lease + assert record.status == RunStatus.pending + + +@pytest.mark.anyio +async def test_heartbeat_errors_unattached_pending_after_launch_deadline(): + """An unattached pending run becomes terminal when its initial lease expires.""" + config = _lease_config(lease_seconds=30, heartbeat_enabled=True) + store = MemoryRunStore() + manager = _make_manager(store=store, run_ownership_config=config) + + record = await manager.create_or_reject("thread-1") + record.lease_expires_at = (datetime.now(UTC) - timedelta(seconds=1)).isoformat() + store.update_lease = AsyncMock(wraps=store.update_lease) + + await manager._renew_leases() + + stored = await store.get(record.run_id) + store.update_lease.assert_not_awaited() + assert record.status == RunStatus.error + assert record.abort_event.is_set() + assert stored is not None + assert stored["status"] == "error" + assert "attach a worker" in stored["error"] @pytest.mark.anyio diff --git a/backend/tests/test_run_manager.py b/backend/tests/test_run_manager.py index ac9f0e8eb..9e65c7c09 100644 --- a/backend/tests/test_run_manager.py +++ b/backend/tests/test_run_manager.py @@ -699,6 +699,40 @@ async def test_create_or_reject_does_not_interrupt_old_run_when_new_run_store_wr assert stored_old["status"] == "running" +@pytest.mark.anyio +async def test_create_or_reject_cancellation_after_registration_interrupts_new_run() -> None: + """Cancellation after local admission must not leave the replacement active.""" + store = MemoryRunStore() + manager = RunManager(store=store) + old = await manager.create("thread-1") + await manager.set_status(old.run_id, RunStatus.running) + persist_started = asyncio.Event() + release_persist = asyncio.Event() + original_persist_status = manager._persist_status + + async def blocking_persist_status(record: Any, status: RunStatus, **kwargs: Any) -> None: + persist_started.set() + await release_persist.wait() + await original_persist_status(record, status, **kwargs) + + manager._persist_status = blocking_persist_status + create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt")) + await persist_started.wait() + create_task.cancel() + release_persist.set() + + with pytest.raises(asyncio.CancelledError): + await create_task + + records = await manager.list_by_thread("thread-1") + replacement = next(record for record in records if record.run_id != old.run_id) + stored_replacement = await store.get(replacement.run_id) + assert replacement.status == RunStatus.interrupted + assert replacement.abort_event.is_set() + assert stored_replacement is not None + assert stored_replacement["status"] == "interrupted" + + @pytest.mark.anyio async def test_create_or_reject_rollback_persists_interrupted_status_to_store(): """rollback strategy should persist interrupted status for old runs."""