mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-10 05:58:36 +00:00
feat(mcp): add durable task runtime foundation (#4665)
* feat(mcp): add durable task runtime foundation * fix(chart): sync embedded config version * fix(mcp): isolate task polls during shutdown * feat(mcp): track consecutive poll errors on mcp_tasks poll_attempt_count grows on every claim (successful polls included), so it cannot drive a failure backoff without misjudging normal long tasks. Add consecutive_poll_error_count: incremented when a claim is released after a poll error, reset to zero by any applied snapshot. The backoff/terminal policy that consumes it lands with the first concrete driver. * fix(mcp): harden durable task lifecycle * fix(mcp): preserve tracked task on dedup conflict --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
e5c62cab5a
commit
e9387394bc
@ -416,6 +416,8 @@ Settings > Tools updates one MCP server at a time: an invalid stdio command on o
|
|||||||
Targeted updates accept both DeerFlow's `type` field and the MCP-spec `transport` field for SSE/HTTP servers.
|
Targeted updates accept both DeerFlow's `type` field and the MCP-spec `transport` field for SSE/HTTP servers.
|
||||||
Runtime MCP and skill updates replace `extensions_config.json` atomically, so an interrupted write cannot leave the shared configuration truncated or partially written.
|
Runtime MCP and skill updates replace `extensions_config.json` atomically, so an interrupted write cannot leave the shared configuration truncated or partially written.
|
||||||
MCP routing hints can also prefer a specific MCP tool for matching requests without forbidding other tools. When `tool_search` defers MCP schemas, matching routing metadata can auto-promote up to `tool_search.auto_promote_top_k` deferred schemas before the model call.
|
MCP routing hints can also prefer a specific MCP tool for matching requests without forbidding other tools. When `tool_search` defers MCP schemas, matching routing metadata can auto-promote up to `tool_search.auto_promote_top_k` deferred schemas before the model call.
|
||||||
|
|
||||||
|
The Gateway also includes a disabled-by-default, protocol-neutral foundation for durable long-running MCP tasks. It stores remote task handles outside model context, polls them under cross-worker leases, rejects results returned after their lease expires, schedules the next attempt from the time a remote status call finishes, isolates unexpected failures between claimed tasks, cancels in-flight polling during Gateway shutdown, and makes expired claims recoverable after restart. If remote submission succeeds but the handle cannot be persisted, the runtime makes a best-effort cancellation so an untracked task is not silently left running. The exact scoped duplicate-handle conflict is surfaced without cancellation because an existing durable row already owns that remote task. Durable recovery requires a SQL database backend (`sqlite` or `postgres`); the in-memory backend does not initialize this task repository. This foundation does not make existing MCP tools asynchronous by itself: `mcp_tasks.enabled` should remain `false` until a compatible task driver is configured. Ordinary `submit/status/cancel` tools and the future MCP Tasks extension can share the same runtime without making the model remember remote task IDs.
|
||||||
See the [MCP Server Guide](backend/docs/MCP_SERVER.md) for detailed instructions.
|
See the [MCP Server Guide](backend/docs/MCP_SERVER.md) for detailed instructions.
|
||||||
|
|
||||||
Security: pass per-request MCP credentials only through `config.context.secrets`;
|
Security: pass per-request MCP credentials only through `config.context.secrets`;
|
||||||
|
|||||||
@ -17,6 +17,7 @@ DeerFlow is a LangGraph-based AI super agent system with a full-stack architectu
|
|||||||
- Gateway streams `write_file` and `str_replace` argument deltas in bounded batches when clients also subscribe to `values`; messages-only consumers retain the original per-chunk contract, while `values` preserves the complete tool call.
|
- Gateway streams `write_file` and `str_replace` argument deltas in bounded batches when clients also subscribe to `values`; messages-only consumers retain the original per-chunk contract, while `values` preserves the complete tool call.
|
||||||
- With `stream_subgraphs`, subgraph frames keep their namespace in the SSE event name (`values|<ns>`, LangGraph Platform style) instead of impersonating root frames — a delegated subagent inherits the parent checkpoint namespace, so publishing its `values` snapshot as bare `values` replaces the whole thread view in SDK clients (#4399). Root-only consumers (file-tool chunk batcher, subagent event persistence, LLM error-fallback detection) ignore namespaced frames. The web frontend does not request subgraph streaming; subtask progress rides root-namespace `task_*` custom events.
|
- With `stream_subgraphs`, subgraph frames keep their namespace in the SSE event name (`values|<ns>`, LangGraph Platform style) instead of impersonating root frames — a delegated subagent inherits the parent checkpoint namespace, so publishing its `values` snapshot as bare `values` replaces the whole thread view in SDK clients (#4399). Root-only consumers (file-tool chunk batcher, subagent event persistence, LLM error-fallback detection) ignore namespaced frames. The web frontend does not request subgraph streaming; subtask progress rides root-namespace `task_*` custom events.
|
||||||
- Scheduled-task executions must reuse that same Gateway run lifecycle. The scheduler may decide *when* work runs, but it must dispatch through the existing run path rather than introducing a parallel execution stack.
|
- Scheduled-task executions must reuse that same Gateway run lifecycle. The scheduler may decide *when* work runs, but it must dispatch through the existing run path rather than introducing a parallel execution stack.
|
||||||
|
- Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. `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 must be discarded even when the owner token still matches. The database is the source of truth. `ThreadState` may receive only a bounded projection in later integration work, never the sole recoverable copy.
|
||||||
- Scheduled-task dispatch enforces "at most one active run per task when `overlap_policy=skip`" at the DB layer via the partial unique index `uq_scheduled_task_run_active` (`scheduled_task_runs.task_id WHERE status IN ('queued','running')`). `ScheduledTaskService.dispatch_task`'s `has_active_runs` check is a non-atomic fast path (its own session, separated from the `create()` insert by `await` points), so two concurrent dispatches — a manual `POST /scheduled-tasks/{id}/trigger` racing the poller, a double-click, or a client retry — can both pass it; the index is the atomic arbiter, and the losing `create` surfaces as `ActiveScheduledRunConflict` (translated from `IntegrityError` in the repository) and collapses to the same outcome as the fast path (manual → 409 conflict, scheduled → a `"skipped"` tombstone). The scheduled-skip tombstone is created directly as terminal `"skipped"` (not a transient `"queued"`) so it never occupies the active slot the pre-existing run still holds. Sibling of the `runs` table's `uq_runs_thread_active` (PR #4003), which keys on `thread_id` and so does not cover the default `fresh_thread_per_run` context where every dispatch gets a new thread. Index is status-only, not `overlap_policy`-conditional (the policy is fixed to `"skip"` in the MVP).
|
- Scheduled-task dispatch enforces "at most one active run per task when `overlap_policy=skip`" at the DB layer via the partial unique index `uq_scheduled_task_run_active` (`scheduled_task_runs.task_id WHERE status IN ('queued','running')`). `ScheduledTaskService.dispatch_task`'s `has_active_runs` check is a non-atomic fast path (its own session, separated from the `create()` insert by `await` points), so two concurrent dispatches — a manual `POST /scheduled-tasks/{id}/trigger` racing the poller, a double-click, or a client retry — can both pass it; the index is the atomic arbiter, and the losing `create` surfaces as `ActiveScheduledRunConflict` (translated from `IntegrityError` in the repository) and collapses to the same outcome as the fast path (manual → 409 conflict, scheduled → a `"skipped"` tombstone). The scheduled-skip tombstone is created directly as terminal `"skipped"` (not a transient `"queued"`) so it never occupies the active slot the pre-existing run still holds. Sibling of the `runs` table's `uq_runs_thread_active` (PR #4003), which keys on `thread_id` and so does not cover the default `fresh_thread_per_run` context where every dispatch gets a new thread. Index is status-only, not `overlap_policy`-conditional (the policy is fixed to `"skip"` in the MVP).
|
||||||
|
|
||||||
**Project Structure**:
|
**Project Structure**:
|
||||||
@ -502,7 +503,7 @@ Setup: Copy `config.example.yaml` to `config.yaml` in the **project root** direc
|
|||||||
|
|
||||||
**Config Hot-Reload Boundary**: Gateway dependencies route through `get_app_config()` on every request, so per-run fields like `models[*].max_tokens`, `summarization.*`, `title.*`, `memory.*`, `subagents.*`, `tools[*]`, and the agent system prompt pick up `config.yaml` edits on the next message. `AppConfig` is intentionally **not** cached on `app.state` — `lifespan()` keeps a local `startup_config` variable for one-shot bootstrap work and passes it to `langgraph_runtime(app, startup_config)`.
|
**Config Hot-Reload Boundary**: Gateway dependencies route through `get_app_config()` on every request, so per-run fields like `models[*].max_tokens`, `summarization.*`, `title.*`, `memory.*`, `subagents.*`, `tools[*]`, and the agent system prompt pick up `config.yaml` edits on the next message. `AppConfig` is intentionally **not** cached on `app.state` — `lifespan()` keeps a local `startup_config` variable for one-shot bootstrap work and passes it to `langgraph_runtime(app, startup_config)`.
|
||||||
|
|
||||||
Infrastructure fields are **restart-required**. The authoritative list lives in `packages/harness/deerflow/config/reload_boundary.py::STARTUP_ONLY_FIELDS` and is mirrored by the standardised `"startup-only:"` prefix on the corresponding `Field(description=...)` in `AppConfig`, so IDE hover on those fields surfaces the reason inline (no need to context-switch into this table). Currently registered: `plugins`, `database`, `checkpointer`, `run_events`, `stream_bridge`, `sandbox`, `log_level`, `logging`, `channels`, `channel_connections`, `scheduler`, `run_ownership`. Adding a new restart-required field requires updating the registry; drift is pinned by `tests/test_reload_boundary.py`.
|
Infrastructure fields are **restart-required**. The authoritative list lives in `packages/harness/deerflow/config/reload_boundary.py::STARTUP_ONLY_FIELDS` and is mirrored by the standardised `"startup-only:"` prefix on the corresponding `Field(description=...)` in `AppConfig`, so IDE hover on those fields surfaces the reason inline (no need to context-switch into this table). Currently registered: `plugins`, `database`, `checkpointer`, `run_events`, `stream_bridge`, `sandbox`, `log_level`, `logging`, `channels`, `channel_connections`, `scheduler`, `mcp_tasks`, `run_ownership`. Adding a new restart-required field requires updating the registry; drift is pinned by `tests/test_reload_boundary.py`.
|
||||||
|
|
||||||
**Persistence backend resolution**: the unified `database` section selects the
|
**Persistence backend resolution**: the unified `database` section selects the
|
||||||
Gateway's LangGraph checkpointer, LangGraph Store, and DeerFlow SQL repositories.
|
Gateway's LangGraph checkpointer, LangGraph Store, and DeerFlow SQL repositories.
|
||||||
@ -837,6 +838,7 @@ E2B output sync records remote file versions and actual host file metadata in a
|
|||||||
### MCP System (`packages/harness/deerflow/mcp/`)
|
### MCP System (`packages/harness/deerflow/mcp/`)
|
||||||
|
|
||||||
- Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management
|
- 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`). `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and a consecutive poll-error counter (incremented on failed polls, reset on any applied snapshot — the total `poll_attempt_count` grows on every claim and cannot distinguish failure streaks) for a later driver-layer backoff/terminal-failure policy; `app/mcp_tasks/McpTaskService` performs status calls outside the Agent/LLM loop. 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, the service best-effort cancels the remote task and preserves the original persistence error 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.
|
||||||
- **Lazy initialization**: Tools loaded on first use via `get_cached_mcp_tools()`
|
- **Lazy initialization**: Tools loaded on first use via `get_cached_mcp_tools()`
|
||||||
- **Cache invalidation**: Detects extensions-config changes by comparing the resolved config path and a `(mtime, size, sha256)` content signature against the values recorded at initialization, not a strict mtime `>` comparison. This catches same-second edits, mtime that stays put or moves backward (`git checkout`, `cp -p` / backup restore, `tar` / `rsync`, object-store / network mounts), and a switch to a different config file with an equal-or-older mtime. The signature helper (`config/file_signature.py::get_config_signature`) is shared with `config/app_config.py::get_app_config()` for the sibling runtime-editable config file, rather than each maintaining its own copy. `ExtensionsConfig.resolve_config_path()` raises `FileNotFoundError` for an explicit `config_path`/`DEER_FLOW_EXTENSIONS_CONFIG_PATH` that points at a missing file — an operator-asserted path going missing is a real misconfiguration, so this is intentionally loud for callers that load the config for actual use (e.g. `from_file()` via `get_mcp_tools()`); only the fallback search mode returns `None`. The MCP cache's own path resolution (`mcp/cache.py::_resolve_config_path`) is narrower: it catches that specific `FileNotFoundError` locally and treats it the same as "unconfigured", so this staleness check degrades to "not stale" instead of propagating an exception when a previously-valid explicit/env-var config disappears mid-run
|
- **Cache invalidation**: Detects extensions-config changes by comparing the resolved config path and a `(mtime, size, sha256)` content signature against the values recorded at initialization, not a strict mtime `>` comparison. This catches same-second edits, mtime that stays put or moves backward (`git checkout`, `cp -p` / backup restore, `tar` / `rsync`, object-store / network mounts), and a switch to a different config file with an equal-or-older mtime. The signature helper (`config/file_signature.py::get_config_signature`) is shared with `config/app_config.py::get_app_config()` for the sibling runtime-editable config file, rather than each maintaining its own copy. `ExtensionsConfig.resolve_config_path()` raises `FileNotFoundError` for an explicit `config_path`/`DEER_FLOW_EXTENSIONS_CONFIG_PATH` that points at a missing file — an operator-asserted path going missing is a real misconfiguration, so this is intentionally loud for callers that load the config for actual use (e.g. `from_file()` via `get_mcp_tools()`); only the fallback search mode returns `None`. The MCP cache's own path resolution (`mcp/cache.py::_resolve_config_path`) is narrower: it catches that specific `FileNotFoundError` locally and treats it the same as "unconfigured", so this staleness check degrades to "not stale" instead of propagating an exception when a previously-valid explicit/env-var config disappears mid-run
|
||||||
- **Transports**: stdio (command-based), SSE, HTTP
|
- **Transports**: stdio (command-based), SSE, HTTP
|
||||||
|
|||||||
@ -329,6 +329,26 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
|||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("Failed to initialize scheduled task service")
|
logger.exception("Failed to initialize scheduled task service")
|
||||||
|
|
||||||
|
try:
|
||||||
|
from app.mcp_tasks import McpTaskService
|
||||||
|
from deerflow.mcp.tasks import McpTaskDriverRegistry
|
||||||
|
|
||||||
|
if getattr(app.state, "mcp_task_repo", None) is not None:
|
||||||
|
mcp_task_drivers = McpTaskDriverRegistry()
|
||||||
|
mcp_task_service = McpTaskService(
|
||||||
|
repository=app.state.mcp_task_repo,
|
||||||
|
drivers=mcp_task_drivers,
|
||||||
|
poll_interval_seconds=startup_config.mcp_tasks.poll_interval_seconds,
|
||||||
|
lease_seconds=startup_config.mcp_tasks.lease_seconds,
|
||||||
|
max_concurrent_polls=startup_config.mcp_tasks.max_concurrent_polls,
|
||||||
|
)
|
||||||
|
app.state.mcp_task_drivers = mcp_task_drivers
|
||||||
|
app.state.mcp_task_service = mcp_task_service
|
||||||
|
if startup_config.mcp_tasks.enabled:
|
||||||
|
await mcp_task_service.start()
|
||||||
|
except Exception:
|
||||||
|
logger.exception("Failed to initialize MCP task service")
|
||||||
|
|
||||||
yield
|
yield
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@ -358,6 +378,12 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
|||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("Failed to stop scheduled task service")
|
logger.exception("Failed to stop scheduled task service")
|
||||||
|
|
||||||
|
if getattr(app.state, "mcp_task_service", None) is not None:
|
||||||
|
try:
|
||||||
|
await app.state.mcp_task_service.stop()
|
||||||
|
except Exception:
|
||||||
|
logger.exception("Failed to stop MCP task service")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from deerflow.community.browser_automation import get_browser_session_manager
|
from deerflow.community.browser_automation import get_browser_session_manager
|
||||||
|
|
||||||
|
|||||||
@ -416,14 +416,17 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen
|
|||||||
|
|
||||||
app.state.thread_store = make_thread_store(sf, app.state.store)
|
app.state.thread_store = make_thread_store(sf, app.state.store)
|
||||||
if sf is not None:
|
if sf is not None:
|
||||||
|
from deerflow.persistence.mcp_tasks import McpTaskRepository
|
||||||
from deerflow.persistence.scheduled_task_runs import (
|
from deerflow.persistence.scheduled_task_runs import (
|
||||||
ScheduledTaskRunRepository,
|
ScheduledTaskRunRepository,
|
||||||
)
|
)
|
||||||
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
|
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
|
||||||
|
|
||||||
|
app.state.mcp_task_repo = McpTaskRepository(sf)
|
||||||
app.state.scheduled_task_repo = ScheduledTaskRepository(sf)
|
app.state.scheduled_task_repo = ScheduledTaskRepository(sf)
|
||||||
app.state.scheduled_task_run_repo = ScheduledTaskRunRepository(sf)
|
app.state.scheduled_task_run_repo = ScheduledTaskRunRepository(sf)
|
||||||
else:
|
else:
|
||||||
|
app.state.mcp_task_repo = None
|
||||||
app.state.scheduled_task_repo = None
|
app.state.scheduled_task_repo = None
|
||||||
app.state.scheduled_task_run_repo = None
|
app.state.scheduled_task_run_repo = None
|
||||||
|
|
||||||
@ -577,6 +580,20 @@ def get_scheduled_task_service(request: Request):
|
|||||||
return val
|
return val
|
||||||
|
|
||||||
|
|
||||||
|
def get_mcp_task_repo(request: Request):
|
||||||
|
val = getattr(request.app.state, "mcp_task_repo", None)
|
||||||
|
if val is None:
|
||||||
|
raise HTTPException(status_code=503, detail="MCP task repo not available")
|
||||||
|
return val
|
||||||
|
|
||||||
|
|
||||||
|
def get_mcp_task_service(request: Request):
|
||||||
|
val = getattr(request.app.state, "mcp_task_service", None)
|
||||||
|
if val is None:
|
||||||
|
raise HTTPException(status_code=503, detail="MCP task service not available")
|
||||||
|
return val
|
||||||
|
|
||||||
|
|
||||||
def get_run_context(request: Request) -> RunContext:
|
def get_run_context(request: Request) -> RunContext:
|
||||||
"""Build a :class:`RunContext` from ``app.state`` singletons.
|
"""Build a :class:`RunContext` from ``app.state`` singletons.
|
||||||
|
|
||||||
|
|||||||
3
backend/app/mcp_tasks/__init__.py
Normal file
3
backend/app/mcp_tasks/__init__.py
Normal file
@ -0,0 +1,3 @@
|
|||||||
|
from app.mcp_tasks.service import McpTaskService
|
||||||
|
|
||||||
|
__all__ = ["McpTaskService"]
|
||||||
213
backend/app/mcp_tasks/service.py
Normal file
213
backend/app/mcp_tasks/service.py
Normal file
@ -0,0 +1,213 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import socket
|
||||||
|
import uuid
|
||||||
|
from dataclasses import replace
|
||||||
|
from datetime import UTC, datetime, timedelta
|
||||||
|
|
||||||
|
from deerflow.mcp.tasks import McpTaskDriverRegistry, TaskReference, TaskSnapshot, TaskSubmitRequest
|
||||||
|
from deerflow.persistence.mcp_tasks import DuplicateMcpRemoteTaskError
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_MAX_POLL_ERROR_CHARS = 4000
|
||||||
|
|
||||||
|
|
||||||
|
class McpTaskService:
|
||||||
|
"""Persist and poll long-running MCP tasks outside the Agent loop."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
repository,
|
||||||
|
drivers: McpTaskDriverRegistry,
|
||||||
|
poll_interval_seconds: int,
|
||||||
|
lease_seconds: int,
|
||||||
|
max_concurrent_polls: int,
|
||||||
|
) -> None:
|
||||||
|
self._repository = repository
|
||||||
|
self._drivers = drivers
|
||||||
|
self._poll_interval_seconds = poll_interval_seconds
|
||||||
|
self._lease_seconds = lease_seconds
|
||||||
|
self._max_concurrent_polls = max_concurrent_polls
|
||||||
|
self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}"
|
||||||
|
self._task: asyncio.Task[None] | None = None
|
||||||
|
self._stop = asyncio.Event()
|
||||||
|
|
||||||
|
@property
|
||||||
|
def drivers(self) -> McpTaskDriverRegistry:
|
||||||
|
return self._drivers
|
||||||
|
|
||||||
|
async def submit(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
driver_name: str,
|
||||||
|
request: TaskSubmitRequest,
|
||||||
|
now: datetime | None = None,
|
||||||
|
) -> dict:
|
||||||
|
"""Submit through one driver and persist the remote handle before returning."""
|
||||||
|
driver = self._drivers.get(driver_name)
|
||||||
|
if driver is None:
|
||||||
|
raise LookupError(f"No MCP task driver registered as {driver_name!r}")
|
||||||
|
|
||||||
|
submitted_at = now or datetime.now(UTC)
|
||||||
|
local_task_id = request.local_task_id or f"mcp-task-{uuid.uuid4().hex}"
|
||||||
|
driver_request = replace(request, local_task_id=local_task_id)
|
||||||
|
submission = await driver.submit(driver_request)
|
||||||
|
snapshot = submission.snapshot
|
||||||
|
next_poll_at = self._next_poll_at(snapshot, now=submitted_at)
|
||||||
|
driver_data = {**request.driver_data, **submission.driver_data}
|
||||||
|
task_reference = TaskReference(
|
||||||
|
local_task_id=local_task_id,
|
||||||
|
user_id=request.user_id,
|
||||||
|
thread_id=request.thread_id,
|
||||||
|
server_name=request.server_name,
|
||||||
|
remote_task_id=submission.remote_task_id,
|
||||||
|
driver_data=driver_data,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
return await self._repository.create(
|
||||||
|
task_id=local_task_id,
|
||||||
|
user_id=request.user_id,
|
||||||
|
thread_id=request.thread_id,
|
||||||
|
run_id=request.run_id,
|
||||||
|
tool_call_id=request.tool_call_id,
|
||||||
|
server_name=request.server_name,
|
||||||
|
driver_name=driver_name,
|
||||||
|
remote_task_id=submission.remote_task_id,
|
||||||
|
task_name=request.task_name,
|
||||||
|
status=snapshot.status.value,
|
||||||
|
result=snapshot.result,
|
||||||
|
error=snapshot.error,
|
||||||
|
input_required=snapshot.input_required,
|
||||||
|
next_poll_at=next_poll_at,
|
||||||
|
driver_data=driver_data,
|
||||||
|
)
|
||||||
|
except DuplicateMcpRemoteTaskError:
|
||||||
|
# This handle already has a durable owner. Cancelling it as
|
||||||
|
# compensation would terminate the pre-existing tracked task.
|
||||||
|
raise
|
||||||
|
except Exception:
|
||||||
|
try:
|
||||||
|
await driver.cancel(task_reference)
|
||||||
|
except Exception: # noqa: BLE001 - preserve the original persistence failure
|
||||||
|
logger.exception(
|
||||||
|
"Failed to cancel untracked MCP task after persistence failure (task_id=%s, driver=%s, remote_task_id=%s)",
|
||||||
|
local_task_id,
|
||||||
|
driver_name,
|
||||||
|
submission.remote_task_id,
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
async def run_once(self, *, now: datetime) -> None:
|
||||||
|
claimed = await self._repository.claim_due_tasks(
|
||||||
|
now=now,
|
||||||
|
lease_owner=self._lease_owner,
|
||||||
|
lease_seconds=self._lease_seconds,
|
||||||
|
limit=self._max_concurrent_polls,
|
||||||
|
)
|
||||||
|
if not claimed:
|
||||||
|
return
|
||||||
|
results = await asyncio.gather(
|
||||||
|
*(self._poll_one(task, now=now) for task in claimed),
|
||||||
|
return_exceptions=True,
|
||||||
|
)
|
||||||
|
for record, result in zip(claimed, results, strict=True):
|
||||||
|
if isinstance(result, BaseException):
|
||||||
|
logger.error(
|
||||||
|
"Unexpected MCP task poll failure (task_id=%s); the lease will expire for recovery",
|
||||||
|
record.get("id"),
|
||||||
|
exc_info=(type(result), result, result.__traceback__),
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _poll_one(self, record: dict, *, now: datetime) -> None:
|
||||||
|
driver_name = str(record.get("driver_name") or "")
|
||||||
|
driver = self._drivers.get(driver_name)
|
||||||
|
if driver is None:
|
||||||
|
await self._release_after_error(
|
||||||
|
record,
|
||||||
|
now=now,
|
||||||
|
error=f"No MCP task driver registered as {driver_name!r}",
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
snapshot = await driver.get_status(TaskReference.from_record(record))
|
||||||
|
except Exception as exc: # noqa: BLE001 - driver boundary; retry on the next poll
|
||||||
|
polled_at = datetime.now(UTC)
|
||||||
|
logger.warning(
|
||||||
|
"MCP task status poll failed (task_id=%s, driver=%s); retrying",
|
||||||
|
record.get("id"),
|
||||||
|
driver_name,
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
await self._release_after_error(record, now=polled_at, error=str(exc) or type(exc).__name__)
|
||||||
|
return
|
||||||
|
|
||||||
|
polled_at = datetime.now(UTC)
|
||||||
|
applied = await self._repository.apply_snapshot(
|
||||||
|
record["id"],
|
||||||
|
lease_owner=self._lease_owner,
|
||||||
|
status=snapshot.status.value,
|
||||||
|
result=snapshot.result,
|
||||||
|
error=snapshot.error,
|
||||||
|
input_required=snapshot.input_required,
|
||||||
|
next_poll_at=self._next_poll_at(snapshot, now=polled_at),
|
||||||
|
polled_at=polled_at,
|
||||||
|
)
|
||||||
|
if not applied:
|
||||||
|
logger.info(
|
||||||
|
"Discarded MCP task poll result after lease ownership changed or expired (task_id=%s)",
|
||||||
|
record.get("id"),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _next_poll_at(self, snapshot: TaskSnapshot, *, now: datetime) -> datetime | None:
|
||||||
|
if not snapshot.is_pollable:
|
||||||
|
return None
|
||||||
|
interval = snapshot.poll_after_seconds or self._poll_interval_seconds
|
||||||
|
return now + timedelta(seconds=interval)
|
||||||
|
|
||||||
|
async def _release_after_error(self, record: dict, *, now: datetime, error: str) -> None:
|
||||||
|
await self._repository.release_claim(
|
||||||
|
record["id"],
|
||||||
|
lease_owner=self._lease_owner,
|
||||||
|
next_poll_at=now + timedelta(seconds=self._poll_interval_seconds),
|
||||||
|
error=error[:_MAX_POLL_ERROR_CHARS],
|
||||||
|
)
|
||||||
|
|
||||||
|
async def start(self) -> None:
|
||||||
|
if self._task is not None:
|
||||||
|
return
|
||||||
|
self._stop.clear()
|
||||||
|
self._task = asyncio.create_task(self._run_loop(), name="deerflow-mcp-task-poller")
|
||||||
|
|
||||||
|
async def stop(self) -> None:
|
||||||
|
task = self._task
|
||||||
|
if task is None:
|
||||||
|
return
|
||||||
|
self._stop.set()
|
||||||
|
task.cancel()
|
||||||
|
try:
|
||||||
|
await task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
self._task = None
|
||||||
|
|
||||||
|
async def _run_loop(self) -> None:
|
||||||
|
while not self._stop.is_set():
|
||||||
|
try:
|
||||||
|
# The first pass runs immediately. Expired leases therefore
|
||||||
|
# recover at startup without a separate destructive sweep.
|
||||||
|
await self.run_once(now=datetime.now(UTC))
|
||||||
|
except Exception:
|
||||||
|
logger.exception("MCP task poll failed; retrying next interval")
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(
|
||||||
|
self._stop.wait(),
|
||||||
|
timeout=self._poll_interval_seconds,
|
||||||
|
)
|
||||||
|
except TimeoutError:
|
||||||
|
continue
|
||||||
@ -24,6 +24,7 @@ from deerflow.config.file_signature import get_config_signature as _get_config_s
|
|||||||
from deerflow.config.guardrails_config import GuardrailsConfig, load_guardrails_config_from_dict
|
from deerflow.config.guardrails_config import GuardrailsConfig, load_guardrails_config_from_dict
|
||||||
from deerflow.config.input_polish_config import InputPolishConfig
|
from deerflow.config.input_polish_config import InputPolishConfig
|
||||||
from deerflow.config.loop_detection_config import LoopDetectionConfig
|
from deerflow.config.loop_detection_config import LoopDetectionConfig
|
||||||
|
from deerflow.config.mcp_tasks_config import McpTasksConfig
|
||||||
from deerflow.config.memory_config import MemoryConfig, load_memory_config_from_dict
|
from deerflow.config.memory_config import MemoryConfig, load_memory_config_from_dict
|
||||||
from deerflow.config.model_config import ModelConfig
|
from deerflow.config.model_config import ModelConfig
|
||||||
from deerflow.config.read_before_write_config import ReadBeforeWriteConfig
|
from deerflow.config.read_before_write_config import ReadBeforeWriteConfig
|
||||||
@ -290,6 +291,13 @@ class AppConfig(BaseModel):
|
|||||||
field_doc="Scheduled task runtime configuration (background poller for one-time and cron agent runs).",
|
field_doc="Scheduled task runtime configuration (background poller for one-time and cron agent runs).",
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
mcp_tasks: McpTasksConfig = Field(
|
||||||
|
default_factory=McpTasksConfig,
|
||||||
|
description=format_field_description(
|
||||||
|
"mcp_tasks",
|
||||||
|
field_doc="Long-running MCP task persistence and background polling runtime.",
|
||||||
|
),
|
||||||
|
)
|
||||||
checkpointer: CheckpointerConfig | None = Field(
|
checkpointer: CheckpointerConfig | None = Field(
|
||||||
default=None,
|
default=None,
|
||||||
description=format_field_description(
|
description=format_field_description(
|
||||||
|
|||||||
10
backend/packages/harness/deerflow/config/mcp_tasks_config.py
Normal file
10
backend/packages/harness/deerflow/config/mcp_tasks_config.py
Normal file
@ -0,0 +1,10 @@
|
|||||||
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
|
|
||||||
|
class McpTasksConfig(BaseModel):
|
||||||
|
"""Startup configuration for the protocol-neutral MCP task poller."""
|
||||||
|
|
||||||
|
enabled: bool = Field(default=False)
|
||||||
|
poll_interval_seconds: int = Field(default=5, ge=1, le=300)
|
||||||
|
lease_seconds: int = Field(default=120, ge=5, le=3600)
|
||||||
|
max_concurrent_polls: int = Field(default=8, ge=1, le=64)
|
||||||
@ -70,6 +70,10 @@ STARTUP_ONLY_FIELDS: dict[str, str] = {
|
|||||||
"ScheduledTaskService is constructed and started once during Gateway lifespan startup; enabled, poll_interval_seconds, lease_seconds, "
|
"ScheduledTaskService is constructed and started once during Gateway lifespan startup; enabled, poll_interval_seconds, lease_seconds, "
|
||||||
"and max_concurrent_runs are captured into the service instance and the background poller task is not rebuilt on config.yaml edits."
|
"and max_concurrent_runs are captured into the service instance and the background poller task is not rebuilt on config.yaml edits."
|
||||||
),
|
),
|
||||||
|
"mcp_tasks": (
|
||||||
|
"McpTaskService is constructed and started once during Gateway lifespan startup; enabled, poll_interval_seconds, lease_seconds, "
|
||||||
|
"and max_concurrent_polls are captured into the service instance and the background poller task is not rebuilt on config.yaml edits."
|
||||||
|
),
|
||||||
"run_ownership": (
|
"run_ownership": (
|
||||||
"RunOwnershipConfig is captured once into RunManager at langgraph_runtime() startup; the lease heartbeat background task is created and "
|
"RunOwnershipConfig is captured once into RunManager at langgraph_runtime() startup; the lease heartbeat background task is created and "
|
||||||
"started there, and heartbeat_enabled / lease_seconds / grace_seconds are not re-read on config.yaml edits."
|
"started there, and heartbeat_enabled / lease_seconds / grace_seconds are not re-read on config.yaml edits."
|
||||||
|
|||||||
24
backend/packages/harness/deerflow/mcp/tasks/__init__.py
Normal file
24
backend/packages/harness/deerflow/mcp/tasks/__init__.py
Normal file
@ -0,0 +1,24 @@
|
|||||||
|
from deerflow.mcp.tasks.driver import McpTaskDriver, McpTaskDriverRegistry
|
||||||
|
from deerflow.mcp.tasks.models import (
|
||||||
|
ATTENTION_TASK_STATUSES,
|
||||||
|
POLLABLE_TASK_STATUSES,
|
||||||
|
TERMINAL_TASK_STATUSES,
|
||||||
|
TaskReference,
|
||||||
|
TaskSnapshot,
|
||||||
|
TaskStatus,
|
||||||
|
TaskSubmission,
|
||||||
|
TaskSubmitRequest,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"ATTENTION_TASK_STATUSES",
|
||||||
|
"McpTaskDriver",
|
||||||
|
"McpTaskDriverRegistry",
|
||||||
|
"POLLABLE_TASK_STATUSES",
|
||||||
|
"TERMINAL_TASK_STATUSES",
|
||||||
|
"TaskReference",
|
||||||
|
"TaskSnapshot",
|
||||||
|
"TaskStatus",
|
||||||
|
"TaskSubmission",
|
||||||
|
"TaskSubmitRequest",
|
||||||
|
]
|
||||||
36
backend/packages/harness/deerflow/mcp/tasks/driver.py
Normal file
36
backend/packages/harness/deerflow/mcp/tasks/driver.py
Normal file
@ -0,0 +1,36 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Protocol
|
||||||
|
|
||||||
|
from deerflow.mcp.tasks.models import TaskReference, TaskSnapshot, TaskSubmission, TaskSubmitRequest
|
||||||
|
|
||||||
|
|
||||||
|
class McpTaskDriver(Protocol):
|
||||||
|
"""Transport/protocol adapter used by the protocol-neutral task runtime."""
|
||||||
|
|
||||||
|
async def submit(self, request: TaskSubmitRequest) -> TaskSubmission: ...
|
||||||
|
|
||||||
|
async def get_status(self, task: TaskReference) -> TaskSnapshot: ...
|
||||||
|
|
||||||
|
async def cancel(self, task: TaskReference) -> TaskSnapshot: ...
|
||||||
|
|
||||||
|
|
||||||
|
class McpTaskDriverRegistry:
|
||||||
|
"""Process-local driver catalog wired at Gateway startup."""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._drivers: dict[str, McpTaskDriver] = {}
|
||||||
|
|
||||||
|
def register(self, name: str, driver: McpTaskDriver) -> None:
|
||||||
|
normalized = name.strip()
|
||||||
|
if not normalized:
|
||||||
|
raise ValueError("driver name must not be empty")
|
||||||
|
if normalized in self._drivers:
|
||||||
|
raise ValueError(f"MCP task driver {normalized!r} is already registered")
|
||||||
|
self._drivers[normalized] = driver
|
||||||
|
|
||||||
|
def get(self, name: str) -> McpTaskDriver | None:
|
||||||
|
return self._drivers.get(name)
|
||||||
|
|
||||||
|
def names(self) -> tuple[str, ...]:
|
||||||
|
return tuple(sorted(self._drivers))
|
||||||
115
backend/packages/harness/deerflow/mcp/tasks/models.py
Normal file
115
backend/packages/harness/deerflow/mcp/tasks/models.py
Normal file
@ -0,0 +1,115 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from enum import StrEnum
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
|
||||||
|
class TaskStatus(StrEnum):
|
||||||
|
"""Protocol-neutral lifecycle states for long-running MCP work."""
|
||||||
|
|
||||||
|
SUBMITTED = "submitted"
|
||||||
|
WORKING = "working"
|
||||||
|
INPUT_REQUIRED = "input_required"
|
||||||
|
COMPLETED = "completed"
|
||||||
|
FAILED = "failed"
|
||||||
|
CANCELLED = "cancelled"
|
||||||
|
|
||||||
|
|
||||||
|
POLLABLE_TASK_STATUSES: frozenset[TaskStatus] = frozenset(
|
||||||
|
{
|
||||||
|
TaskStatus.SUBMITTED,
|
||||||
|
TaskStatus.WORKING,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
TERMINAL_TASK_STATUSES: frozenset[TaskStatus] = frozenset(
|
||||||
|
{
|
||||||
|
TaskStatus.COMPLETED,
|
||||||
|
TaskStatus.FAILED,
|
||||||
|
TaskStatus.CANCELLED,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
ATTENTION_TASK_STATUSES: frozenset[TaskStatus] = frozenset(
|
||||||
|
{
|
||||||
|
TaskStatus.INPUT_REQUIRED,
|
||||||
|
*TERMINAL_TASK_STATUSES,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class TaskSnapshot:
|
||||||
|
"""One normalized status response returned by a task driver."""
|
||||||
|
|
||||||
|
status: TaskStatus
|
||||||
|
result: Any | None = None
|
||||||
|
error: str | None = None
|
||||||
|
input_required: dict[str, Any] | None = None
|
||||||
|
poll_after_seconds: float | None = None
|
||||||
|
|
||||||
|
def __post_init__(self) -> None:
|
||||||
|
if not isinstance(self.status, TaskStatus):
|
||||||
|
object.__setattr__(self, "status", TaskStatus(self.status))
|
||||||
|
if self.poll_after_seconds is not None and self.poll_after_seconds <= 0:
|
||||||
|
raise ValueError("poll_after_seconds must be positive")
|
||||||
|
if self.status == TaskStatus.INPUT_REQUIRED and self.input_required is None:
|
||||||
|
raise ValueError("input_required status requires an input_required payload")
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_pollable(self) -> bool:
|
||||||
|
return self.status in POLLABLE_TASK_STATUSES
|
||||||
|
|
||||||
|
@property
|
||||||
|
def needs_attention(self) -> bool:
|
||||||
|
return self.status in ATTENTION_TASK_STATUSES
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class TaskReference:
|
||||||
|
"""Stable data a driver needs after the originating Agent run has ended."""
|
||||||
|
|
||||||
|
local_task_id: str
|
||||||
|
user_id: str
|
||||||
|
thread_id: str
|
||||||
|
server_name: str
|
||||||
|
remote_task_id: str
|
||||||
|
driver_data: dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_record(cls, record: dict[str, Any]) -> TaskReference:
|
||||||
|
return cls(
|
||||||
|
local_task_id=record["id"],
|
||||||
|
user_id=record["user_id"],
|
||||||
|
thread_id=record["thread_id"],
|
||||||
|
server_name=record["server_name"],
|
||||||
|
remote_task_id=record["remote_task_id"],
|
||||||
|
driver_data=dict(record.get("driver_data") or {}),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class TaskSubmitRequest:
|
||||||
|
"""Protocol-neutral request passed to a driver by an MCP tool wrapper."""
|
||||||
|
|
||||||
|
user_id: str
|
||||||
|
thread_id: str
|
||||||
|
run_id: str | None
|
||||||
|
tool_call_id: str | None
|
||||||
|
server_name: str
|
||||||
|
task_name: str
|
||||||
|
arguments: dict[str, Any]
|
||||||
|
driver_data: dict[str, Any] = field(default_factory=dict)
|
||||||
|
local_task_id: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class TaskSubmission:
|
||||||
|
"""A durable remote handle plus its initial normalized state."""
|
||||||
|
|
||||||
|
remote_task_id: str
|
||||||
|
snapshot: TaskSnapshot
|
||||||
|
driver_data: dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
def __post_init__(self) -> None:
|
||||||
|
if not self.remote_task_id.strip():
|
||||||
|
raise ValueError("remote_task_id must not be empty")
|
||||||
@ -0,0 +1,4 @@
|
|||||||
|
from deerflow.persistence.mcp_tasks.model import McpTaskRow
|
||||||
|
from deerflow.persistence.mcp_tasks.sql import DuplicateMcpRemoteTaskError, McpTaskRepository
|
||||||
|
|
||||||
|
__all__ = ["DuplicateMcpRemoteTaskError", "McpTaskRepository", "McpTaskRow"]
|
||||||
@ -0,0 +1,55 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from sqlalchemy import JSON, DateTime, Index, Integer, String, Text, UniqueConstraint
|
||||||
|
from sqlalchemy.orm import Mapped, mapped_column
|
||||||
|
|
||||||
|
from deerflow.persistence.base import Base
|
||||||
|
|
||||||
|
|
||||||
|
class McpTaskRow(Base):
|
||||||
|
__tablename__ = "mcp_tasks"
|
||||||
|
|
||||||
|
id: Mapped[str] = mapped_column(String(64), primary_key=True)
|
||||||
|
user_id: Mapped[str] = mapped_column(String(64), index=True)
|
||||||
|
thread_id: Mapped[str] = mapped_column(String(64), index=True)
|
||||||
|
run_id: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||||
|
tool_call_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||||
|
server_name: Mapped[str] = mapped_column(String(128))
|
||||||
|
driver_name: Mapped[str] = mapped_column(String(64))
|
||||||
|
remote_task_id: Mapped[str] = mapped_column(String(255))
|
||||||
|
task_name: Mapped[str] = mapped_column(String(255))
|
||||||
|
status: Mapped[str] = mapped_column(String(32), index=True)
|
||||||
|
result: Mapped[Any | None] = mapped_column(JSON, nullable=True)
|
||||||
|
error: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||||
|
input_required: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True)
|
||||||
|
driver_data: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
|
||||||
|
notification_status: Mapped[str] = mapped_column(String(16), default="none", index=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)
|
||||||
|
poll_attempt_count: Mapped[int] = mapped_column(Integer, default=0)
|
||||||
|
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)
|
||||||
|
cancel_requested_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||||
|
completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||||
|
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
|
||||||
|
updated_at: Mapped[datetime] = mapped_column(
|
||||||
|
DateTime(timezone=True),
|
||||||
|
default=lambda: datetime.now(UTC),
|
||||||
|
onupdate=lambda: datetime.now(UTC),
|
||||||
|
)
|
||||||
|
|
||||||
|
__table_args__ = (
|
||||||
|
UniqueConstraint(
|
||||||
|
"user_id",
|
||||||
|
"server_name",
|
||||||
|
"remote_task_id",
|
||||||
|
name="uq_mcp_tasks_user_server_remote",
|
||||||
|
),
|
||||||
|
Index("ix_mcp_tasks_thread_created", "thread_id", "created_at"),
|
||||||
|
Index("ix_mcp_tasks_due", "status", "next_poll_at"),
|
||||||
|
)
|
||||||
240
backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py
Normal file
240
backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py
Normal file
@ -0,0 +1,240 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from datetime import UTC, datetime, timedelta
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from sqlalchemy import or_, select, update
|
||||||
|
from sqlalchemy.exc import IntegrityError
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||||
|
|
||||||
|
from deerflow.mcp.tasks import ATTENTION_TASK_STATUSES, POLLABLE_TASK_STATUSES, TERMINAL_TASK_STATUSES
|
||||||
|
from deerflow.persistence.mcp_tasks.model import McpTaskRow
|
||||||
|
from deerflow.utils.time import coerce_iso
|
||||||
|
|
||||||
|
_POLLABLE_STATUS_VALUES = tuple(status.value for status in POLLABLE_TASK_STATUSES)
|
||||||
|
_ATTENTION_STATUS_VALUES = frozenset(status.value for status in ATTENTION_TASK_STATUSES)
|
||||||
|
_TERMINAL_STATUS_VALUES = frozenset(status.value for status in TERMINAL_TASK_STATUSES)
|
||||||
|
_TIMESTAMP_FIELDS = (
|
||||||
|
"next_poll_at",
|
||||||
|
"last_polled_at",
|
||||||
|
"lease_expires_at",
|
||||||
|
"cancel_requested_at",
|
||||||
|
"completed_at",
|
||||||
|
"created_at",
|
||||||
|
"updated_at",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class DuplicateMcpRemoteTaskError(RuntimeError):
|
||||||
|
"""The current user already tracks this server's remote task handle."""
|
||||||
|
|
||||||
|
|
||||||
|
def _is_remote_task_unique_conflict(exc: IntegrityError) -> bool:
|
||||||
|
original = exc.orig
|
||||||
|
diagnostic = getattr(original, "diag", None)
|
||||||
|
if getattr(diagnostic, "constraint_name", None) == "uq_mcp_tasks_user_server_remote":
|
||||||
|
return True
|
||||||
|
message = str(original)
|
||||||
|
return "uq_mcp_tasks_user_server_remote" in message or "mcp_tasks.user_id, mcp_tasks.server_name, mcp_tasks.remote_task_id" in message
|
||||||
|
|
||||||
|
|
||||||
|
class McpTaskRepository:
|
||||||
|
"""Durable source of truth for long-running MCP task lifecycle state."""
|
||||||
|
|
||||||
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
||||||
|
self._sf = session_factory
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _row_to_dict(row: McpTaskRow) -> dict[str, Any]:
|
||||||
|
data = row.to_dict()
|
||||||
|
for key in _TIMESTAMP_FIELDS:
|
||||||
|
if data.get(key) is not None:
|
||||||
|
data[key] = coerce_iso(data[key])
|
||||||
|
return data
|
||||||
|
|
||||||
|
async def create(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
task_id: str,
|
||||||
|
user_id: str,
|
||||||
|
thread_id: str,
|
||||||
|
run_id: str | None,
|
||||||
|
tool_call_id: str | None,
|
||||||
|
server_name: str,
|
||||||
|
driver_name: str,
|
||||||
|
remote_task_id: str,
|
||||||
|
task_name: str,
|
||||||
|
status: str,
|
||||||
|
result: Any | None,
|
||||||
|
error: str | None,
|
||||||
|
input_required: dict[str, Any] | None,
|
||||||
|
next_poll_at: datetime | None,
|
||||||
|
driver_data: dict[str, Any] | None = None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
needs_attention = status in _ATTENTION_STATUS_VALUES
|
||||||
|
row = McpTaskRow(
|
||||||
|
id=task_id,
|
||||||
|
user_id=user_id,
|
||||||
|
thread_id=thread_id,
|
||||||
|
run_id=run_id,
|
||||||
|
tool_call_id=tool_call_id,
|
||||||
|
server_name=server_name,
|
||||||
|
driver_name=driver_name,
|
||||||
|
remote_task_id=remote_task_id,
|
||||||
|
task_name=task_name,
|
||||||
|
status=status,
|
||||||
|
result=result,
|
||||||
|
error=error,
|
||||||
|
input_required=input_required,
|
||||||
|
driver_data=dict(driver_data or {}),
|
||||||
|
notification_status="pending" if needs_attention else "none",
|
||||||
|
next_poll_at=next_poll_at,
|
||||||
|
completed_at=now if status in _TERMINAL_STATUS_VALUES else None,
|
||||||
|
created_at=now,
|
||||||
|
updated_at=now,
|
||||||
|
)
|
||||||
|
async with self._sf() as session:
|
||||||
|
session.add(row)
|
||||||
|
try:
|
||||||
|
await session.commit()
|
||||||
|
except IntegrityError as exc:
|
||||||
|
await session.rollback()
|
||||||
|
if _is_remote_task_unique_conflict(exc):
|
||||||
|
raise DuplicateMcpRemoteTaskError(f"Remote MCP task {remote_task_id!r} is already tracked for server {server_name!r} by this user") from exc
|
||||||
|
raise
|
||||||
|
await session.refresh(row)
|
||||||
|
return self._row_to_dict(row)
|
||||||
|
|
||||||
|
async def get(self, task_id: str, *, user_id: str) -> dict[str, Any] | None:
|
||||||
|
async with self._sf() as session:
|
||||||
|
row = await session.get(McpTaskRow, task_id)
|
||||||
|
if row is None or row.user_id != user_id:
|
||||||
|
return None
|
||||||
|
return self._row_to_dict(row)
|
||||||
|
|
||||||
|
async def list_by_thread(
|
||||||
|
self,
|
||||||
|
thread_id: str,
|
||||||
|
*,
|
||||||
|
user_id: str,
|
||||||
|
limit: int = 50,
|
||||||
|
active_only: bool = False,
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
stmt = select(McpTaskRow).where(
|
||||||
|
McpTaskRow.thread_id == thread_id,
|
||||||
|
McpTaskRow.user_id == user_id,
|
||||||
|
)
|
||||||
|
if active_only:
|
||||||
|
stmt = stmt.where(McpTaskRow.status.in_(_POLLABLE_STATUS_VALUES))
|
||||||
|
stmt = stmt.order_by(McpTaskRow.created_at.desc(), McpTaskRow.id.desc()).limit(limit)
|
||||||
|
async with self._sf() as session:
|
||||||
|
result = await session.execute(stmt)
|
||||||
|
return [self._row_to_dict(row) for row in result.scalars()]
|
||||||
|
|
||||||
|
async def claim_due_tasks(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
now: datetime,
|
||||||
|
lease_owner: str,
|
||||||
|
lease_seconds: int,
|
||||||
|
limit: int,
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
lease_expires_at = now + timedelta(seconds=lease_seconds)
|
||||||
|
stmt = (
|
||||||
|
select(McpTaskRow)
|
||||||
|
.where(
|
||||||
|
McpTaskRow.status.in_(_POLLABLE_STATUS_VALUES),
|
||||||
|
McpTaskRow.next_poll_at.is_not(None),
|
||||||
|
McpTaskRow.next_poll_at <= now,
|
||||||
|
or_(
|
||||||
|
McpTaskRow.lease_expires_at.is_(None),
|
||||||
|
McpTaskRow.lease_expires_at < now,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.order_by(McpTaskRow.next_poll_at.asc(), McpTaskRow.id.asc())
|
||||||
|
.limit(limit)
|
||||||
|
.with_for_update(skip_locked=True)
|
||||||
|
)
|
||||||
|
async with self._sf() as session:
|
||||||
|
result = await session.execute(stmt)
|
||||||
|
rows = list(result.scalars())
|
||||||
|
for row in rows:
|
||||||
|
row.lease_owner = lease_owner
|
||||||
|
row.lease_expires_at = lease_expires_at
|
||||||
|
row.poll_attempt_count += 1
|
||||||
|
row.updated_at = now
|
||||||
|
await session.commit()
|
||||||
|
return [self._row_to_dict(row) for row in rows]
|
||||||
|
|
||||||
|
async def apply_snapshot(
|
||||||
|
self,
|
||||||
|
task_id: str,
|
||||||
|
*,
|
||||||
|
lease_owner: str,
|
||||||
|
status: str,
|
||||||
|
result: Any | None,
|
||||||
|
error: str | None,
|
||||||
|
input_required: dict[str, Any] | None,
|
||||||
|
next_poll_at: datetime | None,
|
||||||
|
polled_at: datetime,
|
||||||
|
) -> bool:
|
||||||
|
values: dict[str, Any] = {
|
||||||
|
"status": status,
|
||||||
|
"result": result,
|
||||||
|
"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,
|
||||||
|
"updated_at": polled_at,
|
||||||
|
}
|
||||||
|
if status in _ATTENTION_STATUS_VALUES:
|
||||||
|
values["notification_status"] = "pending"
|
||||||
|
if status in _TERMINAL_STATUS_VALUES:
|
||||||
|
values["completed_at"] = polled_at
|
||||||
|
|
||||||
|
stmt = (
|
||||||
|
update(McpTaskRow)
|
||||||
|
.where(
|
||||||
|
McpTaskRow.id == task_id,
|
||||||
|
McpTaskRow.lease_owner == lease_owner,
|
||||||
|
McpTaskRow.lease_expires_at >= polled_at,
|
||||||
|
)
|
||||||
|
.values(**values)
|
||||||
|
)
|
||||||
|
async with self._sf() as session:
|
||||||
|
result_proxy = await session.execute(stmt)
|
||||||
|
await session.commit()
|
||||||
|
return bool(result_proxy.rowcount)
|
||||||
|
|
||||||
|
async def release_claim(
|
||||||
|
self,
|
||||||
|
task_id: str,
|
||||||
|
*,
|
||||||
|
lease_owner: str,
|
||||||
|
next_poll_at: datetime,
|
||||||
|
error: str,
|
||||||
|
) -> bool:
|
||||||
|
stmt = (
|
||||||
|
update(McpTaskRow)
|
||||||
|
.where(
|
||||||
|
McpTaskRow.id == task_id,
|
||||||
|
McpTaskRow.lease_owner == lease_owner,
|
||||||
|
)
|
||||||
|
.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,
|
||||||
|
updated_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
async with self._sf() as session:
|
||||||
|
result_proxy = await session.execute(stmt)
|
||||||
|
await session.commit()
|
||||||
|
return bool(result_proxy.rowcount)
|
||||||
@ -0,0 +1,73 @@
|
|||||||
|
"""durable long-running MCP tasks.
|
||||||
|
|
||||||
|
Revision ID: 0011_mcp_tasks
|
||||||
|
Revises: 0010_run_cancel_request
|
||||||
|
Create Date: 2026-08-04
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from collections.abc import Sequence
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
|
||||||
|
revision: str = "0011_mcp_tasks"
|
||||||
|
down_revision: str | Sequence[str] | None = "0010_run_cancel_request"
|
||||||
|
branch_labels: str | Sequence[str] | None = None
|
||||||
|
depends_on: str | Sequence[str] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
bind = op.get_bind()
|
||||||
|
if sa.inspect(bind).has_table("mcp_tasks"):
|
||||||
|
return
|
||||||
|
|
||||||
|
op.create_table(
|
||||||
|
"mcp_tasks",
|
||||||
|
sa.Column("id", sa.String(length=64), nullable=False),
|
||||||
|
sa.Column("user_id", sa.String(length=64), nullable=False),
|
||||||
|
sa.Column("thread_id", sa.String(length=64), nullable=False),
|
||||||
|
sa.Column("run_id", sa.String(length=64), nullable=True),
|
||||||
|
sa.Column("tool_call_id", sa.String(length=128), nullable=True),
|
||||||
|
sa.Column("server_name", sa.String(length=128), nullable=False),
|
||||||
|
sa.Column("driver_name", sa.String(length=64), nullable=False),
|
||||||
|
sa.Column("remote_task_id", sa.String(length=255), nullable=False),
|
||||||
|
sa.Column("task_name", sa.String(length=255), nullable=False),
|
||||||
|
sa.Column("status", sa.String(length=32), nullable=False),
|
||||||
|
sa.Column("result", sa.JSON(), nullable=True),
|
||||||
|
sa.Column("error", sa.Text(), nullable=True),
|
||||||
|
sa.Column("input_required", sa.JSON(), nullable=True),
|
||||||
|
sa.Column("driver_data", sa.JSON(), nullable=False),
|
||||||
|
sa.Column("notification_status", sa.String(length=16), nullable=False),
|
||||||
|
sa.Column("next_poll_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("last_polled_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("last_poll_error", sa.Text(), nullable=True),
|
||||||
|
sa.Column("poll_attempt_count", sa.Integer(), nullable=False),
|
||||||
|
sa.Column("consecutive_poll_error_count", sa.Integer(), nullable=False),
|
||||||
|
sa.Column("lease_owner", sa.String(length=128), nullable=True),
|
||||||
|
sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("cancel_requested_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True),
|
||||||
|
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||||
|
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
||||||
|
sa.PrimaryKeyConstraint("id"),
|
||||||
|
sa.UniqueConstraint(
|
||||||
|
"user_id",
|
||||||
|
"server_name",
|
||||||
|
"remote_task_id",
|
||||||
|
name="uq_mcp_tasks_user_server_remote",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
with op.batch_alter_table("mcp_tasks", schema=None) as batch_op:
|
||||||
|
batch_op.create_index("ix_mcp_tasks_user_id", ["user_id"], unique=False)
|
||||||
|
batch_op.create_index("ix_mcp_tasks_thread_id", ["thread_id"], unique=False)
|
||||||
|
batch_op.create_index("ix_mcp_tasks_status", ["status"], unique=False)
|
||||||
|
batch_op.create_index("ix_mcp_tasks_notification_status", ["notification_status"], unique=False)
|
||||||
|
batch_op.create_index("ix_mcp_tasks_next_poll_at", ["next_poll_at"], unique=False)
|
||||||
|
batch_op.create_index("ix_mcp_tasks_thread_created", ["thread_id", "created_at"], unique=False)
|
||||||
|
batch_op.create_index("ix_mcp_tasks_due", ["status", "next_poll_at"], unique=False)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_table("mcp_tasks")
|
||||||
@ -22,6 +22,7 @@ from deerflow.persistence.channel_connections.model import (
|
|||||||
ChannelOAuthStateRow,
|
ChannelOAuthStateRow,
|
||||||
)
|
)
|
||||||
from deerflow.persistence.feedback.model import FeedbackRow
|
from deerflow.persistence.feedback.model import FeedbackRow
|
||||||
|
from deerflow.persistence.mcp_tasks.model import McpTaskRow
|
||||||
from deerflow.persistence.models.run_event import RunEventRow
|
from deerflow.persistence.models.run_event import RunEventRow
|
||||||
from deerflow.persistence.run.model import RunRow
|
from deerflow.persistence.run.model import RunRow
|
||||||
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
|
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
|
||||||
@ -37,6 +38,7 @@ __all__ = [
|
|||||||
"ChannelCredentialRow",
|
"ChannelCredentialRow",
|
||||||
"ChannelOAuthStateRow",
|
"ChannelOAuthStateRow",
|
||||||
"FeedbackRow",
|
"FeedbackRow",
|
||||||
|
"McpTaskRow",
|
||||||
"RunEventRow",
|
"RunEventRow",
|
||||||
"RunRow",
|
"RunRow",
|
||||||
"ScheduledTaskRow",
|
"ScheduledTaskRow",
|
||||||
|
|||||||
25
backend/tests/test_mcp_task_config.py
Normal file
25
backend/tests/test_mcp_task_config.py
Normal file
@ -0,0 +1,25 @@
|
|||||||
|
import pytest
|
||||||
|
from pydantic import ValidationError
|
||||||
|
|
||||||
|
from deerflow.config.app_config import AppConfig
|
||||||
|
from deerflow.config.mcp_tasks_config import McpTasksConfig
|
||||||
|
from deerflow.config.reload_boundary import STARTUP_ONLY_FIELDS, STARTUP_ONLY_PREFIX
|
||||||
|
|
||||||
|
|
||||||
|
def test_mcp_task_runtime_is_disabled_by_default_and_bounded():
|
||||||
|
config = McpTasksConfig()
|
||||||
|
assert config.enabled is False
|
||||||
|
assert config.poll_interval_seconds == 5
|
||||||
|
assert config.lease_seconds == 120
|
||||||
|
assert config.max_concurrent_polls == 8
|
||||||
|
|
||||||
|
with pytest.raises(ValidationError):
|
||||||
|
McpTasksConfig(poll_interval_seconds=0)
|
||||||
|
with pytest.raises(ValidationError):
|
||||||
|
McpTasksConfig(max_concurrent_polls=0)
|
||||||
|
|
||||||
|
|
||||||
|
def test_mcp_task_runtime_is_registered_as_startup_only():
|
||||||
|
assert "mcp_tasks" in STARTUP_ONLY_FIELDS
|
||||||
|
field = AppConfig.model_fields["mcp_tasks"]
|
||||||
|
assert (field.description or "").startswith(STARTUP_ONLY_PREFIX)
|
||||||
27
backend/tests/test_mcp_task_models.py
Normal file
27
backend/tests/test_mcp_task_models.py
Normal file
@ -0,0 +1,27 @@
|
|||||||
|
import pytest
|
||||||
|
|
||||||
|
from deerflow.mcp.tasks import McpTaskDriverRegistry, TaskSnapshot, TaskStatus, TaskSubmission
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_snapshot_normalizes_string_statuses():
|
||||||
|
snapshot = TaskSnapshot(status="working") # type: ignore[arg-type]
|
||||||
|
assert snapshot.status is TaskStatus.WORKING
|
||||||
|
assert snapshot.is_pollable is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_input_required_snapshot_requires_payload():
|
||||||
|
with pytest.raises(ValueError, match="requires an input_required payload"):
|
||||||
|
TaskSnapshot(status=TaskStatus.INPUT_REQUIRED)
|
||||||
|
|
||||||
|
|
||||||
|
def test_submission_rejects_empty_remote_id():
|
||||||
|
with pytest.raises(ValueError, match="remote_task_id must not be empty"):
|
||||||
|
TaskSubmission(remote_task_id=" ", snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED))
|
||||||
|
|
||||||
|
|
||||||
|
def test_driver_registry_rejects_duplicate_names():
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
driver = object()
|
||||||
|
registry.register("ordinary", driver) # type: ignore[arg-type]
|
||||||
|
with pytest.raises(ValueError, match="already registered"):
|
||||||
|
registry.register("ordinary", driver) # type: ignore[arg-type]
|
||||||
305
backend/tests/test_mcp_task_repository.py
Normal file
305
backend/tests/test_mcp_task_repository.py
Normal file
@ -0,0 +1,305 @@
|
|||||||
|
from datetime import UTC, datetime, timedelta
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import pytest_asyncio
|
||||||
|
from sqlalchemy.exc import IntegrityError
|
||||||
|
|
||||||
|
from deerflow.config.database_config import DatabaseConfig
|
||||||
|
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
|
||||||
|
from deerflow.persistence.mcp_tasks import DuplicateMcpRemoteTaskError, McpTaskRepository
|
||||||
|
|
||||||
|
|
||||||
|
@pytest_asyncio.fixture(autouse=True)
|
||||||
|
async def _close_persistence_engine():
|
||||||
|
yield
|
||||||
|
await close_engine()
|
||||||
|
|
||||||
|
|
||||||
|
async def _make_repo(tmp_path) -> McpTaskRepository:
|
||||||
|
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||||
|
session_factory = get_session_factory()
|
||||||
|
assert session_factory is not None
|
||||||
|
return McpTaskRepository(session_factory)
|
||||||
|
|
||||||
|
|
||||||
|
async def _create_working_task(
|
||||||
|
repo: McpTaskRepository,
|
||||||
|
*,
|
||||||
|
task_id: str,
|
||||||
|
now: datetime,
|
||||||
|
user_id: str = "user-1",
|
||||||
|
remote_task_id: str | None = None,
|
||||||
|
) -> dict:
|
||||||
|
return await repo.create(
|
||||||
|
task_id=task_id,
|
||||||
|
user_id=user_id,
|
||||||
|
thread_id="thread-1",
|
||||||
|
run_id="run-1",
|
||||||
|
tool_call_id="call-1",
|
||||||
|
server_name="reports",
|
||||||
|
driver_name="fake",
|
||||||
|
remote_task_id=remote_task_id or f"remote-{task_id}",
|
||||||
|
task_name="Generate report",
|
||||||
|
status="working",
|
||||||
|
result=None,
|
||||||
|
error=None,
|
||||||
|
input_required=None,
|
||||||
|
next_poll_at=now - timedelta(seconds=1),
|
||||||
|
driver_data={"status_tool": "status"},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_remote_task_id_is_unique_per_user_and_server(tmp_path):
|
||||||
|
repo = await _make_repo(tmp_path)
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
await _create_working_task(
|
||||||
|
repo,
|
||||||
|
task_id="task-remote-1",
|
||||||
|
now=now,
|
||||||
|
remote_task_id="shared-remote-id",
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(DuplicateMcpRemoteTaskError, match="already tracked"):
|
||||||
|
await _create_working_task(
|
||||||
|
repo,
|
||||||
|
task_id="task-remote-2",
|
||||||
|
now=now,
|
||||||
|
remote_task_id="shared-remote-id",
|
||||||
|
)
|
||||||
|
|
||||||
|
other_user = await _create_working_task(
|
||||||
|
repo,
|
||||||
|
task_id="task-remote-3",
|
||||||
|
now=now,
|
||||||
|
user_id="user-2",
|
||||||
|
remote_task_id="shared-remote-id",
|
||||||
|
)
|
||||||
|
assert other_user["remote_task_id"] == "shared-remote-id"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_other_integrity_errors_are_not_duplicate_remote_tasks(tmp_path):
|
||||||
|
repo = await _make_repo(tmp_path)
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
await _create_working_task(repo, task_id="shared-local-id", now=now)
|
||||||
|
|
||||||
|
with pytest.raises(IntegrityError):
|
||||||
|
await _create_working_task(
|
||||||
|
repo,
|
||||||
|
task_id="shared-local-id",
|
||||||
|
now=now,
|
||||||
|
remote_task_id="different-remote-id",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_claim_due_tasks_skips_live_leases_and_reclaims_expired_ones(tmp_path):
|
||||||
|
repo = await _make_repo(tmp_path)
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
await _create_working_task(repo, task_id="task-1", now=now)
|
||||||
|
|
||||||
|
first = await repo.claim_due_tasks(
|
||||||
|
now=now,
|
||||||
|
lease_owner="worker-1",
|
||||||
|
lease_seconds=60,
|
||||||
|
limit=10,
|
||||||
|
)
|
||||||
|
assert [task["id"] for task in first] == ["task-1"]
|
||||||
|
|
||||||
|
while_live = await repo.claim_due_tasks(
|
||||||
|
now=now + timedelta(seconds=10),
|
||||||
|
lease_owner="worker-2",
|
||||||
|
lease_seconds=60,
|
||||||
|
limit=10,
|
||||||
|
)
|
||||||
|
assert while_live == []
|
||||||
|
|
||||||
|
reclaimed = await repo.claim_due_tasks(
|
||||||
|
now=now + timedelta(seconds=61),
|
||||||
|
lease_owner="worker-2",
|
||||||
|
lease_seconds=60,
|
||||||
|
limit=10,
|
||||||
|
)
|
||||||
|
assert [task["id"] for task in reclaimed] == ["task-1"]
|
||||||
|
assert reclaimed[0]["lease_owner"] == "worker-2"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task(tmp_path):
|
||||||
|
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(
|
||||||
|
now=now,
|
||||||
|
lease_owner="worker-new",
|
||||||
|
lease_seconds=60,
|
||||||
|
limit=10,
|
||||||
|
)
|
||||||
|
|
||||||
|
stale_applied = await repo.apply_snapshot(
|
||||||
|
"task-2",
|
||||||
|
lease_owner="worker-old",
|
||||||
|
status="failed",
|
||||||
|
result=None,
|
||||||
|
error="stale result",
|
||||||
|
input_required=None,
|
||||||
|
next_poll_at=None,
|
||||||
|
polled_at=now,
|
||||||
|
)
|
||||||
|
assert stale_applied is False
|
||||||
|
|
||||||
|
applied = await repo.apply_snapshot(
|
||||||
|
"task-2",
|
||||||
|
lease_owner="worker-new",
|
||||||
|
status="completed",
|
||||||
|
result={"report": "ready"},
|
||||||
|
error=None,
|
||||||
|
input_required=None,
|
||||||
|
next_poll_at=None,
|
||||||
|
polled_at=now,
|
||||||
|
)
|
||||||
|
assert applied is True
|
||||||
|
|
||||||
|
stored = await repo.get("task-2", user_id="user-1")
|
||||||
|
assert stored is not None
|
||||||
|
assert stored["status"] == "completed"
|
||||||
|
assert stored["result"] == {"report": "ready"}
|
||||||
|
assert stored["notification_status"] == "pending"
|
||||||
|
assert stored["lease_owner"] is None
|
||||||
|
|
||||||
|
assert (
|
||||||
|
await repo.claim_due_tasks(
|
||||||
|
now=now + timedelta(hours=1),
|
||||||
|
lease_owner="worker-3",
|
||||||
|
lease_seconds=60,
|
||||||
|
limit=10,
|
||||||
|
)
|
||||||
|
== []
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_apply_snapshot_rejects_result_after_same_workers_lease_expires(tmp_path):
|
||||||
|
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(
|
||||||
|
now=now,
|
||||||
|
lease_owner="worker-1",
|
||||||
|
lease_seconds=60,
|
||||||
|
limit=10,
|
||||||
|
)
|
||||||
|
|
||||||
|
applied = await repo.apply_snapshot(
|
||||||
|
"task-expired",
|
||||||
|
lease_owner="worker-1",
|
||||||
|
status="completed",
|
||||||
|
result={"report": "stale"},
|
||||||
|
error=None,
|
||||||
|
input_required=None,
|
||||||
|
next_poll_at=None,
|
||||||
|
polled_at=now + timedelta(seconds=61),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert applied is False
|
||||||
|
stored = await repo.get("task-expired", user_id="user-1")
|
||||||
|
assert stored is not None
|
||||||
|
assert stored["status"] == "working"
|
||||||
|
assert stored["result"] is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_input_required_is_persisted_and_paused_until_future_resume(tmp_path):
|
||||||
|
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(
|
||||||
|
now=now,
|
||||||
|
lease_owner="worker-1",
|
||||||
|
lease_seconds=60,
|
||||||
|
limit=10,
|
||||||
|
)
|
||||||
|
|
||||||
|
applied = await repo.apply_snapshot(
|
||||||
|
"task-3",
|
||||||
|
lease_owner="worker-1",
|
||||||
|
status="input_required",
|
||||||
|
result=None,
|
||||||
|
error=None,
|
||||||
|
input_required={"prompt": "Approve deployment?"},
|
||||||
|
next_poll_at=None,
|
||||||
|
polled_at=now,
|
||||||
|
)
|
||||||
|
assert applied is True
|
||||||
|
|
||||||
|
stored = await repo.get("task-3", user_id="user-1")
|
||||||
|
assert stored is not None
|
||||||
|
assert stored["input_required"] == {"prompt": "Approve deployment?"}
|
||||||
|
assert stored["notification_status"] == "pending"
|
||||||
|
assert stored["next_poll_at"] is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
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(
|
||||||
|
now=now,
|
||||||
|
lease_owner="worker-1",
|
||||||
|
lease_seconds=60,
|
||||||
|
limit=10,
|
||||||
|
)
|
||||||
|
retry_at = now + timedelta(seconds=30)
|
||||||
|
|
||||||
|
released = await repo.release_claim(
|
||||||
|
"task-4",
|
||||||
|
lease_owner="worker-1",
|
||||||
|
next_poll_at=retry_at,
|
||||||
|
error="temporary network failure",
|
||||||
|
)
|
||||||
|
assert released is True
|
||||||
|
|
||||||
|
stored = await repo.get("task-4", user_id="user-1")
|
||||||
|
assert stored is not None
|
||||||
|
assert stored["status"] == "working"
|
||||||
|
assert stored["last_poll_error"] == "temporary network failure"
|
||||||
|
assert datetime.fromisoformat(stored["next_poll_at"]) == retry_at
|
||||||
|
assert stored["lease_owner"] is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp_path):
|
||||||
|
repo = await _make_repo(tmp_path)
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
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)
|
||||||
|
await repo.release_claim(
|
||||||
|
"task-6",
|
||||||
|
lease_owner="worker-1",
|
||||||
|
next_poll_at=now - timedelta(seconds=1),
|
||||||
|
error="temporary network failure",
|
||||||
|
)
|
||||||
|
stored = await repo.get("task-6", user_id="user-1")
|
||||||
|
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)
|
||||||
|
applied = await repo.apply_snapshot(
|
||||||
|
"task-6",
|
||||||
|
lease_owner="worker-1",
|
||||||
|
status="working",
|
||||||
|
result=None,
|
||||||
|
error=None,
|
||||||
|
input_required=None,
|
||||||
|
next_poll_at=now + timedelta(seconds=5),
|
||||||
|
polled_at=now,
|
||||||
|
)
|
||||||
|
assert applied is True
|
||||||
|
|
||||||
|
stored = await repo.get("task-6", user_id="user-1")
|
||||||
|
assert stored is not None
|
||||||
|
assert stored["consecutive_poll_error_count"] == 0
|
||||||
458
backend/tests/test_mcp_task_service.py
Normal file
458
backend/tests/test_mcp_task_service.py
Normal file
@ -0,0 +1,458 @@
|
|||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
from datetime import UTC, datetime, timedelta
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from app.mcp_tasks.service import McpTaskService
|
||||||
|
from deerflow.mcp.tasks import (
|
||||||
|
McpTaskDriverRegistry,
|
||||||
|
TaskSnapshot,
|
||||||
|
TaskStatus,
|
||||||
|
TaskSubmission,
|
||||||
|
TaskSubmitRequest,
|
||||||
|
)
|
||||||
|
from deerflow.persistence.mcp_tasks import DuplicateMcpRemoteTaskError
|
||||||
|
|
||||||
|
|
||||||
|
class FakeRepository:
|
||||||
|
def __init__(self, rows=None):
|
||||||
|
self.rows = list(rows or [])
|
||||||
|
self.claimed = False
|
||||||
|
self.applied = []
|
||||||
|
self.released = []
|
||||||
|
self.created = []
|
||||||
|
|
||||||
|
async def create(self, **kwargs):
|
||||||
|
self.created.append(kwargs)
|
||||||
|
return {"id": kwargs["task_id"], **kwargs}
|
||||||
|
|
||||||
|
async def claim_due_tasks(self, **_kwargs):
|
||||||
|
if self.claimed:
|
||||||
|
return []
|
||||||
|
self.claimed = True
|
||||||
|
return [dict(row) for row in self.rows]
|
||||||
|
|
||||||
|
async def apply_snapshot(self, task_id, **kwargs):
|
||||||
|
self.applied.append((task_id, kwargs))
|
||||||
|
return True
|
||||||
|
|
||||||
|
async def release_claim(self, task_id, **kwargs):
|
||||||
|
self.released.append((task_id, kwargs))
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
class FailingApplyRepository(FakeRepository):
|
||||||
|
async def apply_snapshot(self, task_id, **kwargs):
|
||||||
|
if task_id == "task-1":
|
||||||
|
raise RuntimeError("database unavailable")
|
||||||
|
return await super().apply_snapshot(task_id, **kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
class FailingCreateRepository(FakeRepository):
|
||||||
|
async def create(self, **kwargs):
|
||||||
|
self.created.append(kwargs)
|
||||||
|
raise RuntimeError("database unavailable")
|
||||||
|
|
||||||
|
|
||||||
|
class DuplicateCreateRepository(FakeRepository):
|
||||||
|
async def create(self, **kwargs):
|
||||||
|
self.created.append(kwargs)
|
||||||
|
raise DuplicateMcpRemoteTaskError("already tracked")
|
||||||
|
|
||||||
|
|
||||||
|
class FakeDriver:
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
snapshots=None,
|
||||||
|
*,
|
||||||
|
submission=None,
|
||||||
|
error: Exception | None = None,
|
||||||
|
cancel_error: Exception | None = None,
|
||||||
|
):
|
||||||
|
self.snapshots = list(snapshots or [])
|
||||||
|
self.submission = submission
|
||||||
|
self.error = error
|
||||||
|
self.cancel_error = cancel_error
|
||||||
|
self.status_calls = []
|
||||||
|
self.submit_calls = []
|
||||||
|
self.cancel_calls = []
|
||||||
|
|
||||||
|
async def submit(self, request):
|
||||||
|
self.submit_calls.append(request)
|
||||||
|
if self.submission is None:
|
||||||
|
raise AssertionError(f"unexpected submit: {request}")
|
||||||
|
return self.submission
|
||||||
|
|
||||||
|
async def get_status(self, task):
|
||||||
|
self.status_calls.append(task)
|
||||||
|
if self.error is not None:
|
||||||
|
raise self.error
|
||||||
|
return self.snapshots.pop(0)
|
||||||
|
|
||||||
|
async def cancel(self, task):
|
||||||
|
self.cancel_calls.append(task)
|
||||||
|
if self.cancel_error is not None:
|
||||||
|
raise self.cancel_error
|
||||||
|
return TaskSnapshot(status=TaskStatus.CANCELLED)
|
||||||
|
|
||||||
|
|
||||||
|
class HangingDriver(FakeDriver):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.started = asyncio.Event()
|
||||||
|
self.cancelled = False
|
||||||
|
|
||||||
|
async def get_status(self, task):
|
||||||
|
self.status_calls.append(task)
|
||||||
|
self.started.set()
|
||||||
|
try:
|
||||||
|
await asyncio.Event().wait()
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
self.cancelled = True
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
def _claimed_row(*, driver_name="fake"):
|
||||||
|
return {
|
||||||
|
"id": "task-1",
|
||||||
|
"user_id": "user-1",
|
||||||
|
"thread_id": "thread-1",
|
||||||
|
"run_id": "run-1",
|
||||||
|
"tool_call_id": "call-1",
|
||||||
|
"server_name": "reports",
|
||||||
|
"driver_name": driver_name,
|
||||||
|
"remote_task_id": "remote-1",
|
||||||
|
"task_name": "Generate report",
|
||||||
|
"status": "working",
|
||||||
|
"driver_data": {"status_tool": "status"},
|
||||||
|
"lease_owner": "ignored-by-service-fixture",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_submit_persists_remote_handle_before_returning():
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
repo = FakeRepository()
|
||||||
|
driver = FakeDriver(
|
||||||
|
submission=TaskSubmission(
|
||||||
|
remote_task_id="remote-1",
|
||||||
|
snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED, poll_after_seconds=9),
|
||||||
|
driver_data={"status_tool": "status"},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("fake", driver)
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=5,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
request = TaskSubmitRequest(
|
||||||
|
user_id="user-1",
|
||||||
|
thread_id="thread-1",
|
||||||
|
run_id="run-1",
|
||||||
|
tool_call_id="call-1",
|
||||||
|
server_name="reports",
|
||||||
|
task_name="Generate report",
|
||||||
|
arguments={"topic": "MCP"},
|
||||||
|
driver_data={"submit_tool": "submit"},
|
||||||
|
)
|
||||||
|
|
||||||
|
created = await service.submit(driver_name="fake", request=request, now=now)
|
||||||
|
|
||||||
|
assert created["remote_task_id"] == "remote-1"
|
||||||
|
persisted = repo.created[0]
|
||||||
|
assert persisted["next_poll_at"] == now + timedelta(seconds=9)
|
||||||
|
assert persisted["driver_data"] == {"submit_tool": "submit", "status_tool": "status"}
|
||||||
|
assert driver.submit_calls[0].local_task_id == created["id"]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_submit_cancels_remote_task_when_persistence_fails():
|
||||||
|
repo = FailingCreateRepository()
|
||||||
|
driver = FakeDriver(
|
||||||
|
submission=TaskSubmission(
|
||||||
|
remote_task_id="remote-1",
|
||||||
|
snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED),
|
||||||
|
driver_data={"status_tool": "status", "cancel_tool": "cancel"},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("fake", driver)
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=5,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
request = TaskSubmitRequest(
|
||||||
|
user_id="user-1",
|
||||||
|
thread_id="thread-1",
|
||||||
|
run_id="run-1",
|
||||||
|
tool_call_id="call-1",
|
||||||
|
server_name="reports",
|
||||||
|
task_name="Generate report",
|
||||||
|
arguments={"topic": "MCP"},
|
||||||
|
driver_data={"submit_tool": "submit"},
|
||||||
|
local_task_id="task-1",
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="database unavailable"):
|
||||||
|
await service.submit(driver_name="fake", request=request)
|
||||||
|
|
||||||
|
assert len(driver.cancel_calls) == 1
|
||||||
|
cancelled = driver.cancel_calls[0]
|
||||||
|
assert cancelled.local_task_id == "task-1"
|
||||||
|
assert cancelled.remote_task_id == "remote-1"
|
||||||
|
assert cancelled.driver_data == {
|
||||||
|
"submit_tool": "submit",
|
||||||
|
"status_tool": "status",
|
||||||
|
"cancel_tool": "cancel",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_duplicate_remote_handle_is_rejected_without_cancelling_existing_task():
|
||||||
|
repo = DuplicateCreateRepository()
|
||||||
|
driver = FakeDriver(
|
||||||
|
submission=TaskSubmission(
|
||||||
|
remote_task_id="remote-1",
|
||||||
|
snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED),
|
||||||
|
driver_data={"cancel_tool": "cancel"},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("fake", driver)
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=5,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(DuplicateMcpRemoteTaskError, match="already tracked"):
|
||||||
|
await service.submit(
|
||||||
|
driver_name="fake",
|
||||||
|
request=TaskSubmitRequest(
|
||||||
|
user_id="user-1",
|
||||||
|
thread_id="thread-2",
|
||||||
|
run_id="run-2",
|
||||||
|
tool_call_id="call-2",
|
||||||
|
server_name="reports",
|
||||||
|
task_name="Generate report",
|
||||||
|
arguments={},
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert driver.cancel_calls == []
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_submit_preserves_persistence_error_when_compensation_cancel_fails(caplog):
|
||||||
|
repo = FailingCreateRepository()
|
||||||
|
driver = FakeDriver(
|
||||||
|
submission=TaskSubmission(
|
||||||
|
remote_task_id="remote-1",
|
||||||
|
snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED),
|
||||||
|
),
|
||||||
|
cancel_error=RuntimeError("cancel unavailable"),
|
||||||
|
)
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("fake", driver)
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=5,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
request = TaskSubmitRequest(
|
||||||
|
user_id="user-1",
|
||||||
|
thread_id="thread-1",
|
||||||
|
run_id="run-1",
|
||||||
|
tool_call_id="call-1",
|
||||||
|
server_name="reports",
|
||||||
|
task_name="Generate report",
|
||||||
|
arguments={"topic": "MCP"},
|
||||||
|
local_task_id="task-1",
|
||||||
|
)
|
||||||
|
|
||||||
|
with caplog.at_level(logging.ERROR), pytest.raises(RuntimeError, match="database unavailable"):
|
||||||
|
await service.submit(driver_name="fake", request=request)
|
||||||
|
|
||||||
|
assert "Failed to cancel untracked MCP task" in caplog.text
|
||||||
|
assert "cancel unavailable" in caplog.text
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_once_polls_without_an_llm_and_schedules_next_poll():
|
||||||
|
repo = FakeRepository([_claimed_row()])
|
||||||
|
driver = FakeDriver([TaskSnapshot(status=TaskStatus.WORKING, poll_after_seconds=12)])
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("fake", driver)
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=5,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
scan_started_at = datetime(2000, 1, 1, tzinfo=UTC)
|
||||||
|
|
||||||
|
await service.run_once(now=scan_started_at)
|
||||||
|
|
||||||
|
assert driver.status_calls[0].remote_task_id == "remote-1"
|
||||||
|
_, update = repo.applied[0]
|
||||||
|
assert update["status"] == "working"
|
||||||
|
assert update["next_poll_at"] == update["polled_at"] + timedelta(seconds=12)
|
||||||
|
assert update["polled_at"] > scan_started_at
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_once_schedules_driver_error_retry_from_poll_completion_time():
|
||||||
|
repo = FakeRepository([_claimed_row()])
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("fake", FakeDriver(error=RuntimeError("network down")))
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=5,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
scan_started_at = datetime(2000, 1, 1, tzinfo=UTC)
|
||||||
|
|
||||||
|
await service.run_once(now=scan_started_at)
|
||||||
|
|
||||||
|
_, released = repo.released[0]
|
||||||
|
retry_started_at = released["next_poll_at"] - timedelta(seconds=5)
|
||||||
|
assert retry_started_at > scan_started_at
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_once_stops_polling_terminal_and_input_required_snapshots():
|
||||||
|
rows = [_claimed_row(), {**_claimed_row(), "id": "task-2", "remote_task_id": "remote-2"}]
|
||||||
|
repo = FakeRepository(rows)
|
||||||
|
driver = FakeDriver(
|
||||||
|
[
|
||||||
|
TaskSnapshot(status=TaskStatus.COMPLETED, result={"report": "ready"}),
|
||||||
|
TaskSnapshot(status=TaskStatus.INPUT_REQUIRED, input_required={"prompt": "Approve?"}),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("fake", driver)
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=5,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
|
||||||
|
await service.run_once(now=datetime.now(UTC))
|
||||||
|
|
||||||
|
updates = {task_id: update for task_id, update in repo.applied}
|
||||||
|
assert updates["task-1"]["status"] == "completed"
|
||||||
|
assert updates["task-1"]["next_poll_at"] is None
|
||||||
|
assert updates["task-2"]["status"] == "input_required"
|
||||||
|
assert updates["task-2"]["input_required"] == {"prompt": "Approve?"}
|
||||||
|
assert updates["task-2"]["next_poll_at"] is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_once_releases_claim_when_driver_is_missing_or_fails():
|
||||||
|
rows = [_claimed_row(driver_name="missing"), {**_claimed_row(), "id": "task-2", "remote_task_id": "remote-2", "driver_name": "broken"}]
|
||||||
|
repo = FakeRepository(rows)
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("broken", FakeDriver(error=RuntimeError("network down")))
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=5,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
|
||||||
|
await service.run_once(now=now)
|
||||||
|
|
||||||
|
released = {task_id: update for task_id, update in repo.released}
|
||||||
|
assert "No MCP task driver registered" in released["task-1"]["error"]
|
||||||
|
assert released["task-2"]["error"] == "network down"
|
||||||
|
assert released["task-1"]["next_poll_at"] == now + timedelta(seconds=5)
|
||||||
|
assert released["task-2"]["next_poll_at"] > now + timedelta(seconds=5)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_run_once_isolates_unexpected_failure_to_its_claimed_task(caplog):
|
||||||
|
rows = [_claimed_row(), {**_claimed_row(), "id": "task-2", "remote_task_id": "remote-2"}]
|
||||||
|
repo = FailingApplyRepository(rows)
|
||||||
|
driver = FakeDriver(
|
||||||
|
[
|
||||||
|
TaskSnapshot(status=TaskStatus.COMPLETED, result={"report": "first"}),
|
||||||
|
TaskSnapshot(status=TaskStatus.COMPLETED, result={"report": "second"}),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("fake", driver)
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=5,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
|
||||||
|
with caplog.at_level(logging.ERROR):
|
||||||
|
await service.run_once(now=datetime.now(UTC))
|
||||||
|
|
||||||
|
assert [task_id for task_id, _update in repo.applied] == ["task-2"]
|
||||||
|
assert "task_id=task-1" in caplog.text
|
||||||
|
assert "database unavailable" in caplog.text
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_start_runs_recovery_poll_immediately_and_stop_is_clean():
|
||||||
|
repo = FakeRepository([])
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=McpTaskDriverRegistry(),
|
||||||
|
poll_interval_seconds=60,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
|
||||||
|
await service.start()
|
||||||
|
for _ in range(20):
|
||||||
|
if repo.claimed:
|
||||||
|
break
|
||||||
|
await __import__("asyncio").sleep(0)
|
||||||
|
await service.stop()
|
||||||
|
|
||||||
|
assert repo.claimed is True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_stop_cancels_a_hung_driver_poll():
|
||||||
|
repo = FakeRepository([_claimed_row()])
|
||||||
|
driver = HangingDriver()
|
||||||
|
registry = McpTaskDriverRegistry()
|
||||||
|
registry.register("fake", driver)
|
||||||
|
service = McpTaskService(
|
||||||
|
repository=repo,
|
||||||
|
drivers=registry,
|
||||||
|
poll_interval_seconds=60,
|
||||||
|
lease_seconds=120,
|
||||||
|
max_concurrent_polls=3,
|
||||||
|
)
|
||||||
|
|
||||||
|
await service.start()
|
||||||
|
await asyncio.wait_for(driver.started.wait(), timeout=1)
|
||||||
|
await asyncio.wait_for(service.stop(), timeout=1)
|
||||||
|
|
||||||
|
assert driver.cancelled is True
|
||||||
@ -157,7 +157,7 @@ async def test_migration_dedupes_duplicate_active_rows_before_unique_index(tmp_p
|
|||||||
with sqlite3.connect(db_path) as raw:
|
with sqlite3.connect(db_path) as raw:
|
||||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||||
# Bootstrap upgrades through the later revisions after 0004.
|
# Bootstrap upgrades through the later revisions after 0004.
|
||||||
assert version_row[0] == "0010_run_cancel_request"
|
assert version_row[0] == "0011_mcp_tasks"
|
||||||
|
|
||||||
# Sanity: the invariant the index enforces is now true — at most one
|
# Sanity: the invariant the index enforces is now true — at most one
|
||||||
# active row per thread.
|
# active row per thread.
|
||||||
|
|||||||
@ -169,7 +169,7 @@ async def test_migration_supersedes_duplicate_active_runs_before_unique_index(tm
|
|||||||
|
|
||||||
with sqlite3.connect(db_path) as raw:
|
with sqlite3.connect(db_path) as raw:
|
||||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||||
assert version_row[0] == "0010_run_cancel_request"
|
assert version_row[0] == "0011_mcp_tasks"
|
||||||
|
|
||||||
# Sanity: the invariant the index enforces now holds — at most one
|
# Sanity: the invariant the index enforces now holds — at most one
|
||||||
# active row per task_id.
|
# active row per task_id.
|
||||||
|
|||||||
@ -48,7 +48,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default
|
|||||||
asyncio_test = pytest.mark.asyncio
|
asyncio_test = pytest.mark.asyncio
|
||||||
|
|
||||||
|
|
||||||
HEAD = "0010_run_cancel_request"
|
HEAD = "0011_mcp_tasks"
|
||||||
BASELINE = "0001_baseline"
|
BASELINE = "0001_baseline"
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@ -28,7 +28,7 @@ from deerflow.persistence.bootstrap import bootstrap_schema
|
|||||||
pytestmark = pytest.mark.asyncio
|
pytestmark = pytest.mark.asyncio
|
||||||
|
|
||||||
|
|
||||||
HEAD = "0010_run_cancel_request"
|
HEAD = "0011_mcp_tasks"
|
||||||
|
|
||||||
|
|
||||||
def _url(tmp_path: Path) -> str:
|
def _url(tmp_path: Path) -> str:
|
||||||
|
|||||||
@ -76,7 +76,7 @@ async def test_legacy_database_recovers_token_usage_column(tmp_path: Path) -> No
|
|||||||
cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()}
|
cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()}
|
||||||
assert "token_usage_by_model" in cols
|
assert "token_usage_by_model" in cols
|
||||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||||
assert version_row[0] == "0010_run_cancel_request"
|
assert version_row[0] == "0011_mcp_tasks"
|
||||||
|
|
||||||
# And the read path that originally 500'd must now succeed.
|
# And the read path that originally 500'd must now succeed.
|
||||||
sf = get_session_factory()
|
sf = get_session_factory()
|
||||||
@ -116,6 +116,6 @@ async def test_legacy_database_with_manual_alter_still_bootstraps(tmp_path: Path
|
|||||||
# No duplicate column -- list, not set, to catch dupes.
|
# No duplicate column -- list, not set, to catch dupes.
|
||||||
assert cols.count("token_usage_by_model") == 1
|
assert cols.count("token_usage_by_model") == 1
|
||||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||||
assert version_row[0] == "0010_run_cancel_request"
|
assert version_row[0] == "0011_mcp_tasks"
|
||||||
finally:
|
finally:
|
||||||
await close_engine()
|
await close_engine()
|
||||||
|
|||||||
@ -2059,6 +2059,24 @@ scheduler:
|
|||||||
max_concurrent_runs: 3
|
max_concurrent_runs: 3
|
||||||
min_once_delay_seconds: 60
|
min_once_delay_seconds: 60
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# Long-running MCP Tasks Configuration
|
||||||
|
# ============================================================================
|
||||||
|
# Protocol-neutral durable task runtime. This foundation is disabled by
|
||||||
|
# default; an MCP task driver must be configured before tasks can be submitted.
|
||||||
|
# All fields are restart-required (captured at Gateway lifespan startup).
|
||||||
|
#
|
||||||
|
# mcp_tasks:
|
||||||
|
# enabled: false # Master switch for the background status poller
|
||||||
|
# poll_interval_seconds: 5 # Scan interval and default task retry interval
|
||||||
|
# lease_seconds: 120 # Expired claims become recoverable after this delay
|
||||||
|
# max_concurrent_polls: 8 # Maximum status calls started by one worker per scan
|
||||||
|
mcp_tasks:
|
||||||
|
enabled: false
|
||||||
|
poll_interval_seconds: 5
|
||||||
|
lease_seconds: 120
|
||||||
|
max_concurrent_polls: 8
|
||||||
|
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
# Run Ownership Configuration
|
# Run Ownership Configuration
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user