diff --git a/src/cai/cli.py b/src/cai/cli.py index 86960a619..abfe90a2a 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -1508,6 +1508,31 @@ async def process_streamed_response(agent, conversation_input): agent.model.add_to_message_history(tool_msg) return result + except OutputGuardrailTripwireTriggered as e: + # Handle guardrail exception specifically - MUST come before broad Exception handler + # Clean up streaming display before showing error + try: + from cai.util import cleanup_all_streaming_resources + cleanup_all_streaming_resources() + except Exception: + pass + + # Clean up the async generator + if stream_iterator is not None: + try: + await stream_iterator.aclose() + except Exception: + pass + + # Clean up the result object if it has cleanup methods + if result is not None and hasattr(result, '_cleanup_tasks'): + try: + result._cleanup_tasks() + except Exception: + pass + + # Re-raise to be caught by outer handler which shows user-friendly message + raise except (KeyboardInterrupt, asyncio.CancelledError) as e: # Handle interruption specifically @@ -1554,8 +1579,12 @@ async def process_streamed_response(agent, conversation_input): result._cleanup_tasks() except Exception: pass + + # Re-raise OutputGuardrailTripwireTriggered to be handled by outer handler + if isinstance(e, OutputGuardrailTripwireTriggered): + raise - # Log error for debugging + # Log error for debugging (non-guardrail exceptions) logger = logging.getLogger(__name__) logger.error(f"Error occurred during streaming: {str(e)}", exc_info=True) diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index 202ecbcb9..4bb70e8b0 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -673,6 +673,10 @@ async def _run_streamed_impl( try: output_guardrail_results = await streamed_result._output_guardrails_task + except OutputGuardrailTripwireTriggered as e: + # Store the guardrail exception immediately so it's checked during streaming + streamed_result._stored_exception = e + output_guardrail_results = [] except Exception: # Exceptions will be checked in the stream_events loop output_guardrail_results = []