mirror of https://github.com/aliasrobotics/cai.git
Merge branch '0.4.0' of https://gitlab.com/aliasrobotics/alias_research/cai into 0.4.0
This commit is contained in:
Víctor Mayoral Vilches
commit
5ab7f64562
|
|
@ -159,13 +159,12 @@ from cai.repl.commands.parallel import PARALLEL_CONFIGS, ParallelConfig
|
|||
from cai import is_pentestperf_available
|
||||
ctf_global = None
|
||||
messages_ctf = ""
|
||||
ctf_init=1
|
||||
previous_ctf_name = os.getenv('CTF_NAME', None)
|
||||
if is_pentestperf_available() and os.getenv('CTF_NAME', None):
|
||||
ctf, messages_ctf = setup_ctf()
|
||||
ctf_global = ctf
|
||||
if os.getenv('CTF_INSIDE', 'True').lower() == 'false':
|
||||
container_id = ""
|
||||
os.environ['CAI_ACTIVE_CONTAINER'] = container_id
|
||||
ctf_init=0
|
||||
|
||||
# NOTE: This is needed when using LiteLLM Proxy Server
|
||||
#
|
||||
|
|
@ -248,12 +247,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'):
|
||||
|
|
@ -267,6 +266,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
|
|||
global previous_ctf_name
|
||||
global ctf_global
|
||||
global messages_ctf
|
||||
global ctf_init
|
||||
if previous_ctf_name != os.getenv('CTF_NAME', None):
|
||||
if is_pentestperf_available():
|
||||
if ctf_global:
|
||||
|
|
@ -274,7 +274,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
|
|||
ctf, messages_ctf = setup_ctf()
|
||||
ctf_global = ctf
|
||||
previous_ctf_name = os.getenv('CTF_NAME', None)
|
||||
|
||||
ctf_init=0
|
||||
# Check if CAI_MAX_TURNS has been updated via /config
|
||||
current_max_turns = os.getenv('CAI_MAX_TURNS', 'inf')
|
||||
if current_max_turns != str(prev_max_turns):
|
||||
|
|
@ -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'):
|
||||
|
|
@ -332,7 +332,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
|
|||
except Exception as e:
|
||||
console.print(f"[red]Error switching agent: {str(e)}[/red]")
|
||||
|
||||
if not force_until_flag:
|
||||
if not force_until_flag and ctf_init!=0:
|
||||
# Get user input with command completion and history
|
||||
user_input = get_user_input(
|
||||
command_completer,
|
||||
|
|
@ -340,11 +340,11 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
|
|||
history_file,
|
||||
get_toolbar_with_refresh,
|
||||
current_text
|
||||
) + messages_ctf
|
||||
)
|
||||
|
||||
else:
|
||||
user_input = messages_ctf
|
||||
if os.getenv('CTF_INSIDE', 'True').lower() == 'false':
|
||||
user_input += ctf_global.get_ip()
|
||||
ctf_init=1
|
||||
idle_time += time.time() - idle_start_time
|
||||
|
||||
# Stop measuring user idle time and start measuring active time
|
||||
|
|
@ -651,7 +651,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'),
|
|||
history_context.append({"role": "user", "content": user_input})
|
||||
conversation_input = history_context
|
||||
else:
|
||||
conversation_input = user_input
|
||||
conversation_input = messages_ctf + user_input
|
||||
|
||||
# Process the conversation with the agent - with parallel execution if enabled
|
||||
if parallel_count > 1:
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,85 @@ 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
|
||||
elif is_async_session_input:
|
||||
should_display_message = True
|
||||
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 +658,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
|
||||
|
|
@ -3192,15 +3262,41 @@ class _Converter:
|
|||
# Check if we're in streaming mode
|
||||
is_streaming_enabled = os.environ.get('CAI_STREAM', 'false').lower() == 'true'
|
||||
|
||||
# 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
|
||||
)
|
||||
# Check if this output was already displayed during streaming
|
||||
# For async sessions, we always display since they don't have real streaming
|
||||
should_display = True
|
||||
|
||||
# If streaming is enabled, check if this was already shown
|
||||
if is_streaming_enabled and hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls:
|
||||
tool_call_info = cls.recent_tool_calls[call_id]
|
||||
# Check if this tool was executed very recently (within last 5 seconds)
|
||||
# This indicates it was likely shown during streaming
|
||||
if 'start_time' in tool_call_info:
|
||||
time_since_execution = time.time() - tool_call_info['start_time']
|
||||
# For regular commands executed recently in streaming mode, skip display
|
||||
# But always display for async session commands (they have session_id in args)
|
||||
if time_since_execution < 5.0:
|
||||
# Parse arguments to check if this is an async session command
|
||||
try:
|
||||
import json
|
||||
args_dict = json.loads(tool_args) if isinstance(tool_args, str) else tool_args
|
||||
# If it has session_id, it's an async command - always show
|
||||
if not (isinstance(args_dict, dict) and args_dict.get("session_id")):
|
||||
should_display = False
|
||||
except:
|
||||
# If we can't parse args, assume it's a regular command
|
||||
should_display = False
|
||||
|
||||
# Only display if it hasn't been shown during streaming
|
||||
if should_display:
|
||||
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()
|
||||
|
|
|
|||
|
|
@ -174,42 +174,67 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
|
|||
def _read_output(self):
|
||||
"""Read output from the process"""
|
||||
try:
|
||||
# Buffer for incomplete lines
|
||||
partial_line = ""
|
||||
|
||||
while self.is_running and self.master is not None:
|
||||
try:
|
||||
# Check if process has exited before reading
|
||||
if self.process and self.process.poll() is not None:
|
||||
self.is_running = False
|
||||
break
|
||||
# Read the output
|
||||
output = os.read(self.master, 1024).decode()
|
||||
|
||||
# Read the output - increased buffer size to avoid cutting commands
|
||||
output = os.read(self.master, 4096).decode('utf-8', errors='replace')
|
||||
|
||||
if output:
|
||||
self.output_buffer.append(output)
|
||||
# Combine with any partial line from previous read
|
||||
full_output = partial_line + output
|
||||
|
||||
# Split into lines but keep the last partial line if it doesn't end with newline
|
||||
lines = full_output.split('\n')
|
||||
|
||||
# If output doesn't end with newline, the last item is a partial line
|
||||
if not output.endswith('\n'):
|
||||
partial_line = lines[-1]
|
||||
lines = lines[:-1]
|
||||
else:
|
||||
partial_line = ""
|
||||
|
||||
# Add complete lines to buffer
|
||||
for line in lines:
|
||||
if line: # Don't add empty lines
|
||||
self.output_buffer.append(line)
|
||||
|
||||
self.last_activity = time.time()
|
||||
else:
|
||||
# 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.
|
||||
if partial_line:
|
||||
# Add any remaining partial line
|
||||
self.output_buffer.append(partial_line)
|
||||
self.is_running = False
|
||||
break # Exit the while loop
|
||||
break
|
||||
except UnicodeDecodeError as e:
|
||||
# Handle unicode decode errors gracefully
|
||||
self.output_buffer.append(f"[Session {self.session_id}] Unicode decode error in output")
|
||||
continue
|
||||
except Exception as read_err:
|
||||
self.output_buffer.append(f"Error reading output buffer: {str(read_err)}")
|
||||
self.is_running = False
|
||||
break
|
||||
self.output_buffer.append(f"Error reading output buffer: {str(read_err)}")
|
||||
self.is_running = False
|
||||
break
|
||||
|
||||
# Add a small sleep to prevent busy-waiting if no output
|
||||
if self.is_process_running():
|
||||
time.sleep(0.05)
|
||||
time.sleep(0.05)
|
||||
|
||||
except Exception as e:
|
||||
self.output_buffer.append(f"Error in read_output loop: {str(e)}")
|
||||
self.is_running = False
|
||||
|
|
@ -261,6 +286,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 +432,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 +440,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 +450,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 +476,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 +630,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 +760,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,9 +803,21 @@ 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
|
||||
}
|
||||
|
||||
# Only add auto_output if not already present (prevents duplication)
|
||||
if args and isinstance(args, dict):
|
||||
# If args were passed and contain auto_output, use that value
|
||||
if "auto_output" in args:
|
||||
session_args["auto_output"] = args["auto_output"]
|
||||
else:
|
||||
# Otherwise, force it to True for session commands
|
||||
session_args["auto_output"] = True
|
||||
else:
|
||||
# No args provided, force auto_output
|
||||
session_args["auto_output"] = True
|
||||
|
||||
# Determine environment info for display
|
||||
env_type = "Local"
|
||||
if session.container_id:
|
||||
|
|
@ -785,15 +825,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 +855,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 +882,46 @@ 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
|
||||
}
|
||||
|
||||
# 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 +1120,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 +1233,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 +1362,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 +1381,49 @@ 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
|
||||
}
|
||||
|
||||
# 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()
|
||||
|
|
|
|||
|
|
@ -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 <id>")
|
||||
- Send input: Use session_id parameter
|
||||
- End: generic_linux_command("session", "kill <id>")
|
||||
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 <id>, kill <id>"""
|
||||
return "Usage: session list|output <id>|kill <id>"
|
||||
|
||||
# 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
|
||||
|
|
|
|||
|
|
@ -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,12 @@ 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}"
|
||||
# Also include session_id to make it unique per session
|
||||
effective_command_args_str += f":session_{args.get('session_id', '')}"
|
||||
elif isinstance(args, str):
|
||||
# If args is a string, it might be a JSON representation or a plain string.
|
||||
try:
|
||||
|
|
@ -1882,6 +1900,11 @@ 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}"
|
||||
# Also include session_id to make it unique per session
|
||||
effective_command_args_str += f":session_{parsed_json_args.get('session_id', '')}"
|
||||
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 +1920,19 @@ 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)}"
|
||||
|
||||
# Special handling for auto_output commands - they should always display
|
||||
# even if a similar command was shown before
|
||||
if isinstance(args, dict) and args.get("auto_output"):
|
||||
# Add auto_output flag to the key to differentiate from manual commands
|
||||
command_key += ":auto_output"
|
||||
|
||||
# --- End of Command Key Generation ---
|
||||
|
||||
# Check for duplicate display conditions
|
||||
|
|
|
|||
|
|
@ -1,51 +1,103 @@
|
|||
import os
|
||||
import pytest
|
||||
import json
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Set test environment variables to avoid OpenAI client initialization errors
|
||||
os.environ["OPENAI_API_KEY"] = "test_key_for_ci_environment"
|
||||
|
||||
from cai.tools.reconnaissance.generic_linux_command import generic_linux_command
|
||||
from cai.tools.reconnaissance.generic_linux_command import generic_linux_command
|
||||
from cai.sdk.agents import RunContextWrapper
|
||||
|
||||
async def test_generic_linux_command_regular_commands():
|
||||
"""Test the execution of a regular command using the generic Linux command tool."""
|
||||
mock_ctx = MagicMock() # Create a mock context for the command execution
|
||||
params = {
|
||||
"command": "echo", # Command to be executed
|
||||
"args": "'hello'" # Arguments for the command
|
||||
}
|
||||
|
||||
# Invoke the tool with the specified parameters and await the result
|
||||
result = await generic_linux_command.on_invoke_tool(mock_ctx, json.dumps(params))
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_linux_command_echo():
|
||||
"""Test the execution of echo command using generic_linux_command."""
|
||||
args = {"command": "echo 'hello'"}
|
||||
result = await generic_linux_command.on_invoke_tool(
|
||||
RunContextWrapper(None), json.dumps(args)
|
||||
)
|
||||
assert result.strip() == 'hello'
|
||||
|
||||
# Assert that the result matches the expected output
|
||||
assert result.replace("\n", "") == 'hello'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_linux_command_ls():
|
||||
"""Test the execution of the 'ls' command using the generic Linux command tool."""
|
||||
mock_ctx = MagicMock() # Create a mock context for the command execution
|
||||
params = {
|
||||
"command": "ls", # Command to be executed
|
||||
"args": "-l" # Arguments for the command
|
||||
}
|
||||
"""Test the execution of ls command using generic_linux_command."""
|
||||
args = {"command": "ls -l"}
|
||||
result = await generic_linux_command.on_invoke_tool(
|
||||
RunContextWrapper(None), json.dumps(args)
|
||||
)
|
||||
# Check that the output contains typical ls -l indicators
|
||||
assert "total" in result or "drwx" in result or "-rw" in result
|
||||
|
||||
# Invoke the tool with the specified parameters and await the result
|
||||
result = await generic_linux_command.on_invoke_tool(mock_ctx, json.dumps(params))
|
||||
|
||||
# Assert that the output contains 'total', which is typical for 'ls -l'
|
||||
assert "total" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_linux_command_invalid_command():
|
||||
"""Test the handling of an invalid command using the generic Linux command tool."""
|
||||
mock_ctx = MagicMock() # Create a mock context for the command execution
|
||||
params = {
|
||||
"command": "invalid_command", # Invalid command to be executed
|
||||
"args": "" # No arguments for the command
|
||||
}
|
||||
"""Test handling of invalid command using generic_linux_command."""
|
||||
args = {"command": "invalid_command_xyz123"}
|
||||
result = await generic_linux_command.on_invoke_tool(
|
||||
RunContextWrapper(None), json.dumps(args)
|
||||
)
|
||||
# Check for common error indicators
|
||||
assert ("not found" in result.lower() or
|
||||
"command not found" in result.lower() or
|
||||
"no such file" in result.lower())
|
||||
|
||||
# Invoke the tool with the specified parameters and await the result
|
||||
result = await generic_linux_command.on_invoke_tool(mock_ctx, json.dumps(params))
|
||||
|
||||
# Assert that the result indicates the command was not found
|
||||
assert "not found" in result
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_linux_command_empty_command():
|
||||
"""Test handling of empty command using generic_linux_command."""
|
||||
args = {"command": ""}
|
||||
result = await generic_linux_command.on_invoke_tool(
|
||||
RunContextWrapper(None), json.dumps(args)
|
||||
)
|
||||
assert "Error: No command provided" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_linux_command_session_list():
|
||||
"""Test session list functionality using generic_linux_command."""
|
||||
args = {"command": "session list"}
|
||||
result = await generic_linux_command.on_invoke_tool(
|
||||
RunContextWrapper(None), json.dumps(args)
|
||||
)
|
||||
assert "No active sessions" in result or "Active sessions:" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_linux_command_env_info():
|
||||
"""Test environment info functionality using generic_linux_command."""
|
||||
args = {"command": "env info"}
|
||||
result = await generic_linux_command.on_invoke_tool(
|
||||
RunContextWrapper(None), json.dumps(args)
|
||||
)
|
||||
assert "Current Environment:" in result
|
||||
assert "CTF Environment:" in result
|
||||
assert "Container:" in result
|
||||
assert "SSH:" in result
|
||||
assert "Workspace:" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_linux_command_interactive_flag():
|
||||
"""Test interactive flag functionality using generic_linux_command."""
|
||||
# Test with interactive=True but a simple command
|
||||
args = {"command": "echo 'test'", "interactive": True}
|
||||
result = await generic_linux_command.on_invoke_tool(
|
||||
RunContextWrapper(None), json.dumps(args)
|
||||
)
|
||||
# Should still work, just might have different session handling
|
||||
assert "test" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_linux_command_with_session_id():
|
||||
"""Test session_id parameter using generic_linux_command."""
|
||||
# Test with a non-existent session_id
|
||||
args = {"command": "echo 'test'", "session_id": "nonexistent123"}
|
||||
result = await generic_linux_command.on_invoke_tool(
|
||||
RunContextWrapper(None), json.dumps(args)
|
||||
)
|
||||
# Should handle gracefully - either execute or give session error
|
||||
assert isinstance(result, str)
|
||||
Loading…
Reference in New Issue