diff --git a/pyproject.toml b/pyproject.toml index 650ff948..385cbd03 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,6 @@ license = {text = "Dual-licensed MIT and Proprietary"} authors = [{ name = "Alias Robotics", email = "research@aliasrobotics.com" }] dependencies = [ # logging and visualization - "flask>=2.0, <3", "folium>=0.15.0, <1", "matplotlib>=3.0, <4", "numpy>=1.21, <3", @@ -26,7 +25,7 @@ dependencies = [ "prompt_toolkit>=3.0.39", "dotenv>=0.9.9", "litellm>=1.63.7", - "mako>=1.3.9", + "mako>=1.3.8", "mcp; python_version >= '3.10'", "mkdocs>=1.6.0", "mkdocs-material>=9.6.0", diff --git a/src/cai/cli.py b/src/cai/cli.py index 37ed9dc7..7027da7a 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -355,7 +355,7 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= title="[bold]Session Summary[/bold]", title_align="left" ) - console.print(time_panel) + console.print(time_panel, end="") print_session_summary(console, metrics, logging_path) @@ -442,6 +442,13 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= } ) + # Fix message list structure BEFORE sending to the model to prevent errors + try: + from cai.util import fix_message_list + history_context = fix_message_list(history_context) + except Exception as e: + console.print(f"[yellow]Warning: Could not preprocess message history: {e}[/yellow]") + # Append the current user input as the last message in the list. conversation_input: list | str if history_context: @@ -483,14 +490,85 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= else: # Use non-streamed response response = asyncio.run(Runner.run(agent, conversation_input)) + + # Process the response items for item in response.new_items: + # Handle tool call output items (tool results) if isinstance(item, ToolCallOutputItem): + # First, ensure there's a corresponding assistant message with tool_calls + # before adding the tool response to prevent the OpenAI error + assistant_with_tool_call_exists = False + tool_call_id = item.raw_item["call_id"] + + for msg in message_history: + if (msg.get("role") == "assistant" and + msg.get("tool_calls") and + any(tc.get("id") == tool_call_id for tc in msg.get("tool_calls", []))): + assistant_with_tool_call_exists = True + break + + # If no matching assistant message exists, create one first + if not assistant_with_tool_call_exists: + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": tool_call_id, + "type": "function", + "function": { + "name": "unknown_function", + "arguments": "{}" + } + }] + } + add_to_message_history(tool_call_msg) + + # Now add the tool response tool_msg = { "role": "tool", - "tool_call_id": item.raw_item["call_id"], # Use consistent format with streaming + "tool_call_id": tool_call_id, "content": item.output, } add_to_message_history(tool_msg) + + # Make sure that assistant messages with tool calls are also added to message_history + # This is especially important for non-streaming mode + if hasattr(agent, 'model'): + # Access the _Converter directly from the OpenAIChatCompletionsModel implementation + from cai.sdk.agents.models.openai_chatcompletions import _Converter + + # Check if recent_tool_calls exists and process them + if hasattr(_Converter, 'recent_tool_calls'): + for call_id, call_info in _Converter.recent_tool_calls.items(): + # Only process new tool calls that haven't been added to message history yet + tool_call_found = False + for msg in message_history: + if (msg.get("role") == "assistant" and + msg.get("tool_calls") and + any(tc.get("id") == call_id for tc in msg.get("tool_calls", []))): + tool_call_found = True + break + + if not tool_call_found: + # Add the assistant message with the tool call + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": call_id, + "type": "function", + "function": { + "name": call_info.get('name', ''), + "arguments": call_info.get('arguments', '{}') + } + }] + } + add_to_message_history(tool_call_msg) + + # Final validation to ensure message history follows OpenAI's requirements + # Ensure every tool message has a preceding assistant message with matching tool_call_id + from cai.util import fix_message_list + message_history[:] = fix_message_list(message_history) turn_count += 1 # Stop measuring active time and start measuring idle time again diff --git a/src/cai/repl/commands/agent.py b/src/cai/repl/commands/agent.py index dec5b0ee..6be0d2b8 100644 --- a/src/cai/repl/commands/agent.py +++ b/src/cai/repl/commands/agent.py @@ -238,7 +238,7 @@ class AgentCommand(Command): os.environ["CAI_AGENT_TYPE"] = selected_agent_key console.print( - f"[green]Switched to agent: {agent_name}[/green]") + f"[green]Switched to agent: {agent_name}[/green]", end="") visualize_agent_graph(agent) return True diff --git a/src/cai/repl/commands/model.py b/src/cai/repl/commands/model.py index 04d2ad52..07cb0480 100644 --- a/src/cai/repl/commands/model.py +++ b/src/cai/repl/commands/model.py @@ -421,7 +421,7 @@ class ModelCommand(Command): change_message, border_style="green", title="Model Changed" - ) + ), end="" ) return True diff --git a/src/cai/repl/ui/banner.py b/src/cai/repl/ui/banner.py index b1acdb2e..1c52de87 100644 --- a/src/cai/repl/ui/banner.py +++ b/src/cai/repl/ui/banner.py @@ -177,7 +177,7 @@ def display_banner(console: Console): [white] Bug bounty-ready AI[/white] """ - console.print(banner) + console.print(banner, end="") # # Create a table showcasing CAI framework capabilities # # @@ -361,4 +361,4 @@ def display_quick_guide(console: Console): border_style="blue", padding=(1, 2), title_align="center" - )) + ), end="") diff --git a/src/cai/repl/ui/prompt.py b/src/cai/repl/ui/prompt.py index 35de00e0..60358333 100644 --- a/src/cai/repl/ui/prompt.py +++ b/src/cai/repl/ui/prompt.py @@ -96,7 +96,7 @@ def get_user_input( # Get user input with all features return prompt( - [('class:prompt', '\nCAI> ')], + [('class:prompt', 'CAI> ')], completer=command_completer, style=create_prompt_style(), history=FileHistory(str(history_file)), diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index a68057c2..afd1ffdf 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -336,6 +336,15 @@ class OpenAIChatCompletionsModel(Model): # Log the user message if item.get("content"): self.logger.log_user_message(item.get("content")) + + # IMPORTANT: Ensure the message list has valid tool call/result pairs + # This needs to happen before the API call to prevent errors + try: + from cai.util import fix_message_list + converted_messages = fix_message_list(converted_messages) + except Exception as e: + logger.warning(f"Failed to fix message list: {e}") + # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) @@ -411,6 +420,10 @@ class OpenAIChatCompletionsModel(Model): # Only display the agent message if we haven't already shown the tool output if should_display_message: + # Ensure we're in non-streaming mode for proper markdown parsing + previous_stream_setting = os.environ.get('CAI_STREAM', 'false') + os.environ['CAI_STREAM'] = 'false' # Force non-streaming mode for markdown parsing + # Print the agent message for CLI display cli_print_agent_messages( agent_name=getattr(self, 'agent_name', 'Agent'), @@ -433,7 +446,11 @@ class OpenAIChatCompletionsModel(Model): interaction_cost=None, total_cost=None, tool_output=None, # Don't pass tool output here, we're using direct display + suppress_empty=True # Suppress empty panels ) + + # Restore previous streaming setting + os.environ['CAI_STREAM'] = previous_stream_setting # --- Add assistant tool call to message_history if present --- # If the response contains tool_calls, add them to message_history as assistant messages @@ -457,6 +474,22 @@ class OpenAIChatCompletionsModel(Model): } add_to_message_history(tool_call_msg) + + # Save the tool call details for later matching with output + # This is important for non-streaming mode to track tool calls properly + if not hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls = {} + + # Store the tool call by ID for later reference + import time + _Converter.recent_tool_calls[tool_call.id] = { + 'name': tool_call.function.name, + 'arguments': tool_call.function.arguments, + 'start_time': time.time(), + 'execution_info': { + 'start_time': time.time() + } + } # Log the assistant tool call message tool_calls_list = [] @@ -583,7 +616,7 @@ class OpenAIChatCompletionsModel(Model): stop_idle_timer() start_active_timer() - # Check if streaming should be shown in rich panel + # --- 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 @@ -922,6 +955,55 @@ class OpenAIChatCompletionsModel(Model): if tool_call_msg not in streamed_tool_calls: streamed_tool_calls.append(tool_call_msg) add_to_message_history(tool_call_msg) + + # NEW: Display tool call immediately when detected in streaming mode + # But only if it has complete arguments and name + if (state.function_calls[tc_index].name and + state.function_calls[tc_index].arguments and + state.function_calls[tc_index].call_id): + # First, finish any existing streaming context if it exists + if streaming_context: + try: + finish_agent_streaming(streaming_context, None) + streaming_context = None + except Exception: + pass + + # Create a message-like object for displaying the function call + tool_msg = type('ToolCallStreamDisplay', (), { + 'content': None, + 'tool_calls': [ + type('ToolCallDetail', (), { + 'function': type('FunctionDetail', (), { + 'name': state.function_calls[tc_index].name, + 'arguments': state.function_calls[tc_index].arguments + }), + 'id': state.function_calls[tc_index].call_id, + 'type': 'function' + }) + ] + }) + + # Display the tool call during streaming + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), + message=tool_msg, + counter=getattr(self, 'interaction_counter', 0), + model=str(self.model), + debug=False, + interaction_input_tokens=estimated_input_tokens, + interaction_output_tokens=estimated_output_tokens, + interaction_reasoning_tokens=0, # Not available during streaming yet + total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens, + total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens, + total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), + interaction_cost=None, + total_cost=None, + tool_output=None, # Will be shown once tool is executed + suppress_empty=True # Prevent empty panels + ) + # 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 @@ -983,8 +1065,15 @@ class OpenAIChatCompletionsModel(Model): ) # Display the tool call in CLI - from cai.util import cli_print_agent_messages try: + # First, finish any existing streaming context if it exists + if streaming_context: + try: + finish_agent_streaming(streaming_context, None) + streaming_context = None + except Exception: + pass + # Create a message-like object to display the function call tool_msg = type('ToolCallWrapper', (), { 'content': None, @@ -1015,8 +1104,12 @@ class OpenAIChatCompletionsModel(Model): total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), interaction_cost=None, total_cost=None, - tool_output=None # Will be shown once the tool is executed + tool_output=None, # Will be shown once the tool is executed + suppress_empty=True # Suppress empty panels during streaming ) + + # Set flag to suppress final output to avoid duplication + self.suppress_final_output = True except Exception as e: logger.error(f"Error displaying tool call in CLI: {e}") @@ -1220,6 +1313,7 @@ class OpenAIChatCompletionsModel(Model): direct_stats["total_cost"] = float(total_cost) # Use the direct copy with guaranteed float costs finish_agent_streaming(streaming_context, direct_stats) + streaming_context = None # Removed extra newline after streaming completes to avoid blank lines pass @@ -1243,8 +1337,10 @@ class OpenAIChatCompletionsModel(Model): for tool_call in tool_call_msg.get("tool_calls", []): tool_calls_list.append(tool_call) self.logger.log_assistant_message(None, tool_calls_list) - # If there was only text output, add that as an assistant message - if (not streamed_tool_calls) and state.text_content_index_and_output and state.text_content_index_and_output[1].text: + + # If we've already shown the tool calls directly during streaming, + # don't log the text message to avoid duplication + elif (not self.suppress_final_output) and state.text_content_index_and_output and state.text_content_index_and_output[1].text: asst_msg = { "role": "assistant", "content": state.text_content_index_and_output[1].text @@ -1253,6 +1349,9 @@ class OpenAIChatCompletionsModel(Model): # Log the assistant message self.logger.log_assistant_message(state.text_content_index_and_output[1].text) + # Reset the suppress flag for future requests + self.suppress_final_output = False + # Log the complete response self.logger.rec_training_data( { @@ -1341,10 +1440,17 @@ class OpenAIChatCompletionsModel(Model): if tracing.include_data(): span.span_data.input = converted_messages - # Ensure message list has correct structure regardless of error condition + # IMPORTANT: Always sanitize the message list to prevent tool call errors + # This is critical to fix common errors with tool/assistant sequences try: from cai.util import fix_message_list + prev_length = len(converted_messages) converted_messages = fix_message_list(converted_messages) + new_length = len(converted_messages) + + # Log if the message list was changed significantly + if new_length != prev_length: + logger.debug(f"Message list was fixed: {prev_length} -> {new_length} messages") except Exception as e: logger.warning(f"Failed to fix message list: {e}") @@ -1531,15 +1637,49 @@ class OpenAIChatCompletionsModel(Model): elif ("An assistant message with 'tool_calls'" in str(e) or "`tool_use` blocks must be followed by a user message with `tool_result`" in str(e) or # noqa: E501 # pylint: disable=C0301 - "An assistant message with 'tool_calls' must be followed by tool messages" in str(e)): # Añadir esta condición + "An assistant message with 'tool_calls' must be followed by tool messages" in str(e) or + "messages with role 'tool' must be a response to a preceeding message with 'tool_calls'" in str(e)): print(f"Error: {str(e)}") + + # Use the pretty message history printer instead of the simple loop + try: + from cai.util import print_message_history + print("\nCurrent message sequence causing the error:") + print_message_history(kwargs["messages"], title="Message Sequence Error") + except ImportError: + # Fall back to simple printing if the function isn't available + print("\nCurrent message sequence causing the error:") + for i, msg in enumerate(kwargs["messages"]): + role = msg.get("role", "unknown") + content_type = ( + "text" if isinstance(msg.get("content"), str) else + "list" if isinstance(msg.get("content"), list) else + "None" if msg.get("content") is None else + type(msg.get("content")).__name__ + ) + tool_calls = "with tool_calls" if msg.get("tool_calls") else "" + tool_call_id = f", tool_call_id: {msg.get('tool_call_id')}" if msg.get("tool_call_id") else "" + + print(f" [{i}] {role}{tool_call_id} (content: {content_type}) {tool_calls}") + # NOTE: EDGE CASE: Report Agent CTRL C error # # This fix CTRL-C error when message list is incomplete # When a tool is not finished but the LLM generates a tool call try: from cai.util import fix_message_list - kwargs["messages"] = fix_message_list(kwargs["messages"]) + print("Attempting to fix message sequence...") + fixed_messages = fix_message_list(kwargs["messages"]) + + # Show the fixed messages if they're different + if fixed_messages != kwargs["messages"]: + try: + from cai.util import print_message_history + print_message_history(fixed_messages, title="Fixed Message Sequence") + except ImportError: + print("Messages fixed successfully.") + + kwargs["messages"] = fixed_messages except Exception as fix_error: print(f"Failed to fix message sequence: {fix_error}") @@ -2042,7 +2182,11 @@ class _Converter: if current_assistant_msg is not None: # The API doesn't support empty arrays for tool_calls if not current_assistant_msg.get("tool_calls"): - del current_assistant_msg["tool_calls"] + # Ensure content is not None if tool_calls are absent and content is also None + # Some models like Anthropic require some content, even if it's just a placeholder. + if current_assistant_msg.get("content") is None: + current_assistant_msg["content"] = "(No text content in this assistant message)" # Or just an empty string if preferred + current_assistant_msg.pop("tool_calls", None) # Use pop with default to avoid KeyError result.append(current_assistant_msg) current_assistant_msg = None @@ -2079,16 +2223,28 @@ class _Converter: flush_assistant_message() tool_calls_param: list[ChatCompletionMessageToolCallParam] = [] for tc in item["tool_calls"]: + function_details = tc.get("function", {}) + arguments = function_details.get("arguments") + # Ensure arguments is a valid JSON string, defaulting to "{}" if empty or None + if arguments is None or (isinstance(arguments, str) and arguments.strip() == ""): + arguments = "{}" + elif isinstance(arguments, dict): + # Ensure it's a string if it's a dict (should already be string per schema) + arguments = json.dumps(arguments) + tool_calls_param.append( ChatCompletionMessageToolCallParam( id=tc.get("id", ""), type=tc.get("type", "function"), - function=tc.get("function", {}), + function={ + "name": function_details.get("name", "unknown_function"), + "arguments": arguments, # Use sanitized arguments + }, ) ) msg_asst: ChatCompletionAssistantMessageParam = { "role": "assistant", - "content": item.get("content"), + "content": item.get("content"), # Content can be None here "tool_calls": tool_calls_param, } result.append(msg_asst) @@ -2225,12 +2381,19 @@ class _Converter: } } + arguments = func_call.get("arguments") # func_call is a dict here + # Ensure arguments is a valid JSON string, defaulting to "{}" if empty or None + if arguments is None or (isinstance(arguments, str) and arguments.strip() == ""): + arguments = "{}" + elif isinstance(arguments, dict): + arguments = json.dumps(arguments) + new_tool_call = ChatCompletionMessageToolCallParam( id=func_call["call_id"], type="function", function={ "name": func_call["name"], - "arguments": func_call["arguments"], + "arguments": arguments, # Use sanitized arguments }, ) tool_calls.append(new_tool_call) @@ -2244,22 +2407,22 @@ class _Converter: # Update execution timing if we have the start time if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: - tool_call = cls.recent_tool_calls[call_id] - if 'start_time' in tool_call: + tool_call_details = cls.recent_tool_calls[call_id] # Renamed for clarity + if 'start_time' in tool_call_details: end_time = time.time() - tool_execution_time = end_time - tool_call['start_time'] + tool_execution_time = end_time - tool_call_details['start_time'] # Update the execution info - if 'execution_info' in tool_call: - tool_call['execution_info']['end_time'] = end_time - tool_call['execution_info']['tool_time'] = tool_execution_time + if 'execution_info' in tool_call_details: + tool_call_details['execution_info']['end_time'] = end_time + tool_call_details['execution_info']['tool_time'] = tool_execution_time # If this is the first tool being executed, record the total time from conversation start if not hasattr(cls, 'conversation_start_time'): - cls.conversation_start_time = tool_call['start_time'] + cls.conversation_start_time = tool_call_details['start_time'] - total_time = end_time - getattr(cls, 'conversation_start_time', tool_call['start_time']) - tool_call['execution_info']['total_time'] = total_time + total_time = end_time - getattr(cls, 'conversation_start_time', tool_call_details['start_time']) + tool_call_details['execution_info']['total_time'] = total_time # Store the output so it can be accessed later if not hasattr(cls, 'tool_outputs'): @@ -2270,67 +2433,73 @@ class _Converter: # Display the tool output immediately with the matched tool call from cai.util import cli_print_tool_output - # Check if we're in streaming mode - don't show tool output panel in streaming mode + # Look up the original tool call to get the name and arguments + tool_name = "Unknown Tool" + tool_args = {} + execution_info = {} + + if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: + tool_call_details = cls.recent_tool_calls[call_id] # Renamed for clarity + tool_name = tool_call_details.get('name', 'Unknown Tool') + tool_args = tool_call_details.get('arguments', {}) + execution_info = tool_call_details.get('execution_info', {}) + + # Get token counts from the OpenAIChatCompletionsModel if available + model_instance = None + for frame in inspect.stack(): + if 'self' in frame.frame.f_locals: + self_obj = frame.frame.f_locals['self'] + if isinstance(self_obj, OpenAIChatCompletionsModel): + model_instance = self_obj + break + + # Always create a token_info dictionary, even if some values are zero + token_info = { + 'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0), + 'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0), + 'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0), + 'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0), + 'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0), + 'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0), + 'model': str(getattr(model_instance, 'model', '')), + } + + # Calculate costs using standard cost model + if model_instance and hasattr(model_instance, 'model'): + from cai.util import calculate_model_cost + model_name_str = str(model_instance.model) # Ensure model name is string + token_info['interaction_cost'] = calculate_model_cost( + model_name_str, + token_info['interaction_input_tokens'], + token_info['interaction_output_tokens'] + ) + token_info['total_cost'] = calculate_model_cost( + model_name_str, + token_info['total_input_tokens'], + token_info['total_output_tokens'] + ) + + # Check if we're in streaming mode is_streaming_enabled = os.environ.get('CAI_STREAM', 'false').lower() == 'true' - if is_streaming_enabled: - # Don't display tool output in streaming mode - it will be handled elsewhere - pass # Just skip the display, but preserve the tool output - else: - # For non-streaming mode, maintain the original behavior - # Look up the original tool call to get the name and arguments - if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: - tool_call = cls.recent_tool_calls[call_id] - tool_name = tool_call.get('name', 'Unknown Tool') - tool_args = tool_call.get('arguments', {}) - execution_info = tool_call.get('execution_info', {}) - - # Get token counts from the OpenAIChatCompletionsModel if available - model_instance = None - for frame in inspect.stack(): - if 'self' in frame.frame.f_locals: - self_obj = frame.frame.f_locals['self'] - if isinstance(self_obj, OpenAIChatCompletionsModel): - model_instance = self_obj - break - - # Always create a token_info dictionary, even if some values are zero - token_info = { - 'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0), - 'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0), - 'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0), - 'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0), - 'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0), - 'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0), - 'model': str(getattr(model_instance, 'model', '')), - } - - # Calculate costs using standard cost model - if model_instance and hasattr(model_instance, 'model'): - from cai.util import calculate_model_cost - model_name = str(model_instance.model) - token_info['interaction_cost'] = calculate_model_cost( - model_name, - token_info['interaction_input_tokens'], - token_info['interaction_output_tokens'] - ) - token_info['total_cost'] = calculate_model_cost( - model_name, - token_info['total_input_tokens'], - token_info['total_output_tokens'] - ) - - # Use the cli_print_tool_output function with actual token values - cli_print_tool_output( - tool_name=tool_name, - args=tool_args, - output=output_content, - call_id=call_id, # Keep call_id for non-streaming mode - execution_info=execution_info, - token_info=token_info - ) + + # Always display tool output regardless of streaming mode + cli_print_tool_output( + tool_name=tool_name, + args=tool_args, + output=output_content, + call_id=call_id, + execution_info=execution_info, + token_info=token_info + ) # Continue with normal processing flush_assistant_message() + + # REMOVED THE BLOCK THAT CREATED A SYNTHETIC ASSISTANT MESSAGE HERE + # The responsibility for ensuring a preceding assistant message + # is now fully deferred to fix_message_list, called later. + + # Now add the tool message msg: ChatCompletionToolMessageParam = { "role": "tool", "tool_call_id": func_output["call_id"], diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 2b4327da..2650a827 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -13,7 +13,7 @@ import uuid import sys import shlex from wasabi import color # pylint: disable=import-error -from cai.util import format_time, start_active_timer, stop_active_timer, start_idle_timer, stop_idle_timer +from cai.util import format_time, start_active_timer, stop_active_timer, start_idle_timer, stop_idle_timer, cli_print_tool_output # Instead of direct import @@ -434,236 +434,80 @@ def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None): return error_msg -def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None): +def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None, custom_args=None): """Runs command locally in the specified workspace_dir.""" - # If streaming is enabled and we have a call_id - if stream and call_id: - return _run_local_streamed(command, call_id, timeout, tool_name, workspace_dir) - # Make sure we're in active time mode for tool execution stop_idle_timer() start_active_timer() + process_start_time = time.time() # Initialize with current time try: target_dir = workspace_dir or _get_workspace_dir() original_cmd_for_msg = command # For logging context_msg = f"(local:{target_dir})" - try: - result = subprocess.run( + + # If streaming is enabled and we have a call_id + if stream: + # Import the streaming utilities from util + from cai.util import start_tool_streaming, update_tool_streaming, finish_tool_streaming + + # Parse command into parts for display + parts = command.strip().split(' ', 1) + cmd_var = parts[0] if parts else "" + args_param_val = parts[1] if len(parts) > 1 else "" # Renamed to avoid conflict with tool_args dict key + + # For generic Linux commands, standardize the tool_name format + if not tool_name: + tool_name = f"{cmd_var}_command" if cmd_var else "command" + + # Create args dictionary with non-empty values only + tool_args = {} + if cmd_var: + tool_args["command"] = cmd_var + if args_param_val and args_param_val.strip(): + tool_args["args"] = args_param_val + + # Add more context for the command + tool_args["workspace"] = os.path.basename(target_dir) + tool_args["full_command"] = command + + # If custom args were provided, merge them with the default args + if custom_args is not None: + if isinstance(custom_args, dict): + # Merge the dictionaries, with custom args taking precedence + for key, value in custom_args.items(): + tool_args[key] = value + + # For generic commands, ensure we have a unique call_id + if not call_id: + call_id = f"cmd_{cmd_var}_{str(uuid.uuid4())[:8]}" + + # Initialize/use the call_id for this streaming session + call_id = start_tool_streaming(tool_name, tool_args, call_id) + + # Start the process + process = subprocess.Popen( command, shell=True, # nosec B602 - capture_output=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, text=True, - check=False, - timeout=timeout, - cwd=target_dir - ) - output = result.stdout if result.stdout else result.stderr - if stdout: - print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 - - # Skip passing output to cli_print_tool_output when CAI_STREAM=true - # This prevents duplicate output in streaming mode - is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' - if not is_streaming_enabled: - # Optional: Add cli_print_tool_output call here if needed for non-streaming - pass - - return output.strip() - except subprocess.TimeoutExpired as e: - error_output = e.stdout if e.stdout else str(e) - if stdout: - print("\033[32m" + error_output + "\033[0m") - return error_output - except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing local command: {e}" - print(color(error_msg, fg="red")) - return error_msg - finally: - # Always switch back to idle mode when function completes - stop_active_timer() - start_idle_timer() - - -def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace_dir=None): - """Run a local command with streaming output to the Tool output panel.""" - # Make sure we're in active time mode for tool execution - stop_idle_timer() - start_active_timer() - - target_dir = workspace_dir or _get_workspace_dir() - try: - # Try to import Rich for nice display - try: - from rich.console import Console - from rich.live import Live - from rich.panel import Panel - from rich.text import Text - from rich.box import ROUNDED - console = Console() - rich_available = True - except ImportError: - rich_available = False - from cai.util import cli_print_tool_output - - output_buffer = [] - start_time = time.time() - # Start the process - process = subprocess.Popen( - command, - shell=True, # nosec B602 - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - bufsize=1, - cwd=target_dir # Set CWD for local process - ) - - # If tool_name is not provided, derive it from the command - if tool_name is None: - # Just use the first command as the tool name - tool_name = command.strip().split()[0] + "_command" - - # Create panel content for Rich display - if rich_available: - # Parse command into command and args - parts = command.strip().split(' ', 1) - cmd = parts[0] if parts else "" - args = parts[1] if len(parts) > 1 else "" - - # Format clean arguments, following the same rules as cli_print_tool_output - arg_parts = [] - if cmd: - arg_parts.append(f"command={cmd}") - if args and args.strip(): # Only add args if non-empty - arg_parts.append(f"args={args}") - args_str = ", ".join(arg_parts) - - header = Text() - header.append(tool_name, style="#00BCD4") - header.append("(", style="yellow") - header.append(args_str, style="yellow") - header.append(")", style="yellow") - tool_time = 0 - - total_time = 0 - if START_TIME is not None: - total_time = time.time() - START_TIME - timing_info = [] - if total_time: - timing_info.append(f"Total: {format_time(total_time)}") - if tool_time: - timing_info.append(f"Tool: {format_time(tool_time)}") - if timing_info: - header.append(f" [{' | '.join(timing_info)}]", style="cyan") - - content = Text() - - panel = Panel( - Text.assemble(header, "\n\n", content), - title="[bold green]Tool Execution[/bold green]", - subtitle="[bold green]Live Output[/bold green]", - border_style="green", - padding=(1, 2), - box=ROUNDED + bufsize=1, + cwd=target_dir ) - # Start Live display - with Live(panel, console=console, refresh_per_second=4) as live: - # Stream stdout in real-time - start_time = time.time() - for line in iter(process.stdout.readline, ''): - if not line: - break - - # Add to output collection - output_buffer.append(line) - - # Update content with new line - content.append(line, style="bright_white") - - # Update tool_time and header with new timing info - tool_time = time.time() - start_time - total_time = 0 - if START_TIME is not None: - total_time = time.time() - START_TIME - # Remove any previous timing info from header (rebuild header) - timing_info = [] - if total_time: - timing_info.append(f"Total: {format_time(total_time)}") - if tool_time: - timing_info.append(f"Tool: {format_time(tool_time)}") - # Rebuild header to update timing - header = Text() - header.append(tool_name, style="#00BCD4") - header.append("(", style="yellow") - header.append(args_str, style="yellow") - header.append(")", style="yellow") - if timing_info: - header.append(f" [{' | '.join(timing_info)}]", style="cyan") - - panel = Panel( - Text.assemble(header, "\n\n", content), - title="[bold green]Tool Execution[/bold green]", - subtitle="[bold green]Live Output[/bold green]", - border_style="green", - padding=(1, 2), - box=ROUNDED - ) - live.update(panel) - # Check if process is done - process.stdout.close() - return_code = process.wait(timeout=timeout) - - # Get any stderr output - stderr_data = process.stderr.read() - if stderr_data: - content.append("\nERROR OUTPUT:\n", style="red") - content.append(stderr_data, style="red") - output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) - panel = Panel( - Text.assemble(header, "\n\n", content), - title="[bold green]Tool Execution[/bold green]", - subtitle="[bold green]Live Output[/bold green]", - border_style="green", - padding=(1, 2), - box=ROUNDED - ) - live.update(panel) - - # Add completion message - panel = Panel( - Text.assemble(header, "\n\n", content), - title="[bold green]Tool Execution[/bold green]", - border_style="green", - padding=(1, 2), - box=ROUNDED - ) - live.update(panel) - - # Wait a moment for the panel to be displayed properly - time.sleep(0.5) - else: - # Fallback to simpler streaming with cli_print_tool_output - # Parse command into command and args (same as rich mode) - parts = command.strip().split(' ', 1) - cmd = parts[0] if parts else "" - args = parts[1] if len(parts) > 1 else "" - - # Create a dictionary with only non-empty values (following the same rules) - tool_args = {} - if cmd: - tool_args["command"] = cmd - if args and args.strip(): - tool_args["args"] = args - # Note: Omitted empty values and async_mode=False as it's default - - # Initial notification - just once - cli_print_tool_output(tool_name, tool_args, "Command started...", call_id=call_id) - - # Buffer for collecting output + # Begin collecting output + output_buffer = [] buffer_size = 0 - update_interval = 10 # lines + update_interval = 10 # lines - default for most tools + + # Use a smaller interval for generic_linux_command for better responsiveness + if tool_name == "generic_linux_command": + update_interval = 3 # Update more frequently for terminal commands + + # Add refresh rate info to tool_args for cli_print_tool_output + if "refresh_rate" not in tool_args: + tool_args["refresh_rate"] = 2 # Stream stdout in real-time for line in iter(process.stdout.readline, ''): @@ -674,52 +518,163 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace output_buffer.append(line) buffer_size += 1 - # Only update the output periodically to reduce panel refresh rate + # Only update periodically to reduce UI refreshes if buffer_size >= update_interval: current_output = ''.join(output_buffer) - cli_print_tool_output(tool_name, tool_args, current_output, call_id=call_id) + update_tool_streaming(tool_name, tool_args, current_output, call_id) buffer_size = 0 - # Check if process is done + # Finish process process.stdout.close() return_code = process.wait(timeout=timeout) + process_execution_time = time.time() - process_start_time # Get any stderr output stderr_data = process.stderr.read() if stderr_data: output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) - # Final output update - always show the final result + # Final output update final_output = ''.join(output_buffer) if return_code != 0: final_output += f"\nCommand exited with code {return_code}" - cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id) - - # Return the full output - return ''.join(output_buffer) - - except subprocess.TimeoutExpired: - error_msg = f"Command timed out after {timeout} seconds" - output_buffer.append("\n" + error_msg) - final_output = ''.join(output_buffer) - - # Update tool output panel with timeout message - if not rich_available: - tool_args = {"command": command} - cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id) - - return final_output + # Calculate execution info with environment details + execution_info = { + "status": "completed" if return_code == 0 else "error", + "return_code": return_code, + "environment": "Local", + "host": os.path.basename(target_dir), + "tool_time": process_execution_time + } + + # Complete the streaming session with final output + finish_tool_streaming(tool_name, tool_args, final_output, call_id, execution_info) + + return final_output + else: + # Standard non-streaming execution + result = subprocess.run( + command, + shell=True, # nosec B602 + capture_output=True, + text=True, + check=False, + timeout=timeout, + cwd=target_dir + ) + output = result.stdout if result.stdout else result.stderr + + # If this is the same command that was recently streamed, don't display it again + # We'll create a similar tool_args structure to what streaming would use + parts = command.strip().split(' ', 1) + cmd_var = parts[0] if parts else "" + args_param_val = parts[1] if len(parts) > 1 else "" + + # Generate a standard tool name for consistency with streaming + standard_tool_name = tool_name or (f"{cmd_var}_command" if cmd_var else "command") + + # Calculate a consistent command key - must match the format used in cli_print_tool_output + command_key = f"{standard_tool_name}:{args_param_val}" + + if hasattr(cli_print_tool_output, '_displayed_commands') and command_key in cli_print_tool_output._displayed_commands: + # Skip stdout display if already shown through streaming + return output.strip() + + if stdout: + # Create a tool_args dictionary for non-streaming display + # that matches the format used in streaming + tool_display_args = { + "command": cmd_var, + "args": args_param_val, + "full_command": command, + "workspace": os.path.basename(target_dir) + } + + # If custom args were provided, merge them with the default args + if custom_args is not None and isinstance(custom_args, dict): + for key, value in custom_args.items(): + tool_display_args[key] = value + + # Calculate execution info + tool_execution_time = time.time() - process_start_time + exec_info = { + "status": "completed" if result.returncode == 0 else "error", + "return_code": result.returncode, + "environment": "Local", + "host": os.path.basename(target_dir), + "tool_time": tool_execution_time + } + + # Display the command output with rich formatting + cli_print_tool_output( + tool_name=standard_tool_name, + args=tool_display_args, + output=output, + execution_info=exec_info + ) + + return output.strip() + except subprocess.TimeoutExpired as e: + error_output = e.stdout if hasattr(e, 'stdout') and e.stdout else str(e) + error_msg = f"Command timed out after {timeout} seconds\n{error_output}" + # If we're streaming, show the timeout in the tool output panel + if stream and call_id: + from cai.util import finish_tool_streaming + # Parse the command the same way we did for streaming + parts = command.strip().split(' ', 1) + cmd_var = parts[0] if parts else "" + args_var = parts[1] if len(parts) > 1 else "" + + # Ensure tool_args has complete information + tool_args = { + "command": cmd_var, + "args": args_var if args_var.strip() else "", + "full_command": command, + "environment": "Local", + "workspace": os.path.basename(target_dir) + } + execution_info = { + "status": "timeout", + "error": str(e), + "environment": "Local", + "host": os.path.basename(target_dir) + } + finish_tool_streaming(tool_name or f"{cmd_var}_command", tool_args, error_msg, call_id, execution_info) + + if stdout: + print("\033[32m" + error_msg + "\033[0m") + + return error_msg except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing command: {str(e)}" + error_msg = f"Error executing local command: {e}" + + # If we're streaming, show the error in the tool output panel + if stream and call_id: + from cai.util import finish_tool_streaming + # Parse the command the same way we did for streaming + parts = command.strip().split(' ', 1) + cmd_var = parts[0] if parts else "" + args_var = parts[1] if len(parts) > 1 else "" + + # Ensure tool_args has complete information + tool_args = { + "command": cmd_var, + "args": args_var if args_var.strip() else "", + "full_command": command, + "environment": "Local", + "workspace": os.path.basename(target_dir) + } + execution_info = { + "status": "error", + "error": str(e), + "environment": "Local", + "host": os.path.basename(target_dir) + } + finish_tool_streaming(tool_name or f"{cmd_var}_command", tool_args, error_msg, call_id, execution_info) + print(color(error_msg, fg="red")) - - # Update tool output panel with error message if simple streaming - if not rich_available: - tool_args = {"command": command} - cli_print_tool_output(tool_name, tool_args, error_msg, call_id=call_id) - return error_msg finally: # Always switch back to idle mode when function completes @@ -729,7 +684,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arguments # noqa: E501 async_mode=False, session_id=None, - timeout=100, stream=False, call_id=None, tool_name=None): + timeout=100, stream=False, call_id=None, tool_name=None, args=None): """ Run command in the appropriate environment (Docker, CTF, SSH, Local) and workspace. @@ -745,6 +700,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg call_id: Unique ID for the command execution (for streaming) tool_name: Name of the tool being executed (for display in streaming output). If None, the tool name will be derived from the command. + args: Additional arguments for the tool (for display and context). Returns: str: Command output, status message, or session ID. @@ -753,6 +709,20 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_idle_timer() start_active_timer() + # Parse command into standard parts to ensure consistent naming + parts = command.strip().split(' ', 1) + cmd_name = parts[0] if parts else "" + cmd_args = parts[1] if len(parts) > 1 else "" + + # Generate a call_id if we're streaming and one wasn't provided + # Use a more specific format that includes the command name for easier tracking + if not call_id and stream: + call_id = f"cmd_{cmd_name}_{str(uuid.uuid4())[:8]}" + + # If no tool_name is provided, derive it from the command in a consistent way + if not tool_name: + tool_name = f"{cmd_name}_command" if cmd_name else "command" + try: # If session_id is provided, send command to that session if session_id: @@ -781,10 +751,6 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg return result # Return the result of sending input ("Input sent..." or error) - # Generate a call_id if we're streaming and one wasn't provided - if stream and not call_id: - call_id = str(uuid.uuid4())[:8] - # 2. Determine Execution Environment (Container > CTF > SSH > Local) active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST']) @@ -817,6 +783,34 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # Handle Streaming Container Execution if stream: + # Import the streaming utilities from util + from cai.util import start_tool_streaming, update_tool_streaming, finish_tool_streaming + + # Create args dictionary with standardized format + tool_args = { + "command": cmd_name, + "args": cmd_args if cmd_args.strip() else "", + "full_command": command, + "container": container_id[:12], + "environment": "Container", + "workspace": container_workspace + } + + # Add refresh rate info for generic_linux_command + if tool_name == "generic_linux_command": + tool_args["refresh_rate"] = 2 + + # Initialize the streaming session with a consistent call_id format + call_id = start_tool_streaming(tool_name, tool_args, call_id) + + # Update with "executing" status + update_tool_streaming( + tool_name, + tool_args, + f"Executing in container {container_id[:12]} at {container_workspace}:\n{command}\n\nPreparing environment...", + call_id + ) + # Ensure workspace directory exists inside the container first mkdir_cmd = [ "docker", "exec", container_id, @@ -829,6 +823,14 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg check=False, timeout=10 ) + + # Update status once environment is prepared + update_tool_streaming( + tool_name, + tool_args, + f"Executing in container {container_id[:12]} at {container_workspace}:\n{command}\n\nRunning command...", + call_id + ) # Build docker exec command as a single shell string for streaming docker_exec_cmd = ( @@ -837,20 +839,114 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg f"{shlex.quote(container_id)} sh -c " f"{shlex.quote(command)}" ) - - # Re-use the local streaming helper to provide real-time output - result = _run_local_streamed( - docker_exec_cmd, - call_id, - timeout, - tool_name, - workspace_dir=_get_workspace_dir() - ) - # Switch back to idle mode after streaming command completes - stop_active_timer() - start_idle_timer() - return result + try: + start_time = time.time() + # Start the process + process = subprocess.Popen( + docker_exec_cmd, + shell=True, # nosec B602 + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + cwd=_get_workspace_dir() + ) + + # Begin collecting output + output_buffer = [] + buffer_size = 0 + update_interval = 10 # lines + + # Stream stdout in real-time + for line in iter(process.stdout.readline, ''): + if not line: + break + + # Add to output collection + output_buffer.append(line) + buffer_size += 1 + + # Only update periodically to reduce UI refreshes + if buffer_size >= update_interval: + current_output = ''.join(output_buffer) + update_tool_streaming(tool_name, tool_args, current_output, call_id) + buffer_size = 0 + + # Finish process + process.stdout.close() + return_code = process.wait(timeout=timeout) + execution_time = time.time() - start_time + + # Get any stderr output + stderr_data = process.stderr.read() + if stderr_data: + output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) + + # Final output update + final_output = ''.join(output_buffer) + if return_code != 0: + final_output += f"\nCommand exited with code {return_code}" + + # Calculate execution info + execution_info = { + "status": "completed" if return_code == 0 else "error", + "return_code": return_code, + "environment": "Container", + "host": container_id[:12], + "tool_time": execution_time + } + + # Complete the streaming session with final output + finish_tool_streaming(tool_name, tool_args, final_output, call_id, execution_info) + + # Switch back to idle mode after streaming command completes + stop_active_timer() + start_idle_timer() + return final_output + + except subprocess.TimeoutExpired as e: + # Handle timeout + error_output = e.stdout if hasattr(e, 'stdout') and e.stdout else str(e) + error_msg = f"Command timed out after {timeout} seconds\n{error_output}" + + execution_info = { + "status": "timeout", + "environment": "Container", + "host": container_id[:12], + "error": str(e) + } + + # Complete with timeout error + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after timeout + stop_active_timer() + start_idle_timer() + # Fallback to local execution on timeout + print(color("Container execution timed out. Attempting execution on host instead.", fg="yellow")) + return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir(), args) + + except Exception as e: + # Handle other errors + error_msg = f"Error executing command in container: {str(e)}" + + execution_info = { + "status": "error", + "environment": "Container", + "host": container_id[:12], + "error": str(e) + } + + # Complete with error + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after error + stop_active_timer() + start_idle_timer() + # Fallback to local execution on error + print(color("Container execution failed. Attempting execution on host instead.", fg="yellow")) + return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir(), args) # Handle Synchronous Execution in Container try: @@ -887,7 +983,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_active_timer() start_idle_timer() # Fallback to local execution, preserving workspace context - return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501 # Switch back to idle mode after command completes stop_active_timer() @@ -903,7 +999,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_active_timer() start_idle_timer() # Fallback to local execution on timeout - return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501 except Exception as e: # pylint: disable=broad-except error_msg = f"Error executing command in container: {str(e)}" print(color(f"{context_msg} {error_msg}", fg="red")) @@ -912,58 +1008,216 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_active_timer() start_idle_timer() # Fallback to local execution on other errors - return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501 # --- CTF Execution --- if ctf: - # Handling streaming for CTF - not fully implemented yet + # If streaming is enabled and we have a call_id, show streaming UI for CTF too if stream: - from cai.util import cli_print_tool_output - if call_id and tool_name: - tool_args = {"command": command, "ctf": True} - cli_print_tool_output( - tool_name, - tool_args, - "Streaming not yet supported for CTF execution. Running normally...", - call_id=call_id - ) + # Import the streaming utilities from util + from cai.util import start_tool_streaming, update_tool_streaming, finish_tool_streaming + + # Create args dictionary with standardized format + tool_args = { + "command": cmd_name, + "args": cmd_args if cmd_args.strip() else "", + "full_command": command, + "environment": "CTF", + "workspace": os.path.basename(_get_workspace_dir()) + } + + # Add refresh rate info for generic_linux_command + if tool_name == "generic_linux_command": + tool_args["refresh_rate"] = 2 + + # Initialize the streaming session with a consistent call_id format + call_id = start_tool_streaming(tool_name, tool_args, call_id) + + target_dir = _get_workspace_dir() + full_command = f"cd '{target_dir}' && {command}" + + # Update with "executing" status + update_tool_streaming( + tool_name, + tool_args, + f"Executing in CTF environment: {full_command}\n\nWaiting for response...", + call_id + ) + + try: + # Execute the command and get the output + start_time = time.time() + output = ctf.get_shell(full_command, timeout=timeout) + execution_time = time.time() - start_time + + # Calculate execution info + execution_info = { + "status": "completed", + "environment": "CTF", + "tool_time": execution_time + } + + # Complete the streaming with final output + finish_tool_streaming(tool_name, tool_args, output, call_id, execution_info) + + # Switch back to idle mode after CTF command completes + stop_active_timer() + start_idle_timer() + return output + + except Exception as e: + # Handle errors in CTF execution + error_msg = f"Error executing CTF command: {str(e)}" + execution_info = { + "status": "error", + "environment": "CTF", + "error": str(e) + } + + # Complete the streaming with error output + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after error + stop_active_timer() + start_idle_timer() + return error_msg + else: + # Standard non-streaming CTF execution + result = _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir - # _run_ctf handles workspace internally using _get_workspace_dir() default - result = _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir - - # Switch back to idle mode after CTF command completes - stop_active_timer() - start_idle_timer() - return result + # Switch back to idle mode after CTF command completes + stop_active_timer() + start_idle_timer() + return result # --- SSH Execution --- if is_ssh_env: - # Async for SSH would require session management via SSH client features - if async_mode: - # Switch back to idle mode before returning message + # If streaming is enabled, show streaming UI for SSH too + if stream: + # Import the streaming utilities from util + from cai.util import start_tool_streaming, update_tool_streaming, finish_tool_streaming + + # Add SSH connection info for display + ssh_user = os.environ.get('SSH_USER', 'user') + ssh_host = os.environ.get('SSH_HOST', 'host') + ssh_connection = f"{ssh_user}@{ssh_host}" + + # Create args dictionary with standardized format + tool_args = { + "command": cmd_name, + "args": cmd_args if cmd_args.strip() else "", + "full_command": command, + "ssh_host": ssh_connection, + "environment": "SSH" + } + + # Add refresh rate info for generic_linux_command + if tool_name == "generic_linux_command": + tool_args["refresh_rate"] = 2 + + # Initialize streaming session with a consistent call_id format + call_id = start_tool_streaming(tool_name, tool_args, call_id) + + # Update with "executing" status + update_tool_streaming( + tool_name, + tool_args, + f"Executing on {ssh_connection}: {command}\n\nWaiting for response...", + call_id + ) + + try: + # Construct SSH command for execution + ssh_pass = os.environ.get('SSH_PASS') + if ssh_pass: + ssh_cmd_list = ["sshpass", "-p", ssh_pass, "ssh", ssh_connection] + else: + ssh_cmd_list = ["ssh", ssh_connection] + ssh_cmd_list.append(command) + + # Execute the command and get the output + start_time = time.time() + result = subprocess.run( + ssh_cmd_list, + capture_output=True, + text=True, + check=False, + timeout=timeout + ) + execution_time = time.time() - start_time + + # Get command output + output = result.stdout if result.stdout else result.stderr + + # Add SSH connection info to the output for clarity + result_with_info = f"Command executed on {ssh_connection}:\n\n{output}" + + # Determine status based on return code + status = "completed" if result.returncode == 0 else "error" + + # Calculate execution info + execution_info = { + "status": status, + "environment": "SSH", + "host": ssh_connection, + "return_code": result.returncode, + "tool_time": execution_time + } + + # Complete the streaming with final output + finish_tool_streaming(tool_name, tool_args, result_with_info, call_id, execution_info) + + # Switch back to idle mode after SSH command completes + stop_active_timer() + start_idle_timer() + return output.strip() + + except subprocess.TimeoutExpired as e: + # Handle timeout errors + error_output = e.stdout if e.stdout else str(e) + error_msg = f"Command timed out after {timeout} seconds\n{error_output}" + + execution_info = { + "status": "timeout", + "environment": "SSH", + "host": ssh_connection, + "error": str(e) + } + + # Complete the streaming with timeout error + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after timeout + stop_active_timer() + start_idle_timer() + return error_msg + + except Exception as e: + # Handle other errors + error_msg = f"Error executing SSH command: {str(e)}" + + execution_info = { + "status": "error", + "environment": "SSH", + "host": ssh_connection, + "error": str(e) + } + + # Complete the streaming with error + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after error + stop_active_timer() + start_idle_timer() + return error_msg + else: + # Standard non-streaming SSH execution + result = _run_ssh(command, stdout, timeout) # Workspace dir less relevant here + + # Switch back to idle mode after SSH command completes stop_active_timer() start_idle_timer() - return "Async mode not fully supported for SSH environment via this function yet." - - # Handling streaming for SSH - not fully implemented yet - if stream: - from cai.util import cli_print_tool_output - if call_id and tool_name: - tool_args = {"command": command, "ssh": True} - cli_print_tool_output( - tool_name, - tool_args, - "Streaming not yet supported for SSH execution. Running normally...", - call_id=call_id - ) - - # _run_ssh handles command execution, workspace is relative to remote home - result = _run_ssh(command, stdout, timeout) # Workspace dir less relevant here - - # Switch back to idle mode after SSH command completes - stop_active_timer() - start_idle_timer() - return result + return result # --- Local Execution (Default Fallback) --- # Let _run_local handle determining the host workspace @@ -989,8 +1243,18 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg start_idle_timer() return f"Started async session {new_session_id} locally. Use this ID to interact." - # Handle Synchronous Execution Locally using _run_local default with streaming support - result = _run_local(command, stdout, timeout, stream, call_id, tool_name, None) + # Handle Synchronous Execution Locally + # Pass stream=True if we're streaming to use streaming functionality + result = _run_local( + command, + stdout, + timeout, + stream=stream, + call_id=call_id, + tool_name=tool_name, + workspace_dir=_get_workspace_dir(), + custom_args=args + ) # Switch back to idle mode after local command completes stop_active_timer() diff --git a/src/cai/tools/reconnaissance/exec_code.py b/src/cai/tools/reconnaissance/exec_code.py index 794ac990..d532719a 100644 --- a/src/cai/tools/reconnaissance/exec_code.py +++ b/src/cai/tools/reconnaissance/exec_code.py @@ -37,62 +37,101 @@ def execute_code(code: str = "", language: str = "python", "python": "py", "php": "php", "bash": "sh", + "shell": "sh", # Add shell as alias for bash "ruby": "rb", "perl": "pl", "golang": "go", + "go": "go", # Add go as alias for golang "javascript": "js", + "js": "js", # Add js as alias for javascript "typescript": "ts", + "ts": "ts", # Add ts as alias for typescript "rust": "rs", "csharp": "cs", + "cs": "cs", # Add cs as alias for csharp "java": "java", - "kotlin": "kt" + "kotlin": "kt", + "c": "c", # Add C language + "cpp": "cpp", # Add C++ language + "c++": "cpp" # Add C++ language alias } - ext = extensions.get(language.lower(), "txt") + # Normalize language to lowercase + language = language.lower() + ext = extensions.get(language, "txt") full_filename = f"{filename}.{ext}" + # Create code file with content create_cmd = f"cat << 'EOF' > {full_filename}\n{code}\nEOF" - result = run_command(create_cmd, ctf=ctf) + result = run_command(create_cmd, ctf=ctf, stream=True, tool_name="execute_code") if "error" in result.lower(): return f"Failed to create code file: {result}" - if language.lower() == "python": + + # Prepare execution command based on language + if language in ["python", "py"]: exec_cmd = f"python3 {full_filename}" - elif language.lower() == "php": + elif language in ["php"]: exec_cmd = f"php {full_filename}" - elif language.lower() in ["bash", "sh"]: + elif language in ["bash", "sh", "shell"]: exec_cmd = f"bash {full_filename}" - elif language.lower() == "ruby": + elif language in ["ruby", "rb"]: exec_cmd = f"ruby {full_filename}" - elif language.lower() == "perl": + elif language in ["perl", "pl"]: exec_cmd = f"perl {full_filename}" - elif language.lower() == "golang" or language.lower() == "go": + elif language in ["golang", "go"]: temp_dir = f"/tmp/go_exec_{filename}" run_command(f"mkdir -p {temp_dir}", ctf=ctf) run_command(f"cp {full_filename} {temp_dir}/main.go", ctf=ctf) run_command(f"cd {temp_dir} && go mod init temp", ctf=ctf) exec_cmd = f"cd {temp_dir} && go run main.go" - elif language.lower() == "javascript": + elif language in ["javascript", "js"]: exec_cmd = f"node {full_filename}" - elif language.lower() == "typescript": + elif language in ["typescript", "ts"]: exec_cmd = f"ts-node {full_filename}" - elif language.lower() == "rust": + elif language in ["rust", "rs"]: # For Rust, we need to compile first run_command(f"rustc {full_filename} -o {filename}", ctf=ctf) exec_cmd = f"./{filename}" - elif language.lower() == "csharp": + elif language in ["csharp", "cs"]: # For C#, compile with dotnet run_command(f"dotnet build {full_filename}", ctf=ctf) exec_cmd = f"dotnet run {full_filename}" - elif language.lower() == "java": + elif language in ["java"]: # For Java, compile first run_command(f"javac {full_filename}", ctf=ctf) exec_cmd = f"java {filename}" - elif language.lower() == "kotlin": + elif language in ["kotlin", "kt"]: # For Kotlin, compile first run_command(f"kotlinc {full_filename} -include-runtime -d {filename}.jar", ctf=ctf) exec_cmd = f"java -jar {filename}.jar" + elif language in ["c"]: + # For C, compile with gcc + run_command(f"gcc {full_filename} -o {filename}", ctf=ctf) + exec_cmd = f"./{filename}" + elif language in ["cpp", "c++"]: + # For C++, compile with g++ + run_command(f"g++ {full_filename} -o {filename}", ctf=ctf) + exec_cmd = f"./{filename}" else: return f"Unsupported language: {language}" - output = run_command(exec_cmd, ctf=ctf, timeout=timeout) + # Execute the code with syntax-highlighted output + # Create a custom tool args dictionary to send language and code info to the tool output function + tool_args = { + "command": "execute", + "language": language, + "filename": filename, + "code": code, # Include the code for syntax highlighting + "timeout": timeout + } + + # Run the command with streaming to get syntax highlighting + output = run_command( + exec_cmd, + ctf=ctf, + timeout=timeout, + stream=True, + tool_name="execute_code", + args=tool_args + ) return output diff --git a/src/cai/util.py b/src/cai/util.py index cc6fc0c7..f786d7a3 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -5,6 +5,7 @@ import os import sys import importlib.resources import pathlib +import json from rich.console import Console from rich.tree import Tree from mako.template import Template # pylint: disable=import-error @@ -21,6 +22,11 @@ from dataclasses import dataclass, field from typing import Dict, Optional import time import threading +from rich.syntax import Syntax # Import Syntax for highlighting +from rich.panel import Panel +from rich.console import Group +from rich.box import ROUNDED +from rich.table import Table # Global timing variables for tracking active and idle time _active_timer_start = None @@ -29,6 +35,9 @@ _idle_timer_start = None _idle_time_total = 0.0 _timing_lock = threading.Lock() +# Set up a global tracker for live streaming panels +_LIVE_STREAMING_PANELS = {} + def start_active_timer(): """ Start measuring active time (when LLM is processing or tool is executing). @@ -539,6 +548,10 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 must have content. 3. If a tool call id appears alone (without a pair), it is removed. 4. There cannot be empty messages. + 5. Each tool_use block (assistant with tool_calls) must be followed by + a tool_result block (tool message with matching tool_call_id). + 6. Each 'tool' message must be immediately preceded by an 'assistant' message + with matching tool_call_id in its tool_calls. Args: messages (List[dict]): List of message dictionaries containing @@ -557,7 +570,13 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 for i, msg in enumerate(messages): # Skip empty messages (considered empty if 'content' is None or only whitespace) - if msg.get("role") in ["user", "system"] and (msg.get("content") is None or not msg.get("content").strip()): + if msg.get("role") in ["user", "system"] and (msg.get("content") is None or not str(msg.get("content", "")).strip()): + # Special case: if it's a system message, set content to empty string instead of skipping + if msg.get("role") == "system": + # Replace None with empty string + msg["content"] = "" + sanitized_messages.append(msg) + # Skip empty user messages entirely continue # Add valid messages to our sanitized list first @@ -595,8 +614,103 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 # Update mapping tool_call_map[tool_id] = {"assistant_idx": len(sanitized_messages) - 2, "tool_idx": len(sanitized_messages) - 1} + # Second pass - ensure correct sequence (tool messages must directly follow their assistant messages) + # This fixes the error "messages with role 'tool' must be a response to a preceeding message with 'tool_calls'" + i = 0 + while i < len(sanitized_messages): + msg = sanitized_messages[i] + + # Check if this is a tool message that might be out of sequence + if msg.get("role") == "tool" and msg.get("tool_call_id"): + tool_id = msg.get("tool_call_id") + + # If this isn't the first message, check if the previous message is a matching assistant message + if i > 0: + prev_msg = sanitized_messages[i-1] + + # Check if the previous message is an assistant message with matching tool_call_id + is_valid_sequence = ( + prev_msg.get("role") == "assistant" and + prev_msg.get("tool_calls") and + any(tc.get("id") == tool_id for tc in prev_msg.get("tool_calls", [])) + ) + + if not is_valid_sequence: + # Find the assistant message with this tool_call_id + assistant_idx = None + for j, assistant_msg in enumerate(sanitized_messages): + if (assistant_msg.get("role") == "assistant" and + assistant_msg.get("tool_calls") and + any(tc.get("id") == tool_id for tc in assistant_msg.get("tool_calls", []))): + assistant_idx = j + break + + # If we found a matching assistant message, move this tool message right after it + if assistant_idx is not None: + # Remember to save the tool message + tool_msg = sanitized_messages.pop(i) + + # Insert right after the assistant message + sanitized_messages.insert(assistant_idx + 1, tool_msg) + + # Adjust i to account for the move + if assistant_idx < i: + # We moved the message backward, so i should point to the next message + # which is now at position i (since we removed a message before it) + continue + else: + # We moved the message forward, so i should now point to the message + # that is now at position i + continue + else: + # No matching assistant message found - create one + assistant_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": tool_id, + "type": "function", + "function": { + "name": "unknown_function", + "arguments": "{}" + } + }] + } + + # Insert the assistant message before the tool message + sanitized_messages.insert(i, assistant_msg) + + # Skip past both messages + i += 2 + continue + else: + # This tool message is at index 0, which means there's no preceding assistant message + # Create a dummy assistant message + assistant_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": tool_id, + "type": "function", + "function": { + "name": "unknown_function", + "arguments": "{}" + } + }] + } + + # Insert the assistant message before the tool message + sanitized_messages.insert(0, assistant_msg) + + # Skip past both messages + i += 2 + continue + + # Move to the next message + i += 1 + # Final validation - ensure all tool calls have responses - for tool_id, indices in tool_call_map.items(): + for tool_id, indices in list(tool_call_map.items()): if indices["tool_idx"] is None: # Tool call without a response - create a synthetic tool message assistant_idx = indices["assistant_idx"] @@ -616,8 +730,59 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 "tool_call_id": tool_id, "content": f"Auto-generated response for {tool_name}" } - # Insert right after the assistant message - sanitized_messages.insert(assistant_idx + 1, tool_msg) + + # Insert immediately after the assistant message + if assistant_idx + 1 < len(sanitized_messages): + # Insert at the position after assistant + sanitized_messages.insert(assistant_idx + 1, tool_msg) + else: + # Just append if we're at the end + sanitized_messages.append(tool_msg) + + # Update the map to note that this tool call now has a response + tool_call_map[tool_id]["tool_idx"] = assistant_idx + 1 + + # Ensure messages have non-null content (required by some providers) + for msg in sanitized_messages: + if msg.get("role") != "tool" and msg.get("content") is None and not msg.get("tool_calls"): + msg["content"] = "" + + # For tool messages, ensure content is never null + if msg.get("role") == "tool" and msg.get("content") is None: + msg["content"] = f"Tool response for {msg.get('tool_call_id', 'unknown')}" + + # Special case for Claude: ensure strict alternating pattern between assistant tool_calls and tool results + # If multiple consecutive assistant messages with tool_calls exist, interleave them with tool responses + i = 0 + while i < len(sanitized_messages) - 1: + current_msg = sanitized_messages[i] + next_msg = sanitized_messages[i + 1] + + # When current message is assistant with tool_calls and next message is NOT a tool response + if (current_msg.get("role") == "assistant" and + current_msg.get("tool_calls") and + (next_msg.get("role") != "tool" or not next_msg.get("tool_call_id"))): + + # Get the first tool call ID + tool_id = current_msg["tool_calls"][0].get("id", "unknown") + tool_name = "unknown_function" + if current_msg["tool_calls"][0].get("function"): + tool_name = current_msg["tool_calls"][0]["function"].get("name", "unknown_function") + + # Create a tool result message + tool_msg = { + "role": "tool", + "tool_call_id": tool_id, + "content": f"Auto-generated response for {tool_name}" + } + + # Insert the tool message after the current assistant message + sanitized_messages.insert(i + 1, tool_msg) + + # Skip over the newly inserted message + i += 2 + else: + i += 1 return sanitized_messages @@ -665,22 +830,23 @@ def get_model_name(model): return model # If not a string, use environment variable return os.environ.get('CAI_MODEL', 'qwen2.5:72b') - # Helper function to format time in a human-readable way + +# Helper function to format time in a human-readable way def format_time(seconds): - if seconds is None: - return "N/A" - - if seconds < 60: - return f"{seconds:.1f}s" - elif seconds < 3600: - minutes = int(seconds / 60) - seconds_remainder = seconds % 60 - return f"{minutes}m {seconds_remainder:.1f}s" - else: - hours = int(seconds / 3600) - minutes = int((seconds % 3600) / 60) - return f"{hours}h {minutes}m" - + if seconds is None: + return "N/A" + + if seconds < 60: + return f"{seconds:.1f}s" + elif seconds < 3600: + minutes = int(seconds / 60) + seconds_remainder = seconds % 60 + return f"{minutes}m {seconds_remainder:.1f}s" + else: + hours = int(seconds / 3600) + minutes = int((seconds % 3600) / 60) + return f"{hours}h {minutes}m" + def get_model_pricing(model_name): """ Get pricing information for a model, using the CostTracker's implementation. @@ -804,28 +970,116 @@ def parse_message_content(message): """ Parse a message object to extract its textual content. Only processes messages that don't have tool calls. + Detects markdown code blocks and applies syntax highlighting in non-streaming mode. + Also formats other markdown elements like headers, lists, and text formatting. Args: message: Can be a string or a Message object with content attribute Returns: - str: The extracted content as a string + str or rich.console.Group: The extracted content as a string or as a rich Group with Syntax highlighting """ - # Check if this is a duplicate print from OpenAIChatCompletionsModel - # If message is already a string, return it + from rich.console import Group + from rich.syntax import Syntax + from rich.text import Text + from rich.markdown import Markdown + import re + + # Extract the raw content + raw_content = "" + + # If message is already a string, use it if isinstance(message, str): - return message - + raw_content = message # If message is a Message object with content attribute - if hasattr(message, 'content') and message.content is not None: - return message.content - + elif hasattr(message, 'content') and message.content is not None: + raw_content = message.content # If message is a dict with content key - if isinstance(message, dict) and 'content' in message: - return message['content'] - + elif isinstance(message, dict) and 'content' in message: + raw_content = message['content'] # If we can't extract content, convert to string - return str(message) + else: + raw_content = str(message) + + # Check if streaming is enabled + streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' + + # Only apply markdown formatting in non-streaming mode + if not streaming_enabled and raw_content: + # Check if content contains markdown code blocks with improved regex + code_block_pattern = r'```(\w*)\s*([\s\S]*?)\s*```' + matches = re.findall(code_block_pattern, raw_content, re.DOTALL) + + if matches: + # Prepare to process markdown with code blocks highlighted + elements = [] + last_end = 0 + + # Find all code blocks with improved regex pattern + for match in re.finditer(r'```(\w*)\s*([\s\S]*?)\s*```', raw_content, re.DOTALL): + # Get text before the code block + start = match.start() + if start > last_end: + text_before = raw_content[last_end:start] + + # Process markdown in the text before the code block + if text_before.strip(): + md = Markdown(text_before) + elements.append(md) + + # Process the code block + lang = match.group(1) or "text" + code = match.group(2) + + # Use the language mapping helper to get proper syntax highlighting + syntax_lang = get_language_from_code_block(lang) + + # Create syntax highlighted code + syntax = Syntax( + code, + syntax_lang, + theme="monokai", + line_numbers=True, + word_wrap=True, + background_color="#272822" + ) + elements.append(syntax) + + last_end = match.end() + + # Add any remaining text after the last code block + if last_end < len(raw_content): + text_after = raw_content[last_end:] + + # Process markdown in the text after the code block + if text_after.strip(): + md = Markdown(text_after) + elements.append(md) + + return Group(*elements) + else: + # If no code blocks, but still contains markdown, use Rich's markdown renderer + # Check for markdown elements (headers, lists, formatting) + has_markdown = any([ + # Headers + re.search(r'^#{1,6}\s+\w+', raw_content, re.MULTILINE), + # Lists + re.search(r'^\s*[-*+]\s+\w+', raw_content, re.MULTILINE), + re.search(r'^\s*\d+\.\s+\w+', raw_content, re.MULTILINE), + # Bold/Italic + '**' in raw_content, + '*' in raw_content and not '**' in raw_content, + '__' in raw_content, + '_' in raw_content and not '__' in raw_content, + # Links + re.search(r'\[.+?\]\(.+?\)', raw_content) + ]) + + if has_markdown: + return Group(Markdown(raw_content)) + + # For streaming mode or no markdown, return the raw content + return raw_content def parse_message_tool_call(message, tool_output=None): """ @@ -946,7 +1200,8 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli total_reasoning_tokens=None, interaction_cost=None, total_cost=None, - tool_output=None): # New parameter for tool output + tool_output=None, # New parameter for tool output + suppress_empty=False): # New parameter to suppress empty panels """Print agent messages/thoughts with enhanced visual formatting.""" # Debug prints to trace the function calls if debug: @@ -978,19 +1233,40 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli else: parsed_message = parse_message_content(message) tool_panels = [] + + # Skip empty panels - THIS IS THE KEY CHANGE + # If suppress_empty is True and there's no parsed message and no tool panels, + # don't create an empty panel to avoid cluttering during streaming + if suppress_empty and not parsed_message and not tool_panels: + return + + # Check if parsed_message is empty or "null" + is_empty_message = (parsed_message == "null" or parsed_message == "" or + (isinstance(parsed_message, str) and not parsed_message.strip())) + + # Also skip if the only message is "null" or empty + if is_empty_message: + if suppress_empty and not tool_panels: + return + + # Check if we have Group content from markdown parsing + is_rich_content = False + from rich.console import Group + if isinstance(parsed_message, Group): + is_rich_content = True # Special handling for Reasoner Agent if agent_name == "Reasoner Agent": text.append(f"[{counter}] ", style="bold red") text.append(f"Agent: {agent_name} ", style="bold yellow") - if parsed_message: + if parsed_message and not is_rich_content: text.append(f">> {parsed_message} ", style="green") text.append(f"[{timestamp}", style="dim") if model: text.append(f" ({os.getenv('CAI_SUPPORT_MODEL')})", style="bold blue") text.append("]", style="dim") - elif not parsed_message: + elif is_empty_message: # When parsed_message is empty, only include timestamp and model info text.append(f"Agent: {agent_name} ", style="bold green") text.append(f"[{timestamp}", style="dim") @@ -1000,7 +1276,7 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli else: text.append(f"[{counter}] ", style="bold cyan") text.append(f"Agent: {agent_name} ", style="bold green") - if parsed_message: + if parsed_message and not is_rich_content: text.append(f">> {parsed_message} ", style="yellow") text.append(f"[{timestamp}", style="dim") if model: @@ -1028,26 +1304,53 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli total_cost ) # Only append token information if there is a parsed message - if parsed_message: + if parsed_message and not is_rich_content: text.append(tokens_text) - panel = Panel( - text, - border_style="red" if agent_name == "Reasoner Agent" else "blue", - box=ROUNDED, - padding=(0, 1), - title=("[bold]Reasoning Analysis[/bold]" - if agent_name == "Reasoner Agent" - else "[bold]Agent Interaction[/bold]"), - title_align="left" - ) + # Create the panel content based on whether we have rich content or not + from rich.panel import Panel + from rich.console import Group + + if is_rich_content: + # For rich content, create a Group with the header, content, and tokens + panel_content = [] + panel_content.append(text) + + # Add spacing between header and content for better readability + panel_content.append(Text("\n")) + + # Add the Group with highlighted content + panel_content.append(parsed_message) + + # Add token information at the bottom with proper spacing + if tokens_text: + panel_content.append(Text("\n")) + panel_content.append(tokens_text) + + panel = Panel( + Group(*panel_content), + border_style="red" if agent_name == "Reasoner Agent" else "blue", + box=ROUNDED, + padding=(1, 1), # Increased padding for better appearance + title="", + title_align="left" + ) + else: + # For regular text content, use the original panel format + panel = Panel( + text, + border_style="red" if agent_name == "Reasoner Agent" else "blue", + box=ROUNDED, + padding=(0, 1), + title="", + title_align="left" + ) #console.print("\n") console.print(panel) # If there are tool panels, print them after the main message panel # But only in non-streaming mode to avoid duplicates - is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' - if tool_panels and not is_streaming_enabled: + if tool_panels: for tool_panel in tool_panels: console.print(tool_panel) @@ -1063,6 +1366,15 @@ def create_agent_streaming_context(agent_name, counter, model): Returns: A dictionary with the streaming context """ + # Add a static variable to track active streaming contexts and prevent duplicates + if not hasattr(create_agent_streaming_context, "_active_streaming"): + create_agent_streaming_context._active_streaming = {} + + # If there's already an active streaming context with the same counter, return it + context_key = f"{agent_name}_{counter}" + if context_key in create_agent_streaming_context._active_streaming: + return create_agent_streaming_context._active_streaming[context_key] + try: from rich.live import Live import shutil @@ -1099,17 +1411,17 @@ def create_agent_streaming_context(agent_name, counter, model): Text.assemble(header, content, footer), border_style="blue", box=ROUNDED, - padding=(1, 2), - title="[bold]Agent Streaming Response[/bold]", + padding=(0, 1), + title="Stream", title_align="left", width=panel_width, expand=True ) # Create Live display object but don't start it until we have content - live = Live(panel, refresh_per_second=20, console=console, auto_refresh=False) + live = Live(panel, refresh_per_second=10, console=console, auto_refresh=True, vertical_overflow="visible") - return { + context = { "live": live, "panel": panel, "header": header, @@ -1122,6 +1434,11 @@ def create_agent_streaming_context(agent_name, counter, model): "is_started": False, # Track if we've started the display "error": None, # Track any errors } + + # Store the context for potential reuse + create_agent_streaming_context._active_streaming[context_key] = context + + return context except Exception as e: # If rich display fails, return None and log the error import sys @@ -1155,8 +1472,8 @@ def update_agent_streaming_content(context, text_delta): Text.assemble(context["header"], context["content"], context["footer"]), border_style="blue", box=ROUNDED, - padding=(1, 2), - title="[bold]Agent Streaming Response[/bold]", + padding=(0, 1), + title="Stream", title_align="left", width=context.get("panel_width", 100), expand=True @@ -1191,6 +1508,13 @@ def finish_agent_streaming(context, final_stats=None): """ if not context: 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 try: # Check if there's actual content to display - don't show empty panels @@ -1253,14 +1577,13 @@ def finish_agent_streaming(context, final_stats=None): Text.assemble( context["header"], context["content"], - Text("\n\n"), - tokens_text if tokens_text else Text(""), + tokens_text if tokens_text else Text(""), context["footer"] ), border_style="blue", box=ROUNDED, - padding=(1, 2), - title="[bold]Agent Streaming Response[/bold]", + padding=(0, 1), + title="Stream", title_align="left", width=context.get("panel_width", 100), expand=True @@ -1293,7 +1616,8 @@ def finish_agent_streaming(context, final_stats=None): return False -def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execution_info=None, token_info=None): + +def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execution_info=None, token_info=None, streaming=False): """ Print a tool call output to the command line. Similar to cli_print_tool_call but for the output of the tool. @@ -1309,44 +1633,181 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut - total_input_tokens, total_output_tokens, total_reasoning_tokens - model: model name string - interaction_cost, total_cost: optional cost values + streaming: Flag indicating if this is part of a streaming output """ - # If it's an empty output, don't print anything - if not output and not call_id: + # If it's an empty output, don't print anything except for streaming sessions + if not output and not call_id and not streaming: return - - # CRITICAL CHECK: When in streaming mode (CAI_STREAM=true), ONLY show output panels - # for streaming updates (those with call_id). This prevents duplicate output panels. - is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' - if is_streaming_enabled and not call_id: - # Skip all non-streaming tool output in streaming mode + # Skip early for execute_code tool in non-streaming mode + if tool_name == "execute_code" and not streaming: return - # Track seen call IDs to prevent duplicate panels + # Set up global tracker for streaming sessions + if not hasattr(cli_print_tool_output, '_streaming_sessions'): + cli_print_tool_output._streaming_sessions = {} + + # Track seen call IDs to prevent duplicate panels for non-streaming outputs if not hasattr(cli_print_tool_output, '_seen_calls'): cli_print_tool_output._seen_calls = {} - - # For streaming updates, only show updates for the same call_id - # but allow the first appearance of each call_id - if call_id: - call_key = f"{call_id}:{output[:20]}" # Use first 20 chars as fingerprint with call_id - # Skip if we've seen this exact output for this call_id before - if call_key in cli_print_tool_output._seen_calls: + # Track all displayed commands to prevent duplicates + if not hasattr(cli_print_tool_output, '_displayed_commands'): + cli_print_tool_output._displayed_commands = set() + + # --- Consistent Command Key Generation --- + effective_command_args_str = "" + if isinstance(args, dict): + # If args is a dictionary, extract the 'args' field. + effective_command_args_str = args.get("args", "") + elif isinstance(args, str): + # If args is a string, it might be a JSON representation or a plain string. + try: + parsed_json_args = json.loads(args) + if isinstance(parsed_json_args, dict): + # Parsed as JSON dict, get the 'args' field. + effective_command_args_str = parsed_json_args.get("args", "") + else: + # Parsed as JSON, but not a dict (e.g., a JSON string literal). + effective_command_args_str = parsed_json_args if isinstance(parsed_json_args, str) else args + except json.JSONDecodeError: + # Not a JSON string, treat 'args' as a plain string. + effective_command_args_str = args + + command_key = f"{tool_name}:{effective_command_args_str}" + # --- End of Command Key Generation --- + + # Check for duplicate display conditions + if streaming: + # For streaming updates, track and update the single streaming session + if call_id: + # If this is a new streaming session, record it + if call_id not in cli_print_tool_output._streaming_sessions: + cli_print_tool_output._streaming_sessions[call_id] = { + 'tool_name': tool_name, + 'args': args, # Store original args for display formatting + 'buffer': output if output else "", + 'start_time': time.time(), + 'last_update': time.time(), + 'command_key': command_key, # Store the generated key + 'is_complete': False + } + # Add the command key to displayed commands + if command_key not in cli_print_tool_output._displayed_commands: + cli_print_tool_output._displayed_commands.add(command_key) + else: + # Update the existing session + session = cli_print_tool_output._streaming_sessions[call_id] + # Always replace buffer with latest output for consistency + session['buffer'] = output + session['last_update'] = time.time() + if execution_info and execution_info.get('is_final', False): + session['is_complete'] = True + + # For streaming outputs, we'll use Rich Live panel if available + try: + from rich.console import Console + from rich.live import Live + from rich.panel import Panel + from rich.text import Text + from rich.box import ROUNDED + + # Access the global live panel dictionary + global _LIVE_STREAMING_PANELS + + # Create the header, content, and panel + # Pass the original 'args' (dict or string) to _create_tool_panel_content for formatting + current_args_for_display = cli_print_tool_output._streaming_sessions[call_id]['args'] + header, content = _create_tool_panel_content( + tool_name, + current_args_for_display, + cli_print_tool_output._streaming_sessions[call_id]['buffer'], + execution_info, + token_info + ) + + # Determine panel style based on status + status = "running" + if execution_info: + status = execution_info.get('status', 'running') + + border_style = "yellow" # Default for running + if status == "completed": + border_style = "green" + elif status in ["error", "timeout"]: + border_style = "red" + + # Create panel title based on status + if status == "running": + title = "[bold yellow]Running[/bold yellow]" + elif status == "completed": + title = "[bold green]Completed[/bold green]" + elif status == "error": + title = "[bold red]Error[/bold red]" + elif status == "timeout": + title = "[bold red]Timeout[/bold red]" + else: + title = "[bold blue]Tool Execution[/bold blue]" + + # Create the panel + panel = Panel( + content, + title=title, + border_style=border_style, + padding=(0, 1), + box=ROUNDED, + title_align="left" + ) + + # 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) + + # 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() + # Remove from the active panel dictionary + 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 + + # Return early for streaming updates + return + + except ImportError: + # Fall back to simple updates without Rich + pass + else: + # For non-streaming outputs, check if we've already seen this command + if command_key in cli_print_tool_output._displayed_commands: + # Command has already been displayed (likely through streaming), skip duplicate display return - # Mark as seen - cli_print_tool_output._seen_calls[call_key] = True + # Add to displayed commands since we're going to show it + # This handles the case where a command is non-streaming from the start + cli_print_tool_output._displayed_commands.add(command_key) + + # For non-streaming updates with call_id, check if already seen + # This _seen_calls logic is an additional layer for non-streaming calls that might have call_ids + # but might be distinct from the primary _displayed_commands check based on command_key. + if call_id and not streaming: + # Create a more specific key for _seen_calls if needed, possibly including output fingerprint + seen_call_key = f"{call_id}:{command_key}:{output[:20]}" - # Limit cache size to prevent memory growth - if len(cli_print_tool_output._seen_calls) > 1000: - # Keep only the most recent 500 entries - cli_print_tool_output._seen_calls = { - k: cli_print_tool_output._seen_calls[k] - for k in list(cli_print_tool_output._seen_calls.keys())[-500:] - } - - # Try to use Rich for better formatting if available + if seen_call_key in cli_print_tool_output._seen_calls: + return + + cli_print_tool_output._seen_calls[seen_call_key] = True + + # Standard tool output display for non-streaming or when rich is not available try: from rich.console import Console from rich.panel import Panel @@ -1357,304 +1818,655 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Create a console for output console = Console(theme=theme) - # Format arguments for display - # Parse JSON string if args is a string - if isinstance(args, str) and args.strip().startswith('{'): - try: - import json - args = json.loads(args) - except: - # Keep as is if not valid JSON - pass + # Get the panel content - with syntax highlighting + header, content = _create_tool_panel_content(tool_name, args, output, execution_info, token_info) - # Format arguments as a clean string - if isinstance(args, dict): - # Only include non-empty values and exclude async_mode=false - arg_parts = [] - for key, value in args.items(): - # Skip empty values - if value == "" or value == {} or value is None: - continue - # Skip async_mode=false (default) - if key == "async_mode" and value is False: - continue - # Format the value - if isinstance(value, str): - arg_parts.append(f"{key}={value}") - else: - arg_parts.append(f"{key}={value}") - args_str = ", ".join(arg_parts) - else: - args_str = str(args) + # Format args for the title display + args_str = _format_tool_args(args, tool_name=tool_name) - # Get session timing information - try: - from cai.cli import START_TIME - total_time = time.time() - START_TIME if START_TIME else None - except ImportError: - total_time = None - + # Determine border style based on status + border_style = "blue" # Default for non-streaming - # Extract execution timing info - - tool_time = None - status = None if execution_info: - # Prefer 'tool_time' if present, else fallback to 'time_taken' - - tool_time = execution_info.get('tool_time') status = execution_info.get('status', 'completed') + if status == "completed": + border_style = "green" + title = f"[bold green]{tool_name}({args_str}) [Completed][/bold green]" + elif status == "error": + border_style = "red" + title = f"[bold red]{tool_name}({args_str}) [Error][/bold red]" + elif status == "timeout": + border_style = "red" + title = f"[bold red]{tool_name}({args_str}) [Timeout][/bold red]" + else: + title = f"[bold blue]{tool_name}({args_str})[/bold blue]" + else: + title = f"[bold blue]{tool_name}({args_str})[/bold blue]" - # Create header for all panel displays (both streaming and non-streaming) - header = Text() - header.append(tool_name, style="#00BCD4") - header.append("(", style="yellow") - header.append(args_str, style="yellow") - header.append(")", style="yellow") - - # Add timing information directly in the header - timing_info = [] - if total_time: - timing_info.append(f"Total: {format_time(total_time)}") - if tool_time: - timing_info.append(f"Tool: {format_time(tool_time)}") - if timing_info: - header.append(f" [{' | '.join(timing_info)}]", style="cyan") - - # Add completion status if available - REMOVED, just showing timing now - # if status: - # if status == 'completed': - # header.append(f" [Completed]", style="green") - # elif status == 'running': - # header.append(f" [Running]", style="yellow") - # elif status == 'error': - # header.append(f" [Error]", style="red") - # elif status == 'timeout': - # header.append(f" [Timeout]", style="red") - # else: - # header.append(f" [{status.title()}]", style="dim") - - # For streaming mode with call_id, use Rich Live display - if call_id: - # Create token information if available - token_content = None - if token_info: - model = token_info.get('model', '') - interaction_input_tokens = token_info.get('interaction_input_tokens', 0) - interaction_output_tokens = token_info.get('interaction_output_tokens', 0) - interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) - total_input_tokens = token_info.get('total_input_tokens', 0) - total_output_tokens = token_info.get('total_output_tokens', 0) - total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) - - if (interaction_input_tokens > 0 or total_input_tokens > 0): - token_text = _create_token_display( - interaction_input_tokens, - interaction_output_tokens, - interaction_reasoning_tokens, - total_input_tokens, - total_output_tokens, - total_reasoning_tokens, - model, - token_info.get('interaction_cost'), - token_info.get('total_cost') - ) - token_content = Text("\n\n") - token_content.append(token_text) - - # Create content text with the output - content = Text(output) - - # Create the panel for display, including token info if available - panel_content = [header, Text("\n\n"), content] - if token_content: - panel_content.append(token_content) - - # Create title - simple title with no timing info - title = "[bold blue]Tool Output[/bold blue]" - - panel = Panel( - Text.assemble(*panel_content), - title=title, - border_style="blue", - padding=(1, 2), - box=ROUNDED, - title_align="left" - ) - - # Display using Rich - console.print(panel) - return - - # For non-streaming output, also use a blue panel with the same format - # Create token information if available - token_text = None - if token_info: - model = token_info.get('model', '') - interaction_input_tokens = token_info.get('interaction_input_tokens', 0) - interaction_output_tokens = token_info.get('interaction_output_tokens', 0) - interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) - total_input_tokens = token_info.get('total_input_tokens', 0) - total_output_tokens = token_info.get('total_output_tokens', 0) - total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) - interaction_cost = token_info.get('interaction_cost') - total_cost = token_info.get('total_cost') - - # Generate token display with CostTracker - if (interaction_input_tokens > 0 or total_input_tokens > 0): - token_text = _create_token_display( - interaction_input_tokens, - interaction_output_tokens, - interaction_reasoning_tokens, - total_input_tokens, - total_output_tokens, - total_reasoning_tokens, - model, - interaction_cost, - total_cost - ) - - # Now create the panel content, starting with the header - panel_content = [header, Text("\n")] - - # Add token display if available - if token_text: - panel_content.append(token_text) - panel_content.append(Text("\n")) # Add spacing after token display - - # Add the output - if output: - output_text = Text(output) - panel_content.append(output_text) - - # If no content was added but we have output, add it directly - if len(panel_content) == 2 and output: # Only header and newline - panel_content.append(Text(output)) - - # Create title - simple title with no timing info - title = "[bold blue]Tool Output[/bold blue]" - - # Create the final panel - always blue now + # Create the panel panel = Panel( - Group(*panel_content), + content, title=title, - border_style="blue", - padding=(1, 2), - box=ROUNDED + border_style=border_style, + padding=(0, 1), + box=ROUNDED, + title_align="left" ) # Display the panel console.print(panel) except ImportError: - # Fall back to simple formatting if Rich is not available - # Format arguments in the cleaner format - # Parse JSON string if args is a string - if isinstance(args, str) and args.strip().startswith('{'): - try: - import json - args = json.loads(args) - except: - # Keep as is if not valid JSON - pass - - # Format arguments as a clean string - if isinstance(args, dict): - # Only include non-empty values and exclude async_mode=false - arg_parts = [] - for key, value in args.items(): - # Skip empty values - if value == "" or value == {} or value is None: - continue - # Skip async_mode=false (default) - if key == "async_mode" and value is False: - continue - # Format the value - if isinstance(value, str): - arg_parts.append(f"{key}={value}") - else: - arg_parts.append(f"{key}={value}") - args_str = ", ".join(arg_parts) + pass + + +# Helper function to create tool panel content +def _create_tool_panel_content(tool_name, args, output, execution_info=None, token_info=None): + """Create the header and content for a tool output panel.""" + from rich.text import Text + from rich.syntax import Syntax # Import Syntax for highlighting + from rich.panel import Panel + from rich.console import Group + from rich.box import ROUNDED + + # Format arguments for display, passing tool_name for specific formatting + args_str = _format_tool_args(args, tool_name=tool_name) + + # Get timing information + timing_info, tool_time = _get_timing_info(execution_info) + + # Create header + header = Text() + header.append(tool_name, style="#00BCD4") + header.append("(", style="yellow") + header.append(args_str, style="yellow") + header.append(")", style="yellow") + + # Add timing information + if timing_info: + header.append(f" [{' | '.join(timing_info)}]", style="cyan") + + # Add environment info if available + if execution_info and execution_info.get('environment'): + env = execution_info.get('environment') + host = execution_info.get('host', '') + if host: + header.append(f" [{env}:{host}]", style="magenta") else: - args_str = str(args) - # Get session timing information + header.append(f" [{env}]", style="magenta") + + # Add status information if available + if execution_info: + status = execution_info.get('status', None) + if status == "completed": + header.append(" [Completed]", style="green") + elif status == "running": + header.append(" [Running]", style="yellow") + elif status == "error": + header.append(" [Error]", style="red") + elif status == "timeout": + header.append(" [Timeout]", style="red") + + # Create token information if available + token_content = _create_token_info_display(token_info) + + # Determine if we need specialized content formatting + group_content = [header] + + # Special handling for execute_code tool + if tool_name == "execute_code" and isinstance(args, dict): + command_type = args.get("command") + actual_code = args.get("code") + language = args.get("language", "python") # Default to python + + # For execute command, show code and output panels + if command_type == "execute" and actual_code: + # Ensure language is a string for Syntax + language_str = get_language_from_code_block(str(language)) + + code_syntax = Syntax(actual_code, language_str, theme="monokai", + line_numbers=True, background_color="#272822", + indent_guides=True, word_wrap=True) + code_panel = Panel( + code_syntax, + title=f"Code ({language_str})", + border_style="cyan", + title_align="left", + box=ROUNDED, + padding=(0,1) + ) + group_content.extend([Text("\n"), code_panel]) + + # Panel for the output of the executed code + if output: + # Try to highlight output as text, or specific language if known (e.g. json) + output_lang = "text" + try: + json.loads(output) # Check if output is JSON + output_lang = "json" + except json.JSONDecodeError: + pass # Not JSON, keep as text + + output_syntax = Syntax(output, output_lang, theme="monokai", + background_color="#272822", word_wrap=True) + output_panel = Panel( + output_syntax, + title="Output", + border_style="green", + title_align="left", + box=ROUNDED, + padding=(0,1) + ) + group_content.extend([Text("\n"), output_panel]) + # For other commands (like cat) or if no code, just show output if any + elif output: + output_syntax = Syntax(output, "text", theme="monokai", + background_color="#272822", word_wrap=True) + output_panel = Panel( + output_syntax, + title="Output", + border_style="green", + title_align="left", + box=ROUNDED, + padding=(0,1) + ) + group_content.extend([Text("\n"), output_panel]) + + # Special handling for generic_linux_command or any command containing 'command' + elif "command" in tool_name.lower() or "shell" in tool_name.lower(): try: - from cai.cli import START_TIME - total_time = time.time() - START_TIME if START_TIME else None - except ImportError: - total_time = None + # Highlight the output as bash + output_syntax = Syntax(output, "bash", theme="monokai", + background_color="#272822", word_wrap=True) - - # For non-streaming output, use the original formatting - tool_call = f"{tool_name}({args_str})" - - # Get tool execution time if available - tool_time_str = "" - execution_status = "" - if execution_info: - time_taken = execution_info.get('time_taken', 0) - status = execution_info.get('status', 'completed') + # Create a panel for the formatted output + output_panel = Panel( + output_syntax, + title="Command Output", + border_style="green", + title_align="left", + box=ROUNDED, + padding=(0, 1) + ) - # Add execution info to the tool call display - if time_taken: - tool_time_str = f"Tool: {format_time(time_taken)}" - execution_status = f" [{status} in {time_taken:.2f}s]" + # Assemble content with highlighted output + group_content.extend([Text("\n"), output_panel]) + + except Exception: + # Fallback if syntax highlighting fails, just add raw output + group_content.extend([Text("\n"), Text(output)]) + + # Add token info if available + if token_content: + group_content.extend([Text("\n"), token_content]) + + return header, Group(*group_content) + +# Helper function to format tool arguments +def _format_tool_args(args, tool_name=None): + """Format tool arguments as a clean string.""" + # If args is already a string, it might be pre-formatted or a simple arg string + if isinstance(args, str): + # If it looks like a JSON dict string, try to parse and format nicely + if args.strip().startswith('{') and args.strip().endswith('}'): + try: + + parsed_dict = json.loads(args) + # Recursively call with the parsed dict for consistent formatting + return _format_tool_args(parsed_dict, tool_name=tool_name) + except json.JSONDecodeError: + # Not valid JSON, or not a dict; return as is + return args + else: + # Simple string arg, return as is + return args + + # Format arguments from a dictionary + if isinstance(args, dict): + # Only include non-empty values and exclude special flags + arg_parts = [] + for key, value in args.items(): + # Skip empty values + if value == "" or value == {} or value is None: + continue + # Skip special flags + if key in ["async_mode", "streaming"] and not value: + continue + + value_str = str(value) + + # Format the value + if isinstance(value, str): + # Truncate long string values that are not specifically handled above + if len(value_str) > 70 and key not in ["code", "args"]: + value_str = value_str[:67] + "..." + arg_parts.append(f"{key}={value_str}") else: - execution_status = f" [{status}]" + arg_parts.append(f"{key}={value_str}") + return ", ".join(arg_parts) + else: + return str(args) + +# Helper function to get timing information +def _get_timing_info(execution_info=None): + """Get timing information for display.""" + import time + + # Get session timing information + try: + from cai.cli import START_TIME + total_time = time.time() - START_TIME if START_TIME else None + except ImportError: + total_time = None + + # Extract execution timing info + tool_time = None + if execution_info: + tool_time = execution_info.get('tool_time') + + # Format timing info for display + timing_info = [] + if total_time: + timing_info.append(f"Total: {format_time(total_time)}") + if tool_time: + timing_info.append(f"Tool: {format_time(tool_time)}") + + return timing_info, tool_time +# Helper function to create token info display +def _create_token_info_display(token_info=None): + """Create token information display text.""" + if not token_info: + return None + + from rich.text import Text + + model = token_info.get('model', '') + interaction_input_tokens = token_info.get('interaction_input_tokens', 0) + interaction_output_tokens = token_info.get('interaction_output_tokens', 0) + interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) + total_input_tokens = token_info.get('total_input_tokens', 0) + total_output_tokens = token_info.get('total_output_tokens', 0) + total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) + + # Only continue if we have actual token information + if not (interaction_input_tokens > 0 or total_input_tokens > 0): + return None + + # Create token display + return _create_token_display( + interaction_input_tokens, + interaction_output_tokens, + interaction_reasoning_tokens, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + model, + token_info.get('interaction_cost'), + token_info.get('total_cost') + ) + +# Helper function for simple tool output without Rich +def _print_simple_tool_output(tool_name, args, output, execution_info=None, token_info=None): + """Print tool output without Rich formatting.""" + # Format arguments + args_str = _format_tool_args(args) + + # Get tool execution time if available + tool_time_str = "" + execution_status = "" + if execution_info: + time_taken = execution_info.get('time_taken', 0) or execution_info.get('tool_time', 0) + status = execution_info.get('status', 'completed') - # Create timing display string - timing_info = [] - if total_time: - timing_info.append(f"Total: {format_time(total_time)}") - if tool_time: - timing_info.append(f"Tool: {format_time(tool_time)}") - if timing_info: - header.append(f" [{' | '.join(timing_info)}]", style="cyan") + # Add execution info to the tool call display + if time_taken: + tool_time_str = f"Tool: {format_time(time_taken)}" + execution_status = f" [{status} in {time_taken:.2f}s]" + else: + execution_status = f" [{status}]" + + # Create timing display string + timing_info, _ = _get_timing_info(execution_info) + timing_display = f" [{' | '.join(timing_info)}]" if timing_info else "" + + # Show tool name, args, execution status and timing display + tool_call = f"{tool_name}({args_str})" + print(color(f"Tool Output: {tool_call}{timing_display}{execution_status}", fg="blue")) + + # If we have token info, display it + if token_info: + model = token_info.get('model', '') + interaction_input_tokens = token_info.get('interaction_input_tokens', 0) + interaction_output_tokens = token_info.get('interaction_output_tokens', 0) + interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) + total_input_tokens = token_info.get('total_input_tokens', 0) + total_output_tokens = token_info.get('total_output_tokens', 0) + total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) + + # If we have complete token information, display it + if (interaction_input_tokens > 0 or total_input_tokens > 0): + # Manually create formatted output similar to _create_token_display + print(color(f" Current: I:{interaction_input_tokens} O:{interaction_output_tokens} R:{interaction_reasoning_tokens}", fg="cyan")) - timing_display = f" [{' | '.join(timing_info)}]" if timing_info else "" - - # Show tool name, args, execution status and timing display - print(color(f"Tool Output: {tool_call}{timing_display}{execution_status}", fg="blue")) - - # If we have token info, display it using the consistent format from _create_token_display - if token_info: - model = token_info.get('model', '') - interaction_input_tokens = token_info.get('interaction_input_tokens', 0) - interaction_output_tokens = token_info.get('interaction_output_tokens', 0) - interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) - total_input_tokens = token_info.get('total_input_tokens', 0) - total_output_tokens = token_info.get('total_output_tokens', 0) - total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) - interaction_cost = token_info.get('interaction_cost') - total_cost = token_info.get('total_cost') + # Calculate or use provided costs + current_cost = COST_TRACKER.process_interaction_cost( + model, + interaction_input_tokens, + interaction_output_tokens, + interaction_reasoning_tokens, + token_info.get('interaction_cost') + ) + total_cost_value = COST_TRACKER.process_total_cost( + model, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + token_info.get('total_cost') + ) + print(color(f" Cost: Current ${current_cost:.4f} | Total ${total_cost_value:.4f} | Session ${COST_TRACKER.session_total_cost:.4f}", fg="cyan")) - # If we have complete token information, display it - if (interaction_input_tokens > 0 or total_input_tokens > 0): - # Manually create formatted output similar to _create_token_display - print(color(f" Current: I:{interaction_input_tokens} O:{interaction_output_tokens} R:{interaction_reasoning_tokens}", fg="cyan")) - - # Calculate or use provided costs - current_cost = COST_TRACKER.process_interaction_cost( - model, - interaction_input_tokens, - interaction_output_tokens, - interaction_reasoning_tokens, - interaction_cost - ) - total_cost_value = COST_TRACKER.process_total_cost( - model, - total_input_tokens, - total_output_tokens, - total_reasoning_tokens, - total_cost - ) - print(color(f" Cost: Current ${current_cost:.4f} | Total ${total_cost_value:.4f} | Session ${COST_TRACKER.session_total_cost:.4f}", fg="cyan")) - - # Show context usage - context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100 - indicator = "🟩" if context_pct < 50 else "🟨" if context_pct < 80 else "🟥" - print(color(f" Context: {context_pct:.1f}% {indicator}", fg="cyan")) + # Show context usage + context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100 + indicator = "🟩" if context_pct < 50 else "🟨" if context_pct < 80 else "🟥" + print(color(f" Context: {context_pct:.1f}% {indicator}", fg="cyan")) + + # Print the actual output + print(output) + print() + +# Add a new function to start a streaming tool execution +def start_tool_streaming(tool_name, args, call_id=None): + """ + Start a streaming tool execution session. + This allows for progressive updates during tool execution. + + Args: + tool_name: Name of the tool being executed + args: Arguments to the tool (dictionary or string) + call_id: Optional call ID for this execution. If not provided, one will be generated. - # Print the actual output - print(output) - print() \ No newline at end of file + Returns: + call_id: The call ID for this streaming session (can be used for updates) + """ + import time + import uuid + + # Generate a command key to check for duplicates - match format used in cli_print_tool_output + if isinstance(args, dict): + cmd = args.get("command", "") + cmd_args = args.get("args", "") + command_key = f"{tool_name}:{cmd_args}" + else: + command_key = f"{tool_name}:{args}" + + # Check if we've already seen this exact command recently + if not hasattr(start_tool_streaming, '_recent_commands'): + start_tool_streaming._recent_commands = {} + + # If we have an existing active streaming session for this command, reuse its call_id + # This prevents duplicate panels when the same command runs multiple times + for existing_call_id, info in list(start_tool_streaming._recent_commands.items()): + # Only consider recent commands (last 10 seconds) + timestamp = info.get('timestamp', 0) + if time.time() - timestamp < 10.0: + existing_command_key = info.get('command_key', '') + # Get the existing session info if available + if (hasattr(cli_print_tool_output, '_streaming_sessions') and + existing_call_id in cli_print_tool_output._streaming_sessions): + session = cli_print_tool_output._streaming_sessions[existing_call_id] + # If this is the same command and not complete, reuse the call_id + if existing_command_key == command_key and not session.get('is_complete', False): + return existing_call_id + + # Generate a call_id if not provided + if not call_id: + cmd_part = "" + if isinstance(args, dict) and "command" in args: + cmd_part = f"{args['command']}_" + call_id = f"cmd_{cmd_part}{str(uuid.uuid4())[:8]}" + + # Track this call_id with command key for better duplicate detection + start_tool_streaming._recent_commands[call_id] = { + 'timestamp': time.time(), + 'command_key': command_key + } + + # Cleanup old entries to prevent memory growth + current_time = time.time() + start_tool_streaming._recent_commands = { + k: v for k, v in start_tool_streaming._recent_commands.items() + if current_time - v.get('timestamp', 0) < 30 # Keep entries from last 30 seconds + } + + # Show initial message with "Starting..." output + cli_print_tool_output( + tool_name=tool_name, + args=args, + output="Starting tool execution...", + call_id=call_id, + execution_info={"status": "running", "start_time": time.time()}, + streaming=True + ) + + return call_id + +# Add a function to update a streaming tool execution +def update_tool_streaming(tool_name, args, output, call_id): + """ + Update a streaming tool execution with new output. + + Args: + tool_name: Name of the tool being executed + args: Arguments to the tool (dictionary or string) + output: New output to display + call_id: The call ID for this streaming session + + Returns: + None + """ + # Update the streaming output + cli_print_tool_output( + tool_name=tool_name, + args=args, + output=output, + call_id=call_id, + execution_info={"status": "running", "replace_buffer": True}, + streaming=True + ) + +# Add a function to complete a streaming tool execution +def finish_tool_streaming(tool_name, args, output, call_id, execution_info=None, token_info=None): + """ + Complete a streaming tool execution. + + Args: + tool_name: Name of the tool being executed + args: Arguments to the tool (dictionary or string) + output: Final output to display + call_id: The call ID for this streaming session + execution_info: Optional execution information + token_info: Optional token information + + Returns: + None + """ + import time + + # Prepare execution info with completion status + if execution_info is None: + execution_info = {} + + # Add completion markers + execution_info["status"] = execution_info.get("status", "completed") + execution_info["is_final"] = True + execution_info["replace_buffer"] = True + + # Calculate execution time if start_time is in the streaming session + if hasattr(cli_print_tool_output, '_streaming_sessions') and call_id in cli_print_tool_output._streaming_sessions: + session = cli_print_tool_output._streaming_sessions[call_id] + if 'start_time' in session and 'tool_time' not in execution_info: + execution_info["tool_time"] = time.time() - session['start_time'] + + # Show the final output + cli_print_tool_output( + tool_name=tool_name, + args=args, + output=output, + call_id=call_id, + execution_info=execution_info, + token_info=token_info, + streaming=True + ) + + # Mark the streaming session as complete + if hasattr(cli_print_tool_output, '_streaming_sessions') and call_id in cli_print_tool_output._streaming_sessions: + cli_print_tool_output._streaming_sessions[call_id]['is_complete'] = True + +def print_message_history(messages, title="Message History"): + """ + Pretty-print a sequence of messages with enhanced debug information. + + Args: + messages (List[dict]): List of message dictionaries to display + title (str, optional): Title to display above the message history + """ + from rich.console import Console + from rich.panel import Panel + from rich.text import Text + from rich.table import Table + + console = Console() + + # Create a table for displaying messages + table = Table(show_header=True, header_style="bold magenta", expand=True) + table.add_column("#", style="dim", width=3) + table.add_column("Role", style="cyan", width=10) + table.add_column("Content", width=40) + table.add_column("Metadata", width=30) + + # Process each message + for i, msg in enumerate(messages): + # Get role with color based on type + role = msg.get("role", "unknown") + role_style = { + "user": "green", + "assistant": "blue", + "system": "yellow", + "tool": "magenta" + }.get(role, "white") + + # Get content preview + content = msg.get("content") + content_preview = "" + if content is None: + content_preview = "[dim]None[/dim]" + elif isinstance(content, str): + # Truncate and escape long content + content_preview = (content[:37] + "...") if len(content) > 40 else content + content_preview = content_preview.replace("\n", "\\n") + elif isinstance(content, list): + content_preview = f"[list with {len(content)} items]" + else: + content_preview = f"[{type(content).__name__}]" + + # Gather metadata + metadata = [] + if msg.get("tool_calls"): + tc_count = len(msg["tool_calls"]) + tc_info = [] + for tc in msg["tool_calls"]: + tc_id = tc.get("id", "unknown") + tc_name = tc.get("function", {}).get("name", "unknown") if "function" in tc else "unknown" + tc_info.append(f"{tc_name}({tc_id})") + metadata.append(f"tool_calls[{tc_count}]: {', '.join(tc_info)}") + + if msg.get("tool_call_id"): + metadata.append(f"tool_call_id: {msg['tool_call_id']}") + + metadata_str = ", ".join(metadata) + + # Add row to table + table.add_row( + str(i), + f"[{role_style}]{role}[/{role_style}]", + content_preview, + metadata_str + ) + + # Create the panel with the table + panel = Panel( + table, + title=f"[bold]{title}[/bold]", + expand=False + ) + + # Display the panel + console.print(panel) + + return len(messages) # Return message count for convenience + +def get_language_from_code_block(lang_identifier): + """ + Maps a language identifier from a markdown code block to a proper syntax + highlighting language name. Handles common aliases and defaults. + + Args: + lang_identifier (str): Language identifier from markdown code block + + Returns: + str: Proper language name for syntax highlighting + """ + # Convert to lowercase and strip whitespace + lang = lang_identifier.lower().strip() if lang_identifier else "" + + # Map common language aliases to their proper names + lang_map = { + # Empty strings or unknown + "": "text", + # Python variants + "py": "python", + "python3": "python", + # JavaScript variants + "js": "javascript", + "jsx": "jsx", + "ts": "typescript", + "tsx": "tsx", + "typescript": "typescript", + # Shell variants + "sh": "bash", + "shell": "bash", + "console": "bash", + "terminal": "bash", + # Web languages + "html": "html", + "css": "css", + "json": "json", + "xml": "xml", + "yml": "yaml", + "yaml": "yaml", + # C family + "c": "c", + "cpp": "cpp", + "c++": "cpp", + "csharp": "csharp", + "cs": "csharp", + "java": "java", + # Other common languages + "go": "go", + "golang": "go", + "ruby": "ruby", + "rb": "ruby", + "rust": "rust", + "php": "php", + "sql": "sql", + "diff": "diff", + "markdown": "markdown", + "md": "markdown", + # Default fallback + "text": "text", + "plaintext": "text", + "txt": "text", + } + + # Return mapped language or default to the original if not in map + return lang_map.get(lang, lang or "text")