diff --git a/backend/app/gateway/routers/mcp.py b/backend/app/gateway/routers/mcp.py index 94506f0e0e..6e91c1260e 100644 --- a/backend/app/gateway/routers/mcp.py +++ b/backend/app/gateway/routers/mcp.py @@ -458,6 +458,7 @@ class McpServerConfigResponse(BaseModel): url: str | None = Field(default=None, description="URL of the MCP server (for sse or http type)") headers: dict[str, str] = Field(default_factory=dict, description="HTTP headers to send (for sse or http type)") oauth: McpOAuthConfigResponse | None = Field(default=None, description="OAuth configuration for MCP HTTP/SSE servers") + agent_auth: bool = Field(default=False, description="Require an administrator-granted agent credential on this deployment HTTP server") user_auth: McpUserScopedAuthConfigResponse | None = Field(default=None, description="Per-user credential injection for MCP HTTP/SSE servers") headers_from_context: McpContextHeadersConfigResponse | None = Field(default=None, description="Per-request credential injection for MCP HTTP/SSE servers: map header names to config.context.secrets keys") description: str = Field(default="", description="Human-readable description of what this MCP server provides") @@ -1106,6 +1107,9 @@ def _merge_preserving_secrets( } if preserve_omitted_fields and "user_auth" not in incoming.model_fields_set: update["user_auth"] = existing.user_auth + # Changing an ordinary field must not silently remove credential enforcement. + if "agent_auth" not in incoming.model_fields_set: + update["agent_auth"] = existing.agent_auth if preserve_omitted_fields and "headers_from_context" not in incoming.model_fields_set: update["headers_from_context"] = existing.headers_from_context if preserve_omitted_fields and "routing" not in incoming.model_fields_set: diff --git a/backend/app/gateway/routers/personal_mcp.py b/backend/app/gateway/routers/personal_mcp.py index 5d82540052..c0bad49cc9 100644 --- a/backend/app/gateway/routers/personal_mcp.py +++ b/backend/app/gateway/routers/personal_mcp.py @@ -45,6 +45,8 @@ def _read_personal_config(user_id: str) -> dict: def _validate_personal_server(server: mcp.McpServerConfigResponse, *, admin: bool) -> None: + if server.agent_auth: + raise HTTPException(422, "agent_auth is available only for deployment MCP servers") # A personal store must not grant ordinary users the deployment operator's # ability to launch arbitrary host packages or query internal services. if not admin: diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index ba588e37a1..a6f859485b 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -92,7 +92,7 @@ from deerflow.runtime.checkpoint_mode import ( ) from deerflow.runtime.checkpoint_state import graph_state_schema from deerflow.runtime.client_context import sanitize_client_context -from deerflow.runtime.context_keys import PROJECT_CONTEXT_KEY +from deerflow.runtime.context_keys import AGENT_AUTH_CONTEXT_KEY, PROJECT_CONTEXT_KEY from deerflow.runtime.events.message_identity import MESSAGE_SEQ_KEY from deerflow.runtime.goal import goal_thread_lock from deerflow.runtime.journal import build_checkpoint_history_seed_events @@ -1152,6 +1152,7 @@ _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: frozenset[str] = ( # at admission from threads_meta; a client-supplied value must # never survive in either run-config section. PROJECT_CONTEXT_KEY, + AGENT_AUTH_CONTEXT_KEY, KNOWLEDGE_SCOPE_KEY, KNOWLEDGE_SCOPE_RUNTIME_KEY, } @@ -2696,6 +2697,20 @@ async def start_run( request_context=getattr(body, "context", None), ) + if agent_config is not None and getattr(run_ctx.app_config, "mcp_agent_auth", None): + from deerflow.mcp.agent_auth import McpAgentAuthorizationError, capture_agent_authority + + try: + agent_authority = await run_assembly( + capture_agent_authority, + run_ctx.app_config, + user_id=owner_user_id or str(user.id), + agent_config=agent_config, + ) + except McpAgentAuthorizationError as exc: + raise HTTPException(status_code=403, detail=str(exc)) from None + run_ctx = replace(run_ctx, agent_mcp_authority=agent_authority) + if scheduled_task_runtime is not None: run_ctx = replace(run_ctx, scheduled_task_runtime=dict(scheduled_task_runtime)) conversation_references = list(getattr(body, "conversation_references", None) or []) diff --git a/backend/docs/MCP_SERVER.md b/backend/docs/MCP_SERVER.md index 72cd8bbcc3..6b63a1a306 100644 --- a/backend/docs/MCP_SERVER.md +++ b/backend/docs/MCP_SERVER.md @@ -530,6 +530,129 @@ token and refresh token in the `mcp_oauth_tokens` table (migration so that startup check cannot see them; set the variable on every multi-instance deployment that uses them. +## Per-agent credentials (deployment HTTP servers) + +For a shared HTTP MCP server, an administrator can grant a different OAuth +service account to each custom agent. The Gateway admits an owner and an +approved agent definition; changing a client-supplied `agent_name` or creating +a same-named agent in another account does not grant credentials. + +Enable enforcement on the deployment server in `extensions_config.json`: + +```json +{ + "mcpServers": { + "ticketing": { + "type": "http", + "url": "https://ticketing.example/mcp", + "headers": {"Authorization": "Bearer $MCP_DISCOVERY_TOKEN"}, + "agent_auth": true + } + } +} +``` + +Tool discovery still uses the server's static headers or server-level OAuth. +Use a discovery credential limited to listing tools. Calls to a protected +server always require an agent grant; there is no passthrough to discovery +credentials. Default Lead, bootstrap runs and owners without a matching grant +cannot call its tools. Personal MCP connections cannot enable this feature. + +Enabling protection also denies calls through tools issued before the change; +start a new run to discover tools with the current authentication chain. Editing +a server preserves `agent_auth` when the field is omitted, including on the +single-server replacement endpoint. Disabling protection requires explicit +`agent_auth: false`. + +Keep the grants in operator-owned `config.yaml`, separately from the +API-writable server configuration: + +```yaml +mcp_agent_auth: + research-tickets: # Stable binding ID; durable tasks retain this reference. + server: ticketing + agent: research-bot + owners: ["alice-user-id", "bob-user-id"] + definition_digests: [""] + oauth: + grant_type: client_credentials + token_url: https://idp.example/realms/work/protocol/openid-connect/token + client_id: research-bot + client_secret: $RESEARCH_BOT_CLIENT_SECRET + scope: tickets.read +``` + +This first version supports explicit owner IDs and OAuth `client_credentials` +on deployment-managed HTTP servers. An owners list lets a GitOps job maintain +membership for identical per-user copies. Wildcards, group matching, static +per-agent tokens and refresh-token grants are not supported. Enabled bindings +must not overlap for the same server, agent and owner. + +### Approve a definition + +Generate the digest in the Gateway environment with the same DeerFlow version +and data root, after reviewing the stored agent: + +```bash +uv run --no-sync --project backend python - alice-user-id research-bot <<'PY' +import sys +from deerflow.config.agents_config import load_agent_config, load_agent_soul +from deerflow.mcp.agent_auth import agent_definition_digest + +owner, name = sys.argv[1:] +config = load_agent_config(name, user_id=owner) +if config is None: + raise SystemExit("Agent not found") +print(agent_definition_digest(config, load_agent_soul(name, user_id=owner))) +PY +``` + +The versioned digest covers the parsed `AgentConfig` and the exact SOUL text. +The Gateway verifies these at run admission and assembles the run from that +same in-memory snapshot. Editing either invalidates admission until an +administrator approves the new digest. Editing an agent does not edit its +grant. For rolling syncs, the binding can temporarily allow both old and new +digests, then remove the old one after copies have been updated. + +The digest does not attest dynamic tool schemas, skill file contents, remote +data or the model's behavior. Review the approved tool/skill selections and +enforce service-account permissions at the MCP service. Operator-controlled +custom interceptors remain trusted code. + +### Delegation, tasks and revocation + +Ordinary subagents inherit the calling agent's grant, even if a child has its +own name or binding. Durable subagent batches and ordinary MCP background +tasks store only nonsecret binding/definition references, then restore that +identity after worker recovery. A new definition rollout does not switch an +existing task's service account or break its status/cancel calls. + +References also bind the MCP transport and exact URL. Changing that destination +invalidates existing references, including after a Gateway restart; old tasks +cannot send their account's credential to the replacement endpoint. References +without a destination fingerprint are refused rather than upgraded implicitly. + +Each call checks the current operator grant, including when the access token +is cached. Disable/delete the binding or remove an owner to deny subsequent +calls, polls and cancellations. A call already sent cannot be recalled. +Changing the OAuth endpoint, client, scope or audience also invalidates old +references; rotating the same account's client secret keeps them valid and +refreshes the cache. Missing grants or failed token acquisition never fall +back to user or server credentials. + +Grant edits are hot-reloaded independently of the server connection snapshot. +As with other task-tool connection settings, changing the server's +`agent_auth` flag requires a Gateway restart when it has `task_toolsets`. +Environment secrets are resolved at configuration load; roll out a config +reload or restart when rotating them. Access tokens are cached in process by +server, binding and credential configuration, using the existing OAuth expiry +and refresh logic. Logs identify the server, agent, owner and binding without +logging credential values. + +Config loading rejects enabled `user_auth` on protected servers and any enabled +`headers_from_context` mapping of `Authorization`, regardless of header casing. +Non-auth context headers can still be forwarded on foreground calls. + ## Request-Scoped Headers (HTTP/SSE MCP Servers) When the credential is chosen by the *caller* rather than by the operator — @@ -600,15 +723,16 @@ The caller supplies the values on each run request: way a value reaches these headers: `user_auth`, the OAuth token returned by the token endpoint, and the static `headers` in the config file. - Precedence for a server declaring several sources: static `headers` < - `oauth` < `user_auth` < `headers_from_context`. The value chosen for this one - request is the most specific, so it wins. + `oauth` < `user_auth` < `headers_from_context`. Protected `agent_auth` servers + reject competing built-in auth sources as described above. - `sse`/`http` only. A stdio server has no HTTP headers; declaring the block there logs a warning and is ignored. - Durable background tasks are the one exception, and only half of one: a `task_toolsets` submit is awaited inside the Agent run and carries these headers, but the status and cancel polls run after that run ends, so they skip them and use configured static/OAuth credentials or the persisted owner's - `user_auth` credential. See *Durable Background + `user_auth` credential, or the admitted agent binding on `agent_auth` servers. + See *Durable Background Tasks* above. Use `user_auth` instead when the credential belongs to a configured DeerFlow @@ -635,7 +759,7 @@ Each entry is a Python import path in `module:variable` format (resolved via `re Tool discovery is server-scoped, so every server gets its own interceptor instance built from a single-server configuration view, and a stateful interceptor (a counter, a cache, a rate limiter) only ever observes that server's calls. Do not rely on one instance being -shared across servers. The built-in OAuth, per-user and context-header interceptors are +shared across servers. The built-in OAuth, agent, user and context-header interceptors are unaffected: they resolve everything per `request.server_name`. Servers that are not discovered never run their builder. diff --git a/backend/packages/harness/deerflow/agents/lead_agent/agent.py b/backend/packages/harness/deerflow/agents/lead_agent/agent.py index 3f96cf4519..208f233819 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/agent.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/agent.py @@ -65,6 +65,7 @@ from deerflow.config.subagents_config import ( effective_subagent_concurrency, effective_total_subagents_per_run, ) +from deerflow.mcp.agent_auth import AGENT_AUTH_CONTEXT_KEY, AgentMcpAuthority from deerflow.models import create_chat_model from deerflow.models.reasoning import resolve_reasoning_contract, resolve_reasoning_request from deerflow.runtime.checkpoint_mode import ( @@ -1028,7 +1029,14 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le scheduler_capability = None agent_name = validate_agent_name(cfg.get("agent_name")) - agent_config = load_agent_config(agent_name, user_id=resolved_user_id) if not is_bootstrap else None + agent_authority = runtime_context.get(AGENT_AUTH_CONTEXT_KEY) if isinstance(runtime_context, Mapping) else None + pinned_agent = isinstance(agent_authority, AgentMcpAuthority) and not is_bootstrap and not cfg.get("is_subagent") + if pinned_agent: + agent_config = agent_authority.agent_config() + if agent_authority.user_id != resolved_user_id or agent_config is None or agent_config.name != agent_name: + raise ValueError("Admitted MCP agent definition does not match this run") + else: + agent_config = load_agent_config(agent_name, user_id=resolved_user_id) if not is_bootstrap else None memory_enabled = getattr(agent_config, "memory_enabled", True) is not False # Keep compatibility with lightweight AgentConfig-shaped objects used by # integrations that predate caller-level subagent restrictions. @@ -1424,6 +1432,7 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le max_concurrent_subagents=max_concurrent_subagents, max_total_subagents=max_total_subagents, agent_name=agent_name, + agent_soul=agent_authority.soul if pinned_agent else None, available_skills=available_skills, app_config=resolved_app_config, deferred_names=setup.deferred_names, diff --git a/backend/packages/harness/deerflow/agents/lead_agent/prompt.py b/backend/packages/harness/deerflow/agents/lead_agent/prompt.py index faffa356c6..4394b3dd00 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/prompt.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/prompt.py @@ -1002,7 +1002,10 @@ def get_skills_prompt_section( def get_agent_soul(agent_name: str | None, *, user_id: str | None = None) -> str: # Append SOUL.md (agent personality) if present - soul = load_agent_soul(agent_name, user_id=user_id) + return _format_agent_soul(load_agent_soul(agent_name, user_id=user_id)) + + +def _format_agent_soul(soul: str | None) -> str: if soul: # SOUL.md is agent-editable (setup_agent / update_agent persist it) and is # rendered into the block of the lead-agent system prompt. Escape it @@ -1132,6 +1135,7 @@ def apply_prompt_template( max_total_subagents: int | None = None, *, agent_name: str | None = None, + agent_soul: str | None = None, available_skills: set[str] | None = None, app_config: AppConfig | None = None, deferred_names: frozenset[str] = frozenset(), @@ -1243,7 +1247,7 @@ def apply_prompt_template( clarification_system=interaction_policy.clarification_system, clarification_reminder=interaction_policy.clarification_reminder, agent_name=agent_name or "DeerFlow 2.0", - soul=get_agent_soul(agent_name, user_id=user_id), + soul=get_agent_soul(agent_name, user_id=user_id) if agent_soul is None else _format_agent_soul(agent_soul), self_update_section=_build_self_update_section(agent_name), skills_section=skills_section, deferred_tools_section=deferred_tools_section, diff --git a/backend/packages/harness/deerflow/config/app_config.py b/backend/packages/harness/deerflow/config/app_config.py index 12068793e6..631b6582a7 100644 --- a/backend/packages/harness/deerflow/config/app_config.py +++ b/backend/packages/harness/deerflow/config/app_config.py @@ -32,6 +32,7 @@ from deerflow.config.hitl_approval_config import HitlApprovalConfig from deerflow.config.input_polish_config import InputPolishConfig from deerflow.config.knowledge_base_config import KnowledgeBaseConfig from deerflow.config.loop_detection_config import LoopDetectionConfig +from deerflow.config.mcp_agent_auth_config import BindingId, McpAgentCredentialBinding from deerflow.config.mcp_tasks_config import McpTasksConfig from deerflow.config.memory_config import MemoryConfig, load_memory_config_from_dict from deerflow.config.model_config import ModelConfig @@ -286,6 +287,22 @@ class AppConfig(BaseModel): skill_scan: SkillScanConfig = Field(default_factory=SkillScanConfig, description="Native deterministic skill safety scanning configuration") skill_evolution: SkillEvolutionConfig = Field(default_factory=SkillEvolutionConfig, description="Agent-managed skill evolution configuration") extensions: ExtensionsConfig = Field(default_factory=ExtensionsConfig, description="Extensions configuration (MCP servers and skills state)") + mcp_agent_auth: dict[BindingId, McpAgentCredentialBinding] = Field(default_factory=dict, description="Administrator-granted MCP agent credentials, checked on each call") + + @field_validator("mcp_agent_auth") + @classmethod + def _unique_mcp_agent_grants(cls, bindings: dict[str, McpAgentCredentialBinding]) -> dict[str, McpAgentCredentialBinding]: + subjects: set[tuple[str, str, str]] = set() + for binding in bindings.values(): + if not binding.enabled: + continue + for owner in binding.owners: + subject = (binding.server, binding.agent, owner) + if subject in subjects: + raise ValueError("MCP agent bindings overlap for the same server, agent and owner") + subjects.add(subject) + return bindings + tool_output: ToolOutputConfig = Field(default_factory=ToolOutputConfig, description="Tool output budget protection configuration") tool_artifacts: ToolArtifactConfig = Field(default_factory=ToolArtifactConfig, description="Tool-artifact handle registry configuration (issue #4676)") tool_search: ToolSearchConfig = Field(default_factory=ToolSearchConfig, description="Tool search / deferred loading configuration") @@ -799,7 +816,7 @@ def _load_and_cache_app_config(config_path: str | None = None, *, extensions_sou return _app_config -def get_app_config() -> AppConfig: +def get_app_config(*, use_runtime_override: bool = True) -> AppConfig: """Get the DeerFlow config instance. Returns a cached singleton instance and automatically reloads it when the @@ -810,10 +827,13 @@ def get_app_config() -> AppConfig: ``reload_app_config(config_path)``, that path is preferred over ambient resolution. Use `reload_app_config()` to force a reload, or `reset_app_config()` to clear the cache. + + Credential grant checks set `use_runtime_override=False` so revocation is + read from the operator configuration even while a run keeps its snapshot. """ global _app_config, _app_config_path, _app_config_mtime, _app_config_signature - runtime_override = _current_app_config.get() + runtime_override = _current_app_config.get() if use_runtime_override else None if runtime_override is not None: return runtime_override @@ -927,10 +947,11 @@ def peek_current_app_config() -> AppConfig | None: return _current_app_config.get() -def peek_loaded_app_config() -> AppConfig | None: +def peek_loaded_app_config(*, use_runtime_override: bool = True) -> AppConfig | None: """Return the configuration this process has loaded, without touching the filesystem. - The runtime-scoped override wins when one is active; otherwise this is the + The runtime-scoped override wins when one is active unless + ``use_runtime_override=False``; otherwise this is the cached singleton ``get_app_config()`` last loaded (or ``set_app_config()`` installed). ``None`` means no configuration has ever been loaded in this process, which is what a host without a ``config.yaml`` looks like — as @@ -938,7 +959,7 @@ def peek_loaded_app_config() -> AppConfig | None: file, where callers need the loaded value to tell "unavailable" from "never configured". """ - override = _current_app_config.get() + override = _current_app_config.get() if use_runtime_override else None if override is not None: return override return _app_config diff --git a/backend/packages/harness/deerflow/config/extensions_config.py b/backend/packages/harness/deerflow/config/extensions_config.py index dff0f0a71d..a835789cb1 100644 --- a/backend/packages/harness/deerflow/config/extensions_config.py +++ b/backend/packages/harness/deerflow/config/extensions_config.py @@ -227,6 +227,7 @@ class McpServerConfig(BaseModel): url: str | None = Field(default=None, description="URL of the MCP server (for sse or http type)") headers: dict[str, str] = Field(default_factory=dict, description="HTTP headers to send (for sse or http type)") oauth: McpOAuthConfig | None = Field(default=None, description="OAuth configuration (for sse or http type)") + agent_auth: bool = Field(default=False, description="Require an administrator-granted agent identity for calls to this deployment HTTP server") user_auth: McpUserScopedAuthConfig | None = Field( default=None, description="Per-user credential injection (for sse or http type): map DeerFlow user ids to per-user credential header values", @@ -271,6 +272,18 @@ class McpServerConfig(BaseModel): ) model_config = ConfigDict(extra="allow") + @model_validator(mode="after") + def _validate_agent_credentials(self) -> "McpServerConfig": + if self.agent_auth: + if self.type != "http": + raise ValueError("agent_auth requires an HTTP MCP server") + if self.user_auth is not None and self.user_auth.enabled: + raise ValueError("agent_auth and user_auth are mutually exclusive") + context_headers = self.headers_from_context + if context_headers is not None and context_headers.enabled and any(name.lower() == "authorization" for name in context_headers.headers): + raise ValueError("agent_auth conflicts with headers_from_context for Authorization") + return self + @field_validator("headers") @classmethod def _validate_header_names(cls, value: dict[str, str]) -> dict[str, str]: diff --git a/backend/packages/harness/deerflow/config/mcp_agent_auth_config.py b/backend/packages/harness/deerflow/config/mcp_agent_auth_config.py new file mode 100644 index 0000000000..935b36adb5 --- /dev/null +++ b/backend/packages/harness/deerflow/config/mcp_agent_auth_config.py @@ -0,0 +1,46 @@ +"""Operator-owned grants for named agents on deployment MCP servers.""" + +from typing import Annotated, Self + +from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_validator, model_validator + +from deerflow.config.extensions_config import McpOAuthConfig + +DefinitionDigest = Annotated[str, StringConstraints(pattern=r"^[0-9a-f]{64}$")] +BindingId = Annotated[str, StringConstraints(pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")] + + +class McpAgentCredentialBinding(BaseModel): + model_config = ConfigDict(extra="forbid", hide_input_in_errors=True) + + enabled: bool = True + server: str = Field(min_length=1) + agent: str = Field(min_length=1) + owners: list[str] + definition_digests: list[DefinitionDigest] = Field(min_length=1) + oauth: McpOAuthConfig = Field(repr=False) + + @field_validator("agent") + @classmethod + def _agent_name(cls, value: str) -> str: + from deerflow.config.agents_config import validate_agent_name + + validate_agent_name(value) + return value + + @field_validator("owners") + @classmethod + def _explicit_owners(cls, value: list[str]) -> list[str]: + if any(not owner.strip() or owner != owner.strip() or owner == "*" for owner in value): + raise ValueError("MCP agent bindings require explicit, nonblank owner IDs") + if len(value) != len(set(value)): + raise ValueError("MCP agent binding owner IDs must be unique") + return value + + @model_validator(mode="after") + def _client_credentials(self) -> Self: + if self.oauth.grant_type != "client_credentials": + raise ValueError("MCP agent bindings support OAuth client_credentials") + if self.enabled and (not self.oauth.enabled or not self.oauth.client_id or not self.oauth.client_secret): + raise ValueError("An enabled MCP agent binding requires enabled OAuth with a client ID and secret") + return self diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index f120c5c21e..473da33cec 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -43,14 +43,18 @@ make test-blocking-io - **Read the run context from `request.runtime`** in an interceptor. Context variables are not reliable there, because the adapter may call from another task. -- **Three credential sources** exist for HTTP and SSE servers: deployment OAuth, - per-user `user_auth`, and per-request `headers_from_context`. A missing +- **Credential sources** include deployment OAuth, agent grants, per-user + `user_auth`, and per-request `headers_from_context`. A missing credential fails the call. It never falls back to another user's or to the deployment credential. - **The OAuth store keeps refresh tokens only**, encrypted and keyed by scope, owner and server. Personal managers pass `owner_user_id`. - **Background task calls have no request.** `McpTaskToolCaller` binds the task owner's identity for the call and restores the previous value afterwards. +- **Agent grants are live policy, not run snapshots.** Every protected call + rechecks the operator binding even with a cached token. Durable references + retain the admitted definition revision; definition rollouts preserve them, + while grant revocation denies them. See the per-agent section in the MCP guide. ## Results diff --git a/backend/packages/harness/deerflow/mcp/agent_auth.py b/backend/packages/harness/deerflow/mcp/agent_auth.py new file mode 100644 index 0000000000..b697e16aff --- /dev/null +++ b/backend/packages/harness/deerflow/mcp/agent_auth.py @@ -0,0 +1,256 @@ +"""Host-admitted agent identities for deployment MCP calls.""" + +from __future__ import annotations + +import asyncio +import logging +import threading +from collections import OrderedDict +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Any, Literal + +from deerflow_extension_api import canonical_hash, canonical_json +from langchain_core.tools import ToolException +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from deerflow.config.agents_config import AgentConfig, load_agent_soul +from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig, get_extensions_config +from deerflow.config.mcp_agent_auth_config import BindingId, DefinitionDigest, McpAgentCredentialBinding +from deerflow.mcp.headers import apply_header_overrides, header_spellings +from deerflow.mcp.oauth import OAuthTokenManager +from deerflow.runtime.context_keys import AGENT_AUTH_CONTEXT_KEY +from deerflow.runtime.user_context import resolve_runtime_user_id + +logger = logging.getLogger(__name__) + + +class McpAgentAuthorizationError(ToolException): + """An agent's administrator grant is absent, revoked or replaced.""" + + +class AgentCredentialReference(BaseModel): + """Nonsecret identity retained by a durable task.""" + + model_config = ConfigDict(extra="forbid", frozen=True, hide_input_in_errors=True) + version: Literal[1] = 1 + server: str + binding: BindingId + agent: str + definition_digest: DefinitionDigest + credential_digest: DefinitionDigest + destination_digest: DefinitionDigest + + +@dataclass(frozen=True) +class AgentMcpAuthority: + user_id: str + references: tuple[AgentCredentialReference, ...] + definition_json: str | None = None + soul: str = "" + + def agent_config(self) -> AgentConfig | None: + return AgentConfig.model_validate_json(self.definition_json) if self.definition_json is not None else None + + +_task_authority: ContextVar[AgentMcpAuthority | None] = ContextVar("deerflow_mcp_task_agent_authority", default=None) +_reference_list = TypeAdapter(list[AgentCredentialReference]) +_token_managers: OrderedDict[tuple[str, str, str], OAuthTokenManager] = OrderedDict() +_token_managers_lock = threading.Lock() + + +def agent_definition_digest(config: AgentConfig, soul: str | None) -> str: + """Hash the stored agent settings and SOUL text that admission will pin.""" + return canonical_hash({"version": 1, "config": config.model_dump(mode="json"), "soul": soul or ""}) + + +def _credential_digest(binding: McpAgentCredentialBinding) -> str: + # Rotating a secret preserves the account. Replacing the endpoint, client, + # scope or other request settings requires a new admission. + return canonical_hash(binding.oauth.model_dump(mode="json", exclude={"client_secret", "refresh_token", "refresh_skew_seconds"})) + + +def _destination_digest(server: McpServerConfig) -> str: + return canonical_hash({"transport": server.type, "url": server.url}) + + +def _current_server(server_name: str) -> McpServerConfig | None: + from deerflow.config.app_config import get_app_config, peek_loaded_app_config + + try: + # Standalone MCP clients need only extensions_config.json. A Gateway + # also applies its operator-owned YAML overrides, never a run snapshot. + config = get_app_config(use_runtime_override=False).extensions if peek_loaded_app_config(use_runtime_override=False) is not None else get_extensions_config() + return config.mcp_servers.get(server_name) + except Exception: + raise ToolException("MCP agent credential policy is unavailable") from None + + +def capture_agent_authority(app_config: Any, *, user_id: str, agent_config: AgentConfig | None) -> AgentMcpAuthority | None: + """Admit a host-loaded definition once, before an agent run is assembled.""" + if agent_config is None: + return None + matches = [] + for binding_id, binding in app_config.mcp_agent_auth.items(): + server = app_config.extensions.mcp_servers.get(binding.server) + if binding.enabled and binding.agent == agent_config.name and user_id in binding.owners and server is not None and server.enabled and server.agent_auth: + matches.append((binding_id, binding)) + if not matches: + return None + soul = load_agent_soul(agent_config.name, user_id=user_id) or "" + digest = agent_definition_digest(agent_config, soul) + references = [] + for binding_id, binding in matches: + if digest not in binding.definition_digests: + raise McpAgentAuthorizationError(f"Agent definition is not approved for MCP binding '{binding_id}'") + references.append( + AgentCredentialReference( + server=binding.server, + binding=binding_id, + agent=binding.agent, + definition_digest=digest, + credential_digest=_credential_digest(binding), + destination_digest=_destination_digest(app_config.extensions.mcp_servers[binding.server]), + ) + ) + return AgentMcpAuthority(user_id, tuple(references), canonical_json(agent_config.model_dump(mode="json")), soul) + + +def runtime_agent_authority(runtime: Any) -> AgentMcpAuthority | None: + context = getattr(runtime, "context", None) + authority = context.get(AGENT_AUTH_CONTEXT_KEY) if isinstance(context, Mapping) else None + if isinstance(authority, AgentMcpAuthority) and authority.user_id == resolve_runtime_user_id(runtime): + return authority + return None + + +def export_agent_authority(authority: AgentMcpAuthority | None) -> list[dict[str, Any]] | None: + """Persist references captured by the host, without definitions or credentials.""" + if not isinstance(authority, AgentMcpAuthority): + return None + return [reference.model_dump(mode="json") for reference in authority.references] + + +def restore_agent_authority(user_id: str, references: Any) -> AgentMcpAuthority | None: + """Restore references from host-owned durable rows; each use rechecks the grant.""" + if references is None: + return None + try: + parsed = _reference_list.validate_python(references) + except ValidationError: + raise McpAgentAuthorizationError("Invalid persisted MCP agent identity") from None + if not parsed or len({ref.server for ref in parsed}) != len(parsed) or len({ref.agent for ref in parsed}) != 1: + raise McpAgentAuthorizationError("Invalid persisted MCP agent identity") + return AgentMcpAuthority(user_id, tuple(parsed)) + + +def task_agent_reference(runtime: Any, server_name: str, *, required: bool) -> dict[str, Any] | None: + authority = runtime_agent_authority(runtime) + if authority is not None: + for reference in authority.references: + if reference.server == server_name: + return reference.model_dump(mode="json") + if required: + raise McpAgentAuthorizationError(f"No admitted agent credential for MCP server '{server_name}'") + return None + + +def _current_binding(reference: AgentCredentialReference, user_id: str) -> McpAgentCredentialBinding: + from deerflow.config.app_config import get_app_config + + try: + config = get_app_config(use_runtime_override=False) + except Exception: + raise ToolException("MCP agent credential policy is unavailable") from None + binding = config.mcp_agent_auth.get(reference.binding) + server = config.extensions.mcp_servers.get(reference.server) + if ( + binding is None + or not binding.enabled + or user_id not in binding.owners + or binding.server != reference.server + or binding.agent != reference.agent + or server is None + or not server.enabled + or not server.agent_auth + or _destination_digest(server) != reference.destination_digest + or _credential_digest(binding) != reference.credential_digest + ): + raise McpAgentAuthorizationError(f"MCP agent credential binding '{reference.binding}' is unavailable or revoked") + # Definition rollout changes admission for new runs. Existing calls and + # durable tasks retain the definition version recorded at their admission. + return binding + + +def _token_manager(reference: AgentCredentialReference, binding: McpAgentCredentialBinding) -> OAuthTokenManager: + key = (reference.server, reference.binding, canonical_hash(binding.oauth.model_dump(mode="json"))) + with _token_managers_lock: + manager = _token_managers.get(key) + if manager is None: + manager = OAuthTokenManager({reference.server: binding.oauth}) + _token_managers[key] = manager + _token_managers.move_to_end(key) + if len(_token_managers) > 128: + _token_managers.popitem(last=False) + return manager + + +@contextmanager +def bind_task_agent_authority(authority: AgentMcpAuthority | None) -> Iterator[None]: + token = _task_authority.set(authority) + try: + yield + finally: + _task_authority.reset(token) + + +async def validate_task_agent_reference(user_id: str, server_name: str, reference: Any) -> AgentMcpAuthority | None: + authority = restore_agent_authority(user_id, [reference]) if reference is not None else None + if authority is not None: + admitted = authority.references[0] + if admitted.server != server_name: + raise McpAgentAuthorizationError("MCP task credential belongs to a different server") + await asyncio.to_thread(_current_binding, admitted, user_id) + return authority + + +def build_agent_auth_interceptor(extensions_config: ExtensionsConfig) -> Any | None: + servers = {name: server for name, server in extensions_config.get_enabled_mcp_servers().items() if server.type in ("http", "sse")} + if not servers: + return None + + async def agent_auth_interceptor(request: Any, handler: Any) -> Any: + server = servers.get(request.server_name) + if server is None: + return await handler(request) + if not server.agent_auth: + current = await asyncio.to_thread(_current_server, request.server_name) + if current is not None and current.agent_auth: + # The old chain may still contain user/context auth writers. + # Require rediscovery instead of injecting into that chain. + raise McpAgentAuthorizationError(f"MCP server '{request.server_name}' now requires agent credentials; start a new run to reload its tools") + return await handler(request) + runtime = getattr(request, "runtime", None) + authority = runtime_agent_authority(runtime) if runtime is not None else _task_authority.get() + reference = next((ref for ref in authority.references if ref.server == request.server_name), None) if authority is not None else None + if reference is None: + raise McpAgentAuthorizationError(f"No admitted agent credential for MCP server '{request.server_name}'") + if _destination_digest(server) != reference.destination_digest: + raise McpAgentAuthorizationError(f"MCP destination no longer matches binding '{reference.binding}'") + binding = await asyncio.to_thread(_current_binding, reference, authority.user_id) + manager = _token_manager(reference, binding) + try: + authorization = await manager.get_authorization_header(reference.server) + except Exception: + raise ToolException(f"Could not acquire the agent credential for MCP server '{reference.server}'") from None + if not authorization: + raise ToolException(f"Missing agent credential for MCP server '{reference.server}'") + # Token acquisition may await an external identity provider. + await asyncio.to_thread(_current_binding, reference, authority.user_id) + logger.info("MCP agent credential selected for server=%s agent=%s owner=%s binding=%s", reference.server, reference.agent, authority.user_id, reference.binding) + headers = apply_header_overrides(request.headers, {"Authorization": authorization}, spellings=header_spellings(server.headers)) + return await handler(request.override(headers=headers)) + + return agent_auth_interceptor diff --git a/backend/packages/harness/deerflow/mcp/interceptors.py b/backend/packages/harness/deerflow/mcp/interceptors.py index 3863c04fe9..fcccb03b28 100644 --- a/backend/packages/harness/deerflow/mcp/interceptors.py +++ b/backend/packages/harness/deerflow/mcp/interceptors.py @@ -3,9 +3,10 @@ from __future__ import annotations import logging -from typing import Any +from typing import Any, Literal from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.mcp.agent_auth import build_agent_auth_interceptor from deerflow.mcp.config_normalization import normalize_mcp_interceptor_paths from deerflow.mcp.context_headers import build_context_headers_interceptor from deerflow.mcp.oauth import build_oauth_tool_interceptor @@ -21,15 +22,20 @@ def build_mcp_tool_interceptors( oauth_builder: Any = build_oauth_tool_interceptor, user_auth_builder: Any = build_user_scoped_auth_interceptor, context_headers_builder: Any = build_context_headers_interceptor, + connection_scope: Literal["deployment", "personal"] = "deployment", resolver: Any = resolve_variable, target_logger: logging.Logger = logger, ) -> list[Any]: - """Build OAuth, user-scoped auth, context headers, then custom MCP interceptors.""" + """Build OAuth, agent/user auth, context headers, then custom interceptors.""" interceptors: list[Any] = [] oauth_interceptor = oauth_builder(extensions_config) if oauth_interceptor is not None: interceptors.append(oauth_interceptor) + agent_auth_interceptor = None if connection_scope == "personal" else build_agent_auth_interceptor(extensions_config) + if agent_auth_interceptor is not None: + interceptors.append(agent_auth_interceptor) + # After OAuth so a server declaring both gets the per-user credential: # interceptors wrap outermost-first, so the later-registered user-scoped # override runs closer to the transport and wins the final header value. diff --git a/backend/packages/harness/deerflow/mcp/oauth.py b/backend/packages/harness/deerflow/mcp/oauth.py index 1ec847c44a..46960dc675 100644 --- a/backend/packages/harness/deerflow/mcp/oauth.py +++ b/backend/packages/harness/deerflow/mcp/oauth.py @@ -530,6 +530,9 @@ def build_oauth_tool_interceptor( spellings_by_server = {server_name: header_spellings(server_config.headers) for server_name, server_config in extensions_config.get_enabled_mcp_servers().items()} async def oauth_interceptor(request: Any, handler: Any) -> Any: + server = extensions_config.mcp_servers.get(request.server_name) + if server is not None and server.agent_auth: + return await handler(request) header = await token_manager.get_authorization_header(request.server_name) if not header: return await handler(request) diff --git a/backend/packages/harness/deerflow/mcp/task_tool_caller.py b/backend/packages/harness/deerflow/mcp/task_tool_caller.py index bc583842de..fa5a5850b7 100644 --- a/backend/packages/harness/deerflow/mcp/task_tool_caller.py +++ b/backend/packages/harness/deerflow/mcp/task_tool_caller.py @@ -14,6 +14,7 @@ from typing import Any, Literal from deerflow.config.extensions_config import ExtensionsConfig from deerflow.config.paths import get_paths from deerflow.constants import MCP_TMP_SUBDIR +from deerflow.mcp.agent_auth import McpAgentAuthorizationError, bind_task_agent_authority, validate_task_agent_reference from deerflow.mcp.client import build_server_params from deerflow.mcp.context_headers import build_context_headers_interceptor from deerflow.mcp.headers import apply_header_overrides @@ -76,6 +77,7 @@ class McpTaskToolCaller: extensions_config: ExtensionsConfig, *, oauth_token_manager: OAuthTokenManager | None = None, + connection_scope: Literal["deployment", "personal"] = "deployment", ) -> None: self._extensions_config = extensions_config self._personal_callers: OrderedDict[str, tuple[PersonalMcpConfigSnapshot, McpTaskToolCaller | None]] = OrderedDict() @@ -86,6 +88,7 @@ class McpTaskToolCaller: # custom ``mcpInterceptors`` builder is invoked exactly once. self._submit_interceptors = build_mcp_tool_interceptors( extensions_config, + connection_scope=connection_scope, oauth_builder=lambda config: build_oauth_tool_interceptor( config, token_manager=self._oauth_token_manager, @@ -117,6 +120,7 @@ class McpTaskToolCaller: thread_incarnation: str | None = None, request_scoped_headers: bool = False, connection_scope: Literal["deployment", "personal"] = "deployment", + agent_auth: dict[str, Any] | None = None, ) -> Any: """Call a raw MCP tool. @@ -126,31 +130,29 @@ class McpTaskToolCaller: run after that run ended. ``connection_scope`` is captured in the task binding, so an equal deployment server name cannot change its owner. """ - if connection_scope == "personal": - caller = await asyncio.to_thread(self._personal_caller_for, user_id, server_name) - await require_personal_mcp_access(user_id, caller._extensions_config.mcp_servers[server_name]) - return await caller._call_configured_tool( - server_name=server_name, - tool_name=tool_name, - arguments=arguments, - user_id=user_id, - 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") - return await self._call_configured_tool( - server_name=server_name, - tool_name=tool_name, - arguments=arguments, - user_id=user_id, - thread_id=thread_id, - thread_incarnation=thread_incarnation, - request_scoped_headers=request_scoped_headers, - connection_scope="deployment", - ) + try: + if connection_scope == "personal" and agent_auth is not None: + raise McpAgentAuthorizationError("Personal MCP tasks cannot use deployment agent credentials") + authority = await validate_task_agent_reference(user_id, server_name, agent_auth) + with bind_task_agent_authority(authority): + caller = self + if connection_scope == "personal": + caller = await asyncio.to_thread(self._personal_caller_for, user_id, server_name) + await require_personal_mcp_access(user_id, caller._extensions_config.mcp_servers[server_name]) + elif connection_scope != "deployment": + raise ValueError("Invalid MCP task connection scope") + return await caller._call_configured_tool( + server_name=server_name, + tool_name=tool_name, + arguments=arguments, + user_id=user_id, + thread_id=thread_id, + thread_incarnation=thread_incarnation, + request_scoped_headers=request_scoped_headers, + connection_scope=connection_scope, + ) + except McpAgentAuthorizationError as exc: + raise McpTaskConnectionUnavailableError(str(exc)) from None def _personal_caller_for(self, user_id: str, server_name: str) -> McpTaskToolCaller: with self._personal_callers_lock: @@ -167,7 +169,7 @@ class McpTaskToolCaller: if caller is None: # 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, owner_user_id=user_id)) + caller = McpTaskToolCaller(snapshot.config, oauth_token_manager=OAuthTokenManager.from_extensions_config(snapshot.config, owner_user_id=user_id), connection_scope="personal") self._personal_callers[user_id] = (snapshot, caller) return caller @@ -258,13 +260,13 @@ class McpTaskToolCaller: connection_scope=connection_scope, ) - authorization = await self._oauth_token_manager.get_authorization_header(server_name) + authorization = None if server_config.agent_auth else await self._oauth_token_manager.get_authorization_header(server_name) if authorization: connection["headers"] = apply_header_overrides( connection.get("headers") or {}, {"Authorization": authorization}, ) - # Only HTTP/SSE servers with enabled user_auth need an ambient owner. + # Agent/user-auth background calls need the persisted owner's identity. # Leave other custom-interceptor contexts unchanged. user_auth = server_config.user_auth return await self._invoke( @@ -275,7 +277,7 @@ class McpTaskToolCaller: server_name=server_name, tool_name=tool_name, arguments=arguments, - background_user_id=user_id if is_background_call and user_auth is not None and user_auth.enabled else None, + background_user_id=user_id if is_background_call and (server_config.agent_auth or (user_auth is not None and user_auth.enabled)) else None, timeout_seconds=server_config.tool_call_timeout, session_init_timeout_seconds=server_config.session_init_timeout, persistent_session=False, diff --git a/backend/packages/harness/deerflow/mcp/tasks/ordinary.py b/backend/packages/harness/deerflow/mcp/tasks/ordinary.py index 79ef9c30a9..678728a54b 100644 --- a/backend/packages/harness/deerflow/mcp/tasks/ordinary.py +++ b/backend/packages/harness/deerflow/mcp/tasks/ordinary.py @@ -48,6 +48,7 @@ class McpTaskToolCaller(Protocol): thread_incarnation: str | None = None, request_scoped_headers: bool = False, connection_scope: Literal["deployment", "personal"] = "deployment", + agent_auth: dict[str, Any] | None = None, ) -> Any: ... @@ -188,6 +189,7 @@ class OrdinaryMcpTaskDriver: # credentials; status and cancel run after that run ended. request_scoped_headers=True, connection_scope=_connection_scope(request.driver_data), + **({"agent_auth": request.driver_data["agent_auth"]} if "agent_auth" in request.driver_data else {}), ) payload = _parse( _SubmitPayload, @@ -211,6 +213,7 @@ class OrdinaryMcpTaskDriver: thread_id=task.thread_id, thread_incarnation=task.thread_incarnation, connection_scope=_connection_scope(task.driver_data), + **({"agent_auth": task.driver_data["agent_auth"]} if "agent_auth" in task.driver_data else {}), ) payload = _parse( _StatusPayload, @@ -231,6 +234,7 @@ class OrdinaryMcpTaskDriver: thread_id=task.thread_id, thread_incarnation=task.thread_incarnation, connection_scope=_connection_scope(task.driver_data), + **({"agent_auth": task.driver_data["agent_auth"]} if "agent_auth" in task.driver_data else {}), ) payload = _parse( _CancelPayload, diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index 9ded2b0c5c..5762c98025 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -883,6 +883,7 @@ def _make_background_submit_tool( status_tool: str, cancel_tool: str, connection_scope: str, + agent_auth: bool = False, ) -> BaseTool: background_contract = f"Submitted as durable background task {task_name!r}; returns a DeerFlow task ID immediately and status polling is handled automatically." @@ -897,6 +898,9 @@ def _make_background_submit_tool( context = deerflow_runtime.context if deerflow_runtime is not None and deerflow_runtime.context else {} run_id = context.get("run_id") tool_call_id = getattr(deerflow_runtime, "tool_call_id", None) if deerflow_runtime is not None else None + from deerflow.mcp.agent_auth import task_agent_reference + + agent_reference = task_agent_reference(deerflow_runtime, server_name, required=agent_auth) created = await submitter.submit( driver_name=ORDINARY_MCP_TASK_DRIVER, request=TaskSubmitRequest( @@ -913,6 +917,7 @@ def _make_background_submit_tool( "status_tool": status_tool, "cancel_tool": cancel_tool, "connection_scope": connection_scope, + **({"agent_auth": agent_reference} if agent_reference is not None else {}), }, ), ) @@ -989,6 +994,7 @@ def _configure_task_tools_for_server( status_tool=toolset.status_tool, cancel_tool=toolset.cancel_tool, connection_scope=connection_scope, + agent_auth=server_config.agent_auth, ) ) return configured @@ -1097,7 +1103,7 @@ class _PreparedServerDiscovery: pool: MCPSessionPool bindings: dict[str, ServerBinding] ownership_domain: MCPPoolDomain - connection_scope: str + connection_scope: MCPPoolDomain personal_user_id: str | None = None @@ -1109,6 +1115,7 @@ def _prepare_server_client( client_cls: Any, initial_oauth_header: str | None, oauth_token_manager: OAuthTokenManager, + connection_scope: MCPPoolDomain = "deployment", ) -> _PreparedServerClient: """Build the client/interceptors for one selected server.""" server_cfg = extensions_config.mcp_servers.get(server_name) @@ -1120,6 +1127,7 @@ def _prepare_server_client( ) tool_interceptors = build_mcp_tool_interceptors( server_view, + connection_scope=connection_scope, oauth_builder=lambda config: build_oauth_tool_interceptor(config, token_manager=oauth_token_manager), resolver=resolve_variable, target_logger=logger, @@ -1415,6 +1423,7 @@ async def _discover_and_wrap_prepared( client_cls=client_cls, initial_oauth_header=initial_oauth_headers.get(server_name), oauth_token_manager=oauth_token_manager, + connection_scope=prepared.connection_scope, ) for server_name in prepared.server_names ] diff --git a/backend/packages/harness/deerflow/mcp/user_config.py b/backend/packages/harness/deerflow/mcp/user_config.py index 0a98687144..c9d4b2edcc 100644 --- a/backend/packages/harness/deerflow/mcp/user_config.py +++ b/backend/packages/harness/deerflow/mcp/user_config.py @@ -60,7 +60,10 @@ def load_user_mcp_config(user_id: str) -> ExtensionsConfig: for name, definition in raw.get("mcpServers", {}).items(): runtime_name = personal_server_name(user_id, name, definition) servers[runtime_name] = {**definition, "tool_name_prefix": True} - return ExtensionsConfig.model_validate({"mcpServers": servers}) + config = ExtensionsConfig.model_validate({"mcpServers": servers}) + if any(server.agent_auth for server in config.mcp_servers.values()): + raise ValueError("agent_auth is available only for deployment MCP servers") + return config def _file_signature(path: Path) -> _FileSignature: diff --git a/backend/packages/harness/deerflow/runtime/context_keys.py b/backend/packages/harness/deerflow/runtime/context_keys.py index b87f292746..f7fb781b4a 100644 --- a/backend/packages/harness/deerflow/runtime/context_keys.py +++ b/backend/packages/harness/deerflow/runtime/context_keys.py @@ -17,6 +17,7 @@ DEFAULT_AGENT_NAME_METADATA_VALUE: Final[str] = "__default__" # Clients can never supply it: it is popped from both run-config sections at # Gateway admission and refused by the run worker's runtime-context merge. PROJECT_CONTEXT_KEY: Final[str] = "__deerflow_project_context" +AGENT_AUTH_CONTEXT_KEY: Final[str] = "__deerflow_mcp_agent_authority" def checkpoint_agent_binding_metadata(metadata: object) -> dict[str, str]: diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 8bfbf23fd1..b28599e0e1 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -62,6 +62,7 @@ from deerflow.runtime.checkpoint_state import ( graph_writable_channels, ) from deerflow.runtime.context_keys import ( + AGENT_AUTH_CONTEXT_KEY, CHECKPOINT_AGENT_NAME_METADATA_KEY, CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY, DEFAULT_AGENT_NAME_METADATA_VALUE, @@ -624,6 +625,7 @@ _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: Final[frozenset[str]] = ( CONVERSATION_READER_CONTEXT_KEY, SCHEDULER_CAPABILITY_CONTEXT_KEY, AGENT_RUNS_CONTEXT_KEY, + AGENT_AUTH_CONTEXT_KEY, "__run_input_control", "__tool_approval_runtime", "_hitl_disabled_resume_escape", @@ -657,6 +659,7 @@ def _build_runtime_context( *, thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET, scheduler_capability: SchedulerRunCapability | None = None, + agent_mcp_authority: Any | None = None, ) -> dict[str, Any]: """Build the dict that becomes ``ToolRuntime.context`` for the run. @@ -686,6 +689,8 @@ def _build_runtime_context( runtime_ctx[CONVERSATION_READER_CONTEXT_KEY] = conversation_reader if scheduler_capability is not None: runtime_ctx[SCHEDULER_CAPABILITY_CONTEXT_KEY] = scheduler_capability + if agent_mcp_authority is not None: + runtime_ctx[AGENT_AUTH_CONTEXT_KEY] = agent_mcp_authority if task_store is not None: from deerflow_extension_api import EXTENSION_TASK_STORE_KEY @@ -751,6 +756,7 @@ class RunContext: conversation_reader: Any | None = field(default=None) agent_runs: Any | None = field(default=None) scheduler_capability: SchedulerRunCapability | None = field(default=None) + agent_mcp_authority: Any | None = field(default=None) # Server-selected occurrence identity. Metadata alone never activates a goal. scheduled_task_runtime: Mapping[str, Any] | None = field(default=None) # Host-only preflight under durable run ownership; failures reject execution. @@ -765,6 +771,7 @@ def _install_runtime_context(config: dict, runtime_context: dict[str, Any]) -> N configurable.pop(CONVERSATION_READER_CONTEXT_KEY, None) configurable.pop(SCHEDULER_CAPABILITY_CONTEXT_KEY, None) configurable.pop(AGENT_RUNS_CONTEXT_KEY, None) + configurable.pop(AGENT_AUTH_CONTEXT_KEY, None) existing_context = config.get("context") if isinstance(existing_context, dict): existing_context.setdefault("thread_id", runtime_context["thread_id"]) @@ -1376,6 +1383,7 @@ async def run_agent( ctx.agent_runs, thread_incarnation=thread_incarnation, scheduler_capability=ctx.scheduler_capability, + agent_mcp_authority=ctx.agent_mcp_authority, ) # A disabled all-reject escape still needs the HITL hook to consume the # pending interrupt. The assembly marker makes that hook deny rather diff --git a/backend/packages/harness/deerflow/subagents/batch_service.py b/backend/packages/harness/deerflow/subagents/batch_service.py index 1fc38fc6c5..554ccbef76 100644 --- a/backend/packages/harness/deerflow/subagents/batch_service.py +++ b/backend/packages/harness/deerflow/subagents/batch_service.py @@ -513,6 +513,10 @@ class SubagentBatchService: ) return executor_kwargs = {} + if "mcp_agent_auth" in spec: + from deerflow.mcp.agent_auth import restore_agent_authority + + executor_kwargs["agent_mcp_authority"] = restore_agent_authority(batch["user_id"], spec["mcp_agent_auth"]) if THREAD_INCARNATION_CONTEXT_KEY in spec: executor_kwargs[THREAD_INCARNATION_CONTEXT_KEY] = spec[THREAD_INCARNATION_CONTEXT_KEY] executor = SubagentExecutor( diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index dbc82ac598..04e6af7958 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -959,6 +959,7 @@ class SubagentExecutor: context_snapshot: ParentContextSnapshot | None = None, thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET, input_gate: Any | None = None, + agent_mcp_authority: Any | None = None, ): """Initialize the executor. @@ -1039,6 +1040,7 @@ class SubagentExecutor: self.uploaded_files = deepcopy(uploaded_files) if uploaded_files is not None else None self.context_snapshot = context_snapshot self.input_gate = input_gate + self.agent_mcp_authority = agent_mcp_authority self.thread_id = thread_id self.thread_incarnation = thread_incarnation # Generate trace_id if not provided (for top-level calls) @@ -1932,6 +1934,10 @@ class SubagentExecutor: # audit). user_id reuses the resolved tracing id; on every # authenticated/IM path this equals the parent context value. context["user_id"] = self.user_id + if self.agent_mcp_authority is not None: + from deerflow.runtime.context_keys import AGENT_AUTH_CONTEXT_KEY + + context[AGENT_AUTH_CONTEXT_KEY] = self.agent_mcp_authority context["user_role"] = self.user_role context["oauth_provider"] = self.oauth_provider context["oauth_id"] = self.oauth_id diff --git a/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py b/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py index 5f32e03d25..4495c8b783 100644 --- a/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py @@ -266,6 +266,11 @@ async def batch_task( execution_spec[THREAD_INCARNATION_CONTEXT_KEY] = context[THREAD_INCARNATION_CONTEXT_KEY] if KNOWLEDGE_SCOPE_RUNTIME_KEY in context: execution_spec["knowledge_scope"] = execution_scope(context[KNOWLEDGE_SCOPE_RUNTIME_KEY]) + from deerflow.mcp.agent_auth import export_agent_authority, runtime_agent_authority + + agent_authority = export_agent_authority(runtime_agent_authority(runtime)) + if agent_authority is not None: + execution_spec["mcp_agent_auth"] = agent_authority try: batch = await submitter.submit( BatchSubmitRequest( diff --git a/backend/packages/harness/deerflow/tools/builtins/task_tool.py b/backend/packages/harness/deerflow/tools/builtins/task_tool.py index 4de3be1630..41cea9019c 100644 --- a/backend/packages/harness/deerflow/tools/builtins/task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/task_tool.py @@ -1110,6 +1110,8 @@ async def task_tool( tools = await run_assembly(get_available_tools, **available_tools_kwargs) # Create executor + from deerflow.mcp.agent_auth import runtime_agent_authority + executor_kwargs = { "config": config, "tools": tools, @@ -1129,6 +1131,7 @@ async def task_tool( "authz_attributes": authz_attributes, "deerflow_trace_id": deerflow_trace_id, "knowledge_scope": knowledge_scope, + "agent_mcp_authority": runtime_agent_authority(runtime), # RFC #4651 PR3: lead-supplied acceptance criteria are handed to the # executor, which appends them to the subagent's task HumanMessage as # untrusted data (sanitized and boundary-framed by diff --git a/backend/tests/blocking_io/test_mcp_agent_auth.py b/backend/tests/blocking_io/test_mcp_agent_auth.py new file mode 100644 index 0000000000..9d5f8ab20b --- /dev/null +++ b/backend/tests/blocking_io/test_mcp_agent_auth.py @@ -0,0 +1,83 @@ +"""Admission and per-call policy reloads must keep filesystem I/O off-loop.""" + +import asyncio +import json + +import pytest +from langchain_core.tools import ToolException +from langchain_mcp_adapters.interceptors import MCPToolCallRequest +from langgraph.runtime import Runtime +from support.mcp_agent_auth import agent_auth_env as agent_auth_env +from support.mcp_agent_auth import agent_mcp_transport as agent_mcp_transport + +from deerflow.config.app_config import get_app_config, reset_app_config +from deerflow.config.extensions_config import reset_extensions_config +from deerflow.mcp.agent_auth import AGENT_AUTH_CONTEXT_KEY, build_agent_auth_interceptor, capture_agent_authority, validate_task_agent_reference +from deerflow.utils.assembly_io import run_assembly + + +@pytest.mark.asyncio +async def test_cold_admission_and_live_policy_reload(agent_auth_env, agent_mcp_transport, tmp_path, monkeypatch): + env = agent_auth_env + path = tmp_path / "config.yaml" + extensions_path = tmp_path / "extensions_config.json" + data = env.config.model_dump(mode="json", exclude={"extensions"}) + await asyncio.to_thread(path.write_text, json.dumps(data), encoding="utf-8") + await asyncio.to_thread(extensions_path.write_text, env.config.extensions.model_dump_json(by_alias=True), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(path)) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(extensions_path)) + reset_app_config() + reset_extensions_config() + config = await run_assembly(get_app_config) + authority = await run_assembly(capture_agent_authority, config, user_id="alice", agent_config=env.agent) + interceptor = build_agent_auth_interceptor(config.extensions) + request = MCPToolCallRequest(name="search", args={}, server_name="tickets", runtime=Runtime(context={"user_id": "alice", AGENT_AUTH_CONTEXT_KEY: authority})) + headers = [] + + async def handler(request): + headers.append(request.headers) + + try: + await interceptor(request, handler) + assert {key.lower(): value for key, value in headers[0].items()}["authorization"] == "Bearer research-client-token" + reference = authority.references[0].model_dump(mode="json") + assert await validate_task_agent_reference("alice", "tickets", reference) is not None + data["mcp_agent_auth"]["research-tickets"]["owners"] = [] + await asyncio.to_thread(path.write_text, json.dumps(data), encoding="utf-8") + with pytest.raises(ToolException): + await interceptor(request, handler) + with pytest.raises(ToolException): + await validate_task_agent_reference("alice", "tickets", reference) + assert len(headers) == 1 + finally: + reset_extensions_config() + reset_app_config() + + +@pytest.mark.asyncio +async def test_existing_standalone_tool_observes_enabled_protection_off_loop(agent_auth_env, tmp_path, monkeypatch): + config = agent_auth_env.config.extensions.model_copy(deep=True) + config.mcp_servers["tickets"].agent_auth = False + path = tmp_path / "extensions.json" + await asyncio.to_thread(path.write_text, config.model_dump_json(by_alias=True), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(path)) + reset_app_config() # MCP-only embeddings need no config.yaml. + reset_extensions_config() + interceptor = build_agent_auth_interceptor(config) + request = MCPToolCallRequest(name="search", args={}, server_name="tickets") + calls = [] + + async def handler(request): + calls.append(request) + + try: + await interceptor(request, handler) + updated = config.model_copy(deep=True) + updated.mcp_servers["tickets"].agent_auth = True + await asyncio.to_thread(path.write_text, updated.model_dump_json(by_alias=True), encoding="utf-8") + with pytest.raises(ToolException): + await interceptor(request, handler) + assert len(calls) == 1 + finally: + reset_extensions_config() + reset_app_config() diff --git a/backend/tests/support/mcp_agent_auth.py b/backend/tests/support/mcp_agent_auth.py new file mode 100644 index 0000000000..2e14322591 --- /dev/null +++ b/backend/tests/support/mcp_agent_auth.py @@ -0,0 +1,125 @@ +"""Offline OAuth/MCP endpoints and isolated operator policy for auth tests.""" + +import json +from collections import OrderedDict +from types import SimpleNamespace + +import httpx +import pytest +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.outputs import ChatGeneration, ChatResult + +from deerflow.config.agents_config import AgentConfig +from deerflow.config.app_config import AppConfig, reset_app_config, set_app_config + + +class AgentAuthModel(BaseChatModel): + """A configured model boundary for the real Gateway subprocess test.""" + + @property + def _llm_type(self): + return "offline-agent-auth" + + def bind_tools(self, tools, **kwargs): + return self + + def _generate(self, messages, stop=None, run_manager=None, **kwargs): + answer = AIMessage(content="done") + if not any(isinstance(message, ToolMessage) for message in messages): + answer = AIMessage(content="", tool_calls=[{"name": "tickets_submit_job", "args": {}, "id": "submit-1", "type": "tool_call"}]) + return ChatResult(generations=[ChatGeneration(message=answer)]) + + +@pytest.fixture +def agent_auth_env(tmp_path, monkeypatch): + from deerflow.mcp import agent_auth + from deerflow.mcp.agent_auth import agent_definition_digest + + monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path)) + monkeypatch.setattr("deerflow.config.paths._paths", None) + monkeypatch.setattr(agent_auth, "_token_managers", OrderedDict()) + agent = AgentConfig(name="research-bot", tool_groups=["mcp"], skills=[]) + soul = "Research only." + directory = tmp_path / "users" / "alice" / "agents" / agent.name + directory.mkdir(parents=True) + (directory / "config.yaml").write_text(json.dumps(agent.model_dump(mode="json")), encoding="utf-8") + (directory / "SOUL.md").write_text(soul, encoding="utf-8") + config = AppConfig.model_validate( + { + "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, + "extensions": { + "mcpServers": { + "tickets": { + "type": "http", + "url": "https://mcp.example", + "agent_auth": True, + "headers": {"authorization": "Bearer discovery-secret"}, + } + } + }, + "mcp_agent_auth": { + "research-tickets": { + "server": "tickets", + "agent": agent.name, + "owners": ["alice"], + "definition_digests": [agent_definition_digest(agent, soul)], + "oauth": { + "grant_type": "client_credentials", + "token_url": "https://idp.example/token", + "client_id": "research-client", + "client_secret": "private-client-secret", + }, + } + }, + } + ) + set_app_config(config) + yield SimpleNamespace(agent=agent, soul=soul, directory=directory, config=config) + reset_app_config() + + +@pytest.fixture +def agent_mcp_transport(monkeypatch): + from urllib.parse import parse_qs + + calls, token_clients, token_forms, discoveries = [], [], [], [] + + def respond(request): + if request.url.path == "/token": + form = parse_qs(request.content.decode()) + token_forms.append(form) + client = form["client_id"][0] + token_clients.append(client) + return httpx.Response(200, json={"access_token": client + "-token", "expires_in": 3600}) + if request.method != "POST": + return httpx.Response(405) + message = json.loads(request.content) + if "id" not in message: + return httpx.Response(202) + method = message["method"] + if method == "initialize": + result = {"protocolVersion": message["params"]["protocolVersion"], "capabilities": {"tools": {}}, "serverInfo": {"name": "tickets", "version": "1"}} + elif method == "tools/list": + discoveries.append(request.headers.get("authorization")) + result = {"tools": [{"name": name, "description": name, "inputSchema": {"type": "object", "properties": {}}} for name in ("submit_job", "poll_job", "cancel_job", "search")]} + elif method == "tools/call": + assert len(request.headers.get_list("authorization")) == 1 + name, authorization = message["params"]["name"], request.headers["authorization"] + calls.append((name, authorization)) + assert authorization.startswith("Bearer research-client") or authorization == "Bearer discovery-secret" + result = {"isError": False, "content": [{"type": "text", "text": "ok"}], "structuredContent": {"task_id": "remote-1", "status": "cancelled" if name == "cancel_job" else "running"}} + else: + raise AssertionError(method) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": message["id"], "result": result}) + + original = httpx.AsyncClient + + class OfflineClient(original): + def __init__(self, *args, **kwargs): + kwargs.setdefault("transport", httpx.MockTransport(respond)) + kwargs["trust_env"] = False + super().__init__(*args, **kwargs) + + monkeypatch.setattr(httpx, "AsyncClient", OfflineClient) + return SimpleNamespace(calls=calls, token_clients=token_clients, token_forms=token_forms, discoveries=discoveries) diff --git a/backend/tests/test_mcp_agent_auth.py b/backend/tests/test_mcp_agent_auth.py new file mode 100644 index 0000000000..e48f975e34 --- /dev/null +++ b/backend/tests/test_mcp_agent_auth.py @@ -0,0 +1,621 @@ +"""Administrator-granted MCP identities stay owner-bound and revocable.""" + +import asyncio +import importlib.util +import json +import sys +from datetime import UTC, datetime, timedelta +from pathlib import Path + +import httpx +import pytest +from langchain_core.tools import ToolException +from langchain_mcp_adapters.interceptors import MCPToolCallRequest +from langgraph.runtime import Runtime +from pydantic import ValidationError +from support.mcp_agent_auth import agent_auth_env as agent_auth_env +from support.mcp_agent_auth import agent_mcp_transport as agent_mcp_transport + +from deerflow.config.agents_config import AgentConfig +from deerflow.config.app_config import AppConfig, pop_current_app_config, push_current_app_config, reset_app_config, set_app_config +from deerflow.config.extensions_config import McpServerConfig + + +@pytest.mark.parametrize( + "options", + [ + {"type": "stdio"}, + {"type": "sse"}, + {"user_auth": {"users": {"alice": "secret"}}}, + {"headers_from_context": {"headers": {"aUtHoRiZaTiOn": "token"}}}, + ], +) +def test_agent_auth_rejects_competing_sources_and_unsupported_transports(options): + with pytest.raises(ValidationError): + McpServerConfig.model_validate({"type": "http", "url": "https://mcp.example", "agent_auth": True, **options}) + + +@pytest.fixture +def real_subagents(monkeypatch): + # conftest stubs this module to break collection-time import cycles. + # Restore production execution, including its persistent loop and registry. + path = Path(__file__).parents[1] / "packages/harness/deerflow/subagents/executor.py" + spec = importlib.util.spec_from_file_location("_mcp_auth_executor", path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + for name in ("deerflow.tools.builtins.task_tool", "deerflow.subagents.batch_service"): + target = importlib.import_module(name) + for attribute in ("SubagentExecutor", "SubagentStatus", "cleanup_background_task", "force_cleanup_background_task", "get_background_task_result", "request_cancel_background_task", "run_on_isolated_subagent_loop"): + monkeypatch.setattr(target, attribute, getattr(module, attribute)) + yield module + module._shutdown_isolated_subagent_loop() + + +def test_admission_pins_the_approved_definition(agent_auth_env): + from deerflow.mcp.agent_auth import capture_agent_authority + + env = agent_auth_env + authority = capture_agent_authority(env.config, user_id="alice", agent_config=env.agent) + assert authority is not None + (env.directory / "SOUL.md").write_text("Changed after admission.", encoding="utf-8") + assert authority.soul == env.soul + assert authority.agent_config().model_dump() == env.agent.model_dump() + env.agent.tool_groups.append("unreviewed-group") + assert authority.agent_config().tool_groups == ["mcp"] + with pytest.raises(ToolException): + capture_agent_authority(env.config, user_id="alice", agent_config=env.agent) + + +def test_owners_and_agent_names_both_bound_admission(agent_auth_env): + from deerflow.mcp.agent_auth import capture_agent_authority + + env = agent_auth_env + assert capture_agent_authority(env.config, user_id="bob", agent_config=env.agent) is None + assert capture_agent_authority(env.config, user_id="alice", agent_config=AgentConfig(name="other")) is None + assert capture_agent_authority(env.config, user_id="alice", agent_config=None) is None + + +@pytest.mark.asyncio +async def test_agent_auth_uses_oauth_and_rechecks_revocation_with_a_cached_token(agent_auth_env, monkeypatch): + from deerflow.mcp.agent_auth import AGENT_AUTH_CONTEXT_KEY, build_agent_auth_interceptor, capture_agent_authority + + env = agent_auth_env + authority = capture_agent_authority(env.config, user_id="alice", agent_config=env.agent) + runtime = Runtime(context={"user_id": "alice", AGENT_AUTH_CONTEXT_KEY: authority}) + interceptor = build_agent_auth_interceptor(env.config.extensions) + token_requests, calls = [], [] + + def token_response(request): + token_requests.append(request) + return httpx.Response(200, json={"access_token": "agent-access-token", "expires_in": 3600}) + + real_client = httpx.AsyncClient + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: real_client(transport=httpx.MockTransport(token_response), **kwargs)) + + async def handler(request): + calls.append(request) + return "ok" + + request = MCPToolCallRequest(name="search", args={}, server_name="tickets", runtime=runtime) + assert await interceptor(request, handler) == "ok" + assert await interceptor(request, handler) == "ok" + assert len(token_requests) == 1 + assert b"client_id=research-client" in token_requests[0].content + assert {k.lower(): v for k, v in calls[0].headers.items()}["authorization"] == "Bearer agent-access-token" + revoked = env.config.model_copy(deep=True) + revoked.mcp_agent_auth["research-tickets"].enabled = False + set_app_config(revoked) + push_current_app_config(env.config) + try: + with pytest.raises(ToolException): + await interceptor(request, handler) + finally: + pop_current_app_config() + assert len(calls) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("forged", [None, {"binding_id": "research-tickets"}, {"agent_name": "research-bot", "user_id": "alice"}]) +async def test_context_data_never_grants_an_agent_credential(agent_auth_env, forged): + from deerflow.mcp.agent_auth import AGENT_AUTH_CONTEXT_KEY, build_agent_auth_interceptor + + interceptor = build_agent_auth_interceptor(agent_auth_env.config.extensions) + + async def handler(request): + pytest.fail("An unauthorized request reached the MCP transport") + + runtime = Runtime(context={"user_id": "alice", "agent_name": "research-bot", AGENT_AUTH_CONTEXT_KEY: forged}) + with pytest.raises(ToolException): + await interceptor(MCPToolCallRequest(name="search", args={}, server_name="tickets", runtime=runtime), handler) + + +@pytest.mark.parametrize( + "updates", + [ + {"owners": ["*"]}, + {"owners": "alice"}, + {"owners": ["alice", "alice"]}, + {"owners": [" alice"]}, + {"agent": "../other"}, + {"definition_digests": []}, + {"definition_digests": ["unreviewed"]}, + {"oauth": {"token_url": "https://idp.example/token", "grant_type": "refresh_token", "refresh_token": "not-supported"}}, + {"oauth": {"token_url": "https://idp.example/token", "client_id": "account"}}, + ], +) +def test_invalid_operator_grants_fail_configuration(agent_auth_env, updates): + from deerflow.config.mcp_agent_auth_config import McpAgentCredentialBinding + + binding = agent_auth_env.config.mcp_agent_auth["research-tickets"].model_dump() + with pytest.raises(ValidationError): + McpAgentCredentialBinding.model_validate({**binding, **updates}) + + +def test_ambiguous_bindings_fail_configuration(agent_auth_env): + data = agent_auth_env.config.model_dump() + data["mcp_agent_auth"]["second-grant"] = data["mcp_agent_auth"]["research-tickets"] + with pytest.raises(ValidationError, match="overlap"): + AppConfig.model_validate(data) + + +def test_personal_file_cannot_claim_deployment_credentials(agent_auth_env): + from deerflow.mcp.user_config import load_user_mcp_config, user_mcp_config_path + + path = user_mcp_config_path("alice") + path.parent.mkdir(parents=True) + path.write_text(agent_auth_env.config.extensions.model_dump_json(by_alias=True), encoding="utf-8") + with pytest.raises(ValueError, match="deployment"): + load_user_mcp_config("alice") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("admitted", [False, True]) +async def test_existing_tool_rejects_newly_enabled_protection(agent_auth_env, agent_mcp_transport, admitted): + 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 McpUserScopedAuthConfig + from deerflow.mcp.agent_auth import AGENT_AUTH_CONTEXT_KEY, capture_agent_authority + from deerflow.mcp.tools import get_mcp_tools + + old = agent_auth_env.config.model_copy(deep=True) + old.extensions.mcp_servers["tickets"].agent_auth = False + old.extensions.mcp_servers["tickets"].user_auth = McpUserScopedAuthConfig(users={"alice": "Bearer research-client-old-user"}) + set_app_config(old) + graph = StateGraph(MessagesState, context_schema=dict) + graph.add_node("tools", ToolNode(await get_mcp_tools(old.extensions), handle_tool_errors=False)) + graph.add_edge(START, "tools") + graph.add_edge("tools", END) + compiled = graph.compile() + value = {"messages": [AIMessage(content="", tool_calls=[{"name": "tickets_search", "args": {}, "id": "search-1", "type": "tool_call"}])]} + context = {"user_id": "alice"} + await compiled.ainvoke(value, context=context) + assert agent_mcp_transport.calls == [("search", "Bearer research-client-old-user")] + set_app_config(agent_auth_env.config) + if admitted: + context[AGENT_AUTH_CONTEXT_KEY] = capture_agent_authority(agent_auth_env.config, user_id="alice", agent_config=agent_auth_env.agent) + push_current_app_config(old) + try: + with pytest.raises(ToolException): + await compiled.ainvoke(value, context=context) + finally: + pop_current_app_config() + assert len(agent_mcp_transport.calls) == 1 + assert agent_mcp_transport.token_clients == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("live_destination_changed", [False, True]) +async def test_task_reference_rejects_changed_live_or_captured_destination(agent_auth_env, agent_mcp_transport, live_destination_changed): + from deerflow.mcp.agent_auth import capture_agent_authority + from deerflow.mcp.task_tool_caller import McpTaskConnectionUnavailableError, McpTaskToolCaller + + env = agent_auth_env + reference = capture_agent_authority(env.config, user_id="alice", agent_config=env.agent).references[0].model_dump(mode="json") + replacement = env.config.model_copy(deep=True) + replacement.extensions.mcp_servers["tickets"].url = "https://replacement.example/mcp" + if live_destination_changed: + set_app_config(replacement) + caller = McpTaskToolCaller(replacement.extensions) + with pytest.raises(McpTaskConnectionUnavailableError): + await caller.call_tool(server_name="tickets", tool_name="poll_job", arguments={}, user_id="alice", thread_id="thread-1", agent_auth=reference) + assert agent_mcp_transport.calls == [] + assert agent_mcp_transport.token_clients == [] + + +@pytest.mark.asyncio +async def test_task_reference_without_destination_fails_closed(agent_auth_env, agent_mcp_transport): + from deerflow.mcp.agent_auth import capture_agent_authority + from deerflow.mcp.task_tool_caller import McpTaskConnectionUnavailableError, McpTaskToolCaller + + env = agent_auth_env + reference = capture_agent_authority(env.config, user_id="alice", agent_config=env.agent).references[0].model_dump(mode="json") + del reference["destination_digest"] + with pytest.raises(McpTaskConnectionUnavailableError): + await McpTaskToolCaller(env.config.extensions).call_tool(server_name="tickets", tool_name="poll_job", arguments={}, user_id="alice", thread_id="thread-1", agent_auth=reference) + assert agent_mcp_transport.calls == [] + assert agent_mcp_transport.token_clients == [] + + +@pytest.mark.asyncio +async def test_concurrent_owners_never_share_tokens_and_secret_rotation_keeps_identity(agent_auth_env, agent_mcp_transport, caplog): + from deerflow.mcp.agent_auth import capture_agent_authority + from deerflow.mcp.task_tool_caller import McpTaskConnectionUnavailableError, McpTaskToolCaller + + env = agent_auth_env + bob_dir = env.directory.parents[2] / "bob" / "agents" / env.agent.name + bob_dir.mkdir(parents=True) + (bob_dir / "SOUL.md").write_text(env.soul, encoding="utf-8") + binding = env.config.mcp_agent_auth["research-tickets"].model_copy(deep=True) + binding.owners = ["bob"] + binding.oauth.client_id = "research-client-bob" + env.config.mcp_agent_auth["bob-tickets"] = binding + references = {user: capture_agent_authority(env.config, user_id=user, agent_config=env.agent).references[0].model_dump(mode="json") for user in ("alice", "bob")} + caller = McpTaskToolCaller(env.config.extensions) + + async def call(user): + return await caller.call_tool(server_name="tickets", tool_name="search", arguments={}, user_id=user, thread_id="thread-" + user, agent_auth=references[user]) + + with caplog.at_level("INFO", logger="deerflow.mcp.agent_auth"): + await asyncio.gather(call("alice"), call("bob"), call("alice"), call("bob")) + assert sorted(agent_mcp_transport.token_clients) == ["research-client", "research-client-bob"] + assert sorted(agent_mcp_transport.calls) == sorted([("search", "Bearer research-client-token")] * 2 + [("search", "Bearer research-client-bob-token")] * 2) + assert "server=tickets agent=research-bot owner=alice binding=research-tickets" in caplog.text + assert "private-client-secret" not in caplog.text and "research-client-token" not in caplog.text + with pytest.raises(McpTaskConnectionUnavailableError): + await caller.call_tool(server_name="tickets", tool_name="search", arguments={}, user_id="bob", thread_id="thread-bob", agent_auth=references["alice"]) + assert len(agent_mcp_transport.calls) == 4 + env.config.mcp_agent_auth["research-tickets"].oauth.client_secret = "rotated-secret" + await call("alice") + assert agent_mcp_transport.token_forms[-1]["client_secret"] == ["rotated-secret"] + # Retargeting the binding to a different client must not impersonate it + # for tasks admitted under the original account. + env.config.mcp_agent_auth["research-tickets"].oauth.client_id = "replacement-account" + with pytest.raises(McpTaskConnectionUnavailableError): + await call("alice") + assert len(agent_mcp_transport.calls) == 5 + assert len(agent_mcp_transport.token_clients) == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revoke_during_fetch", [False, True]) +async def test_token_failure_and_concurrent_revocation_never_fall_back(agent_auth_env, monkeypatch, revoke_during_fetch): + from deerflow.mcp.agent_auth import AGENT_AUTH_CONTEXT_KEY, build_agent_auth_interceptor, capture_agent_authority + + env = agent_auth_env + authority = capture_agent_authority(env.config, user_id="alice", agent_config=env.agent) + started, release = asyncio.Event(), asyncio.Event() + + async def token_endpoint(request): + started.set() + await release.wait() + if revoke_during_fetch: + return httpx.Response(200, json={"access_token": "never-send-this-token", "expires_in": 3600}) + return httpx.Response(401, json={"error": "private-client-secret"}) + + original = httpx.AsyncClient + monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: original(transport=httpx.MockTransport(token_endpoint), **kwargs)) + interceptor = build_agent_auth_interceptor(env.config.extensions) + + async def handler(request): + pytest.fail("Credential failure must not fall back to discovery auth") + + request = MCPToolCallRequest(name="search", args={}, server_name="tickets", runtime=Runtime(context={"user_id": "alice", AGENT_AUTH_CONTEXT_KEY: authority})) + pending = asyncio.create_task(interceptor(request, handler)) + try: + await asyncio.wait_for(started.wait(), timeout=5) + if revoke_during_fetch: + env.config.mcp_agent_auth["research-tickets"].enabled = False + release.set() + with pytest.raises(ToolException) as caught: + await pending + assert "private-client-secret" not in str(caught.value) + assert "never-send-this-token" not in str(caught.value) + finally: + release.set() + await asyncio.gather(pending, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["poll", "cancel"]) +@pytest.mark.parametrize("revoked", [False, True]) +async def test_durable_task_keeps_agent_identity_after_restart_and_definition_rollout(agent_auth_env, agent_mcp_transport, tmp_path, monkeypatch, operation, revoked): + from langchain_core.messages import AIMessage + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.prebuilt import ToolNode + + from app.mcp_tasks import McpTaskService + from deerflow.config.database_config import DatabaseConfig + from deerflow.config.extensions_config import McpTaskToolsetConfig, reset_extensions_config + from deerflow.mcp import agent_auth + from deerflow.mcp.agent_auth import AGENT_AUTH_CONTEXT_KEY, agent_definition_digest, capture_agent_authority + from deerflow.mcp.task_tool_caller import McpTaskToolCaller + from deerflow.mcp.tasks import ORDINARY_MCP_TASK_DRIVER, McpTaskDriverRegistry, OrdinaryMcpTaskDriver + from deerflow.mcp.tasks.runtime import set_mcp_task_config_snapshot, set_mcp_task_submitter, validate_mcp_task_config_snapshot + from deerflow.mcp.tools import get_mcp_tools + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config + from deerflow.persistence.mcp_tasks import McpTaskRepository + from deerflow.persistence.thread_meta.model import ThreadMetaRow + + env = agent_auth_env + server = env.config.extensions.mcp_servers["tickets"] + server.task_toolsets = [McpTaskToolsetConfig(name="job", submit_tool="submit_job", status_tool="poll_job", cancel_tool="cancel_job")] + config_path = tmp_path / "extensions_config.json" + config_path.write_text(env.config.extensions.model_dump_json(by_alias=True), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(config_path)) + reset_extensions_config() + startup = env.config.extensions + set_mcp_task_config_snapshot(startup) + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path / "database"))) + sf = get_session_factory() + repo = McpTaskRepository(sf) + async with sf() as session: + session.add(ThreadMetaRow(thread_id="t-agent", incarnation="inc-1", user_id="alice", metadata_json={})) + await session.commit() + + def service(): + drivers = McpTaskDriverRegistry() + drivers.register(ORDINARY_MCP_TASK_DRIVER, OrdinaryMcpTaskDriver(McpTaskToolCaller(startup))) + return McpTaskService(repository=repo, drivers=drivers, poll_interval_seconds=1, lease_seconds=120, max_concurrent_polls=1) + + try: + first = service() + set_mcp_task_submitter(first) + tools = await get_mcp_tools() + 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) + authority = capture_agent_authority(env.config, user_id="alice", agent_config=env.agent) + result = await graph.compile().ainvoke( + {"messages": [AIMessage(content="", tool_calls=[{"name": "tickets_submit_job", "args": {}, "id": "submit-1", "type": "tool_call"}])]}, + context={"user_id": "alice", "thread_id": "t-agent", "thread_incarnation": "inc-1", AGENT_AUTH_CONTEXT_KEY: authority}, + ) + task_id = json.loads(result["messages"][-1].content)["task_id"] + saved = await repo.get(task_id, user_id="alice", thread_id="t-agent", thread_incarnation="inc-1") + reference = saved["driver_data"]["agent_auth"] + assert reference["definition_digest"] == authority.references[0].definition_digest + assert "private-client-secret" not in json.dumps(saved, default=str) + assert "research-client-token" not in json.dumps(saved, default=str) + + updated = env.config.model_copy(deep=True) + updated.mcp_agent_auth["research-tickets"].definition_digests = [agent_definition_digest(env.agent, "New reviewed definition.")] + if revoked: + updated.mcp_agent_auth["research-tickets"].owners = [] + set_app_config(updated) + validate_mcp_task_config_snapshot(updated.extensions) + agent_auth._token_managers.clear() + restarted = service() + if operation == "cancel": + await restarted.cancel_task(task_id=task_id, thread_id="t-agent", user_id="alice", thread_incarnation="inc-1") + await restarted.run_once(now=datetime.now(UTC) + timedelta(seconds=2)) + saved = await repo.get(task_id, user_id="alice", thread_id="t-agent", thread_incarnation="inc-1") + if revoked: + assert saved["status"] == "failed" + assert "revoked" in saved["error"] + assert agent_mcp_transport.calls == [("submit_job", "Bearer research-client-token")] + assert agent_mcp_transport.token_clients == ["research-client"] + else: + assert agent_mcp_transport.calls == [("submit_job", "Bearer research-client-token"), (operation + "_job", "Bearer research-client-token")] + assert saved["status"] == ("cancelled" if operation == "cancel" else "working") + assert agent_mcp_transport.token_clients == ["research-client", "research-client"] + finally: + set_mcp_task_submitter(None) + set_mcp_task_config_snapshot(None) + reset_extensions_config() + await close_engine() + + +@pytest.mark.no_auto_user +def test_gateway_admits_owner_and_assembles_the_verified_snapshot(agent_auth_env, agent_mcp_transport, tmp_path, monkeypatch): + from _agent_e2e_helpers import FakeToolCallingModel + from langchain_core.messages import AIMessage + from starlette.testclient import TestClient + + from app.gateway.app import create_app + from deerflow.config.agents_config import load_agent_config + from deerflow.config.extensions_config import reset_extensions_config + from deerflow.mcp.agent_auth import AGENT_AUTH_CONTEXT_KEY, agent_definition_digest + + env = agent_auth_env + data = env.config.model_dump(mode="json", exclude={"extensions"}) + data.update( + models=[{"name": "offline", "use": "langchain_openai:ChatOpenAI", "model": "offline", "api_key": "unused"}], + agents_api={"enabled": True}, + database={"backend": "sqlite", "sqlite_dir": str(tmp_path / "gateway-db")}, + title={"enabled": False}, + memory={"enabled": False}, + summarization={"enabled": False}, + tool_search={"enabled": False}, + ) + config_path = tmp_path / "config.yaml" + extensions_path = tmp_path / "extensions_config.json" + config_path.write_text(json.dumps(data), encoding="utf-8") + extensions_path.write_text(env.config.extensions.model_dump_json(by_alias=True), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(config_path)) + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(extensions_path)) + monkeypatch.setenv("AUTH_JWT_SECRET", "isolated-agent-auth-test-secret-1234567890") + monkeypatch.setattr("app.gateway.auth.config._auth_config", None) + monkeypatch.setattr("app.gateway.deps._cached_local_provider", None) + monkeypatch.setattr("app.gateway.deps._cached_repo", None) + reset_app_config() + reset_extensions_config() + prompts = [] + mutate_definition = [] + + class Model(FakeToolCallingModel): + def _generate(self, messages, *args, **kwargs): + prompts.append(messages) + return super()._generate(messages, *args, **kwargs) + + def model_factory(*args, **kwargs): + if mutate_definition: + mutate_definition.pop()() + return Model(responses=[AIMessage(content="", tool_calls=[{"name": "tickets_search", "args": {}, "id": "call-search", "type": "tool_call"}]), AIMessage(content="done")]) + + monkeypatch.setattr("deerflow.agents.lead_agent.agent.create_chat_model", model_factory) + try: + with TestClient(create_app()) as client: + + def register(email): + client.cookies.clear() + response = client.post("/api/v1/auth/register", json={"email": email, "password": "strong-test-password-123"}) + assert response.status_code == 201, response.text + client.headers["X-CSRF-Token"] = client.cookies["csrf_token"] + created = client.post("/api/agents", json={"name": env.agent.name, "soul": env.soul, "skills": [], "tool_groups": ["mcp"]}) + assert created.status_code == 201, created.text + return response.json()["id"] + + def run(context): + thread = client.post("/api/threads", json={}) + assert thread.status_code == 200, thread.text + return client.post( + f"/api/threads/{thread.json()['thread_id']}/runs/wait", + json={ + "assistant_id": "lead_agent", + "input": {"messages": [{"role": "user", "content": "Search tickets."}]}, + "context": {"mode": "flash", "subagent_enabled": False, **context}, + "config": {"configurable": {AGENT_AUTH_CONTEXT_KEY: {"binding": "research-tickets", "user_id": "alice"}}}, + }, + ) + + owner = register("owner@example.com") + config = load_agent_config(env.agent.name, user_id=owner) + binding = data["mcp_agent_auth"]["research-tickets"] + binding["owners"] = [owner] + binding["definition_digests"] = [agent_definition_digest(config, env.soul)] + config_path.write_text(json.dumps(data), encoding="utf-8") + directory = tmp_path / "users" / owner / "agents" / env.agent.name + mutate_definition.append(lambda: (directory / "SOUL.md").write_text("Unreviewed replacement.", encoding="utf-8")) + response = run({"agent_name": env.agent.name}) + assert response.status_code == 200, response.text + assert agent_mcp_transport.calls == [("search", "Bearer research-client-token")], response.text + assert "Bearer discovery-secret" in agent_mcp_transport.discoveries + system = str(prompts[0][0].content) + assert env.soul in system and "Unreviewed replacement." not in system + response = run({"agent_name": env.agent.name}) + assert response.status_code == 403, response.text + assert len(agent_mcp_transport.calls) == 1 + + (directory / "SOUL.md").write_text(env.soul, encoding="utf-8") + for context in ({"agent_name": env.agent.name, "is_bootstrap": True}, {}): + response = run(context) + assert response.status_code == 200, response.text + assert len(agent_mcp_transport.calls) == 1 + register("other-owner@example.com") + denied = client.post("/api/mcp/personal/config/servers", json={"mcp_servers": {"tickets": {"type": "http", "url": "https://mcp.example", "agent_auth": True}}}) + assert denied.status_code == 422, denied.text + response = run({"agent_name": env.agent.name, AGENT_AUTH_CONTEXT_KEY: {"binding": "research-tickets", "user_id": owner}}) + assert response.status_code == 200, response.text + assert len(agent_mcp_transport.calls) == 1 + assert agent_mcp_transport.token_clients == ["research-client"] + finally: + reset_extensions_config() + reset_app_config() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delegation", ["task", "batch_task"]) +@pytest.mark.parametrize("revoked", [False, True]) +async def test_delegation_inherits_parent_identity_including_durable_recovery(agent_auth_env, agent_mcp_transport, real_subagents, tmp_path, monkeypatch, delegation, revoked): + from _agent_e2e_helpers import FakeToolCallingModel, build_single_tool_call_model + from langchain_core.messages import AIMessage + from langgraph.graph import END, START, StateGraph + from langgraph.prebuilt import ToolNode + + from deerflow.agents.thread_state import ThreadState + from deerflow.config.database_config import DatabaseConfig + from deerflow.config.extensions_config import reset_extensions_config + from deerflow.config.model_config import ModelConfig + from deerflow.mcp.agent_auth import AGENT_AUTH_CONTEXT_KEY, capture_agent_authority + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config + from deerflow.persistence.subagent_batches import SubagentBatchRepository, SubagentBatchRow + from deerflow.subagents.batch_runtime import set_subagent_batch_submitter + from deerflow.subagents.batch_service import SubagentBatchService + from deerflow.tools.builtins.batch_task_tool import batch_task + from deerflow.tools.builtins.task_tool import task_tool + + env = agent_auth_env + env.config.models = [ModelConfig(name="offline", use="langchain_openai:ChatOpenAI", model="offline", api_key="unused")] + env.config.tool_search.enabled = False + env.config.memory.enabled = False + env.config.summarization.enabled = False + env.config.title.enabled = False + env.config.subagent_batches.enabled = True + # A child with its own admin binding must still use its caller's identity. + child_binding = env.config.mcp_agent_auth["research-tickets"].model_copy(deep=True) + child_binding.agent = "general-purpose" + child_binding.oauth.client_id = "other-child-account" + env.config.mcp_agent_auth["child-tickets"] = child_binding + authority = capture_agent_authority(env.config, user_id="alice", agent_config=env.agent) + config_path = tmp_path / "extensions_config.json" + config_path.write_text(env.config.extensions.model_dump_json(by_alias=True), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(config_path)) + reset_extensions_config() + observed = [] + + class ChildModel(FakeToolCallingModel): + def _generate(self, messages, *args, **kwargs): + observed.extend(message for message in messages if message.type == "tool") + return super()._generate(messages, *args, **kwargs) + + monkeypatch.setattr(real_subagents, "create_chat_model", lambda **kwargs: ChildModel(responses=build_single_tool_call_model(tool_name="tickets_search", tool_args={}).responses)) + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path / "batch-db"))) + repo = SubagentBatchRepository(get_session_factory()) + + def service(): + return SubagentBatchService(repository=repo, config=env.config.subagent_batches, runtime_config=env.config.subagent_runtime, app_config=env.config) + + first, restarted = service(), None + set_subagent_batch_submitter(first) + try: + graph = StateGraph(ThreadState, context_schema=dict) + graph.add_node("tools", ToolNode([task_tool, batch_task], handle_tool_errors=False)) + graph.add_edge(START, "tools") + graph.add_edge("tools", END) + args = {"subagent_type": "general-purpose"} + if delegation == "task": + args.update(description="Search", prompt="Search tickets.") + if revoked: + env.config.mcp_agent_auth["research-tickets"].enabled = False + else: + args.update(title="Search", items=[{"key": "item-1", "prompt": "Search tickets."}]) + result = await asyncio.wait_for( + graph.compile().ainvoke( + {"messages": [AIMessage(content="", tool_calls=[{"name": delegation, "args": args, "id": "delegate-1", "type": "tool_call"}])]}, + config={"metadata": {"model_name": "offline", "tool_groups": ["mcp"], "available_skills": []}}, + context={"user_id": "alice", "thread_id": "t-delegate", "run_id": "r-delegate", AGENT_AUTH_CONTEXT_KEY: authority}, + ), + timeout=30, + ) + message = result["messages"][-1] + if delegation == "batch_task": + batch_id = message.additional_kwargs["subagent_batch_id"] + async with get_session_factory()() as session: + saved_spec = (await session.get(SubagentBatchRow, batch_id)).execution_spec + references = saved_spec["mcp_agent_auth"] + assert references[0]["binding"] == "research-tickets" + assert "private-client-secret" not in json.dumps(saved_spec) + if revoked: + env.config.mcp_agent_auth["research-tickets"].enabled = False + restarted = service() + await restarted.run_once(now=datetime.now(UTC)) + await asyncio.wait_for(asyncio.gather(*list(restarted._executions.values())), timeout=30) + saved = await repo.get_batch(batch_id, user_id="alice") + assert saved["counts"]["succeeded"] == 1, saved + if revoked: + assert any("revoked" in str(message.content) for message in observed), observed + assert agent_mcp_transport.calls == [] + assert agent_mcp_transport.token_clients == [] + else: + assert agent_mcp_transport.calls == [("search", "Bearer research-client-token")], message + assert agent_mcp_transport.token_clients == ["research-client"] + finally: + set_subagent_batch_submitter(None) + await first.stop() + if restarted is not None: + await restarted.stop() + reset_extensions_config() + await close_engine() diff --git a/backend/tests/test_mcp_agent_auth_restart.py b/backend/tests/test_mcp_agent_auth_restart.py new file mode 100644 index 0000000000..5844815957 --- /dev/null +++ b/backend/tests/test_mcp_agent_auth_restart.py @@ -0,0 +1,267 @@ +"""Kill a real Gateway during a task call and recover through its HTTP API. + +Only the model and remote OAuth/MCP service are deterministic fixtures. Auth, +admission, task submission, leases, SQLite recovery and both Gateway lifespans +run in production code, in distinct OS processes over loopback TCP. +""" + +import json +import os +import signal +import socket +import sqlite3 +import subprocess +import sys +import threading +import time +from contextlib import closing, contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from types import SimpleNamespace +from urllib.parse import parse_qs + +import httpx +import pytest +from support.mcp_agent_auth import agent_auth_env as agent_auth_env + +from deerflow.config.agents_config import load_agent_config +from deerflow.mcp.agent_auth import agent_definition_digest + +ROOT = Path(__file__).resolve().parents[2] +pytestmark = pytest.mark.skipif(os.name != "posix", reason="Verifies SIGKILL and OS-process lease recovery") + + +def _until(check, *, seconds=30): + deadline = time.monotonic() + seconds + while time.monotonic() < deadline: + result = check() + if result: + return result + # This polls observable HTTP state; correctness never depends on a delay. + time.sleep(0.05) + raise AssertionError(f"Timed out waiting for {check}") + + +@contextmanager +def _remote_service(operation): + state = SimpleNamespace(generation=1, reached=threading.Event(), release=threading.Event(), calls=[], tokens=[]) + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *_args): + pass + + def reply(self, status, payload=None): + body = json.dumps(payload).encode() if payload is not None else b"" + try: + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + except (BrokenPipeError, ConnectionResetError): + pass # The killed Gateway cannot receive the parked response. + + def do_GET(self): + self.reply(405) + + def do_POST(self): + body = self.rfile.read(int(self.headers["Content-Length"])) + if self.path == "/token": + form = parse_qs(body.decode()) + if form.get("client_id") != ["research-client"] or form.get("client_secret") != ["test-client-secret"]: + return self.reply(401) + token = f"research-client-{len(state.tokens) + 1}" + state.tokens.append(token) + return self.reply(200, {"access_token": token, "expires_in": 3600}) + message = json.loads(body) + if "id" not in message: + return self.reply(202) + method = message["method"] + if method == "initialize": + result = {"protocolVersion": message["params"]["protocolVersion"], "capabilities": {"tools": {}}, "serverInfo": {"name": "restart-test", "version": "1"}} + elif method == "tools/list": + result = {"tools": [{"name": name, "description": name, "inputSchema": {"type": "object", "properties": {} if name == "submit_job" else {"task_id": {"type": "string"}}}} for name in ("submit_job", "poll_job", "cancel_job")]} + elif method == "tools/call": + name = message["params"]["name"] + headers = self.headers.get_all("Authorization") or [] + state.calls.append({"generation": state.generation, "tool": name, "authorization": headers}) + if len(headers) != 1 or headers[0] not in ["Bearer " + token for token in state.tokens]: + return self.reply(401) + if state.generation == 1 and name == operation + "_job": + state.reached.set() + if not state.release.wait(timeout=45): + return self.reply(500, {"error": "Test barrier timed out"}) + result = {"isError": False, "content": [], "structuredContent": {"task_id": "remote-1", "status": "cancelled" if name == "cancel_job" else "running", "poll_after_seconds": 1}} + else: + return self.reply(400) + self.reply(200, {"jsonrpc": "2.0", "id": message["id"], "result": result}) + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}", state + finally: + state.release.set() + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +@contextmanager +def _gateway(tmp_path, client, port, generation): + # Exclude ambient provider credentials and point dotenv/config/home at + # temporary files. The child resolves its model through normal config. + env = {key: os.environ[key] for key in ("PATH", "HOME", "TMPDIR", "LANG") if key in os.environ} + env.update( + PYTHONPATH=os.pathsep.join((str(ROOT / "backend"), str(ROOT / "backend/tests"))), + DEER_FLOW_HOME=str(tmp_path), + DEER_FLOW_PROJECT_ROOT=str(ROOT), + DEER_FLOW_CONFIG_PATH=str(tmp_path / "config.yaml"), + DEER_FLOW_EXTENSIONS_CONFIG_PATH=str(tmp_path / "extensions.json"), + DEER_FLOW_ENV_FILE=str(tmp_path / "empty.env"), + DEER_FLOW_AUTH_DISABLED="0", + GATEWAY_WORKERS="1", + WEB_CONCURRENCY="1", + AUTH_JWT_SECRET="isolated-gateway-restart-test-secret-1234567890", + NO_PROXY="127.0.0.1,localhost", # Also overrides macOS system proxies. + PYTHONIOENCODING="utf-8", + ) + log_path = tmp_path / f"gateway-{generation}.log" + with log_path.open("w", encoding="utf-8") as log: + process = subprocess.Popen([sys.executable, "-m", "uvicorn", "app.gateway.app:app", "--host", "127.0.0.1", "--port", str(port), "--log-level", "warning"], cwd=tmp_path, env=env, stdout=log, stderr=subprocess.STDOUT) + try: + + def healthy(): + assert process.poll() is None, log_path.read_text(encoding="utf-8")[-8000:] + try: + return client.get("/health").status_code == 200 + except httpx.TransportError: + return False + + _until(healthy, seconds=45) + yield process + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=10) + + +def _write_config(path, data): + replacement = path.with_suffix(".new") + replacement.write_text(json.dumps(data), encoding="utf-8") + replacement.replace(path) + + +@pytest.mark.parametrize("operation", ["poll", "cancel"]) +@pytest.mark.parametrize("policy_change", ["definition", "revocation", "destination"]) +def test_gateway_process_restart_keeps_task_identity_and_enforces_revocation(agent_auth_env, tmp_path, operation, policy_change): + revoked = policy_change == "revocation" + denied = policy_change != "definition" + env = agent_auth_env + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + with _remote_service(operation) as (remote_url, remote), httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=30, trust_env=False) as client: + data = env.config.model_dump(mode="json", exclude={"extensions"}) + data.update( + models=[{"name": "offline", "use": "support.mcp_agent_auth:AgentAuthModel", "model": "offline", "supports_tool_calling": True}], + agents_api={"enabled": True}, + database={"backend": "sqlite", "sqlite_dir": str(tmp_path / "db")}, + title={"enabled": False}, + memory={"enabled": False}, + summarization={"enabled": False}, + tool_search={"enabled": False}, + mcp_tasks={"enabled": True, "poll_interval_seconds": 1, "lease_seconds": 5}, + ) + binding = data["mcp_agent_auth"]["research-tickets"] + binding["oauth"].update(token_url=remote_url + "/token", client_secret="test-client-secret") + config_path = tmp_path / "config.yaml" + _write_config(config_path, data) + server = env.config.extensions.mcp_servers["tickets"].model_dump(mode="json") + server.update(url=remote_url + "/mcp", task_toolsets=[{"name": "job", "submit_tool": "submit_job", "status_tool": "poll_job", "cancel_tool": "cancel_job"}]) + (tmp_path / "extensions.json").write_text(json.dumps({"mcpServers": {"tickets": server}}), encoding="utf-8") + (tmp_path / "empty.env").touch() + + def request(method, path, *, expected=200, **kwargs): + response = client.request(method, path, **kwargs) + assert response.status_code == expected, response.text + return response.json() + + with _gateway(tmp_path, client, port, 1) as first: + owner = request("POST", "/api/v1/auth/register", expected=201, json={"email": "restart@example.com", "password": "isolated-test-password-123"})["id"] + client.headers["X-CSRF-Token"] = client.cookies["csrf_token"] + request("POST", "/api/agents", expected=201, json={"name": env.agent.name, "soul": env.soul, "tool_groups": ["mcp"], "skills": []}) + agent = load_agent_config(env.agent.name, user_id=owner) + original_digest = agent_definition_digest(agent, env.soul) + binding.update(owners=[owner], definition_digests=[original_digest]) + _write_config(config_path, data) + thread_id = request("POST", "/api/threads", json={})["thread_id"] + request( + "POST", + f"/api/threads/{thread_id}/runs/wait", + json={"assistant_id": "lead_agent", "input": {"messages": [{"role": "user", "content": "Submit a job."}]}, "context": {"agent_name": env.agent.name, "mode": "flash", "subagent_enabled": False}}, + ) + tasks = request("GET", f"/api/threads/{thread_id}/mcp-tasks") + assert len(tasks) == 1, tasks + task_url = f"/api/threads/{thread_id}/mcp-tasks/{tasks[0]['task_id']}" + if operation == "cancel": + request("POST", task_url + "/cancel") + assert remote.reached.wait(timeout=20), remote.calls + first.kill() + first.wait(timeout=10) + assert first.returncode == -signal.SIGKILL + + # Inspect production writes only; never manufacture a recoverable row. + with closing(sqlite3.connect(tmp_path / "db/deerflow.db")) as connection: + driver_data = json.loads(connection.execute("SELECT driver_data FROM mcp_tasks WHERE id = ?", (tasks[0]["task_id"],)).fetchone()[0]) + assert driver_data["agent_auth"]["definition_digest"] == original_digest + assert "test-client-secret" not in json.dumps(driver_data) + + new_soul = "Reviewed definition revision two." + (tmp_path / "users" / owner / "agents" / env.agent.name / "SOUL.md").write_text(new_soul, encoding="utf-8") + binding["definition_digests"] = [agent_definition_digest(agent, new_soul)] + binding["enabled"] = not revoked + _write_config(config_path, data) + if policy_change == "destination": + server["url"] = remote_url + "/replacement-mcp" + _write_config(tmp_path / "extensions.json", {"mcpServers": {"tickets": server}}) + remote.generation = 2 + remote.release.set() + with _gateway(tmp_path, client, port, 2) as second: + assert second.pid != first.pid + target_status = "failed" if denied else "cancelled" if operation == "cancel" else "working" + + def recovered(): + record = request("GET", task_url) + sent = [call for call in remote.calls if call["generation"] == 2 and call["tool"] == operation + "_job"] + return record if record["status"] == target_status and (denied or sent) else None + + result = _until(recovered) + if operation == "cancel": + assert result["cancel_attempt_count"] >= 2, result + later_calls = [call for call in remote.calls if call["generation"] == 2] + if denied: + assert result["error"], result + assert later_calls == [] + assert remote.tokens == ["research-client-1"] + else: + assert any(call["tool"] == operation + "_job" for call in later_calls) + assert all(call["authorization"] == ["Bearer research-client-2"] for call in later_calls) + assert remote.tokens == ["research-client-1", "research-client-2"] + (tmp_path / "evidence.json").write_text( + json.dumps( + { + "operation": operation, + "policy_change": policy_change, + "old_pid": first.pid, + "new_pid": second.pid, + "old_exit": first.returncode, + "task_status": result["status"], + "token_acquisitions": len(remote.tokens), + "calls": remote.calls, + }, + indent=2, + ), + encoding="utf-8", + ) diff --git a/backend/tests/test_mcp_config_secrets.py b/backend/tests/test_mcp_config_secrets.py index 9e80cc2757..318e905be0 100644 --- a/backend/tests/test_mcp_config_secrets.py +++ b/backend/tests/test_mcp_config_secrets.py @@ -220,6 +220,33 @@ def test_merge_preserves_masked_env_values(): assert merged.env["KEY"] == "real_secret" +def test_bulk_merge_preserves_agent_auth_until_explicitly_changed(): + existing = McpServerConfigResponse(type="http", url="https://mcp.example", agent_auth=True) + assert _merge_preserving_secrets(McpServerConfigResponse(description="Renamed"), existing).agent_auth is True + assert _merge_preserving_secrets(McpServerConfigResponse(agent_auth=False), existing).agent_auth is False + + +@pytest.mark.asyncio +async def test_targeted_edit_preserves_agent_auth_until_explicitly_changed(monkeypatch, tmp_path): + from deerflow.config.extensions_config import reset_extensions_config + + path = tmp_path / "extensions_config.json" + path.write_text(json.dumps({"mcpServers": {"target": {"type": "http", "url": "https://mcp.example", "agent_auth": True}}}), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(path)) + _track_mcp_reconciliation_handoff(monkeypatch) + reset_extensions_config() + try: + edit = {"type": "http", "url": "https://mcp.example", "description": "Updated description"} + response = await update_mcp_server(_request_with_role("admin"), McpServerConfigUpdateRequest(server_name="target", server=McpServerConfigResponse.model_validate(edit))) + assert response.mcp_servers["target"].agent_auth is True + assert json.loads(path.read_text(encoding="utf-8"))["mcpServers"]["target"]["agent_auth"] is True + edit["agent_auth"] = False + await update_mcp_server(_request_with_role("admin"), McpServerConfigUpdateRequest(server_name="target", server=McpServerConfigResponse.model_validate(edit))) + assert json.loads(path.read_text(encoding="utf-8"))["mcpServers"]["target"]["agent_auth"] is False + finally: + reset_extensions_config() + + def test_merge_preserves_masked_header_values(): """Incoming '***' header values should be replaced with existing secrets.""" incoming = McpServerConfigResponse(headers={"Authorization": "***"}) diff --git a/backend/tests/test_mcp_context_headers.py b/backend/tests/test_mcp_context_headers.py index cad590ab8f..29a771d64e 100644 --- a/backend/tests/test_mcp_context_headers.py +++ b/backend/tests/test_mcp_context_headers.py @@ -503,6 +503,7 @@ def test_registered_last_so_request_secrets_win(): interceptors = build_mcp_tool_interceptors(config, oauth_builder=lambda _cfg: oauth) assert [getattr(i, "__name__", type(i).__name__) for i in interceptors] == [ "oauth", + "agent_auth_interceptor", "user_scoped_auth_interceptor", "context_headers_interceptor", ] @@ -513,7 +514,7 @@ def test_shared_assembly_skips_when_not_configured(): mcp_servers={"plain": McpServerConfig(enabled=True, type="http", url="https://x.example.com")}, skills={}, ) - assert build_mcp_tool_interceptors(config, oauth_builder=lambda _cfg: None) == [] + assert [i.__name__ for i in build_mcp_tool_interceptors(config, oauth_builder=lambda _cfg: None)] == ["agent_auth_interceptor"] # --------------------------------------------------------------------------- diff --git a/backend/tests/test_mcp_custom_interceptors.py b/backend/tests/test_mcp_custom_interceptors.py index 1c5236f13b..651bf9c321 100644 --- a/backend/tests/test_mcp_custom_interceptors.py +++ b/backend/tests/test_mcp_custom_interceptors.py @@ -80,8 +80,9 @@ def test_custom_interceptor_loaded_and_appended(): asyncio.run(get_mcp_tools()) interceptors = _get_interceptors(mock_cls) - assert len(interceptors) == 1 - assert interceptors[0] is fake_interceptor + assert len(interceptors) == 2 + assert interceptors[0].__name__ == "agent_auth_interceptor" + assert interceptors[1] is fake_interceptor def test_multiple_custom_interceptors(): @@ -111,9 +112,9 @@ def test_multiple_custom_interceptors(): asyncio.run(get_mcp_tools()) interceptors = _get_interceptors(mock_cls) - assert len(interceptors) == 2 - assert interceptors[0] is interceptor_a - assert interceptors[1] is interceptor_b + assert len(interceptors) == 3 + assert interceptors[1] is interceptor_a + assert interceptors[2] is interceptor_b def test_custom_interceptor_builder_returning_none_is_skipped(): @@ -130,7 +131,7 @@ def test_custom_interceptor_builder_returning_none_is_skipped(): ): asyncio.run(get_mcp_tools()) - assert len(_get_interceptors(mock_cls)) == 0 + assert [i.__name__ for i in _get_interceptors(mock_cls)] == ["agent_auth_interceptor"] def test_custom_interceptor_resolve_error_logs_warning_and_continues(): @@ -178,7 +179,7 @@ def test_custom_interceptor_builder_exception_logs_warning_and_continues(): def test_no_mcp_interceptors_field_is_safe(): - """When mcpInterceptors is absent from config, no interceptors are added.""" + """Without custom builders, the live protection guard remains installed.""" p = _make_patches(interceptor_paths=None) with ( @@ -190,7 +191,7 @@ def test_no_mcp_interceptors_field_is_safe(): ): asyncio.run(get_mcp_tools()) - assert len(_get_interceptors(mock_cls)) == 0 + assert [i.__name__ for i in _get_interceptors(mock_cls)] == ["agent_auth_interceptor"] def test_custom_interceptor_coexists_with_oauth_interceptor(): @@ -215,9 +216,9 @@ def test_custom_interceptor_coexists_with_oauth_interceptor(): asyncio.run(get_mcp_tools()) interceptors = _get_interceptors(mock_cls) - assert len(interceptors) == 2 + assert len(interceptors) == 3 assert interceptors[0] is oauth_fn - assert interceptors[1] is custom_fn + assert interceptors[2] is custom_fn def test_mcp_interceptors_single_string_is_normalized(): @@ -238,7 +239,8 @@ def test_mcp_interceptors_single_string_is_normalized(): ): asyncio.run(get_mcp_tools()) - assert len(_get_interceptors(mock_cls)) == 1 + assert len(_get_interceptors(mock_cls)) == 2 + assert _get_interceptors(mock_cls)[-1] is fake_interceptor def test_mcp_interceptors_invalid_type_logs_warning(): @@ -255,7 +257,7 @@ def test_mcp_interceptors_invalid_type_logs_warning(): ): asyncio.run(get_mcp_tools()) - assert len(_get_interceptors(mock_cls)) == 0 + assert [i.__name__ for i in _get_interceptors(mock_cls)] == ["agent_auth_interceptor"] mock_warn.assert_called_once() assert "must be a list" in mock_warn.call_args[0][0] @@ -275,7 +277,7 @@ def test_custom_interceptor_non_callable_return_logs_warning(): ): asyncio.run(get_mcp_tools()) - assert len(_get_interceptors(mock_cls)) == 0 + assert [i.__name__ for i in _get_interceptors(mock_cls)] == ["agent_auth_interceptor"] mock_warn.assert_called_once() assert "non-callable" in mock_warn.call_args[0][0] @@ -325,7 +327,7 @@ def test_custom_interceptor_builder_runs_once_per_enabled_server(): assert len(built) == 2, f"builder ran {len(built)} times, expected once per server" assert len(clients) == 2 - per_client = [interceptors[0] for _names, interceptors in clients] + per_client = [interceptors[-1] for _names, interceptors in clients] assert per_client[0] is built[0] assert per_client[1] is built[1] assert per_client[0] is not per_client[1] diff --git a/backend/tests/test_mcp_user_scoped_auth.py b/backend/tests/test_mcp_user_scoped_auth.py index 53e3787a2c..5d71955361 100644 --- a/backend/tests/test_mcp_user_scoped_auth.py +++ b/backend/tests/test_mcp_user_scoped_auth.py @@ -176,9 +176,10 @@ def test_registered_after_oauth_in_shared_assembly(): return await handler(request) interceptors = build_mcp_tool_interceptors(config, oauth_builder=lambda _cfg: oauth) - assert len(interceptors) == 2 + assert len(interceptors) == 3 assert interceptors[0] is oauth - assert interceptors[1].__name__ == "user_scoped_auth_interceptor" + assert interceptors[1].__name__ == "agent_auth_interceptor" + assert interceptors[2].__name__ == "user_scoped_auth_interceptor" def test_shared_assembly_skips_when_no_user_auth(): @@ -187,7 +188,7 @@ def test_shared_assembly_skips_when_no_user_auth(): skills={}, ) interceptors = build_mcp_tool_interceptors(config, oauth_builder=lambda _cfg: None) - assert interceptors == [] + assert [i.__name__ for i in interceptors] == ["agent_auth_interceptor"] def test_gateway_masks_user_auth_credentials(): diff --git a/config.example.yaml b/config.example.yaml index e2a1060bff..5b787da857 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -2957,6 +2957,23 @@ scheduler: recursion_limit: 1000 # ============================================================================ +# Per-agent MCP credentials (opt-in, deployment-managed HTTP servers only). +# Bindings stay operator-owned here; set agent_auth: true on the matching +# extensions_config.json server. See backend/docs/MCP_SERVER.md for the digest +# command, GitOps rollouts, and durable-task revocation behavior. +mcp_agent_auth: {} +# mcp_agent_auth: +# research-tickets: +# server: ticketing +# agent: research-bot +# owners: ["member-user-id"] # Explicit IDs; one list may cover shared copies. +# definition_digests: [""] +# oauth: +# grant_type: client_credentials +# token_url: https://idp.example/realms/work/protocol/openid-connect/token +# client_id: research-bot +# client_secret: $RESEARCH_BOT_CLIENT_SECRET + # Long-running MCP Tasks Configuration # ============================================================================ # Durable runtime for ordinary MCP submit/status/cancel task toolsets. It is