From 9e90a73c39519e8d70b5c6566784c2328ad355b7 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Thu, 3 Sep 2026 16:04:16 +0100 Subject: [PATCH 1/4] Python: Scope MCP HTTP headers per tool Keep provider header processing associated with the MCP tool that issued each request, including when HTTP clients are shared. Clean up request hooks on close and preserve redirect handling. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 4fca7430-5806-4fb4-a190-729e2473c5bb --- python/packages/core/AGENTS.md | 1 + python/packages/core/agent_framework/_mcp.py | 76 +++++++- python/packages/core/tests/core/test_mcp.py | 172 ++++++++++++++++--- 3 files changed, 221 insertions(+), 28 deletions(-) diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 257782f7c3..f012b6a6dd 100644 --- a/python/packages/core/AGENTS.md +++ b/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: diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 3d1d5717e3..15e758d18c 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/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 @@ -3204,16 +3248,42 @@ async def _inject_headers(request: Request) -> None: # ruff:ignore[unused-async headers = {} 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: + request_hooks.remove(self._inject_headers_hook) + self._header_hook_client = None + + async def close(self) -> None: + """Disconnect and remove the tool-specific hook from its HTTP client.""" + try: + await super().close() + 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. diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 669c5d931a..dbb05013b0 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -31,6 +31,7 @@ ) from agent_framework._feature_stage import _WARNED_FEATURES, ExperimentalFeature, ExperimentalWarning from agent_framework._mcp import ( + _MCP_HEADER_OWNER_EXTENSION, MCPTool, _build_prefixed_mcp_name, _get_input_model_from_mcp_prompt, @@ -64,6 +65,16 @@ def _reset_progressive_mcp_warning_state() -> None: _WARNED_FEATURES.discard((ExperimentalWarning, ExperimentalFeature.PROGRESSIVE_TOOLS.value)) +def _request_for_mcp_tool(tool: MCPStreamableHTTPTool, url: str = "http://example.com/mcp") -> Any: + import httpx + + return httpx.Request( + "POST", + url, + extensions={_MCP_HEADER_OWNER_EXTENSION: tool._header_request_owner}, + ) + + # Helper function tests def test_normalize_mcp_name(): """Test MCP name normalization.""" @@ -3972,7 +3983,7 @@ async def test_load_prompts_prevents_multiple_calls(): async def test_mcp_streamable_http_tool_httpx_client_cleanup(): - """Test that MCPStreamableHTTPTool properly passes through httpx clients.""" + """Test that MCPStreamableHTTPTool delegates to caller-provided httpx clients.""" from unittest.mock import AsyncMock, Mock, patch from agent_framework import MCPStreamableHTTPTool @@ -4022,10 +4033,10 @@ async def test_mcp_streamable_http_tool_httpx_client_cleanup(): # Verify the user-provided client was stored assert tool2._httpx_client is user_client, "User-provided client should be stored" - # Verify streamable_http_client was called with the user's client + # Verify the transport wrapper delegates to the user's client. # Get the last call (should be from tool2.connect()) call_args = mock_client.call_args - assert call_args.kwargs["http_client"] is user_client, "User's client should be passed through" + assert call_args.kwargs["http_client"]._client is user_client async def test_load_tools_with_pagination(): @@ -6037,8 +6048,6 @@ def get_mcp_client(self): # pyrefly: ignore[bad-override] async def test_mcp_streamable_http_tool_header_provider_with_httpx_event_hook(): """Test that the httpx event hook injects headers from the contextvar.""" - import httpx - from agent_framework._mcp import MCP_DEFAULT_SSE_READ_TIMEOUT, MCP_DEFAULT_TIMEOUT, _mcp_call_headers tool = MCPStreamableHTTPTool( @@ -6063,7 +6072,7 @@ async def test_mcp_streamable_http_tool_header_provider_with_httpx_event_hook(): # Simulate what happens during a call_tool: contextvar is set token = _mcp_call_headers.set({"X-Custom": "test-value"}) try: - request = httpx.Request("POST", "http://example.com/mcp") + request = _request_for_mcp_tool(tool) await hooks[0](request) assert request.headers.get("X-Custom") == "test-value" finally: @@ -6081,8 +6090,6 @@ async def test_mcp_streamable_http_tool_header_provider_injects_on_ambient_reque outside call_tool, so the contextvar/snapshot are unset. A static header_provider should still be invoked (with empty kwargs) so these requests carry auth headers. """ - import httpx - tool = MCPStreamableHTTPTool( name="test", url="http://example.com/mcp", @@ -6098,7 +6105,7 @@ async def test_mcp_streamable_http_tool_header_provider_injects_on_ambient_reque assert len(hooks) == 1 # No contextvar set and no active call snapshot: simulates the initialize handshake. - request = httpx.Request("POST", "http://example.com/mcp") + request = _request_for_mcp_tool(tool) await hooks[0](request) assert request.headers.get("Authorization") == "******" finally: @@ -6113,8 +6120,6 @@ async def test_mcp_streamable_http_tool_header_provider_ambient_request_tolerate time raise KeyError. The hook should swallow that specific error and proceed without headers rather than failing the initialize handshake. """ - import httpx - tool = MCPStreamableHTTPTool( name="test", url="http://example.com/mcp", @@ -6130,7 +6135,7 @@ async def test_mcp_streamable_http_tool_header_provider_ambient_request_tolerate assert len(hooks) == 1 # No kwargs available at connect time -> provider raises KeyError -> hook swallows it. - request = httpx.Request("POST", "http://example.com/mcp") + request = _request_for_mcp_tool(tool) await hooks[0](request) assert "Authorization" not in request.headers finally: @@ -6146,8 +6151,6 @@ async def test_mcp_streamable_http_tool_header_provider_empty_active_call_skips_ rather than as "unset", which would re-invoke header_provider({}) mid-call and inject headers the caller deliberately omitted. """ - import httpx - from agent_framework._mcp import _mcp_call_headers call_count = 0 @@ -6173,7 +6176,7 @@ def provider(kw: dict[str, Any]) -> dict[str, str]: tool._active_call_headers = {} try: call_count = 0 - request = httpx.Request("POST", "http://example.com/mcp") + request = _request_for_mcp_tool(tool) await hooks[0](request) assert "Authorization" not in request.headers assert call_count == 0, "ambient fallback must not run during a set-but-empty call" @@ -6194,8 +6197,6 @@ async def test_mcp_streamable_http_tool_header_provider_ambient_kwarg_error_is_b """ import logging - import httpx - tool = MCPStreamableHTTPTool( name="test", url="http://example.com/mcp", @@ -6212,7 +6213,7 @@ async def test_mcp_streamable_http_tool_header_provider_ambient_kwarg_error_is_b with caplog.at_level(logging.DEBUG, logger="agent_framework._mcp"): for _ in range(3): - request = httpx.Request("POST", "http://example.com/mcp") + request = _request_for_mcp_tool(tool) await hooks[0](request) assert "Authorization" not in request.headers @@ -6230,7 +6231,6 @@ async def test_mcp_streamable_http_tool_header_provider_ambient_non_keyerror_pro failure or a provider bug - propagates so it is not silently converted into unauthenticated traffic, matching the call_tool path which does not catch header_provider exceptions. """ - import httpx class TokenRefreshError(RuntimeError): pass @@ -6249,7 +6249,7 @@ def failing_provider(kw: dict[str, Any]) -> dict[str, str]: assert len(hooks) == 1 with pytest.raises(TokenRefreshError): - await hooks[0](httpx.Request("POST", "http://example.com/mcp")) + await hooks[0](_request_for_mcp_tool(tool)) finally: if getattr(tool, "_httpx_client", None) is not None: await tool._httpx_client.aclose() # type: ignore[union-attr] @@ -6264,7 +6264,10 @@ async def test_mcp_streamable_http_tool_header_provider_skips_cross_origin_redir tool = MCPStreamableHTTPTool( name="test", url="http://example.com/mcp", - header_provider=lambda kw: {"Authorization": f"Bearer {kw.get('token', '')}"}, + header_provider=lambda kw: { + "Authorization": f"Bearer {kw.get('token', '')}", + "X-API-Key": kw.get("api_key", ""), + }, ) try: @@ -6275,15 +6278,22 @@ async def test_mcp_streamable_http_tool_header_provider_skips_cross_origin_redir hooks = tool._httpx_client.event_hooks.get("request", []) assert len(hooks) == 1 - token = _mcp_call_headers.set({"Authorization": "Bearer secret"}) + token = _mcp_call_headers.set({"Authorization": "Bearer secret", "X-API-Key": "api-secret"}) try: - same_origin = httpx.Request("POST", "http://example.com/redirected") + same_origin = _request_for_mcp_tool(tool, "http://example.com/redirected") await hooks[0](same_origin) assert same_origin.headers.get("Authorization") == "Bearer secret" + assert same_origin.headers.get("X-API-Key") == "api-secret" - cross_origin = httpx.Request("POST", "http://attacker.example/capture") + cross_origin = httpx.Request( + "POST", + "http://attacker.example/capture", + headers=same_origin.headers, + extensions=same_origin.extensions, + ) await hooks[0](cross_origin) assert "Authorization" not in cross_origin.headers + assert "X-API-Key" not in cross_origin.headers finally: _mcp_call_headers.reset(token) finally: @@ -6317,7 +6327,7 @@ async def test_mcp_streamable_http_tool_header_provider_with_user_httpx_client() # Verify the hook injects headers token = _mcp_call_headers.set({"X-Dynamic": "per-request"}) try: - request = httpx.Request("POST", "http://example.com/mcp") + request = _request_for_mcp_tool(tool) await hooks[0](request) assert request.headers.get("X-Dynamic") == "per-request" finally: @@ -6326,6 +6336,118 @@ async def test_mcp_streamable_http_tool_header_provider_with_user_httpx_client() await user_client.aclose() +async def test_mcp_streamable_http_tool_header_provider_isolated_on_shared_httpx_client(): + """Each MCP transport must use its own headers when sharing an httpx client.""" + import httpx + + captured_headers: list[dict[str, str]] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + captured_headers.append({key.lower(): value for key, value in request.headers.items()}) + return httpx.Response(200) + + user_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + tool_a = MCPStreamableHTTPTool( + name="a", + url="http://example.com/mcp", + http_client=user_client, + header_provider=lambda _kw: {"Authorization": "Bearer A", "X-Principal-A": "present"}, + ) + tool_b = MCPStreamableHTTPTool( + name="b", + url="http://example.com/mcp", + http_client=user_client, + header_provider=lambda _kw: {"Authorization": "Bearer B", "X-Principal-B": "present"}, + ) + tool_without_provider = MCPStreamableHTTPTool( + name="anonymous", + url="http://example.com/mcp", + http_client=user_client, + ) + + try: + with patch("agent_framework._mcp.streamable_http_client") as mock_transport: + tool_a.get_mcp_client() + client_a = mock_transport.call_args.kwargs["http_client"] + tool_b.get_mcp_client() + client_b = mock_transport.call_args.kwargs["http_client"] + tool_without_provider.get_mcp_client() + client_without_provider = mock_transport.call_args.kwargs["http_client"] + + for client in (client_a, client_b, client_without_provider, client_a): + async with client.stream("POST", "http://example.com/mcp"): + pass + + assert captured_headers[0].get("authorization") == "Bearer A" + assert captured_headers[0].get("x-principal-a") == "present" + assert "x-principal-b" not in captured_headers[0] + assert captured_headers[1].get("authorization") == "Bearer B" + assert captured_headers[1].get("x-principal-b") == "present" + assert "x-principal-a" not in captured_headers[1] + assert "authorization" not in captured_headers[2] + assert "x-principal-a" not in captured_headers[2] + assert "x-principal-b" not in captured_headers[2] + assert captured_headers[3].get("authorization") == "Bearer A" + assert captured_headers[3].get("x-principal-a") == "present" + assert "x-principal-b" not in captured_headers[3] + finally: + await user_client.aclose() + + +async def test_mcp_streamable_http_tool_removes_header_hook_on_close(): + """Closing one tool must remove only its hook, and reconnecting must restore it.""" + import httpx + + user_client = httpx.AsyncClient() + tool_a = MCPStreamableHTTPTool( + name="a", + url="http://example.com/mcp", + http_client=user_client, + header_provider=lambda _kw: {"Authorization": "Bearer A"}, + ) + tool_b = MCPStreamableHTTPTool( + name="b", + url="http://example.com/mcp", + http_client=user_client, + header_provider=lambda _kw: {"Authorization": "Bearer B"}, + ) + + try: + with patch("agent_framework._mcp.streamable_http_client"): + tool_a.get_mcp_client() + tool_b.get_mcp_client() + assert len(user_client.event_hooks["request"]) == 2 + + with patch.object(MCPTool, "close", new_callable=AsyncMock): + await tool_a.close() + assert user_client.event_hooks["request"] == [tool_b._inject_headers_hook] + + # Reconnecting after close re-attaches exactly one hook for tool A. + with patch("agent_framework._mcp.streamable_http_client"): + tool_a.get_mcp_client() + tool_a.get_mcp_client() + assert user_client.event_hooks["request"].count(tool_a._inject_headers_hook) == 1 + assert len(user_client.event_hooks["request"]) == 2 + finally: + await user_client.aclose() + + +async def test_mcp_header_scoped_client_delegates_unwrapped_attributes(): + """The transport wrapper must stay a drop-in for the caller's httpx client.""" + import httpx + + from agent_framework._mcp import _MCPHeaderScopedClient + + user_client = httpx.AsyncClient(headers={"X-Base": "static"}, follow_redirects=True) + try: + wrapper = _MCPHeaderScopedClient(user_client, object()) + assert wrapper.headers["X-Base"] == "static" + assert wrapper.follow_redirects is True + assert wrapper.build_request("POST", "http://example.com/mcp").method == "POST" + finally: + await user_client.aclose() + + async def test_mcp_streamable_http_tool_header_provider_via_invoke_with_context(): """Test that header_provider receives kwargs via FunctionTool.invoke with FunctionInvocationContext. From 65feaff576029f438db5e315e7af0975bfc17ecf Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Thu, 3 Sep 2026 16:13:25 +0100 Subject: [PATCH 2/4] Python: Refresh MCP headers across redirects Remove previously injected provider headers before applying the latest set so redirected requests retain only current provider output. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 4fca7430-5806-4fb4-a190-729e2473c5bb --- python/packages/core/agent_framework/_mcp.py | 2 + python/packages/core/tests/core/test_mcp.py | 48 ++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 15e758d18c..0112637eca 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -3246,6 +3246,8 @@ 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) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index dbb05013b0..7b3fa5bd76 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -6301,6 +6301,54 @@ async def test_mcp_streamable_http_tool_header_provider_skips_cross_origin_redir await tool._httpx_client.aclose() # type: ignore[union-attr] +async def test_mcp_streamable_http_tool_replaces_headers_on_same_origin_redirect(): + """A redirected request must retain only the provider's latest header set.""" + import httpx + + provider_headers = {"X-Previous": "old"} + tool = MCPStreamableHTTPTool( + name="test", + url="http://example.com/mcp", + header_provider=lambda _kw: provider_headers, + ) + + try: + with patch("agent_framework._mcp.streamable_http_client"): + tool.get_mcp_client() + + assert tool._httpx_client is not None + hooks = tool._httpx_client.event_hooks.get("request", []) + assert len(hooks) == 1 + + initial = _request_for_mcp_tool(tool, "http://example.com/start") + await hooks[0](initial) + assert initial.headers.get("X-Previous") == "old" + + provider_headers = {"X-Current": "new"} + same_origin_redirect = httpx.Request( + "POST", + "http://example.com/redirected", + headers=initial.headers, + extensions=initial.extensions, + ) + await hooks[0](same_origin_redirect) + assert "X-Previous" not in same_origin_redirect.headers + assert same_origin_redirect.headers.get("X-Current") == "new" + + cross_origin_redirect = httpx.Request( + "POST", + "http://other.example/redirected", + headers=same_origin_redirect.headers, + extensions=same_origin_redirect.extensions, + ) + await hooks[0](cross_origin_redirect) + assert "X-Previous" not in cross_origin_redirect.headers + assert "X-Current" not in cross_origin_redirect.headers + finally: + if getattr(tool, "_httpx_client", None) is not None: + await tool._httpx_client.aclose() # type: ignore[union-attr] + + async def test_mcp_streamable_http_tool_header_provider_with_user_httpx_client(): """Test that header_provider works when the user provides their own httpx client.""" import httpx From e96bbeee72563df2d8929f43ada05205bbe174f3 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Thu, 3 Sep 2026 17:38:14 +0100 Subject: [PATCH 3/4] Python: Defer MCP hook cleanup until teardown Keep the request hook attached while lifecycle-owner transport cleanup runs, including when the caller awaiting close is cancelled. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 4fca7430-5806-4fb4-a190-729e2473c5bb --- python/packages/core/agent_framework/_mcp.py | 6 +-- python/packages/core/tests/core/test_mcp.py | 54 +++++++++++++++++++- 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 0112637eca..0a027f5390 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -3279,10 +3279,10 @@ def _remove_header_hook(self) -> None: request_hooks.remove(self._inject_headers_hook) self._header_hook_client = None - async def close(self) -> None: - """Disconnect and remove the tool-specific hook from its HTTP client.""" + async def _close_on_owner(self) -> None: + """Disconnect on the lifecycle owner before removing the request hook.""" try: - await super().close() + await super()._close_on_owner() finally: self._remove_header_hook() diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 7b3fa5bd76..aaeecfc19d 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -6466,8 +6466,7 @@ async def test_mcp_streamable_http_tool_removes_header_hook_on_close(): tool_b.get_mcp_client() assert len(user_client.event_hooks["request"]) == 2 - with patch.object(MCPTool, "close", new_callable=AsyncMock): - await tool_a.close() + await tool_a.close() assert user_client.event_hooks["request"] == [tool_b._inject_headers_hook] # Reconnecting after close re-attaches exactly one hook for tool A. @@ -6480,6 +6479,57 @@ async def test_mcp_streamable_http_tool_removes_header_hook_on_close(): await user_client.aclose() +async def test_mcp_streamable_http_tool_keeps_header_hook_until_cancelled_close_finishes(): + """Caller cancellation must not remove the hook while lifecycle teardown continues.""" + import httpx + + user_client = httpx.AsyncClient() + tool = MCPStreamableHTTPTool( + name="test", + url="http://example.com/mcp", + http_client=user_client, + header_provider=lambda _kw: {"Authorization": "Bearer token"}, + ) + close_started = asyncio.Event() + allow_close = asyncio.Event() + # Recorded rather than asserted here: an assertion raised on the lifecycle owner task + # is swallowed by its error handling, so it would pass even when the hook is detached. + hook_attached_during_teardown: list[bool] = [] + + async def delayed_close() -> None: + close_started.set() + await allow_close.wait() + hook_attached_during_teardown.append(tool._inject_headers_hook in user_client.event_hooks["request"]) + + try: + with patch("agent_framework._mcp.streamable_http_client"): + tool.get_mcp_client() + assert tool._inject_headers_hook in user_client.event_hooks["request"] + + with patch.object(MCPTool, "_close_on_owner", side_effect=delayed_close): + close_task = asyncio.create_task(tool.close()) + await close_started.wait() + owner_task = tool._lifecycle_owner_task + assert owner_task is not None + + close_task.cancel() + with pytest.raises(asyncio.CancelledError): + await close_task + assert tool._inject_headers_hook in user_client.event_hooks["request"] + + allow_close.set() + await owner_task + + assert hook_attached_during_teardown == [True] + assert tool._inject_headers_hook not in user_client.event_hooks["request"] + finally: + allow_close.set() + owner_task = tool._lifecycle_owner_task + if owner_task is not None: + await owner_task + await user_client.aclose() + + async def test_mcp_header_scoped_client_delegates_unwrapped_attributes(): """The transport wrapper must stay a drop-in for the caller's httpx client.""" import httpx From bbbfa5f30574de0b76a9512be1afd06bd5d4f938 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Fri, 4 Sep 2026 10:46:15 +0100 Subject: [PATCH 4/4] Python: Preserve active MCP hook iteration Replace the shared request-hook list during cleanup so in-flight HTTPX hook iteration is not disrupted. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 4fca7430-5806-4fb4-a190-729e2473c5bb --- python/packages/core/agent_framework/_mcp.py | 4 +- python/packages/core/tests/core/test_mcp.py | 64 ++++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 0a027f5390..27d1591795 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -3276,7 +3276,9 @@ def _remove_header_hook(self) -> None: return request_hooks = self._header_hook_client.event_hooks["request"] if self._inject_headers_hook in request_hooks: - request_hooks.remove(self._inject_headers_hook) + 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: diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index aaeecfc19d..871d37c48b 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -6479,6 +6479,70 @@ async def test_mcp_streamable_http_tool_removes_header_hook_on_close(): await user_client.aclose() +async def test_mcp_streamable_http_tool_removes_hook_without_mutating_active_hook_list(): + """Closing one tool must not disrupt an in-progress iteration over shared hooks.""" + import httpx + + from agent_framework._mcp import _MCP_INJECTED_HEADER_KEYS_EXTENSION + + captured_headers: list[httpx.Headers] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + captured_headers.append(request.headers) + return httpx.Response(200) + + user_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + tool_a = MCPStreamableHTTPTool( + name="a", + url="http://example.com/mcp", + http_client=user_client, + header_provider=lambda _kw: {"Authorization": "Bearer token-a"}, + ) + tool_b = MCPStreamableHTTPTool( + name="b", + url="http://example.com/mcp", + http_client=user_client, + header_provider=lambda _kw: {"X-API-Key": "current"}, + ) + hook_started = asyncio.Event() + allow_hook = asyncio.Event() + + async def delayed_hook(_request: httpx.Request) -> None: + hook_started.set() + await allow_hook.wait() + + send_task: asyncio.Task[httpx.Response] | None = None + try: + with patch("agent_framework._mcp.streamable_http_client"): + tool_a.get_mcp_client() + tool_b.get_mcp_client() + user_client.event_hooks["request"].insert(1, delayed_hook) + + request = httpx.Request( + "POST", + "http://other.example/redirected", + headers={"X-API-Key": "previous"}, + extensions={ + _MCP_HEADER_OWNER_EXTENSION: tool_b._header_request_owner, + _MCP_INJECTED_HEADER_KEYS_EXTENSION: ("X-API-Key",), + }, + ) + send_task = asyncio.create_task(user_client.send(request)) + await hook_started.wait() + + await tool_a.close() + allow_hook.set() + await send_task + + assert len(captured_headers) == 1 + assert "X-API-Key" not in captured_headers[0] + finally: + allow_hook.set() + if send_task is not None: + await send_task + await user_client.aclose() + + async def test_mcp_streamable_http_tool_keeps_header_hook_until_cancelled_close_finishes(): """Caller cancellation must not remove the hook while lifecycle teardown continues.""" import httpx