Python: Add origin-scoped headers for MCP connect authentication

python/packages/core/AGENTS.md

@@ -103,6 +103,16 @@ agent_framework/
 
 - **`MCPTool`** - Base wrapper that owns the MCP `ClientSession` and exposes the remote server's tools as `FunctionTool`s.
 - **`MCPStdioTool`** / **`MCPStreamableHTTPTool`** / **`MCPWebsocketTool`** - Transport-specific subclasses.
+- **`MCPStreamableHTTPTool` headers** - ``headers=`` are static same-origin headers applied to
+  ``connect()`` / initialize, discovery, pings, reconnects, and tool calls (origin-scoped, so they
+  are not leaked on cross-origin redirects). ``header_provider`` overlays per-call headers from
+  runtime kwargs (``FunctionInvocationContext.kwargs``). Ambient requests call
+  ``header_provider({})``; kwargs-only providers that raise ``KeyError`` are tolerated so connect
+  can succeed against servers that do not authenticate the handshake. Servers that *do* require
+  handshake auth need ``headers=`` or a provider that returns credentials for ``{}``. Prefer
+  ``headers=`` over baking tokens into a custom ``http_client`` so the origin check still applies.
+  The request hook records keys it injected and strips those keys from subsequent cross-origin
+  requests (HTTPX already strips ``Authorization``; other secrets such as ``X-API-Key`` are removed here).
 - **Argument allowlist (`_prepare_call_kwargs`)** - Before each `tools/call`, kwargs are filtered to an **allowlist** built from the tool's declared parameters (`inputSchema.properties`) plus any user-configured extras. Framework runtime kwargs injected through the function-invocation pipeline (e.g. `thread`, `conversation_id`, `chat_options`, `options`, `response_format`) are stripped by default rather than forwarded. A tool that declares no usable `properties` (including schemas with `additionalProperties: true`) forwards only the configured extras. The `_MCP_FRAMEWORK_DENYLIST` is a safety net for framework-named params a server *declares* in its schema (those are dropped); names explicitly opted in via `additional_tool_argument_names` always win. The reserved `_meta` key is never forwarded as an argument; trusted caller/runtime `_meta` is validated as MCP request metadata, model-supplied `_meta` is discarded in generated MCP functions, and metadata precedence is caller/runtime < OpenTelemetry < tools/list metadata.
 - **`allowed_tools`** (constructor arg on all `MCPTool` subclasses) - Restricts exposed MCP tools by raw remote MCP tool identity. Prefixed local names remain accepted only when the raw remote name already matches its normalized form; normalized/local aliases do not authorize a different raw remote name. If multiple raw remote tool names map to the same local function name, tool loading raises `ToolExecutionException` instead of first-one-wins shadowing.
 - **Progressive MCP disclosure** (`use_progressive_disclosure`, `always_load`) - When enabled on any `MCPTool` subclass, the initial model-facing surface is loader tools (`list_mcp_tools` / `load_tool` / `unload_tool`, prefixed by `tool_name_prefix` when configured) plus allowed tools selected by `always_load` and tools loaded earlier on the same `MCPTool` instance. `list_mcp_tools` only reports tools that pass `allowed_tools`; filtered tools are not listed or loadable. Loader tool names are reserved in progressive mode: remote MCP tools whose local generated name collides with a loader name are omitted from the initial/listed surface, and explicit `load_tool` calls return a model-visible message pointing callers to `tool_name_prefix` or excluding the colliding tool. `load_tool` accepts one tool name or a list of tool names and uses `FunctionInvocationContext.add_tools(...)` so the selected generated MCP `FunctionTool`s become available on the next function-calling iteration while keeping existing approval mode, argument filtering, header-provider runtime kwargs, result parsing, OTel, and task behavior. `unload_tool` accepts one dynamically loaded tool name or a list of names and removes them from the live tool list and persisted progressive surface, but it does not remove tools configured in `always_load`. Invalid `always_load` entries are ignored like unmatched `allowed_tools` entries.

python/packages/core/agent_framework/_mcp.py

