trending/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-10-02 05:24:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(mcp): reserve notification launches atomically (#5965)

Browse Source 
* fix(mcp): reserve notification launches atomically

Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com>

* fix(mcp): classify busy notification conflicts by type

Signed-off-by: Willem Jiang <willem.jiang@gmail.com>

---------

Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com>
Signed-off-by: Willem Jiang <willem.jiang@gmail.com>
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
RongJie G
2026-09-30 09:12:11 +08:00
committed by GitHub
co-authored by Willem Jiang
parent f9f6a36243
commit d642dfdd2f
8 changed files with 556 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/gateway/AGENTS.md
+1 -1
View File
@@ -53,7 +53,7 @@ hot-reloaded top-level `recursion_limit` setting. A valid request-level value
takes precedence; invalid values fall back to that default, and
`max_recursion_limit` caps both sources.
Durable MCP notifications use internal Agent runs. Keep their trusted delivery instruction outside the user-input boundary, and frame serialized remote events as untrusted before model invocation. Strict thread existence/ownership admission dead-letters events whose task outlives its deleted chat instead of recreating the thread.
Durable MCP notifications use internal Agent runs. Keep their trusted delivery instruction outside the user-input boundary, and frame serialized remote events as untrusted before model invocation. Strict thread existence/ownership admission dead-letters events whose task outlives its deleted chat instead of recreating the thread. `start_run` marks run-manager admission conflicts with `BusyThreadConflict`; the notification launcher retries only that typed 409 and treats other 409 responses as permanent.
CORS is same-origin by default when requests enter through nginx on port 2026. Split-origin or port-forwarded browser clients must opt in with `GATEWAY_CORS_ORIGINS` (exact origins); Gateway `CORSMiddleware` and `CSRFMiddleware` both read that variable so browser CORS and auth-origin checks stay aligned. Those clients also need `CORS_EXPOSED_HEADERS` (`csrf_middleware.py`): run-creating routes return the run's id in `Content-Location`, which is not CORS-safelisted, so JS cannot read it unless it is exposed — and the LangGraph SDK resolves run metadata from that header alone, so withholding it breaks `useStream`'s `onCreated` and thread-gated actions.
backend/app/gateway/services.py
+10 -3
View File
@@ -104,6 +104,13 @@ from deerflow.utils.thread_id import validate_thread_id
logger = logging.getLogger(__name__)
class BusyThreadConflict(HTTPException):
"""A retryable run-manager admission conflict exposed as HTTP 409."""
def __init__(self, detail: str) -> None:
super().__init__(status_code=409, detail=detail)
@asynccontextmanager
async def reserve_checkpoint_write(
request: Request,
@@ -2125,7 +2132,7 @@ async def start_run(
)
raise
except ConflictError as exc:
raise HTTPException(status_code=409, detail=str(exc)) from exc
raise BusyThreadConflict(str(exc)) from exc
except UnsupportedStrategyError as exc:
raise HTTPException(status_code=501, detail=str(exc)) from exc
@@ -2286,9 +2293,9 @@ async def launch_mcp_task_notification_run(
require_existing_thread=True,
)
except HTTPException as exc:
if exc.status_code == 409:
if isinstance(exc, BusyThreadConflict):
raise ConflictError(str(exc.detail)) from exc
if exc.status_code == 404:
if exc.status_code in {400, 401, 403, 404, 409, 422, 501}:
raise PermanentNotificationError(str(exc.detail)) from exc
raise
return {"run_id": record.run_id, "thread_id": record.thread_id}
backend/app/mcp_tasks/service.py
+41
View File
@@ -936,6 +936,47 @@ class McpTaskService:
return
source_run = await self._get_run(record.get("run_id"), user_id=record["user_id"]) if record.get("run_id") else None
launch_started_at = _notification_completion_time(not_before=now)
recovering_launch = record.get("notification_status") == "launching"
# Reservation commit outcome is ambiguous under cancellation. Prefer a
# phase-preserving lease release until the await returns definitively.
record["notification_status"] = "launching"
launch_reserved = await self._repository.begin_notification_launch(
task_id,
lease_owner=self._lease_owner,
notification_lease_token=record["notification_lease_token"],
dispatch_version=dispatch_version,
lease_seconds=self._lease_seconds,
now=launch_started_at,
)
if not launch_reserved:
if recovering_launch:
await self._release_ordinary_batch_record(
record,
release=lambda: self._repository.release_notification_lease(
task_id,
lease_owner=self._lease_owner,
notification_lease_token=record["notification_lease_token"],
next_notification_at=launch_started_at,
error=record.get("notification_error"),
count_failure=False,
),
action="release expired notification launch",
)
return
await self._release_ordinary_batch_record(
record,
release=lambda: self._repository.release_notification_claim(
task_id,
lease_owner=self._lease_owner,
notification_lease_token=record["notification_lease_token"],
next_notification_at=launch_started_at,
error=record.get("notification_error"),
replace_with_latest=True,
),
action="release stale notification launch",
)
return
try:
result = await self._launch_notification(
thread_id=record["thread_id"],
backend/packages/harness/deerflow/mcp/AGENTS.md
+2 -2
View File
@@ -62,9 +62,9 @@
**This is defense in depth, not a trust boundary.** `npx`/`uvx` exist to fetch and execute remote packages, so an admin can still point one at a package they published; the boundary is admin authentication plus network reachability. Do not add a check here on the assumption that it makes MCP registration safe for untrusted admins — it does not, and the fix for that is not a bigger denylist.
**Durable MCP task runtime** (`mcp_tasks`, `McpTaskService`; summarized in
[backend/AGENTS.md](../../../../AGENTS.md)): Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. Explicit `task_toolsets` bind raw submit/status/cancel names; only submit remains Agent-visible, and its wrapper persists the remote handle before returning a local ID. `McpTaskService` claims due rows with leases, resolves a protocol-specific `McpTaskDriver`, and writes normalized snapshots back to `mcp_tasks`; expired leases are the restart-recovery mechanism, and a result returned after expiry or after a cancel request must be discarded even when the owner token still matches. The first cancel request fences an in-flight poll lease, while repeats preserve an active cancellation lease so they cannot issue concurrent remote cancels; cancellation backoff starts when the remote attempt finishes, so a slow timeout cannot consume the retry delay. Cancellation, polling, and notification batches isolate per-task exceptions; an unexpected cancellation/poll failure leaves that record's lease to expire, while notification failures release only the affected lease for retry. Input-required and terminal event snapshots are delivered by idempotent Agent runs and marked delivered only after run success; the trusted notification instruction stays outside the input boundary while the serialized remote event is framed as untrusted data. A busy-thread conflict is normalized back to the service boundary so the queued snapshot coalesces to the latest task event. A missing dispatched run becomes a failed delivery attempt, while transient run-store hydration errors stay distinguishable and retry the same lookup. The database is the source of truth; `ThreadState` receives only a bounded current-thread projection, and display names are neutralized at that model-state boundary. The installed process-local submitter is the source of truth for management-tool exposure; hot `mcp_tasks` edits take effect only after restart, and active skills must explicitly declare the list/cancel business tools.
[backend/AGENTS.md](../../../../AGENTS.md)): Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. Explicit `task_toolsets` bind raw submit/status/cancel names; only submit remains Agent-visible, and its wrapper persists the remote handle before returning a local ID. `McpTaskService` claims due rows with leases, resolves a protocol-specific `McpTaskDriver`, and writes normalized snapshots back to `mcp_tasks`; expired leases are the restart-recovery mechanism, and a result returned after expiry or after a cancel request must be discarded even when the owner token still matches. The first cancel request fences an in-flight poll lease, while repeats preserve an active cancellation lease so they cannot issue concurrent remote cancels; cancellation backoff starts when the remote attempt finishes, so a slow timeout cannot consume the retry delay. Cancellation, polling, and notification batches isolate per-task exceptions; an unexpected cancellation/poll failure leaves that record's lease to expire, while notification failures release only the affected lease for retry. Input-required and terminal event snapshots are delivered by idempotent Agent runs and marked delivered only after run success; the trusted notification instruction stays outside the input boundary while the serialized remote event is framed as untrusted data. Before creating an Agent notification run, one owner/token/version/expiry-fenced update reserves `notification_status=launching`; from that point, cancellation, ambiguous launch outcomes, and busy-thread conflicts preserve the reserved snapshot and idempotency key for recovery instead of coalescing to a newer event. A missing dispatched run becomes a failed delivery attempt, while transient run-store hydration errors stay distinguishable and retry the same lookup. The database is the source of truth; `ThreadState` receives only a bounded current-thread projection, and display names are neutralized at that model-state boundary. The installed process-local submitter is the source of truth for management-tool exposure; hot `mcp_tasks` edits take effect only after restart, and active skills must explicitly declare the list/cancel business tools.
**Task notification failure handling**: MCP notification failures use a consecutive counter separate from the idempotency-key `dispatch_attempt`, capped exponential backoff, latest-event rebuilding before a run launches, and a five-attempt budget before `dead_letter`. A permanently missing/mismatched target thread is dead-lettered immediately instead of being recreated or reclaimed. HTTP and Agent cancellation requests return after the durable cancel fence; the background loop alone owns the potentially slow remote call and retry schedule. The HTTP cancel endpoint rejects requests with 503 when the loop is not running (`mcp_tasks_available` false, e.g. `mcp_tasks.enabled=false` with SQL persistence), so a cancellation is never acknowledged without a worker to perform it. The bounded notification error/count/status join poll and cancellation diagnostics in the task detail API and expanded card.
**Task notification failure handling**: MCP notification failures use a consecutive counter separate from the idempotency-key `dispatch_attempt`, capped exponential backoff, latest-event rebuilding before a run launches, and a five-attempt budget for pre-reservation and dispatched-run failures before `dead_letter`. A permanently rejected launch, including a missing/mismatched target thread, is dead-lettered immediately instead of being recreated or reclaimed. An ambiguous launch result or busy-thread conflict preserves `launching`, the snapshot, and the same idempotency key without consuming that budget so a retry can reconcile a run that may already exist. HTTP and Agent cancellation requests return after the durable cancel fence; the background loop alone owns the potentially slow remote call and retry schedule. The HTTP cancel endpoint rejects requests with 503 when the loop is not running (`mcp_tasks_available` false, e.g. `mcp_tasks.enabled=false` with SQL persistence), so a cancellation is never acknowledged without a worker to perform it. The bounded notification error/count/status join poll and cancellation diagnostics in the task detail API and expanded card.
- Personal MCP configs live under the authenticated user root (`user_config.py`), use literal values, and never enter the deployment tool cache. `user_tools.py` guards owner and current revision before every call. Runtime server names bind owner/name/revision; durable task `driver_data.connection_scope` records personal or deployment ownership at submission, and background calls resolve a personal key only in the persisted task owner's file. Older rows without that field are deployment-owned. When a deployment name matches a personal runtime name, Agent assembly omits the personal tool rather than exposing two tools under one name; an already submitted personal task still uses its recorded personal scope. Keep catalog IDs stable across edits; never fall back by display name.
- Personal definitions without `personal_public_network: true` require current administrator authority at discovery and every tool/durable call. `personal_access.py` uses a host-installed async lookup, never a captured role or cached decision; absent/failed lookup denies privileged connections. The Gateway binds the lookup to fresh account records on its database-owning loop for its lifespan. Deployment connections are unaffected. This is admission control, not termination of operations already admitted before a role change.
backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py
+40 -1
View File
@@ -6,7 +6,7 @@ import uuid
from datetime import UTC, datetime, timedelta
from typing import Any
from sqlalchemy import case, or_, select, text, update
from sqlalchemy import and_, case, or_, select, text, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
@@ -675,6 +675,45 @@ class McpTaskRepository:
await session.commit()
return [self._row_to_dict(row) for row in rows]
async def begin_notification_launch(
self,
task_id: str,
*,
lease_owner: str,
notification_lease_token: str,
dispatch_version: int,
lease_seconds: int,
now: datetime,
) -> bool:
"""Reserve one idempotent Agent launch before starting the side effect."""
launchable = or_(
McpTaskRow.notification_status == "launching",
and_(
McpTaskRow.notification_status.in_(("claimed", "retry")),
McpTaskRow.event_version == dispatch_version,
),
)
stmt = (
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,
launchable,
)
.values(
notification_status="launching",
notification_lease_expires_at=now + timedelta(seconds=lease_seconds),
updated_at=now,
)
)
async with self._sf() as session:
result = await session.execute(stmt)
await session.commit()
return bool(result.rowcount)
async def mark_notification_dispatched(
self,
task_id: str,
backend/tests/test_gateway_services.py
+124 -3
View File
@@ -3683,23 +3683,78 @@ def test_launch_mcp_task_notification_run_hides_internal_prompt(_stub_app_config
assert result == {"run_id": "run-notification", "thread_id": "thread-notification"}
def test_start_run_marks_run_manager_conflict_as_busy(_stub_app_config):
import asyncio
from unittest.mock import AsyncMock, patch
from app.gateway.services import BusyThreadConflict, start_run
from deerflow.runtime.runs.manager import ConflictError
async def _scenario():
request, _run_store, _thread_store = _make_start_run_persistence_context()
request.app.state.run_manager.create_or_reject = AsyncMock(side_effect=ConflictError("Thread already has an active run"))
with (
patch("app.gateway.services.resolve_agent_factory", return_value=object()),
patch("app.gateway.services.ensure_checkpoint_history_seeded", new_callable=AsyncMock),
pytest.raises(BusyThreadConflict) as exc_info,
):
await start_run(_run_create_request(), "thread-busy", request)
assert exc_info.value.status_code == 409
assert exc_info.value.detail == "Thread already has an active run"
asyncio.run(_scenario())
def test_launch_mcp_task_notification_run_restores_busy_thread_conflict(_stub_app_config):
import asyncio
from types import SimpleNamespace
from unittest.mock import patch
from app.gateway.services import BusyThreadConflict, launch_mcp_task_notification_run
from deerflow.runtime.runs.manager import ConflictError
async def _scenario():
with (
patch(
"app.gateway.services.start_run",
side_effect=BusyThreadConflict("Thread already has an active run"),
),
pytest.raises(ConflictError, match="Thread already has an active run"),
):
await launch_mcp_task_notification_run(
app=SimpleNamespace(state=SimpleNamespace()),
thread_id="thread-notification",
assistant_id="lead_agent",
owner_user_id="user-1",
task_id="task-1",
dispatch_version=2,
dispatch_attempt=3,
event={"status": "completed", "result": "done"},
)
asyncio.run(_scenario())
def test_launch_mcp_task_notification_run_dead_letters_idempotency_conflict(_stub_app_config):
import asyncio
from types import SimpleNamespace
from unittest.mock import patch
from fastapi import HTTPException
from app.gateway.services import launch_mcp_task_notification_run
from deerflow.runtime.runs.manager import ConflictError
from app.mcp_tasks.errors import PermanentNotificationError
async def _scenario():
with (
patch(
"app.gateway.services.start_run",
side_effect=HTTPException(status_code=409, detail="Thread already has an active run"),
side_effect=HTTPException(
status_code=409,
detail="Idempotency-Key already used with a different request",
),
),
pytest.raises(ConflictError, match="Thread already has an active run"),
pytest.raises(PermanentNotificationError, match="Idempotency-Key"),
):
await launch_mcp_task_notification_run(
app=SimpleNamespace(state=SimpleNamespace()),
@@ -3747,6 +3802,72 @@ def test_launch_mcp_task_notification_run_dead_letters_missing_thread(_stub_app_
asyncio.run(_scenario())
@pytest.mark.parametrize("status_code", [400, 401, 403, 422, 501])
def test_launch_mcp_task_notification_run_dead_letters_deterministic_rejection(_stub_app_config, status_code):
import asyncio
from types import SimpleNamespace
from unittest.mock import patch
from fastapi import HTTPException
from app.gateway.services import launch_mcp_task_notification_run
from app.mcp_tasks.errors import PermanentNotificationError
async def _scenario():
with (
patch(
"app.gateway.services.start_run",
side_effect=HTTPException(status_code=status_code, detail="request rejected"),
),
pytest.raises(PermanentNotificationError, match="request rejected"),
):
await launch_mcp_task_notification_run(
app=SimpleNamespace(state=SimpleNamespace()),
thread_id="thread-notification",
assistant_id="lead_agent",
owner_user_id="user-1",
task_id="task-1",
dispatch_version=2,
dispatch_attempt=3,
event={"status": "completed", "result": "done"},
)
asyncio.run(_scenario())
@pytest.mark.parametrize("status_code", [408, 429, 500])
def test_launch_mcp_task_notification_run_preserves_ambiguous_http_failure(_stub_app_config, status_code):
import asyncio
from types import SimpleNamespace
from unittest.mock import patch
from fastapi import HTTPException
from app.gateway.services import launch_mcp_task_notification_run
async def _scenario():
with (
patch(
"app.gateway.services.start_run",
side_effect=HTTPException(status_code=status_code, detail="launch outcome unknown"),
),
pytest.raises(HTTPException) as exc_info,
):
await launch_mcp_task_notification_run(
app=SimpleNamespace(state=SimpleNamespace()),
thread_id="thread-notification",
assistant_id="lead_agent",
owner_user_id="user-1",
task_id="task-1",
dispatch_version=2,
dispatch_attempt=3,
event={"status": "completed", "result": "done"},
)
assert exc_info.value.status_code == status_code
asyncio.run(_scenario())
def test_start_run_strict_mode_rejects_missing_thread(_stub_app_config):
import asyncio
from types import SimpleNamespace
backend/tests/test_mcp_task_repository.py
+128 -11
View File
@@ -1024,14 +1024,14 @@ async def test_reclaimed_launching_notification_preserves_reserved_snapshot(tmp_
)
first_version = first[0]["dispatch_version"]
first_event = first[0]["dispatch_event"]
async with repo._sf() as session:
row = await session.get(McpTaskRow, "task-launching-compat")
assert row is not None
row.notification_status = "launching"
row.notification_lease_owner = None
row.notification_lease_expires_at = None
row.notification_lease_token = None
await session.commit()
assert await repo.begin_notification_launch(
"task-launching-compat",
lease_owner="notifier-a",
notification_lease_token=first[0]["notification_lease_token"],
dispatch_version=first_version,
lease_seconds=60,
now=now,
)
second_poll = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
await repo.apply_snapshot(
@@ -1049,8 +1049,9 @@ async def test_reclaimed_launching_notification_preserves_reserved_snapshot(tmp_
polled_at=now,
)
reclaimed_at = now + timedelta(seconds=61)
reclaimed = await repo.claim_notification_work(
now=now,
now=reclaimed_at,
lease_owner="notifier-b",
lease_seconds=60,
limit=1,
@@ -1072,9 +1073,9 @@ async def test_reclaimed_launching_notification_preserves_reserved_snapshot(tmp_
get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")),
)
service._lease_owner = "notifier-b"
await service._notify_one_claimed(reclaimed[0], now=now)
await service._notify_one_claimed(reclaimed[0], now=reclaimed_at)
retry_at = now + timedelta(seconds=10)
retry_at = reclaimed_at + timedelta(seconds=10)
after_conflict = await repo.claim_notification_work(
now=retry_at,
lease_owner="notifier-c",
@@ -1097,6 +1098,122 @@ async def test_reclaimed_launching_notification_preserves_reserved_snapshot(tmp_
)
@pytest.mark.asyncio
async def test_begin_notification_launch_rejects_unstarted_stale_snapshot(tmp_path):
repo = await _make_repo(tmp_path)
now = datetime.now(UTC)
await _create_working_task(repo, task_id="task-notify-stale", 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-stale",
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,
)
second_poll = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1)
await repo.apply_snapshot(
"task-notify-stale",
lease_owner="poller",
lease_token=second_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,
)
reserved = await repo.begin_notification_launch(
"task-notify-stale",
lease_owner="notifier",
notification_lease_token=first[0]["notification_lease_token"],
dispatch_version=first[0]["dispatch_version"],
lease_seconds=120,
now=now,
)
assert reserved is False
stored = await repo.get("task-notify-stale", user_id="user-1", thread_id="thread-1", thread_incarnation=None)
assert stored is not None
assert stored["notification_status"] == "claimed"
assert stored["event_version"] > stored["dispatch_version"]
@pytest.mark.asyncio
async def test_begin_notification_launch_is_fenced_by_reclaimed_token(tmp_path, monkeypatch):
repo = await _make_repo(tmp_path)
now = datetime.now(UTC)
await _create_working_task(repo, task_id="task-launch-token-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-launch-token-fence",
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="same-worker",
lease_seconds=60,
limit=1,
tracking_degraded_after_errors=3,
)
operation = repo.begin_notification_launch(
"task-launch-token-fence",
lease_owner="same-worker",
notification_lease_token=first[0]["notification_lease_token"],
dispatch_version=first[0]["dispatch_version"],
lease_seconds=120,
now=now,
)
async with _pause_claim_mutation(monkeypatch, operation) as (stale_launch, resume):
reclaimed = await repo.claim_notification_work(
now=now + timedelta(seconds=61),
lease_owner="same-worker",
lease_seconds=60,
limit=1,
tracking_degraded_after_errors=3,
)
assert len(reclaimed) == 1
new_token = reclaimed[0]["notification_lease_token"]
assert new_token != first[0]["notification_lease_token"]
resume.set()
assert await stale_launch is False
stored = await repo.get("task-launch-token-fence", user_id="user-1", thread_id="thread-1", thread_incarnation=None)
assert stored is not None
assert stored["notification_status"] == "claimed"
assert stored["notification_lease_owner"] == "same-worker"
assert stored["notification_lease_token"] == new_token
@pytest.mark.asyncio
async def test_recovered_launching_at_retry_budget_reconciles_successful_run(tmp_path):
repo = await _make_repo(tmp_path)
backend/tests/test_mcp_task_service.py
+210 -8
View File
@@ -65,6 +65,9 @@ class FakeRepository:
self.released.append((task_id, kwargs))
return True
async def begin_notification_launch(self, *_args, **_kwargs):
return True
class FailingApplyRepository(FakeRepository):
async def apply_snapshot(self, task_id, **kwargs):
@@ -695,6 +698,7 @@ async def test_cancel_recovery_failures_are_isolated_and_later_phases_continue(c
@pytest.mark.asyncio
async def test_notification_delivery_waits_for_successful_agent_run():
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
mark_notification_dispatched=AsyncMock(return_value=True),
finish_notification_run=AsyncMock(return_value=True),
release_notification_claim=AsyncMock(return_value=True),
@@ -893,9 +897,11 @@ async def test_notification_failure_releases_at_its_own_completion_without_waiti
@pytest.mark.asyncio
async def test_notification_busy_thread_replaces_claim_with_latest_event(monkeypatch):
async def test_reserved_notification_busy_thread_preserves_snapshot(monkeypatch):
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
release_notification_claim=AsyncMock(return_value=True),
release_notification_lease=AsyncMock(return_value=True),
)
service = McpTaskService(
repository=repo,
@@ -920,14 +926,16 @@ async def test_notification_busy_thread_replaces_claim_with_latest_event(monkeyp
now=now,
)
released = repo.release_notification_claim.await_args.kwargs
assert released["replace_with_latest"] is True
repo.release_notification_claim.assert_not_awaited()
released = repo.release_notification_lease.await_args.kwargs
assert released["count_failure"] is False
assert released["next_notification_at"] == now + timedelta(seconds=5)
@pytest.mark.asyncio
async def test_reclaimed_launching_busy_thread_preserves_reserved_snapshot(monkeypatch):
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
release_notification_claim=AsyncMock(return_value=True),
release_notification_lease=AsyncMock(return_value=True),
)
@@ -962,9 +970,11 @@ async def test_reclaimed_launching_busy_thread_preserves_reserved_snapshot(monke
@pytest.mark.asyncio
async def test_notification_launch_failure_backs_off_and_replaces_with_latest_event(monkeypatch):
async def test_notification_launch_failure_backs_off_without_replacing_reserved_snapshot(monkeypatch):
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
release_notification_claim=AsyncMock(return_value=True),
release_notification_lease=AsyncMock(return_value=True),
)
service = McpTaskService(
repository=repo,
@@ -991,15 +1001,16 @@ async def test_notification_launch_failure_backs_off_and_replaces_with_latest_ev
now=now,
)
released = repo.release_notification_claim.await_args.kwargs
assert released["replace_with_latest"] is True
assert released["count_failure"] is True
repo.release_notification_claim.assert_not_awaited()
released = repo.release_notification_lease.await_args.kwargs
assert released["count_failure"] is False
assert released["next_notification_at"] == now + timedelta(seconds=40)
@pytest.mark.asyncio
async def test_permanently_rejected_notification_is_dead_lettered():
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
dead_letter_notification=AsyncMock(return_value=True),
release_notification_claim=AsyncMock(return_value=True),
)
@@ -1037,6 +1048,7 @@ async def test_permanently_rejected_notification_is_dead_lettered():
@pytest.mark.asyncio
async def test_permanently_rejected_reclaimed_launching_notification_is_dead_lettered():
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
dead_letter_notification=AsyncMock(return_value=True),
release_notification_claim=AsyncMock(return_value=True),
release_notification_lease=AsyncMock(return_value=True),
@@ -1114,7 +1126,10 @@ async def test_notification_dispatch_uses_launch_completion_time(monkeypatch):
return {"run_id": "notify-run-1"}
monkeypatch.setattr(service_module, "datetime", _MutableDateTime)
repo = SimpleNamespace(mark_notification_dispatched=AsyncMock(return_value=False))
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
mark_notification_dispatched=AsyncMock(return_value=False),
)
service = McpTaskService(
repository=repo,
drivers=McpTaskDriverRegistry(),
@@ -1139,9 +1154,80 @@ async def test_notification_dispatch_uses_launch_completion_time(monkeypatch):
assert repo.mark_notification_dispatched.await_args.kwargs["now"] == launch_completed_at
@pytest.mark.asyncio
async def test_notification_does_not_launch_when_preflight_fence_is_rejected():
launch_notification = AsyncMock(return_value={"run_id": "stale-notify-run"})
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=False),
mark_notification_dispatched=AsyncMock(return_value=False),
release_notification_claim=AsyncMock(return_value=False),
)
service = McpTaskService(
repository=repo,
drivers=McpTaskDriverRegistry(),
poll_interval_seconds=5,
lease_seconds=120,
max_concurrent_polls=3,
launch_notification=launch_notification,
get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")),
)
await service._notify_one_claimed(
{
**_claimed_row(),
"notification_status": "claimed",
"dispatch_version": 2,
"dispatch_attempt": 0,
"dispatch_event": {"status": "completed"},
},
now=datetime.now(UTC),
)
repo.begin_notification_launch.assert_awaited_once()
repo.release_notification_claim.assert_awaited_once()
launch_notification.assert_not_awaited()
repo.mark_notification_dispatched.assert_not_awaited()
@pytest.mark.asyncio
async def test_reclaimed_launching_notification_preserves_snapshot_when_preflight_lease_expires():
launch_notification = AsyncMock(return_value={"run_id": "stale-notify-run"})
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=False),
release_notification_claim=AsyncMock(return_value=False),
release_notification_lease=AsyncMock(return_value=True),
)
service = McpTaskService(
repository=repo,
drivers=McpTaskDriverRegistry(),
poll_interval_seconds=5,
lease_seconds=120,
max_concurrent_polls=3,
launch_notification=launch_notification,
get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")),
)
await service._notify_one_claimed(
{
**_claimed_row(),
"notification_status": "launching",
"dispatch_version": 2,
"dispatch_attempt": 0,
"dispatch_event": {"status": "completed"},
},
now=datetime.now(UTC),
)
repo.begin_notification_launch.assert_awaited_once()
repo.release_notification_lease.assert_awaited_once()
repo.release_notification_claim.assert_not_awaited()
launch_notification.assert_not_awaited()
@pytest.mark.asyncio
async def test_reclaimed_launching_notification_preserves_snapshot_on_uncertain_failure_at_retry_budget():
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
release_notification_claim=AsyncMock(return_value=True),
release_notification_lease=AsyncMock(return_value=True),
dead_letter_notification=AsyncMock(return_value=True),
@@ -2386,6 +2472,9 @@ class NotificationPersistenceRepo:
}
]
async def begin_notification_launch(self, *_args, **_kwargs):
return True
async def mark_notification_dispatched(self, *_args, **_kwargs):
self.mark_started.set()
await asyncio.Event().wait()
@@ -3167,6 +3256,10 @@ class SelfCancellingNotificationRepository(FakeRepository):
self.notification_releases.append((task_id, kwargs))
return True
async def release_notification_lease(self, task_id, **kwargs):
self.notification_releases.append((task_id, kwargs))
return True
@pytest.mark.asyncio
async def test_notification_child_self_cancellation_releases_once():
@@ -3201,6 +3294,10 @@ class BatchNotificationRepository(FakeRepository):
self.notification_releases.append((task_id, kwargs))
return True
async def release_notification_lease(self, task_id, **kwargs):
self.notification_releases.append((task_id, kwargs))
return True
@pytest.mark.asyncio
async def test_notification_outer_cancellation_releases_started_and_never_started_once():
@@ -3537,6 +3634,9 @@ class OrdinaryReleaseBatchRepository:
async def claim_notification_work(self, **_kwargs):
return self._records("notification")
async def begin_notification_launch(self, *_args, **_kwargs):
return True
async def _release(self, task_id, **_kwargs):
self.release_calls.append(task_id)
self.release_started.set()
@@ -3570,6 +3670,9 @@ class OrdinaryReleaseBatchRepository:
async def release_notification_claim(self, task_id, **kwargs):
return await self._release(task_id, **kwargs)
async def release_notification_lease(self, task_id, **kwargs):
return await self._release(task_id, **kwargs)
async def release_poll_claim_after_cancellation(self, task_id, **kwargs):
return await self._release(task_id, **kwargs)
@@ -4382,6 +4485,104 @@ async def test_notification_cancellation_during_source_run_lookup_releases_claim
repo.release_notification_claim.assert_awaited_once()
@pytest.mark.asyncio
async def test_notification_cancellation_during_launch_reservation_preserves_phase():
reservation_started = asyncio.Event()
async def begin_notification_launch(*_args, **_kwargs):
reservation_started.set()
await asyncio.Event().wait()
repo = SimpleNamespace(
begin_notification_launch=begin_notification_launch,
release_notification_claim=AsyncMock(return_value=True),
release_notification_lease=AsyncMock(return_value=True),
)
service = McpTaskService(
repository=repo,
drivers=McpTaskDriverRegistry(),
poll_interval_seconds=5,
lease_seconds=120,
max_concurrent_polls=3,
launch_notification=AsyncMock(),
get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")),
)
record = {
**_claimed_row(),
"notification_status": "claimed",
"dispatch_version": 2,
"dispatch_attempt": 0,
"dispatch_event": {"status": "completed"},
}
now = datetime.now(UTC)
task = asyncio.create_task(
service._run_claimed_batch(
[record],
operation=lambda item: service._notify_one_claimed(item, now=now),
release=service._release_notification_after_cancellation,
action="notification",
)
)
await reservation_started.wait()
assert record["notification_status"] == "launching"
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
repo.release_notification_lease.assert_awaited_once()
repo.release_notification_claim.assert_not_awaited()
@pytest.mark.asyncio
async def test_notification_cancellation_after_launch_reservation_preserves_launching_phase():
launch_started = asyncio.Event()
async def launch_notification(**_kwargs):
launch_started.set()
await asyncio.Event().wait()
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
release_notification_claim=AsyncMock(return_value=True),
release_notification_lease=AsyncMock(return_value=True),
)
service = McpTaskService(
repository=repo,
drivers=McpTaskDriverRegistry(),
poll_interval_seconds=5,
lease_seconds=120,
max_concurrent_polls=3,
launch_notification=launch_notification,
get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")),
)
record = {
**_claimed_row(),
"notification_status": "claimed",
"dispatch_version": 2,
"dispatch_attempt": 0,
"dispatch_event": {"status": "completed"},
}
now = datetime.now(UTC)
task = asyncio.create_task(
service._run_claimed_batch(
[record],
operation=lambda item: service._notify_one_claimed(item, now=now),
release=service._release_notification_after_cancellation,
action="notification",
)
)
await launch_started.wait()
assert record["notification_status"] == "launching"
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
repo.release_notification_lease.assert_awaited_once()
repo.release_notification_claim.assert_not_awaited()
@pytest.mark.asyncio
async def test_reclaimed_launching_notification_cancellation_preserves_phase():
launch_started = asyncio.Event()
@@ -4391,6 +4592,7 @@ async def test_reclaimed_launching_notification_cancellation_preserves_phase():
await asyncio.Event().wait()
repo = SimpleNamespace(
begin_notification_launch=AsyncMock(return_value=True),
release_notification_claim=AsyncMock(return_value=True),
release_notification_lease=AsyncMock(return_value=True),
)