From 25badf6bc2b445ac670a747a9ae350e25546045f Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 15 Sep 2026 15:06:55 +0200 Subject: [PATCH] ref(openai-agents): Stop settings tokens on Invoke Agent spans --- .../openai_agents/patches/agent_run.py | 26 +- .../openai_agents/patches/runner.py | 2 - .../openai_agents/spans/invoke_agent.py | 7 +- .../openai_agents/test_openai_agents.py | 238 ------------------ 4 files changed, 8 insertions(+), 265 deletions(-) diff --git a/sentry_sdk/integrations/openai_agents/patches/agent_run.py b/sentry_sdk/integrations/openai_agents/patches/agent_run.py index 17e825f779..5a1c5deecc 100644 --- a/sentry_sdk/integrations/openai_agents/patches/agent_run.py +++ b/sentry_sdk/integrations/openai_agents/patches/agent_run.py @@ -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") @@ -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) @@ -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) @@ -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") @@ -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") diff --git a/sentry_sdk/integrations/openai_agents/patches/runner.py b/sentry_sdk/integrations/openai_agents/patches/runner.py index df254b12ff..fe225cb9a0 100644 --- a/sentry_sdk/integrations/openai_agents/patches/runner.py +++ b/sentry_sdk/integrations/openai_agents/patches/runner.py @@ -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, ) @@ -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, ) diff --git a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py index f902d867b1..1477dc8457 100644 --- a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py +++ b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py @@ -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 @@ -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"]: diff --git a/tests/integrations/openai_agents/test_openai_agents.py b/tests/integrations/openai_agents/test_openai_agents.py index 6d525ba7ac..cf67c09147 100644 --- a/tests/integrations/openai_agents/test_openai_agents.py +++ b/tests/integrations/openai_agents/test_openai_agents.py @@ -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, @@ -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,