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
- **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])

View File

@ -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.

View File

@ -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

View File

@ -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)))

View File

@ -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()

View File

@ -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

View File

@ -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

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.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)

View File

@ -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(),
)

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.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

View File

@ -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)

View File

@ -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."""

View File

@ -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)

View File

@ -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)

View File

@ -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

View File

@ -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