RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix async streaming broken panels

This commit is contained in:
luijait 2025-05-27 12:32:36 +02:00
parent d9abd4636c
commit df7a5dffaf
2 changed files with 295 additions and 106 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

213
src/cai/sdk/agents/models/openai_chatcompletions.py
View File

@ -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(

188
src/cai/util.py
View File

@ -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)