mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
fix(mcp): share rotating OAuth state with durable task calls (#6342)
* fix(mcp): share rotating OAuth state with durable task calls * refactor(mcp): centralize task server selection --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
7c6fb15f8e
commit
b88d8db5fc
@ -758,6 +758,12 @@ OpenViking memory backend; it does not replace automatic turn capture or recall.
|
||||
|
||||
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. Input-required and terminal updates wake the current chat through idempotent Agent runs, while `list_background_tasks` and `cancel_background_task` let the Agent manage tasks without asking users for remote handles. Current-thread tasks are available through `GET /api/threads/{thread_id}/mcp-tasks`, its detail endpoint, and `POST /api/threads/{thread_id}/mcp-tasks/{task_id}/cancel`; when the task runtime actually starts, the Web UI exposes the same safe local view from the chat header with live status refresh, cancellation, and on-demand result, artifact, input-request, status-error, and cancellation-retry details. Default-disabled and memory-backend deployments hide that UI and do not poll the task endpoints. A failed remote cancellation remains queued with backoff, and its latest bounded error and attempt count stay visible in the expanded task card. 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. `input_required` is notification-only for now: DeerFlow can display the request but cannot yet submit the user's answer back to the remote task.
|
||||
|
||||
For deployment-level HTTP/SSE servers with `task_toolsets`, discovery, ordinary
|
||||
tools, and background task calls share OAuth token state within one Gateway
|
||||
process, including across tool-cache resets. Rotated refresh tokens stay in
|
||||
memory; they are not written back to configuration or shared across processes.
|
||||
After a restart, the configured refresh token must still be valid.
|
||||
|
||||
Notification launch and failed Agent-run deliveries use capped exponential backoff with a visible attempt count and stop after five failed attempts. When a bounded ordinary release exceeds its drain deadline, the service retains ownership until it settles. A permanently rejected target such as a deleted chat is dead-lettered immediately instead of retried forever or recreated. Cancellation endpoints return after durably recording the request; the background service owns the potentially slow remote MCP call and its retry schedule.
|
||||
|
||||
Notification runs keep their trusted delivery instruction separate from the framed, untrusted remote event payload. The process-started task runtime—not a hot config read—controls whether the task-management tools are exposed, so changing `mcp_tasks` requires a Gateway restart. When a skill's `allowed-tools` policy is active, `list_background_tasks` and `cancel_background_task` must be declared explicitly like other business tools.
|
||||
|
||||
@ -396,7 +396,19 @@ 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.
|
||||
For deployment-level HTTP/SSE servers with `task_toolsets`, discovery, ordinary
|
||||
tool calls, and background submit/status/cancel calls share the cached access
|
||||
token, rotated refresh token, and refresh lock for one Gateway process lifetime.
|
||||
Tool-cache resets and rediscovery keep that state. Token rotation does not
|
||||
modify the parsed configuration or trigger the startup-snapshot drift guard;
|
||||
real operator configuration changes still require a restart.
|
||||
|
||||
This state is process-local and is never written back to `extensions_config.json`
|
||||
or an environment variable. A restarted Gateway needs a valid configured
|
||||
refresh token, and separate Gateway workers do not coordinate token rotation.
|
||||
Personal MCP connections remain owner-scoped and do not use this deployment
|
||||
state, even if their runtime names match a deployment server.
|
||||
|
||||
When `user_auth` is enabled on an HTTP/SSE server, background status and
|
||||
cancellation calls use the persisted task owner's configured credential,
|
||||
including after a Gateway restart.
|
||||
|
||||
@ -3,6 +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.
|
||||
- **Task OAuth lifetime**: `tasks/runtime.py` owns shared OAuth state only for enabled deployment servers with `task_toolsets`, from `set_mcp_task_config_snapshot(config)` until its `None` teardown. Use `_task_enabled_servers` for both snapshot and OAuth selection. Validate the frozen snapshot before reusing that state for discovery, ordinary tools, or durable calls. Cache resets must not discard it. Keep tokens and the cross-loop refresh lock together; rotate a private config copy, never the parsed snapshot or config file. Non-task servers remain reloadable and personal callers use independent managers, even on a name collision. This is not cross-process or cross-restart credential storage.
|
||||
- **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.
|
||||
|
||||
@ -5,7 +5,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
@ -24,12 +24,22 @@ class _OAuthToken:
|
||||
expires_at: datetime
|
||||
|
||||
|
||||
@dataclass
|
||||
class _OAuthState:
|
||||
"""One connection's mutable tokens and cross-loop refresh lock."""
|
||||
|
||||
config: McpOAuthConfig
|
||||
token: _OAuthToken | None = None
|
||||
# Sync tool wrappers and the Gateway can refresh from different loops.
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
|
||||
|
||||
class OAuthTokenManager:
|
||||
"""Acquire/cache/refresh OAuth tokens for MCP servers."""
|
||||
|
||||
def __init__(self, oauth_by_server: dict[str, McpOAuthConfig]):
|
||||
self._oauth_by_server = oauth_by_server
|
||||
self._tokens: dict[str, _OAuthToken] = {}
|
||||
# Refresh-token rotation belongs to runtime state, not the parsed config
|
||||
# used by cache/snapshot validation or custom interceptor builders.
|
||||
# A plain threading.Lock, not asyncio.Lock: the embedded/TUI sync tool-call
|
||||
# path (DeerFlowClient.stream() -> LangGraph ToolNode._func -> a
|
||||
# ThreadPoolExecutor -> deerflow.tools.sync.make_sync_tool_wrapper's
|
||||
@ -40,32 +50,47 @@ class OAuthTokenManager:
|
||||
# either deadlocks silently or raises "bound to a different event loop".
|
||||
# threading.Lock has no loop affinity, so it is safe to share across
|
||||
# however many event loops/threads call into the same server's lock.
|
||||
self._locks: dict[str, threading.Lock] = {name: threading.Lock() for name in oauth_by_server}
|
||||
self._states = {name: _OAuthState(config=oauth.model_copy(deep=True)) for name, oauth in oauth_by_server.items()}
|
||||
|
||||
@classmethod
|
||||
def from_extensions_config(cls, extensions_config: ExtensionsConfig) -> OAuthTokenManager:
|
||||
def from_extensions_config(
|
||||
cls,
|
||||
extensions_config: ExtensionsConfig,
|
||||
*,
|
||||
shared_manager: OAuthTokenManager | None = None,
|
||||
) -> OAuthTokenManager:
|
||||
"""Build a manager, optionally reusing already-validated connection state.
|
||||
|
||||
The durable runtime supplies ``shared_manager`` only after validating
|
||||
the deployment snapshot. Personal connections must never use it.
|
||||
"""
|
||||
oauth_by_server: dict[str, McpOAuthConfig] = {}
|
||||
for server_name, server_config in extensions_config.get_enabled_mcp_servers().items():
|
||||
if server_config.oauth and server_config.oauth.enabled:
|
||||
oauth_by_server[server_name] = server_config.oauth
|
||||
return cls(oauth_by_server)
|
||||
manager = cls(oauth_by_server)
|
||||
if shared_manager is not None:
|
||||
for name in manager._states.keys() & shared_manager._states.keys():
|
||||
manager._states[name] = shared_manager._states[name]
|
||||
return manager
|
||||
|
||||
def has_oauth_servers(self) -> bool:
|
||||
return bool(self._oauth_by_server)
|
||||
return bool(self._states)
|
||||
|
||||
def oauth_server_names(self) -> list[str]:
|
||||
return list(self._oauth_by_server.keys())
|
||||
return list(self._states)
|
||||
|
||||
async def get_authorization_header(self, server_name: str) -> str | None:
|
||||
oauth = self._oauth_by_server.get(server_name)
|
||||
if not oauth:
|
||||
state = self._states.get(server_name)
|
||||
if state is None:
|
||||
return None
|
||||
|
||||
token = self._tokens.get(server_name)
|
||||
oauth = state.config
|
||||
token = state.token
|
||||
if token and not self._is_expiring(token, oauth):
|
||||
return self._authorization_value(token, server_name)
|
||||
|
||||
lock = self._locks[server_name]
|
||||
lock = state.lock
|
||||
# Acquire the OS-level lock off-thread so a blocking wait never blocks this
|
||||
# event loop, then release it synchronously (release() never blocks). This
|
||||
# keeps the de-duplication behavior of the old `async with lock:` (only one
|
||||
@ -100,12 +125,12 @@ class OAuthTokenManager:
|
||||
lock.release()
|
||||
raise
|
||||
try:
|
||||
token = self._tokens.get(server_name)
|
||||
token = state.token
|
||||
if token and not self._is_expiring(token, oauth):
|
||||
return self._authorization_value(token, server_name)
|
||||
|
||||
fresh = await self._fetch_token(oauth)
|
||||
self._tokens[server_name] = fresh
|
||||
state.token = fresh
|
||||
logger.info(f"Refreshed OAuth access token for MCP server: {server_name}")
|
||||
return self._authorization_value(fresh, server_name)
|
||||
finally:
|
||||
@ -188,8 +213,8 @@ class OAuthTokenManager:
|
||||
raise ValueError(f"OAuth token response missing '{oauth.token_field}'")
|
||||
|
||||
# Persist a rotated refresh_token so subsequent refreshes use the latest
|
||||
# value. This is an in-process update only — it is intentionally NOT
|
||||
# written back to extensions_config.json. Providers that rotate refresh
|
||||
# value. This updates our private runtime copy only — neither the parsed
|
||||
# extensions config nor extensions_config.json. Providers that rotate refresh
|
||||
# tokens (Auth0, Okta, Google, etc.) return a new refresh_token on each
|
||||
# refresh; discarding it makes the next refresh fail with invalid_grant.
|
||||
if oauth.grant_type == "refresh_token":
|
||||
@ -239,9 +264,13 @@ def build_oauth_tool_interceptor(
|
||||
return oauth_interceptor
|
||||
|
||||
|
||||
async def get_initial_oauth_headers(extensions_config: ExtensionsConfig) -> dict[str, str]:
|
||||
async def get_initial_oauth_headers(
|
||||
extensions_config: ExtensionsConfig,
|
||||
*,
|
||||
token_manager: OAuthTokenManager | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""Get initial OAuth Authorization headers for MCP server connections."""
|
||||
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 {}
|
||||
|
||||
|
||||
@ -26,6 +26,7 @@ from deerflow.mcp.session_pool import (
|
||||
call_pooled_session_tool,
|
||||
get_session_pool,
|
||||
)
|
||||
from deerflow.mcp.tasks.runtime import get_mcp_task_oauth_token_manager
|
||||
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
|
||||
@ -77,7 +78,7 @@ class McpTaskToolCaller:
|
||||
self._extensions_config = extensions_config
|
||||
self._personal_callers: OrderedDict[str, tuple[PersonalMcpConfigSnapshot, McpTaskToolCaller | None]] = OrderedDict()
|
||||
self._personal_callers_lock = threading.Lock()
|
||||
self._oauth_token_manager = oauth_token_manager or OAuthTokenManager.from_extensions_config(extensions_config)
|
||||
self._oauth_token_manager = oauth_token_manager or get_mcp_task_oauth_token_manager(extensions_config)
|
||||
context_headers_interceptor = build_context_headers_interceptor(extensions_config)
|
||||
# Built once so the two chains keep an identical interceptor order and a
|
||||
# custom ``mcpInterceptors`` builder is invoked exactly once.
|
||||
@ -162,7 +163,9 @@ class McpTaskToolCaller:
|
||||
if server is None or not server.enabled:
|
||||
raise LookupError("Personal MCP task connection is missing, disabled or changed")
|
||||
if caller is None:
|
||||
caller = McpTaskToolCaller(snapshot.config)
|
||||
# A personal name can equal a deployment name. Never consult
|
||||
# the deployment OAuth state or startup snapshot for this path.
|
||||
caller = McpTaskToolCaller(snapshot.config, oauth_token_manager=OAuthTokenManager.from_extensions_config(snapshot.config))
|
||||
self._personal_callers[user_id] = (snapshot, caller)
|
||||
return caller
|
||||
|
||||
|
||||
@ -4,8 +4,9 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any, Protocol
|
||||
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig
|
||||
from deerflow.mcp.config_normalization import normalize_mcp_interceptor_paths, normalize_mcp_server_config
|
||||
from deerflow.mcp.oauth import OAuthTokenManager
|
||||
from deerflow.mcp.tasks.models import TaskSubmitRequest
|
||||
|
||||
|
||||
@ -45,13 +46,17 @@ class McpTaskSubmitter(Protocol):
|
||||
_submitter: McpTaskSubmitter | None = None
|
||||
_TaskServerConfigSnapshot = tuple[dict[str, dict[str, Any]], Any]
|
||||
_task_server_config_snapshot: _TaskServerConfigSnapshot | None = None
|
||||
_task_oauth_token_manager: OAuthTokenManager | None = None
|
||||
|
||||
|
||||
def _task_enabled_servers(extensions_config: ExtensionsConfig) -> dict[str, McpServerConfig]:
|
||||
"""Select the deployment connections frozen and shared by the task runtime."""
|
||||
return {name: server for name, server in extensions_config.get_enabled_mcp_servers().items() if server.task_toolsets}
|
||||
|
||||
|
||||
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
|
||||
for server_name, server in _task_enabled_servers(extensions_config).items():
|
||||
servers[server_name] = normalize_mcp_server_config(server, task_runtime=True)
|
||||
raw_interceptors = (extensions_config.model_extra or {}).get("mcpInterceptors")
|
||||
interceptors = normalize_mcp_interceptor_paths(raw_interceptors) if servers else None
|
||||
@ -60,8 +65,15 @@ def _task_server_configs(extensions_config: ExtensionsConfig) -> _TaskServerConf
|
||||
|
||||
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
|
||||
global _task_server_config_snapshot, _task_oauth_token_manager
|
||||
_task_server_config_snapshot = None if extensions_config is None else _task_server_configs(extensions_config)
|
||||
# Keep rotating credentials for task-enabled deployment servers for this
|
||||
# Gateway lifetime, including tool-cache resets. Do not retain OAuth state
|
||||
# for hot-reloadable non-task servers or owner-scoped personal connections.
|
||||
_task_oauth_token_manager = None
|
||||
if extensions_config is not None:
|
||||
task_servers = _task_enabled_servers(extensions_config)
|
||||
_task_oauth_token_manager = OAuthTokenManager.from_extensions_config(extensions_config.model_copy(update={"mcp_servers": task_servers}))
|
||||
|
||||
|
||||
def validate_mcp_task_config_snapshot(extensions_config: ExtensionsConfig) -> None:
|
||||
@ -80,6 +92,12 @@ def validate_mcp_task_config_snapshot(extensions_config: ExtensionsConfig) -> No
|
||||
raise McpTaskConfigurationError(f"MCP task-enabled server configuration changed after Gateway startup ({names}); restart DeerFlow before using durable task tools")
|
||||
|
||||
|
||||
def get_mcp_task_oauth_token_manager(extensions_config: ExtensionsConfig) -> OAuthTokenManager:
|
||||
"""Share OAuth state only for the matching frozen deployment connections."""
|
||||
validate_mcp_task_config_snapshot(extensions_config)
|
||||
return OAuthTokenManager.from_extensions_config(extensions_config, shared_manager=_task_oauth_token_manager)
|
||||
|
||||
|
||||
def set_mcp_task_submitter(submitter: McpTaskSubmitter | None) -> None:
|
||||
"""Install or clear the Gateway-owned submit boundary for this process."""
|
||||
global _submitter
|
||||
|
||||
@ -21,7 +21,7 @@ from deerflow.constants import DEFAULT_MCP_SESSION_INIT_TIMEOUT, MCP_TMP_SUBDIR
|
||||
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.oauth import OAuthTokenManager, build_oauth_tool_interceptor, get_initial_oauth_headers
|
||||
from deerflow.mcp.session_pool import (
|
||||
MCPPoolDomain,
|
||||
call_pooled_session_tool,
|
||||
@ -30,6 +30,7 @@ from deerflow.mcp.session_pool import (
|
||||
from deerflow.mcp.tasks import ORDINARY_MCP_TASK_DRIVER, TaskSubmitRequest
|
||||
from deerflow.mcp.tasks.runtime import (
|
||||
McpTaskConfigurationError,
|
||||
get_mcp_task_oauth_token_manager,
|
||||
get_mcp_task_submitter,
|
||||
validate_mcp_task_config_snapshot,
|
||||
)
|
||||
@ -911,7 +912,8 @@ async def get_mcp_tools(extensions_config: ExtensionsConfig | None = None, *, pe
|
||||
logger.info(f"Initializing MCP client with {len(servers_config)} server(s)")
|
||||
|
||||
# Inject initial OAuth headers for server connections (tool discovery/session init)
|
||||
initial_oauth_headers = await get_initial_oauth_headers(extensions_config)
|
||||
oauth_token_manager = OAuthTokenManager.from_extensions_config(extensions_config) if personal_user_id is not None else get_mcp_task_oauth_token_manager(extensions_config)
|
||||
initial_oauth_headers = await get_initial_oauth_headers(extensions_config, token_manager=oauth_token_manager)
|
||||
for server_name, auth_header in initial_oauth_headers.items():
|
||||
if server_name not in servers_config:
|
||||
continue
|
||||
@ -925,7 +927,7 @@ async def get_mcp_tools(extensions_config: ExtensionsConfig | None = None, *, pe
|
||||
|
||||
tool_interceptors = build_mcp_tool_interceptors(
|
||||
extensions_config,
|
||||
oauth_builder=build_oauth_tool_interceptor,
|
||||
oauth_builder=lambda config: build_oauth_tool_interceptor(config, token_manager=oauth_token_manager),
|
||||
resolver=resolve_variable,
|
||||
target_logger=logger,
|
||||
)
|
||||
|
||||
@ -285,7 +285,7 @@ def test_get_initial_oauth_headers_one_failing_server_does_not_drop_others(monke
|
||||
|
||||
def _client_factory(**kwargs):
|
||||
# The first call is for the failing server, second for the healthy one,
|
||||
# because OAuthTokenManager iterates _oauth_by_server in dict order
|
||||
# because OAuthTokenManager iterates the configured servers in dict order
|
||||
# ('broken-http' < 'secure-http').
|
||||
if not hasattr(_client_factory, "_count"):
|
||||
_client_factory._count = 0 # type: ignore[attr-defined]
|
||||
@ -561,7 +561,7 @@ def test_get_authorization_header_cancelled_while_waiting_does_not_leak_lock(mon
|
||||
)
|
||||
|
||||
manager = OAuthTokenManager.from_extensions_config(config)
|
||||
lock = manager._locks["secure-http"]
|
||||
lock = manager._states["secure-http"].lock
|
||||
|
||||
async def scenario() -> None:
|
||||
# Simulate another in-flight caller already holding the per-server lock
|
||||
|
||||
431
backend/tests/test_mcp_task_oauth.py
Normal file
431
backend/tests/test_mcp_task_oauth.py
Normal file
@ -0,0 +1,431 @@
|
||||
"""Rotating OAuth credentials across discovery and durable deployment tasks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
from urllib.parse import parse_qs
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage
|
||||
from langgraph.graph import END, START, MessagesState, StateGraph
|
||||
from langgraph.prebuilt import ToolNode
|
||||
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
from deerflow.config.paths import Paths
|
||||
from deerflow.mcp.cache import reset_mcp_tools_cache
|
||||
from deerflow.mcp.oauth import get_initial_oauth_headers
|
||||
from deerflow.mcp.task_tool_caller import McpTaskToolCaller
|
||||
from deerflow.mcp.tasks import OrdinaryMcpTaskDriver, TaskReference, TaskStatus
|
||||
from deerflow.mcp.tasks import runtime as task_runtime
|
||||
from deerflow.mcp.tasks.runtime import (
|
||||
McpTaskConfigurationError,
|
||||
set_mcp_task_config_snapshot,
|
||||
set_mcp_task_submitter,
|
||||
validate_mcp_task_config_snapshot,
|
||||
)
|
||||
from deerflow.mcp.tools import get_mcp_tools
|
||||
from deerflow.mcp.user_config import load_user_mcp_config, personal_server_name, user_mcp_config_path
|
||||
|
||||
|
||||
class _Clock:
|
||||
current = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
|
||||
@classmethod
|
||||
def now(cls, _tz):
|
||||
return cls.current
|
||||
|
||||
|
||||
class _RotatingServer:
|
||||
"""Replace only HTTP transport; run real MCP client/adapters and OAuth code."""
|
||||
|
||||
def __init__(self):
|
||||
self.refresh_token = "refresh-0"
|
||||
self.refresh_requests: list[str] = []
|
||||
self.generation = 0
|
||||
self.calls: list[tuple[str, str]] = []
|
||||
|
||||
def handle(self, request: httpx.Request) -> httpx.Response:
|
||||
if request.url.path == "/token":
|
||||
submitted = parse_qs(request.content.decode())["refresh_token"][0]
|
||||
self.refresh_requests.append(submitted)
|
||||
if submitted != self.refresh_token:
|
||||
return httpx.Response(400, json={"error": "invalid_grant"})
|
||||
self.generation += 1
|
||||
self.refresh_token = f"refresh-{self.generation}"
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={"access_token": f"access-{self.generation}", "refresh_token": self.refresh_token, "expires_in": 120},
|
||||
)
|
||||
|
||||
assert request.url.path == "/mcp"
|
||||
authorization = request.headers.get("Authorization")
|
||||
if authorization != f"Bearer access-{self.generation}":
|
||||
return httpx.Response(401)
|
||||
if request.method != "POST":
|
||||
return httpx.Response(405)
|
||||
message = json.loads(request.content)
|
||||
method = message["method"]
|
||||
if "id" not in message:
|
||||
return httpx.Response(202)
|
||||
if method == "initialize":
|
||||
result = {
|
||||
"protocolVersion": message["params"]["protocolVersion"],
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "reports", "version": "1.0"},
|
||||
}
|
||||
elif method == "tools/list":
|
||||
result = {"tools": [{"name": name, "description": name, "inputSchema": {"type": "object", "properties": {}}} for name in ("submit_report", "status_report", "cancel_report", "search")]}
|
||||
elif method == "tools/call":
|
||||
name = message["params"]["name"]
|
||||
self.calls.append((name, authorization))
|
||||
result = {
|
||||
"isError": False,
|
||||
"content": [{"type": "text", "text": "ok"}],
|
||||
"structuredContent": {"task_id": "remote-1", "status": "cancelled" if name == "cancel_report" else "running"},
|
||||
}
|
||||
else:
|
||||
raise AssertionError(method)
|
||||
return httpx.Response(200, json={"jsonrpc": "2.0", "id": message["id"], "result": result})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rotating_server(monkeypatch):
|
||||
server = _RotatingServer()
|
||||
original_client = httpx.AsyncClient
|
||||
|
||||
class OfflineClient(original_client):
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs["transport"] = httpx.MockTransport(server.handle)
|
||||
kwargs["trust_env"] = False
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", OfflineClient)
|
||||
monkeypatch.setattr("deerflow.mcp.oauth.datetime", _Clock)
|
||||
monkeypatch.setattr(_Clock, "current", datetime(2026, 1, 1, tzinfo=UTC))
|
||||
yield server
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def task_config(tmp_path, monkeypatch):
|
||||
path = tmp_path / "extensions_config.json"
|
||||
path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"mcpServers": {
|
||||
"reports": {
|
||||
"type": "http",
|
||||
"url": "https://reports.example.invalid/mcp",
|
||||
"oauth": {
|
||||
"token_url": "https://reports.example.invalid/token",
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": "refresh-0",
|
||||
"refresh_skew_seconds": 0,
|
||||
},
|
||||
"task_toolsets": [{"name": "reports", "submit_tool": "submit_report", "status_tool": "status_report", "cancel_tool": "cancel_report"}],
|
||||
}
|
||||
}
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(path))
|
||||
startup = ExtensionsConfig.from_file()
|
||||
set_mcp_task_config_snapshot(startup)
|
||||
try:
|
||||
yield path, startup
|
||||
finally:
|
||||
set_mcp_task_submitter(None)
|
||||
set_mcp_task_config_snapshot(None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discovery_submit_poll_and_rediscovery_share_rotating_oauth(task_config, rotating_server):
|
||||
path, startup = task_config
|
||||
initial_file = path.read_bytes()
|
||||
initial_config = startup.model_dump()
|
||||
driver = OrdinaryMcpTaskDriver(McpTaskToolCaller(startup))
|
||||
captured_requests = []
|
||||
|
||||
# The Agent-visible wrapper reaches the real driver/caller. Persistence is
|
||||
# outside this test's scope; the submitter only records the request/handle.
|
||||
async def submit(*, driver_name, request):
|
||||
captured_requests.append(request)
|
||||
submission = await driver.submit(request)
|
||||
assert submission.remote_task_id == "remote-1"
|
||||
return {"id": "local-1", "status": submission.snapshot.status.value}
|
||||
|
||||
set_mcp_task_submitter(SimpleNamespace(submit=submit))
|
||||
tools = await get_mcp_tools()
|
||||
assert [tool.name for tool in tools] == ["reports_submit_report", "reports_search"]
|
||||
graph = StateGraph(MessagesState, context_schema=dict)
|
||||
graph.add_node("tools", ToolNode(tools, handle_tool_errors=False))
|
||||
graph.add_edge(START, "tools")
|
||||
graph.add_edge("tools", END)
|
||||
result = await graph.compile().ainvoke(
|
||||
{"messages": [AIMessage(content="", tool_calls=[{"name": "reports_submit_report", "args": {}, "id": "call-1", "type": "tool_call"}])]},
|
||||
context={"user_id": "user-1", "thread_id": "thread-1", "thread_incarnation": "incarnation-1"},
|
||||
)
|
||||
assert json.loads(result["messages"][-1].content)["task_id"] == "local-1"
|
||||
assert len(captured_requests) == 1
|
||||
reference = TaskReference(
|
||||
local_task_id="local-1",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
server_name="reports",
|
||||
remote_task_id="remote-1",
|
||||
driver_data=captured_requests[0].driver_data,
|
||||
)
|
||||
|
||||
_Clock.current += timedelta(seconds=121)
|
||||
assert (await driver.get_status(reference)).status == TaskStatus.WORKING
|
||||
# Ordinary tools on the same task-enabled server must follow the poller's
|
||||
# refresh too, including tools rebuilt from the unchanged on-disk seed.
|
||||
assert await tools[1].ainvoke({})
|
||||
reset_mcp_tools_cache()
|
||||
rediscovered = await get_mcp_tools()
|
||||
assert [tool.name for tool in rediscovered] == ["reports_submit_report", "reports_search"]
|
||||
assert await rediscovered[1].ainvoke({})
|
||||
assert (await driver.cancel(reference)).status == TaskStatus.CANCELLED
|
||||
|
||||
assert rotating_server.refresh_requests == ["refresh-0", "refresh-1"]
|
||||
assert rotating_server.calls == [
|
||||
("submit_report", "Bearer access-1"),
|
||||
("status_report", "Bearer access-2"),
|
||||
("search", "Bearer access-2"),
|
||||
("search", "Bearer access-2"),
|
||||
("cancel_report", "Bearer access-2"),
|
||||
]
|
||||
assert startup.model_dump() == initial_config
|
||||
assert path.read_bytes() == initial_file
|
||||
validate_mcp_task_config_snapshot(startup)
|
||||
validate_mcp_task_config_snapshot(ExtensionsConfig.from_file())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rotation_does_not_mutate_loaded_config(task_config, rotating_server):
|
||||
_path, startup = task_config
|
||||
before = startup.model_dump()
|
||||
caller = McpTaskToolCaller(startup)
|
||||
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 startup.model_dump() == before
|
||||
validate_mcp_task_config_snapshot(startup)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("field", ["url", "refresh_token", "client_secret", "token_url"])
|
||||
async def test_changed_task_credentials_still_require_restart(task_config, rotating_server, field):
|
||||
_path, startup = task_config
|
||||
await get_mcp_tools()
|
||||
changed = ExtensionsConfig.from_file()
|
||||
if field == "url":
|
||||
changed.mcp_servers["reports"].url = "https://other.example.invalid/mcp"
|
||||
else:
|
||||
setattr(changed.mcp_servers["reports"].oauth, field, "changed")
|
||||
with pytest.raises(McpTaskConfigurationError, match="restart DeerFlow"):
|
||||
await get_mcp_tools(changed)
|
||||
assert rotating_server.refresh_requests == ["refresh-0"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("transport", ["http", "sse"])
|
||||
async def test_task_transports_share_discovery_token(task_config, rotating_server, monkeypatch, transport):
|
||||
_path, startup = task_config
|
||||
startup.mcp_servers["reports"].type = transport
|
||||
set_mcp_task_config_snapshot(startup)
|
||||
caller = McpTaskToolCaller(startup)
|
||||
discovery = task_runtime.get_mcp_task_oauth_token_manager(startup.model_copy(deep=True))
|
||||
assert await get_initial_oauth_headers(startup, token_manager=discovery) == {"reports": "Bearer access-1"}
|
||||
connections = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def create_session(connection):
|
||||
connections.append(connection)
|
||||
yield SimpleNamespace(initialize=AsyncMock(), call_tool=AsyncMock(return_value="ok"))
|
||||
|
||||
monkeypatch.setattr("langchain_mcp_adapters.sessions.create_session", create_session)
|
||||
assert await caller.call_tool(server_name="reports", tool_name="status_report", arguments={}, user_id="user-1", thread_id="thread-1") == "ok"
|
||||
assert connections[0]["transport"] == transport
|
||||
assert connections[0]["headers"] == {"Authorization": "Bearer access-1"}
|
||||
assert rotating_server.refresh_requests == ["refresh-0"]
|
||||
|
||||
|
||||
def test_separate_managers_deduplicate_refresh_across_event_loops(task_config, rotating_server, monkeypatch):
|
||||
managers = [task_runtime.get_mcp_task_oauth_token_manager(ExtensionsConfig.from_file()) for _ in range(3)]
|
||||
entered = threading.Event()
|
||||
all_acquiring = threading.Event()
|
||||
release = threading.Event()
|
||||
start = threading.Barrier(3)
|
||||
handle = rotating_server.handle
|
||||
state = managers[0]._states["reports"]
|
||||
original_lock = state.lock
|
||||
arrivals_lock = threading.Lock()
|
||||
arrivals = 0
|
||||
|
||||
class ObservedLock:
|
||||
def acquire(self):
|
||||
nonlocal arrivals
|
||||
with arrivals_lock:
|
||||
arrivals += 1
|
||||
if arrivals == len(managers):
|
||||
all_acquiring.set()
|
||||
return original_lock.acquire()
|
||||
|
||||
def release(self):
|
||||
original_lock.release()
|
||||
|
||||
monkeypatch.setattr(state, "lock", ObservedLock())
|
||||
|
||||
async def held_response(request):
|
||||
if request.url.path == "/token":
|
||||
entered.set()
|
||||
assert await asyncio.to_thread(release.wait, 5), "test did not release token response"
|
||||
return handle(request)
|
||||
|
||||
monkeypatch.setattr(rotating_server, "handle", held_response)
|
||||
|
||||
def fetch(manager):
|
||||
start.wait(timeout=5)
|
||||
return asyncio.run(manager.get_authorization_header("reports"))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=3) as executor:
|
||||
futures = [executor.submit(fetch, manager) for manager in managers]
|
||||
try:
|
||||
assert entered.wait(5), "token request did not start"
|
||||
assert all_acquiring.wait(5), "callers did not contend on the shared refresh lock"
|
||||
finally:
|
||||
release.set()
|
||||
assert [future.result(timeout=5) for future in futures] == ["Bearer access-1"] * 3
|
||||
assert rotating_server.refresh_requests == ["refresh-0"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_waiter_does_not_strand_shared_refresh_lock(task_config, rotating_server, monkeypatch):
|
||||
first = task_runtime.get_mcp_task_oauth_token_manager(ExtensionsConfig.from_file())
|
||||
second = task_runtime.get_mcp_task_oauth_token_manager(ExtensionsConfig.from_file())
|
||||
state = first._states["reports"]
|
||||
original_lock = state.lock
|
||||
entered = threading.Event()
|
||||
|
||||
class ObservedLock:
|
||||
def acquire(self):
|
||||
entered.set()
|
||||
return original_lock.acquire()
|
||||
|
||||
def release(self):
|
||||
original_lock.release()
|
||||
|
||||
monkeypatch.setattr(state, "lock", ObservedLock())
|
||||
original_lock.acquire()
|
||||
waiter = asyncio.create_task(first.get_authorization_header("reports"))
|
||||
try:
|
||||
assert await asyncio.to_thread(entered.wait, 5), "waiter never reached lock acquisition"
|
||||
waiter.cancel()
|
||||
finally:
|
||||
original_lock.release()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(waiter, 5)
|
||||
following = asyncio.create_task(second.get_authorization_header("reports"))
|
||||
try:
|
||||
# asyncio.wait has a bounded return even if cancellation cleanup must
|
||||
# wait for an executor thread stuck acquiring a leaked lock.
|
||||
done, _pending = await asyncio.wait({following}, timeout=5)
|
||||
assert following in done, "cancelled waiter leaked the shared refresh lock"
|
||||
assert following.result() == "Bearer access-1"
|
||||
finally:
|
||||
if original_lock.locked():
|
||||
original_lock.release()
|
||||
await asyncio.gather(following, return_exceptions=True)
|
||||
assert rotating_server.refresh_requests == ["refresh-0"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_task_runtime_reset_drops_rotated_oauth_state(task_config, rotating_server):
|
||||
_path, startup = task_config
|
||||
first = task_runtime.get_mcp_task_oauth_token_manager(ExtensionsConfig.from_file())
|
||||
assert await first.get_authorization_header("reports") == "Bearer access-1"
|
||||
set_mcp_task_config_snapshot(None)
|
||||
# A new lifetime needs a valid operator-supplied seed; retained in-process
|
||||
# credentials must not silently survive teardown or override that seed.
|
||||
startup.mcp_servers["reports"].oauth.refresh_token = "replacement-seed"
|
||||
rotating_server.refresh_token = "replacement-seed"
|
||||
set_mcp_task_config_snapshot(startup)
|
||||
second = task_runtime.get_mcp_task_oauth_token_manager(startup.model_copy(deep=True))
|
||||
assert await second.get_authorization_header("reports") == "Bearer access-2"
|
||||
assert rotating_server.refresh_requests == ["refresh-0", "replacement-seed"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hot_reloadable_servers_do_not_reuse_task_oauth_state(task_config, rotating_server):
|
||||
_path, startup = task_config
|
||||
other = startup.mcp_servers["reports"].model_copy(deep=True)
|
||||
other.task_toolsets = []
|
||||
startup.mcp_servers["other"] = other
|
||||
set_mcp_task_config_snapshot(startup)
|
||||
first = task_runtime.get_mcp_task_oauth_token_manager(startup)
|
||||
assert await first.get_authorization_header("other") == "Bearer access-1"
|
||||
changed = startup.model_copy(deep=True)
|
||||
changed.mcp_servers["other"].oauth.refresh_token = "replacement-seed"
|
||||
rotating_server.refresh_token = "replacement-seed"
|
||||
second = task_runtime.get_mcp_task_oauth_token_manager(changed)
|
||||
assert await second.get_authorization_header("other") == "Bearer access-2"
|
||||
assert rotating_server.refresh_requests == ["refresh-0", "replacement-seed"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_snapshot_and_oauth_sharing_follow_the_same_server_selection(task_config, rotating_server, monkeypatch):
|
||||
_path, startup = task_config
|
||||
startup.mcp_servers["other"] = startup.mcp_servers["reports"].model_copy(deep=True)
|
||||
# Simulate a narrower task-server policy. A server excluded by that policy
|
||||
# must stay outside both the frozen snapshot and the shared OAuth lifetime.
|
||||
monkeypatch.setattr(task_runtime, "_task_enabled_servers", lambda config: {"reports": config.mcp_servers["reports"]}, raising=False)
|
||||
set_mcp_task_config_snapshot(startup)
|
||||
first = task_runtime.get_mcp_task_oauth_token_manager(startup)
|
||||
assert await first.get_authorization_header("reports") == "Bearer access-1"
|
||||
changed = startup.model_copy(deep=True)
|
||||
changed.mcp_servers["other"].oauth.refresh_token = "replacement-seed"
|
||||
validate_mcp_task_config_snapshot(changed)
|
||||
rotating_server.refresh_token = "replacement-seed"
|
||||
second = task_runtime.get_mcp_task_oauth_token_manager(changed)
|
||||
|
||||
assert await second.get_authorization_header("reports") == "Bearer access-1"
|
||||
assert await second.get_authorization_header("other") == "Bearer access-2"
|
||||
assert rotating_server.refresh_requests == ["refresh-0", "replacement-seed"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ["discovery", "durable_call"])
|
||||
async def test_personal_connection_cannot_reuse_same_named_deployment_oauth(task_config, rotating_server, tmp_path, monkeypatch, operation):
|
||||
_path, startup = task_config
|
||||
monkeypatch.setattr("deerflow.mcp.user_config.get_paths", lambda: Paths(base_dir=tmp_path))
|
||||
personal = startup.mcp_servers["reports"].model_dump(mode="json")
|
||||
personal["oauth"]["refresh_token"] = "personal-seed"
|
||||
personal["personal_public_network"] = True
|
||||
path = user_mcp_config_path("user-1")
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps({"mcpServers": {"reports": personal}}), encoding="utf-8")
|
||||
name = personal_server_name("user-1", "reports", personal)
|
||||
startup.mcp_servers = {name: startup.mcp_servers["reports"]}
|
||||
set_mcp_task_config_snapshot(startup)
|
||||
caller = McpTaskToolCaller(startup)
|
||||
assert await get_initial_oauth_headers(startup, token_manager=task_runtime.get_mcp_task_oauth_token_manager(startup)) == {name: "Bearer access-1"}
|
||||
rotating_server.refresh_token = "personal-seed"
|
||||
before = path.read_bytes()
|
||||
|
||||
if operation == "discovery":
|
||||
set_mcp_task_submitter(SimpleNamespace())
|
||||
tools = await get_mcp_tools(load_user_mcp_config("user-1"), personal_user_id="user-1")
|
||||
assert [tool.name for tool in tools] == [f"{name}_submit_report", f"{name}_search"]
|
||||
assert await tools[1].ainvoke({})
|
||||
else:
|
||||
await caller.call_tool(server_name=name, tool_name="status_report", arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", connection_scope="personal")
|
||||
assert rotating_server.refresh_requests == ["refresh-0", "personal-seed"]
|
||||
assert rotating_server.calls[-1][1] == "Bearer access-2"
|
||||
assert path.read_bytes() == before
|
||||
Loading…
x
Reference in New Issue
Block a user