#!/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 # Add src to path sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src")) from cai.sdk.agents.models.openai_chatcompletions import ( get_agent_message_history, get_all_agent_histories, ACTIVE_MODEL_INSTANCES, ) 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.""" # Import here to avoid circular imports from cai.sdk.agents.models.openai_chatcompletions import PERSISTENT_MESSAGE_HISTORIES from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER # Clear all active model instances ACTIVE_MODEL_INSTANCES.clear() # Clear persistent message histories PERSISTENT_MESSAGE_HISTORIES.clear() # Clear AGENT_MANAGER state AGENT_MANAGER.clear_all_histories() AGENT_MANAGER.reset_registry() def setUp(self): """Set up for each test method.""" # AGGRESSIVE cleanup to ensure no state contamination between tests # Clear all active model instances ACTIVE_MODEL_INSTANCES.clear() # Keep a strong reference to prevent garbage collection self._test_model = None # Clear any existing message histories from cai.sdk.agents.models.openai_chatcompletions import ( OpenAIChatCompletionsModel, PERSISTENT_MESSAGE_HISTORIES ) from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER # Clear persistent message histories to ensure clean state PERSISTENT_MESSAGE_HISTORIES.clear() # Clear AGENT_MANAGER state AGENT_MANAGER.clear_all_histories() AGENT_MANAGER.reset_registry() # Ensure we start with clean histories for each test for (name, instance_id), model_ref in list(ACTIVE_MODEL_INSTANCES.items()): model = model_ref() if model_ref else None if model and hasattr(model, 'message_history'): model.message_history.clear() def tearDown(self): """Clean up after each test.""" # Import here to avoid circular imports from cai.sdk.agents.models.openai_chatcompletions import PERSISTENT_MESSAGE_HISTORIES from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER # Clear all active model instances ACTIVE_MODEL_INSTANCES.clear() # Clear persistent message histories PERSISTENT_MESSAGE_HISTORIES.clear() # Clear AGENT_MANAGER state AGENT_MANAGER.clear_all_histories() AGENT_MANAGER.reset_registry() # Clear reference to test model self._test_model = None def get_combined_message_history(self): """Get combined message history from all agents.""" all_messages = [] histories = get_all_agent_histories() for agent_name, history in histories.items(): all_messages.extend(history) return all_messages def add_to_test_message_history(self, msg): """Add a message to the test agent's history.""" # Create a mock model instance for testing from cai.sdk.agents.models.openai_chatcompletions import OpenAIChatCompletionsModel from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER from openai import AsyncOpenAI import os test_agent_name = "test_agent" # Check if we already have a test model instance test_model = None for (name, instance_id), model_ref in ACTIVE_MODEL_INSTANCES.items(): if name == test_agent_name: model = model_ref() if model_ref else None if model: test_model = model break # Create one if it doesn't exist if not test_model: client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY", "test-key")) # Create with explicit agent_id to ensure registration test_model = OpenAIChatCompletionsModel("gpt-4", client, test_agent_name, agent_id="P1") # Store a strong reference to prevent garbage collection self._test_model = test_model # Add the message to the model's history # This will automatically add to AGENT_MANAGER via add_to_message_history test_model.add_to_message_history(msg) # No need to clean up _Converter state since it's now instance-based # 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.""" # No need for _Converter cleanup since it's now instance-based # Simulate the state before CTRL+C during tool execution # 1. User message self.add_to_test_message_history({"role": "user", "content": "Run a long command"}) # 2. Assistant message with tool call tool_call_id = "call_interrupted_123" self.add_to_test_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""" # Get the test model instance test_model = None for (name, instance_id), model_ref in ACTIVE_MODEL_INSTANCES.items(): if name == "test_agent": model = model_ref() if model_ref else None if model: test_model = model break if not test_model: return 0 # Simulate a tool call that was started but interrupted test_model._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(test_model._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 self.get_combined_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 self.get_combined_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", "{}"), }, } ], } self.add_to_test_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)", } self.add_to_test_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( self.get_combined_message_history() ) # TODO: Fix 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(self.get_combined_message_history()) >= 3, ( "Should have user, assistant, tool messages" ) # Check message roles in order roles = [msg["role"] for msg in self.get_combined_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 = self.get_combined_message_history()[1] tool_msg = self.get_combined_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!") # No need to clean up _Converter state since it's instance-based def test_fix_message_list_with_interrupted_tools(self): """Test fix_message_list handles interrupted tool sequences correctly.""" from cai.util import fix_message_list # No need for _Converter cleanup since it's now instance-based # Create an incomplete sequence (tool call without result) self.add_to_test_message_history({"role": "user", "content": "Test command"}) self.add_to_test_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(self.get_combined_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!") # No need to clean up _Converter state since it's instance-based return True except Exception as e: print(f"fix_message_list failed: {e}") # No need to clean up _Converter state since it's instance-based 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.""" # Clear history and check initial state initial_messages = self.get_combined_message_history() print(f"Initial message history (should be empty): {len(initial_messages)} messages") if initial_messages: for i, msg in enumerate(initial_messages): print(f" Unexpected initial message {i}: {msg}") # 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 i, msg in enumerate(test_messages): print(f"Adding message {i}: {msg['role']}") self.add_to_test_message_history(msg) current_count = len(self.get_combined_message_history()) print(f" Total messages after adding: {current_count}") # Verify OpenAI format compliance final_messages = self.get_combined_message_history() assert len(final_messages) == len(test_messages), ( f"Expected {len(test_messages)} messages, got {len(final_messages)}" ) for i, msg in enumerate(self.get_combined_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.""" # Clear history 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 self.add_to_test_message_history({"role": "user", "content": user_input}) # Add assistant message with tool call self.add_to_test_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 self.add_to_test_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(self.get_combined_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(self.get_combined_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.""" # No need for _Converter cleanup since it's now instance-based # Simulate a running tool call that gets interrupted call_id = "call_linux_cmd_123" tool_name = "generic_linux_command" # 1. Add user message self.add_to_test_message_history({"role": "user", "content": "Run a long command"}) # 2. Add assistant message with tool call (simulating qwen format) self.add_to_test_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 # Get the test model instance test_model = None for (name, instance_id), model_ref in ACTIVE_MODEL_INSTANCES.items(): if name == "test_agent": model = model_ref() if model_ref else None if model: test_model = model break if not test_model: self.fail("Could not find test model instance") # Add ONLY our specific tool call to recent_tool_calls test_model._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(test_model._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 self.get_combined_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)", } self.add_to_test_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 messages = self.get_combined_message_history() print(f"Message history after CTRL+C cleanup (total: {len(messages)}):") for i, msg in enumerate(messages): print(f" {i}: {msg.get('role')} - {msg}") # Assertions self.assertEqual(len(messages), 3, f"Expected 3 messages, got {len(messages)}: {messages}") # user + assistant + tool # Check user message self.assertEqual(self.get_combined_message_history()[0]["role"], "user") # Check assistant message has correct tool call self.assertEqual(self.get_combined_message_history()[1]["role"], "assistant") self.assertIsNotNone(self.get_combined_message_history()[1]["tool_calls"]) self.assertEqual(len(self.get_combined_message_history()[1]["tool_calls"]), 1) self.assertEqual(self.get_combined_message_history()[1]["tool_calls"][0]["id"], call_id) self.assertEqual( self.get_combined_message_history()[1]["tool_calls"][0]["function"]["name"], tool_name ) # Check tool response exists and is correct self.assertEqual(self.get_combined_message_history()[2]["role"], "tool") self.assertEqual(self.get_combined_message_history()[2]["tool_call_id"], call_id) self.assertIn("interrupted", self.get_combined_message_history()[2]["content"].lower()) # MOST IMPORTANT: Verify NO duplicate tool calls with unknown_function unknown_function_calls = [] for msg in self.get_combined_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!") # No need to clean up _Converter state since it's instance-based 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(test_instance.get_combined_message_history()) > 0: print("Message history debug:") for i, msg in enumerate(test_instance.get_combined_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)