From 49dbebc76d4dcd4a7a834abfa1138d7e342d9062 Mon Sep 17 00:00:00 2001 From: lidia9 Date: Fri, 25 Apr 2025 10:33:12 +0200 Subject: [PATCH] FIX streaming output tool - still some visualization duplicates --- src/cai/cli.py | 4 +- src/cai/tools/common.py | 200 +++++++++++- .../reconnaissance/generic_linux_command.py | 18 +- src/cai/util.py | 305 ++++++++++++------ 4 files changed, 422 insertions(+), 105 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index c67ec1f1..1d2fed05 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -184,8 +184,8 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= # Function to get the short name of the agent for display def get_agent_short_name(agent): if hasattr(agent, 'name'): - # Split by spaces and take first word to get shortened name - return agent.name.split()[0] + # Return the full agent name instead of just the first word + return agent.name return "Agent" # Prevent the model from using its own rich streaming to avoid conflicts diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 91a9752a..ebf6821e 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -10,6 +10,7 @@ import pty import signal import time import uuid +import sys from wasabi import color # pylint: disable=import-error # Global dictionary to store active sessions @@ -211,7 +212,7 @@ def terminate_session(session_id): return result -def _run_ctf(ctf, command, stdout=False, timeout=100): +def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None): try: # Ensure the command is executed in a shell that supports command # chaining @@ -227,7 +228,11 @@ def _run_ctf(ctf, command, stdout=False, timeout=100): return f"Error executing CTF command: {str(e)}" -def _run_local(command, stdout=False, timeout=100): +def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None): + # If streaming is enabled and we have a call_id + if stream and call_id: + return _run_local_streamed(command, call_id, timeout) + try: # nosec B602 - shell=True is required for command chaining result = subprocess.run( @@ -252,9 +257,184 @@ def _run_local(command, stdout=False, timeout=100): return error_msg +def _run_local_streamed(command, call_id, timeout=100): + """Run a local command with streaming output to the Tool output panel""" + 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 the process + process = subprocess.Popen( + command, + shell=True, # nosec B602 + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1 + ) + + # Create panel content for Rich display + if rich_available: + tool_name = "generic_linux_command" + header = Text() + header.append(tool_name, style="#00BCD4") + header.append("(", style="yellow") + header.append(f"command='{command}'", style="yellow") + header.append(")", style="yellow") + + content = Text() + content.append(f"Executing: {command}\n\n", style="green") + + panel = Panel( + Text.assemble(header, "\n\n", content), + title="[bold blue]Tool Execution[/bold blue]", + subtitle="[bold green]Live Output[/bold green]", + border_style="blue", + padding=(1, 2), + box=ROUNDED + ) + + # Start Live display + with Live(panel, console=console, refresh_per_second=4) as live: + # Stream stdout in real-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") + panel = Panel( + Text.assemble(header, "\n\n", content), + title="[bold blue]Tool Execution[/bold blue]", + subtitle="[bold green]Live Output[/bold green]", + border_style="blue", + 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 blue]Tool Execution[/bold blue]", + subtitle="[bold green]Live Output[/bold green]", + border_style="blue", + padding=(1, 2), + box=ROUNDED + ) + live.update(panel) + + # Add completion message + completion_status = "Completed" if return_code == 0 else f"Failed (code {return_code})" + content.append(f"\nCommand {completion_status}", style="green") + panel = Panel( + Text.assemble(header, "\n\n", content), + title="[bold blue]Tool Execution[/bold blue]", + subtitle=f"[bold green]{completion_status}[/bold green]", + border_style="blue", + 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 + tool_args = {"command": command} + + # Initial notification - just once + cli_print_tool_output("generic_linux_command", tool_args, "Command started...", call_id=call_id) + + # Buffer for collecting output + 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 the output periodically to reduce panel refresh rate + if buffer_size >= update_interval: + current_output = ''.join(output_buffer) + cli_print_tool_output("generic_linux_command", tool_args, current_output, call_id=call_id) + buffer_size = 0 + + # 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: + output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) + + # Final output update - always show the final result + final_output = ''.join(output_buffer) + if return_code != 0: + final_output += f"\nCommand exited with code {return_code}" + + cli_print_tool_output("generic_linux_command", 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("generic_linux_command", tool_args, final_output, call_id=call_id) + + return final_output + + except Exception as e: # pylint: disable=broad-except + error_msg = f"Error executing command: {str(e)}" + 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("generic_linux_command", tool_args, error_msg, call_id=call_id) + + return error_msg + + def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arguments # noqa: E501 async_mode=False, session_id=None, - timeout=100): + timeout=100, stream=False, call_id=None): """ Run command either in CTF container or on the local attacker machine @@ -264,6 +444,9 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stdout: Whether to print output to stdout async_mode: Whether to run the command asynchronously session_id: ID of an existing session to send the command to + timeout: Command timeout in seconds + stream: Whether to stream output in real-time + call_id: Unique ID for the command execution (for streaming) Returns: str: Command output, status message, or session ID @@ -287,10 +470,13 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg time.sleep(0.5) output = get_session_output(session_id, clear=False) print("\033[32m" + output + "\033[0m") - return f"Created session { - session_id}. Use this ID to interact with the session." + return f"Created session {session_id}. Use this ID to interact with the session." + # 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] + # Otherwise, run command normally if ctf: - return _run_ctf(ctf, command, stdout, timeout) - return _run_local(command, stdout, timeout) + return _run_ctf(ctf, command, stdout, timeout, stream, call_id) + return _run_local(command, stdout, timeout, stream, call_id) diff --git a/src/cai/tools/reconnaissance/generic_linux_command.py b/src/cai/tools/reconnaissance/generic_linux_command.py index b14ba4ad..7fa156e7 100644 --- a/src/cai/tools/reconnaissance/generic_linux_command.py +++ b/src/cai/tools/reconnaissance/generic_linux_command.py @@ -1,11 +1,17 @@ """ This is used to create a generic linux command. """ +import os +import time +import uuid +import subprocess +import sys from cai.tools.common import (run_command, list_shell_sessions, get_session_output, terminate_session) # pylint: disable=import-error # noqa E501 from cai.sdk.agents import function_tool +from wasabi import color # pylint: disable=import-error @function_tool @@ -74,13 +80,19 @@ def generic_linux_command(command: str = "", async_commands = ['ssh', 'python -m http.server'] async_mode = any(cmd in full_command for cmd in async_commands) - # NOTE: review this as it's a hack to get - # around the long delays with nc connections + # For SSH sessions or async commands, use different timeout if session_id: timeout = 10 else: timeout = 100 + + # Check if streaming should be enabled + stream = os.getenv('CAI_STREAM', 'false').lower() == 'true' + + # Generate a call_id for streaming if needed + call_id = str(uuid.uuid4())[:8] if stream else None + # Run the command with the appropriate parameters return run_command(full_command, ctf=ctf, async_mode=async_mode, session_id=session_id, - timeout=timeout) + timeout=timeout, stream=stream, call_id=call_id) diff --git a/src/cai/util.py b/src/cai/util.py index ab7107db..6397a230 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -924,9 +924,8 @@ def create_agent_streaming_context(agent_name, counter, model): expand=True ) - # Start the live display - live = Live(panel, refresh_per_second=20, console=console) - live.start() + # 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) return { "live": live, @@ -937,7 +936,8 @@ def create_agent_streaming_context(agent_name, counter, model): "timestamp": timestamp, "model": model, "agent_name": agent_name, - "panel_width": panel_width + "panel_width": panel_width, + "is_started": False # Track if we've started the display } def update_agent_streaming_content(context, text_delta): @@ -945,6 +945,10 @@ def update_agent_streaming_content(context, text_delta): # Parse the text_delta to get just the content if needed parsed_delta = parse_message_content(text_delta) + # Skip empty updates to avoid showing an empty panel + if not parsed_delta or parsed_delta.strip() == "": + return + # Add the parsed text to the content context["content"].append(parsed_delta) @@ -960,12 +964,26 @@ def update_agent_streaming_content(context, text_delta): expand=True ) + # Check if we need to start the display + if not context.get("is_started", False): + context["live"].start() + context["is_started"] = True + # Force an update with the new panel context["live"].update(updated_panel) context["panel"] = updated_panel def finish_agent_streaming(context, final_stats=None): """Finish the streaming session and display final stats if available.""" + # Check if there's actual content to display - don't show empty panels + if not context["content"] or context["content"].plain == "": + # If the display was never started, nothing to do + if not context.get("is_started", False): + return + # Otherwise, stop the display without showing final panel + context["live"].stop() + return + # If we have token stats, add them tokens_text = None if final_stats: @@ -1037,103 +1055,204 @@ def finish_agent_streaming(context, final_stats=None): # Stop the live display context["live"].stop() -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): """ - Print tool execution and output in a single unified panel. + Print a tool call output to the command line. + Similar to cli_print_tool_call but for the output of the tool. Args: - tool_name: Name of the tool that was executed + tool_name: Name of the tool args: Arguments passed to the tool - output: Output from the tool execution - call_id: Optional ID of the tool call - execution_info: Dictionary with execution timing information - token_info: Dictionary with token usage information (ignored - no cost display in tool panels) + output: The output of the tool + call_id: Optional call ID for streaming updates + execution_info: Optional execution information + token_info: Optional token information """ - from rich.panel import Panel - from rich.text import Text - from rich.box import ROUNDED - from rich.console import Group - - # Track which tool calls we've already printed to avoid duplicates - if not hasattr(cli_print_tool_output, '_seen_calls'): - cli_print_tool_output._seen_calls = {} - - # If we have a call_id, check if we've already printed this output - # If no call_id, use a combination of tool_name and args as a key - call_key = call_id if call_id else f"{tool_name}:{args}" - - # Avoid printing duplicates - if call_key in cli_print_tool_output._seen_calls: + # If it's an empty output, don't print anything + if not output and not call_id: return - # Mark this call as seen to avoid duplicates - cli_print_tool_output._seen_calls[call_key] = True - - # Format the tool and arguments - if isinstance(args, dict): #key=value pairs - args_str = ", ".join(f"{k}={v}" for k, v in args.items()) - else: #single string - args_str = args - - # Create content for the tool execution panel - tool_text = Text() - tool_text.append(f"{tool_name}", style="#00BCD4") - - # Add the arguments - if isinstance(args, dict): - args_str = ", ".join(f"{k}={v}" for k, v in args.items()) - tool_text.append("(", style="yellow") - tool_text.append(args_str, style="yellow") - tool_text.append(")", style="yellow") - else: - tool_text.append("(", style="yellow") - tool_text.append(f"{args}", style="yellow") - tool_text.append(")", style="yellow") - - # Create the tool execution panel - tool_panel = Panel( - tool_text, - border_style="blue", - box=ROUNDED, - padding=(1, 2), - title="[bold]Tool Execution[/bold]", - title_align="left", - expand=True - ) - console.print(tool_panel) - - # Create content for the output panel - panel_content = [] - - # Add execution time if available - if execution_info: - total_time = execution_info.get('total_time', 0) - tool_time = execution_info.get('tool_time', 0) - if total_time > 0: - total_time_str = f"{int(total_time // 60)}m {total_time % 60:.1f}s" - tool_time_str = f"{tool_time:.1f}s" - execution_time = f"Total: {total_time_str} | Tool: {tool_time_str}" + # Track seen call IDs to prevent duplicate panels + 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: + return - exec_text = Text() - exec_text.append("[", style="dim") - exec_text.append(f"{execution_time}", style="magenta") - exec_text.append("]", style="dim") - panel_content.append(exec_text) + # Mark as seen + cli_print_tool_output._seen_calls[call_key] = True + + # 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:] + } - # Add the tool output - if output and output.strip(): - output_text = Text("\n" if panel_content else "") - output_text.append(output.strip(), style="#C0C0C0") - panel_content.append(output_text) - - if panel_content: - output_panel = Panel( - Group(*panel_content), - border_style="blue", - box=ROUNDED, + # Try to use Rich for better formatting if available + try: + from rich.console import Console + from rich.panel import Panel + from rich.text import Text + from rich.box import ROUNDED + + # Create a console for output + console = Console() + + # Format arguments as a string if they are a dictionary + if isinstance(args, dict): + args_str = ", ".join([f"{key}='{value}'" for key, value in args.items()]) + else: + args_str = str(args) + + # For streaming mode with call_id, use Rich Live display + if call_id: + # Create the header text + header = Text() + header.append(tool_name, style="#00BCD4") + header.append("(", style="yellow") + header.append(args_str, style="yellow") + header.append(")", style="yellow") + + # Create content text with the output + content = Text(output) + + # Create the panel for display + panel = Panel( + Text.assemble(header, "\n\n", content), + title=f"[bold blue]Tool Output (ID: {call_id})[/bold blue]", + subtitle="[bold green]Live Update[/bold green]", + border_style="blue", + padding=(1, 2), + box=ROUNDED + ) + + # Display using Rich + console.print(panel) + return + + # For non-streaming output, show a standard panel + tool_call = f"{tool_name}({args_str})" + + # Add execution info if available + subtitle = "[bold green]Completed[/bold green]" + if execution_info: + time_taken = execution_info.get('time_taken', 0) + status = execution_info.get('status', 'completed') + + if time_taken: + subtitle = f"[bold green]{status.title()} in {time_taken:.2f}s[/bold green]" + else: + subtitle = f"[bold green]{status.title()}[/bold green]" + + # Create content sections + sections = [] + + # Add token info if available + if token_info: + tokens_in = token_info.get('input_tokens', 0) + tokens_out = token_info.get('output_tokens', 0) + cost = token_info.get('cost', 0) + + token_text = Text() + if tokens_in or tokens_out: + token_text.append(f"Tokens: {tokens_in} in, {tokens_out} out", style="cyan") + if cost: + token_text.append(f"\nCost: ${cost:.6f}", style="cyan") + + if token_text: + sections.append(token_text) + + # Add the main output + if output: + output_text = Text() + if sections: # Add a separator if we have previous sections + output_text.append("\n") + output_text.append(output) + sections.append(output_text) + + # Create the panel with all sections + panel = Panel( + Text.assemble(*sections), + title=f"[bold blue]Tool Output: {tool_call}[/bold blue]", + subtitle=subtitle, + border_style="blue", padding=(1, 2), - title="[bold]Tool Output[/bold]", - title_align="left", - expand=True + box=ROUNDED ) - console.print(output_panel) \ No newline at end of file + + # Display the panel + console.print(panel) + + except ImportError: + # Fall back to simple formatting if Rich is not available + # Format arguments as a string if they are a dictionary + if isinstance(args, dict): + args_str = ", ".join([f"{key}='{value}'" for key, value in args.items()]) + else: + args_str = str(args) + + # Simplify output presentation for streaming mode + if call_id: + # This is a streaming update, so we need to overwrite previous output + # We'll use a basic format that's better suited for streaming + + # Get terminal width for better formatting + try: + term_width = os.get_terminal_size().columns + except: # pylint: disable=bare-except + term_width = 80 + + # Create a header for the tool output + header = f"{tool_name}({args_str})" + header = header[:term_width-4] + + # Clear the screen for the tool output (alternative approach) + # This works better with streamed content as it doesn't scroll the terminal + print(f"\r{header}") + + # Print the content + # For streaming updates, we'll use a simple format + print(output) + + # Force flush to ensure output is displayed immediately + sys.stdout.flush() + return + + # For non-streaming output, use the original formatting + tool_call = f"{tool_name}({args_str})" + + # If there's execution info, add it to the output + if execution_info: + time_taken = execution_info.get('time_taken', 0) + status = execution_info.get('status', 'completed') + + # Add execution info to the tool call display + if time_taken: + tool_call += f" [{status} in {time_taken:.2f}s]" + else: + tool_call += f" [{status}]" + + print(color(f"Tool Output: {tool_call}", fg="blue")) + + # If we have token info, display it + if token_info: + tokens_in = token_info.get('input_tokens', 0) + tokens_out = token_info.get('output_tokens', 0) + cost = token_info.get('cost', 0) + + if tokens_in or tokens_out: + print(color(f" Tokens: {tokens_in} in, {tokens_out} out", fg="cyan")) + if cost: + print(color(f" Cost: ${cost:.6f}", fg="cyan")) + + # Print the actual output + print(output) + print() \ No newline at end of file