Extract nested tool guardrail logic in _run_impl

Author: steven10a

examples/basic/tool_guardrails.py

@@ -32,6 +32,7 @@ def get_user_data(user_id: str) -> dict[str, str]:
         "phone": "555-1234",
     }
 
+
 @function_tool
 def get_contact_info(user_id: str) -> dict[str, str]:
     """Get contact info by ID."""
@@ -122,7 +123,9 @@ async def main():
 
         # Example 2: Input guardrail triggers - function tool call is rejected but execution continues
         print("2. Attempting to send email with suspicious content:")
-        result = await Runner.run(agent, "Send an email to [email protected] introducing the company ACME corp.")
+        result = await Runner.run(
+            agent, "Send an email to [email protected] introducing the company ACME corp."
+        )
         print(f"❌ Guardrail rejected function tool call: {result.final_output}\n")
     except Exception as e:
         print(f"Error: {e}\n")
@@ -136,7 +139,6 @@ async def main():
         print("🚨 Output guardrail triggered: Execution halted for sensitive data")
         print(f"Details: {e.output.output_info}\n")
 
-    
     try:
         # Example 4: Output guardrail triggers - reject returning function tool output but continue execution
         print("4. Rejecting function tool output containing phone numbers:")
@@ -145,6 +147,7 @@ async def main():
     except Exception as e:
         print(f"Error: {e}\n")
 
+
 if __name__ == "__main__":
     asyncio.run(main())
 
@@ -157,12 +160,12 @@ async def main():
 ✅ Successful tool execution: I've sent a welcome email to [email protected] with an appropriate subject and greeting message.
 
 2. Attempting to send email with suspicious content:
-❌ Guardrail rejected function tool call: I'm unable to send the email mentioning ACME Corp as it was blocked by security guardrails.
+❌ Guardrail rejected function tool call: I'm unable to send the email as mentioning ACME Corp. is restricted.
 
 3. Attempting to get user data (contains SSN). Execution blocked:
 🚨 Output guardrail triggered: Execution halted for sensitive data
    Details: {'blocked_pattern': 'SSN', 'tool': 'get_user_data'}
 
 4. Rejecting function tool output containing sensitive data:
