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..27d1591795 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 @@ -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. diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 669c5d931a..871d37c48b 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: @@ -6291,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 @@ -6317,7 +6375,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 +6384,232 @@ 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 + + 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_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 + + 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 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.