From ac6d071954668a3784ef50831df0d828b2cf0d50 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 22 May 2025 20:49:30 +0200 Subject: [PATCH] Async sessions output fixed --- src/cai/tools/common.py | 131 +++++++++++++++--- .../reconnaissance/generic_linux_command.py | 20 ++- src/cai/util.py | 29 +++- 3 files changed, 159 insertions(+), 21 deletions(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 1ce0850b..85923a9f 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -25,6 +25,9 @@ except ImportError: # Global dictionary to store active sessions ACTIVE_SESSIONS = {} +# Global counter for session output commands to ensure they always display +SESSION_OUTPUT_COUNTER = {} + def _get_workspace_dir() -> str: """Determines the target workspace directory based on env vars for host.""" base_dir_env = os.getenv("CAI_WORKSPACE_DIR") @@ -183,8 +186,23 @@ class ShellSession: # pylint: disable=too-many-instance-attributes self.output_buffer.append(output) self.last_activity = time.time() else: - self.is_running = False - break + # os.read() returned empty. This does NOT necessarily mean + # the process itself has exited if self.process.poll() is None. + # It might be idle and waiting for input. + # The self.process.poll() check at the start of the loop iteration + # (or if self.is_process_running() is used in the sleep condition) + # will handle actual termination. + # Thus, we 'pass' here and let the loop continue with its sleep, + # rather than prematurely setting self.is_running = False. + if self.process and self.process.poll() is None: + # Process is alive but PTY read was empty (e.g., idle). + # Do nothing here to change is_running. + pass + else: + # Process is confirmed dead or no process to check, + # and read returned empty. Session is over. + self.is_running = False + break # Exit the while loop except Exception as read_err: self.output_buffer.append(f"Error reading output buffer: {str(read_err)}") self.is_running = False @@ -351,13 +369,56 @@ def send_to_session(session_id, input_data): return session.send_input(input_data) -def get_session_output(session_id, clear=True): +def get_session_output(session_id, clear=True, stdout=True): """Get output from a specific session""" if session_id not in ACTIVE_SESSIONS: return f"Session {session_id} not found" session = ACTIVE_SESSIONS[session_id] - return session.get_output(clear) + output = session.get_output(clear) + + # If stdout is enabled, display using cli_print_tool_output with counter + if stdout: + # Generate unique counter for this session output command + counter_key = f"session_output_{session_id}" + if counter_key not in SESSION_OUTPUT_COUNTER: + SESSION_OUTPUT_COUNTER[counter_key] = 0 + SESSION_OUTPUT_COUNTER[counter_key] += 1 + + # Create args for display + session_args = { + "command": "session", + "args": f"output {session_id}", + "session_id": session_id, + "call_counter": SESSION_OUTPUT_COUNTER[counter_key] # This ensures uniqueness + } + + # Determine environment info for display + env_type = "Local" + if session.container_id: + env_type = f"Container({session.container_id[:12]})" + elif session.ctf: + env_type = "CTF" + + # Create execution info + execution_info = { + "status": "completed", + "environment": env_type, + "host": session.workspace_dir, + "session_id": session_id + } + + # Display the session output using cli_print_tool_output + from cai.util import cli_print_tool_output + cli_print_tool_output( + tool_name="generic_linux_command", + args=session_args, + output=output, + execution_info=execution_info, + streaming=False + ) + + return output def terminate_session(session_id): @@ -741,14 +802,50 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg return f"Session {session_id} not found" session = ACTIVE_SESSIONS[session_id] result = session.send_input(command) # Send the raw command string - if stdout: - output = get_session_output(session_id, clear=False) - env_type = "Local" - if session.container_id: - env_type = f"Container({session.container_id[:12]})" - elif session.ctf: - env_type = "CTF" - print(f"\033[32m(Session {session_id} in {env_type}:{session.workspace_dir}) >> {command}\n{output}\033[0m") # noqa E501 + + # 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}" + if counter_key not in SESSION_OUTPUT_COUNTER: + SESSION_OUTPUT_COUNTER[counter_key] = 0 + SESSION_OUTPUT_COUNTER[counter_key] += 1 + + # Create args for display + session_args = { + "command": command, + "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 + } + + # Determine environment info for display + env_type = "Local" + if session.container_id: + env_type = f"Container({session.container_id[:12]})" + elif session.ctf: + env_type = "CTF" + + # Get the session output to display + output = get_session_output(session_id, clear=False, stdout=False) + + # Create execution info + execution_info = { + "status": "completed", + "environment": env_type, + "host": session.workspace_dir, + "session_id": session_id + } + + # Display the session input and its result using cli_print_tool_output + from cai.util import cli_print_tool_output + cli_print_tool_output( + tool_name="generic_linux_command", + args=session_args, + output=output, + execution_info=execution_info, + streaming=False + ) # For async sessions, we don't switch back to idle mode here # since the session continues to run in the background @@ -770,7 +867,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg context_msg = f"(docker:{container_id[:12]}:{container_workspace})" # Handle Async Session Creation in Container - if async_mode: + # Only create new session if no session_id is provided + if async_mode and not session_id: # Create a session specifically for the container environment new_session_id = create_shell_session(command, container_id=container_id) # noqa E501 if "Failed" in new_session_id: # Check if session creation failed @@ -781,7 +879,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg if stdout: # Wait a moment for initial output time.sleep(0.2) - output = get_session_output(new_session_id, clear=False) + 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 # For async sessions, switch back to idle mode after session creation @@ -1231,7 +1329,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # --- Local Execution (Default Fallback) --- # Let _run_local handle determining the host workspace # Handle Async Session Creation Locally - if async_mode: + # Only create new session if no session_id is provided + if async_mode and not session_id: # create_shell_session uses _get_workspace_dir() when container_id is None new_session_id = create_shell_session(command) if isinstance(new_session_id, str) and "Failed" in new_session_id: # Check failure @@ -1244,7 +1343,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg actual_workspace = session.workspace_dir if session else "unknown" if stdout: time.sleep(0.2) # Allow session buffer to populate - output = get_session_output(new_session_id, clear=False) + 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") # For async, switch back to idle mode after session creation diff --git a/src/cai/tools/reconnaissance/generic_linux_command.py b/src/cai/tools/reconnaissance/generic_linux_command.py index 81fbf0ba..b1ab271f 100644 --- a/src/cai/tools/reconnaissance/generic_linux_command.py +++ b/src/cai/tools/reconnaissance/generic_linux_command.py @@ -63,7 +63,10 @@ def generic_linux_command(command: str = "", if args.startswith("output "): session_id = args.split(" ")[1] - return get_session_output(session_id) + # 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 + get_session_output(session_id, clear=False, stdout=True) + return f"Session {session_id} output displayed above" if args.startswith("kill "): session_id = args.split(" ")[1] @@ -72,12 +75,23 @@ def generic_linux_command(command: str = "", return """Unknown session command. Available: 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 + # Regular command execution full_command = f'{command} {args}'.strip() # 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'] + 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 @@ -99,6 +113,8 @@ def generic_linux_command(command: str = "", if stream and not call_id: 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, diff --git a/src/cai/util.py b/src/cai/util.py index b7223fad..2e77c45a 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -1875,6 +1875,7 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut - interaction_cost, total_cost: optional cost values streaming: Flag indicating if this is part of a streaming output """ + import time # 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 @@ -1891,9 +1892,17 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut if not hasattr(cli_print_tool_output, '_seen_calls'): cli_print_tool_output._seen_calls = {} - # Track all displayed commands to prevent duplicates + # Track all displayed commands to prevent duplicates with cleanup if not hasattr(cli_print_tool_output, '_displayed_commands'): cli_print_tool_output._displayed_commands = set() + cli_print_tool_output._last_cleanup = time.time() + + # Periodic cleanup to prevent memory growth + current_time = time.time() + if current_time - cli_print_tool_output._last_cleanup > 300: # Cleanup every 5 minutes + # Clear the displayed commands set periodically + cli_print_tool_output._displayed_commands.clear() + cli_print_tool_output._last_cleanup = current_time # --- Consistent Command Key Generation --- effective_command_args_str = "" @@ -1915,6 +1924,13 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut effective_command_args_str = args command_key = f"{tool_name}:{effective_command_args_str}" + + # If args contain a call_counter, append it to make the key unique + # This allows commands with counters to always display + if isinstance(args, dict) and "call_counter" in args: + call_counter = args["call_counter"] + command_key += f":counter_{call_counter}" + # --- End of Command Key Generation --- # Check for duplicate display conditions @@ -2058,11 +2074,18 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Create a console for output console = Console(theme=theme) + # Clean args for display (remove internal counters and flags) + display_args = args + if isinstance(args, dict): + # Remove internal tracking fields that shouldn't be shown to the user + display_args = {k: v for k, v in args.items() + if k not in ["call_counter", "input_to_session"]} + # Get the panel content - with syntax highlighting - header, content = _create_tool_panel_content(tool_name, args, output, execution_info, token_info) + header, content = _create_tool_panel_content(tool_name, display_args, output, execution_info, token_info) # Format args for the title display - args_str = _format_tool_args(args, tool_name=tool_name) + args_str = _format_tool_args(display_args, tool_name=tool_name) # Determine border style based on status border_style = "blue" # Default for non-streaming