diff --git a/check_s3_bucket.py b/check_s3_bucket.py new file mode 100644 index 00000000..512ff764 --- /dev/null +++ b/check_s3_bucket.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +import requests +import sys +from urllib.parse import urlparse +import xml.etree.ElementTree as ET + +def check_s3_bucket(bucket_name): + """Check if an S3 bucket is publicly accessible and list its contents if possible.""" + # Remove any s3.amazonaws.com suffix if present + if '.' in bucket_name: + parsed = urlparse('http://' + bucket_name) + bucket_name = parsed.netloc.split('.')[0] + + urls = [ + f"https://{bucket_name}.s3.amazonaws.com", + f"https://s3.amazonaws.com/{bucket_name}" + ] + + for url in urls: + print(f"Checking {url}") + try: + response = requests.get(url, timeout=10) + print(f"Status code: {response.status_code}") + + if response.status_code == 200: + print("Bucket is publicly accessible!") + + # Try to parse XML response to list bucket contents + try: + root = ET.fromstring(response.content) + ns = {'s3': 'http://s3.amazonaws.com/doc/2006-03-01/'} + + print("\nContents:") + for content in root.findall('.//s3:Contents', ns): + key = content.find('s3:Key', ns) + size = content.find('s3:Size', ns) + last_modified = content.find('s3:LastModified', ns) + + if key is not None: + print(f"File: {key.text}", end=" ") + if size is not None: + print(f"(Size: {size.text} bytes)", end=" ") + if last_modified is not None: + print(f"Last Modified: {last_modified.text}", end="") + print() + except Exception as e: + print(f"Error parsing bucket listing: {e}") + print("Raw response:") + print(response.text[:1000]) # Print first 1000 chars + elif response.status_code == 403: + print("Bucket exists but access is forbidden") + elif response.status_code == 404: + print("Bucket not found") + else: + print(f"Unexpected response with status code {response.status_code}") + + except requests.exceptions.RequestException as e: + print(f"Error: {e}") + + # Try alternate spelling as mentioned in the request + alt_url = f"https://trainingaliasrobotics.s3.amazonnaws.com" + print(f"\nChecking alternate spelling: {alt_url}") + try: + response = requests.get(alt_url, timeout=10) + print(f"Status code: {response.status_code}") + print(f"Response text: {response.text[:1000]}") + except requests.exceptions.RequestException as e: + print(f"Error: {e}") + +if __name__ == "__main__": + bucket_name = "trainingaliasrobotics" + check_s3_bucket(bucket_name) diff --git a/check_terrapin.py b/check_terrapin.py new file mode 100644 index 00000000..cb337e1e --- /dev/null +++ b/check_terrapin.py @@ -0,0 +1,46 @@ +import socket +import ssl +import sys + +def check_terrapin(hostname, port=443): + """ + Basic check for Terrapin vulnerability (CVE-2023-45866) + This is a simplified check and may not be 100% accurate + """ + try: + # Create socket and wrap with SSL + context = ssl.create_default_context() + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + + # Check if supports TLS 1.2 (Terrapin affects TLS 1.2 and below) + context.maximum_version = ssl.TLSVersion.TLSv1_2 + + with socket.create_connection((hostname, port), timeout=10) as sock: + with context.wrap_socket(sock, server_hostname=hostname) as ssock: + cipher = ssock.cipher() + if cipher and cipher[0]: + cipher_name = cipher[0] + print(f"Connected using: {cipher_name}") + + # Check if using CBC cipher (Terrapin affects CBC ciphers) + if 'CBC' in cipher_name: + print(f"[POTENTIALLY VULNERABLE] {hostname} might be vulnerable to Terrapin") + print(f"Using CBC cipher: {cipher_name}") + return True + else: + print(f"[LIKELY SAFE] {hostname} is using non-CBC cipher: {cipher_name}") + return False + except ssl.SSLError as e: + print(f"SSL Error: {e}") + except socket.error as e: + print(f"Socket Error: {e}") + except Exception as e: + print(f"Error checking {hostname}: {e}") + + return None + +if __name__ == "__main__": + hostname = "aliasrobotics.com" + print(f"Checking {hostname} for potential Terrapin vulnerability...") + check_terrapin(hostname) diff --git a/check_vulns.py b/check_vulns.py new file mode 100644 index 00000000..a710f9c2 --- /dev/null +++ b/check_vulns.py @@ -0,0 +1,87 @@ +import subprocess +import json +import sys +import socket +import ssl +import requests +from urllib.parse import urlparse + +def check_terrapin(hostname, port=443): + print(f"[*] Checking {hostname}:{port} for Terrapin vulnerability (CVE-2023-48618)") + try: + # Create a socket and wrap it with SSL + context = ssl.create_default_context() + with socket.create_connection((hostname, port)) as sock: + with context.wrap_socket(sock, server_hostname=hostname) as ssock: + # Check if server supports TLS 1.3 + version = ssock.version() + print(f"[*] Server SSL/TLS version: {version}") + + # Terrapin affects TLS 1.3 implementation with specific configurations + if 'TLSv1.3' in version: + print("[!] Server supports TLS 1.3, which may be vulnerable to Terrapin attack") + print("[!] Note: Detailed testing requires specialized tools, as this is a downgrade attack") + print("[!] The vulnerability allows attackers to force downgrade to CBC ciphers in TLS 1.2") + else: + print("[*] Server does not use TLS 1.3, not directly affected by Terrapin") + except Exception as e: + print(f"[-] Error checking for Terrapin: {e}") + +def check_s3_bucket(bucket_name): + print(f"[*] Checking S3 bucket: {bucket_name}") + + # Clean up the bucket name if needed + if bucket_name.startswith("http://") or bucket_name.startswith("https://"): + parsed = urlparse(bucket_name) + bucket_name = parsed.netloc + + if ".s3.amazonaws" in bucket_name: + bucket_name = bucket_name.split(".s3.amazonaws")[0] + + # Direct URL + url = f"http://{bucket_name}.s3.amazonaws.com" + try: + response = requests.get(url) + print(f"[*] Bucket URL: {url}") + print(f"[*] Response status: {response.status_code}") + + if response.status_code == 200: + print("[!] Bucket is publicly accessible") + print("[*] Response content:") + # Print the first 1000 characters of the response content + print(response.text[:1000]) + elif response.status_code == 403: + print("[!] Bucket exists but access is forbidden (listing denied)") + print("[*] This could mean the bucket exists but is not publicly listable") + elif response.status_code == 404: + print("[*] Bucket not found") + else: + print(f"[*] Unexpected response: {response.status_code}") + except Exception as e: + print(f"[-] Error checking bucket: {e}") + + # Try to list bucket objects with aws cli + try: + print("[*] Attempting to list bucket objects with AWS CLI...") + result = subprocess.run( + ["aws", "s3", "ls", f"s3://{bucket_name}", "--no-sign-request"], + capture_output=True, + text=True + ) + if result.returncode == 0 and result.stdout: + print("[!] Bucket is listable without authentication!") + print("[*] Bucket contents:") + print(result.stdout) + else: + print("[*] Bucket is not listable without authentication") + if result.stderr: + print(f"[*] Error: {result.stderr}") + except Exception as e: + print(f"[-] Error using AWS CLI: {e}") + +if __name__ == "__main__": + # Check for Terrapin vulnerability + check_terrapin("aliasrobotics.com") + + # Check for exposed S3 bucket + check_s3_bucket("trainingaliasrobotics") diff --git a/hello_world.py b/hello_world.py new file mode 100644 index 00000000..09907203 --- /dev/null +++ b/hello_world.py @@ -0,0 +1 @@ +print('Hello, World!') diff --git a/helloworld.py b/helloworld.py new file mode 100644 index 00000000..09907203 --- /dev/null +++ b/helloworld.py @@ -0,0 +1 @@ +print('Hello, World!') diff --git a/hola_mundo.py b/hola_mundo.py new file mode 100644 index 00000000..a500c2f4 --- /dev/null +++ b/hola_mundo.py @@ -0,0 +1 @@ +print("Β‘Hola Mundo!") diff --git a/nmap_results b/nmap_results new file mode 100644 index 00000000..e69de29b diff --git a/security_check.py b/security_check.py new file mode 100644 index 00000000..466331f0 --- /dev/null +++ b/security_check.py @@ -0,0 +1,71 @@ +import socket +import ssl + +def test_terrapin(hostname, port=443): + """ + Test if a server is vulnerable to the Terrapin Attack (CVE-2023-48795) + This is a basic test that checks for specific TLS behavior related to CBC ciphers + """ + print(f"Testing {hostname}:{port} for Terrapin vulnerability (CVE-2023-48795)...") + + context = ssl.create_default_context() + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + + # Force the use of CBC ciphers which are affected by the attack + context.set_ciphers('AES256-SHA:AES128-SHA:DES-CBC3-SHA') + + try: + with socket.create_connection((hostname, port)) as sock: + with context.wrap_socket(sock, server_hostname=hostname) as ssock: + cipher = ssock.cipher() + print(f"Connected using: {cipher}") + + # If we can connect with CBC ciphers, the server might be vulnerable + if "CBC" in cipher[0]: + print("Server accepts CBC ciphers which are potentially vulnerable to Terrapin Attack") + print("Note: This is only an indicator and not a definitive test") + print("Full verification would require specialized tools and in-depth testing") + return True + else: + print("Server not using CBC ciphers in this test") + return False + + except Exception as e: + print(f"Error during test: {e}") + return False + +# Also check for the AWS S3 bucket +def check_s3_bucket(bucket_name): + """ + Check if an S3 bucket exists and is accessible + """ + print(f"\nChecking S3 bucket: {bucket_name}") + try: + with socket.create_connection((bucket_name, 443)) as sock: + # S3 should resolve so we can connect, which suggests the bucket exists + print(f"Bucket {bucket_name} exists (DNS resolution successful)") + + # We try to make a simple HTTPS request + context = ssl.create_default_context() + with context.wrap_socket(sock, server_hostname=bucket_name) as ssock: + ssock.sendall(b"GET / HTTP/1.1\r\nHost: " + bucket_name.encode() + b"\r\n\r\n") + data = ssock.recv(4096) + print(f"Response from bucket:\n{data.decode('utf-8', errors='ignore')[:500]}") + + if b"AccessDenied" in data or b"403 Forbidden" in data: + print("Bucket exists but access is denied (this is normal for protected buckets)") + return True + else: + print("Bucket might be publicly accessible!") + return True + + except Exception as e: + print(f"Error checking bucket: {e}") + if "Name or service not known" in str(e): + print("Bucket does not exist or is not accessible") + return False + +# Test both issues +test_terrapin("aliasrobotics.com") +check_s3_bucket("trainingaliasrobotics.s3.amazonaws.com") diff --git a/src/cai/cli.py b/src/cai/cli.py index 37ed9dc7..0a980e27 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -442,6 +442,13 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= } ) + # Fix message list structure BEFORE sending to the model to prevent errors + try: + from cai.util import fix_message_list + history_context = fix_message_list(history_context) + except Exception as e: + console.print(f"[yellow]Warning: Could not preprocess message history: {e}[/yellow]") + # Append the current user input as the last message in the list. conversation_input: list | str if history_context: diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index a68057c2..cd921ac9 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -336,6 +336,15 @@ class OpenAIChatCompletionsModel(Model): # Log the user message if item.get("content"): self.logger.log_user_message(item.get("content")) + + # IMPORTANT: Ensure the message list has valid tool call/result pairs + # This needs to happen before the API call to prevent errors + try: + from cai.util import fix_message_list + converted_messages = fix_message_list(converted_messages) + except Exception as e: + logger.warning(f"Failed to fix message list: {e}") + # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) @@ -433,6 +442,7 @@ class OpenAIChatCompletionsModel(Model): interaction_cost=None, total_cost=None, tool_output=None, # Don't pass tool output here, we're using direct display + suppress_empty=True # Suppress empty panels ) # --- Add assistant tool call to message_history if present --- @@ -457,6 +467,22 @@ class OpenAIChatCompletionsModel(Model): } add_to_message_history(tool_call_msg) + + # Save the tool call details for later matching with output + # This is important for non-streaming mode to track tool calls properly + if not hasattr(_Converter, 'recent_tool_calls'): + _Converter.recent_tool_calls = {} + + # Store the tool call by ID for later reference + import time + _Converter.recent_tool_calls[tool_call.id] = { + 'name': tool_call.function.name, + 'arguments': tool_call.function.arguments, + 'start_time': time.time(), + 'execution_info': { + 'start_time': time.time() + } + } # Log the assistant tool call message tool_calls_list = [] @@ -583,7 +609,7 @@ class OpenAIChatCompletionsModel(Model): stop_idle_timer() start_active_timer() - # Check if streaming should be shown in rich panel + # --- Check if streaming should be shown in rich panel --- should_show_rich_stream = os.getenv('CAI_STREAM', 'false').lower() == 'true' and not self.disable_rich_streaming # Create streaming context if needed @@ -922,6 +948,55 @@ class OpenAIChatCompletionsModel(Model): if tool_call_msg not in streamed_tool_calls: streamed_tool_calls.append(tool_call_msg) add_to_message_history(tool_call_msg) + + # NEW: Display tool call immediately when detected in streaming mode + # But only if it has complete arguments and name + if (state.function_calls[tc_index].name and + state.function_calls[tc_index].arguments and + state.function_calls[tc_index].call_id): + # First, finish any existing streaming context if it exists + if streaming_context: + try: + finish_agent_streaming(streaming_context, None) + streaming_context = None + except Exception: + pass + + # Create a message-like object for displaying the function call + tool_msg = type('ToolCallStreamDisplay', (), { + 'content': None, + 'tool_calls': [ + type('ToolCallDetail', (), { + 'function': type('FunctionDetail', (), { + 'name': state.function_calls[tc_index].name, + 'arguments': state.function_calls[tc_index].arguments + }), + 'id': state.function_calls[tc_index].call_id, + 'type': 'function' + }) + ] + }) + + # Display the tool call during streaming + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), + message=tool_msg, + counter=getattr(self, 'interaction_counter', 0), + model=str(self.model), + debug=False, + interaction_input_tokens=estimated_input_tokens, + interaction_output_tokens=estimated_output_tokens, + interaction_reasoning_tokens=0, # Not available during streaming yet + total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens, + total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens, + total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), + interaction_cost=None, + total_cost=None, + tool_output=None, # Will be shown once tool is executed + suppress_empty=True # Prevent empty panels + ) + # Set flag to suppress final output to avoid duplication + self.suppress_final_output = True except Exception as e: # Ensure streaming context is cleaned up in case of errors @@ -983,8 +1058,15 @@ class OpenAIChatCompletionsModel(Model): ) # Display the tool call in CLI - from cai.util import cli_print_agent_messages try: + # First, finish any existing streaming context if it exists + if streaming_context: + try: + finish_agent_streaming(streaming_context, None) + streaming_context = None + except Exception: + pass + # Create a message-like object to display the function call tool_msg = type('ToolCallWrapper', (), { 'content': None, @@ -1015,8 +1097,12 @@ class OpenAIChatCompletionsModel(Model): total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), interaction_cost=None, total_cost=None, - tool_output=None # Will be shown once the tool is executed + tool_output=None, # Will be shown once the tool is executed + suppress_empty=True # Suppress empty panels during streaming ) + + # Set flag to suppress final output to avoid duplication + self.suppress_final_output = True except Exception as e: logger.error(f"Error displaying tool call in CLI: {e}") @@ -1220,6 +1306,7 @@ class OpenAIChatCompletionsModel(Model): direct_stats["total_cost"] = float(total_cost) # Use the direct copy with guaranteed float costs finish_agent_streaming(streaming_context, direct_stats) + streaming_context = None # Removed extra newline after streaming completes to avoid blank lines pass @@ -1243,8 +1330,10 @@ class OpenAIChatCompletionsModel(Model): for tool_call in tool_call_msg.get("tool_calls", []): tool_calls_list.append(tool_call) self.logger.log_assistant_message(None, tool_calls_list) - # If there was only text output, add that as an assistant message - if (not streamed_tool_calls) and state.text_content_index_and_output and state.text_content_index_and_output[1].text: + + # If we've already shown the tool calls directly during streaming, + # don't log the text message to avoid duplication + elif (not self.suppress_final_output) and state.text_content_index_and_output and state.text_content_index_and_output[1].text: asst_msg = { "role": "assistant", "content": state.text_content_index_and_output[1].text @@ -1253,6 +1342,9 @@ class OpenAIChatCompletionsModel(Model): # Log the assistant message self.logger.log_assistant_message(state.text_content_index_and_output[1].text) + # Reset the suppress flag for future requests + self.suppress_final_output = False + # Log the complete response self.logger.rec_training_data( { @@ -2042,7 +2134,11 @@ class _Converter: if current_assistant_msg is not None: # The API doesn't support empty arrays for tool_calls if not current_assistant_msg.get("tool_calls"): - del current_assistant_msg["tool_calls"] + # Ensure content is not None if tool_calls are absent and content is also None + # Some models like Anthropic require some content, even if it's just a placeholder. + if current_assistant_msg.get("content") is None: + current_assistant_msg["content"] = "(No text content in this assistant message)" # Or just an empty string if preferred + current_assistant_msg.pop("tool_calls", None) # Use pop with default to avoid KeyError result.append(current_assistant_msg) current_assistant_msg = None @@ -2079,16 +2175,28 @@ class _Converter: flush_assistant_message() tool_calls_param: list[ChatCompletionMessageToolCallParam] = [] for tc in item["tool_calls"]: + function_details = tc.get("function", {}) + arguments = function_details.get("arguments") + # Ensure arguments is a valid JSON string, defaulting to "{}" if empty or None + if arguments is None or (isinstance(arguments, str) and arguments.strip() == ""): + arguments = "{}" + elif isinstance(arguments, dict): + # Ensure it's a string if it's a dict (should already be string per schema) + arguments = json.dumps(arguments) + tool_calls_param.append( ChatCompletionMessageToolCallParam( id=tc.get("id", ""), type=tc.get("type", "function"), - function=tc.get("function", {}), + function={ + "name": function_details.get("name", "unknown_function"), + "arguments": arguments, # Use sanitized arguments + }, ) ) msg_asst: ChatCompletionAssistantMessageParam = { "role": "assistant", - "content": item.get("content"), + "content": item.get("content"), # Content can be None here "tool_calls": tool_calls_param, } result.append(msg_asst) @@ -2225,12 +2333,19 @@ class _Converter: } } + arguments = func_call.get("arguments") # func_call is a dict here + # Ensure arguments is a valid JSON string, defaulting to "{}" if empty or None + if arguments is None or (isinstance(arguments, str) and arguments.strip() == ""): + arguments = "{}" + elif isinstance(arguments, dict): + arguments = json.dumps(arguments) + new_tool_call = ChatCompletionMessageToolCallParam( id=func_call["call_id"], type="function", function={ "name": func_call["name"], - "arguments": func_call["arguments"], + "arguments": arguments, # Use sanitized arguments }, ) tool_calls.append(new_tool_call) @@ -2270,64 +2385,64 @@ class _Converter: # Display the tool output immediately with the matched tool call from cai.util import cli_print_tool_output - # Check if we're in streaming mode - don't show tool output panel in streaming mode + # Look up the original tool call to get the name and arguments + tool_name = "Unknown Tool" + tool_args = {} + execution_info = {} + + if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: + tool_call = cls.recent_tool_calls[call_id] + tool_name = tool_call.get('name', 'Unknown Tool') + tool_args = tool_call.get('arguments', {}) + execution_info = tool_call.get('execution_info', {}) + + # Get token counts from the OpenAIChatCompletionsModel if available + model_instance = None + for frame in inspect.stack(): + if 'self' in frame.frame.f_locals: + self_obj = frame.frame.f_locals['self'] + if isinstance(self_obj, OpenAIChatCompletionsModel): + model_instance = self_obj + break + + # Always create a token_info dictionary, even if some values are zero + token_info = { + 'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0), + 'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0), + 'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0), + 'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0), + 'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0), + 'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0), + 'model': str(getattr(model_instance, 'model', '')), + } + + # Calculate costs using standard cost model + if model_instance and hasattr(model_instance, 'model'): + from cai.util import calculate_model_cost + model_name = str(model_instance.model) + token_info['interaction_cost'] = calculate_model_cost( + model_name, + token_info['interaction_input_tokens'], + token_info['interaction_output_tokens'] + ) + token_info['total_cost'] = calculate_model_cost( + model_name, + token_info['total_input_tokens'], + token_info['total_output_tokens'] + ) + + # Check if we're in streaming mode is_streaming_enabled = os.environ.get('CAI_STREAM', 'false').lower() == 'true' - if is_streaming_enabled: - # Don't display tool output in streaming mode - it will be handled elsewhere - pass # Just skip the display, but preserve the tool output - else: - # For non-streaming mode, maintain the original behavior - # Look up the original tool call to get the name and arguments - if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: - tool_call = cls.recent_tool_calls[call_id] - tool_name = tool_call.get('name', 'Unknown Tool') - tool_args = tool_call.get('arguments', {}) - execution_info = tool_call.get('execution_info', {}) - - # Get token counts from the OpenAIChatCompletionsModel if available - model_instance = None - for frame in inspect.stack(): - if 'self' in frame.frame.f_locals: - self_obj = frame.frame.f_locals['self'] - if isinstance(self_obj, OpenAIChatCompletionsModel): - model_instance = self_obj - break - - # Always create a token_info dictionary, even if some values are zero - token_info = { - 'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0), - 'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0), - 'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0), - 'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0), - 'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0), - 'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0), - 'model': str(getattr(model_instance, 'model', '')), - } - - # Calculate costs using standard cost model - if model_instance and hasattr(model_instance, 'model'): - from cai.util import calculate_model_cost - model_name = str(model_instance.model) - token_info['interaction_cost'] = calculate_model_cost( - model_name, - token_info['interaction_input_tokens'], - token_info['interaction_output_tokens'] - ) - token_info['total_cost'] = calculate_model_cost( - model_name, - token_info['total_input_tokens'], - token_info['total_output_tokens'] - ) - - # Use the cli_print_tool_output function with actual token values - cli_print_tool_output( - tool_name=tool_name, - args=tool_args, - output=output_content, - call_id=call_id, # Keep call_id for non-streaming mode - execution_info=execution_info, - token_info=token_info - ) + + # Always display tool output regardless of streaming mode + cli_print_tool_output( + tool_name=tool_name, + args=tool_args, + output=output_content, + call_id=call_id, + execution_info=execution_info, + token_info=token_info + ) # Continue with normal processing flush_assistant_message() diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 2b4327da..fa5d6b8f 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -13,7 +13,7 @@ import uuid import sys import shlex from wasabi import color # pylint: disable=import-error -from cai.util import format_time, start_active_timer, stop_active_timer, start_idle_timer, stop_idle_timer +from cai.util import format_time, start_active_timer, stop_active_timer, start_idle_timer, stop_idle_timer, cli_print_tool_output # Instead of direct import @@ -436,232 +436,61 @@ def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None): def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None): """Runs command locally in the specified workspace_dir.""" - # If streaming is enabled and we have a call_id - if stream and call_id: - return _run_local_streamed(command, call_id, timeout, tool_name, workspace_dir) - # Make sure we're in active time mode for tool execution stop_idle_timer() start_active_timer() + process_start_time = time.time() # Initialize with current time try: target_dir = workspace_dir or _get_workspace_dir() original_cmd_for_msg = command # For logging context_msg = f"(local:{target_dir})" - try: - result = subprocess.run( - command, - shell=True, # nosec B602 - capture_output=True, - text=True, - check=False, - timeout=timeout, - cwd=target_dir - ) - output = result.stdout if result.stdout else result.stderr - if stdout: - print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 - - # Skip passing output to cli_print_tool_output when CAI_STREAM=true - # This prevents duplicate output in streaming mode - is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' - if not is_streaming_enabled: - # Optional: Add cli_print_tool_output call here if needed for non-streaming - pass - - return output.strip() - except subprocess.TimeoutExpired as e: - error_output = e.stdout if e.stdout else str(e) - if stdout: - print("\033[32m" + error_output + "\033[0m") - return error_output - except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing local command: {e}" - print(color(error_msg, fg="red")) - return error_msg - finally: - # Always switch back to idle mode when function completes - stop_active_timer() - start_idle_timer() - - -def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace_dir=None): - """Run a local command with streaming output to the Tool output panel.""" - # Make sure we're in active time mode for tool execution - stop_idle_timer() - start_active_timer() - - target_dir = workspace_dir or _get_workspace_dir() - try: - # Try to import Rich for nice display - try: - from rich.console import Console - from rich.live import Live - from rich.panel import Panel - from rich.text import Text - from rich.box import ROUNDED - console = Console() - rich_available = True - except ImportError: - rich_available = False - from cai.util import cli_print_tool_output + + # If streaming is enabled and we have a call_id + if stream: + # Import the streaming utilities from util + from cai.util import start_tool_streaming, update_tool_streaming, finish_tool_streaming - output_buffer = [] - start_time = time.time() - # Start the process - process = subprocess.Popen( - command, - shell=True, # nosec B602 - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - bufsize=1, - cwd=target_dir # Set CWD for local process - ) - - # If tool_name is not provided, derive it from the command - if tool_name is None: - # Just use the first command as the tool name - tool_name = command.strip().split()[0] + "_command" - - # Create panel content for Rich display - if rich_available: - # Parse command into command and args + # Parse command into parts for display parts = command.strip().split(' ', 1) cmd = parts[0] if parts else "" - args = parts[1] if len(parts) > 1 else "" + args_param_val = parts[1] if len(parts) > 1 else "" # Renamed to avoid conflict with tool_args dict key - # Format clean arguments, following the same rules as cli_print_tool_output - arg_parts = [] - if cmd: - arg_parts.append(f"command={cmd}") - if args and args.strip(): # Only add args if non-empty - arg_parts.append(f"args={args}") - args_str = ", ".join(arg_parts) + # For generic Linux commands, standardize the tool_name format + if not tool_name: + tool_name = f"{cmd}_command" if cmd else "command" - header = Text() - header.append(tool_name, style="#00BCD4") - header.append("(", style="yellow") - header.append(args_str, style="yellow") - header.append(")", style="yellow") - tool_time = 0 - - total_time = 0 - if START_TIME is not None: - total_time = time.time() - START_TIME - timing_info = [] - if total_time: - timing_info.append(f"Total: {format_time(total_time)}") - if tool_time: - timing_info.append(f"Tool: {format_time(tool_time)}") - if timing_info: - header.append(f" [{' | '.join(timing_info)}]", style="cyan") - - content = Text() - - panel = Panel( - Text.assemble(header, "\n\n", content), - title="[bold green]Tool Execution[/bold green]", - subtitle="[bold green]Live Output[/bold green]", - border_style="green", - padding=(1, 2), - box=ROUNDED - ) - - # Start Live display - with Live(panel, console=console, refresh_per_second=4) as live: - # Stream stdout in real-time - start_time = time.time() - for line in iter(process.stdout.readline, ''): - if not line: - break - - # Add to output collection - output_buffer.append(line) - - # Update content with new line - content.append(line, style="bright_white") - - # Update tool_time and header with new timing info - tool_time = time.time() - start_time - total_time = 0 - if START_TIME is not None: - total_time = time.time() - START_TIME - # Remove any previous timing info from header (rebuild header) - timing_info = [] - if total_time: - timing_info.append(f"Total: {format_time(total_time)}") - if tool_time: - timing_info.append(f"Tool: {format_time(tool_time)}") - # Rebuild header to update timing - header = Text() - header.append(tool_name, style="#00BCD4") - header.append("(", style="yellow") - header.append(args_str, style="yellow") - header.append(")", style="yellow") - if timing_info: - header.append(f" [{' | '.join(timing_info)}]", style="cyan") - - panel = Panel( - Text.assemble(header, "\n\n", content), - title="[bold green]Tool Execution[/bold green]", - subtitle="[bold green]Live Output[/bold green]", - border_style="green", - padding=(1, 2), - box=ROUNDED - ) - live.update(panel) - # Check if process is done - process.stdout.close() - return_code = process.wait(timeout=timeout) - - # Get any stderr output - stderr_data = process.stderr.read() - if stderr_data: - content.append("\nERROR OUTPUT:\n", style="red") - content.append(stderr_data, style="red") - output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) - panel = Panel( - Text.assemble(header, "\n\n", content), - title="[bold green]Tool Execution[/bold green]", - subtitle="[bold green]Live Output[/bold green]", - border_style="green", - padding=(1, 2), - box=ROUNDED - ) - live.update(panel) - - # Add completion message - panel = Panel( - Text.assemble(header, "\n\n", content), - title="[bold green]Tool Execution[/bold green]", - border_style="green", - padding=(1, 2), - box=ROUNDED - ) - live.update(panel) - - # Wait a moment for the panel to be displayed properly - time.sleep(0.5) - else: - # Fallback to simpler streaming with cli_print_tool_output - # Parse command into command and args (same as rich mode) - parts = command.strip().split(' ', 1) - cmd = parts[0] if parts else "" - args = parts[1] if len(parts) > 1 else "" - - # Create a dictionary with only non-empty values (following the same rules) + # Create args dictionary with non-empty values only tool_args = {} if cmd: tool_args["command"] = cmd - if args and args.strip(): - tool_args["args"] = args - # Note: Omitted empty values and async_mode=False as it's default + if args_param_val and args_param_val.strip(): + tool_args["args"] = args_param_val - # Initial notification - just once - cli_print_tool_output(tool_name, tool_args, "Command started...", call_id=call_id) + # Add more context for the command + tool_args["workspace"] = os.path.basename(target_dir) + tool_args["full_command"] = command - # Buffer for collecting output + # For generic commands, ensure we have a unique call_id + if not call_id: + call_id = f"cmd_{cmd}_{str(uuid.uuid4())[:8]}" + + # Initialize/use the call_id for this streaming session + call_id = start_tool_streaming(tool_name, tool_args, call_id) + + # Start the process + process = subprocess.Popen( + command, + shell=True, # nosec B602 + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + cwd=target_dir + ) + + # Begin collecting output + output_buffer = [] buffer_size = 0 update_interval = 10 # lines @@ -674,52 +503,133 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace output_buffer.append(line) buffer_size += 1 - # Only update the output periodically to reduce panel refresh rate + # Only update periodically to reduce UI refreshes if buffer_size >= update_interval: current_output = ''.join(output_buffer) - cli_print_tool_output(tool_name, tool_args, current_output, call_id=call_id) + update_tool_streaming(tool_name, tool_args, current_output, call_id) buffer_size = 0 - # Check if process is done + # Finish process process.stdout.close() return_code = process.wait(timeout=timeout) + process_execution_time = time.time() - process_start_time # Get any stderr output stderr_data = process.stderr.read() if stderr_data: output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) - # Final output update - always show the final result + # Final output update final_output = ''.join(output_buffer) if return_code != 0: final_output += f"\nCommand exited with code {return_code}" - cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id) - - # Return the full output - return ''.join(output_buffer) - - except subprocess.TimeoutExpired: - error_msg = f"Command timed out after {timeout} seconds" - output_buffer.append("\n" + error_msg) - final_output = ''.join(output_buffer) - - # Update tool output panel with timeout message - if not rich_available: - tool_args = {"command": command} - cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id) - - return final_output + # Calculate execution info with environment details + execution_info = { + "status": "completed" if return_code == 0 else "error", + "return_code": return_code, + "environment": "Local", + "host": os.path.basename(target_dir), + "tool_time": process_execution_time + } + + # Complete the streaming session with final output + finish_tool_streaming(tool_name, tool_args, final_output, call_id, execution_info) + + return final_output + else: + # Standard non-streaming execution + result = subprocess.run( + command, + shell=True, # nosec B602 + capture_output=True, + text=True, + check=False, + timeout=timeout, + cwd=target_dir + ) + output = result.stdout if result.stdout else result.stderr + + # If this is the same command that was recently streamed, don't display it again + # We'll create a similar tool_args structure to what streaming would use + parts = command.strip().split(' ', 1) + cmd = parts[0] if parts else "" + args_param_val = parts[1] if len(parts) > 1 else "" + + # Generate a standard tool name for consistency with streaming + standard_tool_name = tool_name or (f"{cmd}_command" if cmd else "command") + + # Calculate a consistent command key - must match the format used in cli_print_tool_output + command_key = f"{standard_tool_name}:{args_param_val}" + + if hasattr(cli_print_tool_output, '_displayed_commands') and command_key in cli_print_tool_output._displayed_commands: + # Skip stdout display if already shown through streaming + return output.strip() + + if stdout: + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") + + return output.strip() + except subprocess.TimeoutExpired as e: + error_output = e.stdout if hasattr(e, 'stdout') and e.stdout else str(e) + error_msg = f"Command timed out after {timeout} seconds\n{error_output}" + # If we're streaming, show the timeout in the tool output panel + if stream and call_id: + from cai.util import finish_tool_streaming + # Parse the command the same way we did for streaming + parts = command.strip().split(' ', 1) + cmd = parts[0] if parts else "" + args = parts[1] if len(parts) > 1 else "" + + # Ensure tool_args has complete information + tool_args = { + "command": cmd, + "args": args if args.strip() else "", + "full_command": command, + "environment": "Local", + "workspace": os.path.basename(target_dir) + } + execution_info = { + "status": "timeout", + "error": str(e), + "environment": "Local", + "host": os.path.basename(target_dir) + } + finish_tool_streaming(tool_name or f"{cmd}_command", tool_args, error_msg, call_id, execution_info) + + if stdout: + print("\033[32m" + error_msg + "\033[0m") + + return error_msg except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing command: {str(e)}" + error_msg = f"Error executing local command: {e}" + + # If we're streaming, show the error in the tool output panel + if stream and call_id: + from cai.util import finish_tool_streaming + # Parse the command the same way we did for streaming + parts = command.strip().split(' ', 1) + cmd = parts[0] if parts else "" + args = parts[1] if len(parts) > 1 else "" + + # Ensure tool_args has complete information + tool_args = { + "command": cmd, + "args": args if args.strip() else "", + "full_command": command, + "environment": "Local", + "workspace": os.path.basename(target_dir) + } + execution_info = { + "status": "error", + "error": str(e), + "environment": "Local", + "host": os.path.basename(target_dir) + } + finish_tool_streaming(tool_name or f"{cmd}_command", tool_args, error_msg, call_id, execution_info) + print(color(error_msg, fg="red")) - - # Update tool output panel with error message if simple streaming - if not rich_available: - tool_args = {"command": command} - cli_print_tool_output(tool_name, tool_args, error_msg, call_id=call_id) - return error_msg finally: # Always switch back to idle mode when function completes @@ -753,6 +663,20 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_idle_timer() start_active_timer() + # Parse command into standard parts to ensure consistent naming + parts = command.strip().split(' ', 1) + cmd = parts[0] if parts else "" + args = parts[1] if len(parts) > 1 else "" + + # Generate a call_id if we're streaming and one wasn't provided + # Use a more specific format that includes the command name for easier tracking + if not call_id and stream: + call_id = f"cmd_{cmd}_{str(uuid.uuid4())[:8]}" + + # If no tool_name is provided, derive it from the command in a consistent way + if not tool_name: + tool_name = f"{cmd}_command" if cmd else "command" + try: # If session_id is provided, send command to that session if session_id: @@ -781,10 +705,6 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg return result # Return the result of sending input ("Input sent..." or error) - # Generate a call_id if we're streaming and one wasn't provided - if stream and not call_id: - call_id = str(uuid.uuid4())[:8] - # 2. Determine Execution Environment (Container > CTF > SSH > Local) active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST']) @@ -817,6 +737,30 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # Handle Streaming Container Execution if stream: + # Import the streaming utilities from util + from cai.util import start_tool_streaming, update_tool_streaming, finish_tool_streaming + + # Create args dictionary with standardized format + tool_args = { + "command": cmd, + "args": args if args.strip() else "", + "full_command": command, + "container": container_id[:12], + "environment": "Container", + "workspace": container_workspace + } + + # Initialize the streaming session with a consistent call_id format + call_id = start_tool_streaming(tool_name, tool_args, call_id) + + # Update with "executing" status + update_tool_streaming( + tool_name, + tool_args, + f"Executing in container {container_id[:12]} at {container_workspace}:\n{command}\n\nPreparing environment...", + call_id + ) + # Ensure workspace directory exists inside the container first mkdir_cmd = [ "docker", "exec", container_id, @@ -829,6 +773,14 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg check=False, timeout=10 ) + + # Update status once environment is prepared + update_tool_streaming( + tool_name, + tool_args, + f"Executing in container {container_id[:12]} at {container_workspace}:\n{command}\n\nRunning command...", + call_id + ) # Build docker exec command as a single shell string for streaming docker_exec_cmd = ( @@ -837,20 +789,114 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg f"{shlex.quote(container_id)} sh -c " f"{shlex.quote(command)}" ) - - # Re-use the local streaming helper to provide real-time output - result = _run_local_streamed( - docker_exec_cmd, - call_id, - timeout, - tool_name, - workspace_dir=_get_workspace_dir() - ) - # Switch back to idle mode after streaming command completes - stop_active_timer() - start_idle_timer() - return result + try: + start_time = time.time() + # Start the process + process = subprocess.Popen( + docker_exec_cmd, + shell=True, # nosec B602 + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + cwd=_get_workspace_dir() + ) + + # Begin collecting output + output_buffer = [] + buffer_size = 0 + update_interval = 10 # lines + + # Stream stdout in real-time + for line in iter(process.stdout.readline, ''): + if not line: + break + + # Add to output collection + output_buffer.append(line) + buffer_size += 1 + + # Only update periodically to reduce UI refreshes + if buffer_size >= update_interval: + current_output = ''.join(output_buffer) + update_tool_streaming(tool_name, tool_args, current_output, call_id) + buffer_size = 0 + + # Finish process + process.stdout.close() + return_code = process.wait(timeout=timeout) + execution_time = time.time() - start_time + + # Get any stderr output + stderr_data = process.stderr.read() + if stderr_data: + output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) + + # Final output update + final_output = ''.join(output_buffer) + if return_code != 0: + final_output += f"\nCommand exited with code {return_code}" + + # Calculate execution info + execution_info = { + "status": "completed" if return_code == 0 else "error", + "return_code": return_code, + "environment": "Container", + "host": container_id[:12], + "tool_time": execution_time + } + + # Complete the streaming session with final output + finish_tool_streaming(tool_name, tool_args, final_output, call_id, execution_info) + + # Switch back to idle mode after streaming command completes + stop_active_timer() + start_idle_timer() + return final_output + + except subprocess.TimeoutExpired as e: + # Handle timeout + error_output = e.stdout if hasattr(e, 'stdout') and e.stdout else str(e) + error_msg = f"Command timed out after {timeout} seconds\n{error_output}" + + execution_info = { + "status": "timeout", + "environment": "Container", + "host": container_id[:12], + "error": str(e) + } + + # Complete with timeout error + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after timeout + stop_active_timer() + start_idle_timer() + # Fallback to local execution on timeout + print(color("Container execution timed out. Attempting execution on host instead.", fg="yellow")) + return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir()) + + except Exception as e: + # Handle other errors + error_msg = f"Error executing command in container: {str(e)}" + + execution_info = { + "status": "error", + "environment": "Container", + "host": container_id[:12], + "error": str(e) + } + + # Complete with error + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after error + stop_active_timer() + start_idle_timer() + # Fallback to local execution on error + print(color("Container execution failed. Attempting execution on host instead.", fg="yellow")) + return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir()) # Handle Synchronous Execution in Container try: @@ -916,54 +962,204 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # --- CTF Execution --- if ctf: - # Handling streaming for CTF - not fully implemented yet + # If streaming is enabled and we have a call_id, show streaming UI for CTF too if stream: - from cai.util import cli_print_tool_output - if call_id and tool_name: - tool_args = {"command": command, "ctf": True} - cli_print_tool_output( - tool_name, - tool_args, - "Streaming not yet supported for CTF execution. Running normally...", - call_id=call_id - ) + # Import the streaming utilities from util + from cai.util import start_tool_streaming, update_tool_streaming, finish_tool_streaming + + # Create args dictionary with standardized format + tool_args = { + "command": cmd, + "args": args if args.strip() else "", + "full_command": command, + "environment": "CTF", + "workspace": os.path.basename(_get_workspace_dir()) + } + + # Initialize the streaming session with a consistent call_id format + call_id = start_tool_streaming(tool_name, tool_args, call_id) + + target_dir = _get_workspace_dir() + full_command = f"cd '{target_dir}' && {command}" + + # Update with "executing" status + update_tool_streaming( + tool_name, + tool_args, + f"Executing in CTF environment: {full_command}\n\nWaiting for response...", + call_id + ) + + try: + # Execute the command and get the output + start_time = time.time() + output = ctf.get_shell(full_command, timeout=timeout) + execution_time = time.time() - start_time + + # Calculate execution info + execution_info = { + "status": "completed", + "environment": "CTF", + "tool_time": execution_time + } + + # Complete the streaming with final output + finish_tool_streaming(tool_name, tool_args, output, call_id, execution_info) + + # Switch back to idle mode after CTF command completes + stop_active_timer() + start_idle_timer() + return output + + except Exception as e: + # Handle errors in CTF execution + error_msg = f"Error executing CTF command: {str(e)}" + execution_info = { + "status": "error", + "environment": "CTF", + "error": str(e) + } + + # Complete the streaming with error output + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after error + stop_active_timer() + start_idle_timer() + return error_msg + else: + # Standard non-streaming CTF execution + result = _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir - # _run_ctf handles workspace internally using _get_workspace_dir() default - result = _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir - - # Switch back to idle mode after CTF command completes - stop_active_timer() - start_idle_timer() - return result + # Switch back to idle mode after CTF command completes + stop_active_timer() + start_idle_timer() + return result # --- SSH Execution --- if is_ssh_env: - # Async for SSH would require session management via SSH client features - if async_mode: - # Switch back to idle mode before returning message + # If streaming is enabled, show streaming UI for SSH too + if stream: + # Import the streaming utilities from util + from cai.util import start_tool_streaming, update_tool_streaming, finish_tool_streaming + + # Add SSH connection info for display + ssh_user = os.environ.get('SSH_USER', 'user') + ssh_host = os.environ.get('SSH_HOST', 'host') + ssh_connection = f"{ssh_user}@{ssh_host}" + + # Create args dictionary with standardized format + tool_args = { + "command": cmd, + "args": args if args.strip() else "", + "full_command": command, + "ssh_host": ssh_connection, + "environment": "SSH" + } + + # Initialize streaming session with a consistent call_id format + call_id = start_tool_streaming(tool_name, tool_args, call_id) + + # Update with "executing" status + update_tool_streaming( + tool_name, + tool_args, + f"Executing on {ssh_connection}: {command}\n\nWaiting for response...", + call_id + ) + + try: + # Construct SSH command for execution + ssh_pass = os.environ.get('SSH_PASS') + if ssh_pass: + ssh_cmd_list = ["sshpass", "-p", ssh_pass, "ssh", ssh_connection] + else: + ssh_cmd_list = ["ssh", ssh_connection] + ssh_cmd_list.append(command) + + # Execute the command and get the output + start_time = time.time() + result = subprocess.run( + ssh_cmd_list, + capture_output=True, + text=True, + check=False, + timeout=timeout + ) + execution_time = time.time() - start_time + + # Get command output + output = result.stdout if result.stdout else result.stderr + + # Add SSH connection info to the output for clarity + result_with_info = f"Command executed on {ssh_connection}:\n\n{output}" + + # Determine status based on return code + status = "completed" if result.returncode == 0 else "error" + + # Calculate execution info + execution_info = { + "status": status, + "environment": "SSH", + "host": ssh_connection, + "return_code": result.returncode, + "tool_time": execution_time + } + + # Complete the streaming with final output + finish_tool_streaming(tool_name, tool_args, result_with_info, call_id, execution_info) + + # Switch back to idle mode after SSH command completes + stop_active_timer() + start_idle_timer() + return output.strip() + + except subprocess.TimeoutExpired as e: + # Handle timeout errors + error_output = e.stdout if e.stdout else str(e) + error_msg = f"Command timed out after {timeout} seconds\n{error_output}" + + execution_info = { + "status": "timeout", + "environment": "SSH", + "host": ssh_connection, + "error": str(e) + } + + # Complete the streaming with timeout error + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after timeout + stop_active_timer() + start_idle_timer() + return error_msg + + except Exception as e: + # Handle other errors + error_msg = f"Error executing SSH command: {str(e)}" + + execution_info = { + "status": "error", + "environment": "SSH", + "host": ssh_connection, + "error": str(e) + } + + # Complete the streaming with error + finish_tool_streaming(tool_name, tool_args, error_msg, call_id, execution_info) + + # Switch back to idle mode after error + stop_active_timer() + start_idle_timer() + return error_msg + else: + # Standard non-streaming SSH execution + result = _run_ssh(command, stdout, timeout) # Workspace dir less relevant here + + # Switch back to idle mode after SSH command completes stop_active_timer() start_idle_timer() - return "Async mode not fully supported for SSH environment via this function yet." - - # Handling streaming for SSH - not fully implemented yet - if stream: - from cai.util import cli_print_tool_output - if call_id and tool_name: - tool_args = {"command": command, "ssh": True} - cli_print_tool_output( - tool_name, - tool_args, - "Streaming not yet supported for SSH execution. Running normally...", - call_id=call_id - ) - - # _run_ssh handles command execution, workspace is relative to remote home - result = _run_ssh(command, stdout, timeout) # Workspace dir less relevant here - - # Switch back to idle mode after SSH command completes - stop_active_timer() - start_idle_timer() - return result + return result # --- Local Execution (Default Fallback) --- # Let _run_local handle determining the host workspace @@ -989,8 +1185,16 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg start_idle_timer() return f"Started async session {new_session_id} locally. Use this ID to interact." - # Handle Synchronous Execution Locally using _run_local default with streaming support - result = _run_local(command, stdout, timeout, stream, call_id, tool_name, None) + # Handle Synchronous Execution Locally + # Pass stream=True if we're streaming to use streaming functionality + result = _run_local( + command, + stdout, + timeout, + stream=stream, + call_id=call_id, + tool_name=tool_name + ) # Switch back to idle mode after local command completes stop_active_timer() diff --git a/src/cai/util.py b/src/cai/util.py index cc6fc0c7..ed2d7a22 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -5,6 +5,7 @@ import os import sys import importlib.resources import pathlib +import json from rich.console import Console from rich.tree import Tree from mako.template import Template # pylint: disable=import-error @@ -29,6 +30,9 @@ _idle_timer_start = None _idle_time_total = 0.0 _timing_lock = threading.Lock() +# Set up a global tracker for live streaming panels +_LIVE_STREAMING_PANELS = {} + def start_active_timer(): """ Start measuring active time (when LLM is processing or tool is executing). @@ -539,6 +543,8 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 must have content. 3. If a tool call id appears alone (without a pair), it is removed. 4. There cannot be empty messages. + 5. Each tool_use block (assistant with tool_calls) must be followed by + a tool_result block (tool message with matching tool_call_id). Args: messages (List[dict]): List of message dictionaries containing @@ -557,7 +563,13 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 for i, msg in enumerate(messages): # Skip empty messages (considered empty if 'content' is None or only whitespace) - if msg.get("role") in ["user", "system"] and (msg.get("content") is None or not msg.get("content").strip()): + if msg.get("role") in ["user", "system"] and (msg.get("content") is None or not str(msg.get("content", "")).strip()): + # Special case: if it's a system message, set content to empty string instead of skipping + if msg.get("role") == "system": + # Replace None with empty string + msg["content"] = "" + sanitized_messages.append(msg) + # Skip empty user messages entirely continue # Add valid messages to our sanitized list first @@ -596,7 +608,7 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 tool_call_map[tool_id] = {"assistant_idx": len(sanitized_messages) - 2, "tool_idx": len(sanitized_messages) - 1} # Final validation - ensure all tool calls have responses - for tool_id, indices in tool_call_map.items(): + for tool_id, indices in list(tool_call_map.items()): if indices["tool_idx"] is None: # Tool call without a response - create a synthetic tool message assistant_idx = indices["assistant_idx"] @@ -616,8 +628,59 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 "tool_call_id": tool_id, "content": f"Auto-generated response for {tool_name}" } - # Insert right after the assistant message - sanitized_messages.insert(assistant_idx + 1, tool_msg) + + # Insert immediately after the assistant message + if assistant_idx + 1 < len(sanitized_messages): + # Insert at the position after assistant + sanitized_messages.insert(assistant_idx + 1, tool_msg) + else: + # Just append if we're at the end + sanitized_messages.append(tool_msg) + + # Update the map to note that this tool call now has a response + tool_call_map[tool_id]["tool_idx"] = assistant_idx + 1 + + # Ensure messages have non-null content (required by some providers) + for msg in sanitized_messages: + if msg.get("role") != "tool" and msg.get("content") is None and not msg.get("tool_calls"): + msg["content"] = "" + + # For tool messages, ensure content is never null + if msg.get("role") == "tool" and msg.get("content") is None: + msg["content"] = f"Tool response for {msg.get('tool_call_id', 'unknown')}" + + # Special case for Claude: ensure strict alternating pattern between assistant tool_calls and tool results + # If multiple consecutive assistant messages with tool_calls exist, interleave them with tool responses + i = 0 + while i < len(sanitized_messages) - 1: + current_msg = sanitized_messages[i] + next_msg = sanitized_messages[i + 1] + + # When current message is assistant with tool_calls and next message is NOT a tool response + if (current_msg.get("role") == "assistant" and + current_msg.get("tool_calls") and + (next_msg.get("role") != "tool" or not next_msg.get("tool_call_id"))): + + # Get the first tool call ID + tool_id = current_msg["tool_calls"][0].get("id", "unknown") + tool_name = "unknown_function" + if current_msg["tool_calls"][0].get("function"): + tool_name = current_msg["tool_calls"][0]["function"].get("name", "unknown_function") + + # Create a tool result message + tool_msg = { + "role": "tool", + "tool_call_id": tool_id, + "content": f"Auto-generated response for {tool_name}" + } + + # Insert the tool message after the current assistant message + sanitized_messages.insert(i + 1, tool_msg) + + # Skip over the newly inserted message + i += 2 + else: + i += 1 return sanitized_messages @@ -665,22 +728,23 @@ def get_model_name(model): return model # If not a string, use environment variable return os.environ.get('CAI_MODEL', 'qwen2.5:72b') - # Helper function to format time in a human-readable way + +# Helper function to format time in a human-readable way def format_time(seconds): - if seconds is None: - return "N/A" - - if seconds < 60: - return f"{seconds:.1f}s" - elif seconds < 3600: - minutes = int(seconds / 60) - seconds_remainder = seconds % 60 - return f"{minutes}m {seconds_remainder:.1f}s" - else: - hours = int(seconds / 3600) - minutes = int((seconds % 3600) / 60) - return f"{hours}h {minutes}m" - + if seconds is None: + return "N/A" + + if seconds < 60: + return f"{seconds:.1f}s" + elif seconds < 3600: + minutes = int(seconds / 60) + seconds_remainder = seconds % 60 + return f"{minutes}m {seconds_remainder:.1f}s" + else: + hours = int(seconds / 3600) + minutes = int((seconds % 3600) / 60) + return f"{hours}h {minutes}m" + def get_model_pricing(model_name): """ Get pricing information for a model, using the CostTracker's implementation. @@ -946,7 +1010,8 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli total_reasoning_tokens=None, interaction_cost=None, total_cost=None, - tool_output=None): # New parameter for tool output + tool_output=None, # New parameter for tool output + suppress_empty=False): # New parameter to suppress empty panels """Print agent messages/thoughts with enhanced visual formatting.""" # Debug prints to trace the function calls if debug: @@ -978,6 +1043,17 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli else: parsed_message = parse_message_content(message) tool_panels = [] + + # Skip empty panels - THIS IS THE KEY CHANGE + # If suppress_empty is True and there's no parsed message and no tool panels, + # don't create an empty panel to avoid cluttering during streaming + if suppress_empty and not parsed_message and not tool_panels: + return + + # Also skip if the only message is "null" or empty + if parsed_message == "null" or parsed_message == "": + if suppress_empty and not tool_panels: + return # Special handling for Reasoner Agent if agent_name == "Reasoner Agent": @@ -1046,8 +1122,7 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli # If there are tool panels, print them after the main message panel # But only in non-streaming mode to avoid duplicates - is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' - if tool_panels and not is_streaming_enabled: + if tool_panels: for tool_panel in tool_panels: console.print(tool_panel) @@ -1063,6 +1138,15 @@ def create_agent_streaming_context(agent_name, counter, model): Returns: A dictionary with the streaming context """ + # Add a static variable to track active streaming contexts and prevent duplicates + if not hasattr(create_agent_streaming_context, "_active_streaming"): + create_agent_streaming_context._active_streaming = {} + + # If there's already an active streaming context with the same counter, return it + context_key = f"{agent_name}_{counter}" + if context_key in create_agent_streaming_context._active_streaming: + return create_agent_streaming_context._active_streaming[context_key] + try: from rich.live import Live import shutil @@ -1107,9 +1191,9 @@ def create_agent_streaming_context(agent_name, counter, model): ) # Create Live display object but don't start it until we have content - live = Live(panel, refresh_per_second=20, console=console, auto_refresh=False) + live = Live(panel, refresh_per_second=10, console=console, auto_refresh=True, vertical_overflow="visible") - return { + context = { "live": live, "panel": panel, "header": header, @@ -1122,6 +1206,11 @@ def create_agent_streaming_context(agent_name, counter, model): "is_started": False, # Track if we've started the display "error": None, # Track any errors } + + # Store the context for potential reuse + create_agent_streaming_context._active_streaming[context_key] = context + + return context except Exception as e: # If rich display fails, return None and log the error import sys @@ -1191,6 +1280,13 @@ def finish_agent_streaming(context, final_stats=None): """ if not context: return False + + # Clean up tracking of this context + if hasattr(create_agent_streaming_context, "_active_streaming"): + for key, value in list(create_agent_streaming_context._active_streaming.items()): + if value is context: + del create_agent_streaming_context._active_streaming[key] + break try: # Check if there's actual content to display - don't show empty panels @@ -1293,7 +1389,7 @@ def finish_agent_streaming(context, final_stats=None): return False -def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execution_info=None, token_info=None): +def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execution_info=None, token_info=None, streaming=False): """ Print a tool call output to the command line. Similar to cli_print_tool_call but for the output of the tool. @@ -1309,44 +1405,177 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut - total_input_tokens, total_output_tokens, total_reasoning_tokens - model: model name string - interaction_cost, total_cost: optional cost values + streaming: Flag indicating if this is part of a streaming output """ - # If it's an empty output, don't print anything - if not output and not call_id: + # If it's an empty output, don't print anything except for streaming sessions + if not output and not call_id and not streaming: return + # Set up global tracker for streaming sessions + if not hasattr(cli_print_tool_output, '_streaming_sessions'): + cli_print_tool_output._streaming_sessions = {} - # CRITICAL CHECK: When in streaming mode (CAI_STREAM=true), ONLY show output panels - # for streaming updates (those with call_id). This prevents duplicate output panels. - is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' - if is_streaming_enabled and not call_id: - # Skip all non-streaming tool output in streaming mode - return - - # Track seen call IDs to prevent duplicate panels + # Track seen call IDs to prevent duplicate panels for non-streaming outputs if not hasattr(cli_print_tool_output, '_seen_calls'): cli_print_tool_output._seen_calls = {} - - # For streaming updates, only show updates for the same call_id - # but allow the first appearance of each call_id - if call_id: - call_key = f"{call_id}:{output[:20]}" # Use first 20 chars as fingerprint with call_id - # Skip if we've seen this exact output for this call_id before - if call_key in cli_print_tool_output._seen_calls: + # Track all displayed commands to prevent duplicates + if not hasattr(cli_print_tool_output, '_displayed_commands'): + cli_print_tool_output._displayed_commands = set() + + # --- Consistent Command Key Generation --- + effective_command_args_str = "" + if isinstance(args, dict): + # If args is a dictionary, extract the 'args' field. + effective_command_args_str = args.get("args", "") + elif isinstance(args, str): + # If args is a string, it might be a JSON representation or a plain string. + try: + parsed_json_args = json.loads(args) + if isinstance(parsed_json_args, dict): + # Parsed as JSON dict, get the 'args' field. + effective_command_args_str = parsed_json_args.get("args", "") + else: + # Parsed as JSON, but not a dict (e.g., a JSON string literal). + effective_command_args_str = parsed_json_args if isinstance(parsed_json_args, str) else args + except json.JSONDecodeError: + # Not a JSON string, treat 'args' as a plain string. + effective_command_args_str = args + + command_key = f"{tool_name}:{effective_command_args_str}" + # --- End of Command Key Generation --- + + # Check for duplicate display conditions + if streaming: + # For streaming updates, track and update the single streaming session + if call_id: + # If this is a new streaming session, record it + if call_id not in cli_print_tool_output._streaming_sessions: + cli_print_tool_output._streaming_sessions[call_id] = { + 'tool_name': tool_name, + 'args': args, # Store original args for display formatting + 'buffer': output if output else "", + 'start_time': time.time(), + 'last_update': time.time(), + 'command_key': command_key, # Store the generated key + 'is_complete': False + } + # Add the command key to displayed commands + if command_key not in cli_print_tool_output._displayed_commands: + cli_print_tool_output._displayed_commands.add(command_key) + else: + # Update the existing session + session = cli_print_tool_output._streaming_sessions[call_id] + # Always replace buffer with latest output for consistency + session['buffer'] = output + session['last_update'] = time.time() + if execution_info and execution_info.get('is_final', False): + session['is_complete'] = True + + # For streaming outputs, we'll use Rich Live panel if available + try: + from rich.console import Console + from rich.live import Live + from rich.panel import Panel + from rich.text import Text + from rich.box import ROUNDED + + # Access the global live panel dictionary + global _LIVE_STREAMING_PANELS + + # Create the header, content, and panel + # Pass the original 'args' (dict or string) to _create_tool_panel_content for formatting + current_args_for_display = cli_print_tool_output._streaming_sessions[call_id]['args'] + header, content = _create_tool_panel_content( + tool_name, + current_args_for_display, + cli_print_tool_output._streaming_sessions[call_id]['buffer'], + execution_info, + token_info + ) + + # Determine panel style based on status + status = "running" + if execution_info: + status = execution_info.get('status', 'running') + + border_style = "yellow" # Default for running + if status == "completed": + border_style = "green" + elif status in ["error", "timeout"]: + border_style = "red" + + # Create panel title based on status + if status == "running": + title = "[bold yellow]Tool Execution [Running][/bold yellow]" + elif status == "completed": + title = "[bold green]Tool Execution [Completed][/bold green]" + elif status == "error": + title = "[bold red]Tool Execution [Error][/bold red]" + elif status == "timeout": + title = "[bold red]Tool Execution [Timeout][/bold red]" + else: + title = "[bold blue]Tool Execution[/bold blue]" + + # Create the panel + panel = Panel( + content, + title=title, + border_style=border_style, + padding=(1, 2), + box=ROUNDED, + title_align="left" + ) + + # If we already have a live panel for this call_id, update it + if call_id in _LIVE_STREAMING_PANELS: + live = _LIVE_STREAMING_PANELS[call_id] + live.update(panel) + + # If this is the final update, stop the live panel after a short delay + if execution_info and execution_info.get('is_final', False): + # Give a moment for the final panel to be seen + time.sleep(0.2) + live.stop() + # Remove from the active panel dictionary + del _LIVE_STREAMING_PANELS[call_id] + else: + # Create a new live panel + console = Console(theme=theme) + live = Live(panel, console=console, refresh_per_second=4, auto_refresh=True) + # Start and store the live panel + live.start() + _LIVE_STREAMING_PANELS[call_id] = live + + # Return early for streaming updates + return + + except ImportError: + # Fall back to simple updates without Rich + pass + else: + # For non-streaming outputs, check if we've already seen this command + if command_key in cli_print_tool_output._displayed_commands: + # Command has already been displayed (likely through streaming), skip duplicate display return - # Mark as seen - cli_print_tool_output._seen_calls[call_key] = True + # Add to displayed commands since we're going to show it + # This handles the case where a command is non-streaming from the start + cli_print_tool_output._displayed_commands.add(command_key) + + # For non-streaming updates with call_id, check if already seen + # This _seen_calls logic is an additional layer for non-streaming calls that might have call_ids + # but might be distinct from the primary _displayed_commands check based on command_key. + if call_id and not streaming: + # Create a more specific key for _seen_calls if needed, possibly including output fingerprint + seen_call_key = f"{call_id}:{command_key}:{output[:20]}" - # Limit cache size to prevent memory growth - if len(cli_print_tool_output._seen_calls) > 1000: - # Keep only the most recent 500 entries - cli_print_tool_output._seen_calls = { - k: cli_print_tool_output._seen_calls[k] - for k in list(cli_print_tool_output._seen_calls.keys())[-500:] - } - - # Try to use Rich for better formatting if available + if seen_call_key in cli_print_tool_output._seen_calls: + return + + cli_print_tool_output._seen_calls[seen_call_key] = True + + # Standard tool output display for non-streaming or when rich is not available try: from rich.console import Console from rich.panel import Panel @@ -1357,190 +1586,20 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Create a console for output console = Console(theme=theme) - # Format arguments for display - # Parse JSON string if args is a string - if isinstance(args, str) and args.strip().startswith('{'): - try: - import json - args = json.loads(args) - except: - # Keep as is if not valid JSON - pass + # Get the panel content + header, content = _create_tool_panel_content(tool_name, args, output, execution_info, token_info) - # Format arguments as a clean string - if isinstance(args, dict): - # Only include non-empty values and exclude async_mode=false - arg_parts = [] - for key, value in args.items(): - # Skip empty values - if value == "" or value == {} or value is None: - continue - # Skip async_mode=false (default) - if key == "async_mode" and value is False: - continue - # Format the value - if isinstance(value, str): - arg_parts.append(f"{key}={value}") - else: - arg_parts.append(f"{key}={value}") - args_str = ", ".join(arg_parts) - else: - args_str = str(args) + # Determine border style + border_style = "blue" # Default for non-streaming - # Get session timing information - try: - from cai.cli import START_TIME - total_time = time.time() - START_TIME if START_TIME else None - except ImportError: - total_time = None - - - # Extract execution timing info - - tool_time = None - status = None - if execution_info: - # Prefer 'tool_time' if present, else fallback to 'time_taken' - - tool_time = execution_info.get('tool_time') - status = execution_info.get('status', 'completed') - - # Create header for all panel displays (both streaming and non-streaming) - header = Text() - header.append(tool_name, style="#00BCD4") - header.append("(", style="yellow") - header.append(args_str, style="yellow") - header.append(")", style="yellow") - - # Add timing information directly in the header - timing_info = [] - if total_time: - timing_info.append(f"Total: {format_time(total_time)}") - if tool_time: - timing_info.append(f"Tool: {format_time(tool_time)}") - if timing_info: - header.append(f" [{' | '.join(timing_info)}]", style="cyan") - - # Add completion status if available - REMOVED, just showing timing now - # if status: - # if status == 'completed': - # header.append(f" [Completed]", style="green") - # elif status == 'running': - # header.append(f" [Running]", style="yellow") - # elif status == 'error': - # header.append(f" [Error]", style="red") - # elif status == 'timeout': - # header.append(f" [Timeout]", style="red") - # else: - # header.append(f" [{status.title()}]", style="dim") - - # For streaming mode with call_id, use Rich Live display - if call_id: - # Create token information if available - token_content = None - if token_info: - model = token_info.get('model', '') - interaction_input_tokens = token_info.get('interaction_input_tokens', 0) - interaction_output_tokens = token_info.get('interaction_output_tokens', 0) - interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) - total_input_tokens = token_info.get('total_input_tokens', 0) - total_output_tokens = token_info.get('total_output_tokens', 0) - total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) - - if (interaction_input_tokens > 0 or total_input_tokens > 0): - token_text = _create_token_display( - interaction_input_tokens, - interaction_output_tokens, - interaction_reasoning_tokens, - total_input_tokens, - total_output_tokens, - total_reasoning_tokens, - model, - token_info.get('interaction_cost'), - token_info.get('total_cost') - ) - token_content = Text("\n\n") - token_content.append(token_text) - - # Create content text with the output - content = Text(output) - - # Create the panel for display, including token info if available - panel_content = [header, Text("\n\n"), content] - if token_content: - panel_content.append(token_content) - - # Create title - simple title with no timing info - title = "[bold blue]Tool Output[/bold blue]" - - panel = Panel( - Text.assemble(*panel_content), - title=title, - border_style="blue", - padding=(1, 2), - box=ROUNDED, - title_align="left" - ) - - # Display using Rich - console.print(panel) - return - - # For non-streaming output, also use a blue panel with the same format - # Create token information if available - token_text = None - if token_info: - model = token_info.get('model', '') - interaction_input_tokens = token_info.get('interaction_input_tokens', 0) - interaction_output_tokens = token_info.get('interaction_output_tokens', 0) - interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) - total_input_tokens = token_info.get('total_input_tokens', 0) - total_output_tokens = token_info.get('total_output_tokens', 0) - total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) - interaction_cost = token_info.get('interaction_cost') - total_cost = token_info.get('total_cost') - - # Generate token display with CostTracker - if (interaction_input_tokens > 0 or total_input_tokens > 0): - token_text = _create_token_display( - interaction_input_tokens, - interaction_output_tokens, - interaction_reasoning_tokens, - total_input_tokens, - total_output_tokens, - total_reasoning_tokens, - model, - interaction_cost, - total_cost - ) - - # Now create the panel content, starting with the header - panel_content = [header, Text("\n")] - - # Add token display if available - if token_text: - panel_content.append(token_text) - panel_content.append(Text("\n")) # Add spacing after token display - - # Add the output - if output: - output_text = Text(output) - panel_content.append(output_text) - - # If no content was added but we have output, add it directly - if len(panel_content) == 2 and output: # Only header and newline - panel_content.append(Text(output)) - - # Create title - simple title with no timing info - title = "[bold blue]Tool Output[/bold blue]" - - # Create the final panel - always blue now + # Create the panel panel = Panel( - Group(*panel_content), - title=title, - border_style="blue", + content, + title="[bold blue]Tool Output[/bold blue]", + border_style=border_style, padding=(1, 2), - box=ROUNDED + box=ROUNDED, + title_align="left" ) # Display the panel @@ -1548,113 +1607,378 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut except ImportError: # Fall back to simple formatting if Rich is not available - # Format arguments in the cleaner format - # Parse JSON string if args is a string - if isinstance(args, str) and args.strip().startswith('{'): - try: - import json - args = json.loads(args) - except: - # Keep as is if not valid JSON - pass - - # Format arguments as a clean string - if isinstance(args, dict): - # Only include non-empty values and exclude async_mode=false - arg_parts = [] - for key, value in args.items(): - # Skip empty values - if value == "" or value == {} or value is None: - continue - # Skip async_mode=false (default) - if key == "async_mode" and value is False: - continue - # Format the value - if isinstance(value, str): - arg_parts.append(f"{key}={value}") - else: - arg_parts.append(f"{key}={value}") - args_str = ", ".join(arg_parts) + _print_simple_tool_output(tool_name, args, output, execution_info, token_info) + +# Helper function to create tool panel content +def _create_tool_panel_content(tool_name, args, output, execution_info=None, token_info=None): + """Create the header and content for a tool output panel.""" + from rich.text import Text + + # Format arguments for display + args_str = _format_tool_args(args) + + # Get timing information + timing_info, tool_time = _get_timing_info(execution_info) + + # Create header + header = Text() + header.append(tool_name, style="#00BCD4") + header.append("(", style="yellow") + header.append(args_str, style="yellow") + header.append(")", style="yellow") + + # Add timing information + if timing_info: + header.append(f" [{' | '.join(timing_info)}]", style="cyan") + + # Add environment info if available + if execution_info and execution_info.get('environment'): + env = execution_info.get('environment') + host = execution_info.get('host', '') + if host: + header.append(f" [{env}:{host}]", style="magenta") else: - args_str = str(args) - # Get session timing information - try: - from cai.cli import START_TIME - total_time = time.time() - START_TIME if START_TIME else None - except ImportError: - total_time = None - - - # For non-streaming output, use the original formatting - tool_call = f"{tool_name}({args_str})" - - # Get tool execution time if available - tool_time_str = "" - execution_status = "" - if execution_info: - time_taken = execution_info.get('time_taken', 0) - status = execution_info.get('status', 'completed') - - # Add execution info to the tool call display - if time_taken: - tool_time_str = f"Tool: {format_time(time_taken)}" - execution_status = f" [{status} in {time_taken:.2f}s]" + header.append(f" [{env}]", style="magenta") + + # Add status information if available + if execution_info: + status = execution_info.get('status', None) + if status == "completed": + header.append(" [Completed]", style="green") + elif status == "running": + header.append(" [Running]", style="yellow") + elif status == "error": + header.append(" [Error]", style="red") + elif status == "timeout": + header.append(" [Timeout]", style="red") + + # Create token information if available + token_content = _create_token_info_display(token_info) + + # Assemble the full content + content = Text() + content.append(header) + content.append("\n\n") + content.append(output) + + # Add token info if available + if token_content: + content.append("\n\n") + content.append(token_content) + + return header, content + +# Helper function to format tool arguments +def _format_tool_args(args): + """Format tool arguments as a clean string.""" + # If args is already a string, it might be pre-formatted or a simple arg string + if isinstance(args, str): + # If it looks like a JSON dict string, try to parse and format nicely + if args.strip().startswith('{') and args.strip().endswith('}'): + try: + + parsed_dict = json.loads(args) + # Recursively call with the parsed dict for consistent formatting + return _format_tool_args(parsed_dict) + except json.JSONDecodeError: + # Not valid JSON, or not a dict; return as is + return args + else: + # Simple string arg, return as is + return args + + # Format arguments from a dictionary + if isinstance(args, dict): + # Only include non-empty values and exclude special flags + arg_parts = [] + for key, value in args.items(): + # Skip empty values + if value == "" or value == {} or value is None: + continue + # Skip special flags + if key in ["async_mode", "streaming"] and not value: + continue + # Format the value + if isinstance(value, str): + arg_parts.append(f"{key}={value}") else: - execution_status = f" [{status}]" + arg_parts.append(f"{key}={value}") + return ", ".join(arg_parts) + else: + return str(args) + +# Helper function to get timing information +def _get_timing_info(execution_info=None): + """Get timing information for display.""" + import time + + # Get session timing information + try: + from cai.cli import START_TIME + total_time = time.time() - START_TIME if START_TIME else None + except ImportError: + total_time = None + + # Extract execution timing info + tool_time = None + if execution_info: + tool_time = execution_info.get('tool_time') + + # Format timing info for display + timing_info = [] + if total_time: + timing_info.append(f"Total: {format_time(total_time)}") + if tool_time: + timing_info.append(f"Tool: {format_time(tool_time)}") + + return timing_info, tool_time + +# Helper function to create token info display +def _create_token_info_display(token_info=None): + """Create token information display text.""" + if not token_info: + return None + + from rich.text import Text + + model = token_info.get('model', '') + interaction_input_tokens = token_info.get('interaction_input_tokens', 0) + interaction_output_tokens = token_info.get('interaction_output_tokens', 0) + interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) + total_input_tokens = token_info.get('total_input_tokens', 0) + total_output_tokens = token_info.get('total_output_tokens', 0) + total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) + + # Only continue if we have actual token information + if not (interaction_input_tokens > 0 or total_input_tokens > 0): + return None + + # Create token display + return _create_token_display( + interaction_input_tokens, + interaction_output_tokens, + interaction_reasoning_tokens, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + model, + token_info.get('interaction_cost'), + token_info.get('total_cost') + ) + +# Helper function for simple tool output without Rich +def _print_simple_tool_output(tool_name, args, output, execution_info=None, token_info=None): + """Print tool output without Rich formatting.""" + # Format arguments + args_str = _format_tool_args(args) + + # Get tool execution time if available + tool_time_str = "" + execution_status = "" + if execution_info: + time_taken = execution_info.get('time_taken', 0) or execution_info.get('tool_time', 0) + status = execution_info.get('status', 'completed') - # Create timing display string - timing_info = [] - if total_time: - timing_info.append(f"Total: {format_time(total_time)}") - if tool_time: - timing_info.append(f"Tool: {format_time(tool_time)}") - if timing_info: - header.append(f" [{' | '.join(timing_info)}]", style="cyan") + # Add execution info to the tool call display + if time_taken: + tool_time_str = f"Tool: {format_time(time_taken)}" + execution_status = f" [{status} in {time_taken:.2f}s]" + else: + execution_status = f" [{status}]" + + # Create timing display string + timing_info, _ = _get_timing_info(execution_info) + timing_display = f" [{' | '.join(timing_info)}]" if timing_info else "" + + # Show tool name, args, execution status and timing display + tool_call = f"{tool_name}({args_str})" + print(color(f"Tool Output: {tool_call}{timing_display}{execution_status}", fg="blue")) + + # If we have token info, display it + if token_info: + model = token_info.get('model', '') + interaction_input_tokens = token_info.get('interaction_input_tokens', 0) + interaction_output_tokens = token_info.get('interaction_output_tokens', 0) + interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) + total_input_tokens = token_info.get('total_input_tokens', 0) + total_output_tokens = token_info.get('total_output_tokens', 0) + total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) + + # If we have complete token information, display it + if (interaction_input_tokens > 0 or total_input_tokens > 0): + # Manually create formatted output similar to _create_token_display + print(color(f" Current: I:{interaction_input_tokens} O:{interaction_output_tokens} R:{interaction_reasoning_tokens}", fg="cyan")) - timing_display = f" [{' | '.join(timing_info)}]" if timing_info else "" - - # Show tool name, args, execution status and timing display - print(color(f"Tool Output: {tool_call}{timing_display}{execution_status}", fg="blue")) - - # If we have token info, display it using the consistent format from _create_token_display - if token_info: - model = token_info.get('model', '') - interaction_input_tokens = token_info.get('interaction_input_tokens', 0) - interaction_output_tokens = token_info.get('interaction_output_tokens', 0) - interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) - total_input_tokens = token_info.get('total_input_tokens', 0) - total_output_tokens = token_info.get('total_output_tokens', 0) - total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) - interaction_cost = token_info.get('interaction_cost') - total_cost = token_info.get('total_cost') + # Calculate or use provided costs + current_cost = COST_TRACKER.process_interaction_cost( + model, + interaction_input_tokens, + interaction_output_tokens, + interaction_reasoning_tokens, + token_info.get('interaction_cost') + ) + total_cost_value = COST_TRACKER.process_total_cost( + model, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + token_info.get('total_cost') + ) + print(color(f" Cost: Current ${current_cost:.4f} | Total ${total_cost_value:.4f} | Session ${COST_TRACKER.session_total_cost:.4f}", fg="cyan")) - # If we have complete token information, display it - if (interaction_input_tokens > 0 or total_input_tokens > 0): - # Manually create formatted output similar to _create_token_display - print(color(f" Current: I:{interaction_input_tokens} O:{interaction_output_tokens} R:{interaction_reasoning_tokens}", fg="cyan")) - - # Calculate or use provided costs - current_cost = COST_TRACKER.process_interaction_cost( - model, - interaction_input_tokens, - interaction_output_tokens, - interaction_reasoning_tokens, - interaction_cost - ) - total_cost_value = COST_TRACKER.process_total_cost( - model, - total_input_tokens, - total_output_tokens, - total_reasoning_tokens, - total_cost - ) - print(color(f" Cost: Current ${current_cost:.4f} | Total ${total_cost_value:.4f} | Session ${COST_TRACKER.session_total_cost:.4f}", fg="cyan")) - - # Show context usage - context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100 - indicator = "🟩" if context_pct < 50 else "🟨" if context_pct < 80 else "πŸŸ₯" - print(color(f" Context: {context_pct:.1f}% {indicator}", fg="cyan")) + # Show context usage + context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100 + indicator = "🟩" if context_pct < 50 else "🟨" if context_pct < 80 else "πŸŸ₯" + print(color(f" Context: {context_pct:.1f}% {indicator}", fg="cyan")) + + # Print the actual output + print(output) + print() + +# Add a new function to start a streaming tool execution +def start_tool_streaming(tool_name, args, call_id=None): + """ + Start a streaming tool execution session. + This allows for progressive updates during tool execution. + + Args: + tool_name: Name of the tool being executed + args: Arguments to the tool (dictionary or string) + call_id: Optional call ID for this execution. If not provided, one will be generated. - # Print the actual output - print(output) - print() \ No newline at end of file + Returns: + call_id: The call ID for this streaming session (can be used for updates) + """ + import time + import uuid + + # Generate a command key to check for duplicates - match format used in cli_print_tool_output + if isinstance(args, dict): + cmd = args.get("command", "") + cmd_args = args.get("args", "") + command_key = f"{tool_name}:{cmd_args}" + else: + command_key = f"{tool_name}:{args}" + + # Check if we've already seen this exact command recently + if not hasattr(start_tool_streaming, '_recent_commands'): + start_tool_streaming._recent_commands = {} + + # If we have an existing active streaming session for this command, reuse its call_id + # This prevents duplicate panels when the same command runs multiple times + for existing_call_id, info in list(start_tool_streaming._recent_commands.items()): + # Only consider recent commands (last 10 seconds) + timestamp = info.get('timestamp', 0) + if time.time() - timestamp < 10.0: + existing_command_key = info.get('command_key', '') + # Get the existing session info if available + if (hasattr(cli_print_tool_output, '_streaming_sessions') and + existing_call_id in cli_print_tool_output._streaming_sessions): + session = cli_print_tool_output._streaming_sessions[existing_call_id] + # If this is the same command and not complete, reuse the call_id + if existing_command_key == command_key and not session.get('is_complete', False): + return existing_call_id + + # Generate a call_id if not provided + if not call_id: + cmd_part = "" + if isinstance(args, dict) and "command" in args: + cmd_part = f"{args['command']}_" + call_id = f"cmd_{cmd_part}{str(uuid.uuid4())[:8]}" + + # Track this call_id with command key for better duplicate detection + start_tool_streaming._recent_commands[call_id] = { + 'timestamp': time.time(), + 'command_key': command_key + } + + # Cleanup old entries to prevent memory growth + current_time = time.time() + start_tool_streaming._recent_commands = { + k: v for k, v in start_tool_streaming._recent_commands.items() + if current_time - v.get('timestamp', 0) < 30 # Keep entries from last 30 seconds + } + + # Show initial message with "Starting..." output + cli_print_tool_output( + tool_name=tool_name, + args=args, + output="Starting tool execution...", + call_id=call_id, + execution_info={"status": "running", "start_time": time.time()}, + streaming=True + ) + + return call_id + +# Add a function to update a streaming tool execution +def update_tool_streaming(tool_name, args, output, call_id): + """ + Update a streaming tool execution with new output. + + Args: + tool_name: Name of the tool being executed + args: Arguments to the tool (dictionary or string) + output: New output to display + call_id: The call ID for this streaming session + + Returns: + None + """ + # Update the streaming output + cli_print_tool_output( + tool_name=tool_name, + args=args, + output=output, + call_id=call_id, + execution_info={"status": "running", "replace_buffer": True}, + streaming=True + ) + +# Add a function to complete a streaming tool execution +def finish_tool_streaming(tool_name, args, output, call_id, execution_info=None, token_info=None): + """ + Complete a streaming tool execution. + + Args: + tool_name: Name of the tool being executed + args: Arguments to the tool (dictionary or string) + output: Final output to display + call_id: The call ID for this streaming session + execution_info: Optional execution information + token_info: Optional token information + + Returns: + None + """ + import time + + # Prepare execution info with completion status + if execution_info is None: + execution_info = {} + + # Add completion markers + execution_info["status"] = execution_info.get("status", "completed") + execution_info["is_final"] = True + execution_info["replace_buffer"] = True + + # Calculate execution time if start_time is in the streaming session + if hasattr(cli_print_tool_output, '_streaming_sessions') and call_id in cli_print_tool_output._streaming_sessions: + session = cli_print_tool_output._streaming_sessions[call_id] + if 'start_time' in session and 'tool_time' not in execution_info: + execution_info["tool_time"] = time.time() - session['start_time'] + + # Show the final output + cli_print_tool_output( + tool_name=tool_name, + args=args, + output=output, + call_id=call_id, + execution_info=execution_info, + token_info=token_info, + streaming=True + ) + + # Mark the streaming session as complete + if hasattr(cli_print_tool_output, '_streaming_sessions') and call_id in cli_print_tool_output._streaming_sessions: + cli_print_tool_output._streaming_sessions[call_id]['is_complete'] = True \ No newline at end of file diff --git a/terrapin_check.py b/terrapin_check.py new file mode 100644 index 00000000..5c52e13e --- /dev/null +++ b/terrapin_check.py @@ -0,0 +1,26 @@ +#!/usr/bin/env python3 +import socket +import ssl +import sys + +def check_terrapin(hostname, port=443): + try: + context = ssl.create_default_context() + with socket.create_connection((hostname, port), timeout=10) as sock: + with context.wrap_socket(sock, server_hostname=hostname) as ssock: + cipher = ssock.cipher() + print(f"Connected to {hostname}:{port}") + print(f"Cipher: {cipher}") + + # Check for server features that might indicate Terrapin vulnerability + # Note: This is a simplified check and not a complete test + if cipher[0].startswith(('TLS_RSA_', 'RSA-')): + print(f"[!] Potentially vulnerable to Terrapin: {hostname} uses RSA key exchange") + else: + print(f"[+] Not likely vulnerable to Terrapin: {hostname} doesn't use RSA key exchange") + + except Exception as e: + print(f"Error connecting to {hostname}:{port} - {str(e)}") + +if __name__ == "__main__": + check_terrapin("aliasrobotics.com") diff --git a/terrapin_test.py b/terrapin_test.py new file mode 100644 index 00000000..b144c718 --- /dev/null +++ b/terrapin_test.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python3 +import socket +import ssl +import sys +import re + +def test_terrapin(hostname, port=443): + context = ssl.create_default_context() + + print(f"Testing {hostname}:{port} for Terrapin vulnerability...") + + try: + with socket.create_connection((hostname, port), timeout=10) as sock: + with context.wrap_socket(sock, server_hostname=hostname) as ssock: + # Get cipher info + cipher = ssock.cipher() + print(f"Connected using: {cipher}") + + # Check if using CBC cipher (vulnerable to Terrapin) + if "CBC" in cipher[0]: + print("WARNING: Server is using CBC cipher which may be vulnerable to Terrapin attacks!") + return True + else: + print("Server is not using CBC ciphers, likely not vulnerable to Terrapin.") + return False + except Exception as e: + print(f"Error: {e}") + return None + +if __name__ == "__main__": + hostname = "aliasrobotics.com" + test_terrapin(hostname)