Python: Fix header_provider headers not reaching streamable HTTP transport requests

python/packages/core/agent_framework/_mcp.py

@@ -3029,6 +3029,14 @@ 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. 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()
@@ -3071,7 +3079,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 +3101,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.
@@ -3103,11 +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)
-            try:
-                return await super().call_tool(tool_name, **kwargs)
-            finally:
-                _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)
 
 

python/packages/core/tests/core/test_mcp.py

@@ -6125,6 +6125,184 @@ 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, {k.lower(): v for k, v in request.headers.items()}))
+        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
+
+
+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