Skip to content
Open
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
34 changes: 2 additions & 32 deletions sentry_sdk/integrations/openai_agents/patches/agent_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,14 +36,6 @@ def _get_current_agent(
return getattr(context_wrapper, "_sentry_current_agent", None)


def _close_streaming_workflow_span(agent: "Optional[agents.Agent]") -> None:
"""Close the workflow span for streaming executions if it exists."""
if agent and hasattr(agent, "_sentry_workflow_span"):
workflow_span = agent._sentry_workflow_span
workflow_span.__exit__(*sys.exc_info())
delattr(agent, "_sentry_workflow_span")


def _maybe_start_agent_span(
context_wrapper: "Optional[agents.RunContextWrapper]",
agent: "Optional[agents.Agent]",
Expand Down Expand Up @@ -201,7 +193,6 @@ async def _run_single_turn_streamed(
update_invoke_agent_span(span=span, agent=agent)
span.__exit__(*exc_info)
delattr(context_wrapper, "_sentry_agent_span")
_close_streaming_workflow_span(agent)
reraise(*exc_info)

return result
Expand All @@ -216,7 +207,6 @@ async def _execute_handoffs(
Patched execute_handoffs that
- creates and manages handoff spans.
- ends the agent invocation span.
- ends the workflow span if the response is streamed and an exception is raised in `execute_handoffs()`.
"""

context_wrapper: "Optional[agents.RunContextWrapper]" = kwargs.get(
Expand All @@ -232,23 +222,12 @@ async def _execute_handoffs(
handoff_agent_name = first_handoff.handoff.agent_name
handoff_span(context_wrapper, agent, handoff_agent_name)

if not agent or not context_wrapper or not _has_active_agent_span(context_wrapper):
# Call original method with all parameters
try:
return await original_execute_handoffs(*args, **kwargs)
except Exception:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_close_streaming_workflow_span(agent)
reraise(*exc_info)

# Call original method with all parameters
try:
result = await original_execute_handoffs(*args, **kwargs)
except Exception:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_close_streaming_workflow_span(agent)
span = getattr(context_wrapper, "_sentry_agent_span", None)
if span:
update_invoke_agent_span(span=span, agent=agent)
Expand All @@ -271,9 +250,7 @@ async def _execute_final_output(
**kwargs: "Any",
) -> "SingleStepResult":
"""
Patched execute_final_output that
- ends the agent invocation span.
- ends the workflow span if the response is streamed.
Patched execute_final_output that ends the agent invocation span.
"""

# openai-agents >= 0.14 renamed `agent` to `public_agent`.
Expand All @@ -282,20 +259,13 @@ async def _execute_final_output(
final_output = kwargs.get("final_output")

if not agent or not context_wrapper or not _has_active_agent_span(context_wrapper):
try:
return await original_execute_final_output(*args, **kwargs)
finally:
with capture_internal_exceptions():
# For streaming, close the workflow span (non-streaming uses context manager in _create_run_wrapper)
_close_streaming_workflow_span(agent)
return await original_execute_final_output(*args, **kwargs)

try:
result = await original_execute_final_output(*args, **kwargs)
except Exception:
exc_info = sys.exc_info()
with capture_internal_exceptions():
# For streaming, close the workflow span (non-streaming uses context manager in _create_run_wrapper)
_close_streaming_workflow_span(agent)
span = getattr(context_wrapper, "_sentry_agent_span", None)
if span:
update_invoke_agent_span(span=span, agent=agent, output=final_output)
Expand Down
180 changes: 57 additions & 123 deletions sentry_sdk/integrations/openai_agents/patches/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
)

from ..spans import (
agent_workflow_span,
execute_tool_span,
update_execute_tool_span,
update_invoke_agent_span,
Expand All @@ -28,7 +27,7 @@
from typing import TYPE_CHECKING, TypeVar

if TYPE_CHECKING:
from typing import Any, AsyncIterator, Callable
from typing import Any, Callable

from agents import Agent, Tool, ToolContext

Expand Down Expand Up @@ -126,7 +125,6 @@ def _create_run_wrapper(
) -> "Callable[..., Any]":
"""
Wraps the agents.Runner.run methods to
- create and manage a root span for the agent workflow runs.
- end the agent invocation span if an `AgentsException` is raised in `run()`.

Note agents.Runner.run_sync() is a wrapper around agents.Runner.run(),
Expand All @@ -150,88 +148,73 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
else:
agent = args[0].clone()

with agent_workflow_span(agent) as workflow_span:
# Set conversation ID on workflow span early so it's captured even on errors
conversation_id = kwargs.get("conversation_id")
if conversation_id:
agent._sentry_conversation_id = conversation_id

workflow_span.set_attribute(
SPANDATA.GEN_AI_CONVERSATION_ID, conversation_id
)

if "starting_agent" in kwargs:
kwargs["starting_agent"] = agent
else:
args = (agent, *args[1:])

try:
run_result = await original_func(*args, **kwargs)
except AgentsException as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)

context_wrapper = getattr(exc.run_data, "context_wrapper", None)
if context_wrapper is not None:
invoke_agent_span = getattr(
context_wrapper, "_sentry_agent_span", None
# Set conversation ID on workflow span early so it's captured even on errors
conversation_id = kwargs.get("conversation_id")
if conversation_id:
agent._sentry_conversation_id = conversation_id

if "starting_agent" in kwargs:
kwargs["starting_agent"] = agent
else:
args = (agent, *args[1:])

try:
run_result = await original_func(*args, **kwargs)
except AgentsException as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
_capture_exception(exc)

context_wrapper = getattr(exc.run_data, "context_wrapper", None)
if context_wrapper is not None:
invoke_agent_span = getattr(
context_wrapper, "_sentry_agent_span", None
)

if (
invoke_agent_span is not None
and invoke_agent_span.end_timestamp is None
):
update_invoke_agent_span(
span=invoke_agent_span,
agent=agent,
)

if (
invoke_agent_span is not None
and invoke_agent_span.end_timestamp is None
):
update_invoke_agent_span(
span=invoke_agent_span,
agent=agent,
)

invoke_agent_span.__exit__(*exc_info)
delattr(context_wrapper, "_sentry_agent_span")
reraise(*exc_info)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
# Invoke agent span is not finished in this case.
# This is much less likely to occur than other cases because
# AgentRunner.run() is "just" a while loop around _run_single_turn.
_capture_exception(exc)
reraise(*exc_info)

invoke_agent_span = getattr(
run_result.context_wrapper, "_sentry_agent_span", None
)
if not invoke_agent_span:
return run_result

update_invoke_agent_span(
span=invoke_agent_span,
agent=agent,
)

invoke_agent_span.__exit__(None, None, None)
delattr(run_result.context_wrapper, "_sentry_agent_span")
invoke_agent_span.__exit__(*exc_info)
delattr(context_wrapper, "_sentry_agent_span")
reraise(*exc_info)
except Exception as exc:
exc_info = sys.exc_info()
with capture_internal_exceptions():
# Invoke agent span is not finished in this case.
# This is much less likely to occur than other cases because
# AgentRunner.run() is "just" a while loop around _run_single_turn.
_capture_exception(exc)
reraise(*exc_info)

invoke_agent_span = getattr(
run_result.context_wrapper, "_sentry_agent_span", None
)
if not invoke_agent_span:
return run_result

update_invoke_agent_span(
span=invoke_agent_span,
agent=agent,
)

invoke_agent_span.__exit__(None, None, None)
delattr(run_result.context_wrapper, "_sentry_agent_span")
return run_result

return wrapper


def _create_run_streamed_wrapper(
original_func: "Callable[..., Any]",
) -> "Callable[..., Any]":
"""
Wraps the agents.Runner.run_streamed method to
- create a root span for streaming agent workflow runs.
- end the workflow span if and only if the response stream is consumed or cancelled.

Unlike run(), run_streamed() returns immediately with a RunResultStreaming object
while execution continues in a background task. The workflow span must stay open
throughout the streaming operation and close when streaming completes or is abandoned.

Note: We don't use isolation_scope() here because it uses context variables that
cannot span async boundaries (the __enter__ and __exit__ would be called from
different async contexts, causing ValueError).
Wraps the agents.Runner.run_streamed method to inject run hooks.
"""

@wraps(original_func)
Expand All @@ -247,19 +230,6 @@ def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
if conversation_id:
agent._sentry_conversation_id = conversation_id

# Start workflow span immediately (before run_streamed returns)
workflow_span = agent_workflow_span(agent)
workflow_span.__enter__()

# Set conversation ID on workflow span early so it's captured even on errors
if conversation_id:
workflow_span.set_attribute(
SPANDATA.GEN_AI_CONVERSATION_ID, conversation_id
)

# Store span on agent for cleanup
agent._sentry_workflow_span = workflow_span

if "starting_agent" in kwargs:
kwargs["starting_agent"] = agent
else:
Expand All @@ -276,45 +246,9 @@ def wrapper(*args: "Any", **kwargs: "Any") -> "Any":
# Call original function to get RunResultStreaming
run_result = original_func(*args, **kwargs)
except Exception as exc:
# If run_streamed itself fails (not the background task), clean up immediately
workflow_span.__exit__(*sys.exc_info())
_capture_exception(exc)
raise

def _close_workflow_span() -> None:
if hasattr(agent, "_sentry_workflow_span"):
workflow_span.__exit__(*sys.exc_info())
delattr(agent, "_sentry_workflow_span")

if hasattr(run_result, "stream_events"):
original_stream_events = run_result.stream_events

@wraps(original_stream_events)
async def wrapped_stream_events(
*stream_args: "Any", **stream_kwargs: "Any"
) -> "AsyncIterator[Any]":
try:
async for event in original_stream_events(
*stream_args, **stream_kwargs
):
yield event
finally:
_close_workflow_span()

run_result.stream_events = wrapped_stream_events

if hasattr(run_result, "cancel"):
original_cancel = run_result.cancel

@wraps(original_cancel)
def wrapped_cancel(*cancel_args: "Any", **cancel_kwargs: "Any") -> "Any":
try:
return original_cancel(*cancel_args, **cancel_kwargs)
finally:
_close_workflow_span()

run_result.cancel = wrapped_cancel

return run_result

return wrapper
1 change: 0 additions & 1 deletion sentry_sdk/integrations/openai_agents/spans/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
from .agent_workflow import agent_workflow_span # noqa: F401
from .ai_client import ai_client_context, update_ai_client_span # noqa: F401
from .execute_tool import execute_tool_span, update_execute_tool_span # noqa: F401
from .handoff import handoff_span # noqa: F401
Expand Down
16 changes: 0 additions & 16 deletions sentry_sdk/integrations/openai_agents/spans/agent_workflow.py

This file was deleted.

Loading
Loading