fix: preserve streamed output guardrail tripwires in the run loop (#2758)

Author: seratch

src/agents/run_internal/run_loop.py

@@ -30,6 +30,7 @@
     InputGuardrailTripwireTriggered,
     MaxTurnsExceeded,
     ModelBehaviorError,
+    OutputGuardrailTripwireTriggered,
     RunErrorDetails,
     UserError,
 )
@@ -230,6 +231,17 @@
 ]
 
 
+def _should_attach_generic_agent_error(exc: Exception) -> bool:
+    return not isinstance(
+        exc,
+        (
+            ModelBehaviorError,
+            InputGuardrailTripwireTriggered,
+            OutputGuardrailTripwireTriggered,
+        ),
+    )
+
+
 async def _should_persist_stream_items(
     *,
     session: Session | None,
@@ -344,7 +356,12 @@ async def _run_output_guardrails_for_stream(
 
     try:
         return cast(list[Any], await streamed_result._output_guardrails_task)
+    except OutputGuardrailTripwireTriggered:
+        raise
+    except asyncio.CancelledError:
+        raise
     except Exception:
+        logger.error("Unexpected error in output guardrails", exc_info=True)
         return []
 
 
@@ -1014,7 +1031,7 @@ async def _save_stream_items_without_count(
                         streamed_result._event_queue.put_nowait(QueueCompleteSentinel())
                         break
             except Exception as e:
-                if current_span and not isinstance(e, ModelBehaviorError):
+                if current_span and _should_attach_generic_agent_error(e):
                     _error_tracing.attach_error_to_span(
                         current_span,
                         SpanError(
@@ -1037,7 +1054,7 @@ async def _save_stream_items_without_count(
         )
         raise
     except Exception as e:
-        if current_span and not isinstance(e, ModelBehaviorError):
+        if current_span and _should_attach_generic_agent_error(e):
             _error_tracing.attach_error_to_span(
                 current_span,
                 SpanError(

tests/test_agent_runner_streamed.py

@@ -1290,6 +1290,34 @@ def guardrail_function(
             pass
 
 
+@pytest.mark.asyncio
+async def test_output_guardrail_tripwire_raises_from_run_loop_task_before_stream_consumption():
+    def guardrail_function(
+        context: RunContextWrapper[Any], agent: Agent[Any], agent_output: Any
+    ) -> GuardrailFunctionOutput:
+        return GuardrailFunctionOutput(
+            output_info=None,
+            tripwire_triggered=True,
+        )
+
+    model = FakeModel(initial_output=[get_text_message("first_test")])
+
+    agent = Agent(
+        name="test",
+        output_guardrails=[OutputGuardrail(guardrail_function=guardrail_function)],
+        model=model,
+    )
+
+    result = Runner.run_streamed(agent, input="user_message")
+
+    assert result.run_loop_task is not None
+    with pytest.raises(OutputGuardrailTripwireTriggered):
+        await result.run_loop_task
+
+    assert result.final_output is None
+    assert result.is_complete is True
+
+
 @pytest.mark.asyncio
 async def test_run_input_guardrail_tripwire_triggered_causes_exception_streamed():
     def guardrail_function(