mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-25 14:06:18 +00:00
fix(mcp): make durable task claims cancellation-safe (#4966)
* feat(mcp): re-scope to MCP task claim lifecycle only Keep PR #4966 a small, closed MCP lease/cancellation state-machine change and move RunJournal and Run lifecycle work into dedicated follow-ups. This branch contains only the MCP task claim lifecycle: - mcp task release/snapshot fencing by owner + per-claim lease token - phase-level single-flight poll/cancel/notification owners with retained handoff - routine cancellation no longer persisted as a task failure diagnostic - bounded ordinary release ownership retention past the drain deadline - 0018_mcp_task_lease_tokens migration + migration/bootstrap head assertions - wait_for_task_until helper (MCP uses it); worker-specific capture helper moved to the run-finalization follow-up RunJournal (journal.py + test_run_journal.py) and run lifecycle (manager/worker/store/run sql + run tests) are preserved on backup/cancellation-safety-full and will be raised as separate follow-ups. * fix(mcp): unblock claims after ambiguous handoff resolves A phase-level single-flight owner only guards an ambiguous claim outcome. Once the claim resolves, the phase owner is released immediately; the handoff may continue releasing returned rows as bounded, service-owned background work (transferred to _compensation_tasks on timeout). Per-claim token fencing rejects a late release against a newer claim generation, so a stuck release no longer locks the whole phase until process restart. - README: drop the stale progress-snapshot sentence from the bounded ordinary release description. - service: pop the identity-checked phase owner as soon as the claim outcome is known, then release returned rows with the bounded path; carry the release in _compensation_tasks if it exceeds the drain deadline. - mcp/AGENTS.md: document that only an unresolved claim outcome (not the handoff) blocks later phase scans, and that returned-row releases may continue in the background once the owner is released. - tests: pin that the phase owner is released before a stuck release finishes while the release stays service strong-owned. * refactor(mcp): remove unused single-record claim wrappers _poll_one, _cancel_one, and _notify_one are unreachable in production: the worker always processes claimed records through _run_claimed_batch, so these wrappers preserved a second, dead single-record lifecycle (state is None) whose only observable behavior was a wrapper-specific cancellation release. Remove the three wrappers and migrate the regressions that guarded their cancel/release invariants to exercise the production _run_claimed_batch path (operation=_*_one_claimed, release=_release_*_after_cancellation). The single wrapper-only "state is None" contract (test_poll_release_hang_without_batch) is deleted; all 11 remaining invariants (CancelledError preservation, repeated cancellation, poll-only token-fenced lease release, notification claimed vs dispatched phase release, hung compensation -> service ownership, and background compensation exactly-once observation) are now covered through the real batch lifecycle. * fix(mcp): fence claim-owned mutations against stale generations The per-claim token check in the ORM release/apply paths was only in the SELECT; the final write went out by primary key. On SQLite (where with_for_update() is a no-op) a mutation from an older claim generation could therefore clear a claim that a newer generation had reclaimed after lease expiry — the exact distributed lease-fencing failure the per-claim token was meant to prevent. Make every claim-owned mutation a single atomic conditional UPDATE with the owner and per-claim token in the WHERE clause (rowcount 0 => stale, return False, no mutation): - release_claim: atomic fence; record the poll-failure event after the fence wins (same transaction, holding the write lock). - apply_snapshot / apply_cancel_snapshot: atomic fence; record the event after. - finish_notification_run: atomic fence; use a CASE on event_version >> dispatch_version to keep a newer event pending for redelivery instead of swallowing it as delivered. Add one regression per path: a stale generation's release/apply/finish after a same-worker reclaim is rejected and never clears the newer claim. * test(mcp): pin the migration chain head to the lease-token revision 0026_mcp_task_lease_tokens becomes the alembic head, so the chain-head pin in the 0025 repair test had to move on. Follow the 0023 precedent there (single head plus expected predecessor) instead of pinning a literal head, and give the new revision its own migration test, which owns the pin and covers the nullable claim-token columns on upgrade and their removal on downgrade. * refactor(mcp): close cancellation cleanup leftovers * fix(mcp): retain cancelled release diagnostics * test(mcp): remove obsolete settled compensation case * test(mcp): cover interleaved lease reclaim races --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
dba3967177
commit
906c3d4554
@ -609,7 +609,7 @@ OpenViking memory backend; it does not replace automatic turn capture or recall.
|
||||
|
||||
The Gateway can adapt an MCP server's ordinary `submit` / `status` / `cancel` tools into durable background tasks. The Agent sees only the configured submit tool and a DeerFlow-local task ID; remote IDs are persisted before the submit call returns, while status and cancel stay internal to the runtime. Polling uses cross-worker leases, exponential retry backoff, scoped MCP sessions, bounded result storage, and restart recovery. A status-tool `isError` is retained as a bounded diagnostic and retried; servers report a permanent remote-task outcome through a normal structured result with `status: "failed"`. Remote poll hints are finite positive numbers capped at 24 hours, artifact-reference JSON is limited to 64 KiB, and task/server identifiers are validated against their durable SQL column limits before persistence. Input-required and terminal updates wake the current chat through idempotent Agent runs, while `list_background_tasks` and `cancel_background_task` let the Agent manage tasks without asking users for remote handles. Current-thread tasks are available through `GET /api/threads/{thread_id}/mcp-tasks`, its detail endpoint, and `POST /api/threads/{thread_id}/mcp-tasks/{task_id}/cancel`; when the task runtime actually starts, the Web UI exposes the same safe local view from the chat header with live status refresh, cancellation, and on-demand result, artifact, input-request, status-error, and cancellation-retry details. Default-disabled and memory-backend deployments hide that UI and do not poll the task endpoints. A failed remote cancellation remains queued with backoff, and its latest bounded error and attempt count stay visible in the expanded task card. Enable `mcp_tasks` in `config.yaml`, configure `task_toolsets` with exact raw tool names in `extensions_config.json`, and use a SQL database backend (`sqlite` or `postgres`). Task-enabled server connection, authentication, interceptor, timeout, or binding changes require a Gateway restart so Agent tool discovery and background calls cannot use different configuration versions. `input_required` is notification-only for now: DeerFlow can display the request but cannot yet submit the user's answer back to the remote task.
|
||||
|
||||
Notification launch and failed Agent-run deliveries use capped exponential backoff with a visible attempt count and stop after five failed attempts. A permanently rejected target such as a deleted chat is dead-lettered immediately instead of retried forever or recreated. Cancellation endpoints return after durably recording the request; the background service owns the potentially slow remote MCP call and its retry schedule.
|
||||
Notification launch and failed Agent-run deliveries use capped exponential backoff with a visible attempt count and stop after five failed attempts. When a bounded ordinary release exceeds its drain deadline, the service retains ownership until it settles. A permanently rejected target such as a deleted chat is dead-lettered immediately instead of retried forever or recreated. Cancellation endpoints return after durably recording the request; the background service owns the potentially slow remote MCP call and its retry schedule.
|
||||
|
||||
Notification runs keep their trusted delivery instruction separate from the framed, untrusted remote event payload. The process-started task runtime—not a hot config read—controls whether the task-management tools are exposed, so changing `mcp_tasks` requires a Gateway restart. When a skill's `allowed-tools` policy is active, `list_background_tasks` and `cancel_background_task` must be declared explicitly like other business tools.
|
||||
See the [MCP Server Guide](backend/docs/MCP_SERVER.md) for detailed instructions.
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@ -1,7 +1,7 @@
|
||||
### MCP System (`packages/harness/deerflow/mcp/`)
|
||||
|
||||
- Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management
|
||||
- **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. 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. `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` 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.
|
||||
- **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.
|
||||
- **Long-running ordinary task driver**: `extensions_config.json -> mcpServers.<server>.task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same `(server_name, user_id:thread_id)` stdio session scope; HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task.
|
||||
- **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior.
|
||||
|
||||
@ -48,6 +48,7 @@ class McpTaskRow(Base):
|
||||
next_notification_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
notification_lease_owner: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||
notification_lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
notification_lease_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
next_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), index=True, nullable=True)
|
||||
last_polled_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
last_poll_error: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
@ -55,6 +56,7 @@ class McpTaskRow(Base):
|
||||
consecutive_poll_error_count: Mapped[int] = mapped_column(Integer, default=0)
|
||||
lease_owner: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||
lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
lease_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
cancel_requested_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
cancel_attempt_count: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
|
||||
next_cancel_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
|
||||
@ -2,10 +2,11 @@ from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import or_, select, update
|
||||
from sqlalchemy import case, or_, select, update
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
@ -33,6 +34,11 @@ _TIMESTAMP_FIELDS = (
|
||||
_INFLIGHT_NOTIFICATION_STATUSES = frozenset({"claimed", "dispatched", "retry"})
|
||||
|
||||
|
||||
def _new_claim_token() -> str:
|
||||
"""Return a fresh per-claim token used to fence releases against reclaims."""
|
||||
return uuid.uuid4().hex
|
||||
|
||||
|
||||
def _notification_event(row: McpTaskRow, *, tracking_degraded: bool) -> dict[str, Any] | None:
|
||||
if row.status not in _ATTENTION_STATUS_VALUES and not tracking_degraded:
|
||||
return None
|
||||
@ -235,6 +241,7 @@ class McpTaskRepository:
|
||||
for row in rows:
|
||||
row.lease_owner = lease_owner
|
||||
row.lease_expires_at = lease_expires_at
|
||||
row.lease_token = _new_claim_token()
|
||||
row.poll_attempt_count += 1
|
||||
row.updated_at = now
|
||||
await session.commit()
|
||||
@ -245,6 +252,7 @@ class McpTaskRepository:
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
lease_token: str,
|
||||
status: str,
|
||||
result: Any | None,
|
||||
result_preview: str | None,
|
||||
@ -256,37 +264,44 @@ class McpTaskRepository:
|
||||
polled_at: datetime,
|
||||
) -> bool:
|
||||
async with self._sf() as session:
|
||||
stmt = (
|
||||
select(McpTaskRow)
|
||||
# Atomic fence: a poll result from an older generation must not
|
||||
# overwrite a claim a newer generation reclaimed after lease expiry.
|
||||
values: dict[str, Any] = {
|
||||
"status": status,
|
||||
"result": result,
|
||||
"result_preview": result_preview,
|
||||
"result_truncated": result_truncated,
|
||||
"result_artifact": result_artifact,
|
||||
"error": error,
|
||||
"input_required": input_required,
|
||||
"next_poll_at": next_poll_at,
|
||||
"last_polled_at": polled_at,
|
||||
"last_poll_error": None,
|
||||
"consecutive_poll_error_count": 0,
|
||||
"lease_owner": None,
|
||||
"lease_expires_at": None,
|
||||
"lease_token": None,
|
||||
"updated_at": polled_at,
|
||||
}
|
||||
if status in _TERMINAL_STATUS_VALUES:
|
||||
values["completed_at"] = polled_at
|
||||
update_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 >= polled_at,
|
||||
McpTaskRow.status.not_in(_TERMINAL_STATUS_VALUES),
|
||||
McpTaskRow.cancel_requested_at.is_(None),
|
||||
)
|
||||
.with_for_update()
|
||||
.values(**values)
|
||||
)
|
||||
row = (await session.execute(stmt)).scalar_one_or_none()
|
||||
if row is None:
|
||||
if not update_result.rowcount:
|
||||
return False
|
||||
row.status = status
|
||||
row.result = result
|
||||
row.result_preview = result_preview
|
||||
row.result_truncated = result_truncated
|
||||
row.result_artifact = result_artifact
|
||||
row.error = error
|
||||
row.input_required = input_required
|
||||
row.next_poll_at = next_poll_at
|
||||
row.last_polled_at = polled_at
|
||||
row.last_poll_error = None
|
||||
row.consecutive_poll_error_count = 0
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
row.updated_at = polled_at
|
||||
if status in _TERMINAL_STATUS_VALUES:
|
||||
row.completed_at = polled_at
|
||||
_record_event_if_changed(row, tracking_degraded=False, now=polled_at)
|
||||
row = (await session.execute(select(McpTaskRow).where(McpTaskRow.id == task_id))).scalar_one_or_none()
|
||||
if row is not None:
|
||||
_record_event_if_changed(row, tracking_degraded=False, now=polled_at)
|
||||
await session.commit()
|
||||
return True
|
||||
|
||||
@ -295,30 +310,77 @@ class McpTaskRepository:
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
lease_token: str,
|
||||
next_poll_at: datetime,
|
||||
error: str,
|
||||
tracking_degraded_after_errors: int = 3,
|
||||
) -> bool:
|
||||
async with self._sf() as session:
|
||||
stmt = select(McpTaskRow).where(McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner).with_for_update()
|
||||
row = (await session.execute(stmt)).scalar_one_or_none()
|
||||
if row is None:
|
||||
return False
|
||||
now = datetime.now(UTC)
|
||||
row.next_poll_at = next_poll_at
|
||||
row.last_poll_error = error
|
||||
row.consecutive_poll_error_count = int(row.consecutive_poll_error_count or 0) + 1
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
row.updated_at = now
|
||||
_record_event_if_changed(
|
||||
row,
|
||||
tracking_degraded=row.consecutive_poll_error_count >= tracking_degraded_after_errors,
|
||||
now=now,
|
||||
# Atomic fence: only clear the claim if the owner AND per-claim token
|
||||
# still match. ``with_for_update()`` is a no-op on SQLite, so the old
|
||||
# select-then-write could clear a claim that a newer generation had
|
||||
# reclaimed after lease expiry. A conditional UPDATE makes the fence
|
||||
# atomic: a stale release (rowcount 0) is a no-op and never mutates a
|
||||
# newer claim.
|
||||
update_result = await session.execute(
|
||||
update(McpTaskRow)
|
||||
.where(
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.lease_owner == lease_owner,
|
||||
McpTaskRow.lease_token == lease_token,
|
||||
)
|
||||
.values(
|
||||
next_poll_at=next_poll_at,
|
||||
last_poll_error=error,
|
||||
consecutive_poll_error_count=McpTaskRow.consecutive_poll_error_count + 1,
|
||||
lease_owner=None,
|
||||
lease_expires_at=None,
|
||||
lease_token=None,
|
||||
updated_at=now,
|
||||
)
|
||||
)
|
||||
if not update_result.rowcount:
|
||||
return False
|
||||
# The fence won and we hold the write lock, so read the released row
|
||||
# consistently and record the poll-failure tracking event.
|
||||
row = (await session.execute(select(McpTaskRow).where(McpTaskRow.id == task_id))).scalar_one_or_none()
|
||||
if row is not None:
|
||||
_record_event_if_changed(
|
||||
row,
|
||||
tracking_degraded=int(row.consecutive_poll_error_count or 0) >= tracking_degraded_after_errors,
|
||||
now=now,
|
||||
)
|
||||
await session.commit()
|
||||
return True
|
||||
|
||||
async def release_poll_claim_after_cancellation(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
lease_token: str,
|
||||
) -> bool:
|
||||
"""Release a cancelled poll's lease without recording a poll failure."""
|
||||
stmt = (
|
||||
update(McpTaskRow)
|
||||
.where(
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.lease_owner == lease_owner,
|
||||
McpTaskRow.lease_token == lease_token,
|
||||
)
|
||||
.values(
|
||||
lease_owner=None,
|
||||
lease_expires_at=None,
|
||||
lease_token=None,
|
||||
updated_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(stmt)
|
||||
await session.commit()
|
||||
return bool(result.rowcount)
|
||||
|
||||
async def request_cancel(
|
||||
self,
|
||||
task_id: str,
|
||||
@ -351,6 +413,7 @@ class McpTaskRepository:
|
||||
# lease so it cannot trigger a concurrent remote cancellation.
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
row.lease_token = None
|
||||
row.updated_at = requested_at
|
||||
await session.commit()
|
||||
return self._row_to_dict(row)
|
||||
@ -380,6 +443,7 @@ class McpTaskRepository:
|
||||
for row in rows:
|
||||
row.lease_owner = lease_owner
|
||||
row.lease_expires_at = expires_at
|
||||
row.lease_token = _new_claim_token()
|
||||
row.cancel_attempt_count = int(row.cancel_attempt_count or 0) + 1
|
||||
row.updated_at = now
|
||||
await session.commit()
|
||||
@ -390,6 +454,7 @@ class McpTaskRepository:
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
lease_token: str,
|
||||
status: str,
|
||||
result: Any | None,
|
||||
result_preview: str | None,
|
||||
@ -402,34 +467,38 @@ class McpTaskRepository:
|
||||
if status not in _TERMINAL_STATUS_VALUES:
|
||||
raise ValueError("A cancellation response must report a terminal task status")
|
||||
async with self._sf() as session:
|
||||
stmt = (
|
||||
select(McpTaskRow)
|
||||
update_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 >= completed_at,
|
||||
McpTaskRow.status.not_in(_TERMINAL_STATUS_VALUES),
|
||||
)
|
||||
.with_for_update()
|
||||
.values(
|
||||
status=status,
|
||||
result=result,
|
||||
result_preview=result_preview,
|
||||
result_truncated=result_truncated,
|
||||
result_artifact=result_artifact,
|
||||
error=error,
|
||||
input_required=input_required,
|
||||
next_poll_at=None,
|
||||
next_cancel_at=None,
|
||||
last_cancel_error=None,
|
||||
lease_owner=None,
|
||||
lease_expires_at=None,
|
||||
lease_token=None,
|
||||
completed_at=completed_at,
|
||||
updated_at=completed_at,
|
||||
)
|
||||
)
|
||||
row = (await session.execute(stmt)).scalar_one_or_none()
|
||||
if row is None:
|
||||
if not update_result.rowcount:
|
||||
return False
|
||||
row.status = status
|
||||
row.result = result
|
||||
row.result_preview = result_preview
|
||||
row.result_truncated = result_truncated
|
||||
row.result_artifact = result_artifact
|
||||
row.error = error
|
||||
row.input_required = input_required
|
||||
row.next_poll_at = None
|
||||
row.next_cancel_at = None
|
||||
row.last_cancel_error = None
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
row.completed_at = completed_at
|
||||
row.updated_at = completed_at
|
||||
_record_event_if_changed(row, tracking_degraded=False, now=completed_at)
|
||||
row = (await session.execute(select(McpTaskRow).where(McpTaskRow.id == task_id))).scalar_one_or_none()
|
||||
if row is not None:
|
||||
_record_event_if_changed(row, tracking_degraded=False, now=completed_at)
|
||||
await session.commit()
|
||||
return True
|
||||
|
||||
@ -438,17 +507,23 @@ class McpTaskRepository:
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
lease_token: str,
|
||||
next_cancel_at: datetime,
|
||||
error: str,
|
||||
error: str | None,
|
||||
) -> bool:
|
||||
stmt = (
|
||||
update(McpTaskRow)
|
||||
.where(McpTaskRow.id == task_id, McpTaskRow.lease_owner == lease_owner)
|
||||
.where(
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.lease_owner == lease_owner,
|
||||
McpTaskRow.lease_token == lease_token,
|
||||
)
|
||||
.values(
|
||||
next_cancel_at=next_cancel_at,
|
||||
last_cancel_error=error,
|
||||
lease_owner=None,
|
||||
lease_expires_at=None,
|
||||
lease_token=None,
|
||||
updated_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
@ -485,6 +560,7 @@ class McpTaskRepository:
|
||||
for row in rows:
|
||||
row.notification_lease_owner = lease_owner
|
||||
row.notification_lease_expires_at = expires_at
|
||||
row.notification_lease_token = _new_claim_token()
|
||||
rebuild_snapshot = row.notification_status in ("pending", "claimed") or (row.notification_status == "retry" and row.dispatch_version != row.event_version)
|
||||
if rebuild_snapshot:
|
||||
if row.dispatch_version != row.event_version:
|
||||
@ -506,6 +582,7 @@ class McpTaskRepository:
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
notification_lease_token: str,
|
||||
dispatch_version: int,
|
||||
run_id: str,
|
||||
now: datetime,
|
||||
@ -515,6 +592,7 @@ class McpTaskRepository:
|
||||
.where(
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.notification_lease_owner == lease_owner,
|
||||
McpTaskRow.notification_lease_token == notification_lease_token,
|
||||
McpTaskRow.notification_lease_expires_at >= now,
|
||||
McpTaskRow.dispatch_version == dispatch_version,
|
||||
McpTaskRow.notification_status.in_(("claimed", "retry")),
|
||||
@ -526,6 +604,7 @@ class McpTaskRepository:
|
||||
next_notification_at=now,
|
||||
notification_lease_owner=None,
|
||||
notification_lease_expires_at=None,
|
||||
notification_lease_token=None,
|
||||
updated_at=now,
|
||||
)
|
||||
)
|
||||
@ -539,8 +618,9 @@ class McpTaskRepository:
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
notification_lease_token: str,
|
||||
next_notification_at: datetime,
|
||||
error: str,
|
||||
error: str | None,
|
||||
replace_with_latest: bool,
|
||||
count_failure: bool = False,
|
||||
) -> bool:
|
||||
@ -550,6 +630,7 @@ class McpTaskRepository:
|
||||
"next_notification_at": next_notification_at,
|
||||
"notification_lease_owner": None,
|
||||
"notification_lease_expires_at": None,
|
||||
"notification_lease_token": None,
|
||||
"updated_at": datetime.now(UTC),
|
||||
}
|
||||
if replace_with_latest:
|
||||
@ -559,7 +640,15 @@ class McpTaskRepository:
|
||||
)
|
||||
if count_failure:
|
||||
values["notification_attempt_count"] = McpTaskRow.notification_attempt_count + 1
|
||||
stmt = update(McpTaskRow).where(McpTaskRow.id == task_id, McpTaskRow.notification_lease_owner == lease_owner).values(**values)
|
||||
stmt = (
|
||||
update(McpTaskRow)
|
||||
.where(
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.notification_lease_owner == lease_owner,
|
||||
McpTaskRow.notification_lease_token == notification_lease_token,
|
||||
)
|
||||
.values(**values)
|
||||
)
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(stmt)
|
||||
await session.commit()
|
||||
@ -570,57 +659,67 @@ class McpTaskRepository:
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
notification_lease_token: str,
|
||||
dispatch_version: int,
|
||||
delivered: bool,
|
||||
next_notification_at: datetime | None,
|
||||
error: str | None,
|
||||
now: datetime,
|
||||
) -> bool:
|
||||
if delivered:
|
||||
# A newer event may have arrived after this dispatch was queued; keep
|
||||
# it pending for redelivery instead of swallowing it as delivered.
|
||||
newer = McpTaskRow.event_version > dispatch_version
|
||||
values: dict[str, Any] = {
|
||||
"notified_version": dispatch_version,
|
||||
"notification_status": case((newer, "pending"), else_="delivered"),
|
||||
"dispatch_version": None,
|
||||
"dispatch_attempt": 0,
|
||||
"dispatch_event": None,
|
||||
"notification_run_id": None,
|
||||
"notification_error": None,
|
||||
"notification_attempt_count": 0,
|
||||
"next_notification_at": case((newer, now), else_=None),
|
||||
}
|
||||
else:
|
||||
values = {
|
||||
"notification_status": "retry",
|
||||
"dispatch_attempt": McpTaskRow.dispatch_attempt + 1,
|
||||
"notification_attempt_count": McpTaskRow.notification_attempt_count + 1,
|
||||
"notification_run_id": None,
|
||||
"notification_error": error,
|
||||
"next_notification_at": next_notification_at,
|
||||
}
|
||||
values.update(
|
||||
notification_lease_owner=None,
|
||||
notification_lease_expires_at=None,
|
||||
notification_lease_token=None,
|
||||
updated_at=now,
|
||||
)
|
||||
async with self._sf() as session:
|
||||
stmt = (
|
||||
select(McpTaskRow)
|
||||
result = await session.execute(
|
||||
update(McpTaskRow)
|
||||
.where(
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.notification_lease_owner == lease_owner,
|
||||
McpTaskRow.notification_lease_token == notification_lease_token,
|
||||
McpTaskRow.notification_lease_expires_at >= now,
|
||||
McpTaskRow.dispatch_version == dispatch_version,
|
||||
McpTaskRow.notification_status == "dispatched",
|
||||
)
|
||||
.with_for_update()
|
||||
.values(**values)
|
||||
)
|
||||
row = (await session.execute(stmt)).scalar_one_or_none()
|
||||
if row is None:
|
||||
return False
|
||||
if delivered:
|
||||
row.notified_version = dispatch_version
|
||||
row.notification_status = "pending" if row.event_version > dispatch_version else "delivered"
|
||||
row.dispatch_version = None
|
||||
row.dispatch_attempt = 0
|
||||
row.dispatch_event = None
|
||||
row.notification_run_id = None
|
||||
row.notification_error = None
|
||||
row.notification_attempt_count = 0
|
||||
row.next_notification_at = now if row.event_version > dispatch_version else None
|
||||
else:
|
||||
row.notification_status = "retry"
|
||||
row.dispatch_attempt = int(row.dispatch_attempt or 0) + 1
|
||||
row.notification_attempt_count = int(row.notification_attempt_count or 0) + 1
|
||||
row.notification_run_id = None
|
||||
row.notification_error = error
|
||||
row.next_notification_at = next_notification_at
|
||||
row.notification_lease_owner = None
|
||||
row.notification_lease_expires_at = None
|
||||
row.updated_at = now
|
||||
await session.commit()
|
||||
return True
|
||||
return bool(result.rowcount)
|
||||
|
||||
async def release_notification_lease(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
notification_lease_token: str,
|
||||
next_notification_at: datetime,
|
||||
error: str,
|
||||
error: str | None,
|
||||
count_failure: bool = False,
|
||||
) -> bool:
|
||||
"""Release unexpected notification work without changing its phase."""
|
||||
@ -629,6 +728,7 @@ class McpTaskRepository:
|
||||
"next_notification_at": next_notification_at,
|
||||
"notification_lease_owner": None,
|
||||
"notification_lease_expires_at": None,
|
||||
"notification_lease_token": None,
|
||||
"updated_at": datetime.now(UTC),
|
||||
}
|
||||
if count_failure:
|
||||
@ -638,6 +738,7 @@ class McpTaskRepository:
|
||||
.where(
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.notification_lease_owner == lease_owner,
|
||||
McpTaskRow.notification_lease_token == notification_lease_token,
|
||||
)
|
||||
.values(**values)
|
||||
)
|
||||
@ -651,6 +752,7 @@ class McpTaskRepository:
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
notification_lease_token: str,
|
||||
dispatch_version: int,
|
||||
error: str,
|
||||
count_failure: bool,
|
||||
@ -660,6 +762,7 @@ class McpTaskRepository:
|
||||
base_filters = (
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.notification_lease_owner == lease_owner,
|
||||
McpTaskRow.notification_lease_token == notification_lease_token,
|
||||
McpTaskRow.notification_lease_expires_at >= now,
|
||||
McpTaskRow.dispatch_version == dispatch_version,
|
||||
McpTaskRow.notification_status.in_(("claimed", "retry", "dispatched")),
|
||||
@ -670,6 +773,7 @@ class McpTaskRepository:
|
||||
"next_notification_at": None,
|
||||
"notification_lease_owner": None,
|
||||
"notification_lease_expires_at": None,
|
||||
"notification_lease_token": None,
|
||||
"dispatch_version": None,
|
||||
"dispatch_attempt": 0,
|
||||
"dispatch_event": None,
|
||||
@ -695,6 +799,7 @@ class McpTaskRepository:
|
||||
next_notification_at=now,
|
||||
notification_lease_owner=None,
|
||||
notification_lease_expires_at=None,
|
||||
notification_lease_token=None,
|
||||
dispatch_version=None,
|
||||
dispatch_attempt=0,
|
||||
dispatch_event=None,
|
||||
@ -710,6 +815,7 @@ class McpTaskRepository:
|
||||
task_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
notification_lease_token: str,
|
||||
dispatch_version: int,
|
||||
next_notification_at: datetime,
|
||||
now: datetime,
|
||||
@ -720,6 +826,7 @@ class McpTaskRepository:
|
||||
.where(
|
||||
McpTaskRow.id == task_id,
|
||||
McpTaskRow.notification_lease_owner == lease_owner,
|
||||
McpTaskRow.notification_lease_token == notification_lease_token,
|
||||
McpTaskRow.notification_lease_expires_at >= now,
|
||||
McpTaskRow.dispatch_version == dispatch_version,
|
||||
McpTaskRow.notification_status == "dispatched",
|
||||
@ -728,6 +835,7 @@ class McpTaskRepository:
|
||||
next_notification_at=next_notification_at,
|
||||
notification_lease_owner=None,
|
||||
notification_lease_expires_at=None,
|
||||
notification_lease_token=None,
|
||||
updated_at=now,
|
||||
)
|
||||
)
|
||||
|
||||
@ -24,10 +24,12 @@ The empty-DB path keeps using `create_all` because `Base.metadata` is the only a
|
||||
`0020_threads_meta_project_id` → `0021_batch_acceptance` →
|
||||
`0019_thread_incarnations` → `0022_scheduled_occurrence_seq` →
|
||||
`0023_run_change_seq` → `0023_user_preferences` →
|
||||
`0024_project_documents` → `0025_repair_run_change_seq` (current head). The preference
|
||||
`0024_project_documents` → `0025_repair_run_change_seq` →
|
||||
`0026_mcp_task_lease_tokens` (current head). The preference
|
||||
revision adds a separate owner/key table with a cascading users foreign key and
|
||||
does not alter users; the project-documents revision adds a new owner-scoped
|
||||
shelf table, so the bootstrap forward-compat floor is unchanged.
|
||||
shelf table, and the MCP lease-token revision adds two nullable token columns to
|
||||
`mcp_tasks`, so the bootstrap forward-compat floor is unchanged.
|
||||
The incarnation revision deliberately retains the exact id audited by the
|
||||
rollback-floor binary; Alembic orders revisions by `down_revision`, not by the
|
||||
numeric prefix.
|
||||
@ -152,7 +154,7 @@ on installs that never enabled it. The convention is:
|
||||
- `migrations/versions/0014_managed_subagents.py` — creates the deployment-level managed Subagent catalog table
|
||||
- `migrations/versions/0015_scheduled_task_enqueue.py` — interrupts legacy transient queued rows, adds durable scheduled-run launch leases and attempt counts, expands the one-active-occurrence index to `queued`/`launching`/`running`, and migrates the overlap policy from `skip` to `enqueue`; chains after `0014_managed_subagents`
|
||||
- `migrations/versions/0016_subagent_batches.py` — creates durable native-subagent batch and item tables, including owner/submission idempotency, item identity, lease/recovery state, and result fields
|
||||
- `migrations/versions/0017_personal_access_tokens.py` — creates the personal access token table for programmatic API access
|
||||
- `migrations/versions/0017_personal_access_tokens.py` — creates the personal access token table for programmatic API access; chains after `0016_subagent_batches`
|
||||
- `migrations/versions/0018_oauth_identity_pg_partial.py` — converts `idx_users_oauth_identity` to a partial index on Postgres (`postgresql_where`), matching what `UserRow.__table_args__` already builds via `create_all`; `0001_baseline` never applied the predicate on Postgres, so every `alembic upgrade head`-provisioned deployment carried a full index until this revision. Postgres-only, idempotent (checks `pg_index.indpred` directly), no-op on SQLite (already partial via `sqlite_where`) and on a DB where the index doesn't exist yet. Originally generated as 0017 and renumbered to 0018 after 0017_personal_access_tokens merged first and kept that slot
|
||||
- `migrations/versions/0019_projects.py` — creates the `projects` table (id/user_id/name/instructions/presentation/status + timestamps) for the Projects Phase-1 organization feature; chains after `0018_oauth_identity_pg_partial`
|
||||
- `migrations/versions/0020_threads_meta_project_id.py` — adds nullable `threads_meta.project_id` plus `ix_threads_meta_project_id` (no FK by design: project delete clears membership first, and the reserved `deerflow_project_id` metadata key stays in sync); chains after `0019_projects`
|
||||
@ -162,6 +164,7 @@ on installs that never enabled it. The convention is:
|
||||
- `migrations/versions/0023_run_change_seq.py` — adds `runs.change_seq`, its global singleton allocation clock, and owner-aware cursor indexes. Legacy rows remain at zero and page by run id; lifecycle, cancellation, and model-name mutations allocate monotonically increasing positions in their own transaction. Atomic replacement uses one position for every affected row. Progress snapshots and lease heartbeats do not advance the clock. `0023_user_preferences` follows this revision. Its migration test verifies membership in the single-head chain and the expected predecessor rather than pinning the latest head, so later migrations can extend the chain.
|
||||
- `migrations/versions/0024_project_documents.py` — creates the `project_documents` shelf table (id/project_id/user_id/name/stored_relpath/sha256/size_bytes, nullable promotion provenance and trash fields, timestamps) with indexes on project_id, user_id, sha256 and trashed_at; no DB-level foreign key on project_id by design (project delete trashes the shelf inside its own transaction). New table, so the bootstrap forward-compat floor is unchanged; chains after `0023_user_preferences` (renumbered from 0023 after the rebase)
|
||||
- `migrations/versions/0025_repair_run_change_seq.py` — heals databases that skipped `0023_run_change_seq` because it was inserted ahead of the already-shipped `0023_user_preferences` (#5516): re-applies the guarded `run_change_clock` table, `runs.change_seq` column, and cursor indexes on upgrade; no-ops on healthy shapes; chains after `0024_project_documents`. Its downgrade is intentionally a no-op: the repaired objects belong to ancestor `0023_run_change_seq`, remain required at 0024, and must retain their existing change positions. Only the original 0023 downgrade removes them. `tests/test_run_change_repair_history.py` reconstructs both pre-insertion published descendants and verifies historical upgrade, unchanged healthy positions, and usable run-store writes after downgrade and re-upgrade
|
||||
- `migrations/versions/0026_mcp_task_lease_tokens.py` — chains after `0025_repair_run_change_seq` and adds nullable `mcp_tasks.lease_token` / `notification_lease_token` columns so every poll, cancel, and notification mutation can be fenced to the exact claim generation
|
||||
- `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch provisioning decision, locked revision validation, and the narrow 0019 forward-compatibility exception
|
||||
- `extensions/loader.py::load_extensions` — registers each spec's `table_prefix` with `register_extension_table_prefix()`
|
||||
- Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter, including extension-owned tables), `tests/test_extension_loader.py::TestTablePrefixRegistration` (spec-to-filter wiring), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor), `tests/test_migration_0004_run_ownership_dedupe.py` + `tests/test_migration_0007_scheduled_run_active_dedupe.py` (dedupe-before-unique-index pre-steps), `tests/test_migration_0025_repair_run_change_seq.py` (issue #5516 skipped-revision heal)
|
||||
|
||||
@ -0,0 +1,31 @@
|
||||
"""fence MCP task claims by per-claim lease tokens.
|
||||
|
||||
Revision ID: 0026_mcp_task_lease_tokens
|
||||
Revises: 0025_repair_run_change_seq
|
||||
Create Date: 2026-08-27
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision: str = "0026_mcp_task_lease_tokens"
|
||||
down_revision: str | Sequence[str] | None = "0025_repair_run_change_seq"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
from deerflow.persistence.migrations._helpers import safe_add_column
|
||||
|
||||
safe_add_column("mcp_tasks", sa.Column("lease_token", sa.String(length=64), nullable=True))
|
||||
safe_add_column("mcp_tasks", sa.Column("notification_lease_token", sa.String(length=64), nullable=True))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
from deerflow.persistence.migrations._helpers import safe_drop_column
|
||||
|
||||
safe_drop_column("mcp_tasks", "notification_lease_token")
|
||||
safe_drop_column("mcp_tasks", "lease_token")
|
||||
24
backend/packages/harness/deerflow/runtime/cancellation.py
Normal file
24
backend/packages/harness/deerflow/runtime/cancellation.py
Normal file
@ -0,0 +1,24 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
async def wait_for_task_until( # noqa: UP047
|
||||
task: asyncio.Future[T], *, deadline: float
|
||||
) -> bool:
|
||||
"""Wait through repeated caller cancellation without cancelling task."""
|
||||
loop = asyncio.get_running_loop()
|
||||
while not task.done():
|
||||
remaining = deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
return False
|
||||
try:
|
||||
done, _ = await asyncio.wait({task}, timeout=remaining)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
if task in done:
|
||||
return True
|
||||
return True
|
||||
@ -2,6 +2,16 @@
|
||||
|
||||
Backend tests must preserve the runtime invariants they exercise without changing production execution topology.
|
||||
|
||||
## MCP claim fencing
|
||||
|
||||
`test_mcp_task_repository.py` covers same-worker reclaim during an in-flight
|
||||
release, poll/cancel snapshot, or notification completion. Use explicit events
|
||||
to pause the old operation at the persistence boundary, reclaim via the real
|
||||
repository, then verify the entire new row remains unchanged. Reclaiming before
|
||||
the old operation starts does not catch SQLite SELECT/ORM-flush races. Keep the
|
||||
old completion timestamp within its original lease so expiry cannot mask a
|
||||
missing token fence; always drain paused tasks and restore session patches.
|
||||
|
||||
## Executor starvation tests
|
||||
|
||||
`test_executor_starvation.py` covers the deterministic starvation semantics from RFC #4560:
|
||||
|
||||
@ -173,6 +173,18 @@ def test_local_guidance_files_contain_the_split_original_sections() -> None:
|
||||
assert "Before changing files in this directory" not in text, relative_text
|
||||
|
||||
|
||||
def test_mcp_task_lease_token_migration_is_documented() -> None:
|
||||
guidance = (REPO_ROOT / "backend" / "packages" / "harness" / "deerflow" / "persistence" / "migrations" / "AGENTS.md").read_text(encoding="utf-8")
|
||||
|
||||
for required in (
|
||||
"0026_mcp_task_lease_tokens.py",
|
||||
"0016_subagent_batches",
|
||||
"lease_token",
|
||||
"notification_lease_token",
|
||||
):
|
||||
assert required in guidance
|
||||
|
||||
|
||||
def test_repository_exposes_one_local_and_one_ci_entrypoint() -> None:
|
||||
makefile = (REPO_ROOT / "Makefile").read_text(encoding="utf-8")
|
||||
workflow = (REPO_ROOT / ".github" / "workflows" / "lint-check.yml").read_text(encoding="utf-8")
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import sqlite3
|
||||
from datetime import UTC, datetime, timedelta
|
||||
@ -6,6 +7,7 @@ import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy import event
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config
|
||||
@ -57,6 +59,147 @@ async def _create_working_task(
|
||||
)
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def _pause_claim_mutation(monkeypatch, operation):
|
||||
"""Pause an old mutation while another session reclaims its row.
|
||||
|
||||
The former SELECT/ORM-flush implementation must pause after its ownership
|
||||
read has loaded the old row. The atomic implementation pauses before its
|
||||
conditional UPDATE. Both leave the competing claim free to commit using
|
||||
the production SQLite engine, without replacing any persistence logic.
|
||||
"""
|
||||
entered = asyncio.Event()
|
||||
resume = asyncio.Event()
|
||||
original_execute = AsyncSession.execute
|
||||
intercepted = False
|
||||
|
||||
async def execute(session, statement, *args, **kwargs):
|
||||
nonlocal intercepted
|
||||
if asyncio.current_task() is not task or intercepted:
|
||||
return await original_execute(session, statement, *args, **kwargs)
|
||||
intercepted = True
|
||||
if statement.is_select:
|
||||
result = await original_execute(session, statement, *args, **kwargs)
|
||||
entered.set()
|
||||
await resume.wait()
|
||||
if statement.is_select:
|
||||
return result
|
||||
return await original_execute(session, statement, *args, **kwargs)
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
patch.setattr(AsyncSession, "execute", execute)
|
||||
task = asyncio.create_task(operation)
|
||||
try:
|
||||
await asyncio.wait_for(entered.wait(), timeout=5)
|
||||
yield task, resume
|
||||
finally:
|
||||
resume.set()
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ["release_claim", "apply_snapshot", "apply_cancel_snapshot"])
|
||||
async def test_interleaved_reclaim_fences_inflight_poll_and_cancel_mutations(tmp_path, monkeypatch, operation):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
task_id = "interleaved-claim"
|
||||
await _create_working_task(repo, task_id=task_id, now=now)
|
||||
claim = repo.claim_due_tasks
|
||||
if operation == "apply_cancel_snapshot":
|
||||
await repo.request_cancel(task_id, user_id="user-1", thread_id="thread-1", requested_at=now)
|
||||
claim = repo.claim_cancel_requests
|
||||
first = await claim(now=now, lease_owner="worker-1", lease_seconds=60, limit=1)
|
||||
kwargs = {"lease_owner": "worker-1", "lease_token": first[0]["lease_token"]}
|
||||
if operation == "release_claim":
|
||||
kwargs.update(next_poll_at=now + timedelta(seconds=30), error="old poll failed")
|
||||
else:
|
||||
kwargs.update(
|
||||
status="cancelled" if operation == "apply_cancel_snapshot" else "completed",
|
||||
result={"stale": True},
|
||||
result_preview="old result",
|
||||
result_truncated=False,
|
||||
result_artifact=None,
|
||||
error=None,
|
||||
input_required=None,
|
||||
)
|
||||
if operation == "apply_snapshot":
|
||||
kwargs.update(next_poll_at=None, polled_at=now)
|
||||
else:
|
||||
kwargs.update(completed_at=now)
|
||||
|
||||
async with _pause_claim_mutation(monkeypatch, getattr(repo, operation)(task_id, **kwargs)) as (pending, resume):
|
||||
# Advance only the claim clock, not the stale operation's completion
|
||||
# timestamp: expiry must not reject it before the token fence is tested.
|
||||
second = await asyncio.wait_for(claim(now=now + timedelta(seconds=61), lease_owner="worker-1", lease_seconds=60, limit=1), timeout=5)
|
||||
assert len(second) == 1
|
||||
assert second[0]["lease_token"] != first[0]["lease_token"]
|
||||
before = await repo.get(task_id, user_id="user-1")
|
||||
resume.set()
|
||||
applied = await asyncio.wait_for(pending, timeout=5)
|
||||
|
||||
# Check the entire row, including scheduling, errors, results and event
|
||||
# versions, not just the new lease: stale work must have no side effects.
|
||||
assert await repo.get(task_id, user_id="user-1") == before
|
||||
assert applied is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("delivered", [True, False], ids=["success", "failure"])
|
||||
async def test_interleaved_reclaim_fences_inflight_notification_completion(tmp_path, monkeypatch, delivered):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
task_id = "interleaved-notification"
|
||||
await _create_working_task(repo, task_id=task_id, now=now)
|
||||
poll = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
assert await repo.apply_snapshot(
|
||||
task_id,
|
||||
lease_owner="poller",
|
||||
lease_token=poll[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"done": True},
|
||||
result_preview=None,
|
||||
result_truncated=False,
|
||||
result_artifact=None,
|
||||
error=None,
|
||||
input_required=None,
|
||||
next_poll_at=None,
|
||||
polled_at=now,
|
||||
)
|
||||
claim_kwargs = {"lease_owner": "notifier", "lease_seconds": 60, "limit": 1, "tracking_degraded_after_errors": 3}
|
||||
launch = await repo.claim_notification_work(now=now, **claim_kwargs)
|
||||
assert await repo.mark_notification_dispatched(
|
||||
task_id,
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=launch[0]["notification_lease_token"],
|
||||
dispatch_version=launch[0]["dispatch_version"],
|
||||
run_id="notification-run",
|
||||
now=now,
|
||||
)
|
||||
first = await repo.claim_notification_work(now=now, **claim_kwargs)
|
||||
operation = repo.finish_notification_run(
|
||||
task_id,
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=first[0]["notification_lease_token"],
|
||||
dispatch_version=first[0]["dispatch_version"],
|
||||
delivered=delivered,
|
||||
next_notification_at=None if delivered else now + timedelta(seconds=30),
|
||||
error=None if delivered else "old notification failed",
|
||||
now=now,
|
||||
)
|
||||
async with _pause_claim_mutation(monkeypatch, operation) as (pending, resume):
|
||||
second = await asyncio.wait_for(repo.claim_notification_work(now=now + timedelta(seconds=61), **claim_kwargs), timeout=5)
|
||||
assert len(second) == 1
|
||||
assert second[0]["notification_lease_token"] != first[0]["notification_lease_token"]
|
||||
before = await repo.get(task_id, user_id="user-1")
|
||||
resume.set()
|
||||
applied = await asyncio.wait_for(pending, timeout=5)
|
||||
|
||||
assert await repo.get(task_id, user_id="user-1") == before
|
||||
assert applied is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_task_writer_leaves_thread_incarnation_null(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
@ -301,7 +444,7 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-2", now=now)
|
||||
await repo.claim_due_tasks(
|
||||
claimed = await repo.claim_due_tasks(
|
||||
now=now,
|
||||
lease_owner="worker-new",
|
||||
lease_seconds=60,
|
||||
@ -311,6 +454,7 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task
|
||||
stale_applied = await repo.apply_snapshot(
|
||||
"task-2",
|
||||
lease_owner="worker-old",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
status="failed",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
@ -326,6 +470,7 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task
|
||||
applied = await repo.apply_snapshot(
|
||||
"task-2",
|
||||
lease_owner="worker-new",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"report": "ready"},
|
||||
result_preview=None,
|
||||
@ -365,7 +510,7 @@ async def test_apply_snapshot_rejects_result_after_same_workers_lease_expires(tm
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-expired", now=now)
|
||||
await repo.claim_due_tasks(
|
||||
claimed = await repo.claim_due_tasks(
|
||||
now=now,
|
||||
lease_owner="worker-1",
|
||||
lease_seconds=60,
|
||||
@ -375,6 +520,7 @@ async def test_apply_snapshot_rejects_result_after_same_workers_lease_expires(tm
|
||||
applied = await repo.apply_snapshot(
|
||||
"task-expired",
|
||||
lease_owner="worker-1",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"report": "stale"},
|
||||
result_preview=None,
|
||||
@ -398,7 +544,7 @@ async def test_input_required_is_persisted_and_remains_scheduled_for_slow_pollin
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-3", now=now)
|
||||
await repo.claim_due_tasks(
|
||||
claimed = await repo.claim_due_tasks(
|
||||
now=now,
|
||||
lease_owner="worker-1",
|
||||
lease_seconds=60,
|
||||
@ -408,6 +554,7 @@ async def test_input_required_is_persisted_and_remains_scheduled_for_slow_pollin
|
||||
applied = await repo.apply_snapshot(
|
||||
"task-3",
|
||||
lease_owner="worker-1",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
status="input_required",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
@ -432,7 +579,7 @@ async def test_release_claim_retries_transient_poll_failure(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-4", now=now)
|
||||
await repo.claim_due_tasks(
|
||||
claimed = await repo.claim_due_tasks(
|
||||
now=now,
|
||||
lease_owner="worker-1",
|
||||
lease_seconds=60,
|
||||
@ -443,6 +590,7 @@ async def test_release_claim_retries_transient_poll_failure(tmp_path):
|
||||
released = await repo.release_claim(
|
||||
"task-4",
|
||||
lease_owner="worker-1",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
next_poll_at=retry_at,
|
||||
error="temporary network failure",
|
||||
)
|
||||
@ -456,6 +604,231 @@ async def test_release_claim_retries_transient_poll_failure(tmp_path):
|
||||
assert stored["lease_owner"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_claim_after_same_worker_reclaim_cannot_clear_new_claim(tmp_path):
|
||||
"""A stale release from an older generation must be a no-op once the same
|
||||
worker reclaims the task with a fresh per-claim token (token fencing)."""
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-fence", now=now)
|
||||
claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10)
|
||||
old_token = claimed[0]["lease_token"]
|
||||
|
||||
reclaim_at = now + timedelta(seconds=61) # after the 60s lease expires
|
||||
reclaimed = await repo.claim_due_tasks(now=reclaim_at, lease_owner="worker-1", lease_seconds=61, limit=10)
|
||||
assert reclaimed
|
||||
new_token = reclaimed[0]["lease_token"]
|
||||
assert new_token != old_token
|
||||
|
||||
# The stale release (old owner + old token) must not clear the new claim.
|
||||
released = await repo.release_claim(
|
||||
"task-fence",
|
||||
lease_owner="worker-1",
|
||||
lease_token=old_token,
|
||||
next_poll_at=reclaim_at + timedelta(seconds=30),
|
||||
error="stale release",
|
||||
)
|
||||
assert released is False
|
||||
|
||||
stored = await repo.get("task-fence", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["lease_owner"] == "worker-1"
|
||||
assert stored["lease_token"] == new_token
|
||||
assert stored["lease_expires_at"] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_snapshot_after_same_worker_reclaim_cannot_clear_new_claim(tmp_path):
|
||||
"""A poll snapshot from an older generation must not overwrite a newer claim."""
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-apply-fence", now=now)
|
||||
claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10)
|
||||
old_token = claimed[0]["lease_token"]
|
||||
|
||||
reclaim_at = now + timedelta(seconds=61)
|
||||
reclaimed = await repo.claim_due_tasks(now=reclaim_at, lease_owner="worker-1", lease_seconds=61, limit=10)
|
||||
new_token = reclaimed[0]["lease_token"]
|
||||
assert new_token != old_token
|
||||
|
||||
applied = await repo.apply_snapshot(
|
||||
"task-apply-fence",
|
||||
lease_owner="worker-1",
|
||||
lease_token=old_token,
|
||||
status="completed",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
result_truncated=False,
|
||||
result_artifact=None,
|
||||
error=None,
|
||||
input_required=None,
|
||||
next_poll_at=None,
|
||||
polled_at=reclaim_at,
|
||||
)
|
||||
assert applied is False
|
||||
|
||||
stored = await repo.get("task-apply-fence", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["lease_owner"] == "worker-1"
|
||||
assert stored["lease_token"] == new_token
|
||||
assert stored["status"] == "working"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_cancel_snapshot_after_same_worker_reclaim_cannot_clear_new_claim(tmp_path):
|
||||
"""A cancel snapshot from an older generation must not overwrite a newer claim."""
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-cancel-fence", now=now)
|
||||
claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10)
|
||||
old_token = claimed[0]["lease_token"]
|
||||
|
||||
reclaim_at = now + timedelta(seconds=61)
|
||||
reclaimed = await repo.claim_due_tasks(now=reclaim_at, lease_owner="worker-1", lease_seconds=61, limit=10)
|
||||
new_token = reclaimed[0]["lease_token"]
|
||||
assert new_token != old_token
|
||||
|
||||
applied = await repo.apply_cancel_snapshot(
|
||||
"task-cancel-fence",
|
||||
lease_owner="worker-1",
|
||||
lease_token=old_token,
|
||||
status="cancelled",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
result_truncated=False,
|
||||
result_artifact=None,
|
||||
error=None,
|
||||
input_required=None,
|
||||
completed_at=reclaim_at,
|
||||
)
|
||||
assert applied is False
|
||||
|
||||
stored = await repo.get("task-cancel-fence", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["lease_owner"] == "worker-1"
|
||||
assert stored["lease_token"] == new_token
|
||||
assert stored["status"] == "working"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finish_notification_run_after_reclaim_cannot_clear_new_claim(tmp_path):
|
||||
"""A stale notification finish must not clear a newer notification lease."""
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-notify-fence", now=now)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-notify-fence",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="input_required",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
result_truncated=False,
|
||||
result_artifact=None,
|
||||
error=None,
|
||||
input_required={"prompt": "Approve?"},
|
||||
next_poll_at=now,
|
||||
polled_at=now,
|
||||
)
|
||||
first = await repo.claim_notification_work(
|
||||
now=now,
|
||||
lease_owner="notifier",
|
||||
lease_seconds=60,
|
||||
limit=1,
|
||||
tracking_degraded_after_errors=3,
|
||||
)
|
||||
await repo.mark_notification_dispatched(
|
||||
"task-notify-fence",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=first[0]["notification_lease_token"],
|
||||
dispatch_version=first[0]["dispatch_version"],
|
||||
run_id="notify-run-1",
|
||||
now=now,
|
||||
)
|
||||
reclaimed = await repo.claim_notification_work(
|
||||
now=now,
|
||||
lease_owner="notifier",
|
||||
lease_seconds=60,
|
||||
limit=1,
|
||||
tracking_degraded_after_errors=3,
|
||||
)
|
||||
assert reclaimed
|
||||
new_notify_token = reclaimed[0]["notification_lease_token"]
|
||||
|
||||
finished = await repo.finish_notification_run(
|
||||
"task-notify-fence",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token="stale-notify-token",
|
||||
dispatch_version=reclaimed[0]["dispatch_version"],
|
||||
delivered=True,
|
||||
next_notification_at=None,
|
||||
error=None,
|
||||
now=now,
|
||||
)
|
||||
assert finished is False
|
||||
|
||||
stored = await repo.get("task-notify-fence", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["notification_lease_owner"] == "notifier"
|
||||
assert stored["notification_lease_token"] == new_notify_token
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_poll_claim_after_cancellation_preserves_poll_failure_state(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-cancelled-poll", now=now)
|
||||
claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10)
|
||||
retry_at = now + timedelta(seconds=30)
|
||||
await repo.release_claim(
|
||||
"task-cancelled-poll",
|
||||
lease_owner="worker-1",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
next_poll_at=retry_at,
|
||||
error="temporary network failure",
|
||||
)
|
||||
before = await repo.get("task-cancelled-poll", user_id="user-1")
|
||||
assert before is not None
|
||||
|
||||
reclaimed = await repo.claim_due_tasks(now=retry_at, lease_owner="worker-2", lease_seconds=60, limit=10)
|
||||
released = await repo.release_poll_claim_after_cancellation(
|
||||
"task-cancelled-poll",
|
||||
lease_owner="worker-2",
|
||||
lease_token=reclaimed[0]["lease_token"],
|
||||
)
|
||||
|
||||
assert released is True
|
||||
stored = await repo.get("task-cancelled-poll", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["next_poll_at"] == before["next_poll_at"]
|
||||
assert stored["last_poll_error"] == before["last_poll_error"]
|
||||
assert stored["consecutive_poll_error_count"] == before["consecutive_poll_error_count"]
|
||||
assert stored["poll_attempt_count"] == before["poll_attempt_count"] + 1
|
||||
assert stored["lease_owner"] is None
|
||||
assert stored["lease_expires_at"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_poll_claim_after_cancellation_requires_current_owner(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-stale-cancel", now=now)
|
||||
claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-current", lease_seconds=60, limit=10)
|
||||
|
||||
released = await repo.release_poll_claim_after_cancellation(
|
||||
"task-stale-cancel",
|
||||
lease_owner="worker-stale",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
)
|
||||
|
||||
assert released is False
|
||||
stored = await repo.get("task-stale-cancel", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["lease_owner"] == "worker-current"
|
||||
assert stored["lease_expires_at"] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
@ -463,10 +836,11 @@ async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp
|
||||
await _create_working_task(repo, task_id="task-6", now=now)
|
||||
|
||||
for expected_errors in (1, 2):
|
||||
await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10)
|
||||
claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10)
|
||||
await repo.release_claim(
|
||||
"task-6",
|
||||
lease_owner="worker-1",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
next_poll_at=now - timedelta(seconds=1),
|
||||
error="temporary network failure",
|
||||
)
|
||||
@ -474,10 +848,11 @@ async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp
|
||||
assert stored is not None
|
||||
assert stored["consecutive_poll_error_count"] == expected_errors
|
||||
|
||||
await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10)
|
||||
claimed = await repo.claim_due_tasks(now=now, lease_owner="worker-1", lease_seconds=60, limit=10)
|
||||
applied = await repo.apply_snapshot(
|
||||
"task-6",
|
||||
lease_owner="worker-1",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
status="working",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
@ -500,10 +875,11 @@ async def test_notification_snapshot_is_versioned_and_not_overwritten_in_flight(
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-notify", now=now)
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-notify",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="input_required",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
@ -525,10 +901,11 @@ async def test_notification_snapshot_is_versioned_and_not_overwritten_in_flight(
|
||||
assert first[0]["dispatch_version"] == 1
|
||||
assert first[0]["dispatch_event"]["input_required"] == {"prompt": "Approve?"}
|
||||
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-notify",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"done": True},
|
||||
result_preview=None,
|
||||
@ -548,11 +925,12 @@ async def test_notification_snapshot_is_versioned_and_not_overwritten_in_flight(
|
||||
await repo.mark_notification_dispatched(
|
||||
"task-notify",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=first[0]["notification_lease_token"],
|
||||
dispatch_version=1,
|
||||
run_id="notify-run-1",
|
||||
now=now,
|
||||
)
|
||||
await repo.claim_notification_work(
|
||||
dispatched_claim = await repo.claim_notification_work(
|
||||
now=now,
|
||||
lease_owner="notifier",
|
||||
lease_seconds=60,
|
||||
@ -562,6 +940,7 @@ async def test_notification_snapshot_is_versioned_and_not_overwritten_in_flight(
|
||||
await repo.finish_notification_run(
|
||||
"task-notify",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=dispatched_claim[0]["notification_lease_token"],
|
||||
dispatch_version=1,
|
||||
delivered=True,
|
||||
next_notification_at=None,
|
||||
@ -584,10 +963,11 @@ async def test_notification_retry_rebuilds_a_newer_event_and_resets_its_budget(t
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-retry-latest", now=now)
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-retry-latest",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="input_required",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
@ -608,11 +988,12 @@ async def test_notification_retry_rebuilds_a_newer_event_and_resets_its_budget(t
|
||||
await repo.mark_notification_dispatched(
|
||||
"task-retry-latest",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=first[0]["notification_lease_token"],
|
||||
dispatch_version=first[0]["dispatch_version"],
|
||||
run_id="notify-run-1",
|
||||
now=now,
|
||||
)
|
||||
await repo.claim_notification_work(
|
||||
dispatched_claim = await repo.claim_notification_work(
|
||||
now=now,
|
||||
lease_owner="notifier",
|
||||
lease_seconds=60,
|
||||
@ -623,6 +1004,7 @@ async def test_notification_retry_rebuilds_a_newer_event_and_resets_its_budget(t
|
||||
await repo.finish_notification_run(
|
||||
"task-retry-latest",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=dispatched_claim[0]["notification_lease_token"],
|
||||
dispatch_version=first[0]["dispatch_version"],
|
||||
delivered=False,
|
||||
next_notification_at=retry_at,
|
||||
@ -635,10 +1017,11 @@ async def test_notification_retry_rebuilds_a_newer_event_and_resets_its_budget(t
|
||||
assert failed["dispatch_attempt"] == 1
|
||||
assert failed["notification_attempt_count"] == 1
|
||||
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-retry-latest",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"done": True},
|
||||
result_preview=None,
|
||||
@ -668,10 +1051,11 @@ async def test_unexpected_notification_failure_releases_lease_without_changing_p
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-notify-release", now=now)
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-notify-release",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="input_required",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
@ -682,7 +1066,7 @@ async def test_unexpected_notification_failure_releases_lease_without_changing_p
|
||||
next_poll_at=now,
|
||||
polled_at=now,
|
||||
)
|
||||
await repo.claim_notification_work(
|
||||
claimed = await repo.claim_notification_work(
|
||||
now=now,
|
||||
lease_owner="notifier",
|
||||
lease_seconds=60,
|
||||
@ -694,6 +1078,7 @@ async def test_unexpected_notification_failure_releases_lease_without_changing_p
|
||||
assert await repo.release_notification_lease(
|
||||
"task-notify-release",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=claimed[0]["notification_lease_token"],
|
||||
next_notification_at=retry_at,
|
||||
error="run store unavailable",
|
||||
)
|
||||
@ -711,10 +1096,11 @@ async def test_notification_launch_failure_counts_and_reclaims_latest_snapshot(t
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-launch-retry", now=now)
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-launch-retry",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="input_required",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
@ -737,6 +1123,7 @@ async def test_notification_launch_failure_counts_and_reclaims_latest_snapshot(t
|
||||
assert await repo.release_notification_claim(
|
||||
"task-launch-retry",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=first[0]["notification_lease_token"],
|
||||
next_notification_at=retry_at,
|
||||
error="run store unavailable",
|
||||
replace_with_latest=True,
|
||||
@ -765,10 +1152,11 @@ async def test_permanent_notification_failure_is_not_reclaimed(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-dead-letter", now=now)
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-dead-letter",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"done": True},
|
||||
result_preview=None,
|
||||
@ -790,6 +1178,7 @@ async def test_permanent_notification_failure_is_not_reclaimed(tmp_path):
|
||||
assert await repo.dead_letter_notification(
|
||||
"task-dead-letter",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=claimed[0]["notification_lease_token"],
|
||||
dispatch_version=claimed[0]["dispatch_version"],
|
||||
error="Thread deleted-thread not found",
|
||||
count_failure=True,
|
||||
@ -818,10 +1207,11 @@ async def test_dispatched_notification_can_be_dead_lettered_after_retry_budget(t
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-dispatched-budget", now=now)
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-dispatched-budget",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"done": True},
|
||||
result_preview=None,
|
||||
@ -843,6 +1233,7 @@ async def test_dispatched_notification_can_be_dead_lettered_after_retry_budget(t
|
||||
assert await repo.mark_notification_dispatched(
|
||||
"task-dispatched-budget",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=first[0]["notification_lease_token"],
|
||||
dispatch_version=dispatch_version,
|
||||
run_id="notify-run-1",
|
||||
now=now,
|
||||
@ -859,6 +1250,7 @@ async def test_dispatched_notification_can_be_dead_lettered_after_retry_budget(t
|
||||
assert await repo.dead_letter_notification(
|
||||
"task-dispatched-budget",
|
||||
lease_owner="budget-checker",
|
||||
notification_lease_token=claimed[0]["notification_lease_token"],
|
||||
dispatch_version=dispatch_version,
|
||||
error="Notification delivery stopped after 5 failed attempts",
|
||||
count_failure=False,
|
||||
@ -876,10 +1268,11 @@ async def test_dead_lettering_dispatched_snapshot_preserves_newer_event(tmp_path
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-dispatched-latest", now=now)
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-dispatched-latest",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="input_required",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
@ -901,15 +1294,17 @@ async def test_dead_lettering_dispatched_snapshot_preserves_newer_event(tmp_path
|
||||
assert await repo.mark_notification_dispatched(
|
||||
"task-dispatched-latest",
|
||||
lease_owner="notifier",
|
||||
notification_lease_token=first[0]["notification_lease_token"],
|
||||
dispatch_version=dispatch_version,
|
||||
run_id="notify-run-1",
|
||||
now=now,
|
||||
)
|
||||
|
||||
await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-dispatched-latest",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"done": True},
|
||||
result_preview=None,
|
||||
@ -931,6 +1326,7 @@ async def test_dead_lettering_dispatched_snapshot_preserves_newer_event(tmp_path
|
||||
assert await repo.dead_letter_notification(
|
||||
"task-dispatched-latest",
|
||||
lease_owner="budget-checker",
|
||||
notification_lease_token=claimed[0]["notification_lease_token"],
|
||||
dispatch_version=dispatch_version,
|
||||
error="old snapshot exhausted its retry budget",
|
||||
count_failure=False,
|
||||
@ -958,7 +1354,7 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-cancel", now=now)
|
||||
await repo.claim_due_tasks(now=now, lease_owner="stale-poller", lease_seconds=60, limit=1)
|
||||
stale_poll_claim = await repo.claim_due_tasks(now=now, lease_owner="stale-poller", lease_seconds=60, limit=1)
|
||||
|
||||
requested = await repo.request_cancel(
|
||||
"task-cancel",
|
||||
@ -972,6 +1368,7 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa
|
||||
await repo.apply_snapshot(
|
||||
"task-cancel",
|
||||
lease_owner="stale-poller",
|
||||
lease_token=stale_poll_claim[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"stale": True},
|
||||
result_preview=None,
|
||||
@ -1006,6 +1403,7 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa
|
||||
assert await repo.apply_cancel_snapshot(
|
||||
"task-cancel",
|
||||
lease_owner="canceller",
|
||||
lease_token=claimed[0]["lease_token"],
|
||||
status="cancelled",
|
||||
result=None,
|
||||
result_preview=None,
|
||||
@ -1019,3 +1417,188 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa
|
||||
assert stored is not None
|
||||
assert stored["status"] == "cancelled"
|
||||
assert stored["notification_status"] == "pending"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_late_poll_release_after_same_worker_reclaim_is_fenced(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-late-poll", now=now)
|
||||
|
||||
first = await repo.claim_due_tasks(
|
||||
now=now,
|
||||
lease_owner="worker-same",
|
||||
lease_seconds=1,
|
||||
limit=10,
|
||||
)
|
||||
assert [row["id"] for row in first] == ["task-late-poll"]
|
||||
first_token = first[0]["lease_token"]
|
||||
assert first_token
|
||||
|
||||
reclaimed = await repo.claim_due_tasks(
|
||||
now=now + timedelta(seconds=5),
|
||||
lease_owner="worker-same",
|
||||
lease_seconds=60,
|
||||
limit=10,
|
||||
)
|
||||
assert [row["id"] for row in reclaimed] == ["task-late-poll"]
|
||||
assert reclaimed[0]["lease_token"] != first_token
|
||||
|
||||
released = await repo.release_poll_claim_after_cancellation(
|
||||
"task-late-poll",
|
||||
lease_owner="worker-same",
|
||||
lease_token=first_token,
|
||||
)
|
||||
assert released is False
|
||||
|
||||
stored = await repo.get("task-late-poll", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["lease_owner"] == "worker-same"
|
||||
assert stored["lease_token"] == reclaimed[0]["lease_token"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_late_cancel_release_after_same_worker_reclaim_is_fenced(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-late-cancel", now=now)
|
||||
await repo.request_cancel(
|
||||
"task-late-cancel",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
requested_at=now,
|
||||
)
|
||||
|
||||
first = await repo.claim_cancel_requests(
|
||||
now=now,
|
||||
lease_owner="worker-same",
|
||||
lease_seconds=1,
|
||||
limit=1,
|
||||
)
|
||||
assert [row["id"] for row in first] == ["task-late-cancel"]
|
||||
first_token = first[0]["lease_token"]
|
||||
assert first_token
|
||||
|
||||
reclaimed = await repo.claim_cancel_requests(
|
||||
now=now + timedelta(seconds=5),
|
||||
lease_owner="worker-same",
|
||||
lease_seconds=60,
|
||||
limit=1,
|
||||
)
|
||||
assert [row["id"] for row in reclaimed] == ["task-late-cancel"]
|
||||
assert reclaimed[0]["lease_token"] != first_token
|
||||
|
||||
released = await repo.release_cancel_claim(
|
||||
"task-late-cancel",
|
||||
lease_owner="worker-same",
|
||||
lease_token=first_token,
|
||||
next_cancel_at=now + timedelta(seconds=30),
|
||||
error="cancelled",
|
||||
)
|
||||
assert released is False
|
||||
|
||||
stored = await repo.get("task-late-cancel", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["lease_owner"] == "worker-same"
|
||||
assert stored["lease_token"] == reclaimed[0]["lease_token"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_late_notification_release_after_same_worker_reclaim_is_fenced(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-late-notify", now=now)
|
||||
poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
|
||||
await repo.apply_snapshot(
|
||||
"task-late-notify",
|
||||
lease_owner="poller",
|
||||
lease_token=poll_claim[0]["lease_token"],
|
||||
status="completed",
|
||||
result={"done": True},
|
||||
result_preview=None,
|
||||
result_truncated=False,
|
||||
result_artifact=None,
|
||||
error=None,
|
||||
input_required=None,
|
||||
next_poll_at=None,
|
||||
polled_at=now,
|
||||
)
|
||||
|
||||
first = await repo.claim_notification_work(
|
||||
now=now,
|
||||
lease_owner="notifier-same",
|
||||
lease_seconds=1,
|
||||
limit=1,
|
||||
tracking_degraded_after_errors=3,
|
||||
)
|
||||
assert [row["id"] for row in first] == ["task-late-notify"]
|
||||
first_token = first[0]["notification_lease_token"]
|
||||
assert first_token
|
||||
|
||||
reclaimed = await repo.claim_notification_work(
|
||||
now=now + timedelta(seconds=5),
|
||||
lease_owner="notifier-same",
|
||||
lease_seconds=60,
|
||||
limit=1,
|
||||
tracking_degraded_after_errors=3,
|
||||
)
|
||||
assert [row["id"] for row in reclaimed] == ["task-late-notify"]
|
||||
assert reclaimed[0]["notification_lease_token"] != first_token
|
||||
|
||||
released = await repo.release_notification_lease(
|
||||
"task-late-notify",
|
||||
lease_owner="notifier-same",
|
||||
notification_lease_token=first_token,
|
||||
next_notification_at=now + timedelta(seconds=30),
|
||||
error="cancelled",
|
||||
)
|
||||
assert released is False
|
||||
|
||||
stored = await repo.get("task-late-notify", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["notification_lease_owner"] == "notifier-same"
|
||||
assert stored["notification_lease_token"] == reclaimed[0]["notification_lease_token"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_late_snapshot_apply_after_same_worker_reclaim_is_fenced(tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
await _create_working_task(repo, task_id="task-late-apply", now=now)
|
||||
|
||||
first = await repo.claim_due_tasks(
|
||||
now=now,
|
||||
lease_owner="worker-same",
|
||||
lease_seconds=1,
|
||||
limit=10,
|
||||
)
|
||||
first_token = first[0]["lease_token"]
|
||||
reclaimed = await repo.claim_due_tasks(
|
||||
now=now + timedelta(seconds=5),
|
||||
lease_owner="worker-same",
|
||||
lease_seconds=60,
|
||||
limit=10,
|
||||
)
|
||||
assert [row["id"] for row in reclaimed] == ["task-late-apply"]
|
||||
assert reclaimed[0]["lease_token"] != first_token
|
||||
|
||||
applied = await repo.apply_snapshot(
|
||||
"task-late-apply",
|
||||
lease_owner="worker-same",
|
||||
lease_token=first_token,
|
||||
status="completed",
|
||||
result={"stale": True},
|
||||
result_preview=None,
|
||||
result_truncated=False,
|
||||
result_artifact=None,
|
||||
error=None,
|
||||
input_required=None,
|
||||
next_poll_at=None,
|
||||
polled_at=now + timedelta(seconds=5),
|
||||
)
|
||||
assert applied is False
|
||||
|
||||
stored = await repo.get("task-late-apply", user_id="user-1")
|
||||
assert stored is not None
|
||||
assert stored["lease_owner"] == "worker-same"
|
||||
assert stored["lease_token"] == reclaimed[0]["lease_token"]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@ -15,10 +15,11 @@ import sqlite3
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from alembic import command
|
||||
from alembic.script import ScriptDirectory
|
||||
|
||||
import deerflow.persistence.models # noqa: F401 -- registers ORM models
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.bootstrap import _get_alembic_config, _get_head_revision
|
||||
from deerflow.persistence.bootstrap import _MIGRATIONS_DIR, _get_alembic_config, _get_head_revision
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
|
||||
from deerflow.persistence.run import RunRepository
|
||||
|
||||
@ -66,8 +67,12 @@ def _table_and_column_state(db_path) -> tuple[bool, bool, set[str], str | None]:
|
||||
return "run_change_clock" in tables, "change_seq" in run_columns, run_indexes, version_row[0] if version_row else None
|
||||
|
||||
|
||||
async def test_0025_is_the_chain_head():
|
||||
assert _get_head_revision() == REVISION
|
||||
async def test_0025_chains_into_the_single_head():
|
||||
script = ScriptDirectory(str(_MIGRATIONS_DIR))
|
||||
assert len(script.get_heads()) == 1
|
||||
# Later migrations may advance the head without removing this revision.
|
||||
assert REVISION in {revision.revision for revision in script.walk_revisions()}
|
||||
assert script.get_revision(REVISION).down_revision == PREVIOUS
|
||||
|
||||
|
||||
async def test_0025_repairs_schema_skipped_by_the_0023_insertion(tmp_path):
|
||||
|
||||
51
backend/tests/test_migration_0026_mcp_task_lease_tokens.py
Normal file
51
backend/tests/test_migration_0026_mcp_task_lease_tokens.py
Normal file
@ -0,0 +1,51 @@
|
||||
"""Migration tests for 0026_mcp_task_lease_tokens.
|
||||
|
||||
Adds the nullable per-claim token columns ``McpTaskRepository`` uses to fence
|
||||
poll, cancel, and notification mutations to the exact claim generation. This
|
||||
file owns the chain-head pin, moved on from
|
||||
``test_migration_0025_repair_run_change_seq`` with this revision.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from alembic import command
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from deerflow.persistence import bootstrap
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
REVISION = "0026_mcp_task_lease_tokens"
|
||||
PREVIOUS = "0025_repair_run_change_seq"
|
||||
TOKEN_COLUMNS = {"lease_token", "notification_lease_token"}
|
||||
|
||||
|
||||
async def test_0026_is_the_chain_head():
|
||||
assert bootstrap._get_head_revision() == REVISION
|
||||
|
||||
|
||||
async def test_0026_adds_nullable_claim_tokens_and_downgrades(tmp_path):
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'lease-tokens.db'}")
|
||||
cfg = bootstrap._get_alembic_config(engine)
|
||||
|
||||
async def token_columns() -> dict[str, bool]:
|
||||
async with engine.connect() as conn:
|
||||
columns = await conn.run_sync(lambda sync: sa.inspect(sync).get_columns("mcp_tasks"))
|
||||
return {column["name"]: column["nullable"] for column in columns if column["name"] in TOKEN_COLUMNS}
|
||||
|
||||
try:
|
||||
await asyncio.to_thread(bootstrap._upgrade, cfg, PREVIOUS)
|
||||
assert await token_columns() == {}
|
||||
|
||||
await asyncio.to_thread(bootstrap._upgrade, cfg, "head")
|
||||
# Nullable so rows written before this revision stay valid.
|
||||
assert await token_columns() == dict.fromkeys(TOKEN_COLUMNS, True)
|
||||
|
||||
await asyncio.to_thread(command.downgrade, cfg, PREVIOUS)
|
||||
assert await token_columns() == {}
|
||||
finally:
|
||||
await engine.dispose()
|
||||
91
backend/tests/test_runtime_cancellation.py
Normal file
91
backend/tests/test_runtime_cancellation.py
Normal file
@ -0,0 +1,91 @@
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
import deerflow.runtime.cancellation as cancellation
|
||||
from deerflow.runtime.cancellation import wait_for_task_until
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_task_until_reports_completion():
|
||||
child = asyncio.create_task(asyncio.sleep(0, result="done"))
|
||||
|
||||
completed = await wait_for_task_until(child, deadline=asyncio.get_running_loop().time() + 1)
|
||||
|
||||
assert completed is True
|
||||
assert child.result() == "done"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_task_until_times_out_without_cancelling_child():
|
||||
event = asyncio.Event()
|
||||
child = asyncio.create_task(event.wait())
|
||||
|
||||
completed = await wait_for_task_until(child, deadline=asyncio.get_running_loop().time() + 0.01)
|
||||
|
||||
assert completed is False
|
||||
assert child.done() is False
|
||||
event.set()
|
||||
await child
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_task_until_zero_budget_returns_immediately():
|
||||
event = asyncio.Event()
|
||||
child = asyncio.create_task(event.wait())
|
||||
|
||||
completed = await wait_for_task_until(child, deadline=asyncio.get_running_loop().time())
|
||||
|
||||
assert completed is False
|
||||
assert child.done() is False
|
||||
child.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await child
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_wait_for_task_until_repeated_cancellation_keeps_original_deadline(monkeypatch):
|
||||
clock = iter((0.0, 0.01, 0.02, 0.05))
|
||||
clock_loop = SimpleNamespace(time=lambda: next(clock, 0.05))
|
||||
|
||||
wait_timeouts = []
|
||||
entered_first_wait = asyncio.Event()
|
||||
entered_second_wait = asyncio.Event()
|
||||
|
||||
async def fake_wait(tasks, *, timeout):
|
||||
del tasks
|
||||
wait_timeouts.append(timeout)
|
||||
if len(wait_timeouts) == 1:
|
||||
entered_first_wait.set()
|
||||
await asyncio.Future()
|
||||
if len(wait_timeouts) == 2:
|
||||
entered_second_wait.set()
|
||||
await asyncio.Future()
|
||||
return set(), set()
|
||||
|
||||
monkeypatch.setattr(
|
||||
cancellation,
|
||||
"asyncio",
|
||||
SimpleNamespace(
|
||||
CancelledError=asyncio.CancelledError,
|
||||
get_running_loop=lambda: clock_loop,
|
||||
wait=fake_wait,
|
||||
),
|
||||
)
|
||||
event = asyncio.Event()
|
||||
child = asyncio.create_task(event.wait())
|
||||
waiter = asyncio.create_task(wait_for_task_until(child, deadline=0.05))
|
||||
|
||||
await entered_first_wait.wait()
|
||||
waiter.cancel()
|
||||
await entered_second_wait.wait()
|
||||
waiter.cancel()
|
||||
assert waiter.cancelling() == 2
|
||||
completed = await waiter
|
||||
|
||||
assert completed is False
|
||||
assert wait_timeouts == pytest.approx([0.05, 0.04, 0.03])
|
||||
assert child.done() is False
|
||||
event.set()
|
||||
await child
|
||||
Loading…
x
Reference in New Issue
Block a user