From 47b258ebd796b4aa0446bf9b9f4368c1bcc0a08e Mon Sep 17 00:00:00 2001 From: Aari Date: Sat, 15 Aug 2026 14:26:38 +0800 Subject: [PATCH] feat(mcp): add ordinary durable task driver (#4690) * 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 * feat(mcp): add ordinary durable task driver * test(mcp): address durable task review feedback * fix(mcp): preserve submit tool descriptions * fix(mcp): bound remote task calls * fix(mcp): bound persisted task payloads * fix(mcp): preserve task tool error details * fix(mcp): enforce durable task boundaries * test(mcp): cover task config snapshot lifecycle --------- Co-authored-by: Willem Jiang --- README.md | 5 +- backend/AGENTS.md | 2 +- backend/app/gateway/AGENTS.md | 1 + backend/app/gateway/app.py | 76 +++- backend/app/gateway/routers/mcp.py | 15 +- backend/app/gateway/routers/mcp_tasks.py | 95 +++++ backend/app/mcp_tasks/service.py | 134 ++++++- backend/docs/MCP_SERVER.md | 128 ++++++- .../harness/deerflow/config/AGENTS.md | 2 +- .../deerflow/config/extensions_config.py | 64 +++- .../deerflow/config/mcp_tasks_config.py | 5 + .../packages/harness/deerflow/constants.py | 9 + .../packages/harness/deerflow/mcp/AGENTS.md | 2 + .../harness/deerflow/mcp/interceptors.py | 57 +++ .../packages/harness/deerflow/mcp/oauth.py | 8 +- .../harness/deerflow/mcp/session_pool.py | 13 + .../harness/deerflow/mcp/task_tool_caller.py | 228 ++++++++++++ .../harness/deerflow/mcp/tasks/__init__.py | 8 + .../harness/deerflow/mcp/tasks/models.py | 32 +- .../harness/deerflow/mcp/tasks/ordinary.py | 213 ++++++++++++ .../harness/deerflow/mcp/tasks/runtime.py | 102 ++++++ .../packages/harness/deerflow/mcp/tools.py | 186 ++++++++-- .../deerflow/persistence/mcp_tasks/model.py | 16 +- .../deerflow/persistence/mcp_tasks/sql.py | 12 + .../deerflow/persistence/migrations/AGENTS.md | 2 + .../versions/0012_mcp_task_results.py | 47 +++ .../tests/test_gateway_lifespan_shutdown.py | 61 ++++ backend/tests/test_mcp_session_pool.py | 35 ++ backend/tests/test_mcp_task_config.py | 7 + backend/tests/test_mcp_task_models.py | 18 +- .../tests/test_mcp_task_ordinary_driver.py | 234 +++++++++++++ backend/tests/test_mcp_task_ordinary_e2e.py | 189 ++++++++++ backend/tests/test_mcp_task_repository.py | 28 +- backend/tests/test_mcp_task_runtime_config.py | 139 ++++++++ backend/tests/test_mcp_task_service.py | 327 +++++++++++++++++- backend/tests/test_mcp_task_tool_caller.py | 295 ++++++++++++++++ backend/tests/test_mcp_task_tool_wrapping.py | 203 +++++++++++ backend/tests/test_mcp_task_toolset_config.py | 133 +++++++ backend/tests/test_mcp_tasks_router.py | 140 ++++++++ backend/tests/test_mcp_tool_name_prefix.py | 2 +- ...est_migration_0004_run_ownership_dedupe.py | 2 +- ...ration_0007_scheduled_run_active_dedupe.py | 2 +- backend/tests/test_persistence_bootstrap.py | 2 +- .../test_persistence_bootstrap_concurrency.py | 2 +- .../test_persistence_bootstrap_regression.py | 4 +- config.example.yaml | 15 +- extensions_config.example.json | 16 + 47 files changed, 3217 insertions(+), 99 deletions(-) create mode 100644 backend/app/gateway/routers/mcp_tasks.py create mode 100644 backend/packages/harness/deerflow/mcp/interceptors.py create mode 100644 backend/packages/harness/deerflow/mcp/task_tool_caller.py create mode 100644 backend/packages/harness/deerflow/mcp/tasks/ordinary.py create mode 100644 backend/packages/harness/deerflow/mcp/tasks/runtime.py create mode 100644 backend/packages/harness/deerflow/persistence/migrations/versions/0012_mcp_task_results.py create mode 100644 backend/tests/test_mcp_task_ordinary_driver.py create mode 100644 backend/tests/test_mcp_task_ordinary_e2e.py create mode 100644 backend/tests/test_mcp_task_runtime_config.py create mode 100644 backend/tests/test_mcp_task_tool_caller.py create mode 100644 backend/tests/test_mcp_task_tool_wrapping.py create mode 100644 backend/tests/test_mcp_task_toolset_config.py create mode 100644 backend/tests/test_mcp_tasks_router.py diff --git a/README.md b/README.md index 8444d5ec5..82aff9c22 100644 --- a/README.md +++ b/README.md @@ -420,7 +420,7 @@ See the [Sandbox Configuration Guide](backend/docs/CONFIGURATION.md#sandbox) to DeerFlow supports configurable MCP servers and skills to extend its capabilities. For HTTP/SSE MCP servers, OAuth token flows are supported (`client_credentials`, `refresh_token`). -For stdio MCP servers, per-tool call timeouts can be configured with `tool_call_timeout`. +For stdio MCP servers, per-tool call timeouts can be configured with `tool_call_timeout`; durable background-task calls honor the same setting for HTTP/SSE servers as well. MCP tool names are prefixed with `_` by default to prevent collisions across servers. If a server already namespaces its own tools, set `tool_name_prefix: false` on that server in `extensions_config.json` to keep the original names. Disable the prefix only when the resulting names remain unique across all enabled servers. Settings > Tools updates one MCP server at a time: an invalid stdio command on one server no longer blocks toggling another, while enabling that invalid server remains protected by the command allowlist and surfaces the backend validation message in the UI. Targeted updates accept both DeerFlow's `type` field and the MCP-spec `transport` field for SSE/HTTP servers. @@ -435,7 +435,8 @@ explicit, model-selected MCP tool path can run alongside the separate automatic OpenViking memory backend; it does not replace automatic turn capture or recall. See the [OpenViking MCP tools configuration](backend/docs/MCP_SERVER.md#openviking-mcp-tools). -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. +The Gateway can adapt an MCP server's ordinary `submit` / `status` / `cancel` tools into durable background tasks. The Agent sees only the configured submit tool and a DeerFlow-local task ID; remote IDs are persisted before the submit call returns, while status and cancel stay internal to the runtime. Polling uses cross-worker leases, exponential retry backoff, scoped MCP sessions, bounded result storage, and restart recovery. A status-tool `isError` is retained as a bounded diagnostic and retried; servers report a permanent remote-task outcome through a normal structured result with `status: "failed"`. Remote poll hints are finite positive numbers capped at 24 hours, artifact-reference JSON is limited to 64 KiB, and task/server identifiers are validated against their durable SQL column limits before persistence. Current-thread tasks are available through `GET /api/threads/{thread_id}/mcp-tasks` and its detail endpoint. Enable `mcp_tasks` in `config.yaml`, configure `task_toolsets` with exact raw tool names in `extensions_config.json`, and use a SQL database backend (`sqlite` or `postgres`). Task-enabled server connection, authentication, interceptor, timeout, or binding changes require a Gateway restart so Agent tool discovery and background calls cannot use different configuration versions. This phase does not yet wake the Agent when a task completes or add a frontend task panel. + See the [MCP Server Guide](backend/docs/MCP_SERVER.md) for detailed instructions. Security: pass per-request MCP credentials only through `config.context.secrets`; diff --git a/backend/AGENTS.md b/backend/AGENTS.md index f1a5ae044..fe5bdf7f9 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -19,7 +19,7 @@ DeerFlow is a LangGraph-based AI super agent system with a full-stack architectu - Background subagent identity is deliberately split: the provider `tool_call_id` remains the correlation key for `ToolMessage`, `task_*` SSE events, persisted lifecycle events, frontend cards, and the public `ExtensionData.scope_id` contract (stored as `SubagentResult.external_task_id`), while `SubagentExecutor.execute_async()` generates a full server-side `execution_id` for `SubagentResult.task_id`, the process-wide registry, polling, cancellation, timeout handling, and cleanup. Provider IDs are not globally unique across parent runs, so they must never become registry ownership keys; scheduler closures retain their own `SubagentResult` rather than resolving ownership again through the mutable registry. Terminal subagent token usage travels in the current run's `ToolMessage.additional_kwargs` and is attributed from message state, never through a process-global provider-ID cache. - 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. - The background scheduler is single-instance by default. `scheduler.multi_instance=true` opts into lease-aware recovery across Gateway instances and requires shared Postgres, `run_ownership.heartbeat_enabled=true`, and `run_events.backend=db`; otherwise startup rejects the configuration. Live scheduled runs are preserved when a peer starts; expired leases are atomically taken over, stale post-launch writes are fenced by the dispatch lease owner, and the Postgres advisory-locked budget makes `max_concurrent_runs` a shared global cap (including pre-launch reservations). -- 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. +- Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. Explicit `task_toolsets` bind raw submit/status/cancel names; only submit remains Agent-visible, and its wrapper persists the remote handle before returning a local ID. `McpTaskService` claims due rows with leases, resolves a protocol-specific `McpTaskDriver`, and writes normalized snapshots back to `mcp_tasks`; expired leases are the restart-recovery mechanism, and a result returned after expiry 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). **Project Structure**: diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index b2a182074..84026cb5b 100644 --- a/backend/app/gateway/AGENTS.md +++ b/backend/app/gateway/AGENTS.md @@ -16,6 +16,7 @@ Localhost persistence deliberately reads the direct request `Host` and ignores ` | **Features** (`/api/features`) | `GET /` - report config-gated feature availability (`agents_api.enabled`, `browser_control.enabled`) for frontend UI gating | | **Console** (`/api/console`) | Read-only cross-thread observability for the current user (the data layer for an operations dashboard or external monitoring): `GET /stats` - headline counters (runs/threads/agents/tokens/cost); `GET /runs` - paginated run history joined with thread titles (per-run cost); `GET /usage` - zero-filled daily token series + per-model breakdown with spend. Queries `runs`/`threads_meta` directly as a reporting layer (no new `RunStore` methods); requires a SQL database backend — returns 503 on `database.backend: memory`. Real-cost estimation reads optional `models[*].pricing` (`currency`, `input_per_million`, `output_per_million`, `input_cache_hit_per_million`; `ModelConfig` is `extra="allow"`, so no schema change) and prices each run from its `token_usage_by_model` input/output split. Pricing is **cache-aware**: `RunJournal` accumulates prompt-cache hits from `usage_metadata.input_token_details.cache_read` into a sparse `cache_read_tokens` bucket key (also threaded through `SubagentTokenCollector` → `record_external_llm_usage_records`), and cache-hit input tokens are billed at `input_cache_hit_per_million` (omitted → billed at the miss price, a conservative upper bound). All priced models must use one currency; mixed currencies disable cost reporting and leave cost/currency fields null instead of producing invalid aggregates. Legacy rows fall back to run-level totals at `model_name`; unpriced models yield `cost: null` and cost fields are null when no pricing is configured | | **MCP** (`/api/mcp`) | `GET /config` - get config; `PUT /config` - replace the full config with whole-payload stdio validation; `PATCH /config` - toggle one server while preserving the raw extensions config and validating only an enabled target; both writes reload config and reset the process-local MCP cache | +| **MCP Tasks** (`/api/threads/{id}/mcp-tasks`) | `GET /` - current user's durable tasks for one owned thread; `GET /{task_id}` - bounded result/input/error detail without remote task IDs or driver configuration | | **Skills** (`/api/skills`) | `GET /` - list skills; `GET /{name}` - details; `PUT /{name}` - update enabled; `POST /install` - install from .skill archive (accepts standard optional frontmatter like `version`, `author`, `compatibility`); `POST /reload` - admin-only process-local prompt-cache invalidation after trusted external filesystem changes | | **Integrations** (`/api/integrations`) | `GET /lark/status` - inspect managed Lark/Feishu CLI integration state, including `sandbox_runtime_mode` / `sandbox_runtime_ready` (whether `lark-cli` will actually be present in the sandbox at chat time); `POST /lark/install` - admin-only install of the official `lark-*` managed skill pack; `POST /lark/config/start` and `/lark/config/complete` - internal first-time Lark connection setup; `POST /lark/config/credentials` - atomically switch the caller's per-user Lark app after validating the new `app_id`/`app_secret` through the official CLI's live tenant-token probe, revoke/remove the previous OAuth tokens, and restore the prior credential tree if the switch fails; `POST /lark/auth/start` and `/lark/auth/complete` - browser device-flow user authorization without terminal access, with optional `domains` / exact `scope` for incremental permission grants. Config and auth flows carry a server-issued, per-user generation persisted under the credential lock; a rejected direct switch leaves the current generation unchanged, stale completions return 409, and browser re-registration uses the same token-clearing/revocation transaction as direct credential switches. | | **Memory** (`/api/memory`) | `GET /` - memory data; `POST /reload` - force reload; `GET /config` - config; `GET /status` - config + data | diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index 9cc6f8b81..6ac1149de 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -27,6 +27,7 @@ from app.gateway.routers import ( input_polish, integrations, mcp, + mcp_tasks, memory, models, runs, @@ -331,25 +332,56 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: except Exception: logger.exception("Failed to initialize scheduled task service") - try: - from app.mcp_tasks import McpTaskService - from deerflow.mcp.tasks import McpTaskDriverRegistry + from app.mcp_tasks import McpTaskService + from deerflow.config.extensions_config import ExtensionsConfig + from deerflow.config.mcp_tasks_config import McpTasksConfig + from deerflow.mcp.task_tool_caller import McpTaskToolCaller + from deerflow.mcp.tasks import ( + ORDINARY_MCP_TASK_DRIVER, + McpTaskDriverRegistry, + OrdinaryMcpTaskDriver, + ) + from deerflow.mcp.tasks.runtime import ( + configured_task_toolset_count, + set_mcp_task_config_snapshot, + set_mcp_task_submitter, + validate_mcp_task_runtime_configuration, + ) - 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, + task_extensions_config = ExtensionsConfig.from_file() + mcp_tasks_config = getattr(startup_config, "mcp_tasks", McpTasksConfig()) + mcp_task_repo = getattr(app.state, "mcp_task_repo", None) + set_mcp_task_submitter(None) + set_mcp_task_config_snapshot(task_extensions_config) + validate_mcp_task_runtime_configuration( + mcp_tasks_config=mcp_tasks_config, + extensions_config=task_extensions_config, + repository_available=mcp_task_repo is not None, + ) + if mcp_task_repo is not None: + mcp_task_drivers = McpTaskDriverRegistry() + if configured_task_toolset_count(task_extensions_config): + mcp_task_drivers.register( + ORDINARY_MCP_TASK_DRIVER, + OrdinaryMcpTaskDriver(McpTaskToolCaller(task_extensions_config)), ) - 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") + mcp_task_service = McpTaskService( + repository=mcp_task_repo, + drivers=mcp_task_drivers, + poll_interval_seconds=mcp_tasks_config.poll_interval_seconds, + lease_seconds=mcp_tasks_config.lease_seconds, + max_concurrent_polls=mcp_tasks_config.max_concurrent_polls, + max_poll_backoff_seconds=mcp_tasks_config.max_poll_backoff_seconds, + input_required_poll_interval_seconds=mcp_tasks_config.input_required_poll_interval_seconds, + tracking_degraded_after_errors=mcp_tasks_config.tracking_degraded_after_errors, + max_result_bytes=mcp_tasks_config.max_result_bytes, + result_preview_max_chars=mcp_tasks_config.result_preview_max_chars, + ) + app.state.mcp_task_drivers = mcp_task_drivers + app.state.mcp_task_service = mcp_task_service + if mcp_tasks_config.enabled: + await mcp_task_service.start() + set_mcp_task_submitter(mcp_task_service) yield @@ -385,6 +417,13 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: await app.state.mcp_task_service.stop() except Exception: logger.exception("Failed to stop MCP task service") + finally: + from deerflow.mcp.tasks.runtime import set_mcp_task_submitter + + set_mcp_task_submitter(None) + from deerflow.mcp.tasks.runtime import set_mcp_task_config_snapshot + + set_mcp_task_config_snapshot(None) try: from deerflow.community.browser_automation import get_browser_session_manager @@ -660,6 +699,9 @@ This gateway provides runtime endpoints for agent runs plus custom endpoints for # MCP API is mounted at /api/mcp app.include_router(mcp.router) + # Durable MCP tasks are scoped to their owning thread. + app.include_router(mcp_tasks.router) + # Memory API is mounted at /api/memory app.include_router(memory.router) diff --git a/backend/app/gateway/routers/mcp.py b/backend/app/gateway/routers/mcp.py index 1c97c8cae..29bae875f 100644 --- a/backend/app/gateway/routers/mcp.py +++ b/backend/app/gateway/routers/mcp.py @@ -13,6 +13,7 @@ from app.gateway.deps import require_admin_user from deerflow.config.extensions_config import ( ExtensionsConfig, McpRoutingConfig, + McpTaskToolsetConfig, McpToolOverride, atomic_write_extensions_config, extensions_config_write_lock, @@ -379,11 +380,21 @@ class McpServerConfigResponse(BaseModel): routing: McpRoutingConfig = Field(default_factory=McpRoutingConfig, description="Soft routing hints for tools from this MCP server") tools: dict[str, McpToolOverride] = Field(default_factory=dict, description="Per-original-tool MCP configuration overrides") tool_name_prefix: bool = Field(default=True, description="Whether to prefix discovered tool names with the MCP server name") - tool_call_timeout: float | None = Field(default=None, description="Timeout in seconds for individual stdio MCP tool calls") + tool_call_timeout: float | None = Field( + default=None, + description="Timeout in seconds for individual stdio MCP calls and durable-task calls on every transport", + ) # Default matches McpServerConfig: this model's defaults feed model_dump() # into the persisted extensions config on PUT, so an API-created server that # omits the field must get the same bring-up timeout as a file-created one. - session_init_timeout: float | None = Field(default=DEFAULT_MCP_SESSION_INIT_TIMEOUT, description="Timeout in seconds for MCP server bring-up (tool discovery and persistent stdio session initialization); null means no timeout") + session_init_timeout: float | None = Field( + default=DEFAULT_MCP_SESSION_INIT_TIMEOUT, + description="Timeout in seconds for MCP server bring-up and durable HTTP/SSE task-session initialization; null means no timeout", + ) + task_toolsets: list[McpTaskToolsetConfig] = Field( + default_factory=list, + description="Raw submit/status/cancel tool groups managed as durable background tasks", + ) model_config = ConfigDict(extra="allow") @model_validator(mode="before") diff --git a/backend/app/gateway/routers/mcp_tasks.py b/backend/app/gateway/routers/mcp_tasks.py new file mode 100644 index 000000000..811f562d7 --- /dev/null +++ b/backend/app/gateway/routers/mcp_tasks.py @@ -0,0 +1,95 @@ +"""Thread-scoped read API for durable MCP background tasks.""" + +from __future__ import annotations + +from typing import Any + +from fastapi import APIRouter, HTTPException, Query, Request + +from app.gateway.authz import require_permission +from app.gateway.deps import get_current_user, get_mcp_task_repo, get_mcp_task_service +from deerflow.utils.thread_id import ThreadId + +router = APIRouter(prefix="/api/threads/{thread_id}/mcp-tasks", tags=["mcp-tasks"]) + +_MAX_PUBLIC_ERROR_CHARS = 500 + + +def _short_error(value: Any) -> str | None: + if value is None: + return None + return str(value)[:_MAX_PUBLIC_ERROR_CHARS] + + +def _tracking_degraded(record: dict[str, Any], *, threshold: int) -> bool: + return int(record.get("consecutive_poll_error_count") or 0) >= threshold + + +def _list_item(record: dict[str, Any], *, threshold: int) -> dict[str, Any]: + return { + "task_id": record["id"], + "task_name": record["task_name"], + "status": record["status"], + "created_at": record["created_at"], + "updated_at": record["updated_at"], + "error": _short_error(record.get("error")), + "tracking_degraded": _tracking_degraded(record, threshold=threshold), + } + + +def _detail(record: dict[str, Any], *, threshold: int) -> dict[str, Any]: + return { + **_list_item(record, threshold=threshold), + "last_polled_at": record.get("last_polled_at"), + "last_poll_error": _short_error(record.get("last_poll_error")), + "result": record.get("result"), + "result_preview": record.get("result_preview"), + "result_truncated": bool(record.get("result_truncated")), + "result_artifact": record.get("result_artifact"), + "input_required": record.get("input_required"), + } + + +async def _current_user_id(request: Request) -> str: + user_id = await get_current_user(request) + if user_id is None: + raise HTTPException(status_code=401, detail="Authentication required") + return user_id + + +@router.get("") +@require_permission("threads", "read", owner_check=True) +async def list_mcp_tasks( + thread_id: ThreadId, + request: Request, + limit: int = Query(default=50, ge=1, le=100), +) -> list[dict[str, Any]]: + repository = get_mcp_task_repo(request) + service = get_mcp_task_service(request) + user_id = await _current_user_id(request) + records = await repository.list_by_thread( + thread_id, + user_id=user_id, + limit=limit, + ) + threshold = service.tracking_degraded_after_errors + return [_list_item(record, threshold=threshold) for record in records] + + +@router.get("/{task_id}") +@require_permission("threads", "read", owner_check=True) +async def get_mcp_task( + thread_id: ThreadId, + task_id: str, + request: Request, +) -> dict[str, Any]: + repository = get_mcp_task_repo(request) + service = get_mcp_task_service(request) + user_id = await _current_user_id(request) + record = await repository.get(task_id, user_id=user_id) + if record is None or record["thread_id"] != thread_id: + raise HTTPException(status_code=404, detail="MCP task not found") + return _detail( + record, + threshold=service.tracking_degraded_after_errors, + ) diff --git a/backend/app/mcp_tasks/service.py b/backend/app/mcp_tasks/service.py index b22609c11..f9194b011 100644 --- a/backend/app/mcp_tasks/service.py +++ b/backend/app/mcp_tasks/service.py @@ -1,18 +1,38 @@ from __future__ import annotations import asyncio +import json 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.constants import ( + MCP_TASK_POLL_AFTER_MAX_SECONDS, + MCP_TASK_REMOTE_ID_MAX_LENGTH, + MCP_TASK_RESULT_ARTIFACT_MAX_BYTES, +) +from deerflow.mcp.tasks import ( + McpTaskDriverRegistry, + McpTaskProtocolError, + TaskReference, + TaskSnapshot, + TaskStatus, + TaskSubmitRequest, +) from deerflow.persistence.mcp_tasks import DuplicateMcpRemoteTaskError logger = logging.getLogger(__name__) -_MAX_POLL_ERROR_CHARS = 4000 +_MAX_PERSISTED_ERROR_CHARS = 4_000 +_MAX_INPUT_REQUIRED_BYTES = 65_536 + + +def _bound_error(error: str | None) -> str | None: + if error is None: + return None + return error[:_MAX_PERSISTED_ERROR_CHARS] class McpTaskService: @@ -26,12 +46,22 @@ class McpTaskService: poll_interval_seconds: int, lease_seconds: int, max_concurrent_polls: int, + max_poll_backoff_seconds: int = 300, + input_required_poll_interval_seconds: int = 60, + tracking_degraded_after_errors: int = 3, + max_result_bytes: int = 65_536, + result_preview_max_chars: int = 2_000, ) -> 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._max_poll_backoff_seconds = max_poll_backoff_seconds + self._input_required_poll_interval_seconds = input_required_poll_interval_seconds + self._tracking_degraded_after_errors = tracking_degraded_after_errors + self._max_result_bytes = max_result_bytes + self._result_preview_max_chars = result_preview_max_chars self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}" self._task: asyncio.Task[None] | None = None self._stop = asyncio.Event() @@ -40,6 +70,10 @@ class McpTaskService: def drivers(self) -> McpTaskDriverRegistry: return self._drivers + @property + def tracking_degraded_after_errors(self) -> int: + return self._tracking_degraded_after_errors + async def submit( self, *, @@ -56,8 +90,6 @@ class McpTaskService: 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, @@ -68,6 +100,10 @@ class McpTaskService: driver_data=driver_data, ) try: + if len(submission.remote_task_id) > MCP_TASK_REMOTE_ID_MAX_LENGTH: + raise McpTaskProtocolError(f"MCP task remote_task_id must not exceed {MCP_TASK_REMOTE_ID_MAX_LENGTH} characters") + snapshot = self._normalize_snapshot(submission.snapshot) + next_poll_at = self._next_poll_at(snapshot, now=submitted_at) return await self._repository.create( task_id=local_task_id, user_id=request.user_id, @@ -80,6 +116,9 @@ class McpTaskService: task_name=request.task_name, status=snapshot.status.value, result=snapshot.result, + result_preview=snapshot.result_preview, + result_truncated=snapshot.result_truncated, + result_artifact=snapshot.result_artifact, error=snapshot.error, input_required=snapshot.input_required, next_poll_at=next_poll_at, @@ -134,7 +173,20 @@ class McpTaskService: return try: - snapshot = await driver.get_status(TaskReference.from_record(record)) + snapshot = self._normalize_snapshot(await driver.get_status(TaskReference.from_record(record))) + except McpTaskProtocolError as exc: + logger.error( + "MCP task status contract failed permanently (task_id=%s, driver=%s): %s", + record.get("id"), + driver_name, + exc, + ) + await self._apply_snapshot( + record, + TaskSnapshot(status=TaskStatus.FAILED, error=_bound_error(str(exc))), + polled_at=datetime.now(UTC), + ) + return except Exception as exc: # noqa: BLE001 - driver boundary; retry on the next poll polled_at = datetime.now(UTC) logger.warning( @@ -147,11 +199,23 @@ class McpTaskService: return polled_at = datetime.now(UTC) + await self._apply_snapshot(record, snapshot, polled_at=polled_at) + + async def _apply_snapshot( + self, + record: dict, + snapshot: TaskSnapshot, + *, + polled_at: datetime, + ) -> None: applied = await self._repository.apply_snapshot( record["id"], lease_owner=self._lease_owner, status=snapshot.status.value, result=snapshot.result, + result_preview=snapshot.result_preview, + result_truncated=snapshot.result_truncated, + result_artifact=snapshot.result_artifact, error=snapshot.error, input_required=snapshot.input_required, next_poll_at=self._next_poll_at(snapshot, now=polled_at), @@ -167,16 +231,72 @@ class McpTaskService: if not snapshot.is_pollable: return None interval = snapshot.poll_after_seconds or self._poll_interval_seconds + if snapshot.status == TaskStatus.INPUT_REQUIRED: + interval = max(interval, self._input_required_poll_interval_seconds) + interval = min(interval, MCP_TASK_POLL_AFTER_MAX_SECONDS) return now + timedelta(seconds=interval) async def _release_after_error(self, record: dict, *, now: datetime, error: str) -> None: + consecutive_errors = max(0, int(record.get("consecutive_poll_error_count") or 0)) + retry_seconds = min( + self._poll_interval_seconds * (2 ** min(consecutive_errors, 16)), + self._max_poll_backoff_seconds, + ) + bounded_error = _bound_error(error) + assert bounded_error is not 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], + next_poll_at=now + timedelta(seconds=retry_seconds), + error=bounded_error, ) + def _normalize_snapshot(self, snapshot: TaskSnapshot) -> TaskSnapshot: + """Bound remote payloads without ever storing truncated JSON.""" + snapshot = replace(snapshot, error=_bound_error(snapshot.error)) + if snapshot.result_artifact is not None: + encoded_artifact = self._encode_json_payload( + snapshot.result_artifact, + field_name="result_artifact", + ) + if len(encoded_artifact) > MCP_TASK_RESULT_ARTIFACT_MAX_BYTES: + raise McpTaskProtocolError(f"MCP task result_artifact payload exceeds the {MCP_TASK_RESULT_ARTIFACT_MAX_BYTES}-byte limit") + if snapshot.input_required is not None: + encoded_input = self._encode_json_payload( + snapshot.input_required, + field_name="input_required", + ) + if len(encoded_input) > _MAX_INPUT_REQUIRED_BYTES: + raise McpTaskProtocolError(f"MCP task input_required payload exceeds the {_MAX_INPUT_REQUIRED_BYTES}-byte limit") + if snapshot.result is None: + return snapshot + encoded = self._encode_json_payload(snapshot.result, field_name="result") + if len(encoded) <= self._max_result_bytes: + return snapshot + + if isinstance(snapshot.result, str): + preview_source = snapshot.result + else: + preview_source = encoded.decode("utf-8", errors="replace") + return replace( + snapshot, + result=None, + result_preview=preview_source[: self._result_preview_max_chars], + result_truncated=True, + ) + + @staticmethod + def _encode_json_payload(value, *, field_name: str) -> bytes: + try: + return json.dumps( + value, + ensure_ascii=False, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + except (TypeError, ValueError) as exc: + raise McpTaskProtocolError(f"MCP task {field_name} is not valid JSON: {exc}") from exc + async def start(self) -> None: if self._task is not None: return diff --git a/backend/docs/MCP_SERVER.md b/backend/docs/MCP_SERVER.md index fb47a6503..914dcd8b9 100644 --- a/backend/docs/MCP_SERVER.md +++ b/backend/docs/MCP_SERVER.md @@ -162,14 +162,15 @@ backward compatibility. Disable it only when every resulting tool name remains unique across the enabled servers. Stdio tools continue to use DeerFlow's persistent per-thread session pool regardless of this setting. -## Server Timeouts (Stdio MCP Servers) +## Server Timeouts -Two independent timeouts bound stdio MCP servers. `session_init_timeout` covers -server bring-up — tool discovery (subprocess spawn + `initialize` + -`tools/list`) and persistent-session initialization — and defaults to 60s so a -hung server (e.g. `npx` blocked on a package download, or a server that never -answers `initialize`) cannot block agent construction indefinitely. Set it to -`null` to disable: +Two independent settings bound stdio MCP servers and durable HTTP/SSE task +calls. `session_init_timeout` covers server bring-up — tool discovery +(subprocess spawn + `initialize` + `tools/list`) and persistent-session +initialization — plus ephemeral HTTP/SSE task-session initialization. It +defaults to 60s so a hung server (e.g. `npx` blocked on a package download, or +a server that never answers `initialize`) cannot block agent construction or +the task poller indefinitely. Set it to `null` to disable: ```json { @@ -189,10 +190,11 @@ answers `initialize`) cannot block agent construction indefinitely. Set it to } ``` -`tool_call_timeout` limits each individual tool call in seconds and applies only -to `stdio` servers; `http` and `sse` servers use transport-level timeouts, and -DeerFlow logs a warning if `tool_call_timeout` is configured for those -transports. +`tool_call_timeout` limits each individual stdio tool call in seconds. Ordinary +durable-task submit/status/cancel calls also honor it for `http` and `sse` +servers, independently of transport idle timeouts, so a live connection that +never returns the matching MCP response cannot stall the task poller. Other +`http` and `sse` tools continue to use transport-level timeouts. ## Filesystem MCP Servers @@ -207,6 +209,110 @@ paths such as `/mnt/user-data/...` to paths accepted by `@modelcontextprotocol/server-filesystem`. Use DeerFlow's built-in file tools for DeerFlow workspace files. +## Durable Background Tasks with Ordinary MCP Tools + +An MCP server can expose a fast `submit` tool plus `status` and `cancel` tools +for long-running work. DeerFlow keeps the remote task ID in SQL and polls it +outside the Agent run, so the model does not have to remember or repeatedly +send that ID. + +Enable the restart-required runtime in `config.yaml`: + +```yaml +mcp_tasks: + enabled: true + poll_interval_seconds: 5 + lease_seconds: 120 + max_concurrent_polls: 8 +``` + +Then bind exact remote tool names in `extensions_config.json`. These names are +the server's raw names, before DeerFlow adds any `_` prefix: + +```json +{ + "mcpServers": { + "report-service": { + "enabled": true, + "type": "http", + "url": "https://reports.example.com/mcp", + "session_init_timeout": 60, + "tool_call_timeout": 60, + "task_toolsets": [ + { + "name": "report-generation", + "submit_tool": "submit_report", + "status_tool": "get_report_status", + "cancel_tool": "cancel_report" + } + ] + } + } +} +``` + +The three remote tools must use MCP `structuredContent`; ordinary text blocks +are never parsed as a task protocol: + +- `submit_report()` returns + `{"task_id":"remote-123","status":"running"}` quickly. +- `get_report_status({"task_id":"remote-123"})` returns a status from + `running`, `input_required`, `completed`, `failed`, or `cancelled`. It may + also return `result`, `result_artifact` (`uri` plus `mime_type`), `error`, + `error_code`, `input_required`, and a finite positive + `poll_after_seconds`. DeerFlow caps that remote scheduling hint at 24 hours. +- `cancel_report({"task_id":"remote-123"})` is idempotent and returns the + actual terminal status: `cancelled`, `completed`, or `failed`. + +For the status tool, `isError: true` means that the status call itself failed; +DeerFlow records a bounded snippet of its first text content block and retries +with capped exponential backoff. It does not infer that the remote task failed, +because MCP tool errors do not distinguish transient from permanent conditions. +A server must report a permanent remote-task failure through a normal tool +result (`isError: false` or omitted) whose `structuredContent` contains +`status: "failed"` and an optional `error`. This distinction lets a temporary +server or network outage recover without terminalizing work that may still be +running remotely. + +Persisted task errors are capped at 4,000 characters. An `input_required` +payload must be valid JSON no larger than 64 KiB; an oversized or invalid +payload is treated as a permanent protocol failure instead of being truncated +into a different question. `result_artifact` must likewise serialize as JSON +within 64 KiB; it is a small external reference, not a second result channel. +Remote task IDs and task names are limited to 255 characters, and a task-enabled +server name is limited to 128 characters, matching the durable SQL schema on +both SQLite and PostgreSQL. + +`error_code: "task_not_found"` is a permanent failure. Network and transport +errors remain retryable with capped exponential backoff; the query API reports +`tracking_degraded` after repeated failures. Oversized JSON results are not +cut into invalid JSON: DeerFlow stores a text preview, marks +`result_truncated`, and preserves any external `result_artifact` reference. + +Only submit remains in the Agent's normal tool list. Status and cancel are +runtime-internal. Query the current thread through: + +- `GET /api/threads/{thread_id}/mcp-tasks` +- `GET /api/threads/{thread_id}/mcp-tasks/{task_id}` + +Task toolsets require `database.backend: sqlite` or `postgres`; startup fails +instead of falling back to a synchronous submit when persistence or the task +runtime is disabled. Restart recovery also requires the remote service to keep +the task alive and recognize its ID after DeerFlow reconnects. A stdio server +must therefore persist its own tasks; multi-instance deployments should +normally use an independently running HTTP/SSE service. + +Server-level OAuth works during background polling and refreshes normally. +Request-scoped secrets from a particular Agent run are not durable task +credentials and are unavailable to later background polls; use server-level +authentication for a task toolset. Restart DeerFlow after changing +`mcp_tasks`, `task_toolsets`, `mcpInterceptors`, or any connection, +authentication, transport, or timeout setting on a task-enabled server. +DeerFlow rejects task-tool reloads that no longer match the Gateway's startup +snapshot instead of discovering tools with new settings while the background +poller still calls the old endpoint. Agent-facing description/routing changes +and changes to servers without task toolsets remain hot-reloadable. + ## OAuth Support (HTTP/SSE MCP Servers) For `http` and `sse` MCP servers, DeerFlow supports OAuth token acquisition and automatic token refresh. diff --git a/backend/packages/harness/deerflow/config/AGENTS.md b/backend/packages/harness/deerflow/config/AGENTS.md index 565467536..dffb0cbe4 100644 --- a/backend/packages/harness/deerflow/config/AGENTS.md +++ b/backend/packages/harness/deerflow/config/AGENTS.md @@ -61,7 +61,7 @@ Extensions are optional only in the fallback *search* mode (priority 3-4 above): - `memory` - Memory system (enabled, storage_path, debounce_seconds, shutdown_flush_timeout_seconds, model_name, max_facts, fact_confidence_threshold, injection_enabled, max_injection_tokens, staleness_review_enabled, staleness_age_days, staleness_min_candidates, staleness_max_removals_per_cycle, staleness_protected_categories, staleness_max_lifetime_multiplier, staleness_max_extension_days) **`extensions_config.json`**: -- `mcpServers` - Map of server name → config (enabled, type, command, args, env, url, headers, oauth, description, `routing`, `tools`, `tool_call_timeout`, `session_init_timeout`). `routing.mode="prefer"` emits `` prompt guidance; if `tool_search` defers the hinted tool, `McpRoutingMiddleware` can also auto-promote matching deferred schemas before the model call. It does not hard-disable other tools. `session_init_timeout` (default `DEFAULT_MCP_SESSION_INIT_TIMEOUT` = 60s, `null` to disable) bounds server bring-up: tool discovery and persistent stdio session initialization, so a hung server cannot block agent construction indefinitely. `tool_call_timeout` bounds individual stdio tool calls. +- `mcpServers` - Map of server name → config (enabled, type, command, args, env, url, headers, oauth, description, `routing`, `tools`, `tool_call_timeout`, `session_init_timeout`). `routing.mode="prefer"` emits `` prompt guidance; if `tool_search` defers the hinted tool, `McpRoutingMiddleware` can also auto-promote matching deferred schemas before the model call. It does not hard-disable other tools. `session_init_timeout` (default `DEFAULT_MCP_SESSION_INIT_TIMEOUT` = 60s, `null` to disable) bounds server bring-up: tool discovery and persistent stdio session initialization, so a hung server cannot block agent construction indefinitely; durable HTTP/SSE task calls use it for their ephemeral session initialization too. `tool_call_timeout` bounds individual stdio calls and durable-task calls on every transport; other HTTP/SSE tools use transport-level timeouts. - `tool_search.auto_promote_top_k` - Global MCP routing auto-promote breadth. Default `3`, clamped to `1..5`; applies only when `tool_search.enabled=true` and only to deferred MCP tools with `routing.mode="prefer"` and non-empty keywords. For lead agents the deferred catalog is built from the full configured MCP set; auto-promotion never grants authority because an active skill's runtime policy still filters model-visible schemas, `tool_search` results, and execution. - `skills` - Map of skill name → state (enabled) - `middlewares` - Zero-argument `AgentMiddleware` class paths for lead and subagent runtime extension. `config.yaml -> extensions` can override these fields after validation; overrides are replace-per-field, not list concatenation. diff --git a/backend/packages/harness/deerflow/config/extensions_config.py b/backend/packages/harness/deerflow/config/extensions_config.py index b93e3a73e..cca55a0f4 100644 --- a/backend/packages/harness/deerflow/config/extensions_config.py +++ b/backend/packages/harness/deerflow/config/extensions_config.py @@ -12,7 +12,11 @@ from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from deerflow.config.runtime_paths import existing_project_file -from deerflow.constants import DEFAULT_MCP_SESSION_INIT_TIMEOUT +from deerflow.constants import ( + DEFAULT_MCP_SESSION_INIT_TIMEOUT, + MCP_TASK_NAME_MAX_LENGTH, + MCP_TASK_SERVER_NAME_MAX_LENGTH, +) logger = logging.getLogger(__name__) @@ -62,6 +66,32 @@ class McpToolOverride(BaseModel): model_config = ConfigDict(extra="allow") +class McpTaskToolsetConfig(BaseModel): + """One ordinary submit/status/cancel contract exposed by an MCP server. + + Tool names are the exact raw names advertised by that server. The + presentation prefix added by ``langchain-mcp-adapters`` is deliberately not + part of this durable binding. + """ + + name: str = Field( + min_length=1, + max_length=MCP_TASK_NAME_MAX_LENGTH, + description="Stable local name shown for tasks from this toolset", + ) + submit_tool: str = Field(min_length=1, description="Raw MCP tool name used to submit work") + status_tool: str = Field(min_length=1, description="Raw MCP tool name used to poll work") + cancel_tool: str = Field(min_length=1, description="Raw MCP tool name used to cancel work") + model_config = ConfigDict(extra="forbid") + + @field_validator("name") + @classmethod + def _validate_name_is_not_blank(cls, value: str) -> str: + if not value.strip(): + raise ValueError("MCP task toolset name must not be empty") + return value + + class McpOAuthConfig(BaseModel): """OAuth configuration for an MCP server (HTTP/SSE transports).""" @@ -105,16 +135,21 @@ class McpServerConfig(BaseModel): ) tool_call_timeout: float | None = Field( default=None, - description="Timeout in seconds for individual stdio MCP tool calls. HTTP/SSE servers use transport-level timeouts. None means no timeout.", + description=("Timeout in seconds for individual stdio MCP tool calls and durable-task calls on every transport. Other HTTP/SSE tools use transport-level timeouts. None means no call-level timeout."), ) session_init_timeout: float | None = Field( default=DEFAULT_MCP_SESSION_INIT_TIMEOUT, description=( "Timeout in seconds for MCP server bring-up: tool discovery (subprocess spawn + initialize + tools/list) " - "and persistent stdio session initialization. Defaults to DEFAULT_MCP_SESSION_INIT_TIMEOUT so a hung " - "server cannot block agent construction indefinitely. None means no timeout." + "and persistent stdio session initialization, plus ephemeral HTTP/SSE durable-task session " + "initialization. Defaults to DEFAULT_MCP_SESSION_INIT_TIMEOUT so a hung server cannot block agent " + "construction or the task poller indefinitely. None means no timeout." ), ) + task_toolsets: list[McpTaskToolsetConfig] = Field( + default_factory=list, + description="Ordinary submit/status/cancel tool groups managed by the durable MCP task runtime", + ) model_config = ConfigDict(extra="allow") @model_validator(mode="before") @@ -131,6 +166,18 @@ class McpServerConfig(BaseModel): """ return normalize_mcp_transport_alias(data) + @model_validator(mode="after") + def _validate_task_tool_bindings(self) -> "McpServerConfig": + claimed: dict[str, str] = {} + for toolset in self.task_toolsets: + for role in ("submit_tool", "status_tool", "cancel_tool"): + raw_name = getattr(toolset, role) + previous = claimed.get(raw_name) + if previous is not None: + raise ValueError(f"MCP task tool {raw_name!r} must be unique across task_toolsets and roles; it is configured as both {previous} and {toolset.name}.{role}") + claimed[raw_name] = f"{toolset.name}.{role}" + return self + def resolve_effective_mcp_routing(server_config: McpServerConfig | None, original_tool_name: str) -> dict[str, Any]: """Merge server-level routing with per-tool overrides for one MCP tool.""" @@ -168,6 +215,15 @@ class ExtensionsConfig(BaseModel): ) model_config = ConfigDict(extra="allow", populate_by_name=True) + @model_validator(mode="after") + def _validate_task_server_names_fit_storage(self) -> "ExtensionsConfig": + for server_name, server in self.mcp_servers.items(): + if not server.task_toolsets: + continue + if not server_name.strip() or len(server_name) > MCP_TASK_SERVER_NAME_MAX_LENGTH: + raise ValueError(f"MCP task server name must contain 1 to {MCP_TASK_SERVER_NAME_MAX_LENGTH} characters") + return self + def to_file_dict(self) -> dict[str, Any]: """Serialize in the public extensions_config.json shape.""" return self.model_dump(by_alias=True) diff --git a/backend/packages/harness/deerflow/config/mcp_tasks_config.py b/backend/packages/harness/deerflow/config/mcp_tasks_config.py index 1eeff41af..a20395a80 100644 --- a/backend/packages/harness/deerflow/config/mcp_tasks_config.py +++ b/backend/packages/harness/deerflow/config/mcp_tasks_config.py @@ -8,3 +8,8 @@ class McpTasksConfig(BaseModel): 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) + max_poll_backoff_seconds: int = Field(default=300, ge=1, le=3600) + input_required_poll_interval_seconds: int = Field(default=60, ge=5, le=3600) + tracking_degraded_after_errors: int = Field(default=3, ge=1, le=100) + max_result_bytes: int = Field(default=65_536, ge=1024, le=10_485_760) + result_preview_max_chars: int = Field(default=2_000, ge=64, le=100_000) diff --git a/backend/packages/harness/deerflow/constants.py b/backend/packages/harness/deerflow/constants.py index 13d5e580e..7ee937441 100644 --- a/backend/packages/harness/deerflow/constants.py +++ b/backend/packages/harness/deerflow/constants.py @@ -28,6 +28,15 @@ TOOL_RESULTS_DIRNAME = ".tool-results" # ``mcpServers..session_init_timeout``; ``None`` disables the timeout. DEFAULT_MCP_SESSION_INIT_TIMEOUT = 60.0 +# Durable MCP task storage/protocol limits. The runtime validators and ORM use +# the same constants so SQLite cannot accept values that PostgreSQL later +# rejects at its VARCHAR boundaries. +MCP_TASK_SERVER_NAME_MAX_LENGTH = 128 +MCP_TASK_REMOTE_ID_MAX_LENGTH = 255 +MCP_TASK_NAME_MAX_LENGTH = 255 +MCP_TASK_RESULT_ARTIFACT_MAX_BYTES = 65_536 +MCP_TASK_POLL_AFTER_MAX_SECONDS = 86_400 + # Persisted run-event envelope limits. Runtime definitions and the ORM both # import these from this dependency-free module so lower layers never need to # initialize deerflow.runtime just to validate storage constraints. diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index 76e2039f2..87bdced79 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -2,6 +2,8 @@ - Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management - **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` contract and normalized `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`). A driver-supplied `poll_after_seconds` must be a finite positive number, validated at the `TaskSnapshot` boundary so every driver is held to the same invariant rather than each one guarding the consumer that turns the interval into a `timedelta`. `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and 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. +- **Long-running ordinary task driver**: `extensions_config.json -> mcpServers..task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same `(server_name, user_id:thread_id)` stdio session scope; HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. This stage intentionally has no completion wake-up, natural-language cancellation, `ThreadState` projection, or frontend panel. +- **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior. - **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 - **Transports**: stdio (command-based), SSE, HTTP diff --git a/backend/packages/harness/deerflow/mcp/interceptors.py b/backend/packages/harness/deerflow/mcp/interceptors.py new file mode 100644 index 000000000..e283db6ae --- /dev/null +++ b/backend/packages/harness/deerflow/mcp/interceptors.py @@ -0,0 +1,57 @@ +"""Shared construction of MCP tool-call interceptors.""" + +from __future__ import annotations + +import logging +from typing import Any + +from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.mcp.oauth import build_oauth_tool_interceptor +from deerflow.reflection import resolve_variable + +logger = logging.getLogger(__name__) + + +def build_mcp_tool_interceptors( + extensions_config: ExtensionsConfig, + *, + oauth_builder: Any = build_oauth_tool_interceptor, + resolver: Any = resolve_variable, + target_logger: logging.Logger = logger, +) -> list[Any]: + """Build OAuth followed by configured custom MCP interceptors.""" + interceptors: list[Any] = [] + oauth_interceptor = oauth_builder(extensions_config) + if oauth_interceptor is not None: + interceptors.append(oauth_interceptor) + + raw_paths = (extensions_config.model_extra or {}).get("mcpInterceptors") + if isinstance(raw_paths, str): + raw_paths = [raw_paths] + elif not isinstance(raw_paths, list): + if raw_paths is not None: + target_logger.warning( + "mcpInterceptors must be a list of strings, got %s; skipping", + type(raw_paths).__name__, + ) + raw_paths = [] + + for interceptor_path in raw_paths: + try: + builder = resolver(interceptor_path) + interceptor = builder() + if callable(interceptor): + interceptors.append(interceptor) + target_logger.info("Loaded MCP interceptor: %s", interceptor_path) + elif interceptor is not None: + target_logger.warning( + "Builder %s returned non-callable %s; skipping", + interceptor_path, + type(interceptor).__name__, + ) + except Exception: + target_logger.warning( + f"Failed to load MCP interceptor {interceptor_path}", + exc_info=True, + ) + return interceptors diff --git a/backend/packages/harness/deerflow/mcp/oauth.py b/backend/packages/harness/deerflow/mcp/oauth.py index 19b1d18d3..632b31093 100644 --- a/backend/packages/harness/deerflow/mcp/oauth.py +++ b/backend/packages/harness/deerflow/mcp/oauth.py @@ -175,9 +175,13 @@ class OAuthTokenManager: return _OAuthToken(access_token=access_token, token_type=token_type, expires_at=expires_at) -def build_oauth_tool_interceptor(extensions_config: ExtensionsConfig) -> Any | None: +def build_oauth_tool_interceptor( + extensions_config: ExtensionsConfig, + *, + token_manager: OAuthTokenManager | None = None, +) -> Any | None: """Build a tool interceptor that injects OAuth Authorization headers.""" - token_manager = OAuthTokenManager.from_extensions_config(extensions_config) + token_manager = token_manager or OAuthTokenManager.from_extensions_config(extensions_config) if not token_manager.has_oauth_servers(): return None diff --git a/backend/packages/harness/deerflow/mcp/session_pool.py b/backend/packages/harness/deerflow/mcp/session_pool.py index 37fe9bd98..333505fdb 100644 --- a/backend/packages/harness/deerflow/mcp/session_pool.py +++ b/backend/packages/harness/deerflow/mcp/session_pool.py @@ -351,6 +351,19 @@ class MCPSessionPool: for loop, _ready, task, close_evt in inflight: await self._shutdown_entry(loop, task, close_evt, cancel=True) + async def close_session(self, server_name: str, scope_key: str) -> None: + """Close one exact server/scope session so a retry reconnects cleanly.""" + key = (server_name, scope_key) + with self._lock: + entry = self._entries.pop(key, None) + inflight = self._inflight.pop(key, None) + if entry is not None: + _session, loop, task, close_evt = entry + await self._shutdown_entry(loop, task, close_evt) + if inflight is not None: + loop, _ready, task, close_evt = inflight + await self._shutdown_entry(loop, task, close_evt, cancel=True) + async def close_server(self, server_name: str) -> None: """Close all sessions for a given server.""" with self._lock: diff --git a/backend/packages/harness/deerflow/mcp/task_tool_caller.py b/backend/packages/harness/deerflow/mcp/task_tool_caller.py new file mode 100644 index 000000000..c9581623a --- /dev/null +++ b/backend/packages/harness/deerflow/mcp/task_tool_caller.py @@ -0,0 +1,228 @@ +"""Exact-name MCP calls used by the durable ordinary-task driver.""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Mapping +from datetime import timedelta +from typing import Any + +from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.config.paths import get_paths +from deerflow.mcp.client import build_server_params +from deerflow.mcp.interceptors import build_mcp_tool_interceptors +from deerflow.mcp.oauth import OAuthTokenManager, build_oauth_tool_interceptor +from deerflow.mcp.session_pool import get_session_pool + +logger = logging.getLogger(__name__) + +_MCP_TASK_TMP_SUBDIR = ".mcp/tmp" + + +def mcp_task_session_scope_key(*, user_id: str, thread_id: str) -> str: + """Keep background calls in the same per-user/per-thread session scope.""" + return f"{user_id}:{thread_id}" + + +def _prepare_stdio_connection( + connection: dict[str, Any], + *, + user_id: str, + thread_id: str, +) -> dict[str, Any]: + paths = get_paths() + paths.ensure_thread_dirs(thread_id, user_id=user_id) + work_dir = paths.sandbox_work_dir(thread_id, user_id=user_id) + tmp_dir = work_dir / _MCP_TASK_TMP_SUBDIR + tmp_dir.mkdir(parents=True, exist_ok=True) + tmp_dir.chmod(0o700) + + prepared = dict(connection) + prepared.setdefault("cwd", str(work_dir)) + env = dict(prepared.get("env") or {}) + env.setdefault("TMPDIR", str(tmp_dir)) + env.setdefault("TMP", str(tmp_dir)) + env.setdefault("TEMP", str(tmp_dir)) + prepared["env"] = env + return prepared + + +class McpTaskToolCaller: + """Call configured raw MCP tools without exposing them back to the Agent.""" + + def __init__( + self, + extensions_config: ExtensionsConfig, + *, + oauth_token_manager: OAuthTokenManager | None = None, + ) -> None: + self._extensions_config = extensions_config + self._oauth_token_manager = oauth_token_manager or OAuthTokenManager.from_extensions_config(extensions_config) + self._interceptors = build_mcp_tool_interceptors( + extensions_config, + oauth_builder=lambda config: build_oauth_tool_interceptor( + config, + token_manager=self._oauth_token_manager, + ), + ) + + async def call_tool( + self, + *, + server_name: str, + tool_name: str, + arguments: dict[str, Any], + user_id: str, + thread_id: str, + ) -> Any: + server_config = self._extensions_config.get_enabled_mcp_servers().get(server_name) + if server_config is None: + raise LookupError(f"MCP task server {server_name!r} is missing or disabled in the startup configuration") + connection = build_server_params(server_name, server_config) + transport = connection.get("transport", "stdio") + scope_key = mcp_task_session_scope_key(user_id=user_id, thread_id=thread_id) + + if transport == "stdio": + connection = await asyncio.to_thread( + _prepare_stdio_connection, + connection, + user_id=user_id, + thread_id=thread_id, + ) + pool = get_session_pool() + session_init_timeout = server_config.session_init_timeout + if session_init_timeout is not None: + try: + session = await asyncio.wait_for( + pool.get_session(server_name, scope_key, connection), + timeout=session_init_timeout, + ) + except TimeoutError: + logger.warning( + "MCP task session initialization for server '%s' timed out after %.1fs", + server_name, + session_init_timeout, + ) + raise + else: + session = await pool.get_session(server_name, scope_key, connection) + try: + return await self._invoke( + session=session, + connection=connection, + server_name=server_name, + tool_name=tool_name, + arguments=arguments, + timeout_seconds=server_config.tool_call_timeout, + session_init_timeout_seconds=None, + persistent_session=True, + ) + except Exception: + # A dead pooled subprocess must not poison every later status + # poll. The next retry recreates this exact scoped session. + await pool.close_session(server_name, scope_key) + raise + + authorization = await self._oauth_token_manager.get_authorization_header(server_name) + if authorization: + headers = dict(connection.get("headers") or {}) + headers["Authorization"] = authorization + connection["headers"] = headers + return await self._invoke( + session=None, + connection=connection, + server_name=server_name, + tool_name=tool_name, + arguments=arguments, + timeout_seconds=server_config.tool_call_timeout, + session_init_timeout_seconds=server_config.session_init_timeout, + persistent_session=False, + ) + + async def _invoke( + self, + *, + session: Any | None, + connection: dict[str, Any], + server_name: str, + tool_name: str, + arguments: dict[str, Any], + timeout_seconds: float | None, + session_init_timeout_seconds: float | None, + persistent_session: bool, + ) -> Any: + from langchain_mcp_adapters.interceptors import MCPToolCallRequest + from langchain_mcp_adapters.sessions import create_session + + async def execute(request: MCPToolCallRequest) -> Any: + call_kwargs: dict[str, Any] = {} + if timeout_seconds: + call_kwargs["read_timeout_seconds"] = timedelta(seconds=timeout_seconds) + + if persistent_session: + assert session is not None + if request.headers: + if isinstance(request.headers, Mapping): + call_kwargs["meta"] = {"headers": dict(request.headers)} + else: + logger.warning( + "Ignoring MCP interceptor headers with unsupported type: %s", + type(request.headers).__name__, + ) + return await session.call_tool(request.name, request.args, **call_kwargs) + + effective_connection = dict(connection) + if request.headers: + headers = dict(effective_connection.get("headers") or {}) + headers.update(dict(request.headers)) + effective_connection["headers"] = headers + captured: BaseException | None = None + call_result: Any | None = None + async with create_session(effective_connection) as remote_session: + initialize = remote_session.initialize() + if session_init_timeout_seconds is not None: + await asyncio.wait_for( + initialize, + timeout=session_init_timeout_seconds, + ) + else: + await initialize + try: + call = remote_session.call_tool( + request.name, + request.args, + **call_kwargs, + ) + if timeout_seconds: + call_result = await asyncio.wait_for( + call, + timeout=timeout_seconds, + ) + else: + call_result = await call + except BaseException as exc: # preserve adapter disconnect semantics + captured = exc + if captured is not None: + raise captured + if call_result is None: + raise RuntimeError(f"MCP task tool {request.name!r} returned no result") + return call_result + + handler = execute + for interceptor in reversed(self._interceptors): + inner = handler + + async def wrapped(request: Any, _interceptor: Any = interceptor, _inner: Any = inner) -> Any: + return await _interceptor(request, _inner) + + handler = wrapped + + return await handler( + MCPToolCallRequest( + name=tool_name, + args=arguments, + server_name=server_name, + runtime=None, + ) + ) diff --git a/backend/packages/harness/deerflow/mcp/tasks/__init__.py b/backend/packages/harness/deerflow/mcp/tasks/__init__.py index e42ad10cc..b0d869f56 100644 --- a/backend/packages/harness/deerflow/mcp/tasks/__init__.py +++ b/backend/packages/harness/deerflow/mcp/tasks/__init__.py @@ -9,6 +9,11 @@ from deerflow.mcp.tasks.models import ( TaskSubmission, TaskSubmitRequest, ) +from deerflow.mcp.tasks.ordinary import ( + ORDINARY_MCP_TASK_DRIVER, + McpTaskProtocolError, + OrdinaryMcpTaskDriver, +) __all__ = [ "ATTENTION_TASK_STATUSES", @@ -21,4 +26,7 @@ __all__ = [ "TaskStatus", "TaskSubmission", "TaskSubmitRequest", + "ORDINARY_MCP_TASK_DRIVER", + "McpTaskProtocolError", + "OrdinaryMcpTaskDriver", ] diff --git a/backend/packages/harness/deerflow/mcp/tasks/models.py b/backend/packages/harness/deerflow/mcp/tasks/models.py index 5a49e6ddf..bbced4d6f 100644 --- a/backend/packages/harness/deerflow/mcp/tasks/models.py +++ b/backend/packages/harness/deerflow/mcp/tasks/models.py @@ -1,10 +1,22 @@ from __future__ import annotations -import math from dataclasses import dataclass, field from enum import StrEnum +from math import isfinite from typing import Any +from deerflow.constants import ( + MCP_TASK_NAME_MAX_LENGTH, + MCP_TASK_SERVER_NAME_MAX_LENGTH, +) + + +def _validate_storage_text(value: str, *, field_name: str, max_length: int) -> None: + if not value.strip(): + raise ValueError(f"{field_name} must not be empty") + if len(value) > max_length: + raise ValueError(f"{field_name} must not exceed {max_length} characters") + class TaskStatus(StrEnum): """Protocol-neutral lifecycle states for long-running MCP work.""" @@ -21,6 +33,7 @@ POLLABLE_TASK_STATUSES: frozenset[TaskStatus] = frozenset( { TaskStatus.SUBMITTED, TaskStatus.WORKING, + TaskStatus.INPUT_REQUIRED, } ) TERMINAL_TASK_STATUSES: frozenset[TaskStatus] = frozenset( @@ -44,6 +57,9 @@ class TaskSnapshot: status: TaskStatus result: Any | None = None + result_preview: str | None = None + result_truncated: bool = False + result_artifact: dict[str, str] | None = None error: str | None = None input_required: dict[str, Any] | None = None poll_after_seconds: float | None = None @@ -51,7 +67,7 @@ class TaskSnapshot: 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 (not math.isfinite(self.poll_after_seconds) or self.poll_after_seconds <= 0): + if self.poll_after_seconds is not None and (not isfinite(self.poll_after_seconds) or self.poll_after_seconds <= 0): # NaN and infinity survive a bare `<= 0` check but break the consumer, # which turns this interval into a `timedelta` for the next poll. raise ValueError("poll_after_seconds must be a finite positive number") @@ -104,6 +120,18 @@ class TaskSubmitRequest: driver_data: dict[str, Any] = field(default_factory=dict) local_task_id: str | None = None + def __post_init__(self) -> None: + _validate_storage_text( + self.server_name, + field_name="server_name", + max_length=MCP_TASK_SERVER_NAME_MAX_LENGTH, + ) + _validate_storage_text( + self.task_name, + field_name="task_name", + max_length=MCP_TASK_NAME_MAX_LENGTH, + ) + @dataclass(frozen=True, slots=True) class TaskSubmission: diff --git a/backend/packages/harness/deerflow/mcp/tasks/ordinary.py b/backend/packages/harness/deerflow/mcp/tasks/ordinary.py new file mode 100644 index 000000000..4d8a84fc8 --- /dev/null +++ b/backend/packages/harness/deerflow/mcp/tasks/ordinary.py @@ -0,0 +1,213 @@ +"""Driver for ordinary MCP submit/status/cancel tool contracts.""" + +from __future__ import annotations + +from typing import Any, Literal, Protocol + +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from deerflow.constants import MCP_TASK_REMOTE_ID_MAX_LENGTH +from deerflow.mcp.tasks.models import ( + TaskReference, + TaskSnapshot, + TaskStatus, + TaskSubmission, + TaskSubmitRequest, +) + +ORDINARY_MCP_TASK_DRIVER = "ordinary-tools" +_MAX_TOOL_ERROR_DETAIL_CHARS = 500 + + +class McpTaskProtocolError(RuntimeError): + """A remote task tool returned a deterministic contract violation.""" + + +class McpTaskToolCaller(Protocol): + async def call_tool( + self, + *, + server_name: str, + tool_name: str, + arguments: dict[str, Any], + user_id: str, + thread_id: str, + ) -> Any: ... + + +class _SubmitPayload(BaseModel): + task_id: str = Field(min_length=1) + status: Literal["running"] + model_config = ConfigDict(extra="ignore") + + +class _ResultArtifact(BaseModel): + uri: str = Field(min_length=1) + mime_type: str = Field(min_length=1) + model_config = ConfigDict(extra="ignore") + + +class _StatusPayload(BaseModel): + task_id: str = Field(min_length=1, max_length=MCP_TASK_REMOTE_ID_MAX_LENGTH) + status: Literal["running", "input_required", "completed", "failed", "cancelled"] + result: Any | None = None + result_artifact: _ResultArtifact | None = None + error: str | None = None + error_code: str | None = None + input_required: dict[str, Any] | None = None + poll_after_seconds: float | None = Field(default=None, gt=0, allow_inf_nan=False) + model_config = ConfigDict(extra="ignore") + + +class _CancelPayload(BaseModel): + task_id: str = Field(min_length=1, max_length=MCP_TASK_REMOTE_ID_MAX_LENGTH) + status: Literal["cancelled", "completed", "failed"] + result: Any | None = None + result_artifact: _ResultArtifact | None = None + error: str | None = None + model_config = ConfigDict(extra="ignore") + + +_REMOTE_TO_LOCAL_STATUS = { + "running": TaskStatus.WORKING, + "input_required": TaskStatus.INPUT_REQUIRED, + "completed": TaskStatus.COMPLETED, + "failed": TaskStatus.FAILED, + "cancelled": TaskStatus.CANCELLED, +} + + +def _tool_name(data: dict[str, Any], role: str) -> str: + value = data.get(role) + if not isinstance(value, str) or not value: + raise McpTaskProtocolError(f"Task driver_data is missing required {role!r}") + return value + + +def _first_error_text(call_result: Any) -> str | None: + content = getattr(call_result, "content", None) + if not isinstance(content, (list, tuple)): + return None + for item in content: + if isinstance(item, dict): + if item.get("type") != "text": + continue + text = item.get("text") + else: + if getattr(item, "type", None) != "text": + continue + text = getattr(item, "text", None) + if isinstance(text, str) and (text := text.strip()): + return text[:_MAX_TOOL_ERROR_DETAIL_CHARS] + return None + + +def _structured_content(call_result: Any, *, tool_name: str) -> Any: + if bool(getattr(call_result, "isError", False)): + message = f"MCP task tool {tool_name!r} returned an error" + if detail := _first_error_text(call_result): + message = f"{message}: {detail}" + raise RuntimeError(message) + value = getattr(call_result, "structuredContent", None) + if value is None: + raise McpTaskProtocolError(f"MCP task tool {tool_name!r} must return structuredContent; text content is not parsed") + return value + + +def _parse(model_type: type[BaseModel], value: Any, *, tool_name: str) -> BaseModel: + try: + return model_type.model_validate(value) + except ValidationError as exc: + raise McpTaskProtocolError(f"Invalid structuredContent from MCP task tool {tool_name!r}: {exc}") from exc + + +def _artifact_dict(artifact: _ResultArtifact | None) -> dict[str, str] | None: + return artifact.model_dump() if artifact is not None else None + + +def _snapshot_from_status(payload: _StatusPayload | _CancelPayload) -> TaskSnapshot: + input_required = getattr(payload, "input_required", None) + if getattr(payload, "error_code", None) == "task_not_found": + return TaskSnapshot( + status=TaskStatus.FAILED, + error=payload.error or "Remote MCP task was not found", + ) + if payload.status == "input_required" and input_required is None: + raise McpTaskProtocolError("Invalid structuredContent: input_required status requires input_required") + return TaskSnapshot( + status=_REMOTE_TO_LOCAL_STATUS[payload.status], + result=payload.result, + result_artifact=_artifact_dict(payload.result_artifact), + error=payload.error, + input_required=input_required, + poll_after_seconds=getattr(payload, "poll_after_seconds", None), + ) + + +class OrdinaryMcpTaskDriver: + """Bind a configured ordinary three-tool contract to normalized task state.""" + + def __init__(self, caller: McpTaskToolCaller) -> None: + self._caller = caller + + async def submit(self, request: TaskSubmitRequest) -> TaskSubmission: + tool_name = _tool_name(request.driver_data, "submit_tool") + result = await self._caller.call_tool( + server_name=request.server_name, + tool_name=tool_name, + arguments=request.arguments, + user_id=request.user_id, + thread_id=request.thread_id, + ) + payload = _parse( + _SubmitPayload, + _structured_content(result, tool_name=tool_name), + tool_name=tool_name, + ) + assert isinstance(payload, _SubmitPayload) + return TaskSubmission( + remote_task_id=payload.task_id, + snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED), + driver_data=dict(request.driver_data), + ) + + async def get_status(self, task: TaskReference) -> TaskSnapshot: + tool_name = _tool_name(task.driver_data, "status_tool") + result = await self._caller.call_tool( + server_name=task.server_name, + tool_name=tool_name, + arguments={"task_id": task.remote_task_id}, + user_id=task.user_id, + thread_id=task.thread_id, + ) + payload = _parse( + _StatusPayload, + _structured_content(result, tool_name=tool_name), + tool_name=tool_name, + ) + assert isinstance(payload, _StatusPayload) + self._require_matching_task_id(payload.task_id, task.remote_task_id) + return _snapshot_from_status(payload) + + async def cancel(self, task: TaskReference) -> TaskSnapshot: + tool_name = _tool_name(task.driver_data, "cancel_tool") + result = await self._caller.call_tool( + server_name=task.server_name, + tool_name=tool_name, + arguments={"task_id": task.remote_task_id}, + user_id=task.user_id, + thread_id=task.thread_id, + ) + payload = _parse( + _CancelPayload, + _structured_content(result, tool_name=tool_name), + tool_name=tool_name, + ) + assert isinstance(payload, _CancelPayload) + self._require_matching_task_id(payload.task_id, task.remote_task_id) + return _snapshot_from_status(payload) + + @staticmethod + def _require_matching_task_id(actual: str, expected: str) -> None: + if actual != expected: + raise McpTaskProtocolError(f"MCP task response task_id does not match the persisted remote task: expected {expected!r}, got {actual!r}") diff --git a/backend/packages/harness/deerflow/mcp/tasks/runtime.py b/backend/packages/harness/deerflow/mcp/tasks/runtime.py new file mode 100644 index 000000000..3403fc05b --- /dev/null +++ b/backend/packages/harness/deerflow/mcp/tasks/runtime.py @@ -0,0 +1,102 @@ +"""Process-local bridge from Agent tool wrappers to the Gateway task service.""" + +from __future__ import annotations + +from typing import Any, Protocol + +from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.mcp.tasks.models import TaskSubmitRequest + + +class McpTaskConfigurationError(RuntimeError): + """The configured long-running MCP contract cannot run safely.""" + + +class McpTaskSubmitter(Protocol): + async def submit( + self, + *, + driver_name: str, + request: TaskSubmitRequest, + now: Any | None = None, + ) -> dict: ... + + +_submitter: McpTaskSubmitter | None = None +_TaskServerConfigSnapshot = tuple[dict[str, dict[str, Any]], Any] +_task_server_config_snapshot: _TaskServerConfigSnapshot | None = None + + +def _task_server_configs(extensions_config: ExtensionsConfig) -> _TaskServerConfigSnapshot: + servers: dict[str, dict[str, Any]] = {} + for server_name, server in extensions_config.get_enabled_mcp_servers().items(): + if not server.task_toolsets: + continue + runtime_config = server.model_dump(mode="json") + for presentation_field in ("description", "routing", "tools", "tool_name_prefix"): + runtime_config.pop(presentation_field, None) + servers[server_name] = runtime_config + interceptors = (extensions_config.model_extra or {}).get("mcpInterceptors") if servers else None + return servers, interceptors + + +def set_mcp_task_config_snapshot(extensions_config: ExtensionsConfig | None) -> None: + """Freeze task-enabled server settings for one Gateway process lifetime.""" + global _task_server_config_snapshot + _task_server_config_snapshot = None if extensions_config is None else _task_server_configs(extensions_config) + + +def validate_mcp_task_config_snapshot(extensions_config: ExtensionsConfig) -> None: + """Reject hot changes that would split tool discovery from background calls.""" + if _task_server_config_snapshot is None: + return + current = _task_server_configs(extensions_config) + if current == _task_server_config_snapshot: + return + current_servers, current_interceptors = current + startup_servers, startup_interceptors = _task_server_config_snapshot + changed = sorted(server_name for server_name in current_servers.keys() | startup_servers.keys() if current_servers.get(server_name) != startup_servers.get(server_name)) + if current_interceptors != startup_interceptors: + changed.append("mcpInterceptors") + names = ", ".join(changed) or "" + raise McpTaskConfigurationError(f"MCP task-enabled server configuration changed after Gateway startup ({names}); restart DeerFlow before using durable task tools") + + +def set_mcp_task_submitter(submitter: McpTaskSubmitter | None) -> None: + """Install or clear the Gateway-owned submit boundary for this process.""" + global _submitter + _submitter = submitter + + +def get_mcp_task_submitter() -> McpTaskSubmitter: + if _submitter is None: + raise McpTaskConfigurationError("The MCP task runtime is not initialized. Run this tool through the Gateway with mcp_tasks.enabled=true and a SQL database backend.") + return _submitter + + +def configured_task_toolset_count(extensions_config: ExtensionsConfig) -> int: + return sum(len(server.task_toolsets) for server in extensions_config.get_enabled_mcp_servers().values()) + + +def validate_mcp_task_runtime_configuration( + *, + mcp_tasks_config: Any, + extensions_config: ExtensionsConfig, + repository_available: bool, +) -> None: + """Fail startup when task toolsets would silently fall back to sync calls.""" + if configured_task_toolset_count(extensions_config) == 0: + return + if not bool(getattr(mcp_tasks_config, "enabled", False)): + raise McpTaskConfigurationError("MCP task_toolsets are configured, so mcp_tasks.enabled=true is required; DeerFlow will not silently expose these tools as synchronous calls.") + if not repository_available: + raise McpTaskConfigurationError("MCP task_toolsets require durable SQL persistence. Set database.backend to 'sqlite' or 'postgres'; the memory backend cannot recover tasks after restart.") + from deerflow.mcp.client import build_server_params + + for server_name, server in extensions_config.get_enabled_mcp_servers().items(): + if not server.task_toolsets: + continue + try: + build_server_params(server_name, server) + except ValueError as exc: + raise McpTaskConfigurationError(str(exc)) from exc diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index c796da1da..6a62cda4c 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -14,12 +14,19 @@ from urllib.parse import unquote, urlparse from langchain_core.tools import BaseTool, StructuredTool from langgraph.config import get_config -from deerflow.config.extensions_config import ExtensionsConfig, resolve_effective_mcp_routing +from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig, resolve_effective_mcp_routing from deerflow.config.paths import VIRTUAL_PATH_PREFIX, Paths, get_paths from deerflow.constants import DEFAULT_MCP_SESSION_INIT_TIMEOUT from deerflow.mcp.client import build_servers_config +from deerflow.mcp.interceptors import build_mcp_tool_interceptors from deerflow.mcp.oauth import build_oauth_tool_interceptor, get_initial_oauth_headers from deerflow.mcp.session_pool import get_session_pool +from deerflow.mcp.tasks import ORDINARY_MCP_TASK_DRIVER, TaskSubmitRequest +from deerflow.mcp.tasks.runtime import ( + McpTaskConfigurationError, + get_mcp_task_submitter, + validate_mcp_task_config_snapshot, +) from deerflow.reflection import resolve_variable from deerflow.runtime.user_context import resolve_runtime_user_id from deerflow.tools.mcp_metadata import tag_mcp_routing, tag_mcp_tool @@ -612,6 +619,132 @@ def _make_session_pool_tool( ) +def _raw_mcp_tool_name( + tool: BaseTool, + *, + server_name: str, + tool_name_prefix: bool, +) -> str: + prefix = f"{server_name}_" + if tool_name_prefix and tool.name.startswith(prefix): + return tool.name[len(prefix) :] + return tool.name + + +def _make_background_submit_tool( + tool: BaseTool, + *, + server_name: str, + task_name: str, + submit_tool: str, + status_tool: str, + cancel_tool: str, +) -> BaseTool: + background_contract = f"Submitted as durable background task {task_name!r}; returns a DeerFlow task ID immediately and status polling is handled automatically." + + async def submit_in_background( + runtime: Runtime | None = None, + **arguments: Any, + ) -> dict[str, Any]: + submitter = get_mcp_task_submitter() + thread_id = _extract_thread_id(runtime) + user_id = resolve_runtime_user_id(runtime) + context = runtime.context if runtime is not None and runtime.context else {} + run_id = context.get("run_id") + tool_call_id = getattr(runtime, "tool_call_id", None) if runtime is not None else None + created = await submitter.submit( + driver_name=ORDINARY_MCP_TASK_DRIVER, + request=TaskSubmitRequest( + user_id=user_id, + thread_id=thread_id, + run_id=str(run_id) if run_id is not None else None, + tool_call_id=str(tool_call_id) if tool_call_id is not None else None, + server_name=server_name, + task_name=task_name, + arguments=arguments, + driver_data={ + "submit_tool": submit_tool, + "status_tool": status_tool, + "cancel_tool": cancel_tool, + }, + ), + ) + return { + "task_id": created["id"], + "task_name": task_name, + "status": created["status"], + "message": "Task is running in the background.", + } + + return StructuredTool( + name=tool.name, + description=(f"{tool.description}\n\n{background_contract}" if tool.description else background_contract), + args_schema=tool.args_schema, + coroutine=submit_in_background, + metadata=tool.metadata, + ) + + +def _configure_task_tools_for_server( + tools: list[BaseTool], + *, + server_name: str, + server_config: McpServerConfig, + tool_name_prefix: bool, +) -> list[BaseTool]: + """Hide driver-only tools and replace submit with a durable wrapper.""" + if not server_config.task_toolsets: + return tools + + by_raw_name = { + _raw_mcp_tool_name( + tool, + server_name=server_name, + tool_name_prefix=tool_name_prefix, + ): tool + for tool in tools + } + expected = { + raw_name + for toolset in server_config.task_toolsets + for raw_name in ( + toolset.submit_tool, + toolset.status_tool, + toolset.cancel_tool, + ) + } + missing = sorted(expected - by_raw_name.keys()) + if missing: + raise McpTaskConfigurationError(f"MCP server {server_name!r} task_toolsets reference missing raw tool(s): {', '.join(missing)}") + + hidden = {raw_name for toolset in server_config.task_toolsets for raw_name in (toolset.status_tool, toolset.cancel_tool)} + submit_by_name = {toolset.submit_tool: toolset for toolset in server_config.task_toolsets} + configured: list[BaseTool] = [] + for tool in tools: + raw_name = _raw_mcp_tool_name( + tool, + server_name=server_name, + tool_name_prefix=tool_name_prefix, + ) + if raw_name in hidden: + continue + toolset = submit_by_name.get(raw_name) + if toolset is None: + configured.append(tool) + continue + configured.append( + _make_background_submit_tool( + tool, + server_name=server_name, + task_name=toolset.name, + submit_tool=toolset.submit_tool, + status_tool=toolset.status_tool, + cancel_tool=toolset.cancel_tool, + ) + ) + return configured + + async def get_mcp_tools() -> list[BaseTool]: """Get all tools from enabled MCP servers. @@ -635,6 +768,7 @@ async def get_mcp_tools() -> list[BaseTool]: # made through the Gateway API (which runs in a separate process) are immediately # reflected when initializing MCP tools. extensions_config = ExtensionsConfig.from_file() + validate_mcp_task_config_snapshot(extensions_config) servers_config = build_servers_config(extensions_config) if not servers_config: @@ -655,34 +789,12 @@ async def get_mcp_tools() -> list[BaseTool]: existing_headers["Authorization"] = auth_header servers_config[server_name]["headers"] = existing_headers - tool_interceptors: list[Any] = [] - oauth_interceptor = build_oauth_tool_interceptor(extensions_config) - if oauth_interceptor is not None: - tool_interceptors.append(oauth_interceptor) - - # Load custom interceptors declared in extensions_config.json - # Format: "mcpInterceptors": ["pkg.module:builder_func", ...] - raw_interceptor_paths = (extensions_config.model_extra or {}).get("mcpInterceptors") - if isinstance(raw_interceptor_paths, str): - raw_interceptor_paths = [raw_interceptor_paths] - elif not isinstance(raw_interceptor_paths, list): - if raw_interceptor_paths is not None: - logger.warning(f"mcpInterceptors must be a list of strings, got {type(raw_interceptor_paths).__name__}; skipping") - raw_interceptor_paths = [] - for interceptor_path in raw_interceptor_paths: - try: - builder = resolve_variable(interceptor_path) - interceptor = builder() - if callable(interceptor): - tool_interceptors.append(interceptor) - logger.info(f"Loaded MCP interceptor: {interceptor_path}") - elif interceptor is not None: - logger.warning(f"Builder {interceptor_path} returned non-callable {type(interceptor).__name__}; skipping") - except Exception as e: - logger.warning( - f"Failed to load MCP interceptor {interceptor_path}: {e}", - exc_info=True, - ) + tool_interceptors = build_mcp_tool_interceptors( + extensions_config, + oauth_builder=build_oauth_tool_interceptor, + resolver=resolve_variable, + target_logger=logger, + ) client = MultiServerMCPClient( servers_config, @@ -766,6 +878,7 @@ async def get_mcp_tools() -> list[BaseTool]: transport = servers_config[source_name].get("transport", "stdio") server_cfg = extensions_config.mcp_servers.get(source_name) tool_name_prefix = server_cfg.tool_name_prefix if server_cfg is not None else True + current_server_tools: list[BaseTool] = [] for tool in server_tools: if not _VALID_MCP_TOOL_NAME.fullmatch(tool.name or ""): logger.warning( @@ -784,7 +897,7 @@ async def get_mcp_tools() -> list[BaseTool]: if transport == "stdio": _timeout = server_cfg.tool_call_timeout if server_cfg else None _init_timeout = _resolve_session_init_timeout(server_cfg) - wrapped_tools.append( + current_server_tools.append( _make_session_pool_tool( tool, source_name, @@ -802,7 +915,16 @@ async def get_mcp_tools() -> list[BaseTool]: source_name, transport, ) - wrapped_tools.append(tool) + current_server_tools.append(tool) + + if server_cfg is not None: + current_server_tools = _configure_task_tools_for_server( + current_server_tools, + server_name=source_name, + server_config=server_cfg, + tool_name_prefix=tool_name_prefix, + ) + wrapped_tools.extend(current_server_tools) # Patch tools to support sync invocation, as deerflow client streams synchronously for tool in wrapped_tools: @@ -811,6 +933,8 @@ async def get_mcp_tools() -> list[BaseTool]: return wrapped_tools + except McpTaskConfigurationError: + raise except Exception as e: logger.error(f"Failed to load MCP tools: {e}", exc_info=True) return [] diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py index c60514cef..c14a6042f 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py @@ -3,9 +3,14 @@ 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 import JSON, Boolean, DateTime, Index, Integer, String, Text, UniqueConstraint, false from sqlalchemy.orm import Mapped, mapped_column +from deerflow.constants import ( + MCP_TASK_NAME_MAX_LENGTH, + MCP_TASK_REMOTE_ID_MAX_LENGTH, + MCP_TASK_SERVER_NAME_MAX_LENGTH, +) from deerflow.persistence.base import Base @@ -17,12 +22,15 @@ class McpTaskRow(Base): 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)) + server_name: Mapped[str] = mapped_column(String(MCP_TASK_SERVER_NAME_MAX_LENGTH)) 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)) + remote_task_id: Mapped[str] = mapped_column(String(MCP_TASK_REMOTE_ID_MAX_LENGTH)) + task_name: Mapped[str] = mapped_column(String(MCP_TASK_NAME_MAX_LENGTH)) status: Mapped[str] = mapped_column(String(32), index=True) result: Mapped[Any | None] = mapped_column(JSON, nullable=True) + result_preview: Mapped[str | None] = mapped_column(Text, nullable=True) + result_truncated: Mapped[bool] = mapped_column(Boolean, default=False, server_default=false()) + result_artifact: Mapped[dict[str, str] | 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) diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py index 82535c919..bc23b4138 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py @@ -66,6 +66,9 @@ class McpTaskRepository: task_name: str, status: str, result: Any | None, + result_preview: str | None, + result_truncated: bool, + result_artifact: dict[str, str] | None, error: str | None, input_required: dict[str, Any] | None, next_poll_at: datetime | None, @@ -85,6 +88,9 @@ class McpTaskRepository: task_name=task_name, status=status, result=result, + result_preview=result_preview, + result_truncated=result_truncated, + result_artifact=result_artifact, error=error, input_required=input_required, driver_data=dict(driver_data or {}), @@ -174,6 +180,9 @@ class McpTaskRepository: lease_owner: str, status: str, result: Any | None, + result_preview: str | None, + result_truncated: bool, + result_artifact: dict[str, str] | None, error: str | None, input_required: dict[str, Any] | None, next_poll_at: datetime | None, @@ -182,6 +191,9 @@ class McpTaskRepository: values: dict[str, Any] = { "status": status, "result": result, + "result_preview": result_preview, + "result_truncated": result_truncated, + "result_artifact": result_artifact, "error": error, "input_required": input_required, "next_poll_at": next_poll_at, diff --git a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md index 9fe8b15a3..45a19876a 100644 --- a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md +++ b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md @@ -34,5 +34,7 @@ This invokes `alembic revision --autogenerate` against the live ORM models. Revi - `migrations/versions/0007_scheduled_run_active_index.py` — the `uq_scheduled_task_run_active` partial unique index (at most one queued/running `scheduled_task_runs` row per `task_id`), with a `_dedupe_active_scheduled_runs_per_task()` pre-step (keeps the newest active row per task, supersedes the rest to `interrupted` with an explanatory `error` + `finished_at`) mirroring 0004; chains after `0006_agents` - `migrations/versions/0008_thread_operation_kind.py` — adds `runs.operation_kind` for durable non-run thread reservations; chains after `0007_scheduled_run_active_index` - `migrations/versions/0010_run_cancel_request.py` — adds the nullable `runs.cancel_action` / `cancel_requested_at` handoff used by non-owning workers; chains after `0009_webhook_dedupe` +- `migrations/versions/0011_mcp_tasks.py` — creates the durable long-running MCP task table and its user/server/remote uniqueness constraint +- `migrations/versions/0012_mcp_task_results.py` — adds bounded result preview/truncation/artifact fields for ordinary task drivers - `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch decision + locking - Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor), `tests/test_migration_0004_run_ownership_dedupe.py` + `tests/test_migration_0007_scheduled_run_active_dedupe.py` (dedupe-before-unique-index pre-steps) diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0012_mcp_task_results.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0012_mcp_task_results.py new file mode 100644 index 000000000..c4e756887 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0012_mcp_task_results.py @@ -0,0 +1,47 @@ +"""bounded MCP task result fields. + +Revision ID: 0012_mcp_task_results +Revises: 0011_mcp_tasks +Create Date: 2026-08-05 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa + +revision: str = "0012_mcp_task_results" +down_revision: str | Sequence[str] | None = "0011_mcp_tasks" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + from deerflow.persistence.migrations._helpers import safe_add_column + + safe_add_column( + "mcp_tasks", + sa.Column("result_preview", sa.Text(), nullable=True), + ) + safe_add_column( + "mcp_tasks", + sa.Column( + "result_truncated", + sa.Boolean(), + nullable=False, + server_default=sa.false(), + ), + ) + safe_add_column( + "mcp_tasks", + sa.Column("result_artifact", sa.JSON(), nullable=True), + ) + + +def downgrade() -> None: + from deerflow.persistence.migrations._helpers import safe_drop_column + + safe_drop_column("mcp_tasks", "result_artifact") + safe_drop_column("mcp_tasks", "result_truncated") + safe_drop_column("mcp_tasks", "result_preview") diff --git a/backend/tests/test_gateway_lifespan_shutdown.py b/backend/tests/test_gateway_lifespan_shutdown.py index 9800c614b..2949da588 100644 --- a/backend/tests/test_gateway_lifespan_shutdown.py +++ b/backend/tests/test_gateway_lifespan_shutdown.py @@ -16,6 +16,7 @@ from contextlib import asynccontextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch +import pytest from fastapi import FastAPI @@ -119,6 +120,66 @@ def test_lifespan_sweeps_upload_staging_files_on_startup(): stop_channel_service.assert_awaited_once() +async def _run_lifespan_with_mcp_task_config_snapshot() -> None: + from app.gateway.app import lifespan + from deerflow.config.extensions_config import ExtensionsConfig + from deerflow.mcp.tasks.runtime import McpTaskConfigurationError, validate_mcp_task_config_snapshot + + app = FastAPI() + startup_config = SimpleNamespace( + log_level="INFO", + memory=SimpleNamespace( + token_counting="char", + enabled=False, + shutdown_flush_timeout_seconds=30.0, + ), + ) + startup_extensions = ExtensionsConfig() + changed_extensions = ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "command": "reports-mcp", + "task_toolsets": [ + { + "name": "reports", + "submit_tool": "submit_report", + "status_tool": "status_report", + "cancel_tool": "cancel_report", + } + ], + } + } + } + ) + fake_service = MagicMock() + fake_service.get_status.return_value = {} + + async def fake_start(_startup_config, **_kwargs): + return fake_service + + with ( + patch("app.gateway.app.get_app_config", return_value=startup_config), + patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)), + patch("app.gateway.app.langgraph_runtime", _noop_langgraph_runtime), + patch("app.gateway.app.auth.close_oidc_service", AsyncMock()), + patch("app.channels.service.start_channel_service", side_effect=fake_start), + patch("app.channels.service.stop_channel_service", AsyncMock()), + patch("deerflow.skills.projection.ensure_public_skill_projection"), + patch("deerflow.agents.memory.get_memory_manager", return_value=MagicMock()), + patch("deerflow.config.extensions_config.ExtensionsConfig.from_file", return_value=startup_extensions), + ): + async with lifespan(app): + with pytest.raises(McpTaskConfigurationError, match="reports.*restart"): + validate_mcp_task_config_snapshot(changed_extensions) + + validate_mcp_task_config_snapshot(changed_extensions) + + +def test_lifespan_sets_and_clears_mcp_task_config_snapshot() -> None: + asyncio.run(_run_lifespan_with_mcp_task_config_snapshot()) + + async def _run_lifespan_with_memory_flush( *, enabled: bool, diff --git a/backend/tests/test_mcp_session_pool.py b/backend/tests/test_mcp_session_pool.py index 250a99ef3..892dfe923 100644 --- a/backend/tests/test_mcp_session_pool.py +++ b/backend/tests/test_mcp_session_pool.py @@ -161,6 +161,41 @@ async def test_close_scope(): assert ("s", "t2") in pool._entries +@pytest.mark.asyncio +async def test_close_session_only_evicts_the_exact_server_scope_pair(): + pool = MCPSessionPool() + + class CmFactory: + def __init__(self): + self.closed = False + + async def __aenter__(self): + return AsyncMock() + + async def __aexit__(self, *args): + self.closed = True + return False + + cms: list[CmFactory] = [] + + def make_cm(*_args, **_kwargs): + cm = CmFactory() + cms.append(cm) + return cm + + with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm): + await pool.get_session("s1", "t1", {"transport": "stdio", "command": "x", "args": []}) + await pool.get_session("s2", "t1", {"transport": "stdio", "command": "x", "args": []}) + await pool.get_session("s1", "t2", {"transport": "stdio", "command": "x", "args": []}) + + await pool.close_session("s1", "t1") + + assert cms[0].closed is True + assert cms[1].closed is False + assert cms[2].closed is False + assert set(pool._entries) == {("s2", "t1"), ("s1", "t2")} + + @pytest.mark.asyncio async def test_close_all(): """close_all shuts down every session.""" diff --git a/backend/tests/test_mcp_task_config.py b/backend/tests/test_mcp_task_config.py index 2900ddca4..bf48da372 100644 --- a/backend/tests/test_mcp_task_config.py +++ b/backend/tests/test_mcp_task_config.py @@ -12,11 +12,18 @@ def test_mcp_task_runtime_is_disabled_by_default_and_bounded(): assert config.poll_interval_seconds == 5 assert config.lease_seconds == 120 assert config.max_concurrent_polls == 8 + assert config.max_poll_backoff_seconds == 300 + assert config.input_required_poll_interval_seconds == 60 + assert config.tracking_degraded_after_errors == 3 + assert config.max_result_bytes == 65_536 + assert config.result_preview_max_chars == 2_000 with pytest.raises(ValidationError): McpTasksConfig(poll_interval_seconds=0) with pytest.raises(ValidationError): McpTasksConfig(max_concurrent_polls=0) + with pytest.raises(ValidationError): + McpTasksConfig(max_result_bytes=10) def test_mcp_task_runtime_is_registered_as_startup_only(): diff --git a/backend/tests/test_mcp_task_models.py b/backend/tests/test_mcp_task_models.py index 6538d369e..6486de0cb 100644 --- a/backend/tests/test_mcp_task_models.py +++ b/backend/tests/test_mcp_task_models.py @@ -2,7 +2,7 @@ from datetime import timedelta import pytest -from deerflow.mcp.tasks import McpTaskDriverRegistry, TaskSnapshot, TaskStatus, TaskSubmission +from deerflow.mcp.tasks import McpTaskDriverRegistry, TaskSnapshot, TaskStatus, TaskSubmission, TaskSubmitRequest def test_task_snapshot_normalizes_string_statuses(): @@ -39,6 +39,22 @@ def test_submission_rejects_empty_remote_id(): TaskSubmission(remote_task_id=" ", snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED)) +def test_task_storage_identifiers_reject_values_longer_than_the_database_columns(): + for field_name, request_kwargs in ( + ("server_name", {"server_name": "s" * 129, "task_name": "report"}), + ("task_name", {"server_name": "reports", "task_name": "t" * 256}), + ): + with pytest.raises(ValueError, match=field_name): + TaskSubmitRequest( + user_id="user-1", + thread_id="thread-1", + run_id=None, + tool_call_id=None, + arguments={}, + **request_kwargs, + ) + + def test_driver_registry_rejects_duplicate_names(): registry = McpTaskDriverRegistry() driver = object() diff --git a/backend/tests/test_mcp_task_ordinary_driver.py b/backend/tests/test_mcp_task_ordinary_driver.py new file mode 100644 index 000000000..53be3fd20 --- /dev/null +++ b/backend/tests/test_mcp_task_ordinary_driver.py @@ -0,0 +1,234 @@ +from types import SimpleNamespace + +import pytest + +from deerflow.mcp.tasks import TaskReference, TaskStatus, TaskSubmitRequest +from deerflow.mcp.tasks.ordinary import McpTaskProtocolError, OrdinaryMcpTaskDriver + + +class FakeCaller: + def __init__(self, *results): + self.results = list(results) + self.calls = [] + + async def call_tool(self, **kwargs): + self.calls.append(kwargs) + return self.results.pop(0) + + +def _result(structured_content, *, text="ignored", is_error=False): + return SimpleNamespace( + structuredContent=structured_content, + content=[SimpleNamespace(type="text", text=text)], + isError=is_error, + ) + + +def _request() -> TaskSubmitRequest: + return TaskSubmitRequest( + user_id="user-1", + thread_id="thread-1", + run_id="run-1", + tool_call_id="call-1", + server_name="reports", + task_name="report-generation", + arguments={"topic": "MCP"}, + driver_data={ + "submit_tool": "submit_report", + "status_tool": "get_report_status", + "cancel_tool": "cancel_report", + }, + local_task_id="local-1", + ) + + +def _reference() -> TaskReference: + return TaskReference( + local_task_id="local-1", + user_id="user-1", + thread_id="thread-1", + server_name="reports", + remote_task_id="remote-1", + driver_data={ + "submit_tool": "submit_report", + "status_tool": "get_report_status", + "cancel_tool": "cancel_report", + }, + ) + + +@pytest.mark.asyncio +async def test_submit_uses_structured_content_and_keeps_remote_id_out_of_driver_data() -> None: + caller = FakeCaller(_result({"task_id": "remote-1", "status": "running"}, text='{"task_id":"wrong"}')) + driver = OrdinaryMcpTaskDriver(caller) + + submission = await driver.submit(_request()) + + assert submission.remote_task_id == "remote-1" + assert submission.snapshot.status == TaskStatus.SUBMITTED + assert "remote_task_id" not in submission.driver_data + assert caller.calls == [ + { + "server_name": "reports", + "tool_name": "submit_report", + "arguments": {"topic": "MCP"}, + "user_id": "user-1", + "thread_id": "thread-1", + } + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "structured_content,match", + [ + (None, "structuredContent"), + ({"status": "running"}, "task_id"), + ({"task_id": "remote-1", "status": "queued"}, "status"), + ], +) +async def test_submit_rejects_missing_or_invalid_structured_content(structured_content, match: str) -> None: + driver = OrdinaryMcpTaskDriver(FakeCaller(_result(structured_content, text='{"task_id":"remote-from-text"}'))) + + with pytest.raises(McpTaskProtocolError, match=match): + await driver.submit(_request()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method_name", "argument", "tool_name"), + [ + ("submit", _request(), "submit_report"), + ("get_status", _reference(), "get_report_status"), + ("cancel", _reference(), "cancel_report"), + ], +) +async def test_tool_errors_preserve_only_a_bounded_first_text_detail( + method_name: str, + argument: TaskSubmitRequest | TaskReference, + tool_name: str, +) -> None: + detail = "remote service unavailable: " + "x" * 600 + result = _result(None, text=detail, is_error=True) + result.content.append(SimpleNamespace(type="text", text="second block must not be included")) + driver = OrdinaryMcpTaskDriver(FakeCaller(result)) + + with pytest.raises(RuntimeError) as exc_info: + await getattr(driver, method_name)(argument) + + message = str(exc_info.value) + assert message == f"MCP task tool {tool_name!r} returned an error: {detail[:500]}" + assert not isinstance(exc_info.value, McpTaskProtocolError) + + +@pytest.mark.asyncio +async def test_status_maps_all_protocol_states_and_preserves_artifact() -> None: + caller = FakeCaller( + _result({"task_id": "remote-1", "status": "running", "poll_after_seconds": 7}), + _result( + { + "task_id": "remote-1", + "status": "completed", + "result": {"report": "ready"}, + "result_artifact": {"uri": "s3://reports/1.json", "mime_type": "application/json"}, + } + ), + ) + driver = OrdinaryMcpTaskDriver(caller) + + running = await driver.get_status(_reference()) + completed = await driver.get_status(_reference()) + + assert running.status == TaskStatus.WORKING + assert running.poll_after_seconds == 7 + assert completed.status == TaskStatus.COMPLETED + assert completed.result == {"report": "ready"} + assert completed.result_artifact == { + "uri": "s3://reports/1.json", + "mime_type": "application/json", + } + assert caller.calls[0]["arguments"] == {"task_id": "remote-1"} + + +@pytest.mark.asyncio +async def test_status_keeps_input_required_pollable() -> None: + driver = OrdinaryMcpTaskDriver( + FakeCaller( + _result( + { + "task_id": "remote-1", + "status": "input_required", + "input_required": {"prompt": "Approve deployment?"}, + } + ) + ) + ) + + snapshot = await driver.get_status(_reference()) + + assert snapshot.status == TaskStatus.INPUT_REQUIRED + assert snapshot.input_required == {"prompt": "Approve deployment?"} + assert snapshot.is_pollable is True + + +@pytest.mark.asyncio +async def test_status_rejects_non_finite_poll_after_seconds() -> None: + driver = OrdinaryMcpTaskDriver( + FakeCaller( + _result( + { + "task_id": "remote-1", + "status": "running", + "poll_after_seconds": float("inf"), + } + ) + ) + ) + + with pytest.raises(McpTaskProtocolError, match="poll_after_seconds"): + await driver.get_status(_reference()) + + +@pytest.mark.asyncio +async def test_status_turns_task_not_found_into_permanent_failure() -> None: + driver = OrdinaryMcpTaskDriver( + FakeCaller( + _result( + { + "task_id": "remote-1", + "status": "running", + "error_code": "task_not_found", + "error": "expired", + } + ) + ) + ) + + snapshot = await driver.get_status(_reference()) + + assert snapshot.status == TaskStatus.FAILED + assert snapshot.error == "expired" + + +@pytest.mark.asyncio +async def test_status_rejects_mismatched_remote_id_and_unknown_status() -> None: + caller = FakeCaller( + _result({"task_id": "another-task", "status": "running"}), + _result({"task_id": "remote-1", "status": "paused"}), + ) + driver = OrdinaryMcpTaskDriver(caller) + + with pytest.raises(McpTaskProtocolError, match="task_id does not match"): + await driver.get_status(_reference()) + with pytest.raises(McpTaskProtocolError, match="status"): + await driver.get_status(_reference()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["cancelled", "completed", "failed"]) +async def test_cancel_is_idempotent_and_preserves_actual_terminal_status(status: str) -> None: + driver = OrdinaryMcpTaskDriver(FakeCaller(_result({"task_id": "remote-1", "status": status}))) + + snapshot = await driver.cancel(_reference()) + + assert snapshot.status.value == status diff --git a/backend/tests/test_mcp_task_ordinary_e2e.py b/backend/tests/test_mcp_task_ordinary_e2e.py new file mode 100644 index 000000000..e2ef5f375 --- /dev/null +++ b/backend/tests/test_mcp_task_ordinary_e2e.py @@ -0,0 +1,189 @@ +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace + +import pytest +import pytest_asyncio + +from app.mcp_tasks import McpTaskService +from deerflow.config.database_config import DatabaseConfig +from deerflow.mcp.tasks import ( + ORDINARY_MCP_TASK_DRIVER, + McpTaskDriverRegistry, + OrdinaryMcpTaskDriver, + TaskSubmitRequest, +) +from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config +from deerflow.persistence.mcp_tasks import McpTaskRepository + + +@pytest_asyncio.fixture(autouse=True) +async def _close_persistence_engine(): + yield + await close_engine() + + +class FakeMcpServer: + def __init__(self): + self.status_results = [] + + async def call_tool(self, *, tool_name, arguments, **_scope): + if tool_name == "submit_report": + return SimpleNamespace( + structuredContent={"task_id": arguments["remote_id"], "status": "running"}, + content=[], + isError=False, + ) + if tool_name == "status_report": + status_result = self.status_results.pop(0) + if not isinstance(status_result, dict): + return status_result + return SimpleNamespace( + structuredContent=status_result, + content=[], + isError=False, + ) + if tool_name == "cancel_report": + return SimpleNamespace( + structuredContent={"task_id": arguments["task_id"], "status": "cancelled"}, + content=[], + isError=False, + ) + raise AssertionError(tool_name) + + +def _service(repo, fake_server) -> McpTaskService: + registry = McpTaskDriverRegistry() + registry.register(ORDINARY_MCP_TASK_DRIVER, OrdinaryMcpTaskDriver(fake_server)) + return McpTaskService( + repository=repo, + drivers=registry, + poll_interval_seconds=1, + lease_seconds=120, + max_concurrent_polls=8, + ) + + +def _request(remote_id: str) -> TaskSubmitRequest: + return TaskSubmitRequest( + user_id="user-1", + thread_id="thread-1", + run_id="run-1", + tool_call_id="call-1", + server_name="reports", + task_name="report-generation", + arguments={"remote_id": remote_id}, + driver_data={ + "submit_tool": "submit_report", + "status_tool": "status_report", + "cancel_tool": "cancel_report", + }, + ) + + +@pytest.mark.asyncio +async def test_submit_poll_restart_recovery_complete_and_fail(tmp_path) -> None: + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + session_factory = get_session_factory() + assert session_factory is not None + repo = McpTaskRepository(session_factory) + fake_server = FakeMcpServer() + submitted_at = datetime.now(UTC) + fake_server.status_results.extend( + [ + { + "task_id": "remote-complete", + "status": "running", + "poll_after_seconds": 1, + }, + { + "task_id": "remote-complete", + "status": "completed", + "result": {"report": "ready"}, + }, + ] + ) + + first_process = _service(repo, fake_server) + created = await first_process.submit( + driver_name=ORDINARY_MCP_TASK_DRIVER, + request=_request("remote-complete"), + now=submitted_at, + ) + await first_process.run_once(now=submitted_at + timedelta(seconds=2)) + + # Recreate the service/registry to model a Gateway restart. The only handle + # available to the new process is the row persisted before submit returned. + restarted_process = _service(repo, fake_server) + await restarted_process.run_once(now=datetime.now(UTC) + timedelta(seconds=2)) + + completed = await repo.get(created["id"], user_id="user-1") + assert completed is not None + assert completed["status"] == "completed" + assert completed["result"] == {"report": "ready"} + + fake_server.status_results.append( + { + "task_id": "remote-fail", + "status": "failed", + "error": "report generation failed", + } + ) + failed_created = await restarted_process.submit( + driver_name=ORDINARY_MCP_TASK_DRIVER, + request=_request("remote-fail"), + now=datetime.now(UTC) - timedelta(seconds=2), + ) + await restarted_process.run_once(now=datetime.now(UTC)) + + failed = await repo.get(failed_created["id"], user_id="user-1") + assert failed is not None + assert failed["status"] == "failed" + assert failed["error"] == "report generation failed" + + +@pytest.mark.asyncio +async def test_status_tool_error_retries_with_detail_before_structured_failure_terminalizes(tmp_path) -> None: + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + session_factory = get_session_factory() + assert session_factory is not None + repo = McpTaskRepository(session_factory) + fake_server = FakeMcpServer() + service = _service(repo, fake_server) + submitted_at = datetime.now(UTC) + created = await service.submit( + driver_name=ORDINARY_MCP_TASK_DRIVER, + request=_request("remote-fail"), + now=submitted_at, + ) + fake_server.status_results.extend( + [ + SimpleNamespace( + structuredContent=None, + content=[SimpleNamespace(type="text", text="upstream temporarily unavailable")], + isError=True, + ), + { + "task_id": "remote-fail", + "status": "failed", + "error": "report generation failed", + }, + ] + ) + + await service.run_once(now=submitted_at + timedelta(seconds=2)) + + retrying = await repo.get(created["id"], user_id="user-1") + assert retrying is not None + assert retrying["status"] == "submitted" + assert retrying["consecutive_poll_error_count"] == 1 + assert retrying["last_poll_error"] == ("MCP task tool 'status_report' returned an error: upstream temporarily unavailable") + assert retrying["next_poll_at"] is not None + + await service.run_once(now=datetime.now(UTC) + timedelta(seconds=10)) + + failed = await repo.get(created["id"], user_id="user-1") + assert failed is not None + assert failed["status"] == "failed" + assert failed["error"] == "report generation failed" + assert failed["consecutive_poll_error_count"] == 0 + assert failed["next_poll_at"] is None diff --git a/backend/tests/test_mcp_task_repository.py b/backend/tests/test_mcp_task_repository.py index 1bfae0113..dd2ccd0db 100644 --- a/backend/tests/test_mcp_task_repository.py +++ b/backend/tests/test_mcp_task_repository.py @@ -42,6 +42,9 @@ async def _create_working_task( task_name="Generate report", status="working", result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, error=None, input_required=None, next_poll_at=now - timedelta(seconds=1), @@ -142,6 +145,9 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task lease_owner="worker-old", status="failed", result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, error="stale result", input_required=None, next_poll_at=None, @@ -154,6 +160,9 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task lease_owner="worker-new", status="completed", result={"report": "ready"}, + result_preview=None, + result_truncated=False, + result_artifact={"uri": "s3://reports/2.json", "mime_type": "application/json"}, error=None, input_required=None, next_poll_at=None, @@ -165,6 +174,10 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task assert stored is not None assert stored["status"] == "completed" assert stored["result"] == {"report": "ready"} + assert stored["result_artifact"] == { + "uri": "s3://reports/2.json", + "mime_type": "application/json", + } assert stored["notification_status"] == "pending" assert stored["lease_owner"] is None @@ -196,6 +209,9 @@ async def test_apply_snapshot_rejects_result_after_same_workers_lease_expires(tm lease_owner="worker-1", status="completed", result={"report": "stale"}, + result_preview=None, + result_truncated=False, + result_artifact=None, error=None, input_required=None, next_poll_at=None, @@ -210,7 +226,7 @@ async def test_apply_snapshot_rejects_result_after_same_workers_lease_expires(tm @pytest.mark.asyncio -async def test_input_required_is_persisted_and_paused_until_future_resume(tmp_path): +async def test_input_required_is_persisted_and_remains_scheduled_for_slow_polling(tmp_path): repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-3", now=now) @@ -226,9 +242,12 @@ async def test_input_required_is_persisted_and_paused_until_future_resume(tmp_pa lease_owner="worker-1", status="input_required", result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, error=None, input_required={"prompt": "Approve deployment?"}, - next_poll_at=None, + next_poll_at=now + timedelta(seconds=60), polled_at=now, ) assert applied is True @@ -237,7 +256,7 @@ async def test_input_required_is_persisted_and_paused_until_future_resume(tmp_pa assert stored is not None assert stored["input_required"] == {"prompt": "Approve deployment?"} assert stored["notification_status"] == "pending" - assert stored["next_poll_at"] is None + assert datetime.fromisoformat(stored["next_poll_at"]) == now + timedelta(seconds=60) @pytest.mark.asyncio @@ -293,6 +312,9 @@ async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp lease_owner="worker-1", status="working", result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, error=None, input_required=None, next_poll_at=now + timedelta(seconds=5), diff --git a/backend/tests/test_mcp_task_runtime_config.py b/backend/tests/test_mcp_task_runtime_config.py new file mode 100644 index 000000000..2989efedb --- /dev/null +++ b/backend/tests/test_mcp_task_runtime_config.py @@ -0,0 +1,139 @@ +from types import SimpleNamespace + +import pytest + +from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.mcp.tasks.runtime import ( + McpTaskConfigurationError, + set_mcp_task_config_snapshot, + validate_mcp_task_config_snapshot, + validate_mcp_task_runtime_configuration, +) + + +def _extensions() -> ExtensionsConfig: + return ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "task_toolsets": [ + { + "name": "reports", + "submit_tool": "submit_report", + "status_tool": "status_report", + "cancel_tool": "cancel_report", + } + ] + } + } + } + ) + + +def test_configured_task_toolsets_require_enabled_runtime() -> None: + with pytest.raises(McpTaskConfigurationError, match="mcp_tasks.enabled=true"): + validate_mcp_task_runtime_configuration( + mcp_tasks_config=SimpleNamespace(enabled=False), + extensions_config=_extensions(), + repository_available=True, + ) + + +def test_configured_task_toolsets_require_sql_persistence() -> None: + with pytest.raises(McpTaskConfigurationError, match="database.backend"): + validate_mcp_task_runtime_configuration( + mcp_tasks_config=SimpleNamespace(enabled=True), + extensions_config=_extensions(), + repository_available=False, + ) + + +def test_no_task_toolsets_leave_existing_mcp_runtime_unchanged() -> None: + validate_mcp_task_runtime_configuration( + mcp_tasks_config=SimpleNamespace(enabled=False), + extensions_config=ExtensionsConfig(), + repository_available=False, + ) + + +def test_task_toolset_server_transport_is_validated_at_startup() -> None: + extensions = _extensions() + extensions.mcp_servers["reports"].command = None + + with pytest.raises(McpTaskConfigurationError, match="requires 'command'"): + validate_mcp_task_runtime_configuration( + mcp_tasks_config=SimpleNamespace(enabled=True), + extensions_config=extensions, + repository_available=True, + ) + + +def test_task_enabled_server_changes_require_gateway_restart() -> None: + startup = _extensions() + current = _extensions() + current.mcp_servers["reports"].env["TOKEN"] = "rotated" + set_mcp_task_config_snapshot(startup) + try: + with pytest.raises(McpTaskConfigurationError, match="reports.*restart"): + validate_mcp_task_config_snapshot(current) + finally: + set_mcp_task_config_snapshot(None) + + +def test_unrelated_extension_changes_do_not_invalidate_task_runtime_snapshot() -> None: + startup = _extensions() + current = ExtensionsConfig.model_validate( + { + **startup.model_dump(by_alias=True), + "skills": {"writer": {"enabled": False}}, + "mcpServers": { + **startup.model_dump(by_alias=True)["mcpServers"], + "search": {"command": "search-mcp"}, + }, + } + ) + current.mcp_servers["reports"].description = "Updated Agent-facing description" + set_mcp_task_config_snapshot(startup) + try: + validate_mcp_task_config_snapshot(current) + finally: + set_mcp_task_config_snapshot(None) + + +def test_disabled_task_server_changes_do_not_invalidate_task_runtime_snapshot() -> None: + startup = _extensions() + startup.mcp_servers["reports"].enabled = False + current = _extensions() + current.mcp_servers["reports"].enabled = False + current.mcp_servers["reports"].env["TOKEN"] = "rotated" + set_mcp_task_config_snapshot(startup) + try: + validate_mcp_task_config_snapshot(current) + finally: + set_mcp_task_config_snapshot(None) + + +def test_mcp_interceptor_changes_require_gateway_restart_for_task_tools() -> None: + startup = _extensions() + current = ExtensionsConfig.model_validate( + { + **startup.model_dump(by_alias=True), + "mcpInterceptors": ["example.interceptor:build"], + } + ) + set_mcp_task_config_snapshot(startup) + try: + with pytest.raises(McpTaskConfigurationError, match="mcpInterceptors.*restart"): + validate_mcp_task_config_snapshot(current) + finally: + set_mcp_task_config_snapshot(None) + + +def test_mcp_interceptor_changes_remain_hot_reloadable_without_task_tools() -> None: + startup = ExtensionsConfig() + current = ExtensionsConfig.model_validate({"mcpInterceptors": ["example.interceptor:build"]}) + set_mcp_task_config_snapshot(startup) + try: + validate_mcp_task_config_snapshot(current) + finally: + set_mcp_task_config_snapshot(None) diff --git a/backend/tests/test_mcp_task_service.py b/backend/tests/test_mcp_task_service.py index a5682f8a9..779ed672b 100644 --- a/backend/tests/test_mcp_task_service.py +++ b/backend/tests/test_mcp_task_service.py @@ -12,6 +12,7 @@ from deerflow.mcp.tasks import ( TaskSubmission, TaskSubmitRequest, ) +from deerflow.mcp.tasks.ordinary import McpTaskProtocolError from deerflow.persistence.mcp_tasks import DuplicateMcpRemoteTaskError @@ -215,6 +216,46 @@ async def test_submit_cancels_remote_task_when_persistence_fails(): } +@pytest.mark.asyncio +async def test_submit_cancels_remote_task_when_its_id_exceeds_storage_limit(): + repo = FakeRepository() + driver = FakeDriver( + submission=TaskSubmission( + remote_task_id="r" * 256, + 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(McpTaskProtocolError, match="remote_task_id.*255"): + await service.submit( + driver_name="fake", + 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={}, + local_task_id="task-1", + ), + ) + + assert repo.created == [] + assert len(driver.cancel_calls) == 1 + assert driver.cancel_calls[0].remote_task_id == "r" * 256 + + @pytest.mark.asyncio async def test_duplicate_remote_handle_is_rejected_without_cancelling_existing_task(): repo = DuplicateCreateRepository() @@ -313,6 +354,26 @@ async def test_run_once_polls_without_an_llm_and_schedules_next_poll(): assert update["polled_at"] > scan_started_at +@pytest.mark.asyncio +async def test_run_once_caps_remote_poll_hint_to_one_day(): + repo = FakeRepository([_claimed_row()]) + driver = FakeDriver([TaskSnapshot(status=TaskStatus.WORKING, poll_after_seconds=1e20)]) + 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)) + + _, update = repo.applied[0] + assert update["next_poll_at"] == update["polled_at"] + timedelta(days=1) + + @pytest.mark.asyncio async def test_run_once_schedules_driver_error_retry_from_poll_completion_time(): repo = FakeRepository([_claimed_row()]) @@ -335,7 +396,7 @@ async def test_run_once_schedules_driver_error_retry_from_poll_completion_time() @pytest.mark.asyncio -async def test_run_once_stops_polling_terminal_and_input_required_snapshots(): +async def test_run_once_stops_terminal_tasks_but_keeps_input_required_on_a_slow_poll(): rows = [_claimed_row(), {**_claimed_row(), "id": "task-2", "remote_task_id": "remote-2"}] repo = FakeRepository(rows) driver = FakeDriver( @@ -361,7 +422,269 @@ async def test_run_once_stops_polling_terminal_and_input_required_snapshots(): 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 + assert updates["task-2"]["next_poll_at"] >= updates["task-2"]["polled_at"] + timedelta(seconds=60) + + +@pytest.mark.asyncio +async def test_run_once_uses_exponential_backoff_and_caps_transient_errors(): + rows = [ + {**_claimed_row(), "id": "task-1", "consecutive_poll_error_count": 0}, + {**_claimed_row(), "id": "task-2", "consecutive_poll_error_count": 4}, + ] + repo = FakeRepository(rows) + 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, + max_poll_backoff_seconds=30, + ) + + started_at = datetime.now(UTC) + await service.run_once(now=started_at) + finished_at = datetime.now(UTC) + + released = {task_id: update for task_id, update in repo.released} + assert started_at + timedelta(seconds=5) <= released["task-1"]["next_poll_at"] <= finished_at + timedelta(seconds=5) + assert started_at + timedelta(seconds=30) <= released["task-2"]["next_poll_at"] <= finished_at + timedelta(seconds=30) + + +@pytest.mark.asyncio +async def test_protocol_error_terminalizes_instead_of_retrying_forever(): + repo = FakeRepository([_claimed_row()]) + registry = McpTaskDriverRegistry() + registry.register("fake", FakeDriver(error=McpTaskProtocolError("missing structuredContent"))) + 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)) + + assert repo.released == [] + _, applied = repo.applied[0] + assert applied["status"] == "failed" + assert applied["error"] == "missing structuredContent" + assert applied["next_poll_at"] is None + + +@pytest.mark.asyncio +async def test_protocol_error_message_is_bounded_before_terminal_persistence(): + oversized_error = "e" * 5_000 + repo = FakeRepository([_claimed_row()]) + registry = McpTaskDriverRegistry() + registry.register("fake", FakeDriver(error=McpTaskProtocolError(oversized_error))) + 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)) + + _, applied = repo.applied[0] + assert applied["status"] == "failed" + assert applied["error"] == oversized_error[:4_000] + + +@pytest.mark.asyncio +async def test_persisted_snapshot_errors_are_bounded_on_submit_and_poll(): + oversized_error = "e" * 5_000 + submit_repo = FakeRepository() + submit_driver = FakeDriver( + submission=TaskSubmission( + remote_task_id="remote-1", + snapshot=TaskSnapshot(status=TaskStatus.FAILED, error=oversized_error), + ) + ) + submit_registry = McpTaskDriverRegistry() + submit_registry.register("fake", submit_driver) + submit_service = McpTaskService( + repository=submit_repo, + drivers=submit_registry, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + await submit_service.submit( + driver_name="fake", + 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={}, + ), + ) + + assert submit_repo.created[0]["error"] == oversized_error[:4_000] + + poll_repo = FakeRepository([_claimed_row()]) + poll_registry = McpTaskDriverRegistry() + poll_registry.register( + "fake", + FakeDriver([TaskSnapshot(status=TaskStatus.FAILED, error=oversized_error)]), + ) + poll_service = McpTaskService( + repository=poll_repo, + drivers=poll_registry, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + await poll_service.run_once(now=datetime.now(UTC)) + + _, applied = poll_repo.applied[0] + assert applied["error"] == oversized_error[:4_000] + + +@pytest.mark.asyncio +async def test_oversized_input_required_payload_terminalizes_without_persisting_it(): + repo = FakeRepository([_claimed_row()]) + registry = McpTaskDriverRegistry() + registry.register( + "fake", + FakeDriver( + [ + TaskSnapshot( + status=TaskStatus.INPUT_REQUIRED, + input_required={"prompt": "x" * 65_536}, + ) + ] + ), + ) + 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)) + + assert repo.released == [] + _, applied = repo.applied[0] + assert applied["status"] == "failed" + assert applied["input_required"] is None + assert "input_required payload exceeds the 65536-byte limit" in applied["error"] + assert applied["next_poll_at"] is None + + +@pytest.mark.asyncio +async def test_oversized_result_stores_preview_without_invalid_truncated_json(): + repo = FakeRepository([_claimed_row()]) + registry = McpTaskDriverRegistry() + registry.register( + "fake", + FakeDriver( + [ + TaskSnapshot( + status=TaskStatus.COMPLETED, + result={"report": "x" * 200}, + result_artifact={"uri": "s3://reports/1.json", "mime_type": "application/json"}, + ) + ] + ), + ) + service = McpTaskService( + repository=repo, + drivers=registry, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + max_result_bytes=64, + result_preview_max_chars=24, + ) + + await service.run_once(now=datetime.now(UTC)) + + _, applied = repo.applied[0] + assert applied["result"] is None + assert len(applied["result_preview"]) == 24 + assert applied["result_truncated"] is True + assert applied["result_artifact"] == { + "uri": "s3://reports/1.json", + "mime_type": "application/json", + } + + +@pytest.mark.asyncio +async def test_oversized_result_artifact_terminalizes_without_persisting_it(): + repo = FakeRepository([_claimed_row()]) + registry = McpTaskDriverRegistry() + registry.register( + "fake", + FakeDriver( + [ + TaskSnapshot( + status=TaskStatus.COMPLETED, + result_artifact={ + "uri": "https://example.test/" + "x" * 65_536, + "mime_type": "application/json", + }, + ) + ] + ), + ) + 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)) + + assert repo.released == [] + _, applied = repo.applied[0] + assert applied["status"] == "failed" + assert applied["result_artifact"] is None + assert "result_artifact payload exceeds the 65536-byte limit" in applied["error"] + + +@pytest.mark.asyncio +async def test_non_json_numeric_result_is_a_permanent_protocol_failure(): + repo = FakeRepository([_claimed_row()]) + registry = McpTaskDriverRegistry() + registry.register( + "fake", + FakeDriver( + [ + TaskSnapshot( + status=TaskStatus.COMPLETED, + result={"score": float("nan")}, + ) + ] + ), + ) + 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)) + + assert repo.released == [] + _, applied = repo.applied[0] + assert applied["status"] == "failed" + assert "not valid JSON" in applied["error"] @pytest.mark.asyncio diff --git a/backend/tests/test_mcp_task_tool_caller.py b/backend/tests/test_mcp_task_tool_caller.py new file mode 100644 index 000000000..064aeee9d --- /dev/null +++ b/backend/tests/test_mcp_task_tool_caller.py @@ -0,0 +1,295 @@ +import asyncio +from collections.abc import Coroutine +from contextlib import suppress +from datetime import timedelta +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.mcp.task_tool_caller import McpTaskToolCaller, mcp_task_session_scope_key + + +def _config() -> ExtensionsConfig: + return ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "type": "stdio", + "command": "report-mcp", + "task_toolsets": [ + { + "name": "reports", + "submit_tool": "submit_report", + "status_tool": "status_report", + "cancel_tool": "cancel_report", + } + ], + } + } + } + ) + + +def _remote_config(transport: str = "http") -> ExtensionsConfig: + return ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "type": transport, + "url": "https://reports.example.com/mcp", + "headers": {"X-Static": "configured"}, + } + } + } + ) + + +class _SessionContext: + def __init__(self, session): + self.session = session + + async def __aenter__(self): + return self.session + + async def __aexit__(self, *_args): + return None + + +async def _assert_configured_timeout(awaitable: Coroutine[Any, Any, Any]) -> None: + task = asyncio.create_task(awaitable) + try: + done, _pending = await asyncio.wait({task}, timeout=0.25) + assert task in done, "configured timeout was ignored" + with pytest.raises(TimeoutError): + await task + finally: + if not task.done(): + task.cancel() + with suppress(asyncio.CancelledError): + await task + + +def test_task_session_scope_includes_user_and_thread() -> None: + assert mcp_task_session_scope_key(user_id="user-1", thread_id="thread-1") == "user-1:thread-1" + + +@pytest.mark.asyncio +async def test_stdio_task_call_reuses_exact_scope_and_raw_tool_name() -> None: + result = SimpleNamespace(structuredContent={"task_id": "remote-1", "status": "running"}, isError=False) + session = SimpleNamespace(call_tool=AsyncMock(return_value=result)) + pool = MagicMock() + pool.get_session = AsyncMock(return_value=session) + pool.close_session = AsyncMock() + caller = McpTaskToolCaller(_config()) + + with ( + patch("deerflow.mcp.task_tool_caller.get_session_pool", return_value=pool), + patch( + "deerflow.mcp.task_tool_caller._prepare_stdio_connection", + return_value={"transport": "stdio", "command": "report-mcp"}, + ), + ): + actual = await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + + assert actual is result + pool.get_session.assert_awaited_once_with( + "reports", + "user-1:thread-1", + {"transport": "stdio", "command": "report-mcp"}, + ) + session.call_tool.assert_awaited_once_with("status_report", {"task_id": "remote-1"}) + pool.close_session.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_broken_stdio_task_session_is_evicted_for_next_poll_reconnect() -> None: + session = SimpleNamespace(call_tool=AsyncMock(side_effect=ConnectionError("disconnected"))) + pool = MagicMock() + pool.get_session = AsyncMock(return_value=session) + pool.close_session = AsyncMock() + caller = McpTaskToolCaller(_config()) + + with ( + patch("deerflow.mcp.task_tool_caller.get_session_pool", return_value=pool), + patch( + "deerflow.mcp.task_tool_caller._prepare_stdio_connection", + return_value={"transport": "stdio", "command": "report-mcp"}, + ), + pytest.raises(ConnectionError, match="disconnected"), + ): + await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + + pool.close_session.assert_awaited_once_with("reports", "user-1:thread-1") + + +@pytest.mark.asyncio +async def test_stdio_task_session_initialization_respects_configured_timeout() -> None: + config = _config() + config.mcp_servers["reports"].session_init_timeout = 0.01 + + async def slow_get_session(*_args): + await asyncio.sleep(60) + + pool = MagicMock() + pool.get_session = AsyncMock(side_effect=slow_get_session) + pool.close_session = AsyncMock() + caller = McpTaskToolCaller(config) + + with ( + patch("deerflow.mcp.task_tool_caller.get_session_pool", return_value=pool), + patch( + "deerflow.mcp.task_tool_caller._prepare_stdio_connection", + return_value={"transport": "stdio", "command": "report-mcp"}, + ), + pytest.raises(TimeoutError), + ): + await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + + pool.close_session.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_http_task_call_authenticates_session_initialization() -> None: + result = SimpleNamespace(structuredContent={"task_id": "remote-1", "status": "running"}, isError=False) + session = SimpleNamespace( + initialize=AsyncMock(), + call_tool=AsyncMock(return_value=result), + ) + create_session = MagicMock(return_value=_SessionContext(session)) + caller = McpTaskToolCaller( + _remote_config(), + oauth_token_manager=SimpleNamespace( + has_oauth_servers=lambda: False, + get_authorization_header=AsyncMock(return_value="Bearer task-token"), + ), + ) + + with patch( + "langchain_mcp_adapters.sessions.create_session", + create_session, + ): + actual = await caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + + assert actual is result + create_session.assert_called_once_with( + { + "transport": "http", + "url": "https://reports.example.com/mcp", + "headers": { + "X-Static": "configured", + "Authorization": "Bearer task-token", + }, + } + ) + session.initialize.assert_awaited_once_with() + session.call_tool.assert_awaited_once_with( + "status_report", + {"task_id": "remote-1"}, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "sse"]) +async def test_remote_task_session_initialization_respects_configured_timeout(transport: str) -> None: + config = _remote_config(transport) + config.mcp_servers["reports"].session_init_timeout = 0.01 + + async def slow_initialize(): + await asyncio.sleep(60) + + session = SimpleNamespace( + initialize=AsyncMock(side_effect=slow_initialize), + call_tool=AsyncMock(), + ) + caller = McpTaskToolCaller( + config, + oauth_token_manager=SimpleNamespace( + has_oauth_servers=lambda: False, + get_authorization_header=AsyncMock(return_value=None), + ), + ) + + with patch( + "langchain_mcp_adapters.sessions.create_session", + MagicMock(return_value=_SessionContext(session)), + ): + await _assert_configured_timeout( + caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + ) + + session.call_tool.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport", ["http", "sse"]) +async def test_remote_task_call_respects_configured_timeout(transport: str) -> None: + config = _remote_config(transport) + config.mcp_servers["reports"].tool_call_timeout = 0.01 + + async def slow_call(*_args, **_kwargs): + await asyncio.sleep(60) + + session = SimpleNamespace( + initialize=AsyncMock(), + call_tool=AsyncMock(side_effect=slow_call), + ) + caller = McpTaskToolCaller( + config, + oauth_token_manager=SimpleNamespace( + has_oauth_servers=lambda: False, + get_authorization_header=AsyncMock(return_value=None), + ), + ) + + with patch( + "langchain_mcp_adapters.sessions.create_session", + MagicMock(return_value=_SessionContext(session)), + ): + await _assert_configured_timeout( + caller.call_tool( + server_name="reports", + tool_name="status_report", + arguments={"task_id": "remote-1"}, + user_id="user-1", + thread_id="thread-1", + ) + ) + + session.call_tool.assert_awaited_once_with( + "status_report", + {"task_id": "remote-1"}, + read_timeout_seconds=timedelta(seconds=0.01), + ) diff --git a/backend/tests/test_mcp_task_tool_wrapping.py b/backend/tests/test_mcp_task_tool_wrapping.py new file mode 100644 index 000000000..88e10659a --- /dev/null +++ b/backend/tests/test_mcp_task_tool_wrapping.py @@ -0,0 +1,203 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from langchain_core.tools import StructuredTool +from pydantic import BaseModel + +from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig +from deerflow.mcp.tasks.runtime import ( + McpTaskConfigurationError, + set_mcp_task_config_snapshot, + set_mcp_task_submitter, +) +from deerflow.mcp.tools import _configure_task_tools_for_server, get_mcp_tools + + +class _SubmitArgs(BaseModel): + topic: str + + +def _tool(name: str, *, description: str | None = None) -> StructuredTool: + async def call(topic: str): + return topic + + return StructuredTool( + name=name, + description=description if description is not None else name, + args_schema=_SubmitArgs, + coroutine=call, + ) + + +def _server_config() -> McpServerConfig: + return McpServerConfig.model_validate( + { + "task_toolsets": [ + { + "name": "report-generation", + "submit_tool": "submit_report", + "status_tool": "get_report_status", + "cancel_tool": "cancel_report", + } + ] + } + ) + + +class FakeSubmitter: + def __init__(self): + self.calls = [] + + async def submit(self, **kwargs): + self.calls.append(kwargs) + return { + "id": "mcp-task-local-1", + "status": "submitted", + "remote_task_id": "must-not-leak", + "driver_data": {"must": "not leak"}, + } + + +def test_unconfigured_server_tools_are_returned_unchanged() -> None: + tools = [_tool("reports_search")] + + configured = _configure_task_tools_for_server( + tools, + server_name="reports", + server_config=McpServerConfig(), + tool_name_prefix=True, + ) + + assert configured == tools + assert configured[0] is tools[0] + + +def test_configured_status_and_cancel_tools_are_hidden_from_the_agent() -> None: + tools = [ + _tool("reports_submit_report"), + _tool("reports_get_report_status"), + _tool("reports_cancel_report"), + _tool("reports_search"), + ] + + configured = _configure_task_tools_for_server( + tools, + server_name="reports", + server_config=_server_config(), + tool_name_prefix=True, + ) + + assert [tool.name for tool in configured] == ["reports_submit_report", "reports_search"] + + +def test_submit_wrapper_preserves_server_description_and_appends_background_contract() -> None: + tools = [ + _tool( + "submit_report", + description="Generate a quarterly financial report for the requested topic.", + ), + _tool("get_report_status"), + _tool("cancel_report"), + ] + + configured = _configure_task_tools_for_server( + tools, + server_name="reports", + server_config=_server_config(), + tool_name_prefix=False, + ) + + assert configured[0].description == ( + "Generate a quarterly financial report for the requested topic.\n\nSubmitted as durable background task 'report-generation'; returns a DeerFlow task ID immediately and status polling is handled automatically." + ) + + +def test_configured_task_toolsets_fail_when_a_raw_tool_is_missing() -> None: + with pytest.raises(McpTaskConfigurationError, match="cancel_report"): + _configure_task_tools_for_server( + [_tool("reports_submit_report"), _tool("reports_get_report_status")], + server_name="reports", + server_config=_server_config(), + tool_name_prefix=True, + ) + + +@pytest.mark.asyncio +async def test_submit_wrapper_persists_before_returning_only_the_local_handle() -> None: + submitter = FakeSubmitter() + set_mcp_task_submitter(submitter) + try: + configured = _configure_task_tools_for_server( + [ + _tool("submit_report"), + _tool("get_report_status"), + _tool("cancel_report"), + ], + server_name="reports", + server_config=_server_config(), + tool_name_prefix=False, + ) + submit_tool = configured[0] + runtime = SimpleNamespace( + context={"thread_id": "thread-1", "run_id": "run-1"}, + config={}, + tool_call_id="call-1", + ) + + result = await submit_tool.coroutine(runtime=runtime, topic="MCP") + + assert result == { + "task_id": "mcp-task-local-1", + "task_name": "report-generation", + "status": "submitted", + "message": "Task is running in the background.", + } + call = submitter.calls[0] + request = call["request"] + assert call["driver_name"] == "ordinary-tools" + assert request.user_id == "test-user-autouse" + assert request.thread_id == "thread-1" + assert request.run_id == "run-1" + assert request.tool_call_id == "call-1" + assert request.arguments == {"topic": "MCP"} + assert request.driver_data == { + "submit_tool": "submit_report", + "status_tool": "get_report_status", + "cancel_tool": "cancel_report", + } + finally: + set_mcp_task_submitter(None) + + +@pytest.mark.asyncio +async def test_submit_wrapper_fails_clearly_without_gateway_task_runtime() -> None: + set_mcp_task_submitter(None) + configured = _configure_task_tools_for_server( + [_tool("submit_report"), _tool("get_report_status"), _tool("cancel_report")], + server_name="reports", + server_config=_server_config(), + tool_name_prefix=False, + ) + + with pytest.raises(McpTaskConfigurationError, match="not initialized"): + await configured[0].coroutine(topic="MCP") + + +@pytest.mark.asyncio +async def test_tool_reload_rejects_task_server_runtime_config_drift() -> None: + startup = ExtensionsConfig(mcpServers={"reports": _server_config()}) + current = ExtensionsConfig(mcpServers={"reports": _server_config()}) + current.mcp_servers["reports"].env["TOKEN"] = "rotated" + set_mcp_task_config_snapshot(startup) + try: + with ( + patch( + "deerflow.mcp.tools.ExtensionsConfig.from_file", + return_value=current, + ), + pytest.raises(McpTaskConfigurationError, match="reports.*restart"), + ): + await get_mcp_tools() + finally: + set_mcp_task_config_snapshot(None) diff --git a/backend/tests/test_mcp_task_toolset_config.py b/backend/tests/test_mcp_task_toolset_config.py new file mode 100644 index 000000000..3b342e4d7 --- /dev/null +++ b/backend/tests/test_mcp_task_toolset_config.py @@ -0,0 +1,133 @@ +import pytest +from pydantic import ValidationError + +from app.gateway.routers.mcp import McpServerConfigResponse +from deerflow.config.extensions_config import ExtensionsConfig + + +def test_task_toolsets_preserve_raw_tool_names_and_support_multiple_groups() -> None: + config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "type": "http", + "url": "https://example.test/mcp", + "task_toolsets": [ + { + "name": "report-generation", + "submit_tool": "submit_report", + "status_tool": "get_report_status", + "cancel_tool": "cancel_report", + }, + { + "name": "data-export", + "submit_tool": "start_export", + "status_tool": "get_export_status", + "cancel_tool": "cancel_export", + }, + ], + } + } + } + ) + + toolsets = config.mcp_servers["reports"].task_toolsets + assert [toolset.name for toolset in toolsets] == ["report-generation", "data-export"] + assert toolsets[0].model_dump() == { + "name": "report-generation", + "submit_tool": "submit_report", + "status_tool": "get_report_status", + "cancel_tool": "cancel_report", + } + response = McpServerConfigResponse.model_validate(config.mcp_servers["reports"].model_dump()) + assert response.task_toolsets[0].submit_tool == "submit_report" + + +@pytest.mark.parametrize( + "duplicate_field,duplicate_value", + [ + ("status_tool", "submit_report"), + ("cancel_tool", "submit_report"), + ], +) +def test_task_toolsets_reject_reusing_one_raw_tool_in_multiple_roles( + duplicate_field: str, + duplicate_value: str, +) -> None: + toolset = { + "name": "report-generation", + "submit_tool": "submit_report", + "status_tool": "get_report_status", + "cancel_tool": "cancel_report", + } + toolset[duplicate_field] = duplicate_value + + with pytest.raises(ValidationError, match="must be unique across task_toolsets"): + ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "task_toolsets": [toolset], + } + } + } + ) + + +def test_task_toolsets_reject_reusing_one_raw_tool_across_groups() -> None: + with pytest.raises(ValidationError, match="submit_report.*must be unique"): + ExtensionsConfig.model_validate( + { + "mcpServers": { + "reports": { + "task_toolsets": [ + { + "name": "first", + "submit_tool": "submit_report", + "status_tool": "status_report", + "cancel_tool": "cancel_report", + }, + { + "name": "second", + "submit_tool": "start_export", + "status_tool": "submit_report", + "cancel_tool": "cancel_export", + }, + ] + } + } + } + ) + + +@pytest.mark.parametrize( + ("server_name", "task_name", "match"), + [ + (" ", "reports", "server name.*128"), + ("s" * 129, "reports", "server name.*128"), + ("reports", " ", "task toolset name must not be empty"), + ("reports", "t" * 256, "at most 255"), + ], +) +def test_task_toolsets_reject_names_that_do_not_fit_durable_storage( + server_name: str, + task_name: str, + match: str, +) -> None: + with pytest.raises(ValidationError, match=match): + ExtensionsConfig.model_validate( + { + "mcpServers": { + server_name: { + "task_toolsets": [ + { + "name": task_name, + "submit_tool": "submit_report", + "status_tool": "status_report", + "cancel_tool": "cancel_report", + } + ] + } + } + } + ) diff --git a/backend/tests/test_mcp_tasks_router.py b/backend/tests/test_mcp_tasks_router.py new file mode 100644 index 000000000..154532396 --- /dev/null +++ b/backend/tests/test_mcp_tasks_router.py @@ -0,0 +1,140 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException + +from app.gateway.app import create_app +from app.gateway.routers import mcp_tasks + + +class FakeRepository: + def __init__(self, rows): + self.rows = rows + self.list_calls = [] + self.get_calls = [] + + async def list_by_thread(self, thread_id, *, user_id, limit): + self.list_calls.append((thread_id, user_id, limit)) + return list(self.rows) + + async def get(self, task_id, *, user_id): + self.get_calls.append((task_id, user_id)) + return next((row for row in self.rows if row["id"] == task_id and row["user_id"] == user_id), None) + + +def _record(**overrides): + return { + "id": "mcp-task-1", + "user_id": "user-1", + "thread_id": "thread-1", + "task_name": "report-generation", + "status": "working", + "created_at": "2026-08-05T00:00:00+00:00", + "updated_at": "2026-08-05T00:00:05+00:00", + "last_polled_at": "2026-08-05T00:00:05+00:00", + "error": None, + "last_poll_error": "temporary network failure", + "consecutive_poll_error_count": 3, + "result": None, + "result_preview": None, + "result_truncated": False, + "result_artifact": None, + "input_required": None, + "remote_task_id": "must-not-leak", + "driver_data": {"status_tool": "must-not-leak"}, + "server_name": "must-not-leak", + **overrides, + } + + +def _request(repo): + return SimpleNamespace( + app=SimpleNamespace( + state=SimpleNamespace( + mcp_task_repo=repo, + mcp_task_service=SimpleNamespace(tracking_degraded_after_errors=3), + ) + ) + ) + + +def test_gateway_mounts_thread_scoped_mcp_task_routes() -> None: + paths = {route.path for route in create_app().routes} + assert "/api/threads/{thread_id}/mcp-tasks" in paths + assert "/api/threads/{thread_id}/mcp-tasks/{task_id}" in paths + + +@pytest.mark.asyncio +async def test_list_returns_only_safe_current_user_thread_fields(monkeypatch) -> None: + repo = FakeRepository([_record()]) + monkeypatch.setattr(mcp_tasks, "get_current_user", AsyncMock(return_value="user-1")) + + response = await mcp_tasks.list_mcp_tasks.__wrapped__( + thread_id="thread-1", + request=_request(repo), + limit=25, + ) + + assert repo.list_calls == [("thread-1", "user-1", 25)] + assert response == [ + { + "task_id": "mcp-task-1", + "task_name": "report-generation", + "status": "working", + "created_at": "2026-08-05T00:00:00+00:00", + "updated_at": "2026-08-05T00:00:05+00:00", + "error": None, + "tracking_degraded": True, + } + ] + + +@pytest.mark.asyncio +async def test_detail_exposes_bounded_result_but_not_remote_handle(monkeypatch) -> None: + repo = FakeRepository( + [ + _record( + status="completed", + result={"report": "ready"}, + result_artifact={"uri": "s3://reports/1.json", "mime_type": "application/json"}, + ) + ] + ) + monkeypatch.setattr(mcp_tasks, "get_current_user", AsyncMock(return_value="user-1")) + + response = await mcp_tasks.get_mcp_task.__wrapped__( + thread_id="thread-1", + task_id="mcp-task-1", + request=_request(repo), + ) + + assert response["result"] == {"report": "ready"} + assert response["result_artifact"]["uri"] == "s3://reports/1.json" + assert "remote_task_id" not in response + assert "driver_data" not in response + assert "server_name" not in response + + +@pytest.mark.asyncio +async def test_detail_rejects_cross_user_and_cross_thread_access(monkeypatch) -> None: + repo = FakeRepository([_record()]) + request = _request(repo) + + monkeypatch.setattr(mcp_tasks, "get_current_user", AsyncMock(return_value="user-2")) + with pytest.raises(HTTPException) as cross_user: + await mcp_tasks.get_mcp_task.__wrapped__( + thread_id="thread-1", + task_id="mcp-task-1", + request=request, + ) + assert cross_user.value.status_code == 404 + + monkeypatch.setattr(mcp_tasks, "get_current_user", AsyncMock(return_value="user-1")) + with pytest.raises(HTTPException) as cross_thread: + await mcp_tasks.get_mcp_task.__wrapped__( + thread_id="thread-2", + task_id="mcp-task-1", + request=request, + ) + assert cross_thread.value.status_code == 404 diff --git a/backend/tests/test_mcp_tool_name_prefix.py b/backend/tests/test_mcp_tool_name_prefix.py index 30319192f..933df71f2 100644 --- a/backend/tests/test_mcp_tool_name_prefix.py +++ b/backend/tests/test_mcp_tool_name_prefix.py @@ -117,7 +117,7 @@ async def test_mcp_tool_name_prefix_can_be_disabled_per_server_without_disabling patch("deerflow.mcp.tools.ExtensionsConfig.from_file", return_value=extensions_config), patch("deerflow.mcp.tools.build_servers_config", return_value=servers_config), patch("deerflow.mcp.tools.get_initial_oauth_headers", new_callable=AsyncMock, return_value={}), - patch("deerflow.mcp.tools.build_oauth_tool_interceptor", return_value=None), + patch("deerflow.mcp.tools.build_mcp_tool_interceptors", return_value=[]), patch("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient), patch("langchain_mcp_adapters.tools.load_mcp_tools", side_effect=fake_load_mcp_tools), patch("deerflow.mcp.tools._make_session_pool_tool", side_effect=lambda tool, *_args, **_kwargs: tool) as wrap_tool, diff --git a/backend/tests/test_migration_0004_run_ownership_dedupe.py b/backend/tests/test_migration_0004_run_ownership_dedupe.py index 2f10524ae..72bc978fb 100644 --- a/backend/tests/test_migration_0004_run_ownership_dedupe.py +++ b/backend/tests/test_migration_0004_run_ownership_dedupe.py @@ -157,7 +157,7 @@ async def test_migration_dedupes_duplicate_active_rows_before_unique_index(tmp_p with sqlite3.connect(db_path) as raw: version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() # Bootstrap upgrades through the later revisions after 0004. - assert version_row[0] == "0011_mcp_tasks" + assert version_row[0] == "0012_mcp_task_results" # Sanity: the invariant the index enforces is now true — at most one # active row per thread. diff --git a/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py b/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py index f98780405..fb498b3d7 100644 --- a/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py +++ b/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py @@ -169,7 +169,7 @@ async def test_migration_supersedes_duplicate_active_runs_before_unique_index(tm with sqlite3.connect(db_path) as raw: version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0011_mcp_tasks" + assert version_row[0] == "0012_mcp_task_results" # Sanity: the invariant the index enforces now holds — at most one # active row per task_id. diff --git a/backend/tests/test_persistence_bootstrap.py b/backend/tests/test_persistence_bootstrap.py index 67957e318..d85053277 100644 --- a/backend/tests/test_persistence_bootstrap.py +++ b/backend/tests/test_persistence_bootstrap.py @@ -48,7 +48,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default asyncio_test = pytest.mark.asyncio -HEAD = "0011_mcp_tasks" +HEAD = "0012_mcp_task_results" BASELINE = "0001_baseline" diff --git a/backend/tests/test_persistence_bootstrap_concurrency.py b/backend/tests/test_persistence_bootstrap_concurrency.py index 7b80f8d95..967b283ca 100644 --- a/backend/tests/test_persistence_bootstrap_concurrency.py +++ b/backend/tests/test_persistence_bootstrap_concurrency.py @@ -28,7 +28,7 @@ from deerflow.persistence.bootstrap import bootstrap_schema pytestmark = pytest.mark.asyncio -HEAD = "0011_mcp_tasks" +HEAD = "0012_mcp_task_results" def _url(tmp_path: Path) -> str: diff --git a/backend/tests/test_persistence_bootstrap_regression.py b/backend/tests/test_persistence_bootstrap_regression.py index d8ec9d63a..73e9d7efb 100644 --- a/backend/tests/test_persistence_bootstrap_regression.py +++ b/backend/tests/test_persistence_bootstrap_regression.py @@ -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()} assert "token_usage_by_model" in cols version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0011_mcp_tasks" + assert version_row[0] == "0012_mcp_task_results" # And the read path that originally 500'd must now succeed. 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. assert cols.count("token_usage_by_model") == 1 version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0011_mcp_tasks" + assert version_row[0] == "0012_mcp_task_results" finally: await close_engine() diff --git a/config.example.yaml b/config.example.yaml index 9809fea5b..92802160d 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -2075,8 +2075,9 @@ scheduler: # ============================================================================ # 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. +# Durable runtime for ordinary MCP submit/status/cancel task toolsets. It is +# disabled by default; task_toolsets are configured per server in +# extensions_config.json. # All fields are restart-required (captured at Gateway lifespan startup). # # mcp_tasks: @@ -2084,11 +2085,21 @@ scheduler: # 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 +# max_poll_backoff_seconds: 300 # Cap for exponential retries after transient errors +# input_required_poll_interval_seconds: 60 # Minimum poll interval while waiting for user input +# tracking_degraded_after_errors: 3 # Consecutive errors before query API reports degraded tracking +# max_result_bytes: 65536 # Full JSON result storage limit +# result_preview_max_chars: 2000 # Text preview retained when a result exceeds the limit mcp_tasks: enabled: false poll_interval_seconds: 5 lease_seconds: 120 max_concurrent_polls: 8 + max_poll_backoff_seconds: 300 + input_required_poll_interval_seconds: 60 + tracking_degraded_after_errors: 3 + max_result_bytes: 65536 + result_preview_max_chars: 2000 # ============================================================================ # Run Ownership Configuration diff --git a/extensions_config.example.json b/extensions_config.example.json index 022e7c190..f2c701cdd 100644 --- a/extensions_config.example.json +++ b/extensions_config.example.json @@ -64,6 +64,22 @@ } } } + }, + "long-running-reports": { + "enabled": false, + "type": "http", + "url": "https://reports.example.com/mcp", + "description": "Example ordinary MCP background-task contract", + "session_init_timeout": 60, + "tool_call_timeout": 60, + "task_toolsets": [ + { + "name": "report-generation", + "submit_tool": "submit_report", + "status_tool": "get_report_status", + "cancel_tool": "cancel_report" + } + ] } }, "skills": {}