@@ -166,55 +166,52 @@ async def _run_single_turn_streamed(
166166 return result
167167
168168
169- def _patch_agent_run () -> None :
169+ async def _execute_handoffs (
170+ original_execute_handoffs : "Callable[..., SingleStepResult]" ,
171+ * args : "Any" ,
172+ ** kwargs : "Any" ,
173+ ) -> "SingleStepResult" :
170174 """
171- Patches AgentRunner methods to create agent invocation spans.
172- This directly patches the execution flow to track when agents start and stop.
175+ Patched execute_handoffs that
176+ - creates and manages handoff spans.
177+ - ends the agent invocation span.
178+ - ends the workflow span if the response is streamed and an exception is raised in `execute_handoffs()`.
173179 """
174180
175- # Store original methods
176- original_execute_handoffs = agents . _run_impl . RunImpl . execute_handoffs
177- original_execute_final_output = agents . _run_impl . RunImpl . execute_final_output
181+ context_wrapper = kwargs . get ( "context_wrapper" )
182+ run_handoffs = kwargs . get ( "run_handoffs" )
183+ agent = kwargs . get ( "agent" )
178184
179- @wraps (
180- original_execute_handoffs .__func__
181- if hasattr (original_execute_handoffs , "__func__" )
182- else original_execute_handoffs
183- )
184- async def patched_execute_handoffs (
185- cls : "agents.Runner" , * args : "Any" , ** kwargs : "Any"
186- ) -> "Any" :
187- """
188- Patched execute_handoffs that
189- - creates and manages handoff spans.
190- - ends the agent invocation span.
191- - ends the workflow span if the response is streamed and an exception is raised in `execute_handoffs()`.
192- """
185+ # Create Sentry handoff span for the first handoff (agents library only processes the first one)
186+ if run_handoffs :
187+ first_handoff = run_handoffs [0 ]
188+ handoff_agent_name = first_handoff .handoff .agent_name
189+ handoff_span (context_wrapper , agent , handoff_agent_name )
193190
194- context_wrapper = kwargs .get ("context_wrapper" )
195- run_handoffs = kwargs .get ("run_handoffs" )
196- agent = kwargs .get ("agent" )
191+ # Call original method with all parameters
192+ try :
193+ result = await original_execute_handoffs (* args , ** kwargs )
194+ except Exception :
195+ exc_info = sys .exc_info ()
196+ with capture_internal_exceptions ():
197+ _close_streaming_workflow_span (agent )
198+ reraise (* exc_info )
199+ finally :
200+ # End span for current agent after handoff processing is complete
201+ if agent and context_wrapper and _has_active_agent_span (context_wrapper ):
202+ end_invoke_agent_span (context_wrapper , agent )
197203
198- # Create Sentry handoff span for the first handoff (agents library only processes the first one)
199- if run_handoffs :
200- first_handoff = run_handoffs [0 ]
201- handoff_agent_name = first_handoff .handoff .agent_name
202- handoff_span (context_wrapper , agent , handoff_agent_name )
204+ return result
203205
204- # Call original method with all parameters
205- try :
206- result = await original_execute_handoffs (* args , ** kwargs )
207- except Exception :
208- exc_info = sys .exc_info ()
209- with capture_internal_exceptions ():
210- _close_streaming_workflow_span (agent )
211- reraise (* exc_info )
212- finally :
213- # End span for current agent after handoff processing is complete
214- if agent and context_wrapper and _has_active_agent_span (context_wrapper ):
215- end_invoke_agent_span (context_wrapper , agent )
216206
217- return result
207+ def _patch_agent_run () -> None :
208+ """
209+ Patches AgentRunner methods to create agent invocation spans.
210+ This directly patches the execution flow to track when agents start and stop.
211+ """
212+
213+ # Store original methods
214+ original_execute_final_output = agents ._run_impl .RunImpl .execute_final_output
218215
219216 @wraps (
220217 original_execute_final_output .__func__
@@ -250,7 +247,6 @@ async def patched_execute_final_output(
250247 return result
251248
252249 # Apply patches
253- agents ._run_impl .RunImpl .execute_handoffs = classmethod (patched_execute_handoffs )
254250 agents ._run_impl .RunImpl .execute_final_output = classmethod (
255251 patched_execute_final_output
256252 )
0 commit comments