code review: keep output guardrails

Author: habema

src/agents/run.py

@@ -1085,24 +1085,6 @@ async def _start_streaming(
                     if server_conversation_tracker is not None:
                         server_conversation_tracker.track_server_items(turn_result.model_response)
 
-                    # Check for soft cancel after tool execution completes (before next step)
-                    if streamed_result._cancel_mode == "after_turn":  # type: ignore[comparison-overlap]
-                        # Save session with complete tool execution (tool calls + tool results)
-                        if session is not None:
-                            should_skip_session_save = (
-                                await AgentRunner._input_guardrail_tripwire_triggered_for_stream(
-                                    streamed_result
-                                )
-                            )
-                            if should_skip_session_save is False:
-                                await AgentRunner._save_result_to_session(
-                                    session, [], turn_result.new_step_items
-                                )
-
-                        streamed_result.is_complete = True
-                        streamed_result._event_queue.put_nowait(QueueCompleteSentinel())
-                        break
-
                     if isinstance(turn_result.next_step, NextStepHandoff):
                         # Save the conversation to session if enabled (before handoff)
                         # Note: Non-streaming path doesn't save handoff turns immediately,
@@ -1125,6 +1107,12 @@ async def _start_streaming(
                         streamed_result._event_queue.put_nowait(
                             AgentUpdatedStreamEvent(new_agent=current_agent)
                         )
+
+                        # Check for soft cancel after handoff (before next turn)
+                        if streamed_result._cancel_mode == "after_turn":  # type: ignore[comparison-overlap]
+                            streamed_result.is_complete = True
+                            streamed_result._event_queue.put_nowait(QueueCompleteSentinel())
+                            break
                     elif isinstance(turn_result.next_step, NextStepFinalOutput):
                         streamed_result._output_guardrails_task = asyncio.create_task(
                             cls._run_output_guardrails(
@@ -1170,6 +1158,12 @@ async def _start_streaming(
                                 await AgentRunner._save_result_to_session(
                                     session, [], turn_result.new_step_items
                                 )
+
+                        # Check for soft cancel after tool execution completes (before next turn)
+                        if streamed_result._cancel_mode == "after_turn":  # type: ignore[comparison-overlap]
+                            streamed_result.is_complete = True
+                            streamed_result._event_queue.put_nowait(QueueCompleteSentinel())
+                            break
                 except AgentsException as exc:
                     streamed_result.is_complete = True
                     streamed_result._event_queue.put_nowait(QueueCompleteSentinel())

tests/test_soft_cancel.py

@@ -4,7 +4,8 @@
 
 import pytest
 
-from agents import Agent, Runner, SQLiteSession
+from agents import Agent, OutputGuardrail, Runner, SQLiteSession
+from agents.guardrail import GuardrailFunctionOutput
 
 from .fake_model import FakeModel
 from .test_responses import get_function_tool, get_function_tool_call, get_text_message
@@ -485,3 +486,46 @@ async def test_soft_cancel_with_session_and_multiple_turns():
 
     # Cleanup
     await session.clear_session()
+
+
+@pytest.mark.asyncio
+async def test_soft_cancel_runs_output_guardrails_before_canceling():
+    """Verify output guardrails run even when cancellation happens after final output."""
+    model = FakeModel()
+
+    # Track if guardrail was called
+    guardrail_called = False
+
+    def output_guardrail_fn(context, agent, output):
+        nonlocal guardrail_called
+        guardrail_called = True
+        return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False)
+
+    agent = Agent(
+        name="Assistant",
+        model=model,
+        output_guardrails=[OutputGuardrail(guardrail_function=output_guardrail_fn)],
+    )
+
+    # Setup: agent produces final output
+    model.add_multiple_turn_outputs([[get_text_message("Final answer")]])
+
+    result = Runner.run_streamed(agent, input="What is the answer?")
+
+    # Cancel after seeing the message output event (indicates turn completed)
+    # but before consuming all events
+    async for event in result.stream_events():
+        if event.type == "run_item_stream_event" and event.name == "message_output_created":
+            # Cancel after turn completes - guardrails should still run
+            result.cancel(mode="after_turn")
+            # Don't break - continue consuming to let guardrails complete
+
+    # Guardrail should have been called
+    assert guardrail_called, "Output guardrail should run even when canceling after final output"
+
+    # Final output should be set
+    assert result.final_output is not None, "final_output should be set even when canceling"
+    assert result.final_output == "Final answer"
+
+    # Output guardrail results should be recorded
+    assert len(result.output_guardrail_results) == 1, "Output guardrail results should be recorded"