Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 7 additions & 19 deletions sentry_sdk/integrations/openai_agents/patches/agent_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,9 +66,7 @@ def _maybe_start_agent_span(
if current_agent and current_agent != agent:
span = getattr(context_wrapper, "_sentry_agent_span", None)
if span:
update_invoke_agent_span(
span=span, context=context_wrapper, agent=agent
)
update_invoke_agent_span(span=span, agent=agent)
span.__exit__(None, None, None)
delattr(context_wrapper, "_sentry_agent_span")

Expand Down Expand Up @@ -122,9 +120,7 @@ async def _run_single_turn(
with capture_internal_exceptions():
span = getattr(context_wrapper, "_sentry_agent_span", None)
if span:
update_invoke_agent_span(
span=span, context=context_wrapper, agent=agent
)
update_invoke_agent_span(span=span, agent=agent)
span.__exit__(*exc_info)
delattr(context_wrapper, "_sentry_agent_span")
reraise(*exc_info)
Expand Down Expand Up @@ -202,9 +198,7 @@ async def _run_single_turn_streamed(
with capture_internal_exceptions():
span = getattr(context_wrapper, "_sentry_agent_span", None)
if span:
update_invoke_agent_span(
span=span, context=context_wrapper, agent=agent
)
update_invoke_agent_span(span=span, agent=agent)
span.__exit__(*exc_info)
delattr(context_wrapper, "_sentry_agent_span")
_close_streaming_workflow_span(agent)
Expand Down Expand Up @@ -257,16 +251,14 @@ async def _execute_handoffs(
_close_streaming_workflow_span(agent)
span = getattr(context_wrapper, "_sentry_agent_span", None)
if span:
update_invoke_agent_span(
span=span, context=context_wrapper, agent=agent
)
update_invoke_agent_span(span=span, agent=agent)
span.__exit__(*exc_info)
delattr(context_wrapper, "_sentry_agent_span")
reraise(*exc_info)

span = getattr(context_wrapper, "_sentry_agent_span", None)
if span:
update_invoke_agent_span(span=span, context=context_wrapper, agent=agent)
update_invoke_agent_span(span=span, agent=agent)
span.__exit__(None, None, None)
delattr(context_wrapper, "_sentry_agent_span")

Expand Down Expand Up @@ -306,18 +298,14 @@ async def _execute_final_output(
_close_streaming_workflow_span(agent)
span = getattr(context_wrapper, "_sentry_agent_span", None)
if span:
update_invoke_agent_span(
span=span, context=context_wrapper, agent=agent, output=final_output
)
update_invoke_agent_span(span=span, agent=agent, output=final_output)
span.__exit__(*exc_info)
delattr(context_wrapper, "_sentry_agent_span")
reraise(*exc_info)

span = getattr(context_wrapper, "_sentry_agent_span", None)
if span:
update_invoke_agent_span(
span=span, context=context_wrapper, agent=agent, output=final_output
)
update_invoke_agent_span(span=span, agent=agent, output=final_output)
span.__exit__(None, None, None)
delattr(context_wrapper, "_sentry_agent_span")

Expand Down
2 changes: 0 additions & 2 deletions sentry_sdk/integrations/openai_agents/patches/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,6 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
):
update_invoke_agent_span(
span=invoke_agent_span,
context=context_wrapper,
agent=agent,
)

Expand All @@ -208,7 +207,6 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":

update_invoke_agent_span(
span=invoke_agent_span,
context=run_result.context_wrapper,
agent=agent,
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
from sentry_sdk.utils import has_data_collection_enabled, safe_serialize

from ..consts import SPAN_ORIGIN
from ..utils import _set_agent_data, _set_usage_data
from ..utils import _set_agent_data

if TYPE_CHECKING:
from typing import Any, Optional
Expand Down Expand Up @@ -85,14 +85,9 @@ def invoke_agent_span(

def update_invoke_agent_span(
span: "Span",
context: "Optional[agents.RunContextWrapper]",
agent: "Optional[agents.Agent]",
output: "Any" = None,
) -> None:
# Add aggregated usage data from context_wrapper
if context and hasattr(context, "usage"):
_set_usage_data(span, context.usage)

client = sentry_sdk.get_client()
if has_data_collection_enabled(client.options):
if client.options["data_collection"]["gen_ai"]["outputs"]:
Expand Down
238 changes: 0 additions & 238 deletions tests/integrations/openai_agents/test_openai_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -3275,106 +3275,6 @@ def failing_tool(message: str) -> str:
assert execute_tool_span["status"] == "error"


@pytest.mark.asyncio
async def test_invoke_agent_span_includes_usage_data(
sentry_init,
capture_items,
test_agent,
get_model_response,
):
"""
Test that invoke_agent spans include aggregated usage data from context_wrapper.
This verifies the new functionality added to track token usage in invoke_agent spans.
"""
client = AsyncOpenAI(api_key="z")
model = OpenAIResponsesModel(model="gpt-4", openai_client=client)
agent = test_agent.clone(model=model)

response = get_model_response(
Response(
id="resp_123",
output=[
ResponseOutputMessage(
id="msg_123",
type="message",
status="completed",
content=[
ResponseOutputText(
text="Response with usage",
type="output_text",
annotations=[],
)
],
role="assistant",
)
],
parallel_tool_calls=False,
tool_choice="none",
tools=[],
created_at=10000000,
model="gpt-4.1-2025-04-14",
object="response",
usage=ResponseUsage(
input_tokens=10,
input_tokens_details=InputTokensDetails(
cached_tokens=0,
cache_write_tokens=0,
),
output_tokens=20,
output_tokens_details=OutputTokensDetails(
reasoning_tokens=5,
),
total_tokens=30,
),
),
serialize_pydantic=True,
)
with patch.object(
agent.model._client._client,
"send",
return_value=response,
) as _:
sentry_init(
integrations=[OpenAIAgentsIntegration()],
disabled_integrations=[StdlibIntegration],
traces_sample_rate=1.0,
send_default_pii=True,
)
items = capture_items("span")

result = await agents.Runner.run(
agent, "Test input", run_config=test_run_config
)

assert result is not None

sentry_sdk.flush()
spans = [item.payload for item in items]
invoke_agent_span = next(
span
for span in spans
if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT
)

# Verify invoke_agent span has usage data from context_wrapper
assert invoke_agent_span["name"] == "invoke_agent test_agent"
assert "gen_ai.usage.input_tokens" in invoke_agent_span["attributes"]
assert "gen_ai.usage.output_tokens" in invoke_agent_span["attributes"]
assert "gen_ai.usage.total_tokens" in invoke_agent_span["attributes"]

assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 10
assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20
assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 30
assert (
invoke_agent_span["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]
== 0
)
assert (
invoke_agent_span["attributes"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]
== 5
)


@pytest.mark.asyncio
async def test_ai_client_span_includes_response_model(
sentry_init,
Expand Down Expand Up @@ -3552,144 +3452,6 @@ async def test_ai_client_span_response_model_with_chat_completions(
)


@pytest.mark.asyncio
async def test_multiple_llm_calls_aggregate_usage(
sentry_init,
capture_items,
test_agent,
get_model_response,
):
"""
Test that invoke_agent spans show aggregated usage across multiple LLM calls
(e.g., when tools are used and multiple API calls are made).
"""

@agents.function_tool
def calculator(a: int, b: int) -> int:
"""Add two numbers"""
return a + b

client = AsyncOpenAI(api_key="test-key")
model = OpenAIResponsesModel(model="gpt-4", openai_client=client)
agent_with_tool = test_agent.clone(tools=[calculator], model=model)

tool_call_response = get_model_response(
Response(
id="resp_1",
output=[
ResponseFunctionToolCall(
id="call_123",
call_id="call_123",
name="calculator",
type="function_call",
arguments='{"a": 5, "b": 3}',
)
],
parallel_tool_calls=False,
tool_choice="none",
tools=[],
created_at=10000000,
model="gpt-4.1-2025-04-14",
object="response",
usage=ResponseUsage(
input_tokens=10,
input_tokens_details=InputTokensDetails(
cached_tokens=0,
cache_write_tokens=0,
),
output_tokens=5,
output_tokens_details=OutputTokensDetails(
reasoning_tokens=0,
),
total_tokens=15,
),
),
serialize_pydantic=True,
)

final_response = get_model_response(
Response(
id="resp_2",
output=[
ResponseOutputMessage(
id="msg_final",
type="message",
status="completed",
content=[
ResponseOutputText(
text="The result is 8",
type="output_text",
annotations=[],
)
],
role="assistant",
)
],
parallel_tool_calls=False,
tool_choice="none",
tools=[],
created_at=10000000,
model="gpt-4-0613",
object="response",
usage=ResponseUsage(
input_tokens=20,
input_tokens_details=InputTokensDetails(
cached_tokens=5,
cache_write_tokens=0,
),
output_tokens=15,
output_tokens_details=OutputTokensDetails(
reasoning_tokens=3,
),
total_tokens=35,
),
),
serialize_pydantic=True,
)
with patch.object(
agent_with_tool.model._client._client,
"send",
side_effect=[tool_call_response, final_response],
) as _:
sentry_init(
integrations=[OpenAIAgentsIntegration()],
disabled_integrations=[StdlibIntegration],
traces_sample_rate=1.0,
send_default_pii=True,
)

items = capture_items("span")

result = await agents.Runner.run(
agent_with_tool,
"What is 5 + 3?",
run_config=test_run_config,
)

assert result is not None

sentry_sdk.flush()
spans = [item.payload for item in items]

invoke_agent_span = spans[3]

# Verify invoke_agent span has aggregated usage from both API calls
# Total: 10 + 20 = 30 input tokens, 5 + 15 = 20 output tokens, 15 + 35 = 50 total
assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 30
assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20
assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 50
# Cached tokens should be aggregated: 0 + 5 = 5
assert (
invoke_agent_span["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]
== 5
)
# Reasoning tokens should be aggregated: 0 + 3 = 3
assert (
invoke_agent_span["attributes"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]
== 3
)


@pytest.mark.asyncio
async def test_invoke_agent_span_includes_response_model(
sentry_init,
Expand Down
Loading