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 <willem.jiang@gmail.com>
This commit is contained in:
MiaoRuidx 2026-07-25 23:50:21 +08:00 committed by GitHub
parent 8af760fc30
commit 735f67a5b2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 604 additions and 69 deletions

View File

@ -198,6 +198,10 @@ This section accumulates work toward the **2.1.0** milestone
### Fixed ### 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 - **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 claim so a successful heartbeat after the scan keeps the run active and only
one reconciler reports recovery. ([#4424]) one reconciler reports recovery. ([#4424])

View File

@ -450,6 +450,7 @@ metadata only.
- `RunManager.get()` is async; direct callers must `await` it. - `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. - 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. - 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. - `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. - 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. - 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.

View File

@ -41,6 +41,7 @@ from deerflow.runtime import (
CheckpointStateAccessor, CheckpointStateAccessor,
ConflictError, ConflictError,
DisconnectMode, DisconnectMode,
RunContext,
RunManager, RunManager,
RunRecord, RunRecord,
RunStatus, RunStatus,
@ -92,6 +93,8 @@ _TERMINAL_RUN_STATUSES = {
RunStatus.interrupted, RunStatus.interrupted,
} }
_THREAD_METADATA_SETUP_TIMEOUT_SECONDS = 5.0
_SERVER_OWNED_MESSAGE_METADATA_KEYS = frozenset( _SERVER_OWNED_MESSAGE_METADATA_KEYS = frozenset(
{ {
_DYNAMIC_CONTEXT_REMINDER_KEY, _DYNAMIC_CONTEXT_REMINDER_KEY,
@ -126,6 +129,51 @@ def _run_is_terminal(record: RunRecord) -> bool:
return record.status in _TERMINAL_RUN_STATUSES 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: 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.""" """True when a terminal run has no retained stream on bridges that can tell."""
if not _run_is_terminal(record): 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 owner_context_token = set_current_user(SimpleNamespace(id=owner_user_id)) if owner_user_id else None
try: 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) agent_factory = resolve_agent_factory(body.assistant_id)
is_internal_caller = getattr(getattr(request, "state", None), "auth_source", None) == AUTH_SOURCE_INTERNAL is_internal_caller = getattr(getattr(request, "state", None), "auth_source", None) == AUTH_SOURCE_INTERNAL
command = getattr(body, "command", None) command = getattr(body, "command", None)
@ -1049,8 +1054,59 @@ async def start_run(
request_context=getattr(body, "context", None), request_context=getattr(body, "context", None),
) )
task = asyncio.create_task( async def run_after_metadata(record: RunRecord) -> None:
run_agent( 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, bridge,
run_mgr, run_mgr,
record, record,
@ -1063,8 +1119,43 @@ async def start_run(
interrupt_before=body.interrupt_before, interrupt_before=body.interrupt_before,
interrupt_after=body.interrupt_after, 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 sync is handled by worker.py's finally block which reads the
# title from the checkpoint and calls thread_store.update_display_name # title from the checkpoint and calls thread_store.update_display_name

View File

@ -224,6 +224,20 @@ class RunRepository(RunStore):
await session.commit() await session.commit()
return result.rowcount != 0 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 def update_model_name(self, run_id, model_name):
async with self._sf() as session: 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))) 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)))

View File

@ -167,6 +167,8 @@ class RunRecord:
created_at: str = "" created_at: str = ""
updated_at: str = "" updated_at: str = ""
task: asyncio.Task | None = field(default=None, repr=False) 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_event: asyncio.Event = field(default_factory=asyncio.Event, repr=False)
abort_action: str = "interrupt" abort_action: str = "interrupt"
error: str | None = None error: str | None = None
@ -190,6 +192,17 @@ class RunRecord:
stop_reason: str | None = None 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]] OrphanRecoveryCallback = Callable[[list[RunRecord]], Awaitable[None]]
@ -646,6 +659,69 @@ class RunManager:
sources.add(source) sources.add(source)
return sources 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( async def get_many_by_thread(
self, self,
thread_id: str, thread_id: str,
@ -713,6 +789,7 @@ class RunManager:
run_id: str, run_id: str,
*, *,
poll_interval: float = 0.01, poll_interval: float = 0.01,
abort_event: asyncio.Event | None = None,
) -> None: ) -> None:
"""Wait until older same-thread runs have finished post-cancel cleanup.""" """Wait until older same-thread runs have finished post-cancel cleanup."""
while True: while True:
@ -729,7 +806,14 @@ class RunManager:
if not found_current or not prior_finalizing: if not found_current or not prior_finalizing:
return 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: 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.""" """Return whether a newer in-memory run has been admitted for the thread."""
@ -824,7 +908,7 @@ class RunManager:
record.abort_event.set() record.abort_event.set()
task_active = record.task is not None and not record.task.done() task_active = record.task is not None and not record.task.done()
record.finalizing = task_active record.finalizing = task_active
if task_active: if task_active and record.status == RunStatus.running:
record.task.cancel() record.task.cancel()
record.status = RunStatus.interrupted record.status = RunStatus.interrupted
record.updated_at = _now_iso() record.updated_at = _now_iso()

View File

@ -95,6 +95,14 @@ class RunStore(abc.ABC):
""" """
pass 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 @abc.abstractmethod
async def delete(self, run_id: str) -> None: async def delete(self, run_id: str) -> None:
pass pass

View File

@ -125,6 +125,14 @@ class MemoryRunStore(RunStore):
run["updated_at"] = datetime.now(UTC).isoformat() run["updated_at"] = datetime.now(UTC).isoformat()
return True 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): async def update_model_name(self, run_id, model_name):
if run_id in self._runs: if run_id in self._runs:
self._runs[run_id]["model_name"] = model_name self._runs[run_id]["model_name"] = model_name

View File

@ -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 import capture_workspace_snapshot, record_workspace_changes
from deerflow.workspace_changes.types import WorkspaceSnapshot 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 .naming import resolve_root_run_name
from .schemas import RunStatus 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 # 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. # finally is safe even if an exception fires before streaming begins.
subagent_events: _SubagentEventBuffer | None = None subagent_events: _SubagentEventBuffer | None = None
started = False
try: try:
normalized_stream_modes = normalize_stream_modes(stream_modes) normalized_stream_modes = normalize_stream_modes(stream_modes)
requested_modes: set[str] = set(normalized_stream_modes) requested_modes: set[str] = set(normalized_stream_modes)
lg_modes = to_langgraph_stream_modes(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 mode = ctx.checkpoint_channel_mode
inject_checkpoint_mode(config, mode) inject_checkpoint_mode(config, mode)
checkpoint_config = { checkpoint_config = {
@ -472,9 +489,6 @@ async def run_agent(
progress_reporter=lambda snapshot: run_manager.update_run_progress(run_id, **snapshot), 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: if event_store is not None:
workspace_changes_user_id = get_effective_user_id() workspace_changes_user_id = get_effective_user_id()
try: try:
@ -822,7 +836,7 @@ async def run_agent(
except Exception: except Exception:
logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True) 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: try:
await run_manager.wait_for_prior_finalizing(thread_id, run_id) await run_manager.wait_for_prior_finalizing(thread_id, run_id)
if not await run_manager.has_later_started_run(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) logger.debug("Failed to generate interrupted title for thread %s (non-fatal)", thread_id)
# Sync title from checkpoint to threads_meta.display_name # 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: try:
ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
ckpt_tuple = await checkpointer.aget_tuple(ckpt_config) 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 # Persist run duration to checkpoint metadata so history reads
# don't need to correlate runs and events. # 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: try:
created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00")) created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00"))
updated = datetime.fromisoformat(record.updated_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) 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 # Update threads_meta status based on run outcome
if thread_store is not None: if started and thread_store is not None:
try: try:
final_status = "idle" if record.status == RunStatus.success else record.status.value final_status = "idle" if record.status == RunStatus.success else record.status.value
await thread_store.update_status(thread_id, final_status) await thread_store.update_status(thread_id, final_status)

View File

@ -2,7 +2,10 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import json import json
import logging
from types import SimpleNamespace
import pytest import pytest
@ -18,6 +21,39 @@ def _stub_app_config():
reset_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(): def test_format_sse_basic():
from app.gateway.services import format_sse 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 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(): def test_context_merges_into_configurable():
"""Context values must be merged into config['configurable'] by start_run. """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, CHECKPOINT_MODE_METADATA_KEY,
INTERNAL_CHECKPOINT_MODE_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.schemas import DisconnectMode, RunStatus
from deerflow.runtime.runs.worker import RunContext, run_agent 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(), cleanup=AsyncMock(),
) )
run_manager = SimpleNamespace( run_manager = SimpleNamespace(
try_start=AsyncMock(return_value=RunStartOutcome.started),
wait_for_prior_finalizing=AsyncMock(), wait_for_prior_finalizing=AsyncMock(),
set_status=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 unittest.mock import AsyncMock, MagicMock, call
from deerflow.runtime.checkpoint_mode import CHECKPOINT_MODE_METADATA_KEY 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.schemas import DisconnectMode, RunStatus
from deerflow.runtime.runs.worker import RunContext, run_agent 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(), cleanup=AsyncMock(),
) )
run_manager = SimpleNamespace( run_manager = SimpleNamespace(
try_start=AsyncMock(return_value=RunStartOutcome.started),
wait_for_prior_finalizing=AsyncMock(), wait_for_prior_finalizing=AsyncMock(),
set_status=AsyncMock(), set_status=AsyncMock(),
) )

View File

@ -9,7 +9,7 @@ from langgraph.checkpoint.memory import InMemorySaver
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph 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.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 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.schemas import DisconnectMode, RunStatus
@ -606,6 +606,10 @@ async def test_run_agent_does_not_stream_continuation_after_abort(monkeypatch):
return _gen() return _gen()
class FakeRunManager: 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): async def set_status(self, _run_id, status, **_kwargs):
record.status = status record.status = status
@ -683,6 +687,10 @@ async def test_run_agent_reuses_goal_evaluator_model_for_goal_loop(monkeypatch):
return _gen() return _gen()
class FakeRunManager: 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): async def set_status(self, _run_id, status, **_kwargs):
record.status = status record.status = status
@ -859,6 +867,10 @@ async def test_run_agent_strips_branch_checkpoint_for_goal_continuation(monkeypa
return _gen() return _gen()
class FakeRunManager: 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): async def set_status(self, _run_id, status, **_kwargs):
record.status = status record.status = status

View File

@ -10,7 +10,8 @@ from langgraph.checkpoint.memory import InMemorySaver
import deerflow.runtime.runs.worker as worker import deerflow.runtime.runs.worker as worker
from deerflow.runtime.goal import goal_thread_lock 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 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): async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
yield {"messages": []} 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(run_manager, "set_status", set_status)
monkeypatch.setattr(worker, "_persist_run_duration", persist_duration) monkeypatch.setattr(worker, "_persist_run_duration", persist_duration)

View File

@ -11,7 +11,7 @@ from sqlalchemy.exc import DatabaseError as SQLAlchemyDatabaseError
from deerflow.config.run_ownership_config import RunOwnershipConfig from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime import DisconnectMode, RunManager, RunStatus, ThreadOperationKind 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 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}") 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 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 @pytest.mark.anyio
async def test_completion_persistence_recreates_missing_store_row(): async def test_completion_persistence_recreates_missing_store_row():
"""Completion updates should recreate a missing row and persist final counters.""" """Completion updates should recreate a missing row and persist final counters."""

View File

@ -44,6 +44,9 @@ class _CustomRunStoreWithoutProgress(RunStore):
async def update_status(self, *args, **kwargs): async def update_status(self, *args, **kwargs):
return None return None
async def start_run(self, *args, **kwargs):
return False
async def delete(self, *args, **kwargs): async def delete(self, *args, **kwargs):
return None return None
@ -153,6 +156,22 @@ class TestRunRepository:
assert updated is False assert updated is False
await _cleanup() 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 @pytest.mark.anyio
async def test_update_status_with_error(self, tmp_path): async def test_update_status_with_error(self, tmp_path):
repo = await _make_repo(tmp_path) repo = await _make_repo(tmp_path)

View File

@ -18,7 +18,7 @@ from langgraph.types import Overwrite
from deerflow.agents.thread_state import merge_artifacts, merge_message_writes from deerflow.agents.thread_state import merge_artifacts, merge_message_writes
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY 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.schemas import RunStatus
from deerflow.runtime.runs.worker import ( from deerflow.runtime.runs.worker import (
RollbackPoint, RollbackPoint,
@ -45,6 +45,48 @@ class FakeCheckpointer:
self.aput_writes = AsyncMock() 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=()): def _make_rollback_point(*, checkpoint_id="ckpt-1", messages=("before",), pending_writes=()):
return RollbackPoint( return RollbackPoint(
config={ config={
@ -639,8 +681,6 @@ async def test_run_agent_marks_rollback_unusable_when_capture_fails():
""" """
run_manager = RunManager() run_manager = RunManager()
record = await run_manager.create("thread-1") record = await run_manager.create("thread-1")
record.abort_action = "rollback"
record.abort_event.set()
bridge = SimpleNamespace( bridge = SimpleNamespace(
publish=AsyncMock(), publish=AsyncMock(),
publish_end=AsyncMock(), publish_end=AsyncMock(),
@ -664,6 +704,10 @@ async def test_run_agent_marks_rollback_unusable_when_capture_fails():
class DummyAgent: class DummyAgent:
async def aget_state(self, _config): 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") raise RuntimeError("materialization failed")
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): 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 import deerflow.runtime.runs.worker as worker_module
helper_called = asyncio.Event()
async def _boom(*_args, **_kwargs): async def _boom(*_args, **_kwargs):
helper_called.set()
raise RuntimeError("forced helper failure") raise RuntimeError("forced helper failure")
monkeypatch.setattr(worker_module, "_ensure_interrupted_title", _boom) monkeypatch.setattr(worker_module, "_ensure_interrupted_title", _boom)
run_manager = RunManager() run_manager = RunManager()
record = await run_manager.create("thread-1") record = await run_manager.create("thread-1")
record.status = RunStatus.interrupted
bridge = SimpleNamespace( bridge = SimpleNamespace(
publish=AsyncMock(), 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 # 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 # and ``publish_end`` — i.e. the SSE stream is closed cleanly and the row
# reflects the run outcome. # reflects the run outcome.
assert helper_called.is_set()
assert captured_status.get("status") == ("thread-1", "interrupted") assert captured_status.get("status") == ("thread-1", "interrupted")
bridge.publish_end.assert_awaited_once_with(record.run_id) bridge.publish_end.assert_awaited_once_with(record.run_id)

View File

@ -11,7 +11,7 @@ import asyncio
import pytest 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.schemas import DisconnectMode, RunStatus
from deerflow.runtime.runs.worker import RunContext, run_agent from deerflow.runtime.runs.worker import RunContext, run_agent
from deerflow.trace_context import ( from deerflow.trace_context import (
@ -42,6 +42,9 @@ class _FakeAgent:
class _FakeRunManager: class _FakeRunManager:
async def try_start(self, _run_id: str) -> RunStartOutcome:
return RunStartOutcome.started
async def wait_for_prior_finalizing(self, *_args, **_kwargs) -> None: async def wait_for_prior_finalizing(self, *_args, **_kwargs) -> None:
return None return None

View File

@ -23,7 +23,7 @@ from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from packaging.version import Version from packaging.version import Version
from deerflow.runtime.runs import worker 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.schemas import DisconnectMode, RunStatus
from deerflow.runtime.runs.worker import ( from deerflow.runtime.runs.worker import (
_compose_sse_event, _compose_sse_event,
@ -316,6 +316,10 @@ class _IntegrationRunManager:
def __init__(self, record: RunRecord) -> None: def __init__(self, record: RunRecord) -> None:
self._record = record 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): async def wait_for_prior_finalizing(self, *_args, **_kwargs):
return None return None