Skip to content

Commit a47ef34

Browse files
feat(openai-agents): Inject propagation headers for HostedMCPTool (#5297)
Adds `sentry-trace` and `baggage` to headers field of hosted tool definitions sent to the OpenAI Responses API. - Factor out logic from the `httpx` integration to prevent duplicate Sentry baggage. - Inject headers in the patch for the wrapper of the HTTP call to OpenAI Responses.
1 parent 6db0eb7 commit a47ef34

File tree

4 files changed

+195
-24
lines changed

4 files changed

+195
-24
lines changed

sentry_sdk/integrations/httpx.py

Lines changed: 2 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@
44
from sentry_sdk.integrations import Integration, DidNotEnable
55
from sentry_sdk.tracing import BAGGAGE_HEADER_NAME
66
from sentry_sdk.tracing_utils import (
7-
Baggage,
87
should_propagate_trace,
98
add_http_request_source,
9+
add_sentry_baggage_to_headers,
1010
)
1111
from sentry_sdk.utils import (
1212
SENSITIVE_DATA_SUBSTITUTE,
@@ -19,7 +19,6 @@
1919
from typing import TYPE_CHECKING
2020

2121
if TYPE_CHECKING:
22-
from collections.abc import MutableMapping
2322
from typing import Any
2423

2524

@@ -81,7 +80,7 @@ def send(self: "Client", request: "Request", **kwargs: "Any") -> "Response":
8180
)
8281

8382
if key == BAGGAGE_HEADER_NAME:
84-
_add_sentry_baggage_to_headers(request.headers, value)
83+
add_sentry_baggage_to_headers(request.headers, value)
8584
else:
8685
request.headers[key] = value
8786

@@ -155,22 +154,3 @@ async def send(
155154
return rv
156155

157156
AsyncClient.send = send
158-
159-
160-
def _add_sentry_baggage_to_headers(
161-
headers: "MutableMapping[str, str]", sentry_baggage: str
162-
) -> None:
163-
"""Add the Sentry baggage to the headers.
164-
165-
This function directly mutates the provided headers. The provided sentry_baggage
166-
is appended to the existing baggage. If the baggage already contains Sentry items,
167-
they are stripped out first.
168-
"""
169-
existing_baggage = headers.get(BAGGAGE_HEADER_NAME, "")
170-
stripped_existing_baggage = Baggage.strip_sentry_baggage(existing_baggage)
171-
172-
separator = "," if len(stripped_existing_baggage) > 0 else ""
173-
174-
headers[BAGGAGE_HEADER_NAME] = (
175-
stripped_existing_baggage + separator + sentry_baggage
176-
)

sentry_sdk/integrations/openai_agents/patches/models.py

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,20 +4,57 @@
44
from sentry_sdk.integrations import DidNotEnable
55

66
from ..spans import ai_client_span, update_ai_client_span
7+
8+
import sentry_sdk
79
from sentry_sdk.consts import SPANDATA
10+
from sentry_sdk.utils import logger
11+
from sentry_sdk.tracing import BAGGAGE_HEADER_NAME
12+
from sentry_sdk.tracing_utils import (
13+
should_propagate_trace,
14+
add_sentry_baggage_to_headers,
15+
)
816

917
from typing import TYPE_CHECKING
1018

1119
if TYPE_CHECKING:
1220
from typing import Any, Callable
13-
21+
from sentry_sdk.tracing import Span
1422

1523
try:
1624
import agents
25+
from agents.tool import HostedMCPTool
1726
except ImportError:
1827
raise DidNotEnable("OpenAI Agents not installed")
1928

2029

30+
def _inject_trace_propagation_headers(
31+
hosted_tool: "HostedMCPTool", span: "Span"
32+
) -> None:
33+
headers = hosted_tool.tool_config.get("headers")
34+
if headers is None:
35+
headers = {}
36+
hosted_tool.tool_config["headers"] = headers
37+
38+
mcp_url = hosted_tool.tool_config.get("server_url")
39+
if not mcp_url:
40+
return
41+
42+
if should_propagate_trace(sentry_sdk.get_client(), mcp_url):
43+
for (
44+
key,
45+
value,
46+
) in sentry_sdk.get_current_scope().iter_trace_propagation_headers(span=span):
47+
logger.debug(
48+
"[Tracing] Adding `{key}` header {value} to outgoing request to {mcp_url}.".format(
49+
key=key, value=value, mcp_url=mcp_url
50+
)
51+
)
52+
if key == BAGGAGE_HEADER_NAME:
53+
add_sentry_baggage_to_headers(headers, value)
54+
else:
55+
headers[key] = value
56+
57+
2158
def _create_get_model_wrapper(
2259
original_get_model: "Callable[..., Any]",
2360
) -> "Callable[..., Any]":
@@ -54,7 +91,17 @@ async def wrapped_fetch_response(*args: "Any", **kwargs: "Any") -> "Any":
5491

5592
@wraps(original_get_response)
5693
async def wrapped_get_response(*args: "Any", **kwargs: "Any") -> "Any":
94+
mcp_tools = kwargs.get("tools")
95+
hosted_tools = []
96+
if mcp_tools is not None:
97+
hosted_tools = [
98+
tool for tool in mcp_tools if isinstance(tool, HostedMCPTool)
99+
]
100+
57101
with ai_client_span(agent, kwargs) as span:
102+
for hosted_tool in hosted_tools:
103+
_inject_trace_propagation_headers(hosted_tool, span=span)
104+
58105
result = await original_get_response(*args, **kwargs)
59106

60107
response_model = getattr(agent, "_sentry_raw_response_model", None)

sentry_sdk/tracing_utils.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import os
55
import re
66
import sys
7-
from collections.abc import Mapping
7+
from collections.abc import Mapping, MutableMapping
88
from datetime import timedelta
99
from random import Random
1010
from urllib.parse import quote, unquote
@@ -1285,6 +1285,25 @@ def _should_continue_trace(baggage: "Optional[Baggage]") -> bool:
12851285
return True
12861286

12871287

1288+
def add_sentry_baggage_to_headers(
1289+
headers: "MutableMapping[str, str]", sentry_baggage: str
1290+
) -> None:
1291+
"""Add the Sentry baggage to the headers.
1292+
1293+
This function directly mutates the provided headers. The provided sentry_baggage
1294+
is appended to the existing baggage. If the baggage already contains Sentry items,
1295+
they are stripped out first.
1296+
"""
1297+
existing_baggage = headers.get(BAGGAGE_HEADER_NAME, "")
1298+
stripped_existing_baggage = Baggage.strip_sentry_baggage(existing_baggage)
1299+
1300+
separator = "," if len(stripped_existing_baggage) > 0 else ""
1301+
1302+
headers[BAGGAGE_HEADER_NAME] = (
1303+
stripped_existing_baggage + separator + sentry_baggage
1304+
)
1305+
1306+
12881307
# Circular imports
12891308
from sentry_sdk.tracing import (
12901309
BAGGAGE_HEADER_NAME,

tests/integrations/openai_agents/test_openai_agents.py

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@
1212
from sentry_sdk.integrations.openai_agents.utils import _set_input_data, safe_serialize
1313
from sentry_sdk.utils import parse_version
1414

15+
from openai import AsyncOpenAI
16+
from agents.models.openai_responses import OpenAIResponsesModel
17+
18+
from unittest import mock
19+
from unittest.mock import AsyncMock
20+
1521
import agents
1622
from agents import (
1723
Agent,
@@ -25,16 +31,54 @@
2531
ResponseOutputText,
2632
ResponseFunctionToolCall,
2733
)
34+
from agents.tool import HostedMCPTool
2835
from agents.exceptions import MaxTurnsExceeded, ModelBehaviorError
2936
from agents.version import __version__ as OPENAI_AGENTS_VERSION
3037

38+
from openai.types.responses import Response, ResponseUsage
3139
from openai.types.responses.response_usage import (
3240
InputTokensDetails,
3341
OutputTokensDetails,
3442
)
3543

3644
test_run_config = agents.RunConfig(tracing_disabled=True)
3745

46+
EXAMPLE_RESPONSE = Response(
47+
id="chat-id",
48+
output=[
49+
ResponseOutputMessage(
50+
id="message-id",
51+
content=[
52+
ResponseOutputText(
53+
annotations=[],
54+
text="the model response",
55+
type="output_text",
56+
),
57+
],
58+
role="assistant",
59+
status="completed",
60+
type="message",
61+
),
62+
],
63+
parallel_tool_calls=False,
64+
tool_choice="none",
65+
tools=[],
66+
created_at=10000000,
67+
model="response-model-id",
68+
object="response",
69+
usage=ResponseUsage(
70+
input_tokens=20,
71+
input_tokens_details=InputTokensDetails(
72+
cached_tokens=5,
73+
),
74+
output_tokens=10,
75+
output_tokens_details=OutputTokensDetails(
76+
reasoning_tokens=8,
77+
),
78+
total_tokens=30,
79+
),
80+
)
81+
3882

3983
@pytest.fixture
4084
def mock_usage():
@@ -695,6 +739,87 @@ def simple_test_tool(message: str) -> str:
695739
assert ai_client_span2["data"]["gen_ai.usage.total_tokens"] == 25
696740

697741

742+
@pytest.mark.asyncio
743+
async def test_hosted_mcp_tool_propagation_headers(sentry_init, test_agent):
744+
"""
745+
Test responses API is given trace propagation headers with HostedMCPTool.
746+
"""
747+
748+
hosted_tool = HostedMCPTool(
749+
tool_config={
750+
"type": "mcp",
751+
"server_label": "test_server",
752+
"server_url": "http://example.com/",
753+
"headers": {
754+
"baggage": "custom=data",
755+
},
756+
},
757+
)
758+
759+
client = AsyncOpenAI(api_key="z")
760+
client.responses._post = AsyncMock(return_value=EXAMPLE_RESPONSE)
761+
762+
model = OpenAIResponsesModel(model="gpt-4", openai_client=client)
763+
764+
agent_with_tool = test_agent.clone(
765+
tools=[hosted_tool],
766+
model=model,
767+
)
768+
769+
sentry_init(
770+
integrations=[OpenAIAgentsIntegration()],
771+
traces_sample_rate=1.0,
772+
release="d08ebdb9309e1b004c6f52202de58a09c2268e42",
773+
)
774+
775+
with patch.object(
776+
model._client.responses,
777+
"create",
778+
wraps=model._client.responses.create,
779+
) as create, mock.patch(
780+
"sentry_sdk.tracing_utils.Random.randrange", return_value=500000
781+
):
782+
with sentry_sdk.start_transaction(
783+
name="/interactions/other-dogs/new-dog",
784+
op="greeting.sniff",
785+
trace_id="01234567890123456789012345678901",
786+
) as transaction:
787+
await agents.Runner.run(
788+
agent_with_tool,
789+
"Please use the simple test tool",
790+
run_config=test_run_config,
791+
)
792+
793+
ai_client_span = transaction._span_recorder.spans[-1]
794+
795+
args, kwargs = create.call_args
796+
797+
assert "tools" in kwargs
798+
assert len(kwargs["tools"]) == 1
799+
hosted_mcp_tool = kwargs["tools"][0]
800+
801+
assert hosted_mcp_tool["headers"][
802+
"sentry-trace"
803+
] == "{trace_id}-{parent_span_id}-{sampled}".format(
804+
trace_id=transaction.trace_id,
805+
parent_span_id=ai_client_span.span_id,
806+
sampled=1,
807+
)
808+
809+
expected_outgoing_baggage = (
810+
"custom=data,"
811+
"sentry-trace_id=01234567890123456789012345678901,"
812+
"sentry-sample_rand=0.500000,"
813+
"sentry-environment=production,"
814+
"sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42,"
815+
"sentry-transaction=/interactions/other-dogs/new-dog,"
816+
"sentry-sample_rate=1.0,"
817+
"sentry-sampled=true"
818+
)
819+
820+
assert hosted_mcp_tool["headers"]["baggage"] == expected_outgoing_baggage
821+
822+
698823
@pytest.mark.asyncio
699824
async def test_model_behavior_error(sentry_init, capture_events, test_agent):
700825
"""

0 commit comments

Comments
 (0)