@@ -2940,6 +2940,7 @@ def __init__(
         sampling_max_requests: int | None = _DEFAULT_SAMPLING_MAX_REQUESTS,
         additional_properties: dict[str, Any] | None = None,
         http_client: AsyncClient | None = None,
+        headers: Mapping[str, str] | None = None,
         header_provider: Callable[[dict[str, Any]], dict[str, str]] | None = None,
         task_options: MCPTaskOptions | None = None,
         additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None,
@@ -3013,23 +3014,38 @@ def __init__(
                 and pass your own ``asyncClient`` instance.
                 Security: when you attach sensitive headers (e.g. authentication tokens)
                 via a custom ``http_client``, you are responsible for enforcing the same
-                origin-scoped header policy that the built-in ``header_provider`` hook
-                applies. The framework only injects ``header_provider`` headers on requests
-                whose origin (scheme, host, port) matches the configured ``url``, so tokens
-                are not leaked to third-party origins on cross-origin redirects. A custom
-                client that sets headers unconditionally (e.g. via ``AsyncClient(headers=...)``
-                or ``follow_redirects=True`` without an origin check) can leak those headers
-                to other origins; scope them to the target origin yourself.
+                origin-scoped header policy that the built-in ``headers`` /
+                ``header_provider`` hook applies. The framework only injects those
+                headers on requests whose origin (scheme, host, port) matches the
+                configured ``url``, so tokens are not leaked to third-party origins on
+                cross-origin redirects. A custom client that sets headers unconditionally
+                (e.g. via ``AsyncClient(headers=...)`` or ``follow_redirects=True``
+                without an origin check) can leak those headers to other origins; scope
+                them to the target origin yourself.
+            headers: Optional static HTTP headers attached to every same-origin request,
+                including ``connect()`` / initialize, discovery, pings, reconnects, and
+                tool calls. Use this when the MCP server authenticates the handshake
+                itself. Prefer ``headers`` over baking tokens into a custom
+                ``http_client`` so the origin-scoped injection policy still applies.
+                Per-call ``header_provider`` values overlay these static headers.
+                On a cross-origin redirect, headers previously injected by this hook
+                are stripped from the redirected request (HTTPX already strips
+                ``Authorization``; other secrets such as ``X-API-Key`` are removed here).
             header_provider: Optional callable that receives the runtime keyword arguments
                 (from ``FunctionInvocationContext.kwargs``) and returns a ``dict[str, str]``
-                of HTTP headers to inject into every outbound request to the MCP server.
+                of HTTP headers to inject into outbound requests to the MCP server.
                 Use this to forward per-request context (e.g. authentication tokens set in
                 agent middleware) without creating a separate ``httpx.AsyncClient``.
-                The framework attaches these headers only to requests whose origin (scheme,
-                host, port) matches the configured ``url``, so they are not leaked to other
-                origins on cross-origin redirects. If you instead supply sensitive headers
-                through a custom ``http_client``, you must enforce this same origin-scoped
-                policy yourself.
+                Ambient requests outside ``call_tool`` (the initialize handshake,
+                discovery, and pings) invoke the provider with ``{}``. Providers that
+                authenticate connect must either tolerate empty kwargs and return
+                handshake credentials, or the caller must also supply ``headers``.
+                A ``KeyError`` from a kwargs-only provider is tolerated on ambient
+                requests so connect can still succeed against servers that do not
+                require handshake auth. The framework attaches these headers only to
+                requests whose origin (scheme, host, port) matches the configured
+                ``url``. If you instead supply sensitive headers through a custom
+                ``http_client``, you must enforce this same origin-scoped policy yourself.
             task_options: Options for tools that advertise
                 ``execution.taskSupport == "required"``. See :class:`MCPTaskOptions`.
             additional_tool_argument_names: Extra argument names to forward to the MCP server in
@@ -3073,6 +3089,7 @@ def __init__(
         self.url = url
         self.terminate_on_close = terminate_on_close
         self._httpx_client: AsyncClient | None = http_client
+        self._static_headers: dict[str, str] = dict(headers) if headers else {}
         self._header_provider = header_provider
         # Headers for the in-flight call_tool invocation. The streamable HTTP transport
         # sends requests from tasks spawned at connect time, whose contexts never observe
@@ -3082,6 +3099,10 @@ def __init__(
         # otherwise overwrite each other's snapshot and attach the wrong per-call headers.
         self._active_call_headers: dict[str, str] | None = None
         self._call_headers_lock = asyncio.Lock()
+        # Keys last injected by the request hook. HTTPX may copy non-Authorization secrets
+        # onto a cross-origin redirect; the hook strips these keys when the next request
+        # leaves the configured origin.
+        self._injected_header_keys: set[str] = set()
 
     def _mcp_base_span_attributes(self) -> dict[str, Any]:
         attrs = super()._mcp_base_span_attributes()
@@ -3101,6 +3122,38 @@ def _mcp_base_span_attributes(self) -> dict[str, Any]:
             logger.debug("Failed to parse URL for MCP span transport attributes", exc_info=True)
         return attrs
 
+    def _resolve_outbound_headers(self) -> dict[str, str]:
+        """Resolve origin-scoped headers for the current HTTP request.
+
+        Static ``headers`` always apply. Per-call ``header_provider`` values overlay
+        them when a tool call is in flight. Ambient requests (connect, discovery,
+        pings) call ``header_provider({})`` so static providers can authenticate the
+        handshake; ``KeyError`` is tolerated for kwargs-only providers.
+        """
+        resolved = dict(self._static_headers)
+        call_headers = _mcp_call_headers.get(None)
+        if call_headers is None:
+            call_headers = self._active_call_headers
+        if call_headers is None:
+            # Ambient request made outside call_tool (the initialize handshake,
+            # load_tools/load_prompts discovery, or background pings).
+            if self._header_provider is not None:
+                try:
+                    call_headers = self._header_provider({})
+                except KeyError:
+                    # A kwargs-dependent provider raises on every ambient request.
+                    logger.debug(
+                        "header_provider raised KeyError for MCP server %r on an ambient "
+                        "request (missing per-call kwargs); proceeding with static headers.",
+                        self.name,
+                        exc_info=True,
+                    )
+                    call_headers = {}
+            else:
+                call_headers = {}
+        resolved.update(call_headers)
+        return resolved
+
     def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
         """Get an MCP streamable HTTP client.
 
@@ -3110,7 +3163,7 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
         from httpx import URL, AsyncClient, Request, Timeout
 
         http_client = self._httpx_client
-        if self._header_provider is not None:
+        if self._header_provider is not None or self._static_headers:
             target_origin = _url_origin(URL(self.url))
             if http_client is None:
                 http_client = AsyncClient(
@@ -3123,40 +3176,21 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
 
                 async def _inject_headers(request: Request) -> None:  # ruff:ignore[unused-async]
                     if _url_origin(request.url) != target_origin:
+                        # Strip secrets this hook previously injected. HTTPX removes
+                        # Authorization on cross-origin redirects, but other credentials
+                        # (e.g. X-API-Key) can remain on the redirected request.
+                        for key in self._injected_header_keys:
+                            request.headers.pop(key, None)
                         return
                     # The transport may send this request from a task whose context was
                     # captured before call_tool set the ContextVar; fall back to the
                     # instance-level snapshot of the active call's headers. Both are None
                     # only when this is an ambient request outside call_tool; an active
                     # call that legitimately produced no headers yields an empty dict and
-                    # must not trigger the ambient fallback below.
-                    headers = _mcp_call_headers.get(None)
-                    if headers is None:
-                        headers = self._active_call_headers
-                    if headers is None:
-                        # Ambient request made outside call_tool (the initialize handshake,
-                        # load_tools/load_prompts discovery, or background pings). Invoke the
-                        # provider with empty kwargs so static providers can authenticate these
-                        # requests too. A provider that indexes a required per-call kwarg (e.g.
-                        # kwargs["api_key"]) raises KeyError on the empty dict; that specific
-                        # case is tolerated so connect still succeeds. Any other error is a
-                        # genuine provider failure and is left to propagate, matching the
-                        # call_tool path which does not catch header_provider exceptions.
-                        if self._header_provider is None:
-                            raise RuntimeError("Header injection hook invoked without a header_provider.")
-                        try:
-                            headers = self._header_provider({})
-                        except KeyError:
-                            # A kwargs-dependent provider raises on every ambient request
-                            # (initialize, discovery, and recurring pings).
-                            logger.debug(
-                                "header_provider raised KeyError for MCP server %r on an ambient "
-                                "request (missing per-call kwargs); proceeding without headers.",
-                                self.name,
-                                exc_info=True,
-                            )
-                            headers = {}
-                    for key, value in headers.items():
+                    # must not trigger the ambient header_provider({}) fallback.
+                    outbound = self._resolve_outbound_headers()
+                    self._injected_header_keys = set(outbound)
+                    for key, value in outbound.items():
                         request.headers[key] = value
 
                 self._inject_headers_hook = _inject_headers

python/packages/core/agent_framework/security.py

@@ -3419,29 +3419,15 @@ def __init__(
             raise ValueError("Provide either 'mcp_tool' (an MCPTool instance) or 'url' (a remote MCP server URL).")
 
         if url is not None:
-            from httpx import AsyncClient, Timeout
-
-            from ._mcp import MCP_DEFAULT_SSE_READ_TIMEOUT, MCP_DEFAULT_TIMEOUT, MCPStreamableHTTPTool
-
-            static_headers = dict(headers or {})
-            # Pass headers via an AsyncClient so they are included on ALL requests
-            # (including session.initialize()), not just tool calls. Using
-            # header_provider alone only sets headers via a ContextVar that is
-            # populated during call_tool() and would be empty during initialization,
-            # causing 401s that silently manifest as anyio cancel-scope errors.
-            http_client = (
-                AsyncClient(
-                    headers=static_headers,
-                    follow_redirects=True,
-                    timeout=Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT),
-                )
-                if static_headers
-                else None
-            )
+            from ._mcp import MCPStreamableHTTPTool
+
+            # Prefer origin-scoped ``headers=`` over baking tokens into an
+            # ``AsyncClient(headers=..., follow_redirects=True)``, which can leak
+            # non-Authorization credentials on cross-origin redirects.
             mcp_tool = MCPStreamableHTTPTool(
                 name=name or "mcp",
                 url=url,
-                http_client=http_client,
+                headers=dict(headers) if headers else None,
                 description=description,
             )