Python: Fix HTTP call to MCP losing tracing parent (#3623)#3933
Python: Fix HTTP call to MCP losing tracing parent (#3623)#3933alliscode wants to merge 1 commit intomicrosoft:mainfrom
Conversation
Add OpenTelemetry trace context propagation to MCPStreamableHTTPTool's httpx client via an event hook that injects W3C traceparent/tracestate headers into outgoing HTTP requests. This ensures the remote MCP server's spans are correctly parented under the agent framework's Execute Tool span. - Add _inject_otel_context async event hook using opentelemetry.propagate.inject - Add _get_instrumented_httpx_client to ensure all httpx clients have the hook - Cache auto-created httpx clients; clean up in close() - Update and add tests for trace context injection Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Python Test Coverage Report •
Python Unit Test Overview
|
||||||||||||||||||||||||||||||
There was a problem hiding this comment.
Pull request overview
Adds OpenTelemetry trace context propagation for MCP streamable HTTP calls so the remote MCP server can parent spans under the agent framework’s active span during tool execution.
Changes:
- Introduces an httpx request event hook that injects W3C
traceparent/tracestateheaders viaopentelemetry.propagate.inject. - Ensures
MCPStreamableHTTPToolalways uses an instrumented httpx client (auto-created + cached when not user-supplied) and cleans it up inclose(). - Updates/adds tests around httpx client creation, hook registration, caching, and cleanup.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 6 comments.
| File | Description |
|---|---|
| python/packages/core/agent_framework/_mcp.py | Adds OTel header injection hook, instrumented httpx client creation/caching, and cleanup on close for MCP streamable HTTP transport. |
| python/packages/core/tests/core/test_mcp.py | Updates expectations for the new instrumented auto-client behavior and adds tests for hook registration/caching/cleanup. |
|
|
||
| if self._auto_httpx_client is None: | ||
| self._auto_httpx_client = httpx.AsyncClient( | ||
| event_hooks={"request": [_inject_otel_context]} |
There was a problem hiding this comment.
The auto-created httpx.AsyncClient() uses httpx defaults, including a 5s read timeout. For streamable HTTP/SSE this can break long-lived streams or requests with infrequent data. Consider explicitly configuring the auto client with an appropriate timeout (e.g., disabling read timeout) to match the MCP streamable transport’s expectations.
| event_hooks={"request": [_inject_otel_context]} | |
| event_hooks={"request": [_inject_otel_context]}, | |
| timeout=httpx.Timeout(None), |
| def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: | ||
| """Get an MCP streamable HTTP client. | ||
|
|
||
| Returns: | ||
| An async context manager for the streamable HTTP client transport. | ||
| """ | ||
| # Pass the http_client (which may be None) to streamable_http_client | ||
| return streamable_http_client( | ||
| url=self.url, | ||
| http_client=self._httpx_client, | ||
| http_client=self._get_instrumented_httpx_client(), | ||
| terminate_on_close=self.terminate_on_close if self.terminate_on_close is not None else True, | ||
| ) |
There was a problem hiding this comment.
get_mcp_client() now always passes an internally created httpx client, but the constructor docstring/comments above still state that streamable_http_client will create/manage a default client when http_client isn’t provided. Please update the docs to reflect the new ownership model (tool-managed cached client + close() cleanup).
| # An auto-created httpx client with OTel trace injection should be passed | ||
| mock_http_client.assert_called_once() | ||
| call_kwargs = mock_http_client.call_args.kwargs | ||
| assert call_kwargs["url"] == "http://example.com" | ||
| assert call_kwargs["terminate_on_close"] is True | ||
| # The http_client should be an auto-created instrumented client, not None | ||
| assert call_kwargs["http_client"] is not None | ||
| assert call_kwargs["http_client"] is tool._auto_httpx_client |
There was a problem hiding this comment.
This test now triggers creation of a real httpx.AsyncClient (via get_mcp_client()), but it never gets closed. That can cause unclosed-client ResourceWarnings and flaky test runs. Consider making this test async and calling await tool.close() (or otherwise ensuring the auto-created client is closed).
| async def test_mcp_streamable_http_tool_otel_trace_injection(): | ||
| """Test that MCPStreamableHTTPTool injects OpenTelemetry trace context into outgoing requests.""" | ||
| from agent_framework._mcp import _inject_otel_context | ||
|
|
||
| # Test 1: Auto-created client gets the OTel hook | ||
| tool = MCPStreamableHTTPTool( | ||
| name="test", | ||
| url="http://localhost:8081/mcp", | ||
| ) | ||
| with patch("agent_framework._mcp.streamable_http_client") as mock_client: | ||
| tool.get_mcp_client() | ||
| call_kwargs = mock_client.call_args.kwargs | ||
| client = call_kwargs["http_client"] | ||
| assert _inject_otel_context in client.event_hooks["request"], ( | ||
| "Auto-created client should have OTel trace injection hook" | ||
| ) |
There was a problem hiding this comment.
test_mcp_streamable_http_tool_otel_trace_injection verifies the hook is registered, but it doesn’t assert that traceparent/tracestate headers are actually injected under an active span. Adding an assertion that _inject_otel_context mutates a request’s headers when a span is current would better cover the bug being fixed.
| client = call_kwargs["http_client"] | ||
| assert _inject_otel_context in client.event_hooks["request"], ( | ||
| "Auto-created client should have OTel trace injection hook" | ||
| ) |
There was a problem hiding this comment.
In this test, the first MCPStreamableHTTPTool instance (Test 1) creates an auto-managed httpx.AsyncClient, but the tool is never closed. That can leak resources and generate unclosed-client warnings. Ensure the tool is closed (e.g., await tool.close()) after the assertions.
| ) | |
| ) | |
| await tool.close() |
| async def _inject_otel_context(request: httpx.Request) -> None: # noqa: RUF029 | ||
| """Inject OpenTelemetry trace context into outgoing HTTP request headers. | ||
|
|
||
| Used as an httpx event hook to propagate the active span context | ||
| (W3C ``traceparent`` / ``tracestate``) so that remote MCP servers | ||
| can correlate their spans with the calling agent's trace. | ||
| """ | ||
| from opentelemetry.propagate import inject | ||
|
|
||
| inject(carrier=request.headers) | ||
|
|
There was a problem hiding this comment.
_inject_otel_context is declared async but contains no await (hence the # noqa: RUF029). Since httpx request event hooks can be regular callables too, consider making this a synchronous function and dropping the suppression to avoid unnecessary coroutine overhead and lint ignores.
2 participants
Add OpenTelemetry trace context propagation to MCPStreamableHTTPTool's httpx client via an event hook that injects W3C traceparent/tracestate headers into outgoing HTTP requests. This ensures the remote MCP server's spans are correctly parented under the agent framework's Execute Tool span.
closes #3623
Contribution Checklist