From 735f67a5b27264d6166b45b56687a9896ad10bc3 Mon Sep 17 00:00:00 2001 From: MiaoRuidx <965742329@qq.com> Date: Sat, 25 Jul 2026 23:50:21 +0800 Subject: [PATCH] fix: guard pending run startup cancellation (#4450) * fix: guard pending run startup cancellation * fix(run): address startup review feedback * fix(run): narrow start_run store contract --------- Co-authored-by: MiaoRuidx <12540796+MiaoRuidx@users.noreply.github.com> Co-authored-by: Willem Jiang --- CHANGELOG.md | 4 + backend/AGENTS.md | 1 + backend/app/gateway/services.py | 185 +++++++++++++----- .../harness/deerflow/persistence/run/sql.py | 14 ++ .../harness/deerflow/runtime/runs/manager.py | 88 ++++++++- .../deerflow/runtime/runs/store/base.py | 8 + .../deerflow/runtime/runs/store/memory.py | 8 + .../harness/deerflow/runtime/runs/worker.py | 32 ++- backend/tests/test_gateway_services.py | 166 +++++++++++++++- backend/tests/test_goal_worker.py | 14 +- backend/tests/test_run_duration_checkpoint.py | 8 +- backend/tests/test_run_manager.py | 60 +++++- backend/tests/test_run_repository.py | 19 ++ backend/tests/test_run_worker_rollback.py | 55 +++++- .../tests/test_worker_langfuse_metadata.py | 5 +- .../test_worker_stream_subgraph_namespace.py | 6 +- 16 files changed, 604 insertions(+), 69 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 86c8b6f6d..4c2ffb08d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -198,6 +198,10 @@ This section accumulates work toward the **2.1.0** milestone ### Fixed +- **runtime:** Thread metadata now switches to `running` only after the run passes + the startup barrier, so pending-cancelled runs no longer briefly project + `running`; clients may observe the prior thread status during worker startup. + ([#4450]) - **runtime:** Re-check orphan candidates through an atomic, lease-aware takeover claim so a successful heartbeat after the scan keeps the run active and only one reconciler reports recovery. ([#4424]) diff --git a/backend/AGENTS.md b/backend/AGENTS.md index e0c7e7f2f..f460fc239 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -450,6 +450,7 @@ metadata only. - `RunManager.get()` is async; direct callers must `await` it. - The history batch helpers `list_successful_regenerate_sources()` and `get_many_by_thread()` default to `user_id=AUTO`: they resolve the request user and fail closed when no user context exists. Migration/admin callers that intentionally need an unscoped read must pass `user_id=None` explicitly. - When a persistent `RunStore` is configured, `get()` and `list_by_thread()` hydrate historical runs from the store. In-memory records win for the same `run_id` so task, abort, and stream-control state stays attached to active local runs. +- Thread metadata status switches to `running` only after `RunManager.try_start()` succeeds. Pending-cancelled runs therefore skip the old `running` projection, while clients may observe the prior thread status during the short worker-startup window. - `cancel()` returns a :class:`~deerflow.runtime.CancelOutcome` enum: `cancelled` (local cancel), `taken_over` (non-owning worker claimed the run because the owner's lease expired — marks it as `error`), `lease_valid_elsewhere` (owner's lease is still alive — caller should return 409 + `Retry-After`), `not_active_locally` (heartbeat disabled, preserving the old 409 path), `not_cancellable` (terminal state), or `unknown` (not found in memory or store). `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persists interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions. - Store-only hydrated runs are readable history. In multi-worker mode with heartbeat enabled, cancel on a store-only run can take over (mark `error`) when the owner's lease has expired past the grace window; otherwise it fails with 409 + `Retry-After`. In single-worker mode (heartbeat off), store-only runs still return 409. - Startup/orphan reconciliation must claim stale active rows with `RunStore.claim_for_takeover()`, not a plain `update_status()`. The final claim re-checks `status` and lease expiry atomically, so a heartbeat renewal between the candidate scan and the recovery write keeps the run active. diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index 15844c6ec..a6d7dab2b 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -41,6 +41,7 @@ from deerflow.runtime import ( CheckpointStateAccessor, ConflictError, DisconnectMode, + RunContext, RunManager, RunRecord, RunStatus, @@ -92,6 +93,8 @@ _TERMINAL_RUN_STATUSES = { RunStatus.interrupted, } +_THREAD_METADATA_SETUP_TIMEOUT_SECONDS = 5.0 + _SERVER_OWNED_MESSAGE_METADATA_KEYS = frozenset( { _DYNAMIC_CONTEXT_REMINDER_KEY, @@ -126,6 +129,51 @@ def _run_is_terminal(record: RunRecord) -> bool: return record.status in _TERMINAL_RUN_STATUSES +def _consume_task_result(task: asyncio.Task) -> None: + """Retrieve a detached task's exception without propagating cancellation.""" + if not task.cancelled(): + task.exception() + + +def _log_thread_metadata_task_result(task: asyncio.Task, *, thread_id: str) -> None: + """Log detached metadata setup failures while ignoring cancellation.""" + if task.cancelled(): + return + try: + task.result() + except asyncio.CancelledError: + return + except Exception: + logger.warning( + "Failed to ensure thread_meta for %s after worker detached (non-fatal)", + sanitize_log_param(thread_id), + exc_info=True, + ) + + +async def _ensure_thread_metadata( + run_ctx: RunContext, + record: RunRecord, + *, + owner_user_id: str | None, +) -> None: + """Ensure an admitted run's thread exists without delaying task attachment.""" + thread_store = run_ctx.thread_store + existing = await thread_store.get(record.thread_id) + if existing is None and owner_user_id: + unscoped = await thread_store.get(record.thread_id, user_id=None) + if unscoped is not None: + if unscoped.get("user_id") != owner_user_id: + await thread_store.update_owner(record.thread_id, owner_user_id, user_id=None) + existing = await thread_store.get(record.thread_id) + if existing is None: + await thread_store.create( + record.thread_id, + assistant_id=record.assistant_id, + metadata=record.metadata, + ) + + async def _terminal_record_stream_missing(bridge: StreamBridge, record: RunRecord) -> bool: """True when a terminal run has no retained stream on bridges that can tell.""" if not _run_is_terminal(record): @@ -979,49 +1027,6 @@ 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)) - 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) @@ -1049,8 +1054,59 @@ async def start_run( request_context=getattr(body, "context", None), ) - task = asyncio.create_task( - run_agent( + async def run_after_metadata(record: RunRecord) -> None: + metadata_task = asyncio.create_task( + _ensure_thread_metadata( + run_ctx, + record, + owner_user_id=owner_user_id, + ) + ) + abort_task = asyncio.create_task(record.abort_event.wait()) + metadata_failure_logged = False + try: + done, _ = await asyncio.wait( + (metadata_task, abort_task), + timeout=_THREAD_METADATA_SETUP_TIMEOUT_SECONDS, + return_when=asyncio.FIRST_COMPLETED, + ) + if metadata_task in done: + try: + metadata_task.result() + except asyncio.CancelledError: + pass + except Exception: + metadata_failure_logged = True + logger.warning( + "Failed to ensure thread_meta for %s (non-fatal)", + sanitize_log_param(thread_id), + exc_info=True, + ) + elif abort_task not in done: + logger.warning( + "Timed out ensuring thread_meta for %s after %.1fs", + sanitize_log_param(thread_id), + _THREAD_METADATA_SETUP_TIMEOUT_SECONDS, + ) + finally: + if metadata_task.done(): + if not metadata_failure_logged: + _log_thread_metadata_task_result(metadata_task, thread_id=thread_id) + else: + metadata_task.cancel() + metadata_task.add_done_callback( + lambda task: _log_thread_metadata_task_result( + task, + thread_id=thread_id, + ) + ) + if not abort_task.done(): + abort_task.cancel() + abort_task.add_done_callback(_consume_task_result) + # Continue through run_agent even after metadata abort/timeout: + # its startup barrier is the single path that turns pending + # cancellation into no-agent-construction plus publish_end. + await run_agent( bridge, run_mgr, record, @@ -1063,8 +1119,43 @@ async def start_run( interrupt_before=body.interrupt_before, interrupt_after=body.interrupt_after, ) - ) - record.task = task + + 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 above 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, + ) + + worker = run_after_metadata(record) + try: + # No await is allowed between durable admission and task + # attachment. Metadata setup runs inside the attached + # worker so a pending cancellation can bypass stalled + # thread-store IO and still reach run_agent's startup + # barrier / stream finalization. + record.task = asyncio.create_task(worker) + except Exception as exc: + worker.close() + await run_mgr.fail_start_if_pending( + record.run_id, + error=f"Failed to attach run worker: {exc}", + ) + raise + 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 # 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/persistence/run/sql.py b/backend/packages/harness/deerflow/persistence/run/sql.py index d364d7258..94c148dc8 100644 --- a/backend/packages/harness/deerflow/persistence/run/sql.py +++ b/backend/packages/harness/deerflow/persistence/run/sql.py @@ -224,6 +224,20 @@ class RunRepository(RunStore): await session.commit() return result.rowcount != 0 + async def start_run(self, run_id: str) -> bool: + """Start only a still-pending run; cancelled rows must not be resurrected.""" + async with self._sf() as session: + result = await session.execute( + update(RunRow) + .where( + RunRow.run_id == run_id, + RunRow.status == "pending", + ) + .values(status="running", updated_at=datetime.now(UTC)) + ) + await session.commit() + return result.rowcount != 0 + async def update_model_name(self, run_id, model_name): async with self._sf() as session: await session.execute(update(RunRow).where(RunRow.run_id == run_id).values(model_name=self._normalize_model_name(model_name), updated_at=datetime.now(UTC))) diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index a52739067..835ed0292 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -167,6 +167,8 @@ class RunRecord: created_at: str = "" updated_at: str = "" task: asyncio.Task | None = field(default=None, repr=False) + # Serializes startup if an admitted run is ever handed to more than one worker path. + start_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) abort_event: asyncio.Event = field(default_factory=asyncio.Event, repr=False) abort_action: str = "interrupt" error: str | None = None @@ -190,6 +192,17 @@ class RunRecord: stop_reason: str | None = None +class RunStartOutcome(StrEnum): + """Result of the pending-to-running startup barrier.""" + + started = "started" + cancelled = "cancelled" + + +class RunStartupError(RuntimeError): + """Raised when durable startup cannot be resolved safely.""" + + OrphanRecoveryCallback = Callable[[list[RunRecord]], Awaitable[None]] @@ -646,6 +659,69 @@ class RunManager: sources.add(source) return sources + async def try_start(self, run_id: str) -> RunStartOutcome: + """Transition an uncancelled pending run to running before building the agent.""" + async with self._lock: + record = self._runs.get(run_id) + if record is None: + raise RunStartupError(f"Cannot start unknown run {run_id}") + + async with record.start_lock: + async with self._lock: + if record.abort_event.is_set() or record.status != RunStatus.pending: + return RunStartOutcome.cancelled + + if self._store is not None: + try: + updated = await self._call_store_with_retry( + "start_run", + run_id, + lambda: self._store.start_run(run_id), + ) + except Exception as exc: + raise RunStartupError(f"Failed to start run {run_id}: {exc}") from exc + if updated is False: + async with self._lock: + if record.status == RunStatus.pending: + record.status = RunStatus.interrupted + record.abort_event.set() + record.updated_at = _now_iso() + return RunStartOutcome.cancelled + + async with self._lock: + if record.abort_event.is_set() or record.status != RunStatus.pending: + restore_status = record.status + restore_error = record.error + restore_stop_reason = record.stop_reason + else: + record.status = RunStatus.running + record.updated_at = _now_iso() + logger.info("Run %s -> %s", run_id, RunStatus.running.value) + return RunStartOutcome.started + + if self._store is not None: + await self._persist_status( + record, + restore_status, + error=restore_error, + stop_reason=restore_stop_reason, + ) + return RunStartOutcome.cancelled + + async def fail_start_if_pending(self, run_id: str, *, error: str) -> bool: + """Mark an admitted run as failed if its worker task could not be attached.""" + async with self._lock: + record = self._runs.get(run_id) + if record is None or record.status != RunStatus.pending: + return False + record.status = RunStatus.error + record.error = error + record.abort_event.set() + record.updated_at = _now_iso() + + await self._persist_status(record, RunStatus.error, error=error) + return True + async def get_many_by_thread( self, thread_id: str, @@ -713,6 +789,7 @@ class RunManager: run_id: str, *, poll_interval: float = 0.01, + abort_event: asyncio.Event | None = None, ) -> None: """Wait until older same-thread runs have finished post-cancel cleanup.""" while True: @@ -729,7 +806,14 @@ class RunManager: if not found_current or not prior_finalizing: return - await asyncio.sleep(poll_interval) + if abort_event is None: + await asyncio.sleep(poll_interval) + continue + try: + await asyncio.wait_for(abort_event.wait(), timeout=poll_interval) + except TimeoutError: + continue + return async def has_later_run(self, thread_id: str, run_id: str) -> bool: """Return whether a newer in-memory run has been admitted for the thread.""" @@ -824,7 +908,7 @@ class RunManager: record.abort_event.set() task_active = record.task is not None and not record.task.done() record.finalizing = task_active - if task_active: + if task_active and record.status == RunStatus.running: record.task.cancel() record.status = RunStatus.interrupted record.updated_at = _now_iso() diff --git a/backend/packages/harness/deerflow/runtime/runs/store/base.py b/backend/packages/harness/deerflow/runtime/runs/store/base.py index 677868ffe..b6796ccec 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/base.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/base.py @@ -95,6 +95,14 @@ class RunStore(abc.ABC): """ pass + @abc.abstractmethod + async def start_run(self, run_id: str) -> bool: + """Atomically transition a pending run to running. + + Returns ``False`` when the row is missing or no longer pending. + """ + pass + @abc.abstractmethod async def delete(self, run_id: str) -> None: pass diff --git a/backend/packages/harness/deerflow/runtime/runs/store/memory.py b/backend/packages/harness/deerflow/runtime/runs/store/memory.py index 0bacd8f7d..eac66fabd 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/memory.py @@ -125,6 +125,14 @@ class MemoryRunStore(RunStore): run["updated_at"] = datetime.now(UTC).isoformat() return True + async def start_run(self, run_id) -> bool: + run = self._runs.get(run_id) + if run is None or run["status"] != "pending": + return False + run["status"] = "running" + run["updated_at"] = datetime.now(UTC).isoformat() + return True + async def update_model_name(self, run_id, model_name): if run_id in self._runs: self._runs[run_id]["model_name"] = model_name diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 731d68972..32da098cc 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -75,7 +75,7 @@ from deerflow.utils.messages import message_to_text from deerflow.workspace_changes import capture_workspace_snapshot, record_workspace_changes from deerflow.workspace_changes.types import WorkspaceSnapshot -from .manager import RunManager, RunRecord +from .manager import RunManager, RunRecord, RunStartOutcome from .naming import resolve_root_run_name from .schemas import RunStatus @@ -417,12 +417,29 @@ async def run_agent( # streaming starts and flushed in the finally block. Pre-bound to None so the # finally is safe even if an exception fires before streaming begins. subagent_events: _SubagentEventBuffer | None = None + started = False try: normalized_stream_modes = normalize_stream_modes(stream_modes) requested_modes: set[str] = set(normalized_stream_modes) lg_modes = to_langgraph_stream_modes(normalized_stream_modes) - await run_manager.wait_for_prior_finalizing(thread_id, run_id) + await run_manager.wait_for_prior_finalizing( + thread_id, + run_id, + abort_event=record.abort_event, + ) + + start_outcome = await run_manager.try_start(run_id) + if start_outcome is not RunStartOutcome.started: + return + started = True + + if thread_store is not None: + try: + await thread_store.update_status(thread_id, "running") + except Exception: + logger.debug("Failed to update thread_meta status for %s (non-fatal)", thread_id) + mode = ctx.checkpoint_channel_mode inject_checkpoint_mode(config, mode) checkpoint_config = { @@ -472,9 +489,6 @@ async def run_agent( progress_reporter=lambda snapshot: run_manager.update_run_progress(run_id, **snapshot), ) - # 1. Mark running - await run_manager.set_status(run_id, RunStatus.running) - if event_store is not None: workspace_changes_user_id = get_effective_user_id() try: @@ -822,7 +836,7 @@ async def run_agent( except Exception: logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True) - if checkpointer is not None and record.status == RunStatus.interrupted: + if started and checkpointer is not None and record.status == RunStatus.interrupted: try: await run_manager.wait_for_prior_finalizing(thread_id, run_id) if not await run_manager.has_later_started_run(thread_id, run_id): @@ -831,7 +845,7 @@ async def run_agent( logger.debug("Failed to generate interrupted title for thread %s (non-fatal)", thread_id) # Sync title from checkpoint to threads_meta.display_name - if checkpointer is not None and thread_store is not None: + if started and checkpointer is not None and thread_store is not None: try: ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} ckpt_tuple = await checkpointer.aget_tuple(ckpt_config) @@ -845,7 +859,7 @@ async def run_agent( # Persist run duration to checkpoint metadata so history reads # don't need to correlate runs and events. - if checkpointer is not None and record.status == RunStatus.success: + if started and checkpointer is not None and record.status == RunStatus.success: try: created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00")) updated = datetime.fromisoformat(record.updated_at.replace("Z", "+00:00")) @@ -863,7 +877,7 @@ async def run_agent( logger.debug("Failed to persist run duration for thread %s run %s (non-fatal)", thread_id, run_id) # Update threads_meta status based on run outcome - if thread_store is not None: + if started and thread_store is not None: try: final_status = "idle" if record.status == RunStatus.success else record.status.value await thread_store.update_status(thread_id, final_status) diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index e237ce9e0..457bb5059 100644 --- a/backend/tests/test_gateway_services.py +++ b/backend/tests/test_gateway_services.py @@ -2,7 +2,10 @@ from __future__ import annotations +import asyncio import json +import logging +from types import SimpleNamespace import pytest @@ -18,6 +21,39 @@ def _stub_app_config(): reset_app_config() +def _make_start_run_request(run_manager, *, thread_store=None, auth_source=None): + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.store.memory import InMemoryStore + + from deerflow.persistence.thread_meta.memory import MemoryThreadMetaStore + + store = InMemoryStore() + return SimpleNamespace( + headers={}, + state=SimpleNamespace(auth_source=auth_source), + app=SimpleNamespace( + state=SimpleNamespace( + stream_bridge=SimpleNamespace(), + run_manager=run_manager, + checkpointer=InMemorySaver(), + store=store, + run_event_store=SimpleNamespace(), + run_events_config=None, + thread_store=thread_store or MemoryThreadMetaStore(store), + ) + ), + ) + + +def _run_create_request(content="hello", **kwargs): + from app.gateway.routers.thread_runs import RunCreateRequest + + return RunCreateRequest( + input={"messages": [{"role": "user", "content": content}]}, + **kwargs, + ) + + def test_format_sse_basic(): from app.gateway.services import format_sse @@ -797,6 +833,130 @@ def test_apply_checkpoint_to_run_config_rejects_missing_checkpoint(): assert "missing" in exc.value.detail +@pytest.mark.asyncio +async def test_start_run_checkpoint_validation_failure_does_not_admit_run(_stub_app_config): + from unittest.mock import patch + + from fastapi import HTTPException + + from app.gateway.services import start_run + from deerflow.runtime import RunManager + from deerflow.runtime.runs.store.memory import MemoryRunStore + + thread_id = "thread-invalid-checkpoint" + run_store = MemoryRunStore() + run_manager = RunManager(store=run_store) + request = _make_start_run_request(run_manager) + invalid_body = _run_create_request( + checkpoint_id="missing-checkpoint", + ) + + with ( + patch("app.gateway.services.resolve_agent_factory", return_value=object()), + pytest.raises(HTTPException, match="Checkpoint missing-checkpoint not found"), + ): + await start_run(invalid_body, thread_id, request) + + assert await run_manager.list_by_thread(thread_id, user_id=None) == [] + assert await run_store.list_by_thread(thread_id, user_id=None) == [] + + +@pytest.mark.asyncio +async def test_pending_cancel_bypasses_thread_metadata_and_logs_failure(_stub_app_config, caplog): + from unittest.mock import AsyncMock, patch + + from app.gateway.services import start_run + from deerflow.runtime import RunManager + from deerflow.runtime.runs.store.memory import MemoryRunStore + + metadata_started = asyncio.Event() + + async def get_thread(_thread_id): + metadata_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError as exc: + raise RuntimeError("thread metadata store failed after cancellation") from exc + + async def fake_run_agent(*_args, **_kwargs): + return None + + run_manager = RunManager(store=MemoryRunStore()) + body = _run_create_request() + request = _make_start_run_request( + run_manager, + thread_store=SimpleNamespace( + get=AsyncMock(side_effect=get_thread), + create=AsyncMock(), + update_owner=AsyncMock(), + ), + ) + caplog.set_level(logging.WARNING, logger="app.gateway.services") + 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-cancel-log-meta", request) + await asyncio.wait_for(metadata_started.wait(), timeout=1) + assert record.task is not None + await run_manager.cancel(record.run_id) + await asyncio.wait_for(record.task, timeout=1) + await asyncio.sleep(0) + + assert "thread metadata store failed after cancellation" in caplog.text + + +@pytest.mark.asyncio +async def test_thread_metadata_timeout_logs_and_run_still_starts(_stub_app_config, caplog, monkeypatch): + from unittest.mock import AsyncMock, patch + + import app.gateway.services as services + from app.gateway.services import start_run + from deerflow.runtime import RunManager + from deerflow.runtime.runs.manager import RunStartOutcome + from deerflow.runtime.runs.schemas import RunStatus + from deerflow.runtime.runs.store.memory import MemoryRunStore + + metadata_started = asyncio.Event() + run_agent_called = asyncio.Event() + + async def get_thread(_thread_id): + metadata_started.set() + await asyncio.Event().wait() + + async def fake_run_agent(_bridge, run_manager, record, **_kwargs): + run_agent_called.set() + start_outcome = await run_manager.try_start(record.run_id) + assert start_outcome is RunStartOutcome.started + + monkeypatch.setattr(services, "_THREAD_METADATA_SETUP_TIMEOUT_SECONDS", 0.01) + run_manager = RunManager(store=MemoryRunStore()) + body = _run_create_request() + request = _make_start_run_request( + run_manager, + thread_store=SimpleNamespace( + get=AsyncMock(side_effect=get_thread), + create=AsyncMock(), + update_owner=AsyncMock(), + ), + ) + caplog.set_level(logging.WARNING, logger="app.gateway.services") + + 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-timeout-meta", request) + await asyncio.wait_for(metadata_started.wait(), timeout=1) + assert record.task is not None + await asyncio.wait_for(record.task, timeout=1) + + assert run_agent_called.is_set() + assert record.status == RunStatus.running + assert (await run_manager.get(record.run_id)).status == RunStatus.running + assert "Timed out ensuring thread_meta for thread-timeout-meta" in caplog.text + + def test_context_merges_into_configurable(): """Context values must be merged into config['configurable'] by start_run. @@ -1844,7 +2004,7 @@ async def test_run_agent_full_mode_rejects_delta_before_graph_invocation(): CHECKPOINT_MODE_METADATA_KEY, INTERNAL_CHECKPOINT_MODE_KEY, ) - from deerflow.runtime.runs.manager import RunRecord + from deerflow.runtime.runs.manager import RunRecord, RunStartOutcome from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus from deerflow.runtime.runs.worker import RunContext, run_agent @@ -1859,6 +2019,7 @@ async def test_run_agent_full_mode_rejects_delta_before_graph_invocation(): cleanup=AsyncMock(), ) run_manager = SimpleNamespace( + try_start=AsyncMock(return_value=RunStartOutcome.started), wait_for_prior_finalizing=AsyncMock(), set_status=AsyncMock(), ) @@ -1909,7 +2070,7 @@ async def test_run_agent_full_mode_checks_selected_checkpoint_before_graph(): from unittest.mock import AsyncMock, MagicMock, call from deerflow.runtime.checkpoint_mode import CHECKPOINT_MODE_METADATA_KEY - from deerflow.runtime.runs.manager import RunRecord + from deerflow.runtime.runs.manager import RunRecord, RunStartOutcome from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus from deerflow.runtime.runs.worker import RunContext, run_agent @@ -1930,6 +2091,7 @@ async def test_run_agent_full_mode_checks_selected_checkpoint_before_graph(): cleanup=AsyncMock(), ) run_manager = SimpleNamespace( + try_start=AsyncMock(return_value=RunStartOutcome.started), wait_for_prior_finalizing=AsyncMock(), set_status=AsyncMock(), ) diff --git a/backend/tests/test_goal_worker.py b/backend/tests/test_goal_worker.py index 68f73a7c4..d42312ebb 100644 --- a/backend/tests/test_goal_worker.py +++ b/backend/tests/test_goal_worker.py @@ -9,7 +9,7 @@ from langgraph.checkpoint.memory import InMemorySaver from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph from deerflow.runtime.goal import GoalEvaluation, attach_goal_evaluation, build_goal_state, latest_visible_assistant_signature, read_thread_goal, write_thread_goal from deerflow.runtime.runs import worker -from deerflow.runtime.runs.manager import RunRecord +from deerflow.runtime.runs.manager import RunRecord, RunStartOutcome from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus @@ -606,6 +606,10 @@ async def test_run_agent_does_not_stream_continuation_after_abort(monkeypatch): return _gen() class FakeRunManager: + async def try_start(self, _run_id): + record.status = RunStatus.running + return RunStartOutcome.started + async def set_status(self, _run_id, status, **_kwargs): record.status = status @@ -683,6 +687,10 @@ async def test_run_agent_reuses_goal_evaluator_model_for_goal_loop(monkeypatch): return _gen() class FakeRunManager: + async def try_start(self, _run_id): + record.status = RunStatus.running + return RunStartOutcome.started + async def set_status(self, _run_id, status, **_kwargs): record.status = status @@ -859,6 +867,10 @@ async def test_run_agent_strips_branch_checkpoint_for_goal_continuation(monkeypa return _gen() class FakeRunManager: + async def try_start(self, _run_id): + record.status = RunStatus.running + return RunStartOutcome.started + async def set_status(self, _run_id, status, **_kwargs): record.status = status diff --git a/backend/tests/test_run_duration_checkpoint.py b/backend/tests/test_run_duration_checkpoint.py index 1a6948281..7119de203 100644 --- a/backend/tests/test_run_duration_checkpoint.py +++ b/backend/tests/test_run_duration_checkpoint.py @@ -10,7 +10,8 @@ from langgraph.checkpoint.memory import InMemorySaver import deerflow.runtime.runs.worker as worker from deerflow.runtime.goal import goal_thread_lock -from deerflow.runtime.runs.manager import RunManager +from deerflow.runtime.runs.manager import RunManager, RunStartOutcome +from deerflow.runtime.runs.schemas import RunStatus from deerflow.runtime.runs.worker import RunContext, _persist_run_duration, run_agent @@ -367,6 +368,11 @@ async def test_successful_subsecond_run_persists_zero_duration(monkeypatch: pyte async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): yield {"messages": []} + async def try_start(run_id): + record.status = RunStatus.running + return RunStartOutcome.started + + monkeypatch.setattr(run_manager, "try_start", try_start) monkeypatch.setattr(run_manager, "set_status", set_status) monkeypatch.setattr(worker, "_persist_run_duration", persist_duration) diff --git a/backend/tests/test_run_manager.py b/backend/tests/test_run_manager.py index 13ad3d44c..68e66fddb 100644 --- a/backend/tests/test_run_manager.py +++ b/backend/tests/test_run_manager.py @@ -11,7 +11,7 @@ from sqlalchemy.exc import DatabaseError as SQLAlchemyDatabaseError from deerflow.config.run_ownership_config import RunOwnershipConfig from deerflow.runtime import DisconnectMode, RunManager, RunStatus, ThreadOperationKind -from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, PersistenceRetryPolicy +from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, PersistenceRetryPolicy, RunStartOutcome from deerflow.runtime.runs.store.memory import MemoryRunStore ISO_RE = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}") @@ -384,6 +384,64 @@ async def test_status_persistence_does_not_retry_permanent_sqlalchemy_errors(): assert store.status_update_attempts == 1 +@pytest.mark.anyio +async def test_try_start_respects_durable_and_racing_cancels(): + """Startup must not resurrect durable or locally racing cancels.""" + store = MemoryRunStore() + manager = RunManager(store=store) + record = await manager.create_or_reject("thread-1") + await store.update_status(record.run_id, RunStatus.interrupted.value) + + assert await manager.try_start(record.run_id) == RunStartOutcome.cancelled + assert record.status == RunStatus.interrupted + assert (await store.get(record.run_id))["status"] == RunStatus.interrupted.value + + record = await manager.create_or_reject("thread-2") + original_start_run = store.start_run + + async def start_then_cancel(run_id): + updated = await original_start_run(run_id) + await manager.cancel(record.run_id) + return updated + + store.start_run = start_then_cancel + + assert await manager.try_start(record.run_id) == RunStartOutcome.cancelled + assert record.status == RunStatus.interrupted + assert (await store.get(record.run_id))["status"] == RunStatus.interrupted.value + + +@pytest.mark.anyio +async def test_fail_start_if_pending_marks_pending_run_error_and_persists(): + """Worker attach failures should finalize only runs still pending startup.""" + store = MemoryRunStore() + manager = RunManager(store=store) + record = await manager.create_or_reject("thread-1") + error = "Failed to attach run worker: boom" + + assert await manager.fail_start_if_pending(record.run_id, error=error) is True + + stored = await store.get(record.run_id) + assert record.status == RunStatus.error + assert record.error == error + assert record.abort_event.is_set() + assert stored is not None + assert stored["status"] == RunStatus.error.value + assert stored["error"] == error + + running = await manager.create_or_reject("thread-2") + assert await manager.try_start(running.run_id) == RunStartOutcome.started + + assert await manager.fail_start_if_pending(running.run_id, error="late") is False + + stored_running = await store.get(running.run_id) + assert running.status == RunStatus.running + assert running.error is None + assert stored_running is not None + assert stored_running["status"] == RunStatus.running.value + assert stored_running["error"] is None + + @pytest.mark.anyio async def test_completion_persistence_recreates_missing_store_row(): """Completion updates should recreate a missing row and persist final counters.""" diff --git a/backend/tests/test_run_repository.py b/backend/tests/test_run_repository.py index 95952d463..d1d9e0a83 100644 --- a/backend/tests/test_run_repository.py +++ b/backend/tests/test_run_repository.py @@ -44,6 +44,9 @@ class _CustomRunStoreWithoutProgress(RunStore): async def update_status(self, *args, **kwargs): return None + async def start_run(self, *args, **kwargs): + return False + async def delete(self, *args, **kwargs): return None @@ -153,6 +156,22 @@ class TestRunRepository: assert updated is False await _cleanup() + @pytest.mark.anyio + async def test_start_run_only_updates_pending_rows(self, tmp_path): + repo = await _make_repo(tmp_path) + await repo.put("pending-run", thread_id="t1", status="pending") + await repo.put("cancelled-run", thread_id="t2", status="pending") + await repo.update_status("cancelled-run", "interrupted") + + assert await repo.start_run("pending-run") is True + assert await repo.start_run("cancelled-run") is False + + pending_row = await repo.get("pending-run") + cancelled_row = await repo.get("cancelled-run") + assert pending_row["status"] == "running" + assert cancelled_row["status"] == "interrupted" + await _cleanup() + @pytest.mark.anyio async def test_update_status_with_error(self, tmp_path): repo = await _make_repo(tmp_path) diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index 9082f2c20..dfd1aa493 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -18,7 +18,7 @@ from langgraph.types import Overwrite from deerflow.agents.thread_state import merge_artifacts, merge_message_writes from deerflow.runtime.checkpoint_state import CheckpointStateAccessor from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY -from deerflow.runtime.runs.manager import ConflictError, RunManager +from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, RunManager from deerflow.runtime.runs.schemas import RunStatus from deerflow.runtime.runs.worker import ( RollbackPoint, @@ -45,6 +45,48 @@ class FakeCheckpointer: self.aput_writes = AsyncMock() +@pytest.mark.anyio +async def test_pending_cancel_stops_waiting_for_prior_finalization(): + run_manager = RunManager() + prior = await run_manager.create("thread-cancel-while-waiting") + prior.status = RunStatus.interrupted + prior.finalizing = True + record = await run_manager.create("thread-cancel-while-waiting") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + factory_called = False + + def agent_factory(**_kwargs): + nonlocal factory_called + factory_called = True + raise AssertionError("cancelled pending run must not build an agent") + + record.task = asyncio.create_task( + run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=agent_factory, + graph_input={"messages": []}, + config={}, + ) + ) + await asyncio.sleep(0) + + outcome = await run_manager.cancel(record.run_id) + await asyncio.wait_for(record.task, timeout=0.2) + + assert outcome == CancelOutcome.cancelled + assert prior.finalizing is True + assert factory_called is False + assert record.status == RunStatus.interrupted + bridge.publish_end.assert_awaited_once_with(record.run_id) + + def _make_rollback_point(*, checkpoint_id="ckpt-1", messages=("before",), pending_writes=()): return RollbackPoint( config={ @@ -639,8 +681,6 @@ async def test_run_agent_marks_rollback_unusable_when_capture_fails(): """ run_manager = RunManager() record = await run_manager.create("thread-1") - record.abort_action = "rollback" - record.abort_event.set() bridge = SimpleNamespace( publish=AsyncMock(), publish_end=AsyncMock(), @@ -664,6 +704,10 @@ async def test_run_agent_marks_rollback_unusable_when_capture_fails(): class DummyAgent: async def aget_state(self, _config): + # Cancel after the worker crosses its startup barrier so this test + # exercises the running rollback path, not pending cancellation. + record.abort_action = "rollback" + record.abort_event.set() raise RuntimeError("materialization failed") async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): @@ -2729,14 +2773,16 @@ async def test_worker_finally_block_swallows_helper_exceptions(monkeypatch): """ import deerflow.runtime.runs.worker as worker_module + helper_called = asyncio.Event() + async def _boom(*_args, **_kwargs): + helper_called.set() raise RuntimeError("forced helper failure") monkeypatch.setattr(worker_module, "_ensure_interrupted_title", _boom) run_manager = RunManager() record = await run_manager.create("thread-1") - record.status = RunStatus.interrupted bridge = SimpleNamespace( publish=AsyncMock(), @@ -2792,5 +2838,6 @@ async def test_worker_finally_block_swallows_helper_exceptions(monkeypatch): # The helper raised, but the run still reaches the threads_meta status sync # and ``publish_end`` — i.e. the SSE stream is closed cleanly and the row # reflects the run outcome. + assert helper_called.is_set() assert captured_status.get("status") == ("thread-1", "interrupted") bridge.publish_end.assert_awaited_once_with(record.run_id) diff --git a/backend/tests/test_worker_langfuse_metadata.py b/backend/tests/test_worker_langfuse_metadata.py index ee820cb1d..346e11e45 100644 --- a/backend/tests/test_worker_langfuse_metadata.py +++ b/backend/tests/test_worker_langfuse_metadata.py @@ -11,7 +11,7 @@ import asyncio import pytest -from deerflow.runtime.runs.manager import RunRecord +from deerflow.runtime.runs.manager import RunRecord, RunStartOutcome from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus from deerflow.runtime.runs.worker import RunContext, run_agent from deerflow.trace_context import ( @@ -42,6 +42,9 @@ class _FakeAgent: class _FakeRunManager: + async def try_start(self, _run_id: str) -> RunStartOutcome: + return RunStartOutcome.started + async def wait_for_prior_finalizing(self, *_args, **_kwargs) -> None: return None diff --git a/backend/tests/test_worker_stream_subgraph_namespace.py b/backend/tests/test_worker_stream_subgraph_namespace.py index ce8f6b763..0c1c82046 100644 --- a/backend/tests/test_worker_stream_subgraph_namespace.py +++ b/backend/tests/test_worker_stream_subgraph_namespace.py @@ -23,7 +23,7 @@ from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from packaging.version import Version from deerflow.runtime.runs import worker -from deerflow.runtime.runs.manager import RunRecord +from deerflow.runtime.runs.manager import RunRecord, RunStartOutcome from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus from deerflow.runtime.runs.worker import ( _compose_sse_event, @@ -316,6 +316,10 @@ class _IntegrationRunManager: def __init__(self, record: RunRecord) -> None: self._record = record + async def try_start(self, _run_id): + self._record.status = RunStatus.running + return RunStartOutcome.started + async def wait_for_prior_finalizing(self, *_args, **_kwargs): return None