Python: Scope MCP headers to transport requests

python/packages/core/AGENTS.md

@@ -107,6 +107,7 @@ agent_framework/
 - **`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.
 - **`additional_tool_argument_names`** (constructor arg on all `MCPTool` subclasses) - Opt extra argument names back into the allowlist. Accepts a `Sequence[str]` (applied to every tool) or a `Mapping[str, Sequence[str]]` keyed by **remote tool name**, where the reserved key `"*"` denotes global extras. It is configured only in user code at construction; there is **no per-call/runtime override**, so a model-issued tool call cannot change which names pass through — but note this constrains the *model*, not the *server*, which still widens the effective allowlist through its schema. To use a server that accepts `additionalProperties: true`, list the extra names here and then either (1) manually extend that tool's `inputSchema` (via the `.functions` list after connecting) so the model is prompted to supply them, or (2) supply the values yourself via `function_invocation_kwargs`. If a normal forwarded argument name is supplied by both the model and `function_invocation_kwargs`, the model-supplied value wins; `_meta` is the exception and only trusted runtime/caller metadata is used.
+- **`header_provider` request scoping** - When sharing an `http_client`, keep provider processing scoped to the originating `MCPStreamableHTTPTool`, remove its request hook on `close()`, and strip injected headers from cross-origin redirects.
 - **`function_invocation_kwargs` and MCP servers** - That dict is shared across every tool in the run, including every attached `MCPTool`, and any name in it reaches a server that declares a matching `inputSchema` property. `header_provider` does not mitigate this — it reads the kwargs without consuming them. To keep a credential out of tool arguments, source it outside `function_invocation_kwargs`: read a `ContextVar` inside the provider (this still allows a different value per request), configure a custom `http_client`, or use `env` for `MCPStdioTool`.
 - **Sampling guardrails** (`sampling_callback`) - Passing `client=` advertises `SamplingCapability` so the server can send `sampling/createMessage`. Because remote servers are untrusted (confused-deputy risk), the default `sampling_callback` is **deny-by-default** and applies, in order: a per-session rate limit (`sampling_max_requests`, default `_DEFAULT_SAMPLING_MAX_REQUESTS`), an approval gate (`sampling_approval_callback`), and a `maxTokens` cap (`sampling_max_tokens`, default `_DEFAULT_SAMPLING_MAX_TOKENS`). The approval callback (constructor arg on all subclasses; exported type alias `SamplingApprovalCallback`) receives the raw `CreateMessageRequestParams`, may be sync or async, and must return truthy to approve. When it is `None` (the default) every sampling request is denied; pass `lambda params: True` to restore legacy auto-approve as an explicit opt-in. Requests and denials are logged at WARNING (content is not logged). The per-session counter resets in `_reset_session_state`.
 - **`MCPTaskOptions`** (experimental, `MCP_LONG_RUNNING_TASKS` feature, **frozen**) - Per-tool-instance options controlling the SEP-2663 long-running task lifecycle. When the server advertises a tool with `execution.taskSupport == "required"`, `MCPTool.call_tool` transparently routes through `call_tool_as_task`, which sends an augmented `tools/call`, polls `tasks/get` until terminal, and reinterprets `tasks/result` as a normal `CallToolResult`. Instances are immutable; replace via `MCPTool.task_options = MCPTaskOptions(...)`. Fields:

python/packages/core/agent_framework/_mcp.py

@@ -124,9 +124,43 @@ class MCPSpecificApproval(TypedDict, total=False):
     "_meta",
 })
 _mcp_call_headers: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar("_mcp_call_headers")
+_MCP_HEADER_OWNER_EXTENSION = "agent_framework.mcp_header_owner"
+_MCP_INJECTED_HEADER_KEYS_EXTENSION = "agent_framework.mcp_injected_header_keys"
 MCP_DEFAULT_TIMEOUT = 30
 MCP_DEFAULT_SSE_READ_TIMEOUT = 60 * 5
 
+
+class _MCPHeaderScopedClient:
+    """Attach private tool context to MCP transport requests."""
+
+    def __init__(self, client: AsyncClient, owner: object) -> None:
+        self._client = client
+        self._owner = owner
+
+    def __getattr__(self, name: str) -> Any:
+        # Delegate the rest of the httpx client surface so this wrapper stays a
+        # drop-in for the MCP transport. Only the request-sending methods below
+        # are wrapped; anything the transport reads (timeouts, headers, ...)
+        # comes straight from the caller's client. ``_client`` itself is always a
+        # real instance attribute, so guard against recursing on a partially
+        # initialized wrapper.
+        if name == "_client":
+            raise AttributeError(name)
+        return getattr(self._client, name)
+
+    def _tagged_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+        extensions = dict(kwargs.get("extensions") or {})
+        extensions[_MCP_HEADER_OWNER_EXTENSION] = self._owner
+        kwargs["extensions"] = extensions
+        return kwargs
+
+    def stream(self, *args: Any, **kwargs: Any) -> Any:
+        return self._client.stream(*args, **self._tagged_kwargs(kwargs))
+
+    async def delete(self, *args: Any, **kwargs: Any) -> Any:
+        return await self._client.delete(*args, **self._tagged_kwargs(kwargs))
+
+
 # Default safety limits applied to server-initiated MCP sampling requests
 # (``sampling/createMessage``). MCP servers are untrusted third parties, so the
 # default ``sampling_callback`` denies requests unless an approval callback is
@@ -2981,7 +3015,7 @@ def __init__(
         Note:
             The arguments are used to create a streamable HTTP client using the
             new ``mcp.client.streamable_http.streamable_http_client`` API.
-            If an asyncClient is provided via ``http_client``, it will be used directly.
+            If an asyncClient is provided via ``http_client``, it will be used as the underlying transport client.
             Otherwise, the ``streamable_http_client`` API will create and manage a default client.
 
         Args:
@@ -3058,9 +3092,12 @@ def __init__(
                 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
+                origins on cross-origin redirects; headers injected this way are also removed
+                again if a redirect leaves that origin. If you instead supply sensitive headers
                 through a custom ``http_client``, you must enforce this same origin-scoped
                 policy yourself.
+                Headers returned by the provider are applied only to requests issued by this
+                tool, including when several tools share one ``http_client``.
                 Note that the provider reads these kwargs without consuming them: the same
                 values continue on to the outbound argument filter, so reading a credential
                 here does not withhold it from the server. See
@@ -3128,6 +3165,8 @@ 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()
+        self._header_request_owner = object()
+        self._header_hook_client: AsyncClient | None = None
 
     def _mcp_base_span_attributes(self) -> dict[str, Any]:
         attrs = super()._mcp_base_span_attributes()
@@ -3168,7 +3207,12 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
             if not hasattr(self, "_inject_headers_hook"):
 
                 async def _inject_headers(request: Request) -> None:  # ruff:ignore[unused-async]
+                    request_owner = request.extensions.get(_MCP_HEADER_OWNER_EXTENSION)
+                    if request_owner is not self._header_request_owner:
+                        return
                     if _url_origin(request.url) != target_origin:
+                        for key in request.extensions.pop(_MCP_INJECTED_HEADER_KEYS_EXTENSION, ()):
+                            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
@@ -3202,18 +3246,48 @@ async def _inject_headers(request: Request) -> None:  # ruff:ignore[unused-async
                                 exc_info=True,
                             )
                             headers = {}
+                    for key in request.extensions.pop(_MCP_INJECTED_HEADER_KEYS_EXTENSION, ()):
+                        request.headers.pop(key, None)
                     for key, value in headers.items():
                         request.headers[key] = value
+                    request.extensions[_MCP_INJECTED_HEADER_KEYS_EXTENSION] = tuple(headers)
 
                 self._inject_headers_hook = _inject_headers
+
+            if self._header_hook_client is not http_client:
+                self._remove_header_hook()
+                self._header_hook_client = http_client
+            if self._inject_headers_hook not in http_client.event_hooks["request"]:
                 http_client.event_hooks["request"].append(self._inject_headers_hook)
 
+        transport_http_client = (
+            _MCPHeaderScopedClient(http_client, self._header_request_owner) if http_client is not None else None
+        )
+
         return streamable_http_client(
             url=self.url,
-            http_client=http_client,
+            http_client=transport_http_client,
             terminate_on_close=self.terminate_on_close if self.terminate_on_close is not None else True,
         )
 
+    def _remove_header_hook(self) -> None:
+        """Detach this tool's request hook from its HTTP client."""
+        if self._header_hook_client is None or not hasattr(self, "_inject_headers_hook"):
+            return
+        request_hooks = self._header_hook_client.event_hooks["request"]
+        if self._inject_headers_hook in request_hooks:
+            self._header_hook_client.event_hooks["request"] = [
+                hook for hook in request_hooks if hook is not self._inject_headers_hook
+            ]
+        self._header_hook_client = None
+
+    async def _close_on_owner(self) -> None:
+        """Disconnect on the lifecycle owner before removing the request hook."""
+        try:
+            await super()._close_on_owner()
+        finally:
+            self._remove_header_hook()
+
     async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]:
         """Call a tool, injecting headers from the header_provider if configured.