mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-30 07:00:37 +00:00
fix(mcp): drain compensation work during service shutdown (#7662)
* fix(mcp): drain compensation work during service shutdown * test(mcp): make shutdown drain assertion deterministic * fix(mcp): preserve shutdown drain through caller cancellation * style(mcp): keep test module spacing consistent * fix(mcp): report pending compensation tasks on shutdown timeout * test(mcp): cover compensation observability on poller timeout
This commit is contained in:
NanPan
GitHub
parent
76835c1ff2
commit
5c4ba4d2fe
@ -1689,11 +1689,60 @@ class McpTaskService:
|
||||
_CANCELLATION_DRAIN_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
def _log_pending_compensation_tasks(self, *, reason: str) -> None:
|
||||
pending_count = sum(not task.done() for task in self._compensation_tasks)
|
||||
if pending_count:
|
||||
logger.warning(
|
||||
"%s; %d MCP compensation task(s) remain in the background",
|
||||
reason,
|
||||
pending_count,
|
||||
)
|
||||
|
||||
async def _drain_compensation_tasks(self, *, deadline: float) -> None:
|
||||
"""Drain service-owned compensation work before Gateway dependencies close."""
|
||||
loop = asyncio.get_running_loop()
|
||||
while self._compensation_tasks:
|
||||
tasks = tuple(self._compensation_tasks)
|
||||
remaining = deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
break
|
||||
done, _ = await asyncio.wait(
|
||||
tasks,
|
||||
timeout=remaining,
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
if not done:
|
||||
break
|
||||
# Run completion callbacks before taking the next snapshot. A
|
||||
# handoff may finish one task while registering another.
|
||||
await asyncio.sleep(0)
|
||||
|
||||
self._log_pending_compensation_tasks(
|
||||
reason="MCP compensation drain reached its deadline",
|
||||
)
|
||||
|
||||
async def _drain_compensation_tasks_preserving_cancellation(self, *, deadline: float) -> None:
|
||||
"""Finish the bounded drain before propagating caller cancellation."""
|
||||
drain_task = asyncio.create_task(
|
||||
self._drain_compensation_tasks(deadline=deadline),
|
||||
name="deerflow-mcp-compensation-drain",
|
||||
)
|
||||
try:
|
||||
await asyncio.shield(drain_task)
|
||||
except asyncio.CancelledError:
|
||||
await wait_for_task_until(drain_task, deadline=deadline)
|
||||
if drain_task.done():
|
||||
drain_task.result()
|
||||
raise
|
||||
|
||||
async def stop(self) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
task = self._task
|
||||
if task is None:
|
||||
await self._drain_compensation_tasks_preserving_cancellation(
|
||||
deadline=loop.time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS,
|
||||
)
|
||||
return
|
||||
loop = asyncio.get_running_loop()
|
||||
if self._stopping_task is not task:
|
||||
self._stopping_task = task
|
||||
self._stop_deadline = loop.time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS
|
||||
@ -1710,9 +1759,19 @@ class McpTaskService:
|
||||
except asyncio.CancelledError:
|
||||
if not await wait_for_task_until(task, deadline=deadline):
|
||||
self._log_stop_timeout(task)
|
||||
self._log_pending_compensation_tasks(
|
||||
reason="MCP task poller cleanup reached its deadline before compensation drain could start",
|
||||
)
|
||||
else:
|
||||
await self._drain_compensation_tasks_preserving_cancellation(deadline=deadline)
|
||||
raise
|
||||
if task not in done:
|
||||
self._log_stop_timeout(task)
|
||||
self._log_pending_compensation_tasks(
|
||||
reason="MCP task poller cleanup reached its deadline before compensation drain could start",
|
||||
)
|
||||
return
|
||||
await self._drain_compensation_tasks_preserving_cancellation(deadline=deadline)
|
||||
|
||||
async def _run_loop(self) -> None:
|
||||
while not self._stop.is_set():
|
||||
|
||||
@ -11,7 +11,7 @@ The optional background-task `status` filter flows through `McpTaskSubmitter.lis
|
||||
|
||||
- Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management
|
||||
- **OAuth expiry**: Missing, non-convertible or overflowing `expires_in` values use the one-hour default. Guard both `timedelta` construction and addition to one captured UTC `now`; avoid float-derived range ceilings. `tests/test_mcp_oauth.py` pins both overflow stages and the exact `datetime` boundary with a fixed clock.
|
||||
- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` and `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`); validate finite positive `poll_after_seconds` at that shared boundary. `persistence/mcp_tasks/` owns remote handles, schedules, notification state, leases and separate consecutive poll/delivery error counters. `app/mcp_tasks/McpTaskService` runs outside the Agent/LLM loop. Every claim gets a fresh token; releases and snapshot writes match owner/token, and results require an unexpired lease even before a peer reclaims. Each claim phase has a single-flight owner while its database outcome is ambiguous: later scans skip with a warning. Once it resolves, release that phase owner immediately and hand returned rows to bounded, service-owned release work. Token fencing protects newer generations; a stuck release must not block later claims. Expiry is the cross-process crash fallback, not normal recovery for a live service's late claim. Poll timestamps and retry schedules use remote completion time. Submission persistence failure/cancellation best-effort cancels the remote handle while preserving the original error, except `uq_mcp_tasks_user_server_remote`: an existing row owns that handle and must not be cancelled. Isolate per-task failures; unexpected poll failures recover by expiry. Shutdown cancels the poller with bounded handoff. Cancellation release clears only owner/token and preserves diagnostics, preclaim schedules and poll-failure state; cancellation is not a task failure. Real failures retain backoff and tracking degradation. `input_required` and terminal events become pending notifications; delivery budgets, idempotency and strict target rejection are detailed below. Durable recovery needs SQLite/PostgreSQL; memory leaves the service unavailable. The startup-configured `mcp_tasks` runtime is disabled by default until a concrete driver is registered and does not change ordinary tool behavior.
|
||||
- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` and `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`); validate finite positive `poll_after_seconds` at that shared boundary. `persistence/mcp_tasks/` owns remote handles, schedules, notification state, leases and separate consecutive poll/delivery error counters. `app/mcp_tasks/McpTaskService` runs outside the Agent/LLM loop. Every claim gets a fresh token; releases and snapshot writes match owner/token, and results require an unexpired lease even before a peer reclaims. Each claim phase has a single-flight owner while its database outcome is ambiguous: later scans skip with a warning. Once it resolves, release that phase owner immediately and hand returned rows to bounded, service-owned release work. Token fencing protects newer generations; a stuck release must not block later claims. Expiry is the cross-process crash fallback, not normal recovery for a live service's late claim. Poll timestamps and retry schedules use remote completion time. Submission persistence failure/cancellation best-effort cancels the remote handle while preserving the original error, except `uq_mcp_tasks_user_server_remote`: an existing row owns that handle and must not be cancelled. Isolate per-task failures; unexpected poll failures recover by expiry. Shutdown cancels the poller with bounded handoff, then drains service-owned compensation work within the same deadline; timed-out work remains tracked in the background. Cancellation release clears only owner/token and preserves diagnostics, preclaim schedules and poll-failure state; cancellation is not a task failure. Real failures retain backoff and tracking degradation. `input_required` and terminal events become pending notifications; delivery budgets, idempotency and strict target rejection are detailed below. Durable recovery needs SQLite/PostgreSQL; memory leaves the service unavailable. The startup-configured `mcp_tasks` runtime is disabled by default until a concrete driver is registered and does not change ordinary tool behavior.
|
||||
- **Control-call ownership**: Poll claims recheck eligibility with a conditional UPDATE after SELECT; cancel claims reserve SQLite writers with `BEGIN IMMEDIATE` before SELECT. PostgreSQL retains `SKIP LOCKED`. `renew_control_claim` atomically matches phase, owner, token, nonterminal state and unexpired lease, preserves `updated_at`, and never shortens expiry on a backward clock correction. `McpTaskService._with_control_lease` validates admission, renews every lease/3 with a lease/3 renewal deadline, and supervises the entire operation through persistence. Ownership loss or uncertain renewal cancels the local operation; the task-local loss fence rejects late snapshot/retry writes if the driver swallows cancellation. Bounded handoff retains operation/renewal cleanup through caller cancellation or autonomous completion/loss. Local active IDs remain excluded from poll/cancel claims and consume control capacity until cleanup settles. Notification claims use their existing independent lifecycle. Contract/tests: [control leases](../../../../docs/mcp-control-leases.md), `tests/test_mcp_task_lease_lifecycle.py`.
|
||||
- **Runtime availability boundary**: the installed process-local submitter is the source of truth for durable task-management tool exposure. `mcp_tasks` is startup-only; changing it on disk does not alter the live toolset until the Gateway restarts.
|
||||
- **Task OAuth lifetime**: `tasks/runtime.py` owns shared OAuth state only for enabled deployment servers with `task_toolsets`, from `set_mcp_task_config_snapshot(config)` until its `None` teardown. Use `_task_enabled_servers` for both snapshot and OAuth selection. Validate the frozen snapshot before reusing that state for discovery, ordinary tools, or durable calls. Cache resets must not discard it. Keep tokens and the cross-loop refresh lock together; rotate a private config copy, never the parsed snapshot or config file. Non-task servers remain reloadable and personal callers use independent managers, even on a name collision.
|
||||
|
||||
@ -4845,3 +4845,170 @@ async def test_background_compensation_failure_is_consumed(monkeypatch, caplog):
|
||||
assert len(failures) == 1
|
||||
assert any("release remained unavailable" in failure.getMessage() for failure in failures)
|
||||
assert not service._compensation_tasks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_drains_service_owned_compensation_tasks(monkeypatch):
|
||||
monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.5)
|
||||
service = McpTaskService(
|
||||
repository=FakeRepository(),
|
||||
drivers=McpTaskDriverRegistry(),
|
||||
poll_interval_seconds=60,
|
||||
lease_seconds=120,
|
||||
max_concurrent_polls=3,
|
||||
)
|
||||
release_started = asyncio.Event()
|
||||
release_gate = asyncio.Event()
|
||||
|
||||
async def release_claim():
|
||||
release_started.set()
|
||||
await release_gate.wait()
|
||||
|
||||
compensation = asyncio.create_task(release_claim(), name="test-mcp-compensation")
|
||||
service._track_compensation_task(compensation, action="release poll claim", task_id="task-1")
|
||||
await release_started.wait()
|
||||
await service.start()
|
||||
|
||||
original_drain = service._drain_compensation_tasks
|
||||
drain_started = asyncio.Event()
|
||||
|
||||
async def observed_drain(*, deadline: float):
|
||||
drain_started.set()
|
||||
await original_drain(deadline=deadline)
|
||||
|
||||
monkeypatch.setattr(service, "_drain_compensation_tasks", observed_drain)
|
||||
stopping = asyncio.create_task(service.stop())
|
||||
try:
|
||||
await asyncio.wait_for(drain_started.wait(), timeout=1)
|
||||
assert not stopping.done(), "stop must drain service-owned compensation work"
|
||||
finally:
|
||||
release_gate.set()
|
||||
await asyncio.wait_for(stopping, timeout=1)
|
||||
|
||||
await _wait_for_compensation_tasks_to_clear(service)
|
||||
assert not service._compensation_tasks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_bounds_drain_of_service_owned_compensation_tasks(monkeypatch, caplog):
|
||||
monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05)
|
||||
service = McpTaskService(
|
||||
repository=FakeRepository(),
|
||||
drivers=McpTaskDriverRegistry(),
|
||||
poll_interval_seconds=60,
|
||||
lease_seconds=120,
|
||||
max_concurrent_polls=3,
|
||||
)
|
||||
release_started = asyncio.Event()
|
||||
release_gate = asyncio.Event()
|
||||
|
||||
async def release_claim():
|
||||
release_started.set()
|
||||
await release_gate.wait()
|
||||
|
||||
compensation = asyncio.create_task(release_claim(), name="test-mcp-stuck-compensation")
|
||||
service._track_compensation_task(compensation, action="release poll claim", task_id="task-1")
|
||||
await release_started.wait()
|
||||
await service.start()
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING):
|
||||
await asyncio.wait_for(service.stop(), timeout=0.5)
|
||||
assert "MCP compensation task(s)" in caplog.text
|
||||
assert compensation in service._compensation_tasks
|
||||
finally:
|
||||
release_gate.set()
|
||||
await asyncio.wait_for(compensation, timeout=1)
|
||||
await _wait_for_compensation_tasks_to_clear(service)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_logs_pending_compensation_on_poller_timeout(monkeypatch, caplog):
|
||||
monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.05)
|
||||
service = McpTaskService(
|
||||
repository=FakeRepository(),
|
||||
drivers=McpTaskDriverRegistry(),
|
||||
poll_interval_seconds=60,
|
||||
lease_seconds=120,
|
||||
max_concurrent_polls=3,
|
||||
)
|
||||
poller_started = asyncio.Event()
|
||||
poller_release = asyncio.Event()
|
||||
compensation_started = asyncio.Event()
|
||||
compensation_release = asyncio.Event()
|
||||
|
||||
async def blocked_run_once(*, now):
|
||||
poller_started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
except asyncio.CancelledError:
|
||||
await poller_release.wait()
|
||||
|
||||
async def release_claim():
|
||||
compensation_started.set()
|
||||
await compensation_release.wait()
|
||||
|
||||
monkeypatch.setattr(service, "run_once", blocked_run_once)
|
||||
compensation = asyncio.create_task(release_claim(), name="test-mcp-timeout-compensation")
|
||||
service._track_compensation_task(compensation, action="release poll claim", task_id="task-1")
|
||||
await compensation_started.wait()
|
||||
await service.start()
|
||||
await poller_started.wait()
|
||||
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING):
|
||||
await asyncio.wait_for(service.stop(), timeout=0.5)
|
||||
assert "before compensation drain could start" in caplog.text
|
||||
assert "1 MCP compensation task(s)" in caplog.text
|
||||
assert compensation in service._compensation_tasks
|
||||
finally:
|
||||
poller_release.set()
|
||||
compensation_release.set()
|
||||
if service._task is not None:
|
||||
await asyncio.wait_for(service._task, timeout=1)
|
||||
await _wait_for_compensation_tasks_to_clear(service)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_preserves_cancellation_while_draining_compensation_tasks(monkeypatch):
|
||||
monkeypatch.setattr(service_module, "_CANCELLATION_DRAIN_TIMEOUT_SECONDS", 0.5)
|
||||
service = McpTaskService(
|
||||
repository=FakeRepository(),
|
||||
drivers=McpTaskDriverRegistry(),
|
||||
poll_interval_seconds=60,
|
||||
lease_seconds=120,
|
||||
max_concurrent_polls=3,
|
||||
)
|
||||
release_started = asyncio.Event()
|
||||
release_gate = asyncio.Event()
|
||||
|
||||
async def release_claim():
|
||||
release_started.set()
|
||||
await release_gate.wait()
|
||||
|
||||
compensation = asyncio.create_task(release_claim(), name="test-mcp-cancelled-drain")
|
||||
service._track_compensation_task(compensation, action="release poll claim", task_id="task-1")
|
||||
await release_started.wait()
|
||||
await service.start()
|
||||
|
||||
original_drain = service._drain_compensation_tasks
|
||||
drain_started = asyncio.Event()
|
||||
|
||||
async def observed_drain(*, deadline: float):
|
||||
drain_started.set()
|
||||
await original_drain(deadline=deadline)
|
||||
|
||||
monkeypatch.setattr(service, "_drain_compensation_tasks", observed_drain)
|
||||
stopping = asyncio.create_task(service.stop())
|
||||
try:
|
||||
await asyncio.wait_for(drain_started.wait(), timeout=1)
|
||||
stopping.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not stopping.done(), "stop must keep draining after caller cancellation"
|
||||
release_gate.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(stopping, timeout=1)
|
||||
finally:
|
||||
release_gate.set()
|
||||
if not stopping.done():
|
||||
stopping.cancel()
|
||||
await _wait_for_compensation_tasks_to_clear(service)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user