-✅ Successful tool execution: User data retrieved (phone number redacted for privacy)
+❌ Guardrail rejected function tool output: I'm unable to retrieve the contact info for user456 because it contains restricted information.
 """

src/agents/_run_impl.py

@@ -341,6 +341,8 @@ async def execute_tools_and_side_effects(
                 final_output=check_tool_use.final_output,
                 hooks=hooks,
                 context_wrapper=context_wrapper,
+                tool_input_guardrail_results=tool_input_guardrail_results,
+                tool_output_guardrail_results=tool_output_guardrail_results,
             )
 
         # Now we can check if the model also produced a final output
@@ -574,6 +576,155 @@ def process_model_response(
             mcp_approval_requests=mcp_approval_requests,
         )
 
+    @classmethod
+    async def _execute_input_guardrails(
+        cls,
+        *,
+        func_tool: FunctionTool,
+        tool_context: ToolContext[TContext],
+        agent: Agent[TContext],
+        tool_input_guardrail_results: list[ToolInputGuardrailResult],
+    ) -> str | None:
+        """Execute input guardrails for a tool.
+
+        Args:
+            func_tool: The function tool being executed.
+            tool_context: The tool execution context.
+            agent: The agent executing the tool.
+            tool_input_guardrail_results: List to append guardrail results to.
+
+        Returns:
+            None if tool execution should proceed, or a message string if execution should be
+            skipped.
+
+        Raises:
+            ToolInputGuardrailTripwireTriggered: If a guardrail triggers an exception.
+        """
+        if not func_tool.tool_input_guardrails:
+            return None
+
+        for guardrail in func_tool.tool_input_guardrails:
+            gr_out = await guardrail.run(
+                ToolInputGuardrailData(
+                    context=tool_context,
+                    agent=agent,
+                )
+            )
+
+            # Store the guardrail result
+            tool_input_guardrail_results.append(
+                ToolInputGuardrailResult(
+                    guardrail=guardrail,
+                    output=gr_out,
+                )
+            )
+
+            # Handle different behavior types
+            if gr_out.behavior["type"] == "raise_exception":
+                raise ToolInputGuardrailTripwireTriggered(guardrail=guardrail, output=gr_out)
+            elif gr_out.behavior["type"] == "reject_content":
+                # Set final_result to the message and skip tool execution
+                return gr_out.behavior["message"]
+            elif gr_out.behavior["type"] == "allow":
+                # Continue to next guardrail or tool execution
+                continue
+
+        return None
+
+    @classmethod
+    async def _execute_output_guardrails(
+        cls,
+        *,
+        func_tool: FunctionTool,
+        tool_context: ToolContext[TContext],
+        agent: Agent[TContext],
+        real_result: Any,
+        tool_output_guardrail_results: list[ToolOutputGuardrailResult],
+    ) -> Any:
+        """Execute output guardrails for a tool.
+
+        Args:
+            func_tool: The function tool being executed.
+            tool_context: The tool execution context.
+            agent: The agent executing the tool.
+            real_result: The actual result from the tool execution.
+            tool_output_guardrail_results: List to append guardrail results to.
+
+        Returns:
+            The final result after guardrail processing (may be modified).
+
+        Raises:
+            ToolOutputGuardrailTripwireTriggered: If a guardrail triggers an exception.
+        """
+        if not func_tool.tool_output_guardrails:
+            return real_result
+
+        final_result = real_result
+        for output_guardrail in func_tool.tool_output_guardrails:
+            gr_out = await output_guardrail.run(
+                ToolOutputGuardrailData(
+                    context=tool_context,
+                    agent=agent,
+                    output=real_result,
+                )
+            )
+
+            # Store the guardrail result
+            tool_output_guardrail_results.append(
+                ToolOutputGuardrailResult(
+                    guardrail=output_guardrail,
+                    output=gr_out,
+                )
+            )
+
+            # Handle different behavior types
+            if gr_out.behavior["type"] == "raise_exception":
+                raise ToolOutputGuardrailTripwireTriggered(
+                    guardrail=output_guardrail, output=gr_out
+                )
+            elif gr_out.behavior["type"] == "reject_content":
+                # Override the result with the guardrail message
+                final_result = gr_out.behavior["message"]
+                break
+            elif gr_out.behavior["type"] == "allow":
+                # Continue to next guardrail
+                continue
+
+        return final_result
+
+    @classmethod
+    async def _execute_tool_with_hooks(
+        cls,
+        *,
+        func_tool: FunctionTool,
+        tool_context: ToolContext[TContext],
+        agent: Agent[TContext],
+        hooks: RunHooks[TContext],
+        tool_call: ResponseFunctionToolCall,
+    ) -> Any:
+        """Execute the core tool function with before/after hooks.
+
+        Args:
+            func_tool: The function tool being executed.
+            tool_context: The tool execution context.
+            agent: The agent executing the tool.
+            hooks: The run hooks to execute.
+            tool_call: The tool call details.
+
+        Returns:
+            The result from the tool execution.
+        """
+        await asyncio.gather(
+            hooks.on_tool_start(tool_context, agent, func_tool),
+            (
+                agent.hooks.on_tool_start(tool_context, agent, func_tool)
+                if agent.hooks
+                else _coro.noop_coroutine()
+            ),
+        )
+
+        return await func_tool.on_invoke_tool(tool_context, tool_call.arguments)
+
     @classmethod
     async def execute_function_tool_calls(
         cls,
@@ -603,83 +754,35 @@ async def run_single_tool(
                     span_fn.span_data.input = tool_call.arguments
                 try:
                     # 1) Run input tool guardrails, if any
-                    final_result: Any | None = None
-                    if func_tool.tool_input_guardrails:
-                        for guardrail in func_tool.tool_input_guardrails:
-                            gr_out = await guardrail.run(
-                                ToolInputGuardrailData(
-                                    context=tool_context,
-                                    agent=agent,
-                                )
-                            )
-                            
-                            # Store the guardrail result
-                            tool_input_guardrail_results.append(
-                                ToolInputGuardrailResult(
-                                    guardrail=guardrail,
-                                    output=gr_out,
-                                )
-                            )
-                            
-                            # Handle different behavior types
-                            if gr_out.behavior["type"] == "raise_exception":
-                                raise ToolInputGuardrailTripwireTriggered(
-                                    guardrail=guardrail, output=gr_out
-                                )
-                            elif gr_out.behavior["type"] == "reject_content":
-                                # Set final_result to the message and skip tool execution
-                                final_result = gr_out.behavior["message"]
-                                break
-                            elif gr_out.behavior["type"] == "allow":
-                                # Continue to next guardrail or tool execution
-                                continue
-
-                    if final_result is None:
+                    rejected_message = await cls._execute_input_guardrails(
+                        func_tool=func_tool,
+                        tool_context=tool_context,
+                        agent=agent,
+                        tool_input_guardrail_results=tool_input_guardrail_results,
+                    )
+
+                    if rejected_message is not None:
+                        # Input guardrail rejected the tool call
+                        final_result = rejected_message
+                    else:
                         # 2) Actually run the tool
-                        await asyncio.gather(
-                            hooks.on_tool_start(tool_context, agent, func_tool),
-                            (
-                                agent.hooks.on_tool_start(tool_context, agent, func_tool)
-                                if agent.hooks
-                                else _coro.noop_coroutine()
-                            ),
-                        )
-                        real_result = await func_tool.on_invoke_tool(
-                            tool_context, tool_call.arguments
+                        real_result = await cls._execute_tool_with_hooks(
+                            func_tool=func_tool,
+                            tool_context=tool_context,
+                            agent=agent,
+                            hooks=hooks,
+                            tool_call=tool_call,
                         )
 
                         # 3) Run output tool guardrails, if any
-                        final_result = real_result
-                        if func_tool.tool_output_guardrails:
-                            for output_guardrail in func_tool.tool_output_guardrails:
-                                gr_out = await output_guardrail.run(
-                                    ToolOutputGuardrailData(
-                                        context=tool_context,
-                                        agent=agent,
-                                        output=real_result,
-                                    )
-                                )
+                        final_result = await cls._execute_output_guardrails(
+                            func_tool=func_tool,
+                            tool_context=tool_context,
+                            agent=agent,
+                            real_result=real_result,
+                            tool_output_guardrail_results=tool_output_guardrail_results,
+                        )
 
-                                # Store the guardrail result
-                                tool_output_guardrail_results.append(
-                                    ToolOutputGuardrailResult(
-                                        guardrail=output_guardrail,
-                                        output=gr_out,
-                                    )
-                                )
-                                
-                                # Handle different behavior types
-                                if gr_out.behavior["type"] == "raise_exception":
-                                    raise ToolOutputGuardrailTripwireTriggered(
-                                        guardrail=output_guardrail, output=gr_out
-                                    )
-                                elif gr_out.behavior["type"] == "reject_content":
-                                    # Override the result with the guardrail message
-                                    final_result = gr_out.behavior["message"]
-                                    break
-                                elif gr_out.behavior["type"] == "allow":
-                                    # Continue to next guardrail
-                                    continue
                         # 4) Tool end hooks (with final result, which may have been overridden)
                         await asyncio.gather(
                             hooks.on_tool_end(tool_context, agent, func_tool, final_result),
@@ -932,6 +1035,8 @@ async def execute_handoffs(
             pre_step_items=pre_step_items,
             new_step_items=new_step_items,
             next_step=NextStepHandoff(new_agent),
+            tool_input_guardrail_results=[],
+            tool_output_guardrail_results=[],
         )
 
     @classmethod

src/agents/run.py

@@ -715,6 +715,8 @@ def run_streamed(
             max_turns=max_turns,
             input_guardrail_results=[],
             output_guardrail_results=[],
+            tool_input_guardrail_results=[],
+            tool_output_guardrail_results=[],
             _current_agent_output_schema=output_schema,
             trace=new_trace,
             context_wrapper=context_wrapper,