From 04053b004d6501c88e473cb7863f4e92cb97fae7 Mon Sep 17 00:00:00 2001 From: robbiebusinessacc <65429016+robbiebusinessacc@users.noreply.github.com> Date: Mon, 20 Jul 2026 16:59:53 +0200 Subject: [PATCH 1/2] Python: fix header_provider headers not reaching streamable HTTP requests MCPStreamableHTTPTool.call_tool stores header_provider output in a ContextVar, but the streamable HTTP transport sends requests from tasks spawned at connect time, whose contexts never observe values set later. The request hook therefore always read an empty dict on real connections and the per-call headers (e.g. Authorization) were silently dropped. Keep the ContextVar for in-context reads and add an instance-level snapshot of the active call's headers that the request hook falls back to across tasks. --- python/packages/core/agent_framework/_mcp.py | 15 ++- python/packages/core/tests/core/test_mcp.py | 117 +++++++++++++++++++ 2 files changed, 130 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index abdb0aa5d0..a0abf11d1c 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -3029,6 +3029,11 @@ def __init__( self.terminate_on_close = terminate_on_close self._httpx_client: AsyncClient | None = http_client 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 + # ContextVar values set later inside call_tool, so the request hook needs this + # instance-level snapshot as a cross-task fallback. + self._active_call_headers: dict[str, str] | None = None def _mcp_base_span_attributes(self) -> dict[str, Any]: attrs = super()._mcp_base_span_attributes() @@ -3071,7 +3076,10 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: async def _inject_headers(request: Request) -> None: # noqa: RUF029 if _url_origin(request.url) != target_origin: return - headers = _mcp_call_headers.get({}) + # 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. + headers = _mcp_call_headers.get({}) or self._active_call_headers or {} for key, value in headers.items(): request.headers[key] = value @@ -3090,7 +3098,7 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: When a ``header_provider`` was supplied at construction time, the runtime *kwargs* (originating from ``FunctionInvocationContext.kwargs``) are passed to the provider. The returned headers are attached to every HTTP request - made during this tool call via a ``contextvars.ContextVar``. + made during this tool call via a request hook on the underlying HTTP client. Args: tool_name: The name of the tool to call. @@ -3104,9 +3112,12 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: if self._header_provider is not None: headers = self._header_provider(kwargs) token = _mcp_call_headers.set(headers) + previous_headers = self._active_call_headers + self._active_call_headers = headers try: return await super().call_tool(tool_name, **kwargs) finally: + self._active_call_headers = previous_headers _mcp_call_headers.reset(token) return await super().call_tool(tool_name, **kwargs) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 4b61cd2451..0f8eae4f5b 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -6125,6 +6125,123 @@ def provider(kwargs): assert call_args.kwargs.get("arguments", {}).get("name") == "Alice" +async def test_mcp_streamable_http_tool_header_provider_applies_across_transport_tasks(): + """Regression test for #7161: header_provider headers must reach tools/call requests. + + The streamable HTTP transport sends requests from tasks spawned at connect time, + whose contexts never observe the ContextVar value set later inside call_tool. This + drives the real transport against an in-process mock server and asserts the + per-call Authorization header arrives on the tools/call HTTP request. + """ + import httpx + + captured_requests: list[tuple[str, str, dict[str, str]]] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + if request.method == "DELETE": + return httpx.Response(200) + if request.method == "GET": + return httpx.Response(405) + body = json.loads(request.content.decode()) + method = body.get("method", "") + captured_requests.append((request.method, method, dict(request.headers))) + if method == "initialize": + result = { + "protocolVersion": body["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "mock-server", "version": "1.0.0"}, + } + return httpx.Response( + 200, + headers={"mcp-session-id": "test-session"}, + json={"jsonrpc": "2.0", "id": body["id"], "result": result}, + ) + if method == "tools/list": + result = { + "tools": [ + { + "name": "greet", + "description": "Says hello", + "inputSchema": {"type": "object", "properties": {"name": {"type": "string"}}}, + } + ] + } + return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) + if method == "tools/call": + result = {"content": [{"type": "text", "text": "Hello!"}], "isError": False} + return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) + if "id" in body: + # Any other request (e.g. ping) gets an empty result so the session doesn't block on it. + return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": {}}) + # Notifications (e.g. notifications/initialized) + return httpx.Response(202) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + tool = MCPStreamableHTTPTool( + name="test", + url="http://127.0.0.1:8000/mcp", + load_prompts=False, + http_client=http_client, + header_provider=lambda kw: {"Authorization": f"Bearer {kw.get('api_key', '')}"}, + ) + try: + async with tool: + await tool.call_tool("greet", name="Alice", api_key="secret-token") + finally: + await http_client.aclose() + + call_headers = [headers for _, method, headers in captured_requests if method == "tools/call"] + assert len(call_headers) == 1 + assert call_headers[0].get("authorization") == "Bearer secret-token" + + +async def test_mcp_streamable_http_tool_header_provider_snapshot_restored_after_call(): + """Test that the instance-level header snapshot is set during a call and cleared after.""" + observed_snapshots: list[dict[str, str] | None] = [] + original_call_tool = MCPTool.call_tool + + async def spy_call_tool(self, tool_name, **kwargs): + # Capture the snapshot value during the super call + observed_snapshots.append(self._active_call_headers) + return await original_call_tool(self, tool_name, **kwargs) + + class _TestServer(MCPStreamableHTTPTool): + async def connect(self): # type: ignore[override] # pyrefly: ignore[bad-override] # ty: ignore[invalid-method-override] + self.session = Mock(spec=ClientSession) + self.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="greet", + description="Says hello", + inputSchema={"type": "object", "properties": {"name": {"type": "string"}}}, + ) + ] + ) + ) + self.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="Hello!")]) + ) + self.session.send_ping = AsyncMock() + self.is_connected = True + + def get_mcp_client(self): # pyrefly: ignore[bad-override] + return None + + server = _TestServer( + name="test", + url="http://example.com/mcp", + header_provider=lambda kw: {"X-Auth": kw.get("auth_token", "")}, + ) + async with server: + await server.load_tools() + with patch.object(MCPTool, "call_tool", spy_call_tool): + await server.call_tool("greet", name="Alice", auth_token="bearer-xyz") + + assert observed_snapshots == [{"X-Auth": "bearer-xyz"}] + assert server._active_call_headers is None + + # endregion From d7d29a330e0e590002496384945e4d0dac40e234 Mon Sep 17 00:00:00 2001 From: robbiebusinessacc <65429016+robbiebusinessacc@users.noreply.github.com> Date: Mon, 20 Jul 2026 23:48:27 +0200 Subject: [PATCH 2/2] Python: serialize header_provider tool calls to prevent cross-call header mixing Parallel tool invocations run concurrently per function-invocation batch, so two call_tool invocations on the same MCPStreamableHTTPTool could overwrite each other's active-header snapshot while requests were still in flight, attaching the wrong per-call credentials. Hold a per-instance lock for the duration of a header-bearing call, add a regression test that fails without the lock, and normalize captured header casing in the transport-task test. --- python/packages/core/agent_framework/_mcp.py | 21 ++++--- python/packages/core/tests/core/test_mcp.py | 63 +++++++++++++++++++- 2 files changed, 74 insertions(+), 10 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index a0abf11d1c..1d54537458 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -3032,8 +3032,11 @@ def __init__( # Headers for the in-flight call_tool invocation. The streamable HTTP transport # sends requests from tasks spawned at connect time, whose contexts never observe # ContextVar values set later inside call_tool, so the request hook needs this - # instance-level snapshot as a cross-task fallback. + # instance-level snapshot as a cross-task fallback. The lock serializes tool calls + # when a header_provider is set: parallel invocations on the same instance would + # 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() def _mcp_base_span_attributes(self) -> dict[str, Any]: attrs = super()._mcp_base_span_attributes() @@ -3111,14 +3114,14 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: """ if self._header_provider is not None: headers = self._header_provider(kwargs) - token = _mcp_call_headers.set(headers) - previous_headers = self._active_call_headers - self._active_call_headers = headers - try: - return await super().call_tool(tool_name, **kwargs) - finally: - self._active_call_headers = previous_headers - _mcp_call_headers.reset(token) + async with self._call_headers_lock: + token = _mcp_call_headers.set(headers) + self._active_call_headers = headers + try: + return await super().call_tool(tool_name, **kwargs) + finally: + self._active_call_headers = None + _mcp_call_headers.reset(token) return await super().call_tool(tool_name, **kwargs) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 0f8eae4f5b..a419205afd 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -6144,7 +6144,7 @@ async def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(405) body = json.loads(request.content.decode()) method = body.get("method", "") - captured_requests.append((request.method, method, dict(request.headers))) + captured_requests.append((request.method, method, {k.lower(): v for k, v in request.headers.items()})) if method == "initialize": result = { "protocolVersion": body["params"]["protocolVersion"], @@ -6242,6 +6242,67 @@ def get_mcp_client(self): # pyrefly: ignore[bad-override] assert server._active_call_headers is None +async def test_mcp_streamable_http_tool_header_provider_serializes_concurrent_calls(): + """Concurrent call_tool invocations on the same tool must not mix per-call headers. + + The framework executes parallel tool invocations concurrently, so a second call + must not overwrite the active header snapshot while the first call's requests + are still in flight (that would attach the wrong credentials). + """ + release = asyncio.Event() + in_call = asyncio.Event() + + async def blocking_call_tool(tool_name, *, arguments=None, meta=None): + in_call.set() + await release.wait() + return types.CallToolResult(content=[types.TextContent(type="text", text="Hello!")]) + + class _TestServer(MCPStreamableHTTPTool): + async def connect(self): # type: ignore[override] # pyrefly: ignore[bad-override] # ty: ignore[invalid-method-override] + self.session = Mock(spec=ClientSession) + self.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="greet", + description="Says hello", + inputSchema={"type": "object", "properties": {"name": {"type": "string"}}}, + ) + ] + ) + ) + self.session.call_tool = AsyncMock(side_effect=blocking_call_tool) + self.session.send_ping = AsyncMock() + self.is_connected = True + + def get_mcp_client(self): # pyrefly: ignore[bad-override] + return None + + server = _TestServer( + name="test", + url="http://example.com/mcp", + header_provider=lambda kw: {"X-Auth": kw.get("auth_token", "")}, + ) + async with server: + await server.load_tools() + + first = asyncio.create_task(server.call_tool("greet", name="A", auth_token="token-a")) + await in_call.wait() + assert server._active_call_headers == {"X-Auth": "token-a"} + + second = asyncio.create_task(server.call_tool("greet", name="B", auth_token="token-b")) + for _ in range(10): + await asyncio.sleep(0) + # The second call must be waiting on the first; its headers must not have + # replaced the snapshot the request hook reads for the in-flight call. + assert server._active_call_headers == {"X-Auth": "token-a"} + + release.set() + await asyncio.gather(first, second) + + assert server._active_call_headers is None + + # endregion