mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 05:06:59 +00:00
fix(mcp): isolate pooled sessions by ownership domain (#6133)
* fix(mcp): isolate pooled sessions by ownership domain Separate deployment and personal stdio sessions even when they share the same runtime server name and thread scope. Preserve deployment call compatibility while carrying personal ownership through Agent wrappers, durable task calls, and disconnect eviction. Add focused regression coverage for personal discovery wiring and domain-scoped disconnect cleanup. * test(mcp): address ownership-domain review feedback Accept the new wrapper keyword in the prefix-overlap test double and reuse MCPPoolDomain for internal task-caller ownership parameters. * style(mcp): format ownership-domain changes * style(mcp): apply repository ruff formatting
This commit is contained in:
spud
GitHub
parent
9dd8e3e94b
commit
3f8f91d064
@ -3,7 +3,7 @@
|
||||
- Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management
|
||||
- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` contract and normalized `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`). A driver-supplied `poll_after_seconds` must be a finite positive number, validated at the `TaskSnapshot` boundary so every driver is held to the same invariant rather than each one guarding the consumer that turns the interval into a `timedelta`. `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and separate consecutive poll/delivery error counters; `app/mcp_tasks/McpTaskService` performs status, cancellation, and notification work outside the Agent/LLM loop. Notification retries keep their idempotency attempt separate from the delivery-failure count, use capped exponential backoff, and stop after five failures; strict existing-thread admission dead-letters a deleted/mismatched target immediately. A status result is applied only when the worker still owns an unexpired lease, so a stale result cannot be written after expiry even before another worker reclaims the row. Every claim carries a fresh per-claim token, and every poll/cancel/notification release or snapshot apply is fenced by the owner plus that token, so a release that completes after the same service reclaims the row cannot clear the newer lease. Poll, cancel, and notification claims also have separate phase-level single-flight owners that protect only an ambiguous claim outcome: after an uncancelled timeout or caller cancellation, later scans do not start another claim while the outcome is unknown. Once the claim resolves, the phase owner is released immediately, and any returned rows are released through bounded, service-owned background work; per-claim token fencing prevents a late release from mutating a newer claim generation, so a stuck release does not stall the whole phase. While the claim outcome remains unresolved, later scans skip the phase and emit a warning rather than overlapping an ambiguous database claim. Once the claim resolves, the handoff may continue releasing returned rows in the background without blocking later scans. Lease expiry remains the cross-process crash fallback rather than the normal recovery for a live service's late claim. Poll timestamps and retry schedules are based on the remote call's completion time rather than the scan start. If submission succeeds but persistence fails or the caller is cancelled while persistence is in flight, the service best-effort cancels the remote task and preserves the original error or cancellation if that compensation also fails. The exact `uq_mcp_tasks_user_server_remote` conflict is different: an existing durable row already owns the remote handle, so the conflict surfaces without cancelling that tracked task. Unexpected per-task poll failures are isolated from sibling claims and remain recoverable through lease expiry; Gateway shutdown cancels the poller so a hung external status call cannot block process exit. Routine cancellation releases clear only the lease/token and preserve any pre-existing `last_cancel_error` or `notification_error`; cancellation itself is not persisted as a task failure diagnostic. Cancelling an in-flight poll releases only the owner- and per-claim-token-fenced lease and preserves its preclaim schedule and poll-failure state; real poll failures retain exponential backoff and tracking-degradation behavior. `input_required` and terminal states stop polling and become `notification_status=pending` for later Agent/UI delivery. Durable recovery requires a SQL database backend (`sqlite` or `postgres`); the in-memory backend leaves the repository/service unavailable. The runtime is startup-configured by `mcp_tasks` and disabled by default until a concrete driver is registered; this foundation does not alter ordinary MCP tool behavior on its own.
|
||||
- **Runtime availability boundary**: the installed process-local submitter is the source of truth for durable task-management tool exposure. `mcp_tasks` is startup-only; changing it on disk does not alter the live toolset until the Gateway restarts.
|
||||
- **Long-running ordinary task driver**: `extensions_config.json -> mcpServers.<server>.task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same canonical `(server_name, versioned user/thread/incarnation scope)` stdio session as ordinary calls; non-NULL incarnations use an unambiguous versioned encoding, while an explicitly captured legacy NULL retains the pre-activation `user_id:thread_id` scope. A runtime with a missing/invalid incarnation key fails closed. HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task.
|
||||
- **Long-running ordinary task driver**: `extensions_config.json -> mcpServers.<server>.task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same canonical `(server_name, versioned user/thread/incarnation scope)` stdio session as ordinary calls inside the persisted deployment/personal ownership domain; non-NULL incarnations use an unambiguous versioned encoding, while an explicitly captured legacy NULL retains the pre-activation `user_id:thread_id` scope. A runtime with a missing/invalid incarnation key fails closed. HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task.
|
||||
- **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior.
|
||||
- **Remote task bring-up**: One `session_init_timeout` deadline covers transport context entry and MCP initialization. Report the server and configured bound on expiry; disable it before calling the tool. Enter and exit the adapter's AnyIO context managers in the same task; tests include real SSE endpoint/initialization stalls and connection cleanup.
|
||||
- **Background task identity**: `McpTaskToolCaller` binds an ID-only `_TaskOwner` in the existing user ContextVar only for background HTTP/SSE calls with enabled `user_auth`, restoring the previous value in `finally`. The `CurrentUser` contract guarantees only `.id`, not profile fields or roles. This lets `user_auth` select configured credentials without an active Agent run, including recovery from SQL; concurrent owners remain task-local. Submit, stdio, and calls without enabled `user_auth` keep their original context. Do not persist request secrets or enable `headers_from_context` for background status/cancel calls.
|
||||
@ -11,12 +11,13 @@
|
||||
`asyncio.run` fallback covers only `get_event_loop()` raising (no loop set, or a closed
|
||||
loop); a `RuntimeError` from discovery itself (e.g. `McpTaskConfigurationError`) is
|
||||
logged once and never triggers a second discovery pass.
|
||||
- **Loop-isolated stdio sessions**: Both the live registry and in-flight creations
|
||||
are keyed by `(server_name, scope_key, owning_loop)`. Same-loop callers share
|
||||
initialization and state; live sibling loops must never cancel or replace each
|
||||
other's sessions merely because their server/scope matches. Explicit server,
|
||||
scope, pair, thread-identity and global cleanup span all loops; disconnect
|
||||
eviction matches the exact session identity. Thread-identity cleanup
|
||||
- **Loop- and ownership-isolated stdio sessions**: Both the live registry and in-flight creations
|
||||
are keyed by `(server_name, scope_key, owning_loop, ownership_domain)`. Same-loop
|
||||
callers share initialization and state only inside the same `deployment` or
|
||||
`personal` domain; equal runtime names across those domains never share a
|
||||
session. Server/pair cleanup and disconnect eviction stay inside the caller's
|
||||
domain, while scope, thread-identity and global cleanup intentionally span
|
||||
domains. Thread-identity cleanup
|
||||
(`close_thread_scope`) deliberately matches *every incarnation* of a
|
||||
user/thread pair instead of one exact scope key: a thread-deletion caller
|
||||
cannot rely on reading the current incarnation (the record may predate
|
||||
@ -50,7 +51,7 @@
|
||||
add a parallel routing middleware for PR1-style preference hints.
|
||||
- **Stdio file outputs**: Persistent stdio sessions use the canonical versioned user/thread/incarnation scope (legacy NULL incarnations use the exact `user_id:thread_id` key). For stdio transports only, DeerFlow pins the subprocess default `cwd` to the thread workspace and `TMPDIR`/`TMP`/`TEMP` to `workspace/.mcp/tmp/`, unless the operator explicitly configured `cwd` or temp env values. `.mcp` is a DeerFlow-owned internal namespace: its temporary/debug files remain addressable when returned by a tool but are excluded from run workspace-change summaries — by directory name at any depth, consistent with the other reserved names in `EXCLUDED_DIR_NAMES` (`.git`, `node_modules`, …) and robust if a server ever creates a relative `.mcp` from a different cwd. Both launch paths pin it at the workspace root today. SSE/HTTP transports skip this filesystem prep entirely.
|
||||
- **Configured stdio working directory**: A non-empty `McpServerConfig.cwd` is forwarded by `build_server_params` for discovery, pooled Agent calls, and durable task calls. Omit the connection key when the setting is `None` or empty (including unresolved environment references), so the runtime's thread-workspace default still applies; HTTP/SSE connections must not receive it. `tests/test_mcp_cwd.py` exercises the JSON-config-to-subprocess path with relative entrypoints, relative tool inputs, and discovery/pool defaults for omitted, null, empty, and unresolved values.
|
||||
- **Stdio disconnect recovery**: Ordinary Agent tool calls and durable task submit/status/cancel calls that receive the MCP SDK's explicit `Connection closed` error or an AnyIO closed-stream error evict only that `(server_name, incarnation scope)` session when the registered entry is still the same `ClientSession` that failed. A late error from an old concurrent call cannot evict its replacement or a new in-flight creation. The failing call still surfaces its original error and is never replayed automatically; a later retry creates a fresh subprocess/session. Protocol timeouts, normal `isError=true` tool results, and interceptor failures do not evict a healthy stateful session.
|
||||
- **Stdio disconnect recovery**: Ordinary Agent tool calls and durable task submit/status/cancel calls that receive the MCP SDK's explicit `Connection closed` error or an AnyIO closed-stream error evict only that `(ownership domain, server_name, incarnation scope)` session when the registered entry is still the same `ClientSession` that failed. A late error from an old concurrent call cannot evict its replacement or a new in-flight creation. The failing call still surfaces its original error and is never replayed automatically; a later retry creates a fresh subprocess/session. Protocol timeouts, normal `isError=true` tool results, and interceptor failures do not evict a healthy stateful session.
|
||||
- **Stdio path translation**: MCP-returned local file references are not copied. If a `ResourceLink` or conservative free-text path resolves to an existing file inside the thread's mounted user-data tree, it is translated deterministically to `/mnt/user-data/...`; paths outside that tree remain unchanged. Bare-filename correlation must not replace a segment of a longer path: both `/` and `\` are separators at either boundary, even when the full-path pass leaves the reference unresolved. A trailing `\` immediately before a newline therefore counts as a separator too; that position is more often a Markdown hard break, but a backslash is also a legal filename character on POSIX, so forgoing that single rewrite is the conservative tradeoff. `tests/test_mcp_file_migration.py` covers this through the result converter with real workspace files.
|
||||
- **Runtime updates**: Gateway API saves to extensions_config.json; the Gateway-embedded runtime detects changes via the resolved-path + content-signature check above, so multi-worker / stale-mtime deployments still pick up an added/removed MCP server without a restart (`PUT /api/mcp/config` keeps whole-payload validation, while `PATCH /api/mcp/config` changes only one server's `enabled` field, normalizes the same `type`/MCP-spec `transport` alias as the runtime config model, and validates the target only when enabling it). MCP, skill, and embedded-client writers hold the process-local `extensions_config_write_lock` plus the sidecar advisory `extensions_config_file_lock` for the complete read-modify-write/reload cycle, then share `atomic_write_extensions_config()`, which writes and fsyncs a same-directory temporary file before `os.replace()` and preserves an existing file's mode and symlink target; failed serialization or replacement leaves the prior config intact and cleans up the temporary file. Runtime updates are picked up either by the writer's explicit `reset_mcp_tools_cache()` or by another worker's two-stage staleness check above; a change that leaves the effective MCP slice unchanged (skills-only, middleware-only) no longer retires cached tools or pooled sessions, and when no MCP server is enabled the tool-assembly path calls `refresh_mcp_cache_if_active()`, which retires an already-initialized cache and pool (and voids an in-flight initialization) without performing new tool discovery or config hashing on a never-initialized process. The Gateway admin reset (`POST /api/mcp/cache/reset`) additionally calls `publish_mcp_tools_cache_reset()`: it writes a random generation to a hidden JSON marker beside the resolved config under the same cross-process lock, then retires local state. Every worker mounting that directory records the marker signature when tools publish and invalidates on a later signature, so a remote `tools/list` change converges without a fake config edit; marker drift during discovery fences the stale result. The response reports this transport as `scope=shared_config` rather than claiming an implicit broadcast to replicas with independent filesystems. No resolvable config path keeps the legacy process-local behavior and is reported as `scope=process`.
|
||||
- **Stdio launch policy at the HTTP boundary** (`routers/mcp.py::_validate_mcp_update_request`, shared by `PUT` and the enable branch of `PATCH`): a config file may express anything, but the API is untrusted input, so an API-registered stdio server must (a) name a bare executable from the allowlist — `_DEFAULT_MCP_STDIO_COMMAND_ALLOWLIST` = `{npx, uvx}`, extended by `DEER_FLOW_MCP_STDIO_COMMAND_ALLOWLIST`, with path separators, whitespace, and shell metacharacters rejected in `command`; (b) carry no `args` flag in `_ARBITRARY_EXEC_ARGS`; and (c) set no `env` name in `_CODE_INJECTING_ENV_VARS`. Checks (b) and (c) exist because the command check alone names a binary without constraining what that binary runs. The `env` denylist applies to **every** allowlisted command, and both denylists match `--flag=value` as well as `--flag value`. The `args` denylist's **scope depends on the command**, because where a launcher stops parsing its own flags is what decides whether a token is an exec flag at all:
|
||||
|
||||
@ -5,7 +5,7 @@ each tool call creates a new MCP session. For stateful servers like Playwright,
|
||||
this means browser state (opened pages, filled forms) is lost between calls.
|
||||
|
||||
This module provides a session pool that maintains persistent MCP sessions,
|
||||
scoped by ``(server_name, scope_key, owning_loop)``. Consecutive calls on
|
||||
scoped by ``(server_name, scope_key, owning_loop, ownership_domain)``. Consecutive calls on
|
||||
the same loop share server-side state; independent loops use separate sessions.
|
||||
The sync wrapper uses a fresh loop per call, so it does not preserve that state.
|
||||
Sessions are evicted in LRU order when the pool reaches capacity.
|
||||
@ -45,7 +45,7 @@ import asyncio
|
||||
import logging
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from typing import Any
|
||||
from typing import Any, Literal
|
||||
|
||||
import anyio
|
||||
from mcp import ClientSession
|
||||
@ -56,6 +56,8 @@ from deerflow.mcp_scope import mcp_scope_belongs_to_thread
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MCPPoolDomain = Literal["deployment", "personal"]
|
||||
|
||||
_MCP_CLOSED_STREAM_ERRORS = (
|
||||
anyio.ClosedResourceError,
|
||||
anyio.BrokenResourceError,
|
||||
@ -110,20 +112,22 @@ async def call_pooled_session_tool(
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any],
|
||||
call_kwargs: dict[str, Any],
|
||||
domain: MCPPoolDomain = "deployment",
|
||||
) -> Any:
|
||||
"""Call a pooled session and evict it only after an explicit disconnect."""
|
||||
try:
|
||||
return await session.call_tool(tool_name, arguments, **call_kwargs)
|
||||
except Exception as error:
|
||||
if _is_mcp_transport_disconnect(error):
|
||||
cleanup = asyncio.create_task(pool.close_session_if_current(server_name, scope_key, session))
|
||||
cleanup_call = pool.close_session_if_current(server_name, scope_key, session) if domain == "deployment" else pool.close_session_if_current(server_name, scope_key, session, domain=domain)
|
||||
cleanup = asyncio.create_task(cleanup_call)
|
||||
if await _finish_session_cleanup(cleanup, server_name):
|
||||
raise asyncio.CancelledError
|
||||
raise
|
||||
|
||||
|
||||
class MCPSessionPool:
|
||||
"""Manages persistent MCP sessions scoped by ``(server_name, scope_key, owning_loop)``."""
|
||||
"""Manages persistent MCP sessions scoped by ``(server_name, scope_key, owning_loop, ownership_domain)``."""
|
||||
|
||||
MAX_SESSIONS = 256
|
||||
SESSION_CLOSE_TIMEOUT = 5.0 # seconds to wait when closing a session on a foreign loop
|
||||
@ -131,7 +135,7 @@ class MCPSessionPool:
|
||||
def __init__(self) -> None:
|
||||
# Each entry: (session, owning_loop, owner_task, close_event).
|
||||
self._entries: OrderedDict[
|
||||
tuple[str, str, asyncio.AbstractEventLoop],
|
||||
tuple[str, str, asyncio.AbstractEventLoop, MCPPoolDomain],
|
||||
tuple[
|
||||
ClientSession,
|
||||
asyncio.AbstractEventLoop,
|
||||
@ -139,7 +143,7 @@ class MCPSessionPool:
|
||||
asyncio.Event,
|
||||
],
|
||||
] = OrderedDict()
|
||||
# In-flight creations, keyed by (server, scope, owning_loop). Lets concurrent callers
|
||||
# In-flight creations, keyed by (server, scope, owning_loop, ownership_domain). Lets concurrent callers
|
||||
# on the same loop share a single creation instead of each spawning a
|
||||
# duplicate session. Value: (loop, ready_future, owner_task, close_event).
|
||||
# The owner task promotes the record into ``_entries`` and resolves
|
||||
@ -147,7 +151,7 @@ class MCPSessionPool:
|
||||
# ``_run_session``), so ``ready`` resolving with a *result* always means
|
||||
# the session is registered — never merely handed to one caller.
|
||||
self._inflight: dict[
|
||||
tuple[str, str, asyncio.AbstractEventLoop],
|
||||
tuple[str, str, asyncio.AbstractEventLoop, MCPPoolDomain],
|
||||
tuple[
|
||||
asyncio.AbstractEventLoop,
|
||||
asyncio.Future[ClientSession],
|
||||
@ -169,7 +173,11 @@ class MCPSessionPool:
|
||||
# Session owner task
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _discard_owner(self, key: tuple[str, str, asyncio.AbstractEventLoop], owner: asyncio.Task[Any]) -> None:
|
||||
def _discard_owner(
|
||||
self,
|
||||
key: tuple[str, str, asyncio.AbstractEventLoop, MCPPoolDomain],
|
||||
owner: asyncio.Task[Any],
|
||||
) -> None:
|
||||
"""Retire only this owner, including after asyncio.run shuts its loop down."""
|
||||
with self._lock:
|
||||
entry = self._entries.get(key)
|
||||
@ -181,7 +189,7 @@ class MCPSessionPool:
|
||||
|
||||
async def _run_session(
|
||||
self,
|
||||
key: tuple[str, str, asyncio.AbstractEventLoop],
|
||||
key: tuple[str, str, asyncio.AbstractEventLoop, MCPPoolDomain],
|
||||
connection: dict[str, Any],
|
||||
ready: asyncio.Future[ClientSession],
|
||||
close_evt: asyncio.Event,
|
||||
@ -271,6 +279,8 @@ class MCPSessionPool:
|
||||
server_name: str,
|
||||
scope_key: str,
|
||||
connection: dict[str, Any],
|
||||
*,
|
||||
domain: MCPPoolDomain = "deployment",
|
||||
) -> ClientSession:
|
||||
"""Get or create a persistent MCP session.
|
||||
|
||||
@ -281,12 +291,14 @@ class MCPSessionPool:
|
||||
server_name: MCP server name.
|
||||
scope_key: Isolation key (typically thread_id).
|
||||
connection: Connection configuration for ``create_session``.
|
||||
domain: Ownership namespace for the connection. Deployment and
|
||||
personal resources with the same server/scope remain isolated.
|
||||
|
||||
Returns:
|
||||
An initialized ``ClientSession``.
|
||||
"""
|
||||
current_loop = asyncio.get_running_loop()
|
||||
key = (server_name, scope_key, current_loop)
|
||||
key = (server_name, scope_key, current_loop, domain)
|
||||
|
||||
# Phase 1: inspect/mutate the registry under the thread lock (no awaits).
|
||||
# Decide one of three outcomes atomically: return an existing session,
|
||||
@ -615,7 +627,7 @@ class MCPSessionPool:
|
||||
await self._shutdown_entry(loop, ent_task, ent_close, ready=ent_ready)
|
||||
|
||||
async def close_scope(self, scope_key: str) -> None:
|
||||
"""Close all sessions for a given scope (e.g. thread_id)."""
|
||||
"""Close all sessions for a scope across ownership domains."""
|
||||
with self._lock:
|
||||
keys = [k for k in self._entries if k[1] == scope_key]
|
||||
entries = [(self._entries.pop(k)) for k in keys]
|
||||
@ -624,7 +636,7 @@ class MCPSessionPool:
|
||||
await self._close_owners(entries, inflight)
|
||||
|
||||
async def close_thread_scope(self, *, user_id: str, thread_id: str) -> None:
|
||||
"""Close every session scoped to one user/thread identity.
|
||||
"""Close every session scoped to one user/thread identity across ownership domains.
|
||||
|
||||
Differs from :meth:`close_scope` in matching *all incarnations* of the
|
||||
thread rather than one exact scope key. A thread-deletion path knows the
|
||||
@ -649,12 +661,18 @@ class MCPSessionPool:
|
||||
inflight = [self._inflight.pop(k) for k in inflight_keys]
|
||||
await self._close_owners(entries, inflight)
|
||||
|
||||
async def close_session(self, server_name: str, scope_key: str) -> None:
|
||||
"""Close every session for this server/scope across all owning loops."""
|
||||
async def close_session(
|
||||
self,
|
||||
server_name: str,
|
||||
scope_key: str,
|
||||
*,
|
||||
domain: MCPPoolDomain = "deployment",
|
||||
) -> None:
|
||||
"""Close this ownership domain's session for a server/scope across loops."""
|
||||
with self._lock:
|
||||
keys = [k for k in self._entries if k[:2] == (server_name, scope_key)]
|
||||
keys = [k for k in self._entries if k[:2] == (server_name, scope_key) and k[3] == domain]
|
||||
entries = [self._entries.pop(k) for k in keys]
|
||||
keys = [k for k in self._inflight if k[:2] == (server_name, scope_key)]
|
||||
keys = [k for k in self._inflight if k[:2] == (server_name, scope_key) and k[3] == domain]
|
||||
inflight = [self._inflight.pop(k) for k in keys]
|
||||
await self._close_owners(entries, inflight)
|
||||
|
||||
@ -663,10 +681,15 @@ class MCPSessionPool:
|
||||
server_name: str,
|
||||
scope_key: str,
|
||||
session: ClientSession,
|
||||
*,
|
||||
domain: MCPPoolDomain = "deployment",
|
||||
) -> bool:
|
||||
"""Close *session* only if it is still the registered entry for the key."""
|
||||
"""Close *session* only if it is current in the requested ownership domain."""
|
||||
with self._lock:
|
||||
key = next((k for k, entry in self._entries.items() if k[:2] == (server_name, scope_key) and entry[0] is session), None)
|
||||
key = next(
|
||||
(k for k, entry in self._entries.items() if k[:2] == (server_name, scope_key) and k[3] == domain and entry[0] is session),
|
||||
None,
|
||||
)
|
||||
if key is None:
|
||||
return False
|
||||
entry = self._entries.pop(key)
|
||||
@ -674,12 +697,12 @@ class MCPSessionPool:
|
||||
await self._shutdown_entry(loop, task, close_evt)
|
||||
return True
|
||||
|
||||
async def close_server(self, server_name: str) -> None:
|
||||
"""Close all sessions for a given server."""
|
||||
async def close_server(self, server_name: str, *, domain: MCPPoolDomain = "deployment") -> None:
|
||||
"""Close one ownership domain's sessions for a given server."""
|
||||
with self._lock:
|
||||
keys = [k for k in self._entries if k[0] == server_name]
|
||||
keys = [k for k in self._entries if k[0] == server_name and k[3] == domain]
|
||||
entries = [(self._entries.pop(k)) for k in keys]
|
||||
inflight_keys = [k for k in self._inflight if k[0] == server_name]
|
||||
inflight_keys = [k for k in self._inflight if k[0] == server_name and k[3] == domain]
|
||||
inflight = [self._inflight.pop(k) for k in inflight_keys]
|
||||
await self._close_owners(entries, inflight)
|
||||
|
||||
|
||||
@ -20,7 +20,12 @@ from deerflow.mcp.headers import apply_header_overrides
|
||||
from deerflow.mcp.interceptors import build_mcp_tool_interceptors
|
||||
from deerflow.mcp.oauth import OAuthTokenManager, build_oauth_tool_interceptor
|
||||
from deerflow.mcp.personal_access import require_personal_mcp_access
|
||||
from deerflow.mcp.session_pool import MCPSessionPool, call_pooled_session_tool, get_session_pool
|
||||
from deerflow.mcp.session_pool import (
|
||||
MCPPoolDomain,
|
||||
MCPSessionPool,
|
||||
call_pooled_session_tool,
|
||||
get_session_pool,
|
||||
)
|
||||
from deerflow.mcp.user_config import PersonalMcpConfigSnapshot, load_user_mcp_config_if_changed
|
||||
from deerflow.mcp_scope import mcp_session_scope_key
|
||||
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
||||
@ -129,6 +134,7 @@ class McpTaskToolCaller:
|
||||
thread_id=thread_id,
|
||||
thread_incarnation=thread_incarnation,
|
||||
request_scoped_headers=request_scoped_headers,
|
||||
connection_scope="personal",
|
||||
)
|
||||
if connection_scope != "deployment":
|
||||
raise ValueError("Invalid MCP task connection scope")
|
||||
@ -140,6 +146,7 @@ class McpTaskToolCaller:
|
||||
thread_id=thread_id,
|
||||
thread_incarnation=thread_incarnation,
|
||||
request_scoped_headers=request_scoped_headers,
|
||||
connection_scope="deployment",
|
||||
)
|
||||
|
||||
def _personal_caller_for(self, user_id: str, server_name: str) -> McpTaskToolCaller:
|
||||
@ -169,6 +176,7 @@ class McpTaskToolCaller:
|
||||
thread_id: str,
|
||||
thread_incarnation: str | None,
|
||||
request_scoped_headers: bool,
|
||||
connection_scope: MCPPoolDomain,
|
||||
) -> Any:
|
||||
is_background_call = not request_scoped_headers
|
||||
interceptors = self._interceptors if is_background_call else self._submit_interceptors
|
||||
@ -192,10 +200,11 @@ class McpTaskToolCaller:
|
||||
)
|
||||
pool = get_session_pool()
|
||||
session_init_timeout = server_config.session_init_timeout
|
||||
session_request = pool.get_session(server_name, scope_key, connection) if connection_scope == "deployment" else pool.get_session(server_name, scope_key, connection, domain=connection_scope)
|
||||
if session_init_timeout is not None:
|
||||
try:
|
||||
session = await asyncio.wait_for(
|
||||
pool.get_session(server_name, scope_key, connection),
|
||||
session_request,
|
||||
timeout=session_init_timeout,
|
||||
)
|
||||
except TimeoutError:
|
||||
@ -206,7 +215,7 @@ class McpTaskToolCaller:
|
||||
)
|
||||
raise
|
||||
else:
|
||||
session = await pool.get_session(server_name, scope_key, connection)
|
||||
session = await session_request
|
||||
return await self._invoke(
|
||||
session=session,
|
||||
pool=pool,
|
||||
@ -220,6 +229,7 @@ class McpTaskToolCaller:
|
||||
session_init_timeout_seconds=None,
|
||||
persistent_session=True,
|
||||
interceptors=interceptors,
|
||||
connection_scope=connection_scope,
|
||||
)
|
||||
|
||||
authorization = await self._oauth_token_manager.get_authorization_header(server_name)
|
||||
@ -244,6 +254,7 @@ class McpTaskToolCaller:
|
||||
session_init_timeout_seconds=server_config.session_init_timeout,
|
||||
persistent_session=False,
|
||||
interceptors=interceptors,
|
||||
connection_scope=connection_scope,
|
||||
)
|
||||
|
||||
async def _invoke(
|
||||
@ -261,6 +272,7 @@ class McpTaskToolCaller:
|
||||
session_init_timeout_seconds: float | None,
|
||||
persistent_session: bool,
|
||||
interceptors: list[Any],
|
||||
connection_scope: MCPPoolDomain,
|
||||
) -> Any:
|
||||
from langchain_mcp_adapters.interceptors import MCPToolCallRequest
|
||||
from langchain_mcp_adapters.sessions import create_session
|
||||
@ -280,6 +292,7 @@ class McpTaskToolCaller:
|
||||
"Ignoring MCP interceptor headers with unsupported type: %s",
|
||||
type(request.headers).__name__,
|
||||
)
|
||||
domain_kwargs = {"domain": connection_scope} if connection_scope != "deployment" else {}
|
||||
return await call_pooled_session_tool(
|
||||
session,
|
||||
pool,
|
||||
@ -288,6 +301,7 @@ class McpTaskToolCaller:
|
||||
tool_name=request.name,
|
||||
arguments=request.args,
|
||||
call_kwargs=call_kwargs,
|
||||
**domain_kwargs,
|
||||
)
|
||||
|
||||
effective_connection = dict(connection)
|
||||
|
||||
@ -22,7 +22,11 @@ from deerflow.mcp.client import build_servers_config
|
||||
from deerflow.mcp.headers import apply_header_overrides
|
||||
from deerflow.mcp.interceptors import build_mcp_tool_interceptors, compose_tool_interceptors
|
||||
from deerflow.mcp.oauth import build_oauth_tool_interceptor, get_initial_oauth_headers
|
||||
from deerflow.mcp.session_pool import call_pooled_session_tool, get_session_pool
|
||||
from deerflow.mcp.session_pool import (
|
||||
MCPPoolDomain,
|
||||
call_pooled_session_tool,
|
||||
get_session_pool,
|
||||
)
|
||||
from deerflow.mcp.tasks import ORDINARY_MCP_TASK_DRIVER, TaskSubmitRequest
|
||||
from deerflow.mcp.tasks.runtime import (
|
||||
McpTaskConfigurationError,
|
||||
@ -539,11 +543,12 @@ def _make_session_pool_tool(
|
||||
tool_call_timeout: float | None = None,
|
||||
session_init_timeout: float | None = None,
|
||||
tool_name_prefix: bool = True,
|
||||
ownership_domain: MCPPoolDomain = "deployment",
|
||||
) -> BaseTool:
|
||||
"""Wrap an MCP tool so it reuses a persistent session from the pool.
|
||||
|
||||
Replaces the per-call session creation with pool-managed sessions scoped
|
||||
by ``(server_name, user/thread/incarnation)``. This ensures stateful MCP servers
|
||||
by ``(server_name, user/thread/incarnation, ownership_domain)``. This ensures stateful MCP servers
|
||||
(e.g. Playwright) keep their state across tool calls within the same thread
|
||||
while staying isolated per user.
|
||||
|
||||
@ -607,6 +612,16 @@ def _make_session_pool_tool(
|
||||
session_env.setdefault("TMP", str(tmp_dir))
|
||||
session_env.setdefault("TEMP", str(tmp_dir))
|
||||
session_connection["env"] = session_env
|
||||
session_request = (
|
||||
pool.get_session(server_name, scope_key, session_connection)
|
||||
if ownership_domain == "deployment"
|
||||
else pool.get_session(
|
||||
server_name,
|
||||
scope_key,
|
||||
session_connection,
|
||||
domain=ownership_domain,
|
||||
)
|
||||
)
|
||||
if session_init_timeout is not None:
|
||||
# Cancellation here is safe: MCPSessionPool.get_session owns the
|
||||
# teardown of a session stuck mid-creation (it signals close and
|
||||
@ -614,7 +629,7 @@ def _make_session_pool_tool(
|
||||
# so a hung server cannot leak a session or block the turn.
|
||||
try:
|
||||
session = await asyncio.wait_for(
|
||||
pool.get_session(server_name, scope_key, session_connection),
|
||||
session_request,
|
||||
timeout=session_init_timeout,
|
||||
)
|
||||
except TimeoutError:
|
||||
@ -629,7 +644,9 @@ def _make_session_pool_tool(
|
||||
)
|
||||
raise
|
||||
else:
|
||||
session = await pool.get_session(server_name, scope_key, session_connection)
|
||||
session = await session_request
|
||||
|
||||
domain_kwargs = {"domain": ownership_domain} if ownership_domain != "deployment" else {}
|
||||
|
||||
# Build common call_tool kwargs once — only add keys when needed so
|
||||
# existing call-sites that assert on exact arguments are not affected.
|
||||
@ -657,6 +674,7 @@ def _make_session_pool_tool(
|
||||
tool_name=request.name,
|
||||
arguments=request.args,
|
||||
call_kwargs=kwargs,
|
||||
**domain_kwargs,
|
||||
)
|
||||
|
||||
handler = compose_tool_interceptors(tool_interceptors, base_handler)
|
||||
@ -677,6 +695,7 @@ def _make_session_pool_tool(
|
||||
tool_name=original_name,
|
||||
arguments=arguments,
|
||||
call_kwargs=call_kwargs,
|
||||
**domain_kwargs,
|
||||
)
|
||||
|
||||
# The after-call snapshot diff only feeds bare-filename correlation in
|
||||
@ -1021,6 +1040,7 @@ async def get_mcp_tools(extensions_config: ExtensionsConfig | None = None, *, pe
|
||||
tool_call_timeout=_timeout,
|
||||
session_init_timeout=_init_timeout,
|
||||
tool_name_prefix=tool_name_prefix,
|
||||
ownership_domain=("personal" if personal_user_id is not None else "deployment"),
|
||||
)
|
||||
)
|
||||
else:
|
||||
|
||||
@ -141,6 +141,100 @@ async def test_get_session_reuses_existing():
|
||||
assert mock_cm.__aenter__.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_server_and_scope_are_isolated_by_ownership_domain():
|
||||
"""Deployment and personal resources must never share or retire each other's session."""
|
||||
pool = MCPSessionPool()
|
||||
cms = []
|
||||
|
||||
class Cm:
|
||||
def __init__(self):
|
||||
self.session = AsyncMock()
|
||||
self.closed = False
|
||||
|
||||
async def __aenter__(self):
|
||||
return self.session
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
self.closed = True
|
||||
return False
|
||||
|
||||
def make_cm(*_args, **_kwargs):
|
||||
cm = Cm()
|
||||
cms.append(cm)
|
||||
return cm
|
||||
|
||||
connection = {"transport": "stdio", "command": "x", "args": []}
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", side_effect=make_cm):
|
||||
deployment = await pool.get_session("server", "thread-1", connection)
|
||||
personal = await pool.get_session("server", "thread-1", connection, domain="personal")
|
||||
|
||||
assert deployment is not personal
|
||||
assert await pool.get_session("server", "thread-1", connection) is deployment
|
||||
assert await pool.get_session("server", "thread-1", connection, domain="personal") is personal
|
||||
|
||||
await pool.close_session("server", "thread-1")
|
||||
assert cms[0].closed is True
|
||||
assert cms[1].closed is False
|
||||
assert await pool.get_session("server", "thread-1", connection, domain="personal") is personal
|
||||
|
||||
await pool.close_all()
|
||||
|
||||
assert cms[1].closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_personal_disconnect_evicts_only_personal_domain():
|
||||
"""A dead personal session must not evict the same-named deployment session."""
|
||||
pool = MCPSessionPool()
|
||||
deployment = AsyncMock()
|
||||
personal = AsyncMock()
|
||||
deployment.initialize = AsyncMock()
|
||||
personal.initialize = AsyncMock()
|
||||
personal.call_tool = AsyncMock(side_effect=anyio.ClosedResourceError())
|
||||
closed = {"deployment": False, "personal": False}
|
||||
|
||||
class Cm:
|
||||
def __init__(self, session, name):
|
||||
self.session = session
|
||||
self.name = name
|
||||
|
||||
async def __aenter__(self):
|
||||
return self.session
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
closed[self.name] = True
|
||||
return False
|
||||
|
||||
connection = {"transport": "stdio", "command": "x", "args": []}
|
||||
with patch(
|
||||
"langchain_mcp_adapters.sessions.create_session",
|
||||
side_effect=[Cm(deployment, "deployment"), Cm(personal, "personal")],
|
||||
):
|
||||
assert await pool.get_session("server", "thread-1", connection) is deployment
|
||||
assert await pool.get_session("server", "thread-1", connection, domain="personal") is personal
|
||||
|
||||
with pytest.raises(anyio.ClosedResourceError):
|
||||
await call_pooled_session_tool(
|
||||
personal,
|
||||
pool,
|
||||
server_name="server",
|
||||
scope_key="thread-1",
|
||||
tool_name="probe",
|
||||
arguments={},
|
||||
call_kwargs={},
|
||||
domain="personal",
|
||||
)
|
||||
|
||||
assert closed == {"deployment": False, "personal": True}
|
||||
assert await pool.get_session("server", "thread-1", connection) is deployment
|
||||
assert ("server", "thread-1", asyncio.get_running_loop(), "personal") not in pool._entries
|
||||
|
||||
await pool.close_all()
|
||||
|
||||
assert closed == {"deployment": True, "personal": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_scope_creates_different_session():
|
||||
"""Different scope keys get different sessions."""
|
||||
@ -942,7 +1036,7 @@ async def test_late_disconnect_from_old_session_does_not_evict_replacement(tmp_p
|
||||
with pytest.raises(anyio.ClosedResourceError):
|
||||
await late_call
|
||||
|
||||
assert pool._entries[("srv", "test-user-autouse:default", asyncio.get_running_loop())][0] is replacement
|
||||
assert pool._entries[("srv", "test-user-autouse:default", asyncio.get_running_loop(), "deployment")][0] is replacement
|
||||
await pool.close_all()
|
||||
|
||||
|
||||
@ -2091,7 +2185,7 @@ async def test_get_session_cancelled_during_eviction_teardown_does_not_leak():
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
victim_task = asyncio.create_task(victim_owner())
|
||||
pool._entries[("victim", "scope", asyncio.get_running_loop())] = (MagicMock(), loop, victim_task, asyncio.Event())
|
||||
pool._entries[("victim", "scope", asyncio.get_running_loop(), "deployment")] = (MagicMock(), loop, victim_task, asyncio.Event())
|
||||
|
||||
gate = asyncio.Event()
|
||||
cms: list[_BlockingInitCm] = []
|
||||
@ -2170,7 +2264,7 @@ async def test_cancelled_creator_does_not_close_session_held_by_joiner():
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
victim_task = asyncio.create_task(victim_owner())
|
||||
pool._entries[("victim", "scope", asyncio.get_running_loop())] = (MagicMock(), loop, victim_task, asyncio.Event())
|
||||
pool._entries[("victim", "scope", asyncio.get_running_loop(), "deployment")] = (MagicMock(), loop, victim_task, asyncio.Event())
|
||||
|
||||
gate = asyncio.Event()
|
||||
cms: list[_BlockingInitCm] = []
|
||||
@ -2211,7 +2305,7 @@ async def test_cancelled_creator_does_not_close_session_held_by_joiner():
|
||||
await asyncio.sleep(0.02)
|
||||
assert not cms[0].closed, "cancelling the creator must not close a session already handed to a caller"
|
||||
assert ("srv", "scope-2") in {k[:2] for k in pool._entries}, "committed session must stay registered after creator cancellation"
|
||||
assert pool._entries[("srv", "scope-2", asyncio.get_running_loop())][0] is session, "the joiner's session must be the registered one"
|
||||
assert pool._entries[("srv", "scope-2", asyncio.get_running_loop(), "deployment")][0] is session, "the joiner's session must be the registered one"
|
||||
assert ("srv", "scope-2") not in {k[:2] for k in pool._inflight}
|
||||
|
||||
# Cleanup: close the pool so the parked owner finishes deterministically.
|
||||
@ -2380,8 +2474,8 @@ async def test_cancelling_close_scope_does_not_strand_other_removed_owners():
|
||||
|
||||
first_task = asyncio.create_task(first_owner())
|
||||
second_task = asyncio.create_task(second_owner())
|
||||
pool._entries[("s1", "scope", asyncio.get_running_loop())] = (MagicMock(), loop, first_task, first_close)
|
||||
pool._entries[("s2", "scope", asyncio.get_running_loop())] = (MagicMock(), loop, second_task, second_close)
|
||||
pool._entries[("s1", "scope", asyncio.get_running_loop(), "deployment")] = (MagicMock(), loop, first_task, first_close)
|
||||
pool._entries[("s2", "scope", asyncio.get_running_loop(), "deployment")] = (MagicMock(), loop, second_task, second_close)
|
||||
|
||||
closer = asyncio.create_task(pool.close_scope("scope"))
|
||||
# Let the closer reach its teardown await (both signals are already out).
|
||||
@ -2574,7 +2668,7 @@ async def test_queued_cancel_rechecks_failure_on_owning_loop():
|
||||
gate2 = holder["gate"]
|
||||
|
||||
proxy = _DelayedLoop(loop2)
|
||||
pool._inflight[("s", "t1", proxy)] = (proxy, holder["ready"], holder["task"], holder["close_evt"])
|
||||
pool._inflight[("s", "t1", proxy, "deployment")] = (proxy, holder["ready"], holder["task"], holder["close_evt"])
|
||||
|
||||
# Snapshot moment: the owner has NOT failed yet. close_scope's signal phase
|
||||
# queues the (guarded) cancellation into the holding proxy; its await phase
|
||||
@ -3002,7 +3096,7 @@ async def test_mcp_tools_routed_to_source_server_with_prefix_overlap():
|
||||
|
||||
routed: list[tuple[str, str]] = []
|
||||
|
||||
def fake_wrap(tool, server_name, connection, interceptors, tool_call_timeout=None, session_init_timeout=None, tool_name_prefix=True):
|
||||
def fake_wrap(tool, server_name, connection, interceptors, tool_call_timeout=None, session_init_timeout=None, tool_name_prefix=True, **_kwargs):
|
||||
routed.append((tool.name, server_name))
|
||||
return tool
|
||||
|
||||
|
||||
@ -92,7 +92,7 @@ def test_owner_completion_does_not_remove_a_replacement(loop_pool):
|
||||
|
||||
async def replace():
|
||||
await pool.get_session("s", "u:t", {})
|
||||
key = ("s", "u:t", asyncio.get_running_loop())
|
||||
key = ("s", "u:t", asyncio.get_running_loop(), "deployment")
|
||||
old_owner = pool._entries[key][2]
|
||||
await pool.close_session("s", "u:t")
|
||||
replacement = await pool.get_session("s", "u:t", {})
|
||||
@ -114,7 +114,7 @@ def test_abandoned_closed_loop_entry_can_be_retired(loop_pool, retirement):
|
||||
closed_loop.close()
|
||||
owner = MagicMock(spec=asyncio.Task)
|
||||
owner.done.return_value = False
|
||||
key = ("s", "u:t", closed_loop)
|
||||
key = ("s", "u:t", closed_loop, "deployment")
|
||||
pool._entries[key] = (MagicMock(), closed_loop, owner, asyncio.Event())
|
||||
pool.MAX_SESSIONS = 1
|
||||
if retirement == "explicit":
|
||||
|
||||
@ -144,6 +144,44 @@ async def test_stdio_task_call_reuses_exact_scope_and_raw_tool_name() -> None:
|
||||
pool.close_session.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_personal_stdio_task_call_uses_personal_pool_domain() -> 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)
|
||||
|
||||
caller = McpTaskToolCaller(ExtensionsConfig())
|
||||
personal_caller = McpTaskToolCaller(_config())
|
||||
|
||||
with (
|
||||
patch.object(caller, "_personal_caller_for", return_value=personal_caller),
|
||||
patch("deerflow.mcp.task_tool_caller.require_personal_mcp_access", new_callable=AsyncMock),
|
||||
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",
|
||||
connection_scope="personal",
|
||||
)
|
||||
|
||||
assert actual is result
|
||||
pool.get_session.assert_awaited_once_with(
|
||||
"reports",
|
||||
"user-1:thread-1",
|
||||
{"transport": "stdio", "command": "report-mcp"},
|
||||
domain="personal",
|
||||
)
|
||||
session.call_tool.assert_awaited_once_with("status_report", {"task_id": "remote-1"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.no_auto_user
|
||||
@pytest.mark.parametrize("has_ambient_user", [False, True], ids=["no-user", "existing-user"])
|
||||
|
||||
@ -137,6 +137,52 @@ async def test_mcp_tool_name_prefix_can_be_disabled_per_server_without_disabling
|
||||
assert wrap_tool.call_args.kwargs["tool_name_prefix"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_personal_discovery_routes_stdio_wrapper_to_personal_pool_domain() -> None:
|
||||
extensions_config = ExtensionsConfig.model_validate(
|
||||
{
|
||||
"mcpServers": {
|
||||
"github": {
|
||||
"type": "stdio",
|
||||
"command": "uvx",
|
||||
"args": ["github-mcp"],
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
servers_config = {
|
||||
"github": {
|
||||
"transport": "stdio",
|
||||
"command": "uvx",
|
||||
"args": ["github-mcp"],
|
||||
}
|
||||
}
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, _connections, *, callbacks=None, tool_interceptors=None, **_kwargs) -> None:
|
||||
self.callbacks = callbacks
|
||||
self.tool_interceptors = tool_interceptors or []
|
||||
|
||||
async def get_tools(self, *, server_name=None):
|
||||
assert server_name == "github"
|
||||
return [_tool("github_search")]
|
||||
|
||||
with (
|
||||
patch("deerflow.mcp.user_config.load_user_mcp_config", return_value=extensions_config),
|
||||
patch("deerflow.mcp.personal_access.authorized_personal_config", new_callable=AsyncMock, 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_mcp_tool_interceptors", return_value=[]),
|
||||
patch("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient),
|
||||
patch("deerflow.mcp.tools._make_session_pool_tool", side_effect=lambda tool, *_args, **_kwargs: tool) as wrap_tool,
|
||||
):
|
||||
tools = await get_mcp_tools(extensions_config, personal_user_id="alice")
|
||||
|
||||
assert [tool.name for tool in tools] == ["github_search"]
|
||||
wrap_tool.assert_called_once()
|
||||
assert wrap_tool.call_args.kwargs["ownership_domain"] == "personal"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unprefixed_stdio_tool_keeps_server_like_original_name(tmp_path: Path) -> None:
|
||||
"""Pooling must not strip a prefix that belongs to the MCP tool itself."""
|
||||
@ -163,3 +209,30 @@ async def test_unprefixed_stdio_tool_keeps_server_like_original_name(tmp_path: P
|
||||
await wrapped.coroutine(query="repositories")
|
||||
|
||||
mock_session.call_tool.assert_awaited_once_with("github_search", {"query": "repositories"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_personal_stdio_wrapper_uses_personal_pool_domain(tmp_path: Path) -> None:
|
||||
original_tool = _tool("github_search")
|
||||
mock_session = AsyncMock()
|
||||
mock_session.call_tool = AsyncMock(return_value=MagicMock(content=[], isError=False, structuredContent=None))
|
||||
mock_pool = MagicMock()
|
||||
mock_pool.get_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
with (
|
||||
patch("deerflow.mcp.tools.get_session_pool", return_value=mock_pool),
|
||||
patch("deerflow.mcp.tools.get_paths", return_value=MagicMock()),
|
||||
patch(
|
||||
"deerflow.mcp.tools._prepare_stdio_workspace",
|
||||
return_value=(tmp_path, tmp_path / "tmp", {}),
|
||||
),
|
||||
):
|
||||
wrapped = _make_session_pool_tool(
|
||||
original_tool,
|
||||
"github",
|
||||
{"transport": "stdio", "command": "mcp-server", "args": []},
|
||||
ownership_domain="personal",
|
||||
)
|
||||
await wrapped.coroutine(query="repositories")
|
||||
|
||||
assert mock_pool.get_session.await_args.kwargs["domain"] == "personal"
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user