From 412b0ac88fce12eb039a9ed69697efb17853813c Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Mon, 2 Feb 2026 10:27:02 +0100 Subject: [PATCH] fix(openai-agents): Inject propagation headers for HostedMCPTool when streaming --- .../openai_agents/patches/models.py | 14 ++ .../openai_agents/test_openai_agents.py | 149 +++++++++++++++++- 2 files changed, 162 insertions(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/openai_agents/patches/models.py b/sentry_sdk/integrations/openai_agents/patches/models.py index 9b57a55f1f..7d0b9660b5 100644 --- a/sentry_sdk/integrations/openai_agents/patches/models.py +++ b/sentry_sdk/integrations/openai_agents/patches/models.py @@ -145,7 +145,21 @@ async def wrapped_stream_response(*args: "Any", **kwargs: "Any") -> "Any": if len(args) > 1: span_kwargs["input"] = args[1] + hosted_tools = [] + if len(args) > 3: + mcp_tools = args[3] + + if mcp_tools is not None: + hosted_tools = [ + tool + for tool in mcp_tools + if isinstance(tool, HostedMCPTool) + ] + with ai_client_span(agent, span_kwargs) as span: + for hosted_tool in hosted_tools: + _inject_trace_propagation_headers(hosted_tool, span=span) + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) streaming_response = None diff --git a/tests/integrations/openai_agents/test_openai_agents.py b/tests/integrations/openai_agents/test_openai_agents.py index 4bf212b8f3..7e4f7faeb4 100644 --- a/tests/integrations/openai_agents/test_openai_agents.py +++ b/tests/integrations/openai_agents/test_openai_agents.py @@ -35,7 +35,13 @@ from agents.exceptions import MaxTurnsExceeded, ModelBehaviorError from agents.version import __version__ as OPENAI_AGENTS_VERSION -from openai.types.responses import Response, ResponseUsage +from openai.types.responses import ( + ResponseCreatedEvent, + ResponseTextDeltaEvent, + ResponseCompletedEvent, + Response, + ResponseUsage, +) from openai.types.responses.response_usage import ( InputTokensDetails, OutputTokensDetails, @@ -80,6 +86,63 @@ ) +async def EXAMPLE_STREAMED_RESPONSE(*args, **kwargs): + yield ResponseCreatedEvent( + response=Response( + id="chat-id", + output=[], + parallel_tool_calls=False, + tool_choice="none", + tools=[], + created_at=10000000, + model="response-model-id", + object="response", + ), + type="response.created", + sequence_number=0, + ) + + yield ResponseCompletedEvent( + response=Response( + id="chat-id", + output=[ + ResponseOutputMessage( + id="message-id", + content=[ + ResponseOutputText( + annotations=[], + text="the model response", + type="output_text", + ), + ], + role="assistant", + status="completed", + type="message", + ), + ], + parallel_tool_calls=False, + tool_choice="none", + tools=[], + created_at=10000000, + model="response-model-id", + object="response", + usage=ResponseUsage( + input_tokens=20, + input_tokens_details=InputTokensDetails( + cached_tokens=5, + ), + output_tokens=10, + output_tokens_details=OutputTokensDetails( + reasoning_tokens=8, + ), + total_tokens=30, + ), + ), + type="response.completed", + sequence_number=1, + ) + + @pytest.fixture def mock_usage(): return Usage( @@ -1172,6 +1235,90 @@ def simple_test_tool(message: str) -> str: assert ai_client_span2["data"]["gen_ai.usage.total_tokens"] == 25 +@pytest.mark.asyncio +async def test_hosted_mcp_tool_propagation_header_streamed(sentry_init, test_agent): + """ + Test responses API is given trace propagation headers with HostedMCPTool. + """ + + hosted_tool = HostedMCPTool( + tool_config={ + "type": "mcp", + "server_label": "test_server", + "server_url": "http://example.com/", + "headers": { + "baggage": "custom=data", + }, + }, + ) + + client = AsyncOpenAI(api_key="z") + client.responses._post = AsyncMock(return_value=EXAMPLE_RESPONSE) + + model = OpenAIResponsesModel(model="gpt-4", openai_client=client) + + agent_with_tool = test_agent.clone( + tools=[hosted_tool], + model=model, + ) + + sentry_init( + integrations=[OpenAIAgentsIntegration()], + traces_sample_rate=1.0, + release="d08ebdb9309e1b004c6f52202de58a09c2268e42", + ) + + with patch.object( + model._client.responses, + "create", + side_effect=EXAMPLE_STREAMED_RESPONSE, + ) as create, mock.patch( + "sentry_sdk.tracing_utils.Random.randrange", return_value=500000 + ): + with sentry_sdk.start_transaction( + name="/interactions/other-dogs/new-dog", + op="greeting.sniff", + trace_id="01234567890123456789012345678901", + ) as transaction: + result = agents.Runner.run_streamed( + agent_with_tool, + "Please use the simple test tool", + run_config=test_run_config, + ) + + async for event in result.stream_events(): + pass + + ai_client_span = transaction._span_recorder.spans[-1] + + args, kwargs = create.call_args + + assert "tools" in kwargs + assert len(kwargs["tools"]) == 1 + hosted_mcp_tool = kwargs["tools"][0] + + assert hosted_mcp_tool["headers"][ + "sentry-trace" + ] == "{trace_id}-{parent_span_id}-{sampled}".format( + trace_id=transaction.trace_id, + parent_span_id=ai_client_span.span_id, + sampled=1, + ) + + expected_outgoing_baggage = ( + "custom=data," + "sentry-trace_id=01234567890123456789012345678901," + "sentry-sample_rand=0.500000," + "sentry-environment=production," + "sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42," + "sentry-transaction=/interactions/other-dogs/new-dog," + "sentry-sample_rate=1.0," + "sentry-sampled=true" + ) + + assert hosted_mcp_tool["headers"]["baggage"] == expected_outgoing_baggage + + @pytest.mark.asyncio async def test_hosted_mcp_tool_propagation_headers(sentry_init, test_agent): """