From d6c620f21499cb86828f0cfde7ec5af398164a8a Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 25 May 2025 12:02:20 +0000 Subject: [PATCH] Improve async --- src/cai/cli.py | 10 +- .../agents/models/openai_chatcompletions.py | 91 ++++++- src/cai/tools/common.py | 226 ++++++++++++------ .../reconnaissance/generic_linux_command.py | 166 +++++++------ src/cai/util.py | 26 ++ 5 files changed, 362 insertions(+), 157 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index b2965d35..60491ba7 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -248,12 +248,12 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), return "Agent" # Prevent the model from using its own rich streaming to avoid conflicts - # and suppress final output message to avoid duplicates + # but allow final output message to ensure all agent responses are shown if hasattr(agent, 'model'): if hasattr(agent.model, 'disable_rich_streaming'): agent.model.disable_rich_streaming = False # Now True as the model handles streaming if hasattr(agent.model, 'suppress_final_output'): - agent.model.suppress_final_output = True + agent.model.suppress_final_output = False # Changed to False to show all agent messages # Set the agent name in the model for proper display in streaming panel if hasattr(agent.model, 'set_agent_name'): @@ -320,7 +320,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), if hasattr(agent.model, 'disable_rich_streaming'): agent.model.disable_rich_streaming = False # Now False to let model handle streaming if hasattr(agent.model, 'suppress_final_output'): - agent.model.suppress_final_output = True + agent.model.suppress_final_output = False # Changed to False to show all agent messages # Apply current model to the new agent if hasattr(agent.model, 'model'): @@ -873,9 +873,9 @@ def main(): # Disable rich streaming in the model to avoid conflicts if hasattr(agent.model, 'disable_rich_streaming'): agent.model.disable_rich_streaming = True - # Suppress final output to avoid duplicates + # Allow final output to ensure all agent messages are shown if hasattr(agent.model, 'suppress_final_output'): - agent.model.suppress_final_output = True + agent.model.suppress_final_output = False # Changed to False to show all agent messages # Run the CLI with the selected agent run_cai_cli(agent) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 777d590b..19ea75a3 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -506,6 +506,7 @@ class OpenAIChatCompletionsModel(Model): if _debug.DONT_LOG_MODEL_DATA: logger.debug("Received model response") else: + import json logger.debug( f"LLM resp:\n{json.dumps(response.choices[0].message.model_dump(), indent=2)}\n" ) @@ -552,16 +553,86 @@ class OpenAIChatCompletionsModel(Model): for tool_call in response.choices[0].message.tool_calls: call_id = tool_call.id - # If we're using direct tool output display with cli_print_tool_output, - # and we've already displayed this tool call output, we can skip displaying - # the assistant message to avoid duplication - if (hasattr(_Converter, 'tool_outputs') and call_id in _Converter.tool_outputs and - hasattr(_Converter, 'recent_tool_calls') and call_id in _Converter.recent_tool_calls): - # We've already displayed this tool and its output directly - should_display_message = False + # Check if this tool call has already been displayed + if (hasattr(_Converter, 'tool_outputs') and call_id in _Converter.tool_outputs): + tool_output_content = _Converter.tool_outputs[call_id] + + # Check if this is a command sent to an existing async session + is_async_session_input = False + has_auto_output = False + is_regular_command = False + try: + import json + args = json.loads(tool_call.function.arguments) + # Check if this is a regular command (not a session command) + if (isinstance(args, dict) and + args.get("command") and + not args.get("session_id") and + not args.get("async_mode")): + is_regular_command = True + # Only consider it an async session input if it has session_id AND it's not creating a new session + elif (isinstance(args, dict) and + args.get("session_id") and + not args.get("async_mode") and # Not creating a new session + not args.get("creating_session")): # Not marked as session creation + is_async_session_input = True + # Check if this has auto_output flag + has_auto_output = args.get("auto_output", False) + except: + pass + + # For regular commands that were already shown via streaming, suppress the agent message + if is_regular_command and tool_call.function.name == "generic_linux_command": + # Check if this was executed very recently (likely shown via streaming) + if (hasattr(_Converter, 'recent_tool_calls') and + call_id in _Converter.recent_tool_calls): + tool_call_info = _Converter.recent_tool_calls[call_id] + if 'start_time' in tool_call_info: + import time + time_since_execution = time.time() - tool_call_info['start_time'] + # If executed within last 2 seconds, it was likely shown via streaming + if time_since_execution < 2.0: + should_display_message = False + tool_output = None + # For async session inputs with auto_output, suppress the agent message + elif is_async_session_input and has_auto_output: + should_display_message = False + tool_output = None + # For async session inputs without auto_output, always show the agent message + elif is_async_session_input and not has_auto_output: + should_display_message = True + tool_output = None + # For session creation messages, also show them + elif ("Started async session" in tool_output_content or + "session" in tool_output_content.lower() and "async" in tool_output_content.lower()): + should_display_message = True + tool_output = None + else: + # For other tool calls, check if we should suppress based on timing + # Only suppress if this tool was JUST executed (within last 2 seconds) + if (hasattr(_Converter, 'recent_tool_calls') and + call_id in _Converter.recent_tool_calls): + tool_call_info = _Converter.recent_tool_calls[call_id] + if 'start_time' in tool_call_info: + import time + time_since_execution = time.time() - tool_call_info['start_time'] + # Only suppress if this was executed very recently + if time_since_execution < 2.0: + should_display_message = False + else: + # For older tool calls, show the message + should_display_message = True break + + # Additional check: Always show messages that have text content + # This ensures agent explanations are not suppressed + if (hasattr(response.choices[0].message, 'content') and + response.choices[0].message.content and + str(response.choices[0].message.content).strip()): + # If the message has actual text content, always show it + should_display_message = True - # Only display the agent message if we haven't already shown the tool output + # Display the agent message (this will show the command for async sessions) if should_display_message: # Ensure we're in non-streaming mode for proper markdown parsing previous_stream_setting = os.environ.get('CAI_STREAM', 'false') @@ -588,8 +659,8 @@ class OpenAIChatCompletionsModel(Model): total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), 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 + tool_output=tool_output, # Pass tool_output only when needed + suppress_empty=True # Keep suppress_empty=True as requested ) # Restore previous streaming setting diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 3d5eaa51..74d5dc3f 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -261,6 +261,21 @@ class ShellSession: # pylint: disable=too-many-instance-attributes if clear: self.output_buffer = [] return output + + def get_new_output(self, mark_position=True): + """Get only new output since last marked position""" + if not hasattr(self, '_last_output_position'): + self._last_output_position = 0 + + # Get new output since last position + new_output_lines = self.output_buffer[self._last_output_position:] + new_output = "\n".join(new_output_lines) + + # Update position marker if requested + if mark_position: + self._last_output_position = len(self.output_buffer) + + return new_output def terminate(self): """Terminate the session""" @@ -392,7 +407,7 @@ def terminate_session(session_id): return result -def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None): +def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None, stream=False): """Runs command in CTF env, changing to workspace_dir first.""" target_dir = workspace_dir or _get_workspace_dir() full_command = f"{command}" @@ -400,7 +415,9 @@ def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None): context_msg = f"(ctf:{target_dir})" try: output = ctf.get_shell(full_command, timeout=timeout) - if stdout: + # In streaming mode, don't print to stdout to avoid duplication + # The streaming system will handle the display + if stdout and not stream: print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 return output except Exception as e: # pylint: disable=broad-except @@ -408,7 +425,7 @@ def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None): print(color(error_msg, fg="red")) return error_msg -def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None): +def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None, stream=False): """Runs command via SSH. Assumes SSH agent or passwordless setup unless sshpass is used externally.""" # noqa E501 ssh_user = os.environ.get('SSH_USER') ssh_host = os.environ.get('SSH_HOST') @@ -434,14 +451,16 @@ def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None): timeout=timeout ) output = result.stdout if result.stdout else result.stderr - if stdout: + # In streaming mode, don't print to stdout to avoid duplication + # The streaming system will handle the display + if stdout and not stream: print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 # Return combined output, potentially including errors return output.strip() except subprocess.TimeoutExpired as e: error_output = e.stdout if e.stdout else str(e) timeout_msg = f"Timeout executing SSH command: {error_output}" - if stdout: + if stdout and not stream: print(f"\033[33m{context_msg} $ {original_cmd_for_msg}\nTIMEOUT\n{error_output}\033[0m") # noqa E501 return timeout_msg except FileNotFoundError: @@ -586,55 +605,9 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t ) 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 - ) - + # In non-streaming mode, we should NOT display the output via cli_print_tool_output + # to avoid duplication. The output will be handled by the calling function. + # Only return the raw output for the calling function to handle. return output.strip() except subprocess.TimeoutExpired as e: error_output = e.stdout if hasattr(e, 'stdout') and e.stdout else str(e) @@ -762,6 +735,36 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg session = ACTIVE_SESSIONS[session_id] result = session.send_input(command) # Send the raw command string + # Wait for the command to execute and capture output + # This provides automatic output display for async sessions + wait_time = 3.0 # Wait 3 seconds for command to execute + + # Mark the current position in the output buffer before sending input + session.get_new_output(mark_position=True) # Reset position marker + + # Smart waiting: check for new output every 0.5 seconds, up to max wait time + max_wait = wait_time + check_interval = 0.5 + elapsed = 0.0 + new_output_detected = False + + while elapsed < max_wait: + time.sleep(check_interval) + elapsed += check_interval + + # Check if new output is available + current_new_output = session.get_new_output(mark_position=False) + + # If we detect new output, wait a bit more for it to complete, then break + if current_new_output.strip(): + if not new_output_detected: + new_output_detected = True + # Give it a bit more time to complete the output + time.sleep(0.5) + else: + # We already detected new output and waited, now break + break + # Always show the session output after sending input using the counter mechanism # Generate unique counter for this session input command counter_key = f"session_input_{session_id}" @@ -775,7 +778,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg "args": "", "session_id": session_id, "call_counter": SESSION_OUTPUT_COUNTER[counter_key], # This ensures uniqueness - "input_to_session": True # Flag to identify this as session input + "input_to_session": True, # Flag to identify this as session input + "auto_output": True # Flag to indicate this is automatic output display } # Determine environment info for display @@ -785,15 +789,17 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg elif session.ctf: env_type = "CTF" - # Get the session output to display - output = get_session_output(session_id, clear=False, stdout=False) + # Get only the NEW output to display (not the entire buffer) + output = session.get_new_output(mark_position=True) # Create execution info execution_info = { "status": "completed", "environment": env_type, "host": session.workspace_dir, - "session_id": session_id + "session_id": session_id, + "wait_time": elapsed, + "new_output_detected": new_output_detected } # Display the session input and its result using cli_print_tool_output @@ -813,7 +819,12 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_active_timer() start_idle_timer() - return result # Return the result of sending input ("Input sent..." or error) + # Return the actual output from the session + # The output has already been displayed via cli_print_tool_output + if output and output.strip(): + return output + else: + return f"Command sent to session {session_id}. No output captured." # 2. Determine Execution Environment (Container > CTF > SSH > Local) active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") @@ -835,11 +846,47 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_active_timer() start_idle_timer() return new_session_id - if stdout: - # Wait a moment for initial output - time.sleep(0.2) - output = get_session_output(new_session_id, clear=False, stdout=False) - print(f"\033[32m(Started Session {new_session_id} in {context_msg})\n{output}\033[0m") # noqa E501 + + # Display the command that creates the async session + from cai.util import cli_print_tool_output + + # Create args for display + session_creation_args = { + "command": command, + "args": "", + "session_id": new_session_id, + "async_mode": True, + "creating_session": True # Flag to identify this as session creation + } + + # Create execution info + execution_info = { + "status": "session_created", + "environment": f"Container({container_id[:12]})", + "host": container_workspace, + "session_id": new_session_id + } + + # Get initial output if any + session = ACTIVE_SESSIONS.get(new_session_id) + initial_output = "" + if session: + time.sleep(0.2) # Wait a moment for initial output + initial_output = session.get_new_output(mark_position=True) + + # Format the output message + output_msg = f"Started async session {new_session_id} in container {container_id[:12]}. Use this ID to interact." + if initial_output: + output_msg += f"\n\n{initial_output}" + + # Display the session creation command and initial output + cli_print_tool_output( + tool_name="generic_linux_command", + args=session_creation_args, + output=output_msg, + execution_info=execution_info, + streaming=False + ) # For async sessions, switch back to idle mode after session creation stop_active_timer() @@ -1038,7 +1085,9 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg output = result.stdout if result.stdout else result.stderr output = output.strip() # Clean trailing newline - if stdout: + # In streaming mode, don't print to stdout to avoid duplication + # The streaming system will handle the display + if stdout and not stream: print(f"\033[32m{context_msg} $ {command}\n{output}\033[0m") # noqa E501 # Check if command failed specifically because container isn't running @@ -1149,7 +1198,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg return error_msg else: # Standard non-streaming CTF execution - result = _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir + result = _run_ctf(ctf, command, stdout, timeout, _get_workspace_dir(), stream) # Switch back to idle mode after CTF command completes stop_active_timer() @@ -1278,7 +1327,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg return error_msg else: # Standard non-streaming SSH execution - result = _run_ssh(command, stdout, timeout) # Workspace dir less relevant here + result = _run_ssh(command, stdout, timeout, _get_workspace_dir(), stream) # Switch back to idle mode after SSH command completes stop_active_timer() @@ -1297,13 +1346,50 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_active_timer() start_idle_timer() return new_session_id + + # Display the command that creates the async session + from cai.util import cli_print_tool_output + # Retrieve the actual workspace dir the session is using session = ACTIVE_SESSIONS.get(new_session_id) actual_workspace = session.workspace_dir if session else "unknown" - if stdout: + + # Create args for display + session_creation_args = { + "command": command, + "args": "", + "session_id": new_session_id, + "async_mode": True, + "creating_session": True # Flag to identify this as session creation + } + + # Create execution info + execution_info = { + "status": "session_created", + "environment": "Local", + "host": os.path.basename(actual_workspace), + "session_id": new_session_id + } + + # Get initial output if any + initial_output = "" + if session: time.sleep(0.2) # Allow session buffer to populate - output = get_session_output(new_session_id, clear=False, stdout=False) - print(f"\033[32m(Started Session {new_session_id} in local:{actual_workspace})\n{output}\033[0m") + initial_output = session.get_new_output(mark_position=True) + + # Format the output message + output_msg = f"Started async session {new_session_id} locally. Use this ID to interact." + if initial_output: + output_msg += f"\n\n{initial_output}" + + # Display the session creation command and initial output + cli_print_tool_output( + tool_name="generic_linux_command", + args=session_creation_args, + output=output_msg, + execution_info=execution_info, + streaming=False + ) # For async, switch back to idle mode after session creation stop_active_timer() diff --git a/src/cai/tools/reconnaissance/generic_linux_command.py b/src/cai/tools/reconnaissance/generic_linux_command.py index f6beaede..ef6eabdb 100644 --- a/src/cai/tools/reconnaissance/generic_linux_command.py +++ b/src/cai/tools/reconnaissance/generic_linux_command.py @@ -16,40 +16,53 @@ from wasabi import color # pylint: disable=import-error @function_tool def generic_linux_command(command: str = "", - args: str = "", ctf=None, - async_mode: bool = False, + interactive: bool = False, session_id: str = None) -> str: """ - Execute Linux commands with support for interactive sessions. + Execute commands with session management. - Two modes: - 1. Regular command execution: - - For normal Linux commands - - Example: generic_linux_command("ls", "-la") - - 2. Session management: - - For interactive commands(ssh) - - Workflow: - - Start: generic_linux_command("ssh", "user@host") - → Returns session ID - - List: generic_linux_command("session", "list") - - Get output: generic_linux_command("session", "output ") - - Send input: Use session_id parameter - - End: generic_linux_command("session", "kill ") + Use this tool to run any command. The system automatically detects and handles: + - Regular commands (ls, cat, grep, etc.) + - Interactive commands that need persistent sessions (ssh, nc, python, etc.) + - Session management and output capture + - CTF environments (automatically detected and used when available) + - Container environments (automatically detected and used when available) + - SSH environments (automatically detected and used when available) Args: - command: Command name - args: Command arguments - ctf: CTF environment object - async_mode: Force async session - session_id: Existing session ID + command: The complete command to execute (e.g., "ls -la", "ssh user@host", "cat file.txt") + interactive: Set to True for commands that need persistent sessions (ssh, nc, python, ftp etc.) + Leave False for regular commands + session_id: Use existing session ID to send commands to running interactive sessions. + Get session IDs from previous interactive command outputs. + + Examples: + - Regular command: generic_linux_command("ls -la") + - Interactive command: generic_linux_command("ssh user@host", interactive=True) + - Send to session: generic_linux_command("pwd", session_id="abc12345") + - List sessions: generic_linux_command("session list") + - Kill session: generic_linux_command("session kill abc12345") + - Environment info: generic_linux_command("env info") + + Environment Detection: + The system automatically detects and uses the appropriate execution environment: + - CTF: Commands run in the CTF challenge environment when available + - Container: Commands run in Docker containers when CAI_ACTIVE_CONTAINER is set + - SSH: Commands run via SSH when SSH_USER and SSH_HOST are configured + - Local: Commands run on the local system as fallback Returns: - Command output, session ID, or status message + Command output, session ID for interactive commands, or status message """ - # Special commands for session management - if command == "session": - if args == "list": + # Handle special session management commands + if command.startswith("session "): + parts = command.split(" ", 2) + if len(parts) < 2: + return "Usage: session list|output|kill [session_id]" + + action = parts[1] + + if action == "list": sessions = list_shell_sessions() if not sessions: return "No active sessions" @@ -61,66 +74,75 @@ def generic_linux_command(command: str = "", f"Last activity: {session['last_activity']}\n") return result - if args.startswith("output "): - session_id = args.split(" ")[1] - # Call get_session_output with stdout=True to display via cli_print_tool_output - # The function will handle the display and return a simple confirmation message - output = get_session_output(session_id, clear=False, stdout=True) + elif action == "output" and len(parts) >= 3: + target_session_id = parts[2] + output = get_session_output(target_session_id, clear=False, stdout=True) return output - if args.startswith("kill "): - session_id = args.split(" ")[1] - return terminate_session(session_id) + elif action == "kill" and len(parts) >= 3: + target_session_id = parts[2] + return terminate_session(target_session_id) - return """Unknown session command. - Available: list, output , kill """ + return "Usage: session list|output |kill " - # BUGFIX: Detect if arguments are swapped due to framework issue - # Sometimes async_mode gets the session_id value and session_id gets boolean values - if (isinstance(async_mode, str) and len(async_mode) == 8 and - session_id in [None, True, False]): - # Arguments are swapped - fix them - actual_session_id = async_mode - actual_async_mode = False # Default to False when session_id is provided - session_id = actual_session_id - async_mode = actual_async_mode + # Handle environment information command + if command.strip() == "env info" or command.strip() == "environment info": + env_info = [] + + # Check CTF environment + try: + from cai.cli import ctf_global + if ctf_global and hasattr(ctf_global, 'get_shell'): + env_info.append("🎯 CTF Environment: Active") + else: + env_info.append("🎯 CTF Environment: Not available") + except: + env_info.append("🎯 CTF Environment: Not available") + + # Check Container environment + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + if active_container: + env_info.append(f"🐳 Container: {active_container[:12]}") + else: + env_info.append("🐳 Container: Not active") + + # Check SSH environment + ssh_user = os.getenv('SSH_USER') + ssh_host = os.getenv('SSH_HOST') + if ssh_user and ssh_host: + env_info.append(f"🔗 SSH: {ssh_user}@{ssh_host}") + else: + env_info.append("🔗 SSH: Not configured") + + # Check workspace + try: + from cai.tools.common import _get_workspace_dir + workspace = _get_workspace_dir() + env_info.append(f"📁 Workspace: {workspace}") + except: + env_info.append("📁 Workspace: Unknown") + + return "Current Environment:\n" + "\n".join(env_info) - # Regular command execution - full_command = f'{command} {args}'.strip() + if not command.strip(): + return "Error: No command provided" - # Detect if this should be an async command - # Only auto-detect async commands when no session_id is provided - if not async_mode and not session_id: - async_commands = ['ssh', 'python -m http.server', "ftp"] - async_mode = any(cmd in full_command for cmd in async_commands) - - # For SSH sessions or async commands, use different timeout + # For SSH sessions or interactive 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' + # Command streaming should be independent of LLM streaming + stream = True # Always enable streaming for commands - # Generate a call_id for streaming if needed - # Only use call_id for streaming mode to prevent duplicate panels - # When CAI_STREAM=true, we want the Tool Execution panel but not the Tool Output panel - call_id = str(uuid.uuid4())[:8] if stream else None - - # In streaming mode, prevent displaying both panels by forcing a call_id - # This tricks cli_print_tool_output into thinking it's being called for a streaming update - if stream and not call_id: - call_id = str(uuid.uuid4())[:8] + # Generate a call_id for streaming + call_id = str(uuid.uuid4())[:8] - - - # Run the command with the appropriate parameters - pass the actual tool name! - result = run_command(full_command, ctf=ctf, - async_mode=async_mode, session_id=session_id, + # Run the command with the appropriate parameters + result = run_command(command, ctf=None, + async_mode=interactive, session_id=session_id, timeout=timeout, stream=stream, call_id=call_id, tool_name="generic_linux_command") - # For better output formatting, if we're in streaming mode, we can modify the output - # to include any additional information needed while still preventing the duplicate panel return result diff --git a/src/cai/util.py b/src/cai/util.py index 5b1a3ef7..a116f2b9 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -1343,6 +1343,18 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli parsed_message = parse_message_content(message) tool_panels = [] + # Special handling for async session messages + if tool_output and ("Started async session" in tool_output or "session" in tool_output.lower()): + # For async session creation, show the session message as the main content + if not parsed_message or parsed_message == "null" or parsed_message == "": + parsed_message = tool_output + else: + # If there's already content, append the session message + parsed_message = f"{parsed_message}\n\n{tool_output}" + + # Clear tool_panels to avoid duplication since we're showing the session message as main content + 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 @@ -1875,6 +1887,10 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut if isinstance(args, dict): # If args is a dictionary, extract the 'args' field. effective_command_args_str = args.get("args", "") + # For session commands, also include the actual command to make it unique + if "command" in args and args.get("session_id"): + # For async session commands, include the full command to differentiate + effective_command_args_str = f"{args.get('command', '')}:{effective_command_args_str}" elif isinstance(args, str): # If args is a string, it might be a JSON representation or a plain string. try: @@ -1882,6 +1898,9 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut if isinstance(parsed_json_args, dict): # Parsed as JSON dict, get the 'args' field. effective_command_args_str = parsed_json_args.get("args", "") + # For session commands, also include the actual command + if "command" in parsed_json_args and parsed_json_args.get("session_id"): + effective_command_args_str = f"{parsed_json_args.get('command', '')}:{effective_command_args_str}" 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 @@ -1897,6 +1916,13 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut call_counter = args["call_counter"] command_key += f":counter_{call_counter}" + # For async session inputs, add timestamp to ensure uniqueness + # This prevents duplicate detection for different commands sent to the same session + if isinstance(args, dict) and args.get("session_id") and args.get("input_to_session"): + # Add a timestamp component to make each session input unique + import time + command_key += f":ts_{int(time.time() * 1000)}" + # --- End of Command Key Generation --- # Check for duplicate display conditions