diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index abdb0aa5d0..1d54537458 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/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) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 4b61cd2451..a419205afd 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/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