From 64e75ae27934545985c15f7a1a9322946c3b2761 Mon Sep 17 00:00:00 2001 From: luijait Date: Fri, 23 May 2025 17:43:19 +0200 Subject: [PATCH 1/4] Fix no stream tool calls --- src/cai/cli.py | 82 ++++--------------- .../agents/models/openai_chatcompletions.py | 23 ++++++ 2 files changed, 38 insertions(+), 67 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index 1258b77d..a90c6c33 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -744,80 +744,28 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), # Use non-streamed response response = asyncio.run(Runner.run(agent, conversation_input)) - # Process the response items + # En modo no-streaming, procesamos SOLO los tool outputs de response.new_items + # Los tool calls (assistant messages) ya se añaden correctamente en openai_chatcompletions.py for item in response.new_items: - # Handle tool call output items (tool results) + # Handle ONLY tool call output items (tool results) if isinstance(item, ToolCallOutputItem): - # First, ensure there's a corresponding assistant message with tool_calls - # before adding the tool response to prevent the OpenAI error - assistant_with_tool_call_exists = False tool_call_id = item.raw_item["call_id"] - for msg in message_history: - if (msg.get("role") == "assistant" and - msg.get("tool_calls") and - any(tc.get("id") == tool_call_id for tc in msg.get("tool_calls", []))): - assistant_with_tool_call_exists = True - break + # Verificar si ya existe este tool output en message_history para evitar duplicación + tool_msg_exists = any( + msg.get("role") == "tool" and msg.get("tool_call_id") == tool_call_id + for msg in message_history + ) - # If no matching assistant message exists, create one first - if not assistant_with_tool_call_exists: - tool_call_msg = { - "role": "assistant", - "content": None, - "tool_calls": [{ - "id": tool_call_id, - "type": "function", - "function": { - "name": "unknown_function", - "arguments": "{}" - } - }] + if not tool_msg_exists: + # Añadir solo el tool output al message_history + tool_msg = { + "role": "tool", + "tool_call_id": tool_call_id, + "content": item.output, } - add_to_message_history(tool_call_msg) - - # Now add the tool response - tool_msg = { - "role": "tool", - "tool_call_id": tool_call_id, - "content": item.output, - } - add_to_message_history(tool_msg) + add_to_message_history(tool_msg) - # Make sure that assistant messages with tool calls are also added to message_history - # This is especially important for non-streaming mode - if hasattr(agent, 'model'): - # Access the _Converter directly from the OpenAIChatCompletionsModel implementation - from cai.sdk.agents.models.openai_chatcompletions import _Converter - - # Check if recent_tool_calls exists and process them - if hasattr(_Converter, 'recent_tool_calls'): - for call_id, call_info in _Converter.recent_tool_calls.items(): - # Only process new tool calls that haven't been added to message history yet - tool_call_found = False - for msg in message_history: - if (msg.get("role") == "assistant" and - msg.get("tool_calls") and - any(tc.get("id") == call_id for tc in msg.get("tool_calls", []))): - tool_call_found = True - break - - if not tool_call_found: - # Add the assistant message with the tool call - tool_call_msg = { - "role": "assistant", - "content": None, - "tool_calls": [{ - "id": call_id, - "type": "function", - "function": { - "name": call_info.get('name', ''), - "arguments": call_info.get('arguments', '{}') - } - }] - } - add_to_message_history(tool_call_msg) - # Final validation to ensure message history follows OpenAI's requirements # Ensure every tool message has a preceding assistant message with matching tool_call_id from cai.util import fix_message_list diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 2d7cdca4..777d590b 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -656,6 +656,29 @@ class OpenAIChatCompletionsModel(Model): # Log the assistant message self.logger.log_assistant_message(assistant_msg.content) + # En no-streaming, también necesitamos añadir cualquier tool output al message_history + # Esto se hace procesando los items de output del ModelResponse + items = _Converter.message_to_output_items(response.choices[0].message) + + # Además, necesitamos añadir los tool outputs que se hayan generado + # durante la ejecución de las herramientas + if hasattr(_Converter, 'tool_outputs'): + for call_id, output_content in _Converter.tool_outputs.items(): + # Verificar si ya existe un mensaje tool con este call_id en message_history + tool_msg_exists = any( + msg.get("role") == "tool" and msg.get("tool_call_id") == call_id + for msg in message_history + ) + + if not tool_msg_exists: + # Añadir el mensaje tool al message_history + tool_msg = { + "role": "tool", + "tool_call_id": call_id, + "content": output_content + } + add_to_message_history(tool_msg) + # Log the complete response for the session self.logger.rec_training_data( { From 64fa082a164a3e032a61ea69fd1c54edd88b8bcb Mon Sep 17 00:00:00 2001 From: luijait Date: Fri, 23 May 2025 20:17:41 +0200 Subject: [PATCH 2/4] testing --- ci/test/.test.yml | 55 ++ src/cai/cli.py | 15 +- src/cai/repl/commands/parallel.py | 3 +- tests/cli/__init__.py | 10 + tests/cli/base_cli_test.py | 523 ++++++++++++++++++ tests/cli/test_cli_streaming.py | 701 ++++++++++++++++++++++++ tests/commands/pytest.ini | 20 + tests/commands/test_command_agent.py | 457 +++++++++++++++ tests/commands/test_command_base.py | 528 ++++++++++++++++++ tests/commands/test_command_config.py | 381 +++++++++++++ tests/commands/test_command_help.py | 402 ++++++++++++++ tests/commands/test_command_history.py | 518 +++++++++++++++++ tests/commands/test_command_model.py | 539 ++++++++++++++++++ tests/commands/test_command_parallel.py | 344 ++++++++++++ 14 files changed, 4490 insertions(+), 6 deletions(-) create mode 100644 tests/cli/__init__.py create mode 100644 tests/cli/base_cli_test.py create mode 100644 tests/cli/test_cli_streaming.py create mode 100644 tests/commands/pytest.ini create mode 100644 tests/commands/test_command_agent.py create mode 100644 tests/commands/test_command_base.py create mode 100644 tests/commands/test_command_config.py create mode 100644 tests/commands/test_command_help.py create mode 100644 tests/commands/test_command_history.py create mode 100644 tests/commands/test_command_model.py create mode 100644 tests/commands/test_command_parallel.py diff --git a/ci/test/.test.yml b/ci/test/.test.yml index 1e66a9c6..f5e1a082 100644 --- a/ci/test/.test.yml +++ b/ci/test/.test.yml @@ -293,3 +293,58 @@ <<: *run_test variables: TEST_PATH: tests/others/test_function_schema.py + +💻 cli test_cli_streaming.py: + <<: *run_test + variables: + TEST_PATH: tests/cli/test_cli_streaming.py + +💻 commands test_command_base.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_base.py + +💻 commands test_command_parallel.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_parallel.py + +💻 commands test_command_agent.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_agent.py + +💻 commands test_command_model.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_model.py + +💻 commands test_command_history.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_history.py + +💻 commands test_command_config.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_config.py + +💻 commands test_command_flush.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_flush.py + +💻 commands test_command_help.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_help.py + +💻 commands test_command_virtualization.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_virtualization.py + +💻 commands test_command_workspace.py: + <<: *run_test + variables: + TEST_PATH: tests/commands/test_command_workspace.py diff --git a/src/cai/cli.py b/src/cai/cli.py index a90c6c33..c63482e7 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -413,9 +413,13 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), pending_calls.append(call_info.get('name', 'unknown')) # Apply message list fixes + # NOTE: Commenting this to avoid creating duplicate synthetic tool calls + # The synthetic tool calls we just created are already in the correct format + # if pending_calls: + # from cai.util import fix_message_list + # message_history[:] = fix_message_list(message_history) + # print(f"\033[93mCleaned up {len(pending_calls)} pending tool calls before exit\033[0m") if pending_calls: - from cai.util import fix_message_list - message_history[:] = fix_message_list(message_history) print(f"\033[93mCleaned up {len(pending_calls)} pending tool calls before exit\033[0m") except Exception: pass @@ -830,8 +834,11 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), add_to_message_history(tool_msg) # Apply message list fixes - from cai.util import fix_message_list - message_history[:] = fix_message_list(message_history) + # NOTE: Commenting this to avoid creating duplicate synthetic tool calls + # The synthetic tool calls we just created are already in the correct format + # from cai.util import fix_message_list + # message_history[:] = fix_message_list(message_history) + pass except Exception as cleanup_error: print(f"\033[91mError cleaning up interrupted tools: {str(cleanup_error)}\033[0m") diff --git a/src/cai/repl/commands/parallel.py b/src/cai/repl/commands/parallel.py index 6c6195c3..5582045d 100644 --- a/src/cai/repl/commands/parallel.py +++ b/src/cai/repl/commands/parallel.py @@ -153,9 +153,8 @@ class ParallelCommand(Command): Returns: True if successful """ - global PARALLEL_CONFIGS count = len(PARALLEL_CONFIGS) - PARALLEL_CONFIGS = [] + PARALLEL_CONFIGS.clear() console.print(f"[green]Cleared {count} parallel configurations[/green]") return True diff --git a/tests/cli/__init__.py b/tests/cli/__init__.py new file mode 100644 index 00000000..f9819d84 --- /dev/null +++ b/tests/cli/__init__.py @@ -0,0 +1,10 @@ +""" +CLI Test Package + +This package contains comprehensive tests for the CAI CLI functionality, +including streaming, keyboard interrupts, message flow, and integration tests. +""" + +from .base_cli_test import BaseCLITest + +__all__ = ['BaseCLITest'] \ No newline at end of file diff --git a/tests/cli/base_cli_test.py b/tests/cli/base_cli_test.py new file mode 100644 index 00000000..0c13f433 --- /dev/null +++ b/tests/cli/base_cli_test.py @@ -0,0 +1,523 @@ +#!/usr/bin/env python3 +""" +Base class for CLI testing with comprehensive mocking and utilities. +""" + +import asyncio +import json +import os +import sys +import time +from unittest.mock import AsyncMock, MagicMock, patch, Mock, call +from typing import Any, Dict, List, Optional, Callable +import pytest + +# Add src to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src')) + +from openai.types.chat.chat_completion import ChatCompletion, Choice +from openai.types.chat.chat_completion_message import ChatCompletionMessage +from openai.types.chat.chat_completion_message_tool_call import ( + ChatCompletionMessageToolCall, + Function, +) +from openai.types.completion_usage import CompletionUsage + +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel, ModelResponse, Runner +from cai.sdk.agents.models.openai_chatcompletions import message_history + + +class CLIMessageSimulator: + """Simulates message flow in the CLI with proper timing and state management.""" + + def __init__(self): + self.messages = [] + self.current_index = 0 + self.completion_responses = [] + self.tool_call_responses = {} + self.interrupt_triggers = {} + + def add_user_message(self, content: str, interrupt_after: bool = False): + """Add a user message to the simulation.""" + self.messages.append({ + 'role': 'user', + 'content': content, + 'interrupt_after': interrupt_after + }) + + def add_assistant_response(self, content: str, tool_calls: Optional[List[Dict]] = None): + """Add an expected assistant response.""" + response_data = { + 'role': 'assistant', + 'content': content + } + if tool_calls: + response_data['tool_calls'] = tool_calls + + self.completion_responses.append(response_data) + + def add_tool_response(self, call_id: str, output: str): + """Add a tool call response.""" + self.tool_call_responses[call_id] = output + + def set_interrupt_trigger(self, message_index: int, during_execution: bool = False): + """Set when to trigger a KeyboardInterrupt.""" + self.interrupt_triggers[message_index] = { + 'during_execution': during_execution + } + + def get_next_message(self) -> Optional[Dict]: + """Get the next message in the simulation.""" + if self.current_index < len(self.messages): + msg = self.messages[self.current_index] + self.current_index += 1 + return msg + return None + + def get_completion_response(self, index: int) -> Optional[Dict]: + """Get the completion response for a given index.""" + if index < len(self.completion_responses): + return self.completion_responses[index] + return None + + def should_interrupt(self, index: int, during_execution: bool = False) -> bool: + """Check if an interrupt should be triggered.""" + trigger = self.interrupt_triggers.get(index) + if trigger: + return trigger['during_execution'] == during_execution + return False + + def reset(self): + """Reset the simulator state.""" + self.current_index = 0 + + +class BaseCLITest: + """ + Comprehensive base class for CLI testing with advanced mocking capabilities. + + This class provides: + - Complete CLI environment mocking + - Message flow simulation + - Streaming and non-streaming mode testing + - Keyboard interrupt simulation at various points + - Tool call mocking and verification + - Integration with openai_chatcompletions.py logic + """ + + @classmethod + def setup_class(cls): + """Set up test environment.""" + # Disable external services for testing + os.environ['CAI_TELEMETRY'] = 'false' + os.environ['CAI_TRACING'] = 'false' + os.environ['CAI_STREAM'] = 'false' + os.environ['CAI_MAX_TURNS'] = '5' + + # Ensure we're using a test model + os.environ['CAI_MODEL'] = 'test-model' + + # Disable any CTF components + os.environ.pop('CTF_NAME', None) + + @classmethod + def teardown_class(cls): + """Clean up after tests.""" + message_history.clear() + + def setup_method(self): + """Set up for each test method.""" + message_history.clear() + self.simulator = CLIMessageSimulator() + + def create_mock_completion( + self, + content: str = "Test response", + tool_calls: Optional[List[Dict[str, Any]]] = None, + usage: Optional[Dict[str, int]] = None + ) -> ChatCompletion: + """ + Create a mock ChatCompletion response with proper structure. + + Args: + content: The assistant's response content + tool_calls: List of tool calls to include + usage: Token usage information + + Returns: + Properly formatted ChatCompletion object + """ + message_data = {"role": "assistant", "content": content} + + if tool_calls: + formatted_tool_calls = [] + for tc in tool_calls: + tool_call = ChatCompletionMessageToolCall( + id=tc.get("id", f"call_{int(time.time() * 1000)}"), + type="function", + function=Function( + name=tc.get("function", {}).get("name", "test_function"), + arguments=tc.get("function", {}).get("arguments", "{}") + ) + ) + formatted_tool_calls.append(tool_call) + message_data["tool_calls"] = formatted_tool_calls + + msg = ChatCompletionMessage(**message_data) + choice = Choice(index=0, finish_reason="stop", message=msg) + + # Default usage if not provided + if not usage: + usage = {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30} + + return ChatCompletion( + id=f"test-completion-{int(time.time() * 1000)}", + created=int(time.time()), + model="test-model", + object="chat.completion", + choices=[choice], + usage=CompletionUsage(**usage) + ) + + def create_mock_agent(self, model_name: str = "test-model") -> Agent: + """Create a mock agent with proper configuration.""" + mock_client = AsyncMock() + mock_client.base_url = "http://test-url" + + test_model = OpenAIChatCompletionsModel( + model=model_name, + openai_client=mock_client + ) + + return Agent( + name="TestAgent", + instructions="You are a test assistant", + model=test_model + ) + + def create_mock_model_response( + self, + content: str = "Test response", + items: Optional[List] = None + ) -> ModelResponse: + """Create a mock ModelResponse for Runner.run.""" + from cai.sdk.agents.usage import Usage + return ModelResponse( + output=items or [], + usage=Usage(requests=1, input_tokens=10, output_tokens=20, total_tokens=30), + referenceable_id=None + ) + + def create_input_simulator(self, messages: List[str], interrupts: Optional[Dict[int, str]] = None): + """ + Create an input simulator that provides predefined messages and can trigger interrupts. + + Args: + messages: List of user input messages + interrupts: Dict mapping message index to interrupt type + e.g., {1: "during_input", 2: "during_processing"} + + Returns: + A function that can be used to mock user input + """ + message_index = [0] + + def mock_input_function(*args, **kwargs): + current_index = message_index[0] + + # Check if we should interrupt before providing input + if interrupts and current_index in interrupts: + interrupt_type = interrupts[current_index] + if interrupt_type == "before_input": + raise KeyboardInterrupt(f"Simulated interrupt before message {current_index}") + + # Provide the next message if available + if current_index < len(messages): + message = messages[current_index] + message_index[0] += 1 + + # Check if we should interrupt after providing input + if interrupts and current_index in interrupts: + interrupt_type = interrupts[current_index] + if interrupt_type == "after_input": + # Return the message but arrange for interrupt on next call + return message + + return message + else: + # No more messages, trigger completion interrupt + raise KeyboardInterrupt("Test completed - no more messages") + + return mock_input_function + + def create_litellm_simulator(self, responses: List[ChatCompletion], interrupts: Optional[Dict[int, str]] = None): + """ + Create a LiteLLM simulator that provides predefined responses and can trigger interrupts. + + Args: + responses: List of ChatCompletion responses to return + interrupts: Dict mapping response index to interrupt type + + Returns: + A function that can be used to mock litellm.completion + """ + response_index = [0] + + def mock_litellm_function(*args, **kwargs): + current_index = response_index[0] + + # Check if we should interrupt during processing + if interrupts and current_index in interrupts: + interrupt_type = interrupts[current_index] + if interrupt_type == "during_llm_call": + raise KeyboardInterrupt(f"Simulated interrupt during LLM call {current_index}") + + # Return the next response if available + if current_index < len(responses): + response = responses[current_index] + response_index[0] += 1 + return response + else: + # Return the last response for any additional calls + return responses[-1] if responses else self.create_mock_completion() + + return mock_litellm_function + + def run_cli_simulation( + self, + agent: Agent, + user_inputs: List[str], + expected_responses: List[str], + stream_mode: bool = False, + interrupts: Optional[Dict[int, str]] = None, + tool_calls: Optional[Dict[int, List[Dict]]] = None, + verify_message_flow: bool = True + ) -> Dict[str, Any]: + """ + Run a complete CLI simulation with full control over inputs, outputs, and interrupts. + + Args: + agent: The agent to use for testing + user_inputs: List of user input messages + expected_responses: List of expected assistant responses + stream_mode: Whether to test streaming mode + interrupts: Dict mapping indices to interrupt types + tool_calls: Dict mapping response indices to tool calls + verify_message_flow: Whether to verify message history flow + + Returns: + Dict with simulation results and verification data + """ + # Set streaming mode + os.environ['CAI_STREAM'] = 'true' if stream_mode else 'false' + + # Prepare mock responses + mock_responses = [] + for i, response_content in enumerate(expected_responses): + response_tool_calls = tool_calls.get(i) if tool_calls else None + mock_responses.append( + self.create_mock_completion(response_content, response_tool_calls) + ) + + # Create simulators + input_simulator = self.create_input_simulator(user_inputs, interrupts) + litellm_simulator = self.create_litellm_simulator(mock_responses, interrupts) + + # Track execution results + results = { + 'user_inputs_processed': [], + 'assistant_responses': [], + 'tool_calls_made': [], + 'tool_outputs': [], + 'interrupts_caught': [], + 'message_history_final': [], + 'llm_calls': [], + 'exceptions': [], + 'stream_events': [] if stream_mode else None + } + + # Enhanced mocking for CLI components + mock_patches = [ + # Core CLI input/output + patch('cai.repl.ui.prompt.get_user_input', side_effect=input_simulator), + patch('cai.repl.ui.logging.setup_session_logging', return_value="test_history.txt"), + + # Session recording + patch('cai.sdk.agents.run_to_jsonl.get_session_recorder'), + + # CLI UI components + patch('cai.repl.commands.FuzzyCommandCompleter'), + patch('cai.repl.ui.keybindings.create_key_bindings'), + patch('cai.repl.ui.banner.display_banner'), + patch('cai.repl.ui.banner.display_quick_guide'), + + # LLM calls + patch('litellm.completion', side_effect=litellm_simulator), + patch('litellm.acompletion', side_effect=litellm_simulator), + + # Timing functions + patch('cai.util.start_idle_timer'), + patch('cai.util.stop_idle_timer'), + patch('cai.util.start_active_timer'), + patch('cai.util.stop_active_timer'), + patch('cai.util.get_active_time_seconds', return_value=1.0), + patch('cai.util.get_idle_time_seconds', return_value=2.0), + + # Rich console output + patch('rich.console.Console.print'), + ] + + # Apply all patches and run simulation + from cai.cli import run_cai_cli + + def apply_patches_and_run(): + with patch.multiple( + 'cai.repl.ui.prompt', + get_user_input=input_simulator + ), patch.multiple( + 'litellm', + completion=litellm_simulator, + acompletion=litellm_simulator + ), patch.multiple( + 'cai.repl.ui.logging', + setup_session_logging=Mock(return_value="test_history.txt") + ), patch.multiple( + 'cai.sdk.agents.run_to_jsonl', + get_session_recorder=Mock(return_value=Mock( + filename="test_session.jsonl", + log_user_message=Mock(), + log_assistant_message=Mock(), + log_session_end=Mock(), + rec_training_data=Mock() + )) + ), patch.multiple( + 'cai.repl.commands', + FuzzyCommandCompleter=Mock() + ), patch.multiple( + 'cai.repl.ui.keybindings', + create_key_bindings=Mock() + ), patch.multiple( + 'cai.repl.ui.banner', + display_banner=Mock(), + display_quick_guide=Mock() + ), patch.multiple( + 'cai.util', + start_idle_timer=Mock(), + stop_idle_timer=Mock(), + start_active_timer=Mock(), + stop_active_timer=Mock(), + get_active_time_seconds=Mock(return_value=1.0), + get_idle_time_seconds=Mock(return_value=2.0) + ), patch.multiple( + 'rich.console', + Console=Mock() + ): + try: + run_cai_cli( + starting_agent=agent, + max_turns=len(user_inputs), + force_until_flag=False + ) + except KeyboardInterrupt as e: + results['interrupts_caught'].append(str(e)) + except Exception as e: + results['exceptions'].append(str(e)) + + # Execute the simulation + apply_patches_and_run() + + # Capture final state + results['message_history_final'] = list(message_history) + + # Verify message flow if requested + if verify_message_flow: + results['message_flow_valid'] = self._verify_message_flow( + user_inputs, expected_responses, tool_calls + ) + + return results + + def _verify_message_flow( + self, + user_inputs: List[str], + expected_responses: List[str], + tool_calls: Optional[Dict[int, List[Dict]]] = None + ) -> bool: + """Verify that the message flow in message_history is correct.""" + try: + # Check that we have the expected number of messages + expected_message_count = len(user_inputs) + len(expected_responses) + if tool_calls: + # Add tool call messages and tool result messages + expected_message_count += sum(len(calls) * 2 for calls in tool_calls.values()) + + if len(message_history) < len(user_inputs): + return False + + # Verify message sequence + message_index = 0 + for i in range(len(user_inputs)): + # Check user message + if message_index >= len(message_history): + return False + + user_msg = message_history[message_index] + if user_msg.get('role') != 'user' or user_inputs[i] not in str(user_msg.get('content', '')): + return False + + message_index += 1 + + # Check assistant message if we expect one + if i < len(expected_responses): + if message_index >= len(message_history): + return False + + assistant_msg = message_history[message_index] + if assistant_msg.get('role') != 'assistant': + return False + + message_index += 1 + + return True + + except Exception: + return False + + def assert_message_history_contains(self, role: str, content_substring: str): + """Assert that message history contains a message with the given role and content.""" + for msg in message_history: + if (msg.get('role') == role and + content_substring in str(msg.get('content', ''))): + return True + raise AssertionError( + f"Message history does not contain {role} message with content '{content_substring}'" + ) + + def assert_tool_call_made(self, function_name: str): + """Assert that a tool call was made with the given function name.""" + for msg in message_history: + if msg.get('role') == 'assistant' and msg.get('tool_calls'): + for tool_call in msg['tool_calls']: + if tool_call.get('function', {}).get('name') == function_name: + return True + raise AssertionError(f"No tool call found for function '{function_name}'") + + def assert_keyboard_interrupt_handled(self, results: Dict[str, Any]): + """Assert that keyboard interrupts were properly handled.""" + assert len(results['interrupts_caught']) > 0, "No keyboard interrupts were caught" + + def print_message_history_debug(self): + """Print the current message history for debugging.""" + print("\n=== MESSAGE HISTORY DEBUG ===") + for i, msg in enumerate(message_history): + role = msg.get('role', 'unknown') + content = str(msg.get('content', ''))[:100] + tool_calls = msg.get('tool_calls', []) + tool_call_id = msg.get('tool_call_id', '') + + print(f"[{i}] {role}: {content}") + if tool_calls: + print(f" Tool calls: {len(tool_calls)}") + if tool_call_id: + print(f" Tool call ID: {tool_call_id}") + print("=== END MESSAGE HISTORY ===\n") \ No newline at end of file diff --git a/tests/cli/test_cli_streaming.py b/tests/cli/test_cli_streaming.py new file mode 100644 index 00000000..88beae47 --- /dev/null +++ b/tests/cli/test_cli_streaming.py @@ -0,0 +1,701 @@ +#!/usr/bin/env python3 +""" +Test streaming functionality in the CLI. +Tests streaming mode, streaming interrupts, and streaming vs non-streaming behavior. +""" + +import os +import sys +import time +import unittest +from unittest.mock import patch, Mock, MagicMock + +# Add src to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'src')) + +from openai.types.chat.chat_completion import ChatCompletion, Choice +from openai.types.chat.chat_completion_message import ChatCompletionMessage +from openai.types.chat.chat_completion_message_tool_call import ( + ChatCompletionMessageToolCall, + Function, +) +from openai.types.completion_usage import CompletionUsage + +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel, ModelResponse +from cai.sdk.agents.models.openai_chatcompletions import message_history + + +class TestCLIStreaming(unittest.TestCase): + """Test CLI streaming functionality by testing components directly.""" + + @classmethod + def setUpClass(cls): + """Set up test environment.""" + os.environ['CAI_TELEMETRY'] = 'false' + os.environ['CAI_TRACING'] = 'false' + os.environ['CAI_STREAM'] = 'false' + + @classmethod + def tearDownClass(cls): + """Clean up after tests.""" + message_history.clear() + + def setUp(self): + """Set up for each test method.""" + # AGGRESSIVE cleanup to ensure no state contamination between tests + message_history.clear() + + # Clean up any lingering _Converter state + try: + from cai.sdk.agents.models.openai_chatcompletions import _Converter + # Clear all possible _Converter attributes + for attr in ['recent_tool_calls', 'tool_outputs', 'conversation_start_time']: + if hasattr(_Converter, attr): + if isinstance(getattr(_Converter, attr), dict): + getattr(_Converter, attr).clear() + else: + delattr(_Converter, attr) + except Exception: + pass + + # Also ensure environment is clean + os.environ['CAI_STREAM'] = 'false' + os.environ['CAI_TELEMETRY'] = 'false' + os.environ['CAI_TRACING'] = 'false' + + def test_ctrl_c_cleanup_message_consistency(self): + """Test CTRL+C cleanup logic maintains message consistency.""" + from cai.sdk.agents.models.openai_chatcompletions import ( + add_to_message_history, + _Converter + ) + + # Complete cleanup - clear everything including _Converter state + message_history.clear() + if hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls.clear() + if hasattr(_Converter, 'tool_outputs'): + _Converter.tool_outputs.clear() + + # Simulate the state before CTRL+C during tool execution + # 1. User message + add_to_message_history({"role": "user", "content": "Run a long command"}) + + # 2. Assistant message with tool call + tool_call_id = "call_interrupted_123" + add_to_message_history({ + "role": "assistant", + "content": "I'll run that command for you.", + "tool_calls": [{ + "id": tool_call_id, + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": '{"command": "sleep", "args": "30"}' + } + }] + }) + + # 3. Simulate CTRL+C happening during tool execution + # This is where the real cleanup logic would kick in + def simulate_ctrl_c_cleanup(): + """Simulate the exact cleanup logic from cli.py""" + + # Initialize _Converter attributes if they don't exist + if not hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls = {} + if not hasattr(_Converter, 'tool_outputs'): + _Converter.tool_outputs = {} + + # Simulate a tool call that was started but interrupted + _Converter.recent_tool_calls[tool_call_id] = { + 'name': 'generic_linux_command', + 'arguments': '{"command": "sleep", "args": "30"}', + 'start_time': time.time() - 5 # Started 5 seconds ago + } + + # Simulate the cleanup logic from cli.py lines 603-654 + pending_calls = [] + for call_id, call_info in list(_Converter.recent_tool_calls.items()): + # Check if this tool call has a corresponding response in message_history + tool_response_exists = any( + msg.get("role") == "tool" and msg.get("tool_call_id") == call_id + for msg in message_history + ) + + if not tool_response_exists: + # Add assistant message if needed (should already exist in our case) + assistant_exists = any( + msg.get("role") == "assistant" and + msg.get("tool_calls") and + any(tc.get("id") == call_id for tc in msg.get("tool_calls", [])) + for msg in message_history + ) + + if not assistant_exists: + # This shouldn't happen in our test but add for completeness + assistant_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": call_id, + "type": "function", + "function": { + "name": call_info.get('name', 'unknown_function'), + "arguments": call_info.get('arguments', '{}') + } + }] + } + add_to_message_history(assistant_msg) + + # Add synthetic tool response for interrupted tool + tool_msg = { + "role": "tool", + "tool_call_id": call_id, + "content": "Operation interrupted by user (Keyboard Interrupt)" + } + add_to_message_history(tool_msg) + pending_calls.append(call_info.get('name', 'unknown')) + + # Apply message list fixes like the real system does + from cai.util import fix_message_list + try: + fixed_messages = fix_message_list(message_history) + message_history.clear() + message_history.extend(fixed_messages) + return len(pending_calls) + except Exception as e: + print(f"fix_message_list failed: {e}") + return 0 + + # Execute the cleanup + cleaned_count = simulate_ctrl_c_cleanup() + + # Verify the cleanup worked + assert cleaned_count > 0, "Should have cleaned up at least one pending tool call" + + # Verify message history consistency + self.verify_message_history_openai_compliance() + + # Verify we have the expected sequence + assert len(message_history) >= 3, "Should have user, assistant, tool messages" + + # Check message roles in order + roles = [msg['role'] for msg in message_history] + assert roles[0] == 'user', "First message should be user" + assert roles[1] == 'assistant', "Second message should be assistant" + assert roles[2] == 'tool', "Third message should be tool" + + # Verify tool call/result consistency + assistant_msg = message_history[1] + tool_msg = message_history[2] + assert assistant_msg.get('tool_calls'), "Assistant message should have tool calls" + assert tool_msg['tool_call_id'] == assistant_msg['tool_calls'][0]['id'], \ + "Tool call ID should match" + assert "interrupted" in tool_msg['content'].lower(), \ + "Tool result should indicate interruption" + + print("✅ CTRL+C cleanup message consistency test passed!") + + # Clean up _Converter state after test + if hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls.clear() + if hasattr(_Converter, 'tool_outputs'): + _Converter.tool_outputs.clear() + + def test_fix_message_list_with_interrupted_tools(self): + """Test fix_message_list handles interrupted tool sequences correctly.""" + from cai.sdk.agents.models.openai_chatcompletions import ( + add_to_message_history, + _Converter + ) + from cai.util import fix_message_list + + # Complete cleanup - clear everything including _Converter state + message_history.clear() + if hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls.clear() + if hasattr(_Converter, 'tool_outputs'): + _Converter.tool_outputs.clear() + + # Create an incomplete sequence (tool call without result) + add_to_message_history({"role": "user", "content": "Test command"}) + add_to_message_history({ + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": "call_incomplete_456", + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": '{"command": "test", "args": "--help"}' + } + }] + }) + + # At this point we have incomplete sequence - no tool result + incomplete_messages = list(message_history) + + # Apply fix_message_list + try: + fixed_messages = fix_message_list(incomplete_messages) + + # Verify fix_message_list added the missing tool result + assert len(fixed_messages) > len(incomplete_messages), \ + "fix_message_list should add missing tool result" + + # Find the added tool message + tool_msg = None + for msg in fixed_messages: + if msg.get('role') == 'tool' and msg.get('tool_call_id') == 'call_incomplete_456': + tool_msg = msg + break + + assert tool_msg is not None, "fix_message_list should add tool result message" + + # Verify the fixed messages comply with OpenAI format + for i, msg in enumerate(fixed_messages): + assert 'role' in msg, f"Fixed message {i} missing role" + assert msg['role'] in ['user', 'assistant', 'system', 'tool'], \ + f"Fixed message {i} has invalid role" + + print("✅ fix_message_list with interrupted tools test passed!") + + # Clean up _Converter state after test + if hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls.clear() + if hasattr(_Converter, 'tool_outputs'): + _Converter.tool_outputs.clear() + + return True + + except Exception as e: + print(f"fix_message_list failed: {e}") + # Clean up even in case of failure + if hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls.clear() + if hasattr(_Converter, 'tool_outputs'): + _Converter.tool_outputs.clear() + return False + + def test_generic_linux_command_interrupt_simulation(self): + """Test generic_linux_command behavior during interruption.""" + + # Mock the generic_linux_command function behavior + def mock_interrupted_command(): + """Simulate generic_linux_command being interrupted""" + try: + # Simulate command starting + output = "Command started...\nProcessing files..." + + # Simulate interrupt during execution (like CTRL+C) + raise KeyboardInterrupt("User interrupted command") + + except KeyboardInterrupt: + # Simulate the real behavior - command returns partial output + interrupted_output = f"{output}\nCommand interrupted by user" + return interrupted_output + + # Test the mock + result = mock_interrupted_command() + + # Verify it behaves like the real interrupted command + assert "Command started" in result, "Should include partial output" + assert "interrupted" in result, "Should indicate interruption" + + print("✅ Generic linux command interrupt simulation test passed!") + + def test_message_history_openai_format_compliance(self): + """Test that message_history always maintains OpenAI ChatCompletion format.""" + from cai.sdk.agents.models.openai_chatcompletions import add_to_message_history + + # Clear history + message_history.clear() + + # Test various message types that should maintain OpenAI format + test_messages = [ + # User message + {"role": "user", "content": "Test user message"}, + + # Assistant message with content + {"role": "assistant", "content": "Test assistant response"}, + + # Assistant message with tool calls + { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": "call_test_123", + "type": "function", + "function": { + "name": "test_function", + "arguments": '{"param": "value"}' + } + }] + }, + + # Tool message + { + "role": "tool", + "tool_call_id": "call_test_123", + "content": "Tool execution result" + }, + + # System message + {"role": "system", "content": "You are a helpful assistant"} + ] + + # Add all messages + for msg in test_messages: + add_to_message_history(msg) + + # Verify OpenAI format compliance + assert len(message_history) == len(test_messages), \ + f"Expected {len(test_messages)} messages, got {len(message_history)}" + + for i, msg in enumerate(message_history): + # Required fields + assert 'role' in msg, f"Message {i} missing required 'role' field" + + # Valid roles + valid_roles = ['user', 'assistant', 'system', 'tool', 'developer'] + assert msg['role'] in valid_roles, \ + f"Message {i} has invalid role '{msg['role']}', must be one of {valid_roles}" + + # Role-specific validation + if msg['role'] == 'tool': + assert 'tool_call_id' in msg, f"Tool message {i} missing 'tool_call_id'" + assert 'content' in msg, f"Tool message {i} missing 'content'" + + if msg['role'] == 'assistant' and msg.get('tool_calls'): + assert isinstance(msg['tool_calls'], list), \ + f"Assistant message {i} tool_calls must be a list" + for j, tc in enumerate(msg['tool_calls']): + assert 'id' in tc, f"Tool call {j} in message {i} missing 'id'" + assert 'type' in tc, f"Tool call {j} in message {i} missing 'type'" + assert 'function' in tc, f"Tool call {j} in message {i} missing 'function'" + assert 'name' in tc['function'], \ + f"Tool call {j} function in message {i} missing 'name'" + assert 'arguments' in tc['function'], \ + f"Tool call {j} function in message {i} missing 'arguments'" + + print("✅ Message history OpenAI format compliance test passed!") + + def test_streaming_mode_configuration(self): + """Test streaming mode can be configured and detected.""" + # Test non-streaming mode + os.environ['CAI_STREAM'] = 'false' + assert os.environ['CAI_STREAM'] == 'false' + + # Test streaming mode + os.environ['CAI_STREAM'] = 'true' + assert os.environ['CAI_STREAM'] == 'true' + + print("✅ Streaming mode configuration test passed!") + + def test_multiple_interrupt_scenarios(self): + """Test multiple CTRL+C scenarios maintain consistency.""" + from cai.sdk.agents.models.openai_chatcompletions import add_to_message_history + + # Clear history + message_history.clear() + + scenarios = [ + ("Run first command", "call_1", "First command interrupted"), + ("Run second command", "call_2", "Second command interrupted"), + ("Run third command", "call_3", "Third command completed successfully") + ] + + for user_input, call_id, result_content in scenarios: + # Add user message + add_to_message_history({"role": "user", "content": user_input}) + + # Add assistant message with tool call + add_to_message_history({ + "role": "assistant", + "content": "I'll run that command for you.", + "tool_calls": [{ + "id": call_id, + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": f'{{"command": "test", "args": "{user_input}"}}' + } + }] + }) + + # Add tool result + add_to_message_history({ + "role": "tool", + "tool_call_id": call_id, + "content": result_content + }) + + # Verify consistency after each scenario + self.verify_message_history_openai_compliance() + + # Final verification + assert len(message_history) == len(scenarios) * 3 + print("✅ Multiple interrupt scenarios test passed!") + + def verify_message_history_openai_compliance(self): + """Helper method to verify message_history complies with OpenAI format.""" + for i, msg in enumerate(message_history): + # Basic structure checks + assert isinstance(msg, dict), f"Message {i} must be a dictionary" + assert 'role' in msg, f"Message {i} missing 'role' field" + + # Role validation + valid_roles = ['user', 'assistant', 'system', 'tool', 'developer'] + assert msg['role'] in valid_roles, \ + f"Message {i} role '{msg['role']}' not in valid roles {valid_roles}" + + # Content or tool_calls must exist for most roles + if msg['role'] in ['user', 'system', 'developer']: + assert 'content' in msg, f"Message {i} with role '{msg['role']}' missing content" + + elif msg['role'] == 'assistant': + # Assistant must have content OR tool_calls + has_content = 'content' in msg and msg['content'] is not None + has_tool_calls = 'tool_calls' in msg and msg['tool_calls'] + assert has_content or has_tool_calls, \ + f"Assistant message {i} must have content or tool_calls" + + elif msg['role'] == 'tool': + assert 'tool_call_id' in msg, f"Tool message {i} missing tool_call_id" + assert 'content' in msg, f"Tool message {i} missing content" + + def test_ctrl_c_during_tool_execution_real_behavior(self): + """Test real CTRL+C behavior during tool execution without duplicates.""" + from cai.sdk.agents.models.openai_chatcompletions import ( + add_to_message_history, + message_history, + _Converter + ) + + # COMPLETE cleanup - clear everything + message_history.clear() + if hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls.clear() + if hasattr(_Converter, 'tool_outputs'): + _Converter.tool_outputs.clear() + + # Simulate a running tool call that gets interrupted + call_id = "call_linux_cmd_123" + tool_name = "generic_linux_command" + + # 1. Add user message + add_to_message_history({"role": "user", "content": "Run a long command"}) + + # 2. Add assistant message with tool call (simulating qwen format) + add_to_message_history({ + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": call_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": '{"command": "sleep 10"}' + } + }] + }) + + # 3. Simulate CTRL+C cleanup behavior from cli.py + # Initialize _Converter attributes if they don't exist + if not hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls = {} + if not hasattr(_Converter, 'tool_outputs'): + _Converter.tool_outputs = {} + + # Add ONLY our specific tool call to recent_tool_calls + _Converter.recent_tool_calls[call_id] = { + 'name': tool_name, + 'arguments': '{"command": "sleep 10"}' + } + + # Simulate the KeyboardInterrupt cleanup logic from cli.py + try: + # Check for pending tool calls without responses + for call_id_check, call_info in list(_Converter.recent_tool_calls.items()): + # Check if tool response exists + tool_response_exists = any( + msg.get("role") == "tool" and msg.get("tool_call_id") == call_id_check + for msg in message_history + ) + + if not tool_response_exists: + # Add synthetic tool response (this is what cli.py does now) + tool_msg = { + "role": "tool", + "tool_call_id": call_id_check, + "content": "Operation interrupted by user (Keyboard Interrupt)" + } + add_to_message_history(tool_msg) + + # NOTE: The fix means we DON'T call fix_message_list here anymore + # This prevents duplicate synthetic tool calls + + except Exception as e: + print(f"Error in cleanup: {e}") + + # Verify message consistency after CTRL+C + print("Message history after CTRL+C cleanup:") + for i, msg in enumerate(message_history): + print(f" {i}: {msg.get('role')} - {msg}") + + # Assertions + self.assertEqual(len(message_history), 3) # user + assistant + tool + + # Check user message + self.assertEqual(message_history[0]["role"], "user") + + # Check assistant message has correct tool call + self.assertEqual(message_history[1]["role"], "assistant") + self.assertIsNotNone(message_history[1]["tool_calls"]) + self.assertEqual(len(message_history[1]["tool_calls"]), 1) + self.assertEqual(message_history[1]["tool_calls"][0]["id"], call_id) + self.assertEqual( + message_history[1]["tool_calls"][0]["function"]["name"], + tool_name + ) + + # Check tool response exists and is correct + self.assertEqual(message_history[2]["role"], "tool") + self.assertEqual(message_history[2]["tool_call_id"], call_id) + self.assertIn("interrupted", message_history[2]["content"].lower()) + + # MOST IMPORTANT: Verify NO duplicate tool calls with unknown_function + unknown_function_calls = [] + for msg in message_history: + if (msg.get("role") == "assistant" and + msg.get("tool_calls")): + for tc in msg["tool_calls"]: + if tc.get("function", {}).get("name") == "unknown_function": + unknown_function_calls.append(tc) + + self.assertEqual( + len(unknown_function_calls), 0, + f"Found {len(unknown_function_calls)} duplicate unknown_function calls: {unknown_function_calls}" + ) + + # Verify OpenAI format compliance + self.verify_message_history_openai_compliance() + + print("✓ CTRL+C test passed - no duplicates!") + + # Clean up + if hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls.clear() + if hasattr(_Converter, 'tool_outputs'): + _Converter.tool_outputs.clear() + + +if __name__ == '__main__': + print("🧪 Running simplified CLI streaming tests...") + + # Try to use unittest.main() first + try: + import unittest + if len(sys.argv) == 1: # No command line args, run all tests + unittest.main(verbosity=2, exit=False) + else: + # If there are command line args, run manual tests for debugging + # Create test instance + test_instance = TestCLIStreaming() + test_instance.setUpClass() + + # List of test methods - focused on direct testing without asyncio + test_methods = [ + 'test_streaming_mode_configuration', + 'test_message_history_openai_format_compliance', + 'test_ctrl_c_cleanup_message_consistency', + 'test_fix_message_list_with_interrupted_tools', + 'test_generic_linux_command_interrupt_simulation', + 'test_multiple_interrupt_scenarios', + 'test_ctrl_c_during_tool_execution_real_behavior' + ] + + results = {} + + for method_name in test_methods: + try: + print(f"\n🔬 Running {method_name}...") + # PROPERLY call setUp for each test + test_instance.setUp() + method = getattr(test_instance, method_name) + method() + results[method_name] = 'PASSED' + except Exception as e: + results[method_name] = f'FAILED: {str(e)}' + print(f"❌ {method_name} failed: {e}") + + # Print debug info for failures + if len(message_history) > 0: + print("Message history debug:") + for i, msg in enumerate(message_history): + print(f" [{i}] {msg.get('role', 'unknown')}: {str(msg.get('content', ''))[:50]}") + + # Print summary + print("\n" + "="*60) + print("📊 STREAMING TESTS SUMMARY") + print("="*60) + + passed = sum(1 for r in results.values() if r == 'PASSED') + failed = len(results) - passed + + for test_name, result in results.items(): + status_emoji = "✅" if result == 'PASSED' else "❌" + print(f"{status_emoji} {test_name}: {result}") + + print(f"\n🎯 Results: {passed} passed, {failed} failed") + + if failed == 0: + print("🎉 All simplified streaming tests passed!") + print("\n🔍 These tests verify:") + print("- Streaming mode configuration") + print("- Message history OpenAI format compliance") + print("- CTRL+C cleanup maintains message consistency") + print("- fix_message_list handles interrupted tools") + print("- Generic linux command interrupt simulation") + print("- Multiple interrupt scenarios") + print("- Real CTRL+C behavior during tool execution without duplicates") + else: + print(f"💥 {failed} streaming tests failed!") + sys.exit(1) + + except Exception as unittest_error: + print(f"Error running with unittest: {unittest_error}") + print("Falling back to manual test execution...") + + # Manual fallback + test_instance = TestCLIStreaming() + test_instance.setUpClass() + + test_methods = [ + 'test_streaming_mode_configuration', + 'test_message_history_openai_format_compliance', + 'test_ctrl_c_cleanup_message_consistency', + 'test_fix_message_list_with_interrupted_tools', + 'test_generic_linux_command_interrupt_simulation', + 'test_multiple_interrupt_scenarios', + 'test_ctrl_c_during_tool_execution_real_behavior' + ] + + results = {} + for method_name in test_methods: + try: + print(f"\n🔬 Running {method_name}...") + test_instance.setUp() # CRITICAL: Call setUp for each test + method = getattr(test_instance, method_name) + method() + results[method_name] = 'PASSED' + except Exception as e: + results[method_name] = f'FAILED: {str(e)}' + print(f"❌ {method_name} failed: {e}") + + passed = sum(1 for r in results.values() if r == 'PASSED') + failed = len(results) - passed + print(f"\n🎯 Manual Results: {passed} passed, {failed} failed") + if failed > 0: + sys.exit(1) \ No newline at end of file diff --git a/tests/commands/pytest.ini b/tests/commands/pytest.ini new file mode 100644 index 00000000..c3ff1bdb --- /dev/null +++ b/tests/commands/pytest.ini @@ -0,0 +1,20 @@ +[tool:pytest] +testpaths = . +python_files = test_*.py +python_classes = Test* +python_functions = test_* +addopts = + -v + --tb=short + --strict-markers + --disable-warnings + --color=yes + --durations=10 +markers = + integration: marks tests as integration tests (deselect with '-m "not integration"') + slow: marks tests as slow (deselect with '-m "not slow"') + unit: marks tests as unit tests +filterwarnings = + ignore::DeprecationWarning + ignore::PendingDeprecationWarning + ignore::UserWarning \ No newline at end of file diff --git a/tests/commands/test_command_agent.py b/tests/commands/test_command_agent.py new file mode 100644 index 00000000..5fca7199 --- /dev/null +++ b/tests/commands/test_command_agent.py @@ -0,0 +1,457 @@ +#!/usr/bin/env python3 +""" +Test suite for the agent command functionality. +Tests all handle methods and input possibilities for the agent command. +""" + +import os +import sys +import pytest +from unittest.mock import patch, Mock, MagicMock + +# Add src to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), + '..', '..', 'src')) + +from cai.repl.commands.agent import AgentCommand +from cai.repl.commands.base import Command + + +class TestAgentCommand: + """Test cases for AgentCommand.""" + + @pytest.fixture(autouse=True) + def setup_and_cleanup(self): + """Setup and cleanup for each test.""" + # Set up test environment + os.environ['CAI_TELEMETRY'] = 'false' + os.environ['CAI_TRACING'] = 'false' + + # Clear any agent-related environment variables + env_vars_to_clear = [ + 'CAI_AGENT_TYPE', 'CTF_MODEL', 'CAI_CODE_MODEL', + 'CAI_TEST_MODEL', 'CAI_CUSTOM_MODEL' + ] + for var in env_vars_to_clear: + if var in os.environ: + del os.environ[var] + + yield + + # Cleanup after each test + for var in env_vars_to_clear: + if var in os.environ: + del os.environ[var] + + @pytest.fixture + def agent_command(self): + """Create an AgentCommand instance for testing.""" + return AgentCommand() + + @pytest.fixture + def mock_agents(self): + """Create mock agents for testing.""" + agents = {} + + # Create mock agent objects with required attributes + for name in ['code', 'test', 'custom', 'basic']: + mock_agent = Mock() + mock_agent.name = name + mock_agent.model = f"model-for-{name}" + mock_agent.description = f"Description for {name} agent" + mock_agent.instructions = f"Instructions for {name} agent" + + # Configure properties that need len() to work + mock_agent.functions = [] # Empty list instead of Mock + mock_agent.handoffs = [] # Empty list instead of Mock + mock_agent.tools = [] # Empty list instead of Mock + mock_agent.input_guardrails = [] # Empty list instead of Mock + mock_agent.output_guardrails = [] # Empty list instead of Mock + mock_agent.hooks = [] # Empty list instead of Mock + + # Other optional properties + mock_agent.parallel_tool_calls = False + mock_agent.handoff_description = None + mock_agent.output_type = None + + agents[name] = mock_agent + + return agents + + def test_command_initialization(self, agent_command): + """Test that AgentCommand initializes correctly.""" + assert agent_command.name == "/agent" + assert agent_command.description == "Manage and switch between agents" + assert agent_command.aliases == ["/a"] + + # Check subcommands are available + expected_subcommands = ["list", "select", "info", "multi"] + assert set(agent_command.get_subcommands()) == set(expected_subcommands) + + def test_get_model_display_code_agent(self, agent_command): + """Test model display for code agent.""" + mock_agent = Mock() + mock_agent.model = "gpt-4" + + result = agent_command._get_model_display("code", mock_agent) + assert result == "gpt-4" + + def test_get_model_display_with_ctf_model(self, agent_command): + """Test model display when CTF_MODEL is set.""" + os.environ['CTF_MODEL'] = "claude-3" + + mock_agent = Mock() + mock_agent.model = "claude-3" + + result = agent_command._get_model_display("test", mock_agent) + assert result == "" # Should return empty for non-code agents with CTF_MODEL + + def test_get_model_display_with_env_var(self, agent_command): + """Test model display with agent-specific environment variable.""" + os.environ['CAI_TEST_MODEL'] = "custom-model" + + mock_agent = Mock() + mock_agent.model = "default-model" + + result = agent_command._get_model_display("test", mock_agent) + assert result == "custom-model" + + def test_get_model_display_for_info_code_agent(self, agent_command): + """Test model display for info view with code agent.""" + mock_agent = Mock() + mock_agent.model = "gpt-4" + + result = agent_command._get_model_display_for_info("code", mock_agent) + assert result == "gpt-4" + + def test_get_model_display_for_info_with_ctf_model(self, agent_command): + """Test model display for info view when CTF_MODEL is set.""" + os.environ['CTF_MODEL'] = "claude-3" + + mock_agent = Mock() + mock_agent.model = "claude-3" + + result = agent_command._get_model_display_for_info("test", mock_agent) + assert result == "Default CTF Model" + + @patch('cai.repl.commands.agent.get_available_agents') + @patch('cai.repl.commands.agent.get_agent_module') + def test_handle_list(self, mock_get_module, mock_get_agents, + agent_command, mock_agents): + """Test listing available agents.""" + mock_get_agents.return_value = mock_agents + mock_get_module.return_value = "test_module" + + result = agent_command.handle_list([]) + assert result is True + + # Verify get_available_agents was called + mock_get_agents.assert_called_once() + + @patch('cai.repl.commands.agent.get_available_agents') + @patch('cai.repl.commands.agent.visualize_agent_graph') + def test_handle_select_by_name(self, mock_visualize, mock_get_agents, + agent_command, mock_agents): + """Test selecting an agent by name.""" + mock_get_agents.return_value = mock_agents + + result = agent_command.handle_select(["code"]) + assert result is True + assert os.environ.get('CAI_AGENT_TYPE') == "code" + + # Verify visualization was called + mock_visualize.assert_called_once_with(mock_agents["code"]) + + @patch('cai.repl.commands.agent.get_available_agents') + @patch('cai.repl.commands.agent.visualize_agent_graph') + def test_handle_select_by_number(self, mock_visualize, mock_get_agents, + agent_command, mock_agents): + """Test selecting an agent by number.""" + mock_get_agents.return_value = mock_agents + + result = agent_command.handle_select(["2"]) + assert result is True + + # Should select the second agent in the dict (order may vary) + agent_keys = list(mock_agents.keys()) + expected_key = agent_keys[1] # Second agent (0-indexed) + assert os.environ.get('CAI_AGENT_TYPE') == expected_key + + @patch('cai.repl.commands.agent.get_available_agents') + def test_handle_select_invalid_name(self, mock_get_agents, + agent_command, mock_agents): + """Test selecting an invalid agent name.""" + mock_get_agents.return_value = mock_agents + + result = agent_command.handle_select(["invalid_agent"]) + assert result is False + assert 'CAI_AGENT_TYPE' not in os.environ + + @patch('cai.repl.commands.agent.get_available_agents') + def test_handle_select_invalid_number(self, mock_get_agents, + agent_command, mock_agents): + """Test selecting an invalid agent number.""" + mock_get_agents.return_value = mock_agents + + result = agent_command.handle_select(["99"]) + assert result is False + assert 'CAI_AGENT_TYPE' not in os.environ + + def test_handle_select_no_args(self, agent_command): + """Test select command with no arguments.""" + result = agent_command.handle_select([]) + assert result is False + + @patch('cai.repl.commands.agent.get_available_agents') + def test_handle_info_by_name(self, mock_get_agents, + agent_command, mock_agents): + """Test getting info for an agent by name.""" + mock_get_agents.return_value = mock_agents + + result = agent_command.handle_info(["code"]) + assert result is True + + @patch('cai.repl.commands.agent.get_available_agents') + def test_handle_info_by_number(self, mock_get_agents, + agent_command, mock_agents): + """Test getting info for an agent by number.""" + mock_get_agents.return_value = mock_agents + + result = agent_command.handle_info(["1"]) + assert result is True + + @patch('cai.repl.commands.agent.get_available_agents') + def test_handle_info_invalid_name(self, mock_get_agents, + agent_command, mock_agents): + """Test getting info for an invalid agent name.""" + mock_get_agents.return_value = mock_agents + + result = agent_command.handle_info(["invalid_agent"]) + assert result is False + + @patch('cai.repl.commands.agent.get_available_agents') + def test_handle_info_invalid_number(self, mock_get_agents, + agent_command, mock_agents): + """Test getting info for an invalid agent number.""" + mock_get_agents.return_value = mock_agents + + result = agent_command.handle_info(["99"]) + assert result is False + + def test_handle_info_no_args(self, agent_command): + """Test info command with no arguments.""" + result = agent_command.handle_info([]) + assert result is False + + @patch('cai.repl.commands.agent.get_available_agents') + def test_handle_info_with_complex_agent(self, mock_get_agents, agent_command): + """Test info command with agent that has complex attributes.""" + # Create a more complex mock agent + complex_agent = Mock() + complex_agent.name = "complex_agent" + complex_agent.description = "A complex agent for testing" + complex_agent.instructions = lambda: "Dynamic instructions" + # Use real lists instead of Mocks for len() to work + complex_agent.functions = [Mock(), Mock()] + complex_agent.parallel_tool_calls = True + complex_agent.handoff_description = "Handoff description" + complex_agent.handoffs = [Mock()] + complex_agent.tools = [Mock(), Mock(), Mock()] + complex_agent.input_guardrails = [Mock()] + complex_agent.output_guardrails = [Mock(), Mock()] + complex_agent.output_type = "str" + complex_agent.hooks = [Mock()] + + mock_get_agents.return_value = {"complex": complex_agent} + + result = agent_command.handle_info(["complex"]) + assert result is True + + def test_command_base_functionality(self, agent_command): + """Test that the command inherits from base Command properly.""" + assert isinstance(agent_command, Command) + assert agent_command.name == "/agent" + assert "/a" in agent_command.aliases + + @patch('cai.repl.commands.agent.get_available_agents') + @patch('cai.repl.commands.agent.get_agent_module') + def test_handle_main_command_routing(self, mock_get_module, mock_get_agents, + agent_command, mock_agents): + """Test that main handle method routes to correct subcommands.""" + mock_get_agents.return_value = mock_agents + mock_get_module.return_value = "test_module" + + # Test routing to list (no args defaults to list) + result1 = agent_command.handle([]) + assert result1 is True + + # Test routing to list explicitly + result2 = agent_command.handle(["list"]) + assert result2 is True + + # Test routing to info + result3 = agent_command.handle(["info", "code"]) + assert result3 is True + + # Test direct agent selection (not a subcommand) + result4 = agent_command.handle(["code"]) + assert result4 is True + assert os.environ.get('CAI_AGENT_TYPE') == "code" + + @patch('cai.repl.commands.agent.get_available_agents') + def test_agent_with_callable_instructions(self, mock_get_agents, agent_command): + """Test agent with callable instructions.""" + mock_agent = Mock() + mock_agent.name = "callable_agent" + mock_agent.description = "Agent with callable instructions" + mock_agent.instructions = lambda context_variables=None: "Callable instructions" + # Configure required properties + mock_agent.functions = [] + mock_agent.handoffs = [] + mock_agent.tools = [] + mock_agent.input_guardrails = [] + mock_agent.output_guardrails = [] + mock_agent.hooks = [] + mock_agent.parallel_tool_calls = False + mock_agent.handoff_description = None + mock_agent.output_type = None + + mock_get_agents.return_value = {"callable": mock_agent} + + result = agent_command.handle_info(["callable"]) + assert result is True + + @patch('cai.repl.commands.agent.get_available_agents') + def test_agent_with_multiline_description(self, mock_get_agents, agent_command): + """Test agent with multiline description that should be cleaned.""" + mock_agent = Mock() + mock_agent.name = "multiline_agent" + mock_agent.description = """This is a + multiline description + that should be cleaned""" + mock_agent.instructions = "Simple instructions" + # Configure required properties + mock_agent.functions = [] + mock_agent.handoffs = [] + mock_agent.tools = [] + mock_agent.input_guardrails = [] + mock_agent.output_guardrails = [] + mock_agent.hooks = [] + mock_agent.parallel_tool_calls = False + mock_agent.handoff_description = None + mock_agent.output_type = None + + mock_get_agents.return_value = {"multiline": mock_agent} + + result = agent_command.handle_info(["multiline"]) + assert result is True + + +@pytest.mark.integration +class TestAgentCommandIntegration: + """Integration tests for agent command functionality.""" + + @pytest.fixture(autouse=True) + def setup_integration(self): + """Setup for integration tests.""" + # Clear environment variables + env_vars_to_clear = [ + 'CAI_AGENT_TYPE', 'CTF_MODEL', 'CAI_CODE_MODEL', + 'CAI_TEST_MODEL', 'CAI_CUSTOM_MODEL' + ] + for var in env_vars_to_clear: + if var in os.environ: + del os.environ[var] + + yield + + # Cleanup + for var in env_vars_to_clear: + if var in os.environ: + del os.environ[var] + + @patch('cai.repl.commands.agent.get_available_agents') + @patch('cai.repl.commands.agent.get_agent_module') + @patch('cai.repl.commands.agent.visualize_agent_graph') + def test_full_workflow(self, mock_visualize, mock_get_module, mock_get_agents): + """Test a complete workflow of listing, selecting, and getting info.""" + # Setup mock agents + agents = {} + for name in ['agent1', 'agent2', 'agent3']: + mock_agent = Mock() + mock_agent.name = name + mock_agent.model = f"model-{name}" + mock_agent.description = f"Description for {name}" + mock_agent.instructions = f"Instructions for {name}" + + # Configure properties that need len() to work + mock_agent.functions = [] + mock_agent.handoffs = [] + mock_agent.tools = [] + mock_agent.input_guardrails = [] + mock_agent.output_guardrails = [] + mock_agent.hooks = [] + mock_agent.parallel_tool_calls = False + mock_agent.handoff_description = None + mock_agent.output_type = None + + agents[name] = mock_agent + + mock_get_agents.return_value = agents + mock_get_module.return_value = "test_module" + + cmd = AgentCommand() + + # List agents + result1 = cmd.handle(["list"]) + assert result1 is True + + # Select an agent by name + result2 = cmd.handle(["select", "agent1"]) + assert result2 is True + assert os.environ.get('CAI_AGENT_TYPE') == "agent1" + + # Get info for an agent + result3 = cmd.handle(["info", "agent2"]) + assert result3 is True + + # Select by number + result4 = cmd.handle(["select", "2"]) + assert result4 is True + + # Direct selection (not using select subcommand) + result5 = cmd.handle(["agent3"]) + assert result5 is True + assert os.environ.get('CAI_AGENT_TYPE') == "agent3" + + @patch('cai.repl.commands.agent.get_available_agents') + def test_environment_variable_handling(self, mock_get_agents): + """Test how environment variables affect model display.""" + mock_agent = Mock() + mock_agent.name = "test_agent" + mock_agent.model = "default-model" + + mock_get_agents.return_value = {"test": mock_agent} + + cmd = AgentCommand() + + # Test without environment variables + result1 = cmd._get_model_display("test", mock_agent) + assert result1 == "default-model" + + # Test with agent-specific environment variable + os.environ['CAI_TEST_MODEL'] = "env-specific-model" + result2 = cmd._get_model_display("test", mock_agent) + assert result2 == "env-specific-model" + + # Test with CTF_MODEL + os.environ['CTF_MODEL'] = "default-model" + result3 = cmd._get_model_display("test", mock_agent) + assert result3 == "" # Should be empty for table display + + result4 = cmd._get_model_display_for_info("test", mock_agent) + assert result4 == "Default CTF Model" # Should show this for info display + + +if __name__ == '__main__': + pytest.main([__file__, "-v"]) \ No newline at end of file diff --git a/tests/commands/test_command_base.py b/tests/commands/test_command_base.py new file mode 100644 index 00000000..857b37e2 --- /dev/null +++ b/tests/commands/test_command_base.py @@ -0,0 +1,528 @@ +#!/usr/bin/env python3 +""" +Test suite for the base command system functionality. +Tests the Command class, command registry, and base command handling. +""" + +import os +import sys +import pytest +from unittest.mock import patch, Mock, MagicMock + +# Add src to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), + '..', '..', 'src')) + +from cai.repl.commands.base import ( + Command, COMMANDS, COMMAND_ALIASES, + register_command, get_command, handle_command +) + + +class TestCommand: + """Test cases for the base Command class.""" + + @pytest.fixture(autouse=True) + def setup_and_cleanup(self): + """Setup and cleanup for each test.""" + # Store original command registry + self.original_commands = COMMANDS.copy() + self.original_aliases = COMMAND_ALIASES.copy() + + yield + + # Restore original command registry + COMMANDS.clear() + COMMANDS.update(self.original_commands) + COMMAND_ALIASES.clear() + COMMAND_ALIASES.update(self.original_aliases) + + @pytest.fixture + def sample_command(self): + """Create a sample command for testing.""" + return Command( + name="/test", + description="Test command for unit testing", + aliases=["/t", "/test-cmd"] + ) + + def test_command_initialization(self, sample_command): + """Test that Command initializes correctly.""" + assert sample_command.name == "/test" + assert sample_command.description == "Test command for unit testing" + assert sample_command.aliases == ["/t", "/test-cmd"] + assert isinstance(sample_command.subcommands, dict) + assert len(sample_command.subcommands) == 0 + + def test_command_initialization_without_aliases(self): + """Test Command initialization without aliases.""" + cmd = Command("/test", "Test command") + assert cmd.name == "/test" + assert cmd.description == "Test command" + assert cmd.aliases == [] + + def test_add_subcommand(self, sample_command): + """Test adding a subcommand to a command.""" + def test_handler(args): + return True + + sample_command.add_subcommand( + "test_sub", + "Test subcommand", + test_handler + ) + + assert "test_sub" in sample_command.subcommands + assert sample_command.subcommands["test_sub"]["description"] == "Test subcommand" + assert sample_command.subcommands["test_sub"]["handler"] == test_handler + + def test_get_subcommands(self, sample_command): + """Test getting list of subcommand names.""" + def handler1(args): + return True + def handler2(args): + return True + + sample_command.add_subcommand("sub1", "Description 1", handler1) + sample_command.add_subcommand("sub2", "Description 2", handler2) + + subcommands = sample_command.get_subcommands() + assert set(subcommands) == {"sub1", "sub2"} + + def test_get_subcommand_description(self, sample_command): + """Test getting subcommand description.""" + def test_handler(args): + return True + + sample_command.add_subcommand( + "test_sub", + "Test subcommand description", + test_handler + ) + + description = sample_command.get_subcommand_description("test_sub") + assert description == "Test subcommand description" + + # Test unknown subcommand + unknown_description = sample_command.get_subcommand_description("unknown") + assert unknown_description == "" + + def test_handle_with_subcommand(self, sample_command): + """Test handling a command with a valid subcommand.""" + def test_handler(args): + return True + + sample_command.add_subcommand("test_sub", "Test subcommand", test_handler) + + result = sample_command.handle(["test_sub"]) + assert result is True + + def test_handle_with_subcommand_and_args(self, sample_command): + """Test handling a command with subcommand and additional arguments.""" + def test_handler(args): + assert args == ["arg1", "arg2"] + return True + + sample_command.add_subcommand("test_sub", "Test subcommand", test_handler) + + result = sample_command.handle(["test_sub", "arg1", "arg2"]) + assert result is True + + def test_handle_no_args(self, sample_command): + """Test handling command with no arguments.""" + result = sample_command.handle([]) + assert result is False # Default implementation returns False + + def test_handle_unknown_subcommand(self, sample_command): + """Test handling command with unknown subcommand.""" + result = sample_command.handle(["unknown_subcommand"]) + assert result is False # Default implementation returns False + + def test_handle_no_args_default_behavior(self, sample_command): + """Test the default handle_no_args behavior.""" + result = sample_command.handle_no_args() + assert result is False + + def test_handle_unknown_subcommand_default_behavior(self, sample_command): + """Test the default handle_unknown_subcommand behavior.""" + result = sample_command.handle_unknown_subcommand("unknown") + assert result is False + + +class TestCommandRegistry: + """Test cases for the command registry system.""" + + @pytest.fixture(autouse=True) + def setup_and_cleanup(self): + """Setup and cleanup for each test.""" + # Store original command registry + self.original_commands = COMMANDS.copy() + self.original_aliases = COMMAND_ALIASES.copy() + + # Clear registry for clean tests + COMMANDS.clear() + COMMAND_ALIASES.clear() + + yield + + # Restore original command registry + COMMANDS.clear() + COMMANDS.update(self.original_commands) + COMMAND_ALIASES.clear() + COMMAND_ALIASES.update(self.original_aliases) + + @pytest.fixture + def test_commands(self): + """Create test commands for registry testing.""" + cmd1 = Command("/test1", "First test command", ["/t1"]) + cmd2 = Command("/test2", "Second test command", ["/t2", "/test-two"]) + cmd3 = Command("/test3", "Third test command") + return [cmd1, cmd2, cmd3] + + def test_register_command(self, test_commands): + """Test registering a command.""" + cmd = test_commands[0] + register_command(cmd) + + assert cmd.name in COMMANDS + assert COMMANDS[cmd.name] == cmd + assert "/t1" in COMMAND_ALIASES + assert COMMAND_ALIASES["/t1"] == cmd.name + + def test_register_multiple_commands(self, test_commands): + """Test registering multiple commands.""" + for cmd in test_commands: + register_command(cmd) + + # Check all commands are registered + assert len(COMMANDS) == 3 + for cmd in test_commands: + assert cmd.name in COMMANDS + assert COMMANDS[cmd.name] == cmd + + # Check aliases + assert COMMAND_ALIASES["/t1"] == "/test1" + assert COMMAND_ALIASES["/t2"] == "/test2" + assert COMMAND_ALIASES["/test-two"] == "/test2" + + def test_register_command_with_duplicate_name(self, test_commands): + """Test registering commands with duplicate names (should overwrite).""" + cmd1 = test_commands[0] + cmd2 = Command("/test1", "Different description") + + register_command(cmd1) + register_command(cmd2) # Should overwrite cmd1 + + assert COMMANDS["/test1"] == cmd2 + assert COMMANDS["/test1"].description == "Different description" + + def test_get_command_by_name(self, test_commands): + """Test getting a command by its name.""" + cmd = test_commands[0] + register_command(cmd) + + retrieved_cmd = get_command("/test1") + assert retrieved_cmd == cmd + + def test_get_command_by_alias(self, test_commands): + """Test getting a command by its alias.""" + cmd = test_commands[0] + register_command(cmd) + + retrieved_cmd = get_command("/t1") + assert retrieved_cmd == cmd + + def test_get_command_nonexistent(self): + """Test getting a non-existent command.""" + result = get_command("/nonexistent") + assert result is None + + def test_handle_command_by_name(self, test_commands): + """Test handling a command by its name.""" + cmd = test_commands[0] + cmd.handle = Mock(return_value=True) + register_command(cmd) + + result = handle_command("/test1", ["arg1", "arg2"]) + assert result is True + cmd.handle.assert_called_once_with(["arg1", "arg2"]) + + def test_handle_command_by_alias(self, test_commands): + """Test handling a command by its alias.""" + cmd = test_commands[0] + cmd.handle = Mock(return_value=True) + register_command(cmd) + + result = handle_command("/t1", ["arg1"]) + assert result is True + cmd.handle.assert_called_once_with(["arg1"]) + + def test_handle_command_nonexistent(self): + """Test handling a non-existent command.""" + result = handle_command("/nonexistent", ["args"]) + assert result is False + + def test_handle_command_no_args(self, test_commands): + """Test handling a command with no arguments.""" + cmd = test_commands[0] + cmd.handle = Mock(return_value=True) + register_command(cmd) + + result = handle_command("/test1") + assert result is True + cmd.handle.assert_called_once_with(None) + + +class TestCustomCommand: + """Test cases using custom command implementations.""" + + @pytest.fixture(autouse=True) + def setup_and_cleanup(self): + """Setup and cleanup for each test.""" + # Store original command registry + self.original_commands = COMMANDS.copy() + self.original_aliases = COMMAND_ALIASES.copy() + + # Clear registry for clean tests + COMMANDS.clear() + COMMAND_ALIASES.clear() + + yield + + # Restore original command registry + COMMANDS.clear() + COMMANDS.update(self.original_commands) + COMMAND_ALIASES.clear() + COMMAND_ALIASES.update(self.original_aliases) + + def test_custom_command_with_overridden_methods(self): + """Test a custom command with overridden handle methods.""" + + class CustomCommand(Command): + def __init__(self): + super().__init__("/custom", "Custom test command", ["/c"]) + self.handle_no_args_called = False + self.handle_unknown_subcommand_called = False + + def handle_no_args(self): + self.handle_no_args_called = True + return True + + def handle_unknown_subcommand(self, subcommand): + self.handle_unknown_subcommand_called = True + self.last_unknown_subcommand = subcommand + return True + + cmd = CustomCommand() + register_command(cmd) + + # Test handle_no_args + result1 = handle_command("/custom") + assert result1 is True + assert cmd.handle_no_args_called is True + + # Test handle_unknown_subcommand + result2 = handle_command("/custom", ["unknown"]) + assert result2 is True + assert cmd.handle_unknown_subcommand_called is True + assert cmd.last_unknown_subcommand == "unknown" + + def test_custom_command_with_subcommands(self): + """Test a custom command with predefined subcommands.""" + + class CustomCommandWithSubcommands(Command): + def __init__(self): + super().__init__("/multi", "Multi-subcommand test", ["/m"]) + self.add_subcommand("start", "Start something", self.handle_start) + self.add_subcommand("stop", "Stop something", self.handle_stop) + self.add_subcommand("status", "Check status", self.handle_status) + + self.start_called = False + self.stop_called = False + self.status_called = False + + def handle_start(self, args): + self.start_called = True + self.start_args = args + return True + + def handle_stop(self, args): + self.stop_called = True + self.stop_args = args + return True + + def handle_status(self, args): + self.status_called = True + self.status_args = args + return True + + cmd = CustomCommandWithSubcommands() + register_command(cmd) + + # Test each subcommand + result1 = handle_command("/multi", ["start", "param1"]) + assert result1 is True + assert cmd.start_called is True + assert cmd.start_args == ["param1"] + + result2 = handle_command("/m", ["stop"]) # Test alias + assert result2 is True + assert cmd.stop_called is True + assert cmd.stop_args is None # When no args after subcommand, None is passed + + result3 = handle_command("/multi", ["status", "verbose"]) + assert result3 is True + assert cmd.status_called is True + assert cmd.status_args == ["verbose"] + + def test_command_error_handling(self): + """Test command error handling when handlers raise exceptions.""" + + class ErrorCommand(Command): + def __init__(self): + super().__init__("/error", "Error test command") + self.add_subcommand("crash", "Crash handler", self.handle_crash) + + def handle_crash(self, args): + raise ValueError("Test error") + + cmd = ErrorCommand() + register_command(cmd) + + # The command should propagate the exception + with pytest.raises(ValueError, match="Test error"): + handle_command("/error", ["crash"]) + + +@pytest.mark.integration +class TestCommandIntegration: + """Integration tests for the command system.""" + + @pytest.fixture(autouse=True) + def setup_integration(self): + """Setup for integration tests.""" + # Store original command registry + self.original_commands = COMMANDS.copy() + self.original_aliases = COMMAND_ALIASES.copy() + + # Clear registry for clean tests + COMMANDS.clear() + COMMAND_ALIASES.clear() + + yield + + # Restore original command registry + COMMANDS.clear() + COMMANDS.update(self.original_commands) + COMMAND_ALIASES.clear() + COMMAND_ALIASES.update(self.original_aliases) + + def test_complete_command_lifecycle(self): + """Test the complete lifecycle of commands.""" + # Create multiple commands with various features + class Command1(Command): + def __init__(self): + super().__init__("/cmd1", "First command", ["/c1"]) + self.add_subcommand("action", "Do action", self.handle_action) + self.action_count = 0 + + def handle_action(self, args): + self.action_count += 1 + return True + + class Command2(Command): + def __init__(self): + super().__init__("/cmd2", "Second command", ["/c2", "/second"]) + self.call_count = 0 + + def handle_no_args(self): + self.call_count += 1 + return True + + # Register commands + cmd1 = Command1() + cmd2 = Command2() + register_command(cmd1) + register_command(cmd2) + + # Test command registry state + assert len(COMMANDS) == 2 + assert len(COMMAND_ALIASES) == 3 + + # Test various command executions + assert handle_command("/cmd1", ["action"]) is True + assert cmd1.action_count == 1 + + assert handle_command("/c1", ["action"]) is True # Using alias + assert cmd1.action_count == 2 + + assert handle_command("/cmd2") is True + assert cmd2.call_count == 1 + + assert handle_command("/c2") is True # Using alias + assert cmd2.call_count == 2 + + assert handle_command("/second") is True # Using second alias + assert cmd2.call_count == 3 + + # Test non-existent command + assert handle_command("/nonexistent") is False + + def test_command_alias_conflicts(self): + """Test handling of alias conflicts (last registered wins).""" + cmd1 = Command("/cmd1", "First command", ["/shared"]) + cmd2 = Command("/cmd2", "Second command", ["/shared"]) # Same alias + + register_command(cmd1) + register_command(cmd2) # This should overwrite the alias + + # The alias should point to the last registered command + assert COMMAND_ALIASES["/shared"] == "/cmd2" + + retrieved_cmd = get_command("/shared") + assert retrieved_cmd == cmd2 + + def test_complex_subcommand_routing(self): + """Test complex subcommand routing scenarios.""" + + class ComplexCommand(Command): + def __init__(self): + super().__init__("/complex", "Complex command") + self.add_subcommand("sub1", "Subcommand 1", self.handle_sub1) + self.add_subcommand("sub2", "Subcommand 2", self.handle_sub2) + + self.results = {} + + def handle_sub1(self, args): + self.results["sub1"] = args + return True + + def handle_sub2(self, args): + self.results["sub2"] = args + return False # Return False to test error handling + + def handle_unknown_subcommand(self, subcommand): + self.results["unknown"] = subcommand + return True + + cmd = ComplexCommand() + register_command(cmd) + + # Test subcommand 1 + result1 = handle_command("/complex", ["sub1", "arg1", "arg2"]) + assert result1 is True + assert cmd.results["sub1"] == ["arg1", "arg2"] + + # Test subcommand 2 (returns False) + result2 = handle_command("/complex", ["sub2", "test"]) + assert result2 is False + assert cmd.results["sub2"] == ["test"] + + # Test unknown subcommand + result3 = handle_command("/complex", ["unknown_sub"]) + assert result3 is True + assert cmd.results["unknown"] == "unknown_sub" + + +if __name__ == '__main__': + pytest.main([__file__, "-v"]) \ No newline at end of file diff --git a/tests/commands/test_command_config.py b/tests/commands/test_command_config.py new file mode 100644 index 00000000..b0143cc2 --- /dev/null +++ b/tests/commands/test_command_config.py @@ -0,0 +1,381 @@ +#!/usr/bin/env python3 +""" +Test suite for the config command functionality. +Tests all handle methods and input possibilities for the config command. +""" + +import os +import sys +import pytest +from unittest.mock import patch, Mock, MagicMock + +# Add src to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), + '..', '..', 'src')) + +from cai.repl.commands.config import ( + ConfigCommand, ENV_VARS, get_env_var_value, set_env_var +) +from cai.repl.commands.base import Command + + +class TestConfigCommand: + """Test cases for ConfigCommand.""" + + @pytest.fixture(autouse=True) + def setup_and_cleanup(self): + """Setup and cleanup for each test.""" + # Set up test environment + os.environ['CAI_TELEMETRY'] = 'false' + os.environ['CAI_TRACING'] = 'false' + + # Store original values of environment variables we'll modify + self.original_env_vars = {} + for var_info in ENV_VARS.values(): + var_name = var_info["name"] + if var_name in os.environ: + self.original_env_vars[var_name] = os.environ[var_name] + + yield + + # Restore original environment variables + for var_info in ENV_VARS.values(): + var_name = var_info["name"] + if var_name in self.original_env_vars: + os.environ[var_name] = self.original_env_vars[var_name] + elif var_name in os.environ: + del os.environ[var_name] + + @pytest.fixture + def config_command(self): + """Create a ConfigCommand instance for testing.""" + return ConfigCommand() + + def test_command_initialization(self, config_command): + """Test that ConfigCommand initializes correctly.""" + assert config_command.name == "/config" + assert config_command.description == "Display and configure environment variables" + assert config_command.aliases == ["/cfg"] + + # Check subcommands are registered + expected_subcommands = ["list", "set", "get"] + assert set(config_command.get_subcommands()) == set(expected_subcommands) + + def test_env_vars_structure(self): + """Test that ENV_VARS has the expected structure.""" + assert isinstance(ENV_VARS, dict) + assert len(ENV_VARS) > 0 + + # Check each env var has required fields + for num, var_info in ENV_VARS.items(): + assert isinstance(num, int) + assert "name" in var_info + assert "description" in var_info + assert "default" in var_info + + # Check name is a string + assert isinstance(var_info["name"], str) + assert isinstance(var_info["description"], str) + # default can be None or string + assert var_info["default"] is None or isinstance(var_info["default"], str) + + def test_get_env_var_value_with_set_value(self): + """Test getting environment variable value when it's set.""" + # Set a test value + os.environ["CAI_MODEL"] = "test-model" + + result = get_env_var_value("CAI_MODEL") + assert result == "test-model" + + def test_get_env_var_value_with_default(self): + """Test getting environment variable value when it's not set (returns default).""" + # Make sure the variable is not set + if "CAI_MODEL" in os.environ: + del os.environ["CAI_MODEL"] + + result = get_env_var_value("CAI_MODEL") + # Should return the default value from ENV_VARS + expected_default = None + for var_info in ENV_VARS.values(): + if var_info["name"] == "CAI_MODEL": + expected_default = var_info["default"] + break + + assert result == (expected_default or "Not set") + + def test_get_env_var_value_unknown_variable(self): + """Test getting value for unknown environment variable.""" + result = get_env_var_value("UNKNOWN_VARIABLE") + assert result == "Unknown variable" + + def test_set_env_var(self): + """Test setting an environment variable.""" + result = set_env_var("TEST_VAR", "test_value") + assert result is True + assert os.environ.get("TEST_VAR") == "test_value" + + # Cleanup + if "TEST_VAR" in os.environ: + del os.environ["TEST_VAR"] + + def test_handle_list(self, config_command): + """Test listing all environment variables.""" + result = config_command.handle_list([]) + assert result is True + + def test_handle_no_args(self, config_command): + """Test handling when no arguments provided (should default to list).""" + result = config_command.handle_no_args() + assert result is True + + def test_handle_get_valid_number(self, config_command): + """Test getting a variable by valid number.""" + # Test with variable number 6 (CAI_MODEL) + result = config_command.handle_get(["6"]) + assert result is True + + def test_handle_get_invalid_number(self, config_command): + """Test getting a variable by invalid number.""" + result = config_command.handle_get(["999"]) + assert result is False + + def test_handle_get_non_integer(self, config_command): + """Test getting a variable with non-integer number.""" + result = config_command.handle_get(["not_a_number"]) + assert result is False + + def test_handle_get_no_args(self, config_command): + """Test get command with no arguments.""" + result = config_command.handle_get([]) + assert result is False + + def test_handle_set_valid_number_and_value(self, config_command): + """Test setting a variable by valid number and value.""" + # Test with variable number 6 (CAI_MODEL) + result = config_command.handle_set(["6", "new-test-model"]) + assert result is True + assert os.environ.get("CAI_MODEL") == "new-test-model" + + def test_handle_set_invalid_number(self, config_command): + """Test setting a variable by invalid number.""" + result = config_command.handle_set(["999", "some_value"]) + assert result is False + + def test_handle_set_non_integer(self, config_command): + """Test setting a variable with non-integer number.""" + result = config_command.handle_set(["not_a_number", "some_value"]) + assert result is False + + def test_handle_set_no_args(self, config_command): + """Test set command with no arguments.""" + result = config_command.handle_set([]) + assert result is False + + def test_handle_set_insufficient_args(self, config_command): + """Test set command with insufficient arguments.""" + result = config_command.handle_set(["6"]) # Missing value + assert result is False + + def test_handle_set_with_spaces_in_value(self, config_command): + """Test setting a variable with value containing spaces.""" + result = config_command.handle_set(["6", "model with spaces"]) + assert result is True + assert os.environ.get("CAI_MODEL") == "model with spaces" + + def test_handle_set_empty_value(self, config_command): + """Test setting a variable with empty value.""" + result = config_command.handle_set(["6", ""]) + assert result is True + assert os.environ.get("CAI_MODEL") == "" + + def test_command_base_functionality(self, config_command): + """Test that the command inherits from base Command properly.""" + assert isinstance(config_command, Command) + assert config_command.name == "/config" + assert "/cfg" in config_command.aliases + + def test_handle_main_command_routing(self, config_command): + """Test that main handle method routes to correct subcommands.""" + # Test routing to list (no args defaults to list) + result1 = config_command.handle([]) + assert result1 is True + + # Test routing to list explicitly + result2 = config_command.handle(["list"]) + assert result2 is True + + # Test routing to get + result3 = config_command.handle(["get", "6"]) + assert result3 is True + + # Test routing to set + result4 = config_command.handle(["set", "6", "test-value"]) + assert result4 is True + assert os.environ.get("CAI_MODEL") == "test-value" + + def test_handle_unknown_subcommand(self, config_command): + """Test handling of unknown subcommands.""" + result = config_command.handle(["unknown_subcommand"]) + assert result is False + + def test_specific_env_vars_exist(self): + """Test that specific important environment variables are defined.""" + important_vars = [ + "CAI_MODEL", "CAI_DEBUG", "CAI_BRIEF", "CAI_MAX_TURNS", + "CAI_TRACING", "CAI_AGENT_TYPE", "CTF_NAME", "CTF_CHALLENGE" + ] + + defined_var_names = [var_info["name"] for var_info in ENV_VARS.values()] + + for var_name in important_vars: + assert var_name in defined_var_names, f"{var_name} should be defined in ENV_VARS" + + def test_env_var_defaults_are_reasonable(self): + """Test that environment variable defaults are reasonable.""" + # Check that important variables have reasonable defaults + for var_info in ENV_VARS.values(): + var_name = var_info["name"] + default = var_info["default"] + + # Boolean-like variables should have "true"/"false" defaults or numeric values + # Exclude interval variables as they have numeric values + boolean_keywords = ["debug", "brief", "tracing", "memory", "online", "offline", "inside"] + if (any(keyword in var_name.lower() for keyword in boolean_keywords) and + "interval" not in var_name.lower()): + if default is not None: + # Accept boolean strings or numeric values (for debug levels) + valid_values = ["true", "false", "0", "1", "2"] + assert default.lower() in valid_values, f"{var_name} should have boolean-like or numeric default" + + # Numeric variables should have numeric defaults + if any(keyword in var_name.lower() for keyword in ["turns", "limit", "interval"]): + if default is not None: + # Should be numeric or "inf" + try: + float(default) + except ValueError: + assert default == "inf", f"{var_name} should have numeric default or 'inf'" + + +@pytest.mark.integration +class TestConfigCommandIntegration: + """Integration tests for config command functionality.""" + + @pytest.fixture(autouse=True) + def setup_integration(self): + """Setup for integration tests.""" + # Store original values + self.original_env_vars = {} + for var_info in ENV_VARS.values(): + var_name = var_info["name"] + if var_name in os.environ: + self.original_env_vars[var_name] = os.environ[var_name] + + yield + + # Restore original values + for var_info in ENV_VARS.values(): + var_name = var_info["name"] + if var_name in self.original_env_vars: + os.environ[var_name] = self.original_env_vars[var_name] + elif var_name in os.environ: + del os.environ[var_name] + + def test_full_config_workflow(self): + """Test a complete workflow of listing, getting, and setting variables.""" + cmd = ConfigCommand() + + # List all variables + result1 = cmd.handle(["list"]) + assert result1 is True + + # Get a specific variable (CAI_MODEL - number 6) + result2 = cmd.handle(["get", "6"]) + assert result2 is True + + # Set the variable to a new value + result3 = cmd.handle(["set", "6", "integration-test-model"]) + assert result3 is True + assert os.environ.get("CAI_MODEL") == "integration-test-model" + + # Get the variable again to verify it changed + result4 = cmd.handle(["get", "6"]) + assert result4 is True + + # Set it back to default (if it had a default) + for var_info in ENV_VARS.values(): + if var_info["name"] == "CAI_MODEL" and var_info["default"]: + result5 = cmd.handle(["set", "6", var_info["default"]]) + assert result5 is True + break + + def test_multiple_variable_modifications(self): + """Test modifying multiple variables in sequence.""" + cmd = ConfigCommand() + + # Modify several variables + modifications = [ + ("6", "test-model"), # CAI_MODEL + ("7", "2"), # CAI_DEBUG + ("8", "true"), # CAI_BRIEF + ] + + for var_num, value in modifications: + result = cmd.handle(["set", var_num, value]) + assert result is True + + # Verify the change + get_result = cmd.handle(["get", var_num]) + assert get_result is True + + # Verify all changes are still present + for var_num, expected_value in modifications: + # Find the variable name + var_info = ENV_VARS[int(var_num)] + var_name = var_info["name"] + assert os.environ.get(var_name) == expected_value + + def test_edge_case_values(self): + """Test setting variables with edge case values.""" + cmd = ConfigCommand() + + edge_cases = [ + ("6", ""), # Empty string + ("6", "value with spaces"), # Spaces + ("6", "special!@#$%chars"), # Special characters + ("6", "very_long_value_that_exceeds_normal_length_expectations"), # Long value + ("7", "0"), # Zero + ("7", "999"), # Large number + ] + + for var_num, value in edge_cases: + result = cmd.handle(["set", var_num, value]) + assert result is True + + # Find the variable name and verify + var_info = ENV_VARS[int(var_num)] + var_name = var_info["name"] + assert os.environ.get(var_name) == value + + def test_all_env_vars_can_be_set_and_retrieved(self): + """Test that all defined environment variables can be set and retrieved.""" + cmd = ConfigCommand() + + for var_num, var_info in ENV_VARS.items(): + var_name = var_info["name"] + test_value = f"test_value_for_{var_name}" + + # Set the variable + set_result = cmd.handle(["set", str(var_num), test_value]) + assert set_result is True, f"Failed to set {var_name}" + + # Get the variable + get_result = cmd.handle(["get", str(var_num)]) + assert get_result is True, f"Failed to get {var_name}" + + # Verify the value was set correctly + assert os.environ.get(var_name) == test_value + + +if __name__ == '__main__': + pytest.main([__file__, "-v"]) \ No newline at end of file diff --git a/tests/commands/test_command_help.py b/tests/commands/test_command_help.py new file mode 100644 index 00000000..1dafcd61 --- /dev/null +++ b/tests/commands/test_command_help.py @@ -0,0 +1,402 @@ +#!/usr/bin/env python3 + +""" +Test suite for the help command functionality. +Tests all handle methods and input possibilities for the help command. +""" + +import pytest +from unittest.mock import patch, Mock, MagicMock +from rich.console import Console +from rich.table import Table +from rich.panel import Panel + +from cai.repl.commands.help import HelpCommand +from cai.repl.commands.base import Command, COMMANDS, COMMAND_ALIASES + + +class TestHelpCommand: + """Test class for HelpCommand functionality.""" + + @pytest.fixture + def help_command(self): + """Create a HelpCommand instance for testing.""" + return HelpCommand() + + @pytest.fixture + def mock_console(self): + """Create a mock console for testing output.""" + with patch('cai.repl.commands.help.console') as mock_console: + yield mock_console + + @pytest.fixture + def mock_commands_registry(self): + """Create a mock commands registry for testing.""" + mock_registry = { + '/memory': Mock(name='/memory', description='Memory commands'), + '/help': Mock(name='/help', description='Help commands'), + '/agent': Mock(name='/agent', description='Agent commands') + } + + with patch('cai.repl.commands.help.COMMANDS', mock_registry): + yield mock_registry + + @pytest.fixture + def mock_aliases_registry(self): + """Create a mock aliases registry for testing.""" + mock_aliases = { + '/h': '/help', + '/m': '/memory', + '/a': '/agent' + } + + with patch('cai.repl.commands.help.COMMAND_ALIASES', mock_aliases): + yield mock_aliases + + def test_command_initialization(self, help_command): + """Test that HelpCommand initializes correctly.""" + assert isinstance(help_command, Command) + assert help_command.name == "/help" + assert "Display help information about commands and features" in help_command.description + assert "/h" in help_command.aliases + + # Check that all expected subcommands are registered + expected_subcommands = [ + "memory", "agents", "graph", "platform", "shell", + "env", "aliases", "model", "turns", "config" + ] + for subcommand in expected_subcommands: + assert subcommand in help_command.subcommands + + def test_handle_no_args(self, help_command, mock_console): + """Test handling help command with no arguments.""" + result = help_command.handle_no_args() + + assert result is True + # Should print multiple panels/tables + assert mock_console.print.call_count >= 5 + + def test_handle_memory_subcommand(self, help_command, mock_console): + """Test memory subcommand help.""" + # Mock a memory command + mock_memory_cmd = Mock() + mock_memory_cmd.name = "/memory" + mock_memory_cmd.show_help = Mock(return_value=True) + + with patch('cai.repl.commands.help.COMMANDS', {'/memory': mock_memory_cmd}): + result = help_command.handle_memory() + + assert result is True + # Should call the memory command's show_help if available + mock_memory_cmd.show_help.assert_called_once() + + def test_handle_memory_subcommand_fallback(self, help_command, mock_console): + """Test memory subcommand fallback when show_help not available.""" + # Mock memory command without show_help + mock_memory_cmd = Mock() + mock_memory_cmd.name = "/memory" + # Remove show_help attribute + del mock_memory_cmd.show_help + + with patch('cai.repl.commands.help.COMMANDS', {'/memory': mock_memory_cmd}): + result = help_command.handle_memory() + + assert result is True + # Should print fallback help + assert mock_console.print.call_count >= 1 + + def test_handle_agents_subcommand(self, help_command, mock_console): + """Test agents subcommand help.""" + result = help_command.handle_agents() + + assert result is True + mock_console.print.assert_called_once() + + # Verify the content contains agent-related information + call_args = mock_console.print.call_args[0][0] + assert hasattr(call_args, 'renderable') + panel_content = str(call_args.renderable) + assert "Agent Commands" in panel_content or "agent" in panel_content.lower() + assert "/agent list" in panel_content + + def test_handle_graph_subcommand(self, help_command, mock_console): + """Test graph subcommand help.""" + result = help_command.handle_graph() + + assert result is True + mock_console.print.assert_called_once() + + # Verify graph-related content + call_args = mock_console.print.call_args[0][0] + assert hasattr(call_args, 'renderable') + panel_content = str(call_args.renderable) + assert "Graph" in panel_content or "graph" in panel_content.lower() + assert "/graph show" in panel_content or "graph" in panel_content.lower() + + def test_handle_platform_subcommand_with_show_help(self, help_command, mock_console): + """Test platform subcommand when platform command has show_help.""" + mock_platform_cmd = Mock() + mock_platform_cmd.name = "/platform" + mock_platform_cmd.show_help = Mock(return_value=True) + + with patch('cai.repl.commands.help.COMMANDS', {'/platform': mock_platform_cmd}): + result = help_command.handle_platform() + + assert result is True + mock_platform_cmd.show_help.assert_called_once() + + def test_handle_platform_subcommand_fallback(self, help_command, mock_console): + """Test platform subcommand fallback.""" + with patch('cai.repl.commands.help.COMMANDS', {}): + result = help_command.handle_platform() + + assert result is True + mock_console.print.assert_called_once() + + call_args = mock_console.print.call_args[0][0] + assert hasattr(call_args, 'renderable') + panel_content = str(call_args.renderable) + assert "Platform" in panel_content or "platform" in panel_content.lower() + + def test_handle_shell_subcommand(self, help_command, mock_console): + """Test shell subcommand help.""" + result = help_command.handle_shell() + + assert result is True + mock_console.print.assert_called_once() + + call_args = mock_console.print.call_args[0][0] + assert hasattr(call_args, 'renderable') + panel_content = str(call_args.renderable) + assert "Shell" in panel_content or "shell" in panel_content.lower() + assert "/shell " in panel_content or "shell" in panel_content.lower() + + def test_handle_env_subcommand(self, help_command, mock_console): + """Test env subcommand help.""" + result = help_command.handle_env() + + assert result is True + mock_console.print.assert_called_once() + + call_args = mock_console.print.call_args[0][0] + assert hasattr(call_args, 'renderable') + panel_content = str(call_args.renderable) + assert "Environment" in panel_content or "environment" in panel_content.lower() + assert "CAI_MODEL" in panel_content + + def test_handle_aliases_subcommand(self, help_command, mock_console, + mock_aliases_registry): + """Test aliases subcommand help.""" + result = help_command.handle_aliases() + + assert result is True + # Should print multiple times (header, table, tips) + assert mock_console.print.call_count >= 2 + + def test_handle_model_subcommand(self, help_command, mock_console): + """Test model subcommand help.""" + result = help_command.handle_model() + + assert result is True + # Should print multiple panels/tables + assert mock_console.print.call_count >= 2 + + def test_handle_turns_subcommand(self, help_command, mock_console): + """Test turns subcommand help.""" + result = help_command.handle_turns() + + assert result is True + assert mock_console.print.call_count >= 2 + + def test_handle_config_subcommand(self, help_command, mock_console): + """Test config subcommand help.""" + result = help_command.handle_config() + + assert result is True + assert mock_console.print.call_count >= 2 + + def test_handle_help_aliases(self, help_command, mock_console, + mock_commands_registry, mock_aliases_registry): + """Test handle_help_aliases method directly.""" + result = help_command.handle_help_aliases() + + assert result is True + # Should print header, table, and tips + assert mock_console.print.call_count >= 3 + + def test_handle_help_memory(self, help_command, mock_console): + """Test handle_help_memory method directly.""" + result = help_command.handle_help_memory() + + assert result is True + # Should print multiple panels/tables + assert mock_console.print.call_count >= 4 + + def test_handle_help_model(self, help_command, mock_console): + """Test handle_help_model method directly.""" + result = help_command.handle_help_model() + + assert result is True + # Should print multiple panels/tables + assert mock_console.print.call_count >= 4 + + def test_handle_help_turns(self, help_command, mock_console): + """Test handle_help_turns method directly.""" + result = help_command.handle_help_turns() + + assert result is True + # Should print multiple panels/tables + assert mock_console.print.call_count >= 3 + + def test_handle_help_config(self, help_command, mock_console): + """Test handle_help_config method directly.""" + result = help_command.handle_help_config() + + assert result is True + # Should print header, table, and notes + assert mock_console.print.call_count >= 3 + + def test_handle_help_platform_manager_with_extensions(self, help_command, + mock_console): + """Test platform manager help with extensions available.""" + # Mock platform extensions + mock_platform_manager = Mock() + mock_platform_manager.list_platforms.return_value = ['test_platform'] + mock_platform = Mock() + mock_platform.description = "Test platform" + mock_platform.get_commands.return_value = ['test_command'] + mock_platform_manager.get_platform.return_value = mock_platform + + with patch('cai.repl.commands.help.HAS_PLATFORM_EXTENSIONS', True): + with patch('cai.repl.commands.help.is_caiextensions_platform_available', + return_value=True): + # Mock the platform manager without importing caiextensions + with patch('sys.modules', {'caiextensions.platform.base': Mock(platform_manager=mock_platform_manager)}): + result = help_command.handle_help_platform_manager() + + assert result is True + assert mock_console.print.call_count >= 1 + + def test_handle_help_platform_manager_no_extensions(self, help_command, + mock_console): + """Test platform manager help without extensions.""" + with patch('cai.repl.commands.help.HAS_PLATFORM_EXTENSIONS', False): + result = help_command.handle_help_platform_manager() + + assert result is True + mock_console.print.assert_called_once() + call_args = mock_console.print.call_args[0][0] + panel_content = str(call_args.renderable if hasattr(call_args, 'renderable') else call_args) + assert "No platform extensions available" in panel_content + + def test_print_command_table(self, help_command, mock_console): + """Test _print_command_table helper method.""" + test_commands = [ + ("/test", "/t", "Test command description"), + ("/example", "/e", "Example command description") + ] + + help_command._print_command_table("Test Commands", test_commands) + + mock_console.print.assert_called_once() + + def test_create_styled_table_function(self): + """Test create_styled_table helper function.""" + from cai.repl.commands.help import create_styled_table + + headers = [("Command", "yellow"), ("Description", "white")] + table = create_styled_table("Test Table", headers) + + assert isinstance(table, Table) + assert table.title == "Test Table" + + def test_create_notes_panel_function(self): + """Test create_notes_panel helper function.""" + from cai.repl.commands.help import create_notes_panel + + notes = ["Note 1", "Note 2", "Note 3"] + panel = create_notes_panel(notes, "Test Notes") + + assert isinstance(panel, Panel) + + def test_full_help_workflow(self, help_command, mock_console): + """Test complete help workflow integration.""" + # Test main help + result1 = help_command.handle_no_args() + assert result1 is True + + # Test various subcommands + result2 = help_command.handle_agents() + assert result2 is True + + result3 = help_command.handle_shell() + assert result3 is True + + result4 = help_command.handle_env() + assert result4 is True + + # All should succeed + assert all([result1, result2, result3, result4]) + + def test_handle_memory_no_memory_command(self, help_command, mock_console): + """Test memory subcommand when no memory command exists.""" + with patch('cai.repl.commands.help.COMMANDS', {}): + result = help_command.handle_memory() + + assert result is True + # Should fall back to handle_help_memory + assert mock_console.print.call_count >= 1 + + def test_handle_platform_with_import_error(self, help_command, mock_console): + """Test platform help with import errors.""" + with patch('cai.repl.commands.help.HAS_PLATFORM_EXTENSIONS', True): + with patch('cai.repl.commands.help.is_caiextensions_platform_available', + return_value=False): + result = help_command.handle_help_platform_manager() + + assert result is True + assert mock_console.print.call_count >= 1 + + def test_handle_platform_empty_platforms(self, help_command, mock_console): + """Test platform help with no platforms registered.""" + mock_platform_manager = Mock() + mock_platform_manager.list_platforms.return_value = [] + + with patch('cai.repl.commands.help.HAS_PLATFORM_EXTENSIONS', True): + with patch('cai.repl.commands.help.is_caiextensions_platform_available', + return_value=True): + with patch('sys.modules', {'caiextensions.platform.base': Mock(platform_manager=mock_platform_manager)}): + result = help_command.handle_help_platform_manager() + + assert result is True + mock_console.print.assert_called_once() + call_args = mock_console.print.call_args[0][0] + panel_content = str(call_args.renderable if hasattr(call_args, 'renderable') else call_args) + assert "No platforms registered" in panel_content + + def test_handle_aliases_with_empty_registry(self, help_command, mock_console): + """Test aliases help with empty aliases registry.""" + with patch('cai.repl.commands.help.COMMAND_ALIASES', {}): + with patch('cai.repl.commands.help.COMMANDS', {}): + result = help_command.handle_help_aliases() + + assert result is True + # Should still create the table structure even if empty + assert mock_console.print.call_count >= 2 + + def test_subcommand_with_none_args(self, help_command): + """Test that subcommands handle None arguments correctly.""" + # All subcommands should accept None args and return True + result1 = help_command.handle_memory(None) + result2 = help_command.handle_agents(None) + result3 = help_command.handle_graph(None) + result4 = help_command.handle_shell(None) + result5 = help_command.handle_env(None) + result6 = help_command.handle_aliases(None) + result7 = help_command.handle_model(None) + result8 = help_command.handle_turns(None) + result9 = help_command.handle_config(None) + result10 = help_command.handle_platform(None) + + assert all([result1, result2, result3, result4, result5, + result6, result7, result8, result9, result10]) \ No newline at end of file diff --git a/tests/commands/test_command_history.py b/tests/commands/test_command_history.py new file mode 100644 index 00000000..6199efa3 --- /dev/null +++ b/tests/commands/test_command_history.py @@ -0,0 +1,518 @@ +#!/usr/bin/env python3 +""" +Test suite for the history command functionality. +Tests all handle methods and input possibilities for the history command. +""" + +import os +import sys +import pytest +import json +from unittest.mock import patch, Mock, MagicMock + +# Add src to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), + '..', '..', 'src')) + +from cai.repl.commands.history import HistoryCommand +from cai.repl.commands.base import Command + + +class TestHistoryCommand: + """Test cases for HistoryCommand.""" + + @pytest.fixture(autouse=True) + def setup_and_cleanup(self): + """Setup and cleanup for each test.""" + # Set up test environment + os.environ['CAI_TELEMETRY'] = 'false' + os.environ['CAI_TRACING'] = 'false' + + yield + + @pytest.fixture + def history_command(self): + """Create a HistoryCommand instance for testing.""" + return HistoryCommand() + + @pytest.fixture + def sample_message_history(self): + """Create sample message history for testing.""" + return [ + { + "role": "user", + "content": "Hello, can you help me with Python?" + }, + { + "role": "assistant", + "content": "Of course! I'd be happy to help you with Python. What specific topic or problem would you like assistance with?" + }, + { + "role": "user", + "content": "How do I create a list?" + }, + { + "role": "assistant", + "content": "I'll help you create a tool to demonstrate list creation.", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "create_example", + "arguments": '{"language": "python", "topic": "lists"}' + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call_123", + "content": "# Creating lists in Python\nmy_list = [1, 2, 3]\nprint(my_list)" + }, + { + "role": "assistant", + "content": "Here's how you create a list in Python: you use square brackets and separate items with commas." + } + ] + + @pytest.fixture + def complex_message_history(self): + """Create complex message history with various message types.""" + return [ + { + "role": "system", + "content": "You are a helpful assistant." + }, + { + "role": "user", + "content": "Execute a command for me" + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_cmd_1", + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": '{"command": "ls", "args": "-la"}' + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call_cmd_1", + "content": "total 8\ndrwxr-xr-x 2 user user 4096 Jan 1 12:00 .\ndrwxr-xr-x 3 user user 4096 Jan 1 12:00 .." + }, + { + "role": "assistant", + "content": "The directory listing shows two entries: the current directory (.) and parent directory (..)." + }, + { + "role": "user", + "content": "Now run multiple commands" + }, + { + "role": "assistant", + "content": "I'll run multiple commands for you.", + "tool_calls": [ + { + "id": "call_cmd_2", + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": '{"command": "pwd"}' + } + }, + { + "id": "call_cmd_3", + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": '{"command": "whoami"}' + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call_cmd_2", + "content": "/home/user" + }, + { + "role": "tool", + "tool_call_id": "call_cmd_3", + "content": "user" + } + ] + + def test_command_initialization(self, history_command): + """Test that HistoryCommand initializes correctly.""" + assert history_command.name == "/history" + assert history_command.description == "Display the conversation history" + assert history_command.aliases == ["/his"] + + @patch('cai.sdk.agents.models.openai_chatcompletions.message_history') + def test_handle_no_args_with_history(self, mock_message_history, + history_command, sample_message_history): + """Test handling with no arguments when history exists.""" + mock_message_history.clear() + mock_message_history.extend(sample_message_history) + + result = history_command.handle([]) + assert result is True + + @patch('cai.sdk.agents.models.openai_chatcompletions.message_history') + def test_handle_no_args_empty_history(self, mock_message_history, history_command): + """Test handling with no arguments when history is empty.""" + mock_message_history.clear() + + result = history_command.handle([]) + assert result is True + + @patch('cai.sdk.agents.models.openai_chatcompletions.message_history') + def test_handle_with_args(self, mock_message_history, + history_command, sample_message_history): + """Test handling when arguments are provided (should be ignored).""" + mock_message_history.clear() + mock_message_history.extend(sample_message_history) + + result = history_command.handle(["some", "args"]) + assert result is True + + @patch('cai.sdk.agents.models.openai_chatcompletions.message_history') + def test_handle_complex_history(self, mock_message_history, + history_command, complex_message_history): + """Test handling complex message history with various message types.""" + mock_message_history.clear() + mock_message_history.extend(complex_message_history) + + result = history_command.handle([]) + assert result is True + + def test_format_message_content_simple_text(self, history_command): + """Test formatting simple text content.""" + content = "This is a simple message" + tool_calls = None + + result = history_command._format_message_content(content, tool_calls) + assert result == content + + def test_format_message_content_long_text(self, history_command): + """Test formatting long text content (should be truncated).""" + content = "This is a very long message that should be truncated because it exceeds the 100 character limit that is set in the formatting function for display purposes" + tool_calls = None + + result = history_command._format_message_content(content, tool_calls) + assert len(result) <= 100 + assert result.endswith("...") + + def test_format_message_content_empty(self, history_command): + """Test formatting empty content.""" + content = "" + tool_calls = None + + result = history_command._format_message_content(content, tool_calls) + assert "Empty message" in result + + def test_format_message_content_none(self, history_command): + """Test formatting None content.""" + content = None + tool_calls = None + + result = history_command._format_message_content(content, tool_calls) + assert "Empty message" in result + + def test_format_message_content_with_tool_calls(self, history_command): + """Test formatting content with tool calls.""" + content = "I'll help you with that" + tool_calls = [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "test_function", + "arguments": '{"param1": "value1", "param2": "value2"}' + } + } + ] + + result = history_command._format_message_content(content, tool_calls) + assert "Function:" in result + assert "test_function" in result + assert "Args:" in result + + def test_format_message_content_with_multiple_tool_calls(self, history_command): + """Test formatting content with multiple tool calls.""" + content = "I'll run multiple commands" + tool_calls = [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "function_one", + "arguments": '{"arg": "value1"}' + } + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "function_two", + "arguments": '{"arg": "value2"}' + } + } + ] + + result = history_command._format_message_content(content, tool_calls) + assert "function_one" in result + assert "function_two" in result + assert result.count("Function:") == 2 + + def test_format_message_content_with_invalid_json_args(self, history_command): + """Test formatting content with invalid JSON in tool call arguments.""" + content = "Testing invalid JSON" + tool_calls = [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "test_function", + "arguments": "invalid json string" + } + } + ] + + result = history_command._format_message_content(content, tool_calls) + assert "Function:" in result + assert "test_function" in result + assert "invalid json string" in result + + def test_format_message_content_with_long_tool_args(self, history_command): + """Test formatting content with very long tool call arguments.""" + content = "Testing long arguments" + long_args = json.dumps({ + "very_long_parameter": "This is a very long parameter value that should be truncated when displayed in the history because it exceeds the character limit" + }) + tool_calls = [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "test_function", + "arguments": long_args + } + } + ] + + result = history_command._format_message_content(content, tool_calls) + assert "Function:" in result + assert "test_function" in result + assert "..." in result # Should be truncated + + @patch('cai.sdk.agents.models.openai_chatcompletions.message_history') + def test_handle_import_error(self, mock_message_history, history_command): + """Test handling when import fails.""" + # Create a new command instance and monkey patch its handle_no_args method + # to simulate an import error + test_command = HistoryCommand() + + def mock_handle_no_args_with_import_error(): + try: + # Simulate the import failing + raise ImportError("Mock import error") + except ImportError: + return False + + # Replace the method to simulate import failure + test_command.handle_no_args = mock_handle_no_args_with_import_error + + result = test_command.handle_no_args() + assert result is False + + def test_command_base_functionality(self, history_command): + """Test that the command inherits from base Command properly.""" + assert isinstance(history_command, Command) + assert history_command.name == "/history" + assert "/his" in history_command.aliases + + @patch('cai.sdk.agents.models.openai_chatcompletions.message_history') + def test_handle_with_corrupted_message(self, mock_message_history, history_command): + """Test handling when message history contains corrupted data.""" + # Create message history with missing or corrupted fields + corrupted_history = [ + { + "role": "user", + "content": "Normal message" + }, + { + # Missing role field + "content": "Message without role" + }, + { + "role": "assistant", + # Missing content field + "tool_calls": [{"id": "call_1", "function": {"name": "test"}}] + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": None # None content + } + ] + + mock_message_history.clear() + mock_message_history.extend(corrupted_history) + + # Should handle gracefully without crashing + result = history_command.handle([]) + assert result is True + + @patch('cai.sdk.agents.models.openai_chatcompletions.message_history') + def test_handle_with_messages_parameter(self, mock_message_history, + history_command, sample_message_history): + """Test handle method with explicit messages parameter.""" + mock_message_history.clear() + mock_message_history.extend(sample_message_history) + + # The handle method should work with messages parameter + result = history_command.handle([], messages=sample_message_history) + assert result is True + + +@pytest.mark.integration +class TestHistoryCommandIntegration: + """Integration tests for history command functionality.""" + + @pytest.fixture(autouse=True) + def setup_integration(self): + """Setup for integration tests.""" + yield + + @patch('cai.sdk.agents.models.openai_chatcompletions.message_history') + def test_full_conversation_history_workflow(self, mock_message_history): + """Test a complete conversation workflow and history display.""" + # Simulate a conversation building up over time + conversation_steps = [ + # Step 1: Initial user message + [ + {"role": "user", "content": "Hello"} + ], + # Step 2: Assistant response + [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi! How can I help you?"} + ], + # Step 3: User asks for help + [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi! How can I help you?"}, + {"role": "user", "content": "Can you run a command?"} + ], + # Step 4: Assistant with tool call + [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi! How can I help you?"}, + {"role": "user", "content": "Can you run a command?"}, + { + "role": "assistant", + "content": "I'll run that command for you.", + "tool_calls": [ + { + "id": "call_cmd", + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": '{"command": "ls"}' + } + } + ] + } + ], + # Step 5: Tool response + [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi! How can I help you?"}, + {"role": "user", "content": "Can you run a command?"}, + { + "role": "assistant", + "content": "I'll run that command for you.", + "tool_calls": [ + { + "id": "call_cmd", + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": '{"command": "ls"}' + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call_cmd", + "content": "file1.txt\nfile2.txt\nfolder1/" + } + ] + ] + + cmd = HistoryCommand() + + # Test history at each step + for i, step_history in enumerate(conversation_steps): + mock_message_history.clear() + mock_message_history.extend(step_history) + + result = cmd.handle([]) + assert result is True, f"Failed at conversation step {i + 1}" + + @patch('cai.sdk.agents.models.openai_chatcompletions.message_history') + def test_edge_case_message_combinations(self, mock_message_history): + """Test various edge case message combinations.""" + edge_cases = [ + # Empty history + [], + # Only system message + [{"role": "system", "content": "You are a helpful assistant."}], + # Only user messages + [ + {"role": "user", "content": "First message"}, + {"role": "user", "content": "Second message"} + ], + # Tool call without response + [ + {"role": "user", "content": "Run command"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_incomplete", + "type": "function", + "function": {"name": "test_command", "arguments": "{}"} + } + ] + } + ], + # Tool response without call + [ + {"role": "user", "content": "Test"}, + {"role": "tool", "tool_call_id": "orphan_call", "content": "Result"} + ] + ] + + cmd = HistoryCommand() + + for i, case_history in enumerate(edge_cases): + mock_message_history.clear() + mock_message_history.extend(case_history) + + result = cmd.handle([]) + assert result is True, f"Failed at edge case {i + 1}" + + +if __name__ == '__main__': + pytest.main([__file__, "-v"]) \ No newline at end of file diff --git a/tests/commands/test_command_model.py b/tests/commands/test_command_model.py new file mode 100644 index 00000000..2a15ad8a --- /dev/null +++ b/tests/commands/test_command_model.py @@ -0,0 +1,539 @@ +#!/usr/bin/env python3 +""" +Test suite for the model command functionality. +Tests all handle methods and input possibilities for the model command. +""" + +import os +import sys +import pytest +import datetime +from unittest.mock import patch, Mock, MagicMock + +# Add src to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), + '..', '..', 'src')) + +from cai.repl.commands.model import ModelCommand, ModelShowCommand +from cai.repl.commands.base import Command + + +class TestModelCommand: + """Test cases for ModelCommand.""" + + @pytest.fixture(autouse=True) + def setup_and_cleanup(self): + """Setup and cleanup for each test.""" + # Set up test environment + os.environ['CAI_TELEMETRY'] = 'false' + os.environ['CAI_TRACING'] = 'false' + + # Store original CAI_MODEL if it exists + self.original_model = os.environ.get('CAI_MODEL') + + yield + + # Restore original CAI_MODEL or remove if it didn't exist + if self.original_model is not None: + os.environ['CAI_MODEL'] = self.original_model + elif 'CAI_MODEL' in os.environ: + del os.environ['CAI_MODEL'] + + @pytest.fixture + def model_command(self): + """Create a ModelCommand instance for testing.""" + return ModelCommand() + + @pytest.fixture + def model_show_command(self): + """Create a ModelShowCommand instance for testing.""" + return ModelShowCommand() + + @pytest.fixture + def mock_litellm_response(self): + """Create a mock response for LiteLLM model data.""" + return { + "gpt-4": { + "input_cost_per_token": 0.00003, + "output_cost_per_token": 0.00006, + "max_tokens": 8192, + "supports_function_calling": True, + "supports_vision": False, + "litellm_provider": "openai" + }, + "claude-3-sonnet-20240229": { + "input_cost_per_token": 0.000015, + "output_cost_per_token": 0.000075, + "max_tokens": 200000, + "supports_function_calling": True, + "supports_vision": True, + "litellm_provider": "anthropic" + }, + "deepseek/deepseek-v3": { + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002, + "max_tokens": 128000, + "supports_function_calling": True, + "supports_vision": False, + "litellm_provider": "deepseek" + } + } + + @pytest.fixture + def mock_ollama_response(self): + """Create a mock response for Ollama models.""" + return { + "models": [ + { + "name": "llama3", + "size": 4661211648 # ~4.3 GB + }, + { + "name": "mistral:7b", + "size": 7365960192 # ~6.9 GB + } + ] + } + + def test_command_initialization(self, model_command): + """Test that ModelCommand initializes correctly.""" + assert model_command.name == "/model" + assert model_command.description == "View or change the current LLM model" + assert model_command.aliases == ["/mod"] + + # Check that cached models and numbers are initialized + assert hasattr(model_command, 'cached_models') + assert hasattr(model_command, 'cached_model_numbers') + assert hasattr(model_command, 'last_model_fetch') + + def test_model_show_command_initialization(self, model_show_command): + """Test that ModelShowCommand initializes correctly.""" + assert model_show_command.name == "/model-show" + assert model_show_command.description == "Show all available models from LiteLLM repository" + assert model_show_command.aliases == ["/mod-show"] + + @patch('requests.get') + def test_handle_no_args_with_mock_data(self, mock_get, model_command, + mock_litellm_response, mock_ollama_response): + """Test showing current model and available models with no arguments.""" + # Mock LiteLLM response + mock_litellm = Mock() + mock_litellm.status_code = 200 + mock_litellm.json.return_value = mock_litellm_response + + # Mock Ollama response + mock_ollama = Mock() + mock_ollama.status_code = 200 + mock_ollama.json.return_value = mock_ollama_response + + # Configure the mock to return different responses based on URL + def side_effect(url, timeout=None): + if "litellm" in url: + return mock_litellm + elif "ollama" in url: + return mock_ollama + else: + return Mock(status_code=404) + + mock_get.side_effect = side_effect + + # Set a model first + os.environ['CAI_MODEL'] = 'gpt-4' + + result = model_command.handle([]) + assert result is True + + @patch('requests.get') + def test_handle_select_model_by_name(self, mock_get, model_command, + mock_litellm_response): + """Test selecting a model by name.""" + # Mock LiteLLM response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_litellm_response + mock_get.return_value = mock_response + + result = model_command.handle(["gpt-4"]) + assert result is True + assert os.environ.get('CAI_MODEL') == 'gpt-4' + + @patch('requests.get') + def test_handle_select_model_by_number(self, mock_get, model_command, + mock_litellm_response): + """Test selecting a model by number.""" + # Mock LiteLLM response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_litellm_response + mock_get.return_value = mock_response + + # First call to populate cache + model_command.handle([]) + + # Then select by number + result = model_command.handle(["1"]) + assert result is True + assert 'CAI_MODEL' in os.environ + + def test_handle_select_custom_model(self, model_command): + """Test selecting a custom model name not in the predefined list.""" + result = model_command.handle(["custom-model-name"]) + assert result is True + assert os.environ.get('CAI_MODEL') == 'custom-model-name' + + @patch('requests.get') + def test_handle_with_network_error(self, mock_get, model_command): + """Test handling when network requests fail.""" + # Mock network failure + mock_get.side_effect = Exception("Network error") + + result = model_command.handle([]) + assert result is True # Should still work, just without external data + + @patch('requests.get') + def test_handle_model_pricing_data_error(self, mock_get, model_command): + """Test handling when LiteLLM API returns error.""" + # Mock HTTP error + mock_response = Mock() + mock_response.status_code = 404 + mock_get.return_value = mock_response + + result = model_command.handle([]) + assert result is True # Should still work with built-in models + + def test_command_base_functionality(self, model_command): + """Test that the command inherits from base Command properly.""" + assert isinstance(model_command, Command) + assert model_command.name == "/model" + assert "/mod" in model_command.aliases + + +class TestModelShowCommand: + """Test cases for ModelShowCommand.""" + + @pytest.fixture(autouse=True) + def setup_and_cleanup(self): + """Setup and cleanup for each test.""" + # Set up test environment + os.environ['CAI_TELEMETRY'] = 'false' + os.environ['CAI_TRACING'] = 'false' + + yield + + @pytest.fixture + def model_show_command(self): + """Create a ModelShowCommand instance for testing.""" + return ModelShowCommand() + + @pytest.fixture + def mock_litellm_response(self): + """Create a mock response for LiteLLM model data.""" + return { + "gpt-4": { + "input_cost_per_token": 0.00003, + "output_cost_per_token": 0.00006, + "max_tokens": 8192, + "supports_function_calling": True, + "supports_vision": False, + "litellm_provider": "openai" + }, + "claude-3-sonnet-20240229": { + "input_cost_per_token": 0.000015, + "output_cost_per_token": 0.000075, + "max_tokens": 200000, + "supports_function_calling": True, + "supports_vision": True, + "litellm_provider": "anthropic" + }, + "gpt-3.5-turbo": { + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002, + "max_tokens": 4096, + "supports_function_calling": False, + "supports_vision": False, + "litellm_provider": "openai" + } + } + + @pytest.fixture + def mock_ollama_response(self): + """Create a mock response for Ollama models.""" + return { + "models": [ + { + "name": "llama3", + "size": 4661211648 + }, + { + "name": "mistral:7b", + "size": 7365960192 + } + ] + } + + @patch('requests.get') + def test_handle_no_args(self, mock_get, model_show_command, + mock_litellm_response, mock_ollama_response): + """Test showing all models with no arguments.""" + # Mock LiteLLM response + mock_litellm = Mock() + mock_litellm.status_code = 200 + mock_litellm.json.return_value = mock_litellm_response + + # Mock Ollama response + mock_ollama = Mock() + mock_ollama.status_code = 200 + mock_ollama.json.return_value = mock_ollama_response + + # Configure the mock to return different responses based on URL + def side_effect(url, timeout=None): + if "litellm" in url: + return mock_litellm + elif "ollama" in url: + return mock_ollama + else: + return Mock(status_code=404) + + mock_get.side_effect = side_effect + + result = model_show_command.handle([]) + assert result is True + + @patch('requests.get') + def test_handle_supported_filter(self, mock_get, model_show_command, + mock_litellm_response): + """Test showing only supported models (with function calling).""" + # Mock LiteLLM response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_litellm_response + mock_get.return_value = mock_response + + result = model_show_command.handle(["supported"]) + assert result is True + + @patch('requests.get') + def test_handle_search_filter(self, mock_get, model_show_command, + mock_litellm_response): + """Test filtering models by search term.""" + # Mock LiteLLM response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_litellm_response + mock_get.return_value = mock_response + + result = model_show_command.handle(["gpt"]) + assert result is True + + @patch('requests.get') + def test_handle_supported_and_search(self, mock_get, model_show_command, + mock_litellm_response): + """Test combining supported filter with search term.""" + # Mock LiteLLM response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_litellm_response + mock_get.return_value = mock_response + + result = model_show_command.handle(["supported", "claude"]) + assert result is True + + @patch('requests.get') + def test_handle_network_error(self, mock_get, model_show_command): + """Test handling when network request fails.""" + # Mock network failure + mock_get.side_effect = Exception("Network error") + + result = model_show_command.handle([]) + assert result is True # Should handle gracefully + + @patch('requests.get') + def test_handle_http_error(self, mock_get, model_show_command): + """Test handling when API returns HTTP error.""" + # Mock HTTP error + mock_response = Mock() + mock_response.status_code = 500 + mock_get.return_value = mock_response + + result = model_show_command.handle([]) + assert result is True # Should handle gracefully + + @patch('requests.get') + def test_handle_with_ollama_error(self, mock_get, model_show_command, + mock_litellm_response): + """Test handling when Ollama is not available but LiteLLM works.""" + # Mock LiteLLM success but Ollama failure + def side_effect(url, timeout=None): + if "litellm" in url: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_litellm_response + return mock_response + elif "ollama" in url: + raise Exception("Ollama not available") + else: + return Mock(status_code=404) + + mock_get.side_effect = side_effect + + result = model_show_command.handle([]) + assert result is True + + +@pytest.mark.integration +class TestModelCommandIntegration: + """Integration tests for model command functionality.""" + + @pytest.fixture(autouse=True) + def setup_integration(self): + """Setup for integration tests.""" + # Store original CAI_MODEL if it exists + self.original_model = os.environ.get('CAI_MODEL') + + yield + + # Restore original CAI_MODEL or remove if it didn't exist + if self.original_model is not None: + os.environ['CAI_MODEL'] = self.original_model + elif 'CAI_MODEL' in os.environ: + del os.environ['CAI_MODEL'] + + @patch('requests.get') + def test_full_model_workflow(self, mock_get): + """Test a complete workflow of listing and selecting models.""" + # Mock responses + mock_litellm_response = { + "gpt-4": { + "input_cost_per_token": 0.00003, + "output_cost_per_token": 0.00006, + "max_tokens": 8192, + "supports_function_calling": True + }, + "claude-3-sonnet-20240229": { + "input_cost_per_token": 0.000015, + "output_cost_per_token": 0.000075, + "max_tokens": 200000, + "supports_function_calling": True + } + } + + mock_ollama_response = { + "models": [ + {"name": "llama3", "size": 4661211648} + ] + } + + # Configure mock responses + def side_effect(url, timeout=None): + if "litellm" in url: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_litellm_response + return mock_response + elif "ollama" in url: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_ollama_response + return mock_response + else: + return Mock(status_code=404) + + mock_get.side_effect = side_effect + + model_cmd = ModelCommand() + model_show_cmd = ModelShowCommand() + + # List all models + result1 = model_cmd.handle([]) + assert result1 is True + + # Show detailed model info + result2 = model_show_cmd.handle([]) + assert result2 is True + + # Select a model by name + result3 = model_cmd.handle(["gpt-4"]) + assert result3 is True + assert os.environ.get('CAI_MODEL') == 'gpt-4' + + # Show current model again + result4 = model_cmd.handle([]) + assert result4 is True + + # Select by number (after cache is populated) + result5 = model_cmd.handle(["1"]) + assert result5 is True + + @patch('requests.get') + def test_model_selection_edge_cases(self, mock_get): + """Test edge cases in model selection.""" + # Mock minimal response to avoid network dependency + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = {"gpt-4": {}} + mock_get.return_value = mock_response + + cmd = ModelCommand() + + # Test selecting non-existent number (large number) + result1 = cmd.handle(["999"]) + assert result1 is True + assert os.environ.get('CAI_MODEL') == '999' # Should use as literal + + # Test selecting model with special characters + result2 = cmd.handle(["custom/model:latest"]) + assert result2 is True + assert os.environ.get('CAI_MODEL') == 'custom/model:latest' + + # Test empty model name (edge case) + result3 = cmd.handle([""]) + assert result3 is True + assert os.environ.get('CAI_MODEL') == '' + + @patch('requests.get') + def test_model_show_filters_combination(self, mock_get): + """Test various combinations of filters in model-show command.""" + mock_response = { + "gpt-4": { + "supports_function_calling": True, + "input_cost_per_token": 0.00003, + "output_cost_per_token": 0.00006 + }, + "gpt-3.5-turbo": { + "supports_function_calling": False, + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002 + }, + "claude-3-sonnet": { + "supports_function_calling": True, + "input_cost_per_token": 0.000015, + "output_cost_per_token": 0.000075 + } + } + + mock_http_response = Mock() + mock_http_response.status_code = 200 + mock_http_response.json.return_value = mock_response + mock_get.return_value = mock_http_response + + cmd = ModelShowCommand() + + # Test supported only + result1 = cmd.handle(["supported"]) + assert result1 is True + + # Test search only + result2 = cmd.handle(["gpt"]) + assert result2 is True + + # Test supported + search + result3 = cmd.handle(["supported", "claude"]) + assert result3 is True + + # Test search + supported (different order) + result4 = cmd.handle(["claude", "supported"]) + assert result4 is True + + +if __name__ == '__main__': + pytest.main([__file__, "-v"]) \ No newline at end of file diff --git a/tests/commands/test_command_parallel.py b/tests/commands/test_command_parallel.py new file mode 100644 index 00000000..95410bd7 --- /dev/null +++ b/tests/commands/test_command_parallel.py @@ -0,0 +1,344 @@ +#!/usr/bin/env python3 +""" +Test suite for the parallel command functionality. +Tests all handle methods and input possibilities for the parallel command. +""" + +import os +import sys +import pytest +from unittest.mock import patch, Mock, MagicMock + +# Add src to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), + '..', '..', 'src')) + +from cai.repl.commands.parallel import ParallelCommand, ParallelConfig, PARALLEL_CONFIGS +from cai.repl.commands.base import Command + + +class TestParallelCommand: + """Test cases for ParallelCommand.""" + + @pytest.fixture(autouse=True) + def setup_and_cleanup(self): + """Setup and cleanup for each test.""" + # Clear parallel configs before each test + PARALLEL_CONFIGS.clear() + + # Set up test environment + os.environ['CAI_TELEMETRY'] = 'false' + os.environ['CAI_TRACING'] = 'false' + + yield + + # Cleanup after each test + PARALLEL_CONFIGS.clear() + + @pytest.fixture + def parallel_command(self): + """Create a ParallelCommand instance for testing.""" + return ParallelCommand() + + def test_command_initialization(self, parallel_command): + """Test that ParallelCommand initializes correctly.""" + assert parallel_command.name == "/parallel" + assert parallel_command.description == "Configure multiple agents to run in parallel with different settings" + assert parallel_command.aliases == ["/par", "/p"] + + # Check subcommands are registered + expected_subcommands = ["add", "list", "clear", "remove"] + assert set(parallel_command.get_subcommands()) == set(expected_subcommands) + + def test_parallel_config_initialization(self): + """Test ParallelConfig initialization.""" + config = ParallelConfig("test_agent", "gpt-4", "Test prompt") + assert config.agent_name == "test_agent" + assert config.model == "gpt-4" + assert config.prompt == "Test prompt" + + # Test default values + config_default = ParallelConfig("test_agent") + assert config_default.agent_name == "test_agent" + assert config_default.model is None + assert config_default.prompt is None + + def test_parallel_config_str_representation(self): + """Test ParallelConfig string representation.""" + # Test with all parameters + config = ParallelConfig("test_agent", "gpt-4", "Test prompt") + str_repr = str(config) + assert "Agent: test_agent" in str_repr + assert "model: gpt-4" in str_repr + assert "prompt: 'Test prompt'" in str_repr + + # Test with long prompt (should be truncated) + long_prompt = "This is a very long prompt that should be truncated when displayed" + config_long = ParallelConfig("test_agent", "gpt-4", long_prompt) + str_repr_long = str(config_long) + assert "..." in str_repr_long + + # Test with minimal parameters + config_minimal = ParallelConfig("test_agent") + str_repr_minimal = str(config_minimal) + assert "Agent: test_agent" in str_repr_minimal + assert "model:" not in str_repr_minimal + assert "prompt:" not in str_repr_minimal + + @patch('cai.repl.commands.parallel.get_available_agents') + def test_handle_add_valid_agent(self, mock_get_agents, parallel_command): + """Test adding a valid agent to parallel config.""" + # Mock available agents + mock_get_agents.return_value = { + "test_agent": Mock(), + "another_agent": Mock() + } + + # Test basic add + result = parallel_command.handle_add(["test_agent"]) + assert result is True + assert len(PARALLEL_CONFIGS) == 1 + assert PARALLEL_CONFIGS[0].agent_name == "test_agent" + assert PARALLEL_CONFIGS[0].model is None + assert PARALLEL_CONFIGS[0].prompt is None + + @patch('cai.repl.commands.parallel.get_available_agents') + def test_handle_add_with_model_and_prompt(self, mock_get_agents, parallel_command): + """Test adding agent with model and prompt parameters.""" + mock_get_agents.return_value = {"test_agent": Mock()} + + args = ["test_agent", "--model", "gpt-4", "--prompt", "Custom prompt"] + result = parallel_command.handle_add(args) + + assert result is True + assert len(PARALLEL_CONFIGS) == 1 + config = PARALLEL_CONFIGS[0] + assert config.agent_name == "test_agent" + assert config.model == "gpt-4" + assert config.prompt == "Custom prompt" + + @patch('cai.repl.commands.parallel.get_available_agents') + def test_handle_add_invalid_agent(self, mock_get_agents, parallel_command): + """Test adding an invalid agent name.""" + mock_get_agents.return_value = {"valid_agent": Mock()} + + result = parallel_command.handle_add(["invalid_agent"]) + assert result is False + assert len(PARALLEL_CONFIGS) == 0 + + def test_handle_add_no_args(self, parallel_command): + """Test add command with no arguments.""" + result = parallel_command.handle_add([]) + assert result is False + assert len(PARALLEL_CONFIGS) == 0 + + @patch('cai.repl.commands.parallel.get_available_agents') + def test_handle_add_multiple_agents(self, mock_get_agents, parallel_command): + """Test adding multiple agents.""" + mock_get_agents.return_value = { + "agent1": Mock(), + "agent2": Mock(), + "agent3": Mock() + } + + # Add first agent + result1 = parallel_command.handle_add(["agent1", "--model", "gpt-4"]) + assert result1 is True + + # Add second agent + result2 = parallel_command.handle_add(["agent2", "--prompt", "Second prompt"]) + assert result2 is True + + # Add third agent + result3 = parallel_command.handle_add(["agent3"]) + assert result3 is True + + assert len(PARALLEL_CONFIGS) == 3 + assert PARALLEL_CONFIGS[0].agent_name == "agent1" + assert PARALLEL_CONFIGS[1].agent_name == "agent2" + assert PARALLEL_CONFIGS[2].agent_name == "agent3" + + def test_handle_list_empty(self, parallel_command): + """Test listing when no parallel configs exist.""" + result = parallel_command.handle_list([]) + assert result is True + assert len(PARALLEL_CONFIGS) == 0 + + def test_handle_list_with_configs(self, parallel_command): + """Test listing existing parallel configs.""" + # Add some configs + PARALLEL_CONFIGS.append(ParallelConfig("agent1", "gpt-4", "Prompt 1")) + PARALLEL_CONFIGS.append(ParallelConfig("agent2", None, None)) + PARALLEL_CONFIGS.append(ParallelConfig("agent3", "claude", "Long prompt")) + + result = parallel_command.handle_list([]) + assert result is True + + def test_handle_clear_empty(self, parallel_command): + """Test clearing empty parallel configs.""" + result = parallel_command.handle_clear([]) + assert result is True + assert len(PARALLEL_CONFIGS) == 0 + + def test_handle_clear_with_configs(self, parallel_command): + """Test clearing existing parallel configs.""" + # Add some configs + PARALLEL_CONFIGS.append(ParallelConfig("agent1")) + PARALLEL_CONFIGS.append(ParallelConfig("agent2")) + PARALLEL_CONFIGS.append(ParallelConfig("agent3")) + + assert len(PARALLEL_CONFIGS) == 3 + + result = parallel_command.handle_clear([]) + assert result is True + assert len(PARALLEL_CONFIGS) == 0 + + def test_handle_remove_valid_index(self, parallel_command): + """Test removing a config by valid index.""" + # Add some configs + PARALLEL_CONFIGS.append(ParallelConfig("agent1")) + PARALLEL_CONFIGS.append(ParallelConfig("agent2")) + PARALLEL_CONFIGS.append(ParallelConfig("agent3")) + + # Remove the second config (index 2) + result = parallel_command.handle_remove(["2"]) + assert result is True + assert len(PARALLEL_CONFIGS) == 2 + assert PARALLEL_CONFIGS[0].agent_name == "agent1" + assert PARALLEL_CONFIGS[1].agent_name == "agent3" + + def test_handle_remove_invalid_index(self, parallel_command): + """Test removing with invalid index.""" + PARALLEL_CONFIGS.append(ParallelConfig("agent1")) + + # Test invalid numeric index + result1 = parallel_command.handle_remove(["5"]) + assert result1 is False + assert len(PARALLEL_CONFIGS) == 1 + + # Test negative index + result2 = parallel_command.handle_remove(["-1"]) + assert result2 is False + assert len(PARALLEL_CONFIGS) == 1 + + # Test non-numeric index + result3 = parallel_command.handle_remove(["invalid"]) + assert result3 is False + assert len(PARALLEL_CONFIGS) == 1 + + def test_handle_remove_no_args(self, parallel_command): + """Test remove command with no arguments.""" + PARALLEL_CONFIGS.append(ParallelConfig("agent1")) + + result = parallel_command.handle_remove([]) + assert result is False + assert len(PARALLEL_CONFIGS) == 1 + + def test_command_base_functionality(self, parallel_command): + """Test that the command inherits from base Command properly.""" + assert isinstance(parallel_command, Command) + assert parallel_command.name == "/parallel" + assert "/par" in parallel_command.aliases + assert "/p" in parallel_command.aliases + + @patch('cai.repl.commands.parallel.get_available_agents') + def test_handle_main_command_routing(self, mock_get_agents, parallel_command): + """Test that main handle method routes to correct subcommands.""" + mock_get_agents.return_value = {"test_agent": Mock()} + + # Test routing to add + result1 = parallel_command.handle(["add", "test_agent"]) + assert result1 is True + assert len(PARALLEL_CONFIGS) == 1 + + # Test routing to list + result2 = parallel_command.handle(["list"]) + assert result2 is True + + # Test routing to clear + result3 = parallel_command.handle(["clear"]) + assert result3 is True + assert len(PARALLEL_CONFIGS) == 0 + + def test_handle_unknown_subcommand(self, parallel_command): + """Test handling of unknown subcommands.""" + # This will use the default handle method from base class + # which should route to handle_unknown_subcommand + result = parallel_command.handle(["unknown_subcommand"]) + assert result is False + + def test_handle_no_args(self, parallel_command): + """Test handling when no arguments provided.""" + # The base handle method should route to handle_no_args + result = parallel_command.handle([]) + assert result is False + + +@pytest.mark.integration +class TestParallelCommandIntegration: + """Integration tests for parallel command functionality.""" + + @pytest.fixture(autouse=True) + def setup_integration(self): + """Setup for integration tests.""" + PARALLEL_CONFIGS.clear() + yield + PARALLEL_CONFIGS.clear() + + @patch('cai.repl.commands.parallel.get_available_agents') + def test_full_workflow(self, mock_get_agents): + """Test a complete workflow of adding, listing, and removing configs.""" + mock_get_agents.return_value = { + "agent1": Mock(), + "agent2": Mock(), + "agent3": Mock() + } + + cmd = ParallelCommand() + + # Start with empty configs + assert len(PARALLEL_CONFIGS) == 0 + + # Add multiple configs + cmd.handle(["add", "agent1", "--model", "gpt-4"]) + cmd.handle(["add", "agent2", "--prompt", "Test prompt"]) + cmd.handle(["add", "agent3", "--model", "claude", "--prompt", "Another prompt"]) + + assert len(PARALLEL_CONFIGS) == 3 + + # List configs (should not change count) + cmd.handle(["list"]) + assert len(PARALLEL_CONFIGS) == 3 + + # Remove one config + cmd.handle(["remove", "2"]) + assert len(PARALLEL_CONFIGS) == 2 + + # Clear all configs + cmd.handle(["clear"]) + assert len(PARALLEL_CONFIGS) == 0 + + @patch('cai.repl.commands.parallel.get_available_agents') + def test_edge_case_combinations(self, mock_get_agents): + """Test edge cases and unusual parameter combinations.""" + mock_get_agents.return_value = {"test_agent": Mock()} + + cmd = ParallelCommand() + + # Test partial parameters + result1 = cmd.handle(["add", "test_agent", "--model"]) + assert result1 is True # Should still work with incomplete args + + # Test empty string parameters + result2 = cmd.handle(["add", "test_agent", "--prompt", ""]) + assert result2 is True + + # Test parameters in different order + result3 = cmd.handle(["add", "test_agent", "--prompt", "Test", "--model", "gpt-4"]) + assert result3 is True + + assert len(PARALLEL_CONFIGS) == 3 + + +if __name__ == '__main__': + pytest.main([__file__, "-v"]) \ No newline at end of file From 2324b244a3eed5b5da3706cc10eafe5277c308fd Mon Sep 17 00:00:00 2001 From: luijait Date: Fri, 23 May 2025 21:12:18 +0200 Subject: [PATCH 3/4] Delete unnecesary tests --- ci/test/.test.yml | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/ci/test/.test.yml b/ci/test/.test.yml index f5e1a082..23e3e48c 100644 --- a/ci/test/.test.yml +++ b/ci/test/.test.yml @@ -328,23 +328,8 @@ <<: *run_test variables: TEST_PATH: tests/commands/test_command_config.py - -💻 commands test_command_flush.py: - <<: *run_test - variables: - TEST_PATH: tests/commands/test_command_flush.py - + 💻 commands test_command_help.py: <<: *run_test variables: TEST_PATH: tests/commands/test_command_help.py - -💻 commands test_command_virtualization.py: - <<: *run_test - variables: - TEST_PATH: tests/commands/test_command_virtualization.py - -💻 commands test_command_workspace.py: - <<: *run_test - variables: - TEST_PATH: tests/commands/test_command_workspace.py From 486f88cd218dadf8582a69cae370053b664f444b Mon Sep 17 00:00:00 2001 From: luijait Date: Sat, 24 May 2025 07:37:44 +0000 Subject: [PATCH 4/4] Little fix in session outputs --- src/cai/tools/common.py | 41 ------------------- .../reconnaissance/generic_linux_command.py | 4 +- 2 files changed, 2 insertions(+), 43 deletions(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 85923a9f..3d5eaa51 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -377,47 +377,6 @@ def get_session_output(session_id, clear=True, stdout=True): session = ACTIVE_SESSIONS[session_id] 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 diff --git a/src/cai/tools/reconnaissance/generic_linux_command.py b/src/cai/tools/reconnaissance/generic_linux_command.py index b1ab271f..f6beaede 100644 --- a/src/cai/tools/reconnaissance/generic_linux_command.py +++ b/src/cai/tools/reconnaissance/generic_linux_command.py @@ -65,8 +65,8 @@ def generic_linux_command(command: str = "", 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 - get_session_output(session_id, clear=False, stdout=True) - return f"Session {session_id} output displayed above" + output = get_session_output(session_id, clear=False, stdout=True) + return output if args.startswith("kill "): session_id = args.split(" ")[1]