mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(mcp): retain durable control claims through slow calls and cancellation (#6780)
* fix(mcp): retain durable control claims through slow calls and cancellation * docs(mcp): distinguish poll and cancel claim fencing --------- Co-authored-by: xyi74976-del <290725798+xyi74976-del@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
yi-coding
GitHub
parent
a674132eed
commit
c04946a73f
@ -953,6 +953,7 @@ For HTTP/SSE MCP servers, OAuth token flows are supported (`client_credentials`,
|
||||
Missing, malformed, or out-of-range token response `expires_in` values use a one-hour default lifetime. This includes lifetimes that cannot be added to the current time without overflowing the expiry timestamp.
|
||||
Durable HTTP/SSE task status and cancellation calls select configured `user_auth` credentials using the persisted task owner, including after restart; per-request secrets are not retained for background calls. If a request-scoped credential overrides submit authentication, both credentials must authorize access to the same remote task.
|
||||
For stdio MCP servers, per-tool call timeouts can be configured with `tool_call_timeout`; durable background-task calls honor the same setting for HTTP/SSE servers as well.
|
||||
Durable background-task status and cancellation calls renew their claims while they run, so a slow response can still be saved after the original lease duration. If ownership is lost or renewal cannot be confirmed, the worker stops its local call and leaves recovery to the durable state. See [control-call leases](backend/docs/mcp-control-leases.md) for recovery and shutdown behavior.
|
||||
The persistent stdio session pool evicts only idle sessions at capacity. If every pooled session is executing a tool, a new session request fails with a capacity error and can be retried after a call finishes.
|
||||
For stdio file outputs, a bare filename is linked to a uniquely matching file created or changed by that call. Filenames embedded in unrelated paths, including Windows backslash paths, are left intact. Explicit remote URLs (`scheme://...`) retain their original query parameters and fragments even when these contain a generated file's name or path, including URLs with backslash-escaped quotes or markup delimiters; local references outside the URL still translate normally.
|
||||
For HTTP/SSE background-task calls, `session_init_timeout` separately bounds connection setup (including the SSE endpoint event) and MCP initialization together; it stops applying once the tool call begins. Initialization deadline errors identify the server and configured time limit.
|
||||
|
||||
@ -9,7 +9,7 @@ from collections.abc import Awaitable, Callable
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
from typing import Any, Literal
|
||||
|
||||
from app.mcp_tasks.errors import PermanentNotificationError
|
||||
from deerflow.constants import (
|
||||
@ -42,6 +42,7 @@ _CANCELLATION_DRAIN_TIMEOUT_SECONDS = 5.0
|
||||
@dataclass(slots=True)
|
||||
class _BatchRecordState:
|
||||
record: dict[str, Any]
|
||||
parent_task: asyncio.Task[Any] | None = None
|
||||
ordinary_release_task: asyncio.Future[Any] | None = None
|
||||
ordinary_release_terminal: bool = False
|
||||
cancellation_release_task: asyncio.Future[Any] | None = None
|
||||
@ -63,6 +64,13 @@ _current_batch_record: ContextVar[_BatchRecordState | None] = ContextVar(
|
||||
"mcp_task_current_batch_record",
|
||||
default=None,
|
||||
)
|
||||
_control_lease_lost: ContextVar[asyncio.Event | None] = ContextVar("mcp_control_lease_lost", default=None)
|
||||
|
||||
|
||||
def _require_control_lease() -> None:
|
||||
lost = _control_lease_lost.get()
|
||||
if lost is not None and lost.is_set():
|
||||
raise asyncio.CancelledError("MCP control claim ownership was lost")
|
||||
|
||||
|
||||
def _bound_error(error: str | None) -> str | None:
|
||||
@ -129,6 +137,8 @@ class McpTaskService:
|
||||
self._stop_timeout_logged = False
|
||||
self._compensation_tasks: set[asyncio.Future[Any]] = set()
|
||||
self._claim_owners: dict[str, _ClaimOwner] = {}
|
||||
# Includes children still draining after a bounded cancellation handoff.
|
||||
self._active_control_calls: set[str] = set()
|
||||
self._stop = asyncio.Event()
|
||||
|
||||
@property
|
||||
@ -597,9 +607,9 @@ class McpTaskService:
|
||||
release: Callable[[dict[str, Any]], Awaitable[None]],
|
||||
action: str,
|
||||
) -> tuple[list[_BatchRecordState], list[Any]]:
|
||||
states = [_BatchRecordState(record) for record in records]
|
||||
batch_state = _BatchState()
|
||||
parent_task = asyncio.current_task()
|
||||
states = [_BatchRecordState(record, parent_task=parent_task) for record in records]
|
||||
batch_state = _BatchState()
|
||||
|
||||
async def run_one(state: _BatchRecordState) -> Any:
|
||||
context_token = _current_batch_record.set(state)
|
||||
@ -663,24 +673,150 @@ class McpTaskService:
|
||||
|
||||
return states, results
|
||||
|
||||
async def _with_control_lease(
|
||||
self,
|
||||
record: dict[str, Any],
|
||||
*,
|
||||
phase: Literal["poll", "cancel"],
|
||||
operation: Callable[[], Awaitable[Any]],
|
||||
) -> Any:
|
||||
task_id = record["id"]
|
||||
if task_id in self._active_control_calls or len(self._active_control_calls) >= self._max_concurrent_polls:
|
||||
raise asyncio.CancelledError("MCP control capacity is occupied")
|
||||
self._active_control_calls.add(task_id)
|
||||
lost = asyncio.Event()
|
||||
children: list[asyncio.Task[Any]] = []
|
||||
renewals: set[asyncio.Task[bool]] = set()
|
||||
detach_cleanup = False
|
||||
interval = self._lease_seconds / 3
|
||||
|
||||
def renewal_done(renewal: asyncio.Task[bool]) -> None:
|
||||
renewals.discard(renewal)
|
||||
# Its waiter may have been cancelled on the same turn it failed.
|
||||
_consume_task_error(renewal)
|
||||
|
||||
async def renew() -> bool:
|
||||
renewal = asyncio.create_task(
|
||||
self._repository.renew_control_claim(
|
||||
task_id,
|
||||
phase=phase,
|
||||
lease_owner=self._lease_owner,
|
||||
lease_token=record["lease_token"],
|
||||
now=datetime.now(UTC),
|
||||
lease_seconds=self._lease_seconds,
|
||||
),
|
||||
name=f"mcp-{phase}-renew-{task_id}",
|
||||
)
|
||||
renewals.add(renewal)
|
||||
renewal.add_done_callback(renewal_done)
|
||||
done, _ = await asyncio.wait({renewal}, timeout=interval)
|
||||
if not done:
|
||||
raise TimeoutError("MCP control claim renewal exceeded its deadline")
|
||||
return renewal.result()
|
||||
|
||||
async def owned_operation() -> Any:
|
||||
state = _current_batch_record.get()
|
||||
if state is not None and state.parent_task is not None and state.parent_task.cancelling():
|
||||
raise asyncio.CancelledError("MCP batch cancellation was requested before execution")
|
||||
token = _control_lease_lost.set(lost)
|
||||
try:
|
||||
return await operation()
|
||||
finally:
|
||||
_control_lease_lost.reset(token)
|
||||
|
||||
async def heartbeat(work: asyncio.Task[Any]) -> None:
|
||||
while not work.done():
|
||||
await asyncio.sleep(interval)
|
||||
try:
|
||||
renewed = await renew()
|
||||
except Exception:
|
||||
if work.done():
|
||||
return
|
||||
logger.warning("Cannot confirm MCP %s claim ownership (task_id=%s); stopping the local call", phase, task_id, exc_info=True)
|
||||
lost.set()
|
||||
return
|
||||
if not renewed:
|
||||
if not work.done():
|
||||
lost.set()
|
||||
logger.info("MCP %s claim ownership changed or expired (task_id=%s); stopping the local call", phase, task_id)
|
||||
return
|
||||
|
||||
try:
|
||||
# Claim acquisition can itself outlast its lease. Check before any
|
||||
# remote side effect, rather than waiting for the first heartbeat.
|
||||
if not await renew():
|
||||
detach_cleanup = True
|
||||
raise asyncio.CancelledError("MCP control claim expired before execution")
|
||||
work = asyncio.create_task(owned_operation(), name=f"mcp-{phase}-owned-{task_id}")
|
||||
pulse = asyncio.create_task(heartbeat(work), name=f"mcp-{phase}-heartbeat-{task_id}")
|
||||
children = [work, pulse]
|
||||
await asyncio.wait(children, return_when=asyncio.FIRST_COMPLETED)
|
||||
detach_cleanup = True
|
||||
if work.done():
|
||||
return await work
|
||||
# Includes a heartbeat that cancelled itself: fail closed instead
|
||||
# of continuing an unmonitored remote operation.
|
||||
lost.set()
|
||||
detach_cleanup = True
|
||||
raise asyncio.CancelledError("MCP control claim ownership could not be confirmed")
|
||||
except Exception:
|
||||
detach_cleanup = True
|
||||
raise
|
||||
finally:
|
||||
if children and not children[0].done():
|
||||
lost.set()
|
||||
children.extend(renewals)
|
||||
for child in children:
|
||||
if not child.done():
|
||||
child.cancel()
|
||||
if children:
|
||||
completion = asyncio.gather(*children, return_exceptions=True)
|
||||
completion.add_done_callback(lambda _: self._active_control_calls.discard(task_id))
|
||||
if detach_cleanup:
|
||||
# Ownership loss can happen without caller cancellation.
|
||||
# Keep uncooperative children owned without blocking scans.
|
||||
await self._drain_cancellation_task(
|
||||
completion,
|
||||
action=f"drain {phase} control children",
|
||||
task_id=task_id,
|
||||
deadline=asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS,
|
||||
)
|
||||
else:
|
||||
# Caller cancellation already has a bounded batch handoff.
|
||||
# That owner must retain the children until they settle.
|
||||
while not completion.done():
|
||||
try:
|
||||
await asyncio.shield(completion)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
completion.result()
|
||||
else:
|
||||
self._active_control_calls.discard(task_id)
|
||||
|
||||
async def run_once(self, *, now: datetime) -> None:
|
||||
await self._run_cancellations(now=now)
|
||||
|
||||
claimed = await self._claim_with_cancellation_release(
|
||||
lambda: self._repository.claim_due_tasks(
|
||||
now=now,
|
||||
lease_owner=self._lease_owner,
|
||||
lease_seconds=self._lease_seconds,
|
||||
limit=self._max_concurrent_polls,
|
||||
),
|
||||
phase="poll",
|
||||
action="poll claim",
|
||||
release=self._release_poll_after_cancellation,
|
||||
capacity = self._max_concurrent_polls - len(self._active_control_calls)
|
||||
claimed = (
|
||||
await self._claim_with_cancellation_release(
|
||||
lambda: self._repository.claim_due_tasks(
|
||||
now=now,
|
||||
lease_owner=self._lease_owner,
|
||||
lease_seconds=self._lease_seconds,
|
||||
limit=capacity,
|
||||
**({"exclude_task_ids": tuple(self._active_control_calls)} if self._active_control_calls else {}),
|
||||
),
|
||||
phase="poll",
|
||||
action="poll claim",
|
||||
release=self._release_poll_after_cancellation,
|
||||
)
|
||||
if capacity > 0
|
||||
else []
|
||||
)
|
||||
if claimed:
|
||||
states, results = await self._run_claimed_batch(
|
||||
claimed,
|
||||
operation=lambda record: self._poll_one_claimed(record, now=now),
|
||||
operation=lambda record: self._with_control_lease(record, phase="poll", operation=lambda: self._poll_one_claimed(record, now=now)),
|
||||
release=self._release_poll_after_cancellation,
|
||||
action="poll",
|
||||
)
|
||||
@ -767,14 +903,16 @@ class McpTaskService:
|
||||
|
||||
async def _run_cancellations(self, *, now: datetime) -> None:
|
||||
claim = getattr(self._repository, "claim_cancel_requests", None)
|
||||
if claim is None:
|
||||
capacity = self._max_concurrent_polls - len(self._active_control_calls)
|
||||
if claim is None or capacity <= 0:
|
||||
return
|
||||
records = await self._claim_with_cancellation_release(
|
||||
lambda: claim(
|
||||
now=now,
|
||||
lease_owner=self._lease_owner,
|
||||
lease_seconds=self._lease_seconds,
|
||||
limit=self._max_concurrent_polls,
|
||||
limit=capacity,
|
||||
**({"exclude_task_ids": tuple(self._active_control_calls)} if self._active_control_calls else {}),
|
||||
),
|
||||
phase="cancel",
|
||||
action="cancel claim",
|
||||
@ -783,7 +921,7 @@ class McpTaskService:
|
||||
if records:
|
||||
states, results = await self._run_claimed_batch(
|
||||
records,
|
||||
operation=self._cancel_one_claimed,
|
||||
operation=lambda record: self._with_control_lease(record, phase="cancel", operation=lambda: self._cancel_one_claimed(record)),
|
||||
release=self._release_cancel_after_cancellation,
|
||||
action="cancel",
|
||||
)
|
||||
@ -821,6 +959,7 @@ class McpTaskService:
|
||||
),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - remote cancellation is retryable
|
||||
_require_control_lease()
|
||||
attempts = max(0, int(record.get("cancel_attempt_count") or 1) - 1)
|
||||
retry_seconds = min(self._poll_interval_seconds * (2 ** min(attempts, 16)), self._max_poll_backoff_seconds)
|
||||
failed_at = datetime.now(UTC)
|
||||
@ -838,7 +977,8 @@ class McpTaskService:
|
||||
)
|
||||
|
||||
async def _apply_cancel_snapshot(self, record: dict[str, Any], snapshot: TaskSnapshot) -> None:
|
||||
await self._repository.apply_cancel_snapshot(
|
||||
_require_control_lease()
|
||||
applied = await self._repository.apply_cancel_snapshot(
|
||||
record["id"],
|
||||
lease_owner=self._lease_owner,
|
||||
lease_token=record["lease_token"],
|
||||
@ -851,6 +991,11 @@ class McpTaskService:
|
||||
input_required=snapshot.input_required,
|
||||
completed_at=datetime.now(UTC),
|
||||
)
|
||||
if not applied:
|
||||
logger.info(
|
||||
"Discarded MCP task cancel result after lease ownership changed or expired (task_id=%s)",
|
||||
record.get("id"),
|
||||
)
|
||||
|
||||
async def _run_notifications(self, *, now: datetime) -> None:
|
||||
if self._launch_notification is None or self._get_run is None:
|
||||
@ -1415,6 +1560,7 @@ class McpTaskService:
|
||||
*,
|
||||
polled_at: datetime,
|
||||
) -> None:
|
||||
_require_control_lease()
|
||||
applied = await self._repository.apply_snapshot(
|
||||
record["id"],
|
||||
lease_owner=self._lease_owner,
|
||||
@ -1445,6 +1591,7 @@ class McpTaskService:
|
||||
return now + timedelta(seconds=interval)
|
||||
|
||||
async def _release_after_error(self, record: dict, *, now: datetime, error: str) -> None:
|
||||
_require_control_lease()
|
||||
consecutive_errors = max(0, int(record.get("consecutive_poll_error_count") or 0))
|
||||
retry_seconds = min(
|
||||
self._poll_interval_seconds * (2 ** min(consecutive_errors, 16)),
|
||||
|
||||
@ -392,7 +392,7 @@ runtime-internal. Query the current thread through:
|
||||
Task toolsets require `database.backend: sqlite` or `postgres`; startup fails
|
||||
instead of falling back to a synchronous submit when persistence or the task
|
||||
runtime is disabled. Restart recovery also requires the remote service to keep
|
||||
the task alive and recognize its ID after DeerFlow reconnects. A stdio server
|
||||
the task alive and recognize its ID after DeerFlow reconnects. Status and cancellation calls renew their control lease while executing and saving their result; the lease duration is a crash-recovery window, not a tool-call timeout. See [control-call leases](mcp-control-leases.md) for ownership loss, cancellation cleanup, and SQLite claim serialization. A stdio server
|
||||
must therefore persist its own tasks; multi-instance deployments should
|
||||
normally use an independently running HTTP/SSE service.
|
||||
|
||||
|
||||
72
backend/docs/mcp-control-leases.md
Normal file
72
backend/docs/mcp-control-leases.md
Normal file
@ -0,0 +1,72 @@
|
||||
# Durable MCP control-call leases
|
||||
|
||||
The protocol-neutral task runtime polls remote work and sends cancellation
|
||||
requests outside the Agent loop. `mcp_tasks.lease_seconds` (default 120) is a
|
||||
crash-recovery window; it does not impose that duration as a remote-call timeout.
|
||||
The server's existing `tool_call_timeout` still controls remote-call deadlines.
|
||||
|
||||
## Why the full operation needs ownership
|
||||
|
||||
Without renewal, a status/cancel request taking longer than the claim duration
|
||||
returns an outcome that the database rejects as expired. Repeating that slow
|
||||
request can discard every terminal result. Another worker can also reclaim the
|
||||
row while the original remote request is still running. Separately, SQLite
|
||||
ignores `SELECT FOR UPDATE`: two deferred transactions can read the same due
|
||||
row and both return it as claimed.
|
||||
|
||||
Poll claims retain the upstream conditional UPDATE that rechecks every eligibility
|
||||
predicate after reading candidates, including cancellation and lease ownership.
|
||||
Cancel claim transactions reserve SQLite's writer with `BEGIN IMMEDIATE` before
|
||||
reading candidates. PostgreSQL continues using `FOR UPDATE SKIP LOCKED`.
|
||||
These transactions contain only claim bookkeeping, never remote calls.
|
||||
|
||||
Each control operation verifies and extends its claim before calling the driver,
|
||||
then renews every one third of the lease duration. Each renewal has its own
|
||||
one-third-duration deadline, so a database stall cannot leave a live remote
|
||||
call unmonitored indefinitely. The heartbeat covers the driver, normalization,
|
||||
and the existing snapshot/retry persistence boundary. A renewal matches the
|
||||
row ID, worker, claim token, unexpired lease, nonterminal state and poll/cancel
|
||||
phase. It cannot renew a replaced generation or a poll superseded by cancellation.
|
||||
It changes no task activity timestamp or attempt/error counters. Backward clock
|
||||
corrections cannot shorten the stored expiry.
|
||||
|
||||
## Loss, cancellation and cleanup
|
||||
|
||||
A failed, timed-out, or rejected renewal stops the local operation. The existing
|
||||
token-fenced cancellation release preserves task status and diagnostics for
|
||||
recovery. A task-local ownership-loss fence prevents a driver that swallows
|
||||
cancellation from subsequently saving a snapshot or recording a retry failure.
|
||||
Stopping a local status request does not itself cancel the remote task.
|
||||
|
||||
The batch supervisor retains ownership of operation, heartbeat and renewal
|
||||
cleanup through repeated cancellation. Caller cancellation uses its existing
|
||||
bounded batch handoff; autonomous completion or ownership loss uses the same
|
||||
bounded drain mechanism. Unfinished cleanup remains service-owned in the background. Draining
|
||||
control calls keep their local capacity slots and are excluded from later
|
||||
poll/cancel claims until their children actually settle. Other tasks can proceed
|
||||
when capacity remains. If all control slots are occupied by uncooperative cleanup,
|
||||
new control calls wait rather than accumulating more abandoned calls.
|
||||
|
||||
This is not exactly-once remote execution. Process death, loss of the event loop,
|
||||
or a transport that ignores cancellation can still overlap a peer's recovered
|
||||
request after expiry. Servers must keep status queries safe to repeat and
|
||||
cancellation idempotent. Notification delivery retains its separate existing
|
||||
claim/idempotency state machine. No schema migration or new configuration is
|
||||
required.
|
||||
|
||||
## Verification
|
||||
|
||||
Run from `backend/`:
|
||||
|
||||
```sh
|
||||
PYTHONPATH=packages/harness:packages/extension-api:. uv run --no-sync pytest \
|
||||
tests/test_mcp_task_lease_lifecycle.py \
|
||||
tests/test_mcp_task_repository.py tests/test_mcp_task_service.py -q
|
||||
```
|
||||
|
||||
The lifecycle matrix uses the production SQLite repository with a controlled
|
||||
durable clock and event-gated remote calls. It covers slow poll/cancel results,
|
||||
interleaved claims, expiry/token/phase fences, ownership loss, unavailable or
|
||||
hung renewal, persistence after remote completion, delayed admission, repeated
|
||||
cancellation, late cleanup, spare-capacity progress and concurrent scan admission.
|
||||
No live MCP service or provider key is needed.
|
||||
@ -1,10 +1,16 @@
|
||||
### MCP System (`packages/harness/deerflow/mcp/`)
|
||||
|
||||
Poll and cancel snapshot writes rejected by the ownership fence log the discard;
|
||||
neither path changes retry semantics when a peer owns the task. The shared
|
||||
`McpTaskService._apply_cancel_snapshot` checks the control lease and logs rejected
|
||||
writes for both remote cancellation results and locally synthesized missing-connection failures.
|
||||
|
||||
The optional background-task `status` filter flows through `McpTaskSubmitter.list_tasks` → `McpTaskService.list_tasks` → `McpTaskRepository.list_by_thread`. SQL applies status, active-state, user, thread, and current-incarnation predicates before ordering and LIMIT; never filter an already truncated list. Omitted statuses do not add a keyword argument to legacy implementations. Filtering uses existing states without changing polling, notifications, retries, or cancellation. Regression coverage: `tests/test_background_task_status_filter.py`.
|
||||
|
||||
- Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management
|
||||
- **OAuth expiry**: Missing, non-convertible or overflowing `expires_in` values use the one-hour default. Guard both `timedelta` construction and addition to one captured UTC `now`; avoid float-derived range ceilings. `tests/test_mcp_oauth.py` pins both overflow stages and the exact `datetime` boundary with a fixed clock.
|
||||
- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` contract and normalized `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`). A driver-supplied `poll_after_seconds` must be a finite positive number, validated at the `TaskSnapshot` boundary so every driver is held to the same invariant rather than each one guarding the consumer that turns the interval into a `timedelta`. `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and separate consecutive poll/delivery error counters; `app/mcp_tasks/McpTaskService` performs status, cancellation, and notification work outside the Agent/LLM loop. Notification retries keep their idempotency attempt separate from the delivery-failure count, use capped exponential backoff, and stop after five failures; strict existing-thread admission dead-letters a deleted/mismatched target immediately. A status result is applied only when the worker still owns an unexpired lease, so a stale result cannot be written after expiry even before another worker reclaims the row. Every claim carries a fresh per-claim token, and every poll/cancel/notification release or snapshot apply is fenced by the owner plus that token, so a release that completes after the same service reclaims the row cannot clear the newer lease. Poll, cancel, and notification claims also have separate phase-level single-flight owners that protect only an ambiguous claim outcome: after an uncancelled timeout or caller cancellation, later scans do not start another claim while the outcome is unknown. Once the claim resolves, the phase owner is released immediately, and any returned rows are released through bounded, service-owned background work; per-claim token fencing prevents a late release from mutating a newer claim generation, so a stuck release does not stall the whole phase. While the claim outcome remains unresolved, later scans skip the phase and emit a warning rather than overlapping an ambiguous database claim. Once the claim resolves, the handoff may continue releasing returned rows in the background without blocking later scans. Lease expiry remains the cross-process crash fallback rather than the normal recovery for a live service's late claim. Poll timestamps and retry schedules are based on the remote call's completion time rather than the scan start. If submission succeeds but persistence fails or the caller is cancelled while persistence is in flight, the service best-effort cancels the remote task and preserves the original error or cancellation if that compensation also fails. The exact `uq_mcp_tasks_user_server_remote` conflict is different: an existing durable row already owns the remote handle, so the conflict surfaces without cancelling that tracked task. Unexpected per-task poll failures are isolated from sibling claims and remain recoverable through lease expiry; Gateway shutdown cancels the poller so a hung external status call cannot block process exit. Routine cancellation releases clear only the lease/token and preserve any pre-existing `last_cancel_error` or `notification_error`; cancellation itself is not persisted as a task failure diagnostic. Cancelling an in-flight poll releases only the owner- and per-claim-token-fenced lease and preserves its preclaim schedule and poll-failure state; real poll failures retain exponential backoff and tracking-degradation behavior. `input_required` and terminal states stop polling and become `notification_status=pending` for later Agent/UI delivery. Durable recovery requires a SQL database backend (`sqlite` or `postgres`); the in-memory backend leaves the repository/service unavailable. The runtime is startup-configured by `mcp_tasks` and disabled by default until a concrete driver is registered; this foundation does not alter ordinary MCP tool behavior on its own.
|
||||
- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` and `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`); validate finite positive `poll_after_seconds` at that shared boundary. `persistence/mcp_tasks/` owns remote handles, schedules, notification state, leases and separate consecutive poll/delivery error counters. `app/mcp_tasks/McpTaskService` runs outside the Agent/LLM loop. Every claim gets a fresh token; releases and snapshot writes match owner/token, and results require an unexpired lease even before a peer reclaims. Each claim phase has a single-flight owner while its database outcome is ambiguous: later scans skip with a warning. Once it resolves, release that phase owner immediately and hand returned rows to bounded, service-owned release work. Token fencing protects newer generations; a stuck release must not block later claims. Expiry is the cross-process crash fallback, not normal recovery for a live service's late claim. Poll timestamps and retry schedules use remote completion time. Submission persistence failure/cancellation best-effort cancels the remote handle while preserving the original error, except `uq_mcp_tasks_user_server_remote`: an existing row owns that handle and must not be cancelled. Isolate per-task failures; unexpected poll failures recover by expiry. Shutdown cancels the poller with bounded handoff. Cancellation release clears only owner/token and preserves diagnostics, preclaim schedules and poll-failure state; cancellation is not a task failure. Real failures retain backoff and tracking degradation. `input_required` and terminal events become pending notifications; delivery budgets, idempotency and strict target rejection are detailed below. Durable recovery needs SQLite/PostgreSQL; memory leaves the service unavailable. The startup-configured `mcp_tasks` runtime is disabled by default until a concrete driver is registered and does not change ordinary tool behavior.
|
||||
- **Control-call ownership**: Poll claims recheck eligibility with a conditional UPDATE after SELECT; cancel claims reserve SQLite writers with `BEGIN IMMEDIATE` before SELECT. PostgreSQL retains `SKIP LOCKED`. `renew_control_claim` atomically matches phase, owner, token, nonterminal state and unexpired lease, preserves `updated_at`, and never shortens expiry on a backward clock correction. `McpTaskService._with_control_lease` validates admission, renews every lease/3 with a lease/3 renewal deadline, and supervises the entire operation through persistence. Ownership loss or uncertain renewal cancels the local operation; the task-local loss fence rejects late snapshot/retry writes if the driver swallows cancellation. Bounded handoff retains operation/renewal cleanup through caller cancellation or autonomous completion/loss. Local active IDs remain excluded from poll/cancel claims and consume control capacity until cleanup settles. Notification claims use their existing independent lifecycle. Contract/tests: [control leases](../../../../docs/mcp-control-leases.md), `tests/test_mcp_task_lease_lifecycle.py`.
|
||||
- **Runtime availability boundary**: the installed process-local submitter is the source of truth for durable task-management tool exposure. `mcp_tasks` is startup-only; changing it on disk does not alter the live toolset until the Gateway restarts.
|
||||
- **Task OAuth lifetime**: `tasks/runtime.py` owns shared OAuth state only for enabled deployment servers with `task_toolsets`, from `set_mcp_task_config_snapshot(config)` until its `None` teardown. Use `_task_enabled_servers` for both snapshot and OAuth selection. Validate the frozen snapshot before reusing that state for discovery, ordinary tools, or durable calls. Cache resets must not discard it. Keep tokens and the cross-loop refresh lock together; rotate a private config copy, never the parsed snapshot or config file. Non-task servers remain reloadable and personal callers use independent managers, even on a name collision.
|
||||
- **OAuth token store**: `oauth_store.py` persists only `refresh_token` states in `mcp_oauth_tokens` (encrypted, key `(scope, owner_user_id or '', server_key)`); personal managers must pass `owner_user_id`. `app/gateway/mcp_oauth_tokens.py` installs the store around the runtime so run drains still persist. Under the process lock: re-read (a mismatched fingerprint means the seed wins), adopt a valid peer token (same `token_request_fingerprint` only), else claim via lease + version CAS, fetch while renewing the lease (`REFRESH_MAX_SECONDS` cap), write under the lease, release on failure. An unstored rotation stays pending (`unsaved_base_version`) and is written back before anything older is used. Never hold a row lock across the HTTP call or touch the store on the loop. DB errors fail the refresh; an unloadable key falls back to process memory. `invalid_grant` raises `McpOAuthReauthorizationRequired`, the owner's tool error. Both MCP config writers drop rows of connections a save removes (`forget_unconfigured_mcp_oauth_tokens`; personal toggles keep both runtime-name variants). Tests: `tests/test_mcp_oauth_token_store.py`, `tests/blocking_io/test_mcp_oauth_token_store.py`.
|
||||
|
||||
@ -42,6 +42,10 @@ missing or strictly older leases remain recoverable.
|
||||
|
||||
`mcp_tasks/sql.py::claim_notification_work` never returns a row without a dispatch event. A pending snapshot whose condition resolved before delivery (tracking recovered, or an answered `input_required` prompt returned the task to `working`) is consumed at claim time: `notified_version` advances, the content fingerprint clears so an identical later event is recorded again, stale error/attempt/dispatch state clears, status returns to `none`, and the claim lease is released instead of launching an empty event. A legacy row left without a dispatch event is consumed only when nothing newer exists; a newer recorded event is rebuilt and claimed instead. Consumption is a conditional UPDATE fenced on the event version and phase the claim read, so a concurrent newer event (SQLite has no row locks) is never acknowledged and stays claimable. Coverage: `tests/test_mcp_task_repository.py` (resolved snapshot, legacy retry, expired-claim reclaim, identical-event recurrence, interleaved newer completion, legacy launching/dispatched with and without a newer completion).
|
||||
|
||||
## MCP control leases
|
||||
|
||||
`McpTaskRepository` retains upstream poll claims' eligibility-fenced UPDATE after candidate selection, so an interleaved cancellation or competing claim is not overwritten. Cancel claims reserve the SQLite writer before SELECT; `FOR UPDATE SKIP LOCKED` alone does not serialize SQLite readers. Renewal is one owner/token/phase/expiry-fenced UPDATE and changes only lease expiry, explicitly preserving the ORM `updated_at` column. A stale generation, expired lease, terminal task or changed cancellation phase cannot renew. Keep local draining-row exclusions before LIMIT. See [control-call leases](../../../../docs/mcp-control-leases.md).
|
||||
|
||||
## Run terminal authority and lease expiry
|
||||
|
||||
`RunRepository.update_status(..., owner_worker_id=...)` and
|
||||
|
||||
@ -4,7 +4,7 @@ import hashlib
|
||||
import json
|
||||
import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
from typing import Any, Literal
|
||||
|
||||
from sqlalchemy import and_, case, or_, select, text, update
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
@ -304,6 +304,7 @@ class McpTaskRepository:
|
||||
lease_owner: str,
|
||||
lease_seconds: int,
|
||||
limit: int,
|
||||
exclude_task_ids: tuple[str, ...] = (),
|
||||
) -> list[dict[str, Any]]:
|
||||
lease_expires_at = now + timedelta(seconds=lease_seconds)
|
||||
claimable = (
|
||||
@ -317,6 +318,8 @@ class McpTaskRepository:
|
||||
),
|
||||
)
|
||||
stmt = select(McpTaskRow.id).where(*claimable).order_by(McpTaskRow.next_poll_at.asc(), McpTaskRow.id.asc()).limit(limit).with_for_update(skip_locked=True)
|
||||
if exclude_task_ids:
|
||||
stmt = stmt.where(McpTaskRow.id.not_in(exclude_task_ids))
|
||||
async with self._sf() as session:
|
||||
candidate_ids = list((await session.execute(stmt)).scalars())
|
||||
# Atomic fence: ``with_for_update()`` is a no-op on SQLite, so a
|
||||
@ -350,6 +353,40 @@ class McpTaskRepository:
|
||||
await session.commit()
|
||||
return claimed
|
||||
|
||||
async def renew_control_claim(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
phase: Literal["poll", "cancel"],
|
||||
lease_owner: str,
|
||||
lease_token: str,
|
||||
now: datetime,
|
||||
lease_seconds: int,
|
||||
) -> bool:
|
||||
"""Extend only the still-live control claim; never resurrect a lease."""
|
||||
if phase not in ("poll", "cancel"):
|
||||
raise ValueError("Unknown MCP control claim phase")
|
||||
expires_at = now + timedelta(seconds=lease_seconds)
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(
|
||||
update(McpTaskRow)
|
||||
.where(
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.lease_owner == lease_owner,
|
||||
McpTaskRow.lease_token == lease_token,
|
||||
McpTaskRow.lease_expires_at >= now,
|
||||
McpTaskRow.status.not_in(_TERMINAL_STATUS_VALUES),
|
||||
McpTaskRow.cancel_requested_at.is_(None) if phase == "poll" else McpTaskRow.cancel_requested_at.is_not(None),
|
||||
)
|
||||
.values(
|
||||
# A backward wall-clock correction must not shorten a lease.
|
||||
lease_expires_at=case((McpTaskRow.lease_expires_at > expires_at, McpTaskRow.lease_expires_at), else_=expires_at),
|
||||
updated_at=McpTaskRow.updated_at,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
return bool(result.rowcount)
|
||||
|
||||
async def apply_snapshot(
|
||||
self,
|
||||
task_id: str,
|
||||
@ -543,6 +580,7 @@ class McpTaskRepository:
|
||||
lease_owner: str,
|
||||
lease_seconds: int,
|
||||
limit: int,
|
||||
exclude_task_ids: tuple[str, ...] = (),
|
||||
task_id: str | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
stmt = select(McpTaskRow).where(
|
||||
@ -555,7 +593,13 @@ class McpTaskRepository:
|
||||
if task_id is not None:
|
||||
stmt = stmt.where(McpTaskRow.id == task_id)
|
||||
stmt = stmt.order_by(McpTaskRow.next_cancel_at.asc(), McpTaskRow.id.asc()).limit(limit).with_for_update(skip_locked=True)
|
||||
if exclude_task_ids:
|
||||
stmt = stmt.where(McpTaskRow.id.not_in(exclude_task_ids))
|
||||
async with self._sf() as session:
|
||||
if session.get_bind().dialect.name == "sqlite":
|
||||
# SQLite ignores FOR UPDATE. Reserve its writer before reading
|
||||
# candidates so two workers cannot return the same claim.
|
||||
await session.execute(text("BEGIN IMMEDIATE"))
|
||||
rows = list((await session.execute(stmt)).scalars())
|
||||
expires_at = now + timedelta(seconds=lease_seconds)
|
||||
for row in rows:
|
||||
|
||||
597
backend/tests/test_mcp_task_lease_lifecycle.py
Normal file
597
backend/tests/test_mcp_task_lease_lifecycle.py
Normal file
@ -0,0 +1,597 @@
|
||||
"""Durable ownership must outlive slow MCP control calls, without duplicate claims."""
|
||||
|
||||
import asyncio
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
import app.mcp_tasks.service as service_module
|
||||
from app.mcp_tasks.service import McpTaskService
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.mcp.tasks import McpTaskConnectionUnavailableError, McpTaskDriverRegistry, TaskSnapshot, TaskStatus
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
|
||||
from deerflow.persistence.mcp_tasks import McpTaskRepository
|
||||
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
||||
|
||||
|
||||
class Clock(datetime):
|
||||
current = datetime(2030, 1, 1, tzinfo=UTC)
|
||||
|
||||
@classmethod
|
||||
def now(cls, tz=None):
|
||||
return cls.current
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def task_repo(tmp_path, monkeypatch):
|
||||
Clock.current = datetime(2030, 1, 1, tzinfo=UTC)
|
||||
monkeypatch.setattr(service_module, "datetime", Clock)
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
factory = get_session_factory()
|
||||
async with factory() as session:
|
||||
session.add(ThreadMetaRow(thread_id="thread-1", user_id="user-1", metadata_json={}, created_at=Clock.current, updated_at=Clock.current))
|
||||
await session.commit()
|
||||
repo = McpTaskRepository(factory)
|
||||
await create_working(repo, "task-1")
|
||||
try:
|
||||
yield repo
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def create_working(repo, task_id, driver_name="fake"):
|
||||
await repo.create(
|
||||
task_id=task_id,
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
expected_thread_incarnation=None,
|
||||
run_id="run-1",
|
||||
tool_call_id="call-1",
|
||||
server_name="reports",
|
||||
driver_name=driver_name,
|
||||
remote_task_id="remote-" + task_id,
|
||||
task_name="Report",
|
||||
status="working",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
result_truncated=False,
|
||||
result_artifact=None,
|
||||
error=None,
|
||||
input_required=None,
|
||||
next_poll_at=Clock.current,
|
||||
driver_data={},
|
||||
)
|
||||
|
||||
|
||||
class SlowDriver:
|
||||
def __init__(self):
|
||||
self.calls = 0
|
||||
self.started = asyncio.Event()
|
||||
self.finish = asyncio.Event()
|
||||
self.stopped = asyncio.Event()
|
||||
|
||||
async def get_status(self, _task):
|
||||
self.calls += 1
|
||||
self.started.set()
|
||||
try:
|
||||
await self.finish.wait()
|
||||
return TaskSnapshot(status=TaskStatus.COMPLETED, result={"answer": 42})
|
||||
finally:
|
||||
self.stopped.set()
|
||||
|
||||
async def cancel(self, task):
|
||||
await self.get_status(task)
|
||||
return TaskSnapshot(status=TaskStatus.CANCELLED)
|
||||
|
||||
|
||||
def make_service(repo, driver):
|
||||
drivers = McpTaskDriverRegistry()
|
||||
drivers.register("fake", driver)
|
||||
return McpTaskService(repository=repo, drivers=drivers, poll_interval_seconds=1, lease_seconds=5, max_concurrent_polls=2)
|
||||
|
||||
|
||||
async def read_task(repo):
|
||||
return await repo.get("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("reclaimed", [False, True])
|
||||
@pytest.mark.parametrize("connection_unavailable", [False, True])
|
||||
async def test_cancel_snapshot_discard_is_logged_without_changing_peer_or_retry_state(task_repo, monkeypatch, caplog, reclaimed, connection_unavailable):
|
||||
repo = task_repo
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
apply = repo.apply_cancel_snapshot
|
||||
peer = None
|
||||
|
||||
async def apply_after_ownership_check(*args, **kwargs):
|
||||
nonlocal peer
|
||||
if reclaimed:
|
||||
# Pause at the persistence boundary after _require_control_lease(),
|
||||
# then let another worker reclaim before the real fenced UPDATE.
|
||||
Clock.current += timedelta(seconds=6)
|
||||
await repo.claim_cancel_requests(now=Clock.current, lease_owner="peer", lease_seconds=5, limit=1)
|
||||
peer = await read_task(repo)
|
||||
return await apply(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(repo, "apply_cancel_snapshot", apply_after_ownership_check)
|
||||
driver = SlowDriver()
|
||||
driver.finish.set()
|
||||
if connection_unavailable:
|
||||
|
||||
async def missing_connection(_task):
|
||||
raise McpTaskConnectionUnavailableError("MCP task connection was removed")
|
||||
|
||||
monkeypatch.setattr(driver, "cancel", missing_connection)
|
||||
with caplog.at_level("INFO", logger=service_module.__name__):
|
||||
await make_service(repo, driver)._run_cancellations(now=Clock.current)
|
||||
|
||||
stored = await read_task(repo)
|
||||
messages = [record.getMessage() for record in caplog.records if "Discarded MCP task cancel result" in record.getMessage()]
|
||||
if reclaimed:
|
||||
assert peer is not None and peer["lease_owner"] == "peer"
|
||||
assert stored == peer, "discarding the old snapshot must not mutate or retry the peer's claim"
|
||||
assert len(messages) == 1 and "task_id=task-1" in messages[0]
|
||||
else:
|
||||
assert stored["status"] == ("failed" if connection_unavailable else "cancelled")
|
||||
assert messages == [], "an applied cancellation must not be reported as discarded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase,expected", [("poll", "completed"), ("cancel", "cancelled")])
|
||||
async def test_slow_control_call_keeps_claim_and_persists_terminal_result(task_repo, phase, expected):
|
||||
repo = task_repo
|
||||
if phase == "cancel":
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
driver = SlowDriver()
|
||||
service = make_service(repo, driver)
|
||||
cycle = asyncio.create_task(service.run_once(now=Clock.current))
|
||||
try:
|
||||
await asyncio.wait_for(driver.started.wait(), timeout=3)
|
||||
Clock.current += timedelta(seconds=4)
|
||||
# One real heartbeat interval (5/3s), with a controlled durable clock.
|
||||
await asyncio.sleep(2)
|
||||
Clock.current += timedelta(seconds=2)
|
||||
claim = repo.claim_due_tasks if phase == "poll" else repo.claim_cancel_requests
|
||||
peer = await claim(now=Clock.current, lease_owner="peer", lease_seconds=5, limit=1)
|
||||
assert peer == [], "a live slow control call lost its lease and was reclaimed"
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(cycle, timeout=3)
|
||||
assert (await read_task(repo))["status"] == expected
|
||||
finally:
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(asyncio.gather(cycle, return_exceptions=True), timeout=3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase", ["poll", "cancel"])
|
||||
async def test_sqlite_claims_have_one_winner_when_reads_overlap(task_repo, monkeypatch, phase):
|
||||
repo = task_repo
|
||||
if phase == "cancel":
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
claim = repo.claim_due_tasks if phase == "poll" else repo.claim_cancel_requests
|
||||
read = asyncio.Event()
|
||||
resume = asyncio.Event()
|
||||
execute = AsyncSession.execute
|
||||
first = None
|
||||
|
||||
async def pause_first_select(session, statement, *args, **kwargs):
|
||||
result = await execute(session, statement, *args, **kwargs)
|
||||
if asyncio.current_task() is first and statement.is_select and not read.is_set():
|
||||
read.set()
|
||||
await resume.wait()
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(AsyncSession, "execute", pause_first_select)
|
||||
first = asyncio.create_task(claim(now=Clock.current, lease_owner="first", lease_seconds=5, limit=1))
|
||||
second = None
|
||||
try:
|
||||
await asyncio.wait_for(read.wait(), timeout=3)
|
||||
second = asyncio.create_task(claim(now=Clock.current, lease_owner="second", lease_seconds=5, limit=1))
|
||||
# Let the competitor attempt its claim while the first read is paused.
|
||||
# Serialized cancel claims and fenced poll UPDATEs must each yield one winner.
|
||||
await asyncio.wait({second}, timeout=0.2)
|
||||
resume.set()
|
||||
a, b = await asyncio.wait_for(asyncio.gather(first, second), timeout=3)
|
||||
assert len(a) + len(b) == 1, "both workers received the same durable task"
|
||||
finally:
|
||||
resume.set()
|
||||
await asyncio.gather(*[t for t in (first, second) if t is not None], return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase", ["poll", "cancel"])
|
||||
@pytest.mark.parametrize("drift", ["owner", "token", "expired", "other_phase", "terminal"])
|
||||
async def test_renewal_cannot_revive_or_change_an_ineligible_claim(task_repo, phase, drift):
|
||||
repo = task_repo
|
||||
initial = Clock.current
|
||||
if phase == "cancel":
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=initial)
|
||||
claim = repo.claim_due_tasks if phase == "poll" else repo.claim_cancel_requests
|
||||
row = (await claim(now=initial, lease_owner="worker", lease_seconds=5, limit=1))[0]
|
||||
from sqlalchemy import update
|
||||
|
||||
from deerflow.persistence.mcp_tasks.model import McpTaskRow
|
||||
|
||||
if drift in {"other_phase", "terminal"}:
|
||||
values = {"status": "completed"} if drift == "terminal" else {"cancel_requested_at": initial if phase == "poll" else None}
|
||||
async with repo._sf() as session:
|
||||
await session.execute(update(McpTaskRow).where(McpTaskRow.id == "task-1").values(**values))
|
||||
await session.commit()
|
||||
now = initial + timedelta(seconds=6 if drift == "expired" else 1)
|
||||
renewed = await repo.renew_control_claim(
|
||||
"task-1",
|
||||
phase=phase,
|
||||
lease_owner="wrong" if drift == "owner" else "worker",
|
||||
lease_token="wrong" if drift == "token" else row["lease_token"],
|
||||
now=now,
|
||||
lease_seconds=5,
|
||||
)
|
||||
assert renewed is False
|
||||
async with repo._sf() as session:
|
||||
stored = await session.get(McpTaskRow, "task-1")
|
||||
assert stored.lease_expires_at.replace(tzinfo=UTC) == initial + timedelta(seconds=5)
|
||||
assert stored.lease_token == row["lease_token"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase", ["poll", "cancel"])
|
||||
async def test_renewal_extends_only_lease_and_never_moves_expiry_backwards(task_repo, phase):
|
||||
repo = task_repo
|
||||
initial = Clock.current
|
||||
if phase == "cancel":
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=initial)
|
||||
claim = repo.claim_due_tasks if phase == "poll" else repo.claim_cancel_requests
|
||||
row = (await claim(now=initial, lease_owner="worker", lease_seconds=5, limit=1))[0]
|
||||
for seconds in (4, 2):
|
||||
assert await repo.renew_control_claim("task-1", phase=phase, lease_owner="worker", lease_token=row["lease_token"], now=initial + timedelta(seconds=seconds), lease_seconds=5)
|
||||
assert await claim(now=initial + timedelta(seconds=6), lease_owner="peer", lease_seconds=5, limit=1) == []
|
||||
from deerflow.persistence.mcp_tasks.model import McpTaskRow
|
||||
|
||||
async with repo._sf() as session:
|
||||
stored = await session.get(McpTaskRow, "task-1")
|
||||
assert stored.lease_expires_at.replace(tzinfo=UTC) == initial + timedelta(seconds=9)
|
||||
assert stored.updated_at.replace(tzinfo=UTC) == initial
|
||||
assert stored.poll_attempt_count == row["poll_attempt_count"]
|
||||
assert stored.cancel_attempt_count == row["cancel_attempt_count"]
|
||||
assert stored.last_polled_at is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase", ["poll", "cancel"])
|
||||
async def test_reclaimed_control_call_is_stopped_without_touching_peer_claim(task_repo, phase):
|
||||
repo = task_repo
|
||||
if phase == "cancel":
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
driver = SlowDriver()
|
||||
service = make_service(repo, driver)
|
||||
cycle = asyncio.create_task(service.run_once(now=Clock.current))
|
||||
try:
|
||||
await asyncio.wait_for(driver.started.wait(), timeout=3)
|
||||
Clock.current += timedelta(seconds=6)
|
||||
claim = repo.claim_due_tasks if phase == "poll" else repo.claim_cancel_requests
|
||||
peer = (await claim(now=Clock.current, lease_owner="peer", lease_seconds=5, limit=1))[0]
|
||||
await asyncio.wait_for(driver.stopped.wait(), timeout=3)
|
||||
await asyncio.wait_for(cycle, timeout=3)
|
||||
assert (await read_task(repo))["status"] == "working"
|
||||
assert await repo.renew_control_claim("task-1", phase=phase, lease_owner="peer", lease_token=peer["lease_token"], now=Clock.current, lease_seconds=5)
|
||||
finally:
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(asyncio.gather(cycle, return_exceptions=True), timeout=3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase", ["poll", "cancel"])
|
||||
async def test_failed_renewal_stops_remote_call_without_recording_task_failure(task_repo, monkeypatch, phase):
|
||||
repo = task_repo
|
||||
if phase == "cancel":
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
renew = repo.renew_control_claim
|
||||
admitted = False
|
||||
|
||||
async def unavailable_after_admission(*args, **kwargs):
|
||||
nonlocal admitted
|
||||
if admitted:
|
||||
raise RuntimeError("database unavailable during heartbeat")
|
||||
admitted = True
|
||||
return await renew(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(repo, "renew_control_claim", unavailable_after_admission)
|
||||
driver = SlowDriver()
|
||||
cycle = asyncio.create_task(make_service(repo, driver).run_once(now=Clock.current))
|
||||
try:
|
||||
await asyncio.wait_for(driver.started.wait(), timeout=3)
|
||||
await asyncio.wait_for(driver.stopped.wait(), timeout=3)
|
||||
await asyncio.wait_for(cycle, timeout=3)
|
||||
stored = await read_task(repo)
|
||||
assert stored["status"] == "working"
|
||||
assert stored["last_poll_error"] is None
|
||||
assert stored["last_cancel_error"] is None
|
||||
assert stored["consecutive_poll_error_count"] == 0
|
||||
finally:
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(asyncio.gather(cycle, return_exceptions=True), timeout=3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_covers_snapshot_persistence_after_remote_completion(task_repo, monkeypatch):
|
||||
repo = task_repo
|
||||
applying = asyncio.Event()
|
||||
commit = asyncio.Event()
|
||||
apply = repo.apply_snapshot
|
||||
|
||||
async def delayed_write(*args, **kwargs):
|
||||
applying.set()
|
||||
await commit.wait()
|
||||
return await apply(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(repo, "apply_snapshot", delayed_write)
|
||||
driver = SlowDriver()
|
||||
driver.finish.set()
|
||||
cycle = asyncio.create_task(make_service(repo, driver).run_once(now=Clock.current))
|
||||
try:
|
||||
await asyncio.wait_for(applying.wait(), timeout=3)
|
||||
Clock.current += timedelta(seconds=4)
|
||||
await asyncio.sleep(2)
|
||||
Clock.current += timedelta(seconds=2)
|
||||
assert await repo.claim_due_tasks(now=Clock.current, lease_owner="peer", lease_seconds=5, limit=1) == []
|
||||
commit.set()
|
||||
await asyncio.wait_for(cycle, timeout=3)
|
||||
assert (await read_task(repo))["status"] == "completed"
|
||||
finally:
|
||||
commit.set()
|
||||
await asyncio.wait_for(asyncio.gather(cycle, return_exceptions=True), timeout=3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase", ["poll", "cancel"])
|
||||
async def test_claim_expiring_during_acquisition_does_not_start_remote_call(task_repo, monkeypatch, phase):
|
||||
repo = task_repo
|
||||
if phase == "cancel":
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
method = "claim_due_tasks" if phase == "poll" else "claim_cancel_requests"
|
||||
claim = getattr(repo, method)
|
||||
|
||||
async def delayed_claim(**kwargs):
|
||||
rows = await claim(**kwargs)
|
||||
if rows:
|
||||
Clock.current += timedelta(seconds=6)
|
||||
return rows
|
||||
|
||||
monkeypatch.setattr(repo, method, delayed_claim)
|
||||
driver = SlowDriver()
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(make_service(repo, driver).run_once(now=Clock.current), timeout=3)
|
||||
assert not driver.started.is_set()
|
||||
assert (await read_task(repo))["status"] == "working"
|
||||
|
||||
|
||||
class DrainingDriver(SlowDriver):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.cleanup = asyncio.Event()
|
||||
self.calls = 0
|
||||
|
||||
async def get_status(self, task):
|
||||
self.calls += 1
|
||||
self.started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
self.cleanup.set()
|
||||
await self.finish.wait()
|
||||
return TaskSnapshot(status=TaskStatus.COMPLETED, result={"late": True})
|
||||
finally:
|
||||
self.stopped.set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase", ["poll", "cancel"])
|
||||
async def test_repeated_cancellation_retains_local_owner_until_remote_cleanup(task_repo, monkeypatch, phase):
|
||||
repo = task_repo
|
||||
if phase == "cancel":
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05)
|
||||
driver = DrainingDriver()
|
||||
service = make_service(repo, driver)
|
||||
service._max_concurrent_polls = 1
|
||||
cycle = asyncio.create_task(service.run_once(now=Clock.current))
|
||||
try:
|
||||
await asyncio.wait_for(driver.started.wait(), timeout=3)
|
||||
cycle.cancel("first cancellation")
|
||||
await asyncio.wait_for(driver.cleanup.wait(), timeout=3)
|
||||
cycle.cancel("second cancellation")
|
||||
with pytest.raises(asyncio.CancelledError, match="first cancellation"):
|
||||
await asyncio.wait_for(cycle, timeout=3)
|
||||
Clock.current += timedelta(seconds=6)
|
||||
await asyncio.wait_for(service.run_once(now=Clock.current), timeout=3)
|
||||
assert driver.calls == 1
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(driver.stopped.wait(), timeout=3)
|
||||
# Wait for the existing service-owned batch handoff, not a sleep guess.
|
||||
await asyncio.wait_for(asyncio.gather(*service._compensation_tasks), timeout=3)
|
||||
assert (await read_task(repo))["status"] == "working"
|
||||
assert service._active_control_calls == set()
|
||||
# Once cleanup really completes, ordinary recovery can resume.
|
||||
monkeypatch.setattr(driver, "get_status", SlowDriver.get_status.__get__(driver))
|
||||
await asyncio.wait_for(service.run_once(now=Clock.current), timeout=3)
|
||||
assert (await read_task(repo))["status"] == ("completed" if phase == "poll" else "cancelled")
|
||||
finally:
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(asyncio.gather(cycle, *service._compensation_tasks, return_exceptions=True), timeout=3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("phase", ["poll", "cancel"])
|
||||
async def test_hung_renewal_stops_control_call_before_lease_expires(task_repo, monkeypatch, phase):
|
||||
repo = task_repo
|
||||
if phase == "cancel":
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
renew = repo.renew_control_claim
|
||||
renew_started = asyncio.Event()
|
||||
release_renewal = asyncio.Event()
|
||||
admitted = False
|
||||
|
||||
async def hang_after_admission(*args, **kwargs):
|
||||
nonlocal admitted
|
||||
if admitted:
|
||||
renew_started.set()
|
||||
await release_renewal.wait()
|
||||
admitted = True
|
||||
return await renew(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(repo, "renew_control_claim", hang_after_admission)
|
||||
driver = SlowDriver()
|
||||
cycle = asyncio.create_task(make_service(repo, driver).run_once(now=Clock.current))
|
||||
try:
|
||||
await asyncio.wait_for(driver.started.wait(), timeout=3)
|
||||
await asyncio.wait_for(renew_started.wait(), timeout=3)
|
||||
await asyncio.wait_for(driver.stopped.wait(), timeout=2.5)
|
||||
await asyncio.wait_for(cycle, timeout=3)
|
||||
assert (await read_task(repo))["status"] == "working"
|
||||
finally:
|
||||
release_renewal.set()
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(asyncio.gather(cycle, return_exceptions=True), timeout=3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("uncooperative_child", ["renewal", "driver"])
|
||||
async def test_uncertain_renewal_cleanup_hands_off_without_stalling_other_tasks(task_repo, monkeypatch, uncooperative_child):
|
||||
repo = task_repo
|
||||
renew = repo.renew_control_claim
|
||||
renewal_started = asyncio.Event()
|
||||
cleanup_started = asyncio.Event()
|
||||
finish_cleanup = asyncio.Event()
|
||||
admitted = False
|
||||
|
||||
async def uncooperative_renewal(*args, **kwargs):
|
||||
nonlocal admitted
|
||||
if args[0] != "task-1":
|
||||
return await renew(*args, **kwargs)
|
||||
if not admitted:
|
||||
admitted = True
|
||||
return await renew(*args, **kwargs)
|
||||
renewal_started.set()
|
||||
if uncooperative_child == "driver":
|
||||
return False
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
cleanup_started.set()
|
||||
await finish_cleanup.wait()
|
||||
return False
|
||||
|
||||
class DrainingDriver(SlowDriver):
|
||||
async def get_status(self, task):
|
||||
if uncooperative_child == "renewal":
|
||||
return await super().get_status(task)
|
||||
self.calls += 1
|
||||
self.started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
cleanup_started.set()
|
||||
await finish_cleanup.wait()
|
||||
return TaskSnapshot(status=TaskStatus.COMPLETED, result={"stale": True})
|
||||
finally:
|
||||
self.stopped.set()
|
||||
|
||||
monkeypatch.setattr(repo, "renew_control_claim", uncooperative_renewal)
|
||||
monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05)
|
||||
driver = DrainingDriver()
|
||||
service = make_service(repo, driver)
|
||||
cycle = asyncio.create_task(service.run_once(now=Clock.current))
|
||||
try:
|
||||
await asyncio.wait_for(renewal_started.wait(), timeout=3)
|
||||
await asyncio.wait_for(cleanup_started.wait(), timeout=3)
|
||||
await asyncio.wait_for(asyncio.shield(cycle), timeout=0.5)
|
||||
assert driver.stopped.is_set() == (uncooperative_child == "renewal")
|
||||
assert (await read_task(repo))["status"] == "working"
|
||||
# A draining renewal still owns its local slot, even after its old
|
||||
# durable lease becomes eligible for recovery.
|
||||
Clock.current += timedelta(seconds=6)
|
||||
await create_working(repo, "task-2", driver_name="healthy")
|
||||
healthy = SlowDriver()
|
||||
healthy.finish.set()
|
||||
service._drivers.register("healthy", healthy)
|
||||
await asyncio.wait_for(service.run_once(now=Clock.current), timeout=3)
|
||||
assert driver.calls == 1
|
||||
assert healthy.calls == 1
|
||||
assert (await repo.get("task-2", user_id="user-1", thread_id="thread-1", thread_incarnation=None))["status"] == "completed"
|
||||
assert service._active_control_calls == {"task-1"}
|
||||
finish_cleanup.set()
|
||||
await asyncio.wait_for(asyncio.gather(*service._compensation_tasks), timeout=3)
|
||||
assert service._active_control_calls == set()
|
||||
assert (await read_task(repo))["status"] == "working", "draining driver persisted its stale result"
|
||||
finally:
|
||||
finish_cleanup.set()
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(asyncio.gather(cycle, *service._compensation_tasks, return_exceptions=True), timeout=3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_scans_recheck_control_capacity_at_execution(task_repo, monkeypatch):
|
||||
repo = task_repo
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
await create_working(repo, "task-2")
|
||||
claim = repo.claim_cancel_requests
|
||||
cancel_claimed = asyncio.Event()
|
||||
return_claim = asyncio.Event()
|
||||
|
||||
async def delayed_cancel_claim(**kwargs):
|
||||
rows = await claim(**kwargs)
|
||||
if rows:
|
||||
cancel_claimed.set()
|
||||
await return_claim.wait()
|
||||
return rows
|
||||
|
||||
monkeypatch.setattr(repo, "claim_cancel_requests", delayed_cancel_claim)
|
||||
driver = SlowDriver()
|
||||
service = make_service(repo, driver)
|
||||
service._max_concurrent_polls = 1
|
||||
first = asyncio.create_task(service.run_once(now=Clock.current))
|
||||
second = None
|
||||
try:
|
||||
await asyncio.wait_for(cancel_claimed.wait(), timeout=3)
|
||||
second = asyncio.create_task(service.run_once(now=Clock.current))
|
||||
await asyncio.wait_for(driver.started.wait(), timeout=3)
|
||||
return_claim.set()
|
||||
await asyncio.wait_for(asyncio.shield(first), timeout=1)
|
||||
assert driver.calls == 1, "distinct poll/cancel claims exceeded shared control capacity"
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(second, timeout=3)
|
||||
assert (await repo.get("task-2", user_id="user-1", thread_id="thread-1", thread_incarnation=None))["status"] == "completed"
|
||||
assert (await read_task(repo))["status"] == "working"
|
||||
finally:
|
||||
return_claim.set()
|
||||
driver.finish.set()
|
||||
await asyncio.wait_for(asyncio.gather(*[t for t in (first, second) if t is not None], return_exceptions=True), timeout=3)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_connection_cancel_does_not_persist_after_local_ownership_loss(task_repo, monkeypatch):
|
||||
repo = task_repo
|
||||
await repo.request_cancel("task-1", user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=Clock.current)
|
||||
driver = SlowDriver()
|
||||
applied = []
|
||||
apply = repo.apply_cancel_snapshot
|
||||
|
||||
async def missing_connection(_task):
|
||||
# A transport may finish with an error after ignoring lease-loss cancellation.
|
||||
lost = service_module._control_lease_lost.get()
|
||||
assert lost is not None
|
||||
lost.set()
|
||||
raise McpTaskConnectionUnavailableError("MCP task connection was removed")
|
||||
|
||||
async def record_apply(*args, **kwargs):
|
||||
applied.append(args)
|
||||
return await apply(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(driver, "cancel", missing_connection)
|
||||
monkeypatch.setattr(repo, "apply_cancel_snapshot", record_apply)
|
||||
await make_service(repo, driver)._run_cancellations(now=Clock.current)
|
||||
|
||||
assert applied == []
|
||||
assert (await read_task(repo))["status"] == "working"
|
||||
@ -47,6 +47,9 @@ class FakeRepository:
|
||||
self.created.append(kwargs)
|
||||
return {"id": kwargs["task_id"], **kwargs}
|
||||
|
||||
async def renew_control_claim(self, *_args, **_kwargs):
|
||||
return True
|
||||
|
||||
async def claim_due_tasks(self, **_kwargs):
|
||||
if self.claimed:
|
||||
return []
|
||||
@ -633,6 +636,7 @@ async def test_cancel_task_persists_request_without_calling_remote():
|
||||
async def test_cancel_failure_schedules_retry_from_call_completion_time():
|
||||
record = {**_claimed_row(), "cancel_attempt_count": 1}
|
||||
repo = SimpleNamespace(
|
||||
renew_control_claim=AsyncMock(return_value=True),
|
||||
claim_cancel_requests=AsyncMock(return_value=[record]),
|
||||
release_cancel_claim=AsyncMock(return_value=True),
|
||||
claim_due_tasks=AsyncMock(return_value=[]),
|
||||
@ -660,6 +664,7 @@ async def test_cancel_connection_unavailable_fails_locally_instead_of_retrying_f
|
||||
message = "Personal MCP task connection is missing, disabled or changed"
|
||||
record = {**_claimed_row(), "cancel_attempt_count": 1}
|
||||
repo = SimpleNamespace(
|
||||
renew_control_claim=AsyncMock(return_value=True),
|
||||
claim_cancel_requests=AsyncMock(return_value=[record]),
|
||||
apply_cancel_snapshot=AsyncMock(return_value=True),
|
||||
release_cancel_claim=AsyncMock(return_value=True),
|
||||
@ -699,6 +704,7 @@ async def test_cancel_recovery_failures_are_isolated_and_later_phases_continue(c
|
||||
return True
|
||||
|
||||
repo = SimpleNamespace(
|
||||
renew_control_claim=AsyncMock(return_value=True),
|
||||
claim_cancel_requests=AsyncMock(return_value=records),
|
||||
release_cancel_claim=AsyncMock(side_effect=release_cancel_claim),
|
||||
claim_due_tasks=AsyncMock(return_value=[]),
|
||||
@ -3655,6 +3661,9 @@ async def test_batch_completion_cancel_race_releases_once_for_100_rounds():
|
||||
|
||||
|
||||
class OrdinaryReleaseBatchRepository:
|
||||
async def renew_control_claim(self, *_args, **_kwargs):
|
||||
return True
|
||||
|
||||
def __init__(self, *, phase, outcome):
|
||||
self.phase = phase
|
||||
self.outcome = outcome
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user