diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 4928b717..4a459d8e 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -14,7 +14,7 @@ import asyncio from collections.abc import AsyncIterator, Iterable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast, overload -from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming, calculate_model_cost, COST_TRACKER +from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming, calculate_model_cost, COST_TRACKER, cli_print_tool_output, _LIVE_STREAMING_PANELS, start_claude_thinking_if_applicable, finish_claude_thinking_display from cai.util import start_idle_timer, stop_idle_timer, start_active_timer, stop_active_timer from wasabi import color from cai.sdk.agents.run_to_jsonl import get_session_recorder @@ -501,6 +501,10 @@ class OpenAIChatCompletionsModel(Model): # Let the interrupt propagate up to end the current operation stop_active_timer() start_idle_timer() + + # Add visual feedback for interruption + print("\n[Non-streaming response interrupted - Cleanup completed]", file=sys.stderr) + raise if _debug.DONT_LOG_MODEL_DATA: @@ -817,58 +821,66 @@ class OpenAIChatCompletionsModel(Model): """ Yields a partial message as it is generated, as well as the usage information. """ - # IMPORTANT: Pre-process input to ensure it's in the correct format - # for streaming. This helps prevent errors during stream handling. - if not isinstance(input, str): - # Convert input items to messages and verify structure - try: - input_items = list(input) # Make sure it's a list - # Pre-verify the input messages to avoid errors during streaming - from cai.util import fix_message_list - - # Apply fix_message_list to the input items that are dictionaries - dict_items = [item for item in input_items if isinstance(item, dict)] - if dict_items: - fixed_dict_items = fix_message_list(dict_items) - - # Replace the original dict items with fixed ones while preserving non-dict items - new_input = [] - dict_index = 0 - for item in input_items: - if isinstance(item, dict): - if dict_index < len(fixed_dict_items): - new_input.append(fixed_dict_items[dict_index]) - dict_index += 1 - else: - new_input.append(item) - - # Update input with the fixed version - input = new_input - except Exception as e: - print(f"Warning: Error pre-processing input for streaming: {e}") - # Continue with original input even if pre-processing failed - - # Increment the interaction counter for CLI display - self.interaction_counter += 1 - self._intermediate_logs() - - # Stop idle timer and start active timer to track LLM processing time - stop_idle_timer() - start_active_timer() - - # --- Check if streaming should be shown in rich panel --- - should_show_rich_stream = os.getenv('CAI_STREAM', 'false').lower() == 'true' and not self.disable_rich_streaming - - # Create streaming context if needed + # Initialize streaming contexts as None streaming_context = None - if should_show_rich_stream: - streaming_context = create_agent_streaming_context( - agent_name=self.agent_name, - counter=self.interaction_counter, - model=str(self.model) - ) + thinking_context = None + stream_interrupted = False try: + # IMPORTANT: Pre-process input to ensure it's in the correct format + # for streaming. This helps prevent errors during stream handling. + if not isinstance(input, str): + # Convert input items to messages and verify structure + try: + input_items = list(input) # Make sure it's a list + # Pre-verify the input messages to avoid errors during streaming + from cai.util import fix_message_list + + # Apply fix_message_list to the input items that are dictionaries + dict_items = [item for item in input_items if isinstance(item, dict)] + if dict_items: + fixed_dict_items = fix_message_list(dict_items) + + # Replace the original dict items with fixed ones while preserving non-dict items + new_input = [] + dict_index = 0 + for item in input_items: + if isinstance(item, dict): + if dict_index < len(fixed_dict_items): + new_input.append(fixed_dict_items[dict_index]) + dict_index += 1 + else: + new_input.append(item) + + # Update input with the fixed version + input = new_input + except Exception as e: + print(f"Warning: Error pre-processing input for streaming: {e}") + # Continue with original input even if pre-processing failed + + # Increment the interaction counter for CLI display + self.interaction_counter += 1 + self._intermediate_logs() + + # Stop idle timer and start active timer to track LLM processing time + stop_idle_timer() + start_active_timer() + + # --- Check if streaming should be shown in rich panel --- + should_show_rich_stream = os.getenv('CAI_STREAM', 'false').lower() == 'true' and not self.disable_rich_streaming + + # Create streaming context if needed + if should_show_rich_stream: + try: + streaming_context = create_agent_streaming_context( + agent_name=self.agent_name, + counter=self.interaction_counter, + model=str(self.model) + ) + except Exception as e: + print(f"Warning: Could not create streaming context: {e}") + streaming_context = None + with generation_span( model=str(self.model), model_config=dataclasses.asdict(model_settings) @@ -1017,8 +1029,6 @@ class OpenAIChatCompletionsModel(Model): streamed_tool_calls = [] # Initialize Claude thinking display if applicable - thinking_context = None - from cai.util import start_claude_thinking_if_applicable if should_show_rich_stream: # Only show thinking in rich streaming mode thinking_context = start_claude_thinking_if_applicable( str(self.model), @@ -1044,6 +1054,10 @@ class OpenAIChatCompletionsModel(Model): try: async for chunk in stream: + # Check if we've been interrupted + if stream_interrupted: + break + if not state.started: state.started = True yield ResponseCreatedEvent( @@ -1467,22 +1481,18 @@ class OpenAIChatCompletionsModel(Model): # Set flag to suppress final output to avoid duplication self.suppress_final_output = True - except Exception as e: - # Ensure streaming context is cleaned up in case of errors - if streaming_context: - try: - finish_agent_streaming(streaming_context, None) - except Exception: - pass + except KeyboardInterrupt: + # Handle interruption during streaming + stream_interrupted = True + print("\n[Streaming interrupted by user]", file=sys.stderr) - # Ensure thinking context is cleaned up in case of errors - if thinking_context: - try: - from cai.util import finish_claude_thinking_display - finish_claude_thinking_display(thinking_context) - except Exception: - pass - raise e + # Let the exception propagate after cleanup + raise + + except Exception as e: + # Handle other exceptions during streaming + logger.error(f"Error during streaming: {e}") + raise # Special handling for Ollama - check if accumulated text contains a valid function call if is_ollama and ollama_full_content and len(state.function_calls) == 0: @@ -1884,27 +1894,72 @@ class OpenAIChatCompletionsModel(Model): stop_active_timer() start_idle_timer() + except KeyboardInterrupt: + # Handle keyboard interruption specifically + stream_interrupted = True + + # Make sure to clean up and re-raise + raise + except Exception as e: - # Ensure streaming context is cleaned up in case of errors + # Handle other exceptions + logger.error(f"Error in stream_response: {e}") + raise + + finally: + # Always clean up resources + # This block executes whether the try block succeeds, fails, or is interrupted + + # Clean up streaming context if streaming_context: try: - finish_agent_streaming(streaming_context, None) - except Exception: - pass + # Check if we need to force stop the streaming panel + if streaming_context.get("is_started", False) and streaming_context.get("live"): + streaming_context["live"].stop() - # Ensure thinking context is cleaned up in case of errors - if 'thinking_context' in locals() and thinking_context: + # Remove from active streaming contexts + if hasattr(create_agent_streaming_context, "_active_streaming"): + for key, value in list(create_agent_streaming_context._active_streaming.items()): + if value is streaming_context: + del create_agent_streaming_context._active_streaming[key] + break + except Exception as cleanup_error: + logger.debug(f"Error cleaning up streaming context: {cleanup_error}") + + # Clean up thinking context + if thinking_context: try: + # Force finish the thinking display from cai.util import finish_claude_thinking_display finish_claude_thinking_display(thinking_context) + except Exception as cleanup_error: + logger.debug(f"Error cleaning up thinking context: {cleanup_error}") + + # Clean up any live streaming panels + if hasattr(cli_print_tool_output, '_streaming_sessions'): + # Find any sessions related to this stream + for call_id in list(cli_print_tool_output._streaming_sessions.keys()): + if call_id in _LIVE_STREAMING_PANELS: + try: + live = _LIVE_STREAMING_PANELS[call_id] + live.stop() + del _LIVE_STREAMING_PANELS[call_id] + except Exception: + pass + + # Stop active timer and start idle timer + try: + stop_active_timer() + start_idle_timer() + except Exception: + pass + + # If the stream was interrupted, add a visual indicator + if stream_interrupted: + try: + print("\n[Stream interrupted - Cleanup completed]", file=sys.stderr) except Exception: pass - - # Stop active timer and start idle timer when streaming errors out - stop_active_timer() - start_idle_timer() - - raise e @overload async def _fetch_response( diff --git a/src/cai/util.py b/src/cai/util.py index cfb9cce7..ffa87576 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -33,6 +33,8 @@ import uuid from cai import is_pentestperf_available if is_pentestperf_available(): import pentestperf as ptt +import signal +import weakref # Global timing variables for tracking active and idle time _active_timer_start = None @@ -49,6 +51,83 @@ _LIVE_STREAMING_PANELS = {} # Global tracker for Claude thinking streaming panels _CLAUDE_THINKING_PANELS = {} +# Global flag to track if cleanup is in progress +_cleanup_in_progress = False +_cleanup_lock = threading.Lock() + +def cleanup_all_streaming_resources(): + """ + Clean up all active streaming resources. + This is called when the program is interrupted or exits. + """ + global _cleanup_in_progress + + with _cleanup_lock: + if _cleanup_in_progress: + return + _cleanup_in_progress = True + + try: + # Clean up all active Live streaming panels + for call_id, live in list(_LIVE_STREAMING_PANELS.items()): + try: + if hasattr(live, 'stop'): + live.stop() + except Exception: + pass + _LIVE_STREAMING_PANELS.clear() + + # Clean up all Claude thinking panels + for thinking_id, context in list(_CLAUDE_THINKING_PANELS.items()): + try: + if context and context.get("live") and context.get("is_started"): + context["live"].stop() + except Exception: + pass + _CLAUDE_THINKING_PANELS.clear() + + # Clean up active streaming contexts from create_agent_streaming_context + if hasattr(create_agent_streaming_context, "_active_streaming"): + for context_key, context in list(create_agent_streaming_context._active_streaming.items()): + try: + if context and context.get("live") and context.get("is_started"): + context["live"].stop() + except Exception: + pass + create_agent_streaming_context._active_streaming.clear() + + # Reset any streaming session states + if hasattr(cli_print_tool_output, '_streaming_sessions'): + cli_print_tool_output._streaming_sessions.clear() + + except Exception as e: + print(f"\nError during streaming cleanup: {e}", file=sys.stderr) + finally: + _cleanup_in_progress = False + +def signal_handler(signum, frame): + """ + Handle interrupt signals (CTRL+C) gracefully. + """ + # Stop any active timers + try: + stop_active_timer() + start_idle_timer() + except Exception: + pass + + # Clean up all streaming resources + cleanup_all_streaming_resources() + + # Re-raise KeyboardInterrupt to allow normal interrupt handling + raise KeyboardInterrupt() + +# Register signal handler for CTRL+C +signal.signal(signal.SIGINT, signal_handler) + +# Register cleanup at exit +atexit.register(cleanup_all_streaming_resources) + def start_active_timer(): """ Start measuring active time (when LLM is processing or tool is executing). @@ -1560,6 +1639,7 @@ def create_agent_streaming_context(agent_name, counter, model): "panel_width": panel_width, "is_started": False, # Track if we've started the display "error": None, # Track any errors + "context_key": context_key, # Store the key for cleanup } # Store the context for potential reuse @@ -1583,6 +1663,11 @@ def update_agent_streaming_content(context, text_delta, token_stats=None): """ if not context: return False + + # Check if cleanup is in progress to avoid updating a context being cleaned up + global _cleanup_in_progress + if _cleanup_in_progress: + return False try: # Only parse and add text if we have actual content to add @@ -1675,16 +1760,25 @@ def update_agent_streaming_content(context, text_delta, token_stats=None): context["is_started"] = True except Exception as e: context["error"] = str(e) + # Clean up the context if we can't start it + context_key = context.get("context_key") + if context_key and hasattr(create_agent_streaming_context, "_active_streaming"): + create_agent_streaming_context._active_streaming.pop(context_key, None) return False # Force an update with the new panel - context["live"].update(updated_panel) - context["panel"] = updated_panel - context["live"].refresh() + if context.get("is_started", False): + context["live"].update(updated_panel) + context["panel"] = updated_panel + context["live"].refresh() return True except Exception as e: # If there's an error, set it in the context context["error"] = str(e) + # Try to clean up the context + context_key = context.get("context_key") + if context_key and hasattr(create_agent_streaming_context, "_active_streaming"): + create_agent_streaming_context._active_streaming.pop(context_key, None) return False def finish_agent_streaming(context, final_stats=None): @@ -1698,12 +1792,15 @@ def finish_agent_streaming(context, final_stats=None): if not context: return False + # Check if cleanup is in progress + global _cleanup_in_progress + if _cleanup_in_progress: + return False + # Clean up tracking of this context - if hasattr(create_agent_streaming_context, "_active_streaming"): - for key, value in list(create_agent_streaming_context._active_streaming.items()): - if value is context: - del create_agent_streaming_context._active_streaming[key] - break + context_key = context.get("context_key") + if context_key and hasattr(create_agent_streaming_context, "_active_streaming"): + create_agent_streaming_context._active_streaming.pop(context_key, None) try: # Check if there's actual content to display - don't show empty panels @@ -1813,17 +1910,23 @@ def finish_agent_streaming(context, final_stats=None): expand=True ) - # Update one last time - context["live"].update(final_panel) - - # Ensure updates are displayed before stopping - time.sleep(0.1) - - # Stop the live display - try: - context["live"].stop() - except Exception as e: - context["error"] = str(e) + # Update one last time if display is started + if context.get("is_started", False): + try: + context["live"].update(final_panel) + + # Ensure updates are displayed before stopping + time.sleep(0.1) + + # Stop the live display + context["live"].stop() + except Exception as e: + context["error"] = str(e) + # Try to force stop if update failed + try: + context["live"].stop() + except Exception: + pass return True except Exception as e: @@ -1868,6 +1971,11 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut if tool_name == "execute_code" and not streaming: return + # Check if cleanup is in progress + global _cleanup_in_progress + if _cleanup_in_progress: + return + # Set up global tracker for streaming sessions if not hasattr(cli_print_tool_output, '_streaming_sessions'): cli_print_tool_output._streaming_sessions = {} @@ -2026,29 +2134,55 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # If we already have a live panel for this call_id, update it if call_id in _LIVE_STREAMING_PANELS: live = _LIVE_STREAMING_PANELS[call_id] - live.update(panel) + try: + live.update(panel) + except Exception as e: + # If update fails, try to clean up + try: + live.stop() + except Exception: + pass + del _LIVE_STREAMING_PANELS[call_id] # If this is the final update, stop the live panel after a short delay if execution_info and execution_info.get('is_final', False): # Give a moment for the final panel to be seen time.sleep(0.2) - live.stop() + try: + live.stop() + except Exception: + pass # Remove from the active panel dictionary - del _LIVE_STREAMING_PANELS[call_id] + if call_id in _LIVE_STREAMING_PANELS: + del _LIVE_STREAMING_PANELS[call_id] else: # Create a new live panel console = Console(theme=theme) live = Live(panel, console=console, refresh_per_second=4, auto_refresh=True) # Start and store the live panel - live.start() - _LIVE_STREAMING_PANELS[call_id] = live + try: + live.start() + _LIVE_STREAMING_PANELS[call_id] = live + except Exception as e: + # If we can't start the live panel, fall back to simple output + _print_simple_tool_output(tool_name, args, output, execution_info, token_info) # Return early for streaming updates return - except ImportError: + except (ImportError, Exception) as e: # Fall back to simple updates without Rich - pass + # If we had a live panel, try to clean it up + if call_id in _LIVE_STREAMING_PANELS: + try: + _LIVE_STREAMING_PANELS[call_id].stop() + except Exception: + pass + del _LIVE_STREAMING_PANELS[call_id] + + # Use simple output + _print_simple_tool_output(tool_name, args, output, execution_info, token_info) + return else: # For non-streaming outputs, check if we've already seen this command if command_key in cli_print_tool_output._displayed_commands: @@ -2169,7 +2303,7 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Display the panel console.print(panel) - except ImportError: + except (ImportError, Exception) as e: # Fall back to simple output format without rich _print_simple_tool_output(tool_name, args, output, execution_info, token_info)