From 618d982b2d4ed3c6d30cfaf2c5e8d086a88e623b Mon Sep 17 00:00:00 2001 From: luijait Date: Sat, 10 May 2025 21:21:43 +0200 Subject: [PATCH 01/12] Add a lot of stream --- check_s3_bucket.py | 72 ++ check_terrapin.py | 46 + check_vulns.py | 87 ++ hello_world.py | 1 + helloworld.py | 1 + hola_mundo.py | 1 + nmap_results | 0 security_check.py | 71 ++ src/cai/cli.py | 7 + .../agents/models/openai_chatcompletions.py | 247 +++-- src/cai/tools/common.py | 796 ++++++++------ src/cai/util.py | 998 ++++++++++++------ terrapin_check.py | 26 + terrapin_test.py | 32 + 14 files changed, 1686 insertions(+), 699 deletions(-) create mode 100644 check_s3_bucket.py create mode 100644 check_terrapin.py create mode 100644 check_vulns.py create mode 100644 hello_world.py create mode 100644 helloworld.py create mode 100644 hola_mundo.py create mode 100644 nmap_results create mode 100644 security_check.py create mode 100644 terrapin_check.py create mode 100644 terrapin_test.py 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) From db2263cff89950039e60ce5922e48e6d391cf267 Mon Sep 17 00:00:00 2001 From: luijait Date: Sat, 10 May 2025 21:22:00 +0200 Subject: [PATCH 02/12] Add a lot of stream --- check_s3_bucket.py | 72 -------------------------------------- check_terrapin.py | 46 ------------------------ check_vulns.py | 87 ---------------------------------------------- hello_world.py | 1 - helloworld.py | 1 - hola_mundo.py | 1 - nmap_results | 0 terrapin_check.py | 26 -------------- terrapin_test.py | 32 ----------------- 9 files changed, 266 deletions(-) delete mode 100644 check_s3_bucket.py delete mode 100644 check_terrapin.py delete mode 100644 check_vulns.py delete mode 100644 hello_world.py delete mode 100644 helloworld.py delete mode 100644 hola_mundo.py delete mode 100644 nmap_results delete mode 100644 terrapin_check.py delete mode 100644 terrapin_test.py diff --git a/check_s3_bucket.py b/check_s3_bucket.py deleted file mode 100644 index 512ff764..00000000 --- a/check_s3_bucket.py +++ /dev/null @@ -1,72 +0,0 @@ -#!/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 deleted file mode 100644 index cb337e1e..00000000 --- a/check_terrapin.py +++ /dev/null @@ -1,46 +0,0 @@ -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 deleted file mode 100644 index a710f9c2..00000000 --- a/check_vulns.py +++ /dev/null @@ -1,87 +0,0 @@ -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 deleted file mode 100644 index 09907203..00000000 --- a/hello_world.py +++ /dev/null @@ -1 +0,0 @@ -print('Hello, World!') diff --git a/helloworld.py b/helloworld.py deleted file mode 100644 index 09907203..00000000 --- a/helloworld.py +++ /dev/null @@ -1 +0,0 @@ -print('Hello, World!') diff --git a/hola_mundo.py b/hola_mundo.py deleted file mode 100644 index a500c2f4..00000000 --- a/hola_mundo.py +++ /dev/null @@ -1 +0,0 @@ -print("¡Hola Mundo!") diff --git a/nmap_results b/nmap_results deleted file mode 100644 index e69de29b..00000000 diff --git a/terrapin_check.py b/terrapin_check.py deleted file mode 100644 index 5c52e13e..00000000 --- a/terrapin_check.py +++ /dev/null @@ -1,26 +0,0 @@ -#!/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 deleted file mode 100644 index b144c718..00000000 --- a/terrapin_test.py +++ /dev/null @@ -1,32 +0,0 @@ -#!/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) From bd3dc2b28168f2cf2584fc5316e8081f1f982b52 Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 08:53:01 +0200 Subject: [PATCH 03/12] Rebase --- security_check.py | 71 ----------------------------------------------- 1 file changed, 71 deletions(-) delete mode 100644 security_check.py diff --git a/security_check.py b/security_check.py deleted file mode 100644 index 466331f0..00000000 --- a/security_check.py +++ /dev/null @@ -1,71 +0,0 @@ -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") From f6574eb1124e97705a19295e07f609fb8031b7b3 Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 09:43:02 +0200 Subject: [PATCH 04/12] Fix message list --- src/cai/tools/common.py | 22 +- src/cai/util.py | 476 ++++++++++++++++++++++++++++------------ 2 files changed, 357 insertions(+), 141 deletions(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index fa5d6b8f..a0159c15 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -492,7 +492,15 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t # Begin collecting output output_buffer = [] buffer_size = 0 - update_interval = 10 # lines + update_interval = 10 # lines - default for most tools + + # Use a smaller interval for generic_linux_command for better responsiveness + if tool_name == "generic_linux_command": + update_interval = 3 # Update more frequently for terminal commands + + # Add refresh rate info to tool_args for cli_print_tool_output + if "refresh_rate" not in tool_args: + tool_args["refresh_rate"] = 2 # Stream stdout in real-time for line in iter(process.stdout.readline, ''): @@ -750,6 +758,10 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg "workspace": container_workspace } + # Add refresh rate info for generic_linux_command + if tool_name == "generic_linux_command": + tool_args["refresh_rate"] = 2 + # Initialize the streaming session with a consistent call_id format call_id = start_tool_streaming(tool_name, tool_args, call_id) @@ -976,6 +988,10 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg "workspace": os.path.basename(_get_workspace_dir()) } + # Add refresh rate info for generic_linux_command + if tool_name == "generic_linux_command": + tool_args["refresh_rate"] = 2 + # Initialize the streaming session with a consistent call_id format call_id = start_tool_streaming(tool_name, tool_args, call_id) @@ -1057,6 +1073,10 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg "environment": "SSH" } + # Add refresh rate info for generic_linux_command + if tool_name == "generic_linux_command": + tool_args["refresh_rate"] = 2 + # Initialize streaming session with a consistent call_id format call_id = start_tool_streaming(tool_name, tool_args, call_id) diff --git a/src/cai/util.py b/src/cai/util.py index ed2d7a22..f0487d22 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -535,154 +535,150 @@ def fix_litellm_transcription_annotations(): def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 """ Sanitizes the message list passed as a parameter to align with the - OpenAI API message format. + OpenAI API message format, with special attention to tool call sequencing. Adjusts the message list to comply with the following rules: - 1. A tool call id appears no more than twice. - 2. Each tool call id appears as a pair, and both messages - 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). + 1. Each assistant message with tool_calls must be immediately + followed by a tool message for each tool_call, in order. + 2. If a tool message is missing or misplaced, a synthetic one is inserted. + Actual but misplaced/orphaned tool messages are discarded. + 3. Empty user messages are removed. System messages get "" if content is None. + 4. Content of tool messages is ensured to be a string. + 5. Basic deduplication of consecutive identical messages is performed. Args: - messages (List[dict]): List of message dictionaries containing - role, content, and optionally tool_calls or - tool_call_id fields. + messages (List[dict]): List of message dictionaries. Returns: - List[dict]: Sanitized list of messages with invalid tool calls - and empty messages removed. + List[dict]: Sanitized list of messages. """ - # Deep-copy to ensure we don't modify the input - sanitized_messages = [] - - # First pass - identify tool_call_ids from assistant messages and tool messages - tool_call_map = {} # Map from tool_call_id to (assistant_idx, tool_idx) - - 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 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 + if not messages: + return [] + + # Pass 1: Initial filtering and basic content adjustments (same as before) + temp_messages = [] + for msg in messages: + role = msg.get("role") + content = msg.get("content") + tool_calls = msg.get("tool_calls") + if role == "user" and (content is None or str(content).strip() == ""): continue - - # Add valid messages to our sanitized list first - sanitized_messages.append(msg) - - # Now track tool calls and tool messages for pairing - if msg.get("role") == "assistant" and msg.get("tool_calls"): - for tc in msg["tool_calls"]: - if tc.get("id"): - tool_id = tc.get("id") - if tool_id not in tool_call_map: - tool_call_map[tool_id] = {"assistant_idx": len(sanitized_messages) - 1, "tool_idx": None} - - if msg.get("role") == "tool" and msg.get("tool_call_id"): - tool_id = msg.get("tool_call_id") - if tool_id in tool_call_map: - tool_call_map[tool_id]["tool_idx"] = len(sanitized_messages) - 1 - else: - # Tool response without a matching tool call - create a synthetic pair - # by adding a dummy assistant message with a tool_call - assistant_msg = { - "role": "assistant", - "content": None, - "tool_calls": [{ - "id": tool_id, - "type": "function", - "function": { - "name": "unknown_function", - "arguments": "{}" - } - }] - } - # Insert the assistant message *before* the tool message - sanitized_messages.insert(len(sanitized_messages) - 1, assistant_msg) - # Update mapping - 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 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"] - assistant_msg = sanitized_messages[assistant_idx] - - # Find the relevant tool call - tool_name = "unknown_function" - for tc in assistant_msg["tool_calls"]: - if tc.get("id") == tool_id: - if tc.get("function") and tc["function"].get("name"): - tool_name = tc["function"]["name"] - break - - # Create an automatic tool response message - tool_msg = { - "role": "tool", - "tool_call_id": tool_id, - "content": f"Auto-generated response for {tool_name}" - } - - # 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"): + if role == "system" and content is None: 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')}" + if role == "assistant" and tool_calls and content == "": + msg["content"] = None + if role == "assistant" and tool_calls: + for tc_idx, tc in enumerate(tool_calls): + if not tc.get("id"): + import uuid + generated_id = f"generated_id_{uuid.uuid4().hex[:8]}_{tc_idx}" + tc["id"] = generated_id + tc_name = tc.get("function", {}).get("name", "unknown") + # print(color(f"Warning: Generated missing tool_call_id '{generated_id}' for tool '{tc_name}'. Input messages should have tool_call_ids.", "yellow")) + temp_messages.append(msg) - # 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] + if not temp_messages: + return [] + + # Pass 2: Enforce strict tool_call -> tool_result sequencing and remove orphans. + final_messages_pass2 = [] + temp_msg_idx = 0 + # Keep track of tool_call_ids that are expecting a result from the current assistant batch + expecting_results_for_ids = [] + + while temp_msg_idx < len(temp_messages): + current_msg = temp_messages[temp_msg_idx] + + if current_msg.get("role") == "assistant" and current_msg.get("tool_calls"): + final_messages_pass2.append(current_msg) # Add assistant msg + temp_msg_idx += 1 # Consume assistant msg from temp_messages + + # This assistant message made tool calls, so we now expect their results. + expecting_results_for_ids = [(tc.get("id"), tc.get("function", {}).get("name", "unknown")) for tc in current_msg.get("tool_calls", [])] + + for tc_id, tool_name in expecting_results_for_ids: + # Check if the *next* message in temp_messages is the correct tool result + if (temp_msg_idx < len(temp_messages) and + temp_messages[temp_msg_idx].get("role") == "tool" and + temp_messages[temp_msg_idx].get("tool_call_id") == tc_id): + # Correct tool result found and is next, add it. + final_messages_pass2.append(temp_messages[temp_msg_idx]) + temp_msg_idx += 1 # Consume this tool result from temp_messages + else: + # Expected tool result is not immediately next or is missing. + # Insert a synthetic one. + synthetic_tool_msg = { + "role": "tool", + "tool_call_id": tc_id, + "content": f"Synthetic placeholder: Result for {tool_name} (ID: {tc_id}). Original result missing or misplaced immediately after call." + } + final_messages_pass2.append(synthetic_tool_msg) + expecting_results_for_ids = [] # Reset expectations after processing this assistant's calls + + elif current_msg.get("role") == "tool": + # A tool message is encountered. It should only be here if it was part of an expected sequence handled above. + # If `expecting_results_for_ids` is empty, it means we are not currently inside an assistant->tool sequence. + # Any tool message encountered now is either a duplicate of one already processed (if it matched an expectation) + # or it's orphaned/misplaced. We discard it to prevent sequence errors. + # The previous block (assistant with tool_calls) would have consumed this tool message if it was correctly placed. + # So, if we reach this `elif` for a "tool" message, it means it was NOT consumed as an expected result. + # print(color(f"Warning: Discarding tool message (ID: {current_msg.get('tool_call_id')}) as it was not an expected immediate result.", "yellow")) + temp_msg_idx += 1 # Consume (discard) this tool message from temp_messages - # 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"))): + else: # User, System, or Assistant (text-only, no tool_calls processed in this iteration) + final_messages_pass2.append(current_msg) + temp_msg_idx += 1 + expecting_results_for_ids = [] # Text-only assistant resets expectation of tool results. + + # Pass 3: Content normalization (ensure required content fields are present and appropriately typed) + # ... (This part remains the same as your existing Pass 3) + for msg in final_messages_pass2: + role = msg.get("role") + content = msg.get("content") + tool_calls = msg.get("tool_calls") + + if role == "assistant": + if tool_calls and content is None: + if msg.get("content") == "": msg["content"] = None + elif not tool_calls and content is None: + msg["content"] = "" + elif role == "tool": + if content is None: + msg["content"] = f"No output from tool {msg.get('tool_call_id', 'unknown_id')}" + elif not isinstance(content, str): + msg["content"] = str(content) + elif role in ["user", "system"]: + if content is None: + msg["content"] = "" + elif isinstance(content, list): + for part_idx, part in enumerate(content): + if isinstance(part, dict) and part.get("type") == "text" and part.get("text") is None: + content[part_idx]["text"] = "" + + # Pass 4: Basic deduplication of consecutive identical messages (remains the same) + if not final_messages_pass2: + return [] + deduplicated_messages = [final_messages_pass2[0]] + for i in range(1, len(final_messages_pass2)): + prev_msg = deduplicated_messages[-1] + curr_msg = final_messages_pass2[i] + is_duplicate = False + if prev_msg.get("role") == curr_msg.get("role"): + if curr_msg.get("role") == "tool": + if prev_msg.get("tool_call_id") == curr_msg.get("tool_call_id") and \ + prev_msg.get("content") == curr_msg.get("content"): + is_duplicate = True + elif curr_msg.get("role") == "assistant" and curr_msg.get("tool_calls"): + if (prev_msg.get("role") == "assistant" and prev_msg.get("tool_calls") and + prev_msg.get("tool_calls") == curr_msg.get("tool_calls") and + prev_msg.get("content") == curr_msg.get("content")): + is_duplicate = True + elif prev_msg.get("content") == curr_msg.get("content") and not curr_msg.get("tool_calls") and not prev_msg.get("tool_calls"): + is_duplicate = True + if not is_duplicate: + deduplicated_messages.append(curr_msg) - # 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 + return deduplicated_messages def cli_print_tool_call(tool_name="", args="", output="", prefix=" "): """Print a tool call with pretty formatting""" @@ -1542,7 +1538,21 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut else: # Create a new live panel console = Console(theme=theme) - live = Live(panel, console=console, refresh_per_second=4, auto_refresh=True) + + # Set refresh rate - lower for generic_linux_command + refresh_rate = 4 # Default refresh rate + + # Override the refresh rate for specific tools + if tool_name == "generic_linux_command": + refresh_rate = 2 # Lower refresh rate for terminal commands + + # Custom refresh rate from execution_info takes precedence + if execution_info and "refresh_rate" in execution_info: + refresh_rate = execution_info.get("refresh_rate") + + # Create live panel with the appropriate refresh rate + live = Live(panel, console=console, refresh_per_second=refresh_rate, auto_refresh=True) + # Start and store the live panel live.start() _LIVE_STREAMING_PANELS[call_id] = live @@ -1613,6 +1623,9 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut 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 + from rich.panel import Panel + from rich.box import ROUNDED + from rich.console import Group # Format arguments for display args_str = _format_tool_args(args) @@ -1655,6 +1668,131 @@ def _create_tool_panel_content(tool_name, args, output, execution_info=None, tok # Create token information if available token_content = _create_token_info_display(token_info) + # Prepare the content for the panel + group_content = [header, Text("\n\n")] + + # Calculate title width for simplified display option + title_width = len(header.plain) + max_title_width = 100 # Maximum width before simplifying display + + # Extract filtered args for special tool displays + filtered_args = {} + if isinstance(args, dict): + filtered_args = args + else: + try: + # Try to parse the args as JSON if it's a string + import json + if isinstance(args, str) and args.strip().startswith('{'): + filtered_args = json.loads(args) + except (json.JSONDecodeError, TypeError, ValueError): + pass + + # Special handling for execute_code tool + if tool_name == "execute_code" and "language" in filtered_args: + try: + from rich.syntax import Syntax # pylint: disable=import-outside-toplevel,import-error # noqa: E402,E501 + + language = filtered_args.get("language", "python") + code = filtered_args.get("code", "") + + # Create syntax highlighted code panel + syntax = Syntax(code, language, theme="monokai", line_numbers=True, + background_color="#272822", indent_guides=True) + code_panel = Panel( + syntax, + title="Code", + border_style="arrow", + title_align="left", + box=ROUNDED, + padding=(1, 2) + ) + + # Create output panel + output_panel = Panel( + Text(output, style="yellow"), + title="Output", + border_style="border", + title_align="left", + box=ROUNDED, + padding=(1, 2) + ) + + # Don't show code in arguments for execute_code tool + if title_width > max_title_width: + # Just show the tool name without the code in args + simplified_text = Text() + simplified_text.append(f"{tool_name}(", style="bold cyan") + simplified_text.append("...", style="yellow") + simplified_text.append(")", style="bold cyan") + + # Show timeout and filename in the simplified display + timeout = filtered_args.get("timeout", 100) + filename = filtered_args.get("filename", "exploit") + + # Format the simplified timing display + total_elapsed = "N/A" + tool_elapsed = "N/A" + if timing_info: + for info in timing_info: + if info.startswith("Total:"): + total_elapsed = info.replace("Total: ", "") + elif info.startswith("Tool:"): + tool_elapsed = info.replace("Tool: ", "") + + simplified_text.append( + f" [File: {filename} | Timeout: {timeout}s | " + f"Total: {total_elapsed} | Tool: {tool_elapsed}]", + style="bold magenta") + + group_content[0] = simplified_text + + # Add the code and output panels + group_content.extend([ + code_panel, + output_panel + ]) + + # Add token info if available + if token_content: + group_content.append(token_content) + + # Create the full content + content = Group(*group_content) + return header, content + + except Exception as e: # pylint: disable=broad-exception-caught # noqa: E722,E501 + # Fallback if syntax highlighting fails + # Just continue with standard output display + pass + + # Special handling for generic_linux_command tool in bash output + elif tool_name == "generic_linux_command": + try: + from rich.syntax import Syntax # pylint: disable=import-outside-toplevel,import-error # noqa: E402,E501 + + # Create a syntax highlighted bash output + bash_syntax = Syntax(output, "bash", theme="monokai", + background_color="#272822", word_wrap=True) + + # Add the syntax highlighted output + group_content.append(bash_syntax) + + # Add token info if available + if token_content: + group_content.append(Text("\n")) + group_content.append(token_content) + + # Create the full content + content = Group(*group_content) + return header, content + + except Exception as e: # pylint: disable=broad-exception-caught # noqa: E722,E501 + # Fallback if syntax highlighting fails + # Just continue with standard output display + pass + + # Standard output display for other tools # Assemble the full content content = Text() content.append(header) @@ -1899,13 +2037,31 @@ def start_tool_streaming(tool_name, args, call_id=None): if current_time - v.get('timestamp', 0) < 30 # Keep entries from last 30 seconds } + # Custom starting message for generic_linux_command + initial_output = "Starting tool execution..." + tool_specific_info = {} + + # Special handling for generic_linux_command - show command being executed + if tool_name == "generic_linux_command": + # Extract the command being executed for display + if isinstance(args, str): + cmd_display = args + else: + # If it's a dictionary, try to combine command and args + command_val = args.get("command", "") + args_val = args.get("args", "") + cmd_display = f"{command_val} {args_val}".strip() + + initial_output = f"Executing: {cmd_display}\n\nWaiting for output..." + tool_specific_info = {"refresh_rate": 2} # Lower refresh rate for terminal output + # Show initial message with "Starting..." output cli_print_tool_output( tool_name=tool_name, args=args, - output="Starting tool execution...", + output=initial_output, call_id=call_id, - execution_info={"status": "running", "start_time": time.time()}, + execution_info={"status": "running", "start_time": time.time(), **tool_specific_info}, streaming=True ) @@ -1981,4 +2137,44 @@ def finish_tool_streaming(tool_name, args, output, call_id, execution_info=None, # 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 + cli_print_tool_output._streaming_sessions[call_id]['is_complete'] = True + +# Function to add a message to history if it's not a duplicate +def add_to_message_history(msg): + """Add a message to history with refined deduplication logic.""" + if not message_history: + message_history.append(msg) + return + + is_duplicate = False + last_msg = message_history[-1] + + # General check for exact same message object or deep equality with the last one + if msg == last_msg: + is_duplicate = True + else: + current_role = msg.get("role") + last_role = last_msg.get("role") + + if current_role == last_role: + if current_role in ["system", "user"]: + if msg.get("content") == last_msg.get("content"): + is_duplicate = True + elif current_role == "assistant": + if msg.get("tool_calls") and last_msg.get("tool_calls"): + # Compare entire tool_calls structure and content + if (msg.get("tool_calls") == last_msg.get("tool_calls") and + msg.get("content") == last_msg.get("content")): + is_duplicate = True + elif not msg.get("tool_calls") and not last_msg.get("tool_calls"): + # Assistant text messages + if msg.get("content") == last_msg.get("content"): + is_duplicate = True + # If one has tool_calls and the other doesn't, they are not duplicates + elif current_role == "tool": + if (msg.get("tool_call_id") == last_msg.get("tool_call_id") and + msg.get("content") == last_msg.get("content")): + is_duplicate = True + + if not is_duplicate: + message_history.append(msg) \ No newline at end of file From d59a1e7b11e5c916f6cdc9e7da489003af10f10a Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 10:13:23 +0200 Subject: [PATCH 05/12] Fix --- portscanner.txt | 63 +++++++ src/cai/util.py | 478 ++++++++++++++---------------------------------- 2 files changed, 204 insertions(+), 337 deletions(-) create mode 100644 portscanner.txt diff --git a/portscanner.txt b/portscanner.txt new file mode 100644 index 00000000..de43cf73 --- /dev/null +++ b/portscanner.txt @@ -0,0 +1,63 @@ +package main + +import ( + "fmt" + "net" + "sync" + "time" +) + +func main() { + target := "192.168.1.1" + fmt.Printf("Escaneando puertos en %s...\n", target) + + var wg sync.WaitGroup + + // Semáforo para limitar la cantidad de conexiones concurrentes + // y evitar sobrecargar la red + sem := make(chan struct{}, 100) + + // Variable para almacenar los puertos abiertos con un mutex para + // evitar condiciones de carrera + var openPorts []int + var mutex sync.Mutex + + startTime := time.Now() + + // Escanea los primeros 1024 puertos (los más comunes) + for port := 1; port <= 1024; port++ { + wg.Add(1) + sem <- struct{}{} // Adquirir semáforo + + go func(p int) { + defer wg.Done() + defer func() { <-sem }() // Liberar semáforo + + address := fmt.Sprintf("%s:%d", target, p) + conn, err := net.DialTimeout("tcp", address, 500*time.Millisecond) + + if err == nil { + mutex.Lock() + openPorts = append(openPorts, p) + mutex.Unlock() + conn.Close() + } + }(port) + } + + // Esperar a que terminen todas las goroutines + wg.Wait() + + // Ordenar e imprimir los resultados + fmt.Printf("\nEscaneo completado en %s\n", time.Since(startTime)) + fmt.Printf("Puertos abiertos en %s:\n", target) + + if len(openPorts) == 0 { + fmt.Println("No se encontraron puertos abiertos") + } else { + for _, port := range openPorts { + fmt.Printf("- Puerto %d: abierto\n", port) + } + fmt.Printf("\nTotal de puertos abiertos: %d\n", len(openPorts)) + } +} diff --git a/src/cai/util.py b/src/cai/util.py index f0487d22..ed2d7a22 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -535,150 +535,154 @@ def fix_litellm_transcription_annotations(): def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 """ Sanitizes the message list passed as a parameter to align with the - OpenAI API message format, with special attention to tool call sequencing. + OpenAI API message format. Adjusts the message list to comply with the following rules: - 1. Each assistant message with tool_calls must be immediately - followed by a tool message for each tool_call, in order. - 2. If a tool message is missing or misplaced, a synthetic one is inserted. - Actual but misplaced/orphaned tool messages are discarded. - 3. Empty user messages are removed. System messages get "" if content is None. - 4. Content of tool messages is ensured to be a string. - 5. Basic deduplication of consecutive identical messages is performed. + 1. A tool call id appears no more than twice. + 2. Each tool call id appears as a pair, and both messages + 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. + messages (List[dict]): List of message dictionaries containing + role, content, and optionally tool_calls or + tool_call_id fields. Returns: - List[dict]: Sanitized list of messages. + List[dict]: Sanitized list of messages with invalid tool calls + and empty messages removed. """ - if not messages: - return [] - - # Pass 1: Initial filtering and basic content adjustments (same as before) - temp_messages = [] - for msg in messages: - role = msg.get("role") - content = msg.get("content") - tool_calls = msg.get("tool_calls") - if role == "user" and (content is None or str(content).strip() == ""): - continue - if role == "system" and content is None: - msg["content"] = "" - if role == "assistant" and tool_calls and content == "": - msg["content"] = None - if role == "assistant" and tool_calls: - for tc_idx, tc in enumerate(tool_calls): - if not tc.get("id"): - import uuid - generated_id = f"generated_id_{uuid.uuid4().hex[:8]}_{tc_idx}" - tc["id"] = generated_id - tc_name = tc.get("function", {}).get("name", "unknown") - # print(color(f"Warning: Generated missing tool_call_id '{generated_id}' for tool '{tc_name}'. Input messages should have tool_call_ids.", "yellow")) - temp_messages.append(msg) + # Deep-copy to ensure we don't modify the input + sanitized_messages = [] - if not temp_messages: - return [] - - # Pass 2: Enforce strict tool_call -> tool_result sequencing and remove orphans. - final_messages_pass2 = [] - temp_msg_idx = 0 - # Keep track of tool_call_ids that are expecting a result from the current assistant batch - expecting_results_for_ids = [] - - while temp_msg_idx < len(temp_messages): - current_msg = temp_messages[temp_msg_idx] - - if current_msg.get("role") == "assistant" and current_msg.get("tool_calls"): - final_messages_pass2.append(current_msg) # Add assistant msg - temp_msg_idx += 1 # Consume assistant msg from temp_messages + # First pass - identify tool_call_ids from assistant messages and tool messages + tool_call_map = {} # Map from tool_call_id to (assistant_idx, tool_idx) + + 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 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 - # This assistant message made tool calls, so we now expect their results. - expecting_results_for_ids = [(tc.get("id"), tc.get("function", {}).get("name", "unknown")) for tc in current_msg.get("tool_calls", [])] - - for tc_id, tool_name in expecting_results_for_ids: - # Check if the *next* message in temp_messages is the correct tool result - if (temp_msg_idx < len(temp_messages) and - temp_messages[temp_msg_idx].get("role") == "tool" and - temp_messages[temp_msg_idx].get("tool_call_id") == tc_id): - # Correct tool result found and is next, add it. - final_messages_pass2.append(temp_messages[temp_msg_idx]) - temp_msg_idx += 1 # Consume this tool result from temp_messages - else: - # Expected tool result is not immediately next or is missing. - # Insert a synthetic one. - synthetic_tool_msg = { - "role": "tool", - "tool_call_id": tc_id, - "content": f"Synthetic placeholder: Result for {tool_name} (ID: {tc_id}). Original result missing or misplaced immediately after call." - } - final_messages_pass2.append(synthetic_tool_msg) - expecting_results_for_ids = [] # Reset expectations after processing this assistant's calls - - elif current_msg.get("role") == "tool": - # A tool message is encountered. It should only be here if it was part of an expected sequence handled above. - # If `expecting_results_for_ids` is empty, it means we are not currently inside an assistant->tool sequence. - # Any tool message encountered now is either a duplicate of one already processed (if it matched an expectation) - # or it's orphaned/misplaced. We discard it to prevent sequence errors. - # The previous block (assistant with tool_calls) would have consumed this tool message if it was correctly placed. - # So, if we reach this `elif` for a "tool" message, it means it was NOT consumed as an expected result. - # print(color(f"Warning: Discarding tool message (ID: {current_msg.get('tool_call_id')}) as it was not an expected immediate result.", "yellow")) - temp_msg_idx += 1 # Consume (discard) this tool message from temp_messages + # Add valid messages to our sanitized list first + sanitized_messages.append(msg) - else: # User, System, or Assistant (text-only, no tool_calls processed in this iteration) - final_messages_pass2.append(current_msg) - temp_msg_idx += 1 - expecting_results_for_ids = [] # Text-only assistant resets expectation of tool results. - - # Pass 3: Content normalization (ensure required content fields are present and appropriately typed) - # ... (This part remains the same as your existing Pass 3) - for msg in final_messages_pass2: - role = msg.get("role") - content = msg.get("content") - tool_calls = msg.get("tool_calls") - - if role == "assistant": - if tool_calls and content is None: - if msg.get("content") == "": msg["content"] = None - elif not tool_calls and content is None: - msg["content"] = "" - elif role == "tool": - if content is None: - msg["content"] = f"No output from tool {msg.get('tool_call_id', 'unknown_id')}" - elif not isinstance(content, str): - msg["content"] = str(content) - elif role in ["user", "system"]: - if content is None: - msg["content"] = "" - elif isinstance(content, list): - for part_idx, part in enumerate(content): - if isinstance(part, dict) and part.get("type") == "text" and part.get("text") is None: - content[part_idx]["text"] = "" - - # Pass 4: Basic deduplication of consecutive identical messages (remains the same) - if not final_messages_pass2: - return [] - deduplicated_messages = [final_messages_pass2[0]] - for i in range(1, len(final_messages_pass2)): - prev_msg = deduplicated_messages[-1] - curr_msg = final_messages_pass2[i] - is_duplicate = False - if prev_msg.get("role") == curr_msg.get("role"): - if curr_msg.get("role") == "tool": - if prev_msg.get("tool_call_id") == curr_msg.get("tool_call_id") and \ - prev_msg.get("content") == curr_msg.get("content"): - is_duplicate = True - elif curr_msg.get("role") == "assistant" and curr_msg.get("tool_calls"): - if (prev_msg.get("role") == "assistant" and prev_msg.get("tool_calls") and - prev_msg.get("tool_calls") == curr_msg.get("tool_calls") and - prev_msg.get("content") == curr_msg.get("content")): - is_duplicate = True - elif prev_msg.get("content") == curr_msg.get("content") and not curr_msg.get("tool_calls") and not prev_msg.get("tool_calls"): - is_duplicate = True - if not is_duplicate: - deduplicated_messages.append(curr_msg) + # Now track tool calls and tool messages for pairing + if msg.get("role") == "assistant" and msg.get("tool_calls"): + for tc in msg["tool_calls"]: + if tc.get("id"): + tool_id = tc.get("id") + if tool_id not in tool_call_map: + tool_call_map[tool_id] = {"assistant_idx": len(sanitized_messages) - 1, "tool_idx": None} + + if msg.get("role") == "tool" and msg.get("tool_call_id"): + tool_id = msg.get("tool_call_id") + if tool_id in tool_call_map: + tool_call_map[tool_id]["tool_idx"] = len(sanitized_messages) - 1 + else: + # Tool response without a matching tool call - create a synthetic pair + # by adding a dummy assistant message with a tool_call + assistant_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": tool_id, + "type": "function", + "function": { + "name": "unknown_function", + "arguments": "{}" + } + }] + } + # Insert the assistant message *before* the tool message + sanitized_messages.insert(len(sanitized_messages) - 1, assistant_msg) + # Update mapping + 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 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"] + assistant_msg = sanitized_messages[assistant_idx] - return deduplicated_messages + # Find the relevant tool call + tool_name = "unknown_function" + for tc in assistant_msg["tool_calls"]: + if tc.get("id") == tool_id: + if tc.get("function") and tc["function"].get("name"): + tool_name = tc["function"]["name"] + break + + # Create an automatic tool response message + tool_msg = { + "role": "tool", + "tool_call_id": tool_id, + "content": f"Auto-generated response for {tool_name}" + } + + # 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 def cli_print_tool_call(tool_name="", args="", output="", prefix=" "): """Print a tool call with pretty formatting""" @@ -1538,21 +1542,7 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut else: # Create a new live panel console = Console(theme=theme) - - # Set refresh rate - lower for generic_linux_command - refresh_rate = 4 # Default refresh rate - - # Override the refresh rate for specific tools - if tool_name == "generic_linux_command": - refresh_rate = 2 # Lower refresh rate for terminal commands - - # Custom refresh rate from execution_info takes precedence - if execution_info and "refresh_rate" in execution_info: - refresh_rate = execution_info.get("refresh_rate") - - # Create live panel with the appropriate refresh rate - live = Live(panel, console=console, refresh_per_second=refresh_rate, auto_refresh=True) - + 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 @@ -1623,9 +1613,6 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut 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 - from rich.panel import Panel - from rich.box import ROUNDED - from rich.console import Group # Format arguments for display args_str = _format_tool_args(args) @@ -1668,131 +1655,6 @@ def _create_tool_panel_content(tool_name, args, output, execution_info=None, tok # Create token information if available token_content = _create_token_info_display(token_info) - # Prepare the content for the panel - group_content = [header, Text("\n\n")] - - # Calculate title width for simplified display option - title_width = len(header.plain) - max_title_width = 100 # Maximum width before simplifying display - - # Extract filtered args for special tool displays - filtered_args = {} - if isinstance(args, dict): - filtered_args = args - else: - try: - # Try to parse the args as JSON if it's a string - import json - if isinstance(args, str) and args.strip().startswith('{'): - filtered_args = json.loads(args) - except (json.JSONDecodeError, TypeError, ValueError): - pass - - # Special handling for execute_code tool - if tool_name == "execute_code" and "language" in filtered_args: - try: - from rich.syntax import Syntax # pylint: disable=import-outside-toplevel,import-error # noqa: E402,E501 - - language = filtered_args.get("language", "python") - code = filtered_args.get("code", "") - - # Create syntax highlighted code panel - syntax = Syntax(code, language, theme="monokai", line_numbers=True, - background_color="#272822", indent_guides=True) - code_panel = Panel( - syntax, - title="Code", - border_style="arrow", - title_align="left", - box=ROUNDED, - padding=(1, 2) - ) - - # Create output panel - output_panel = Panel( - Text(output, style="yellow"), - title="Output", - border_style="border", - title_align="left", - box=ROUNDED, - padding=(1, 2) - ) - - # Don't show code in arguments for execute_code tool - if title_width > max_title_width: - # Just show the tool name without the code in args - simplified_text = Text() - simplified_text.append(f"{tool_name}(", style="bold cyan") - simplified_text.append("...", style="yellow") - simplified_text.append(")", style="bold cyan") - - # Show timeout and filename in the simplified display - timeout = filtered_args.get("timeout", 100) - filename = filtered_args.get("filename", "exploit") - - # Format the simplified timing display - total_elapsed = "N/A" - tool_elapsed = "N/A" - if timing_info: - for info in timing_info: - if info.startswith("Total:"): - total_elapsed = info.replace("Total: ", "") - elif info.startswith("Tool:"): - tool_elapsed = info.replace("Tool: ", "") - - simplified_text.append( - f" [File: {filename} | Timeout: {timeout}s | " - f"Total: {total_elapsed} | Tool: {tool_elapsed}]", - style="bold magenta") - - group_content[0] = simplified_text - - # Add the code and output panels - group_content.extend([ - code_panel, - output_panel - ]) - - # Add token info if available - if token_content: - group_content.append(token_content) - - # Create the full content - content = Group(*group_content) - return header, content - - except Exception as e: # pylint: disable=broad-exception-caught # noqa: E722,E501 - # Fallback if syntax highlighting fails - # Just continue with standard output display - pass - - # Special handling for generic_linux_command tool in bash output - elif tool_name == "generic_linux_command": - try: - from rich.syntax import Syntax # pylint: disable=import-outside-toplevel,import-error # noqa: E402,E501 - - # Create a syntax highlighted bash output - bash_syntax = Syntax(output, "bash", theme="monokai", - background_color="#272822", word_wrap=True) - - # Add the syntax highlighted output - group_content.append(bash_syntax) - - # Add token info if available - if token_content: - group_content.append(Text("\n")) - group_content.append(token_content) - - # Create the full content - content = Group(*group_content) - return header, content - - except Exception as e: # pylint: disable=broad-exception-caught # noqa: E722,E501 - # Fallback if syntax highlighting fails - # Just continue with standard output display - pass - - # Standard output display for other tools # Assemble the full content content = Text() content.append(header) @@ -2037,31 +1899,13 @@ def start_tool_streaming(tool_name, args, call_id=None): if current_time - v.get('timestamp', 0) < 30 # Keep entries from last 30 seconds } - # Custom starting message for generic_linux_command - initial_output = "Starting tool execution..." - tool_specific_info = {} - - # Special handling for generic_linux_command - show command being executed - if tool_name == "generic_linux_command": - # Extract the command being executed for display - if isinstance(args, str): - cmd_display = args - else: - # If it's a dictionary, try to combine command and args - command_val = args.get("command", "") - args_val = args.get("args", "") - cmd_display = f"{command_val} {args_val}".strip() - - initial_output = f"Executing: {cmd_display}\n\nWaiting for output..." - tool_specific_info = {"refresh_rate": 2} # Lower refresh rate for terminal output - # Show initial message with "Starting..." output cli_print_tool_output( tool_name=tool_name, args=args, - output=initial_output, + output="Starting tool execution...", call_id=call_id, - execution_info={"status": "running", "start_time": time.time(), **tool_specific_info}, + execution_info={"status": "running", "start_time": time.time()}, streaming=True ) @@ -2137,44 +1981,4 @@ def finish_tool_streaming(tool_name, args, output, call_id, execution_info=None, # 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 - -# Function to add a message to history if it's not a duplicate -def add_to_message_history(msg): - """Add a message to history with refined deduplication logic.""" - if not message_history: - message_history.append(msg) - return - - is_duplicate = False - last_msg = message_history[-1] - - # General check for exact same message object or deep equality with the last one - if msg == last_msg: - is_duplicate = True - else: - current_role = msg.get("role") - last_role = last_msg.get("role") - - if current_role == last_role: - if current_role in ["system", "user"]: - if msg.get("content") == last_msg.get("content"): - is_duplicate = True - elif current_role == "assistant": - if msg.get("tool_calls") and last_msg.get("tool_calls"): - # Compare entire tool_calls structure and content - if (msg.get("tool_calls") == last_msg.get("tool_calls") and - msg.get("content") == last_msg.get("content")): - is_duplicate = True - elif not msg.get("tool_calls") and not last_msg.get("tool_calls"): - # Assistant text messages - if msg.get("content") == last_msg.get("content"): - is_duplicate = True - # If one has tool_calls and the other doesn't, they are not duplicates - elif current_role == "tool": - if (msg.get("tool_call_id") == last_msg.get("tool_call_id") and - msg.get("content") == last_msg.get("content")): - is_duplicate = True - - if not is_duplicate: - message_history.append(msg) \ No newline at end of file + cli_print_tool_output._streaming_sessions[call_id]['is_complete'] = True \ No newline at end of file From fc5a6848f3d6079f84fc926814f0b56d7a9b104e Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 10:13:29 +0200 Subject: [PATCH 06/12] Fix --- portscanner.txt | 63 ------------------------------------------------- 1 file changed, 63 deletions(-) delete mode 100644 portscanner.txt diff --git a/portscanner.txt b/portscanner.txt deleted file mode 100644 index de43cf73..00000000 --- a/portscanner.txt +++ /dev/null @@ -1,63 +0,0 @@ -package main - -import ( - "fmt" - "net" - "sync" - "time" -) - -func main() { - target := "192.168.1.1" - fmt.Printf("Escaneando puertos en %s...\n", target) - - var wg sync.WaitGroup - - // Semáforo para limitar la cantidad de conexiones concurrentes - // y evitar sobrecargar la red - sem := make(chan struct{}, 100) - - // Variable para almacenar los puertos abiertos con un mutex para - // evitar condiciones de carrera - var openPorts []int - var mutex sync.Mutex - - startTime := time.Now() - - // Escanea los primeros 1024 puertos (los más comunes) - for port := 1; port <= 1024; port++ { - wg.Add(1) - sem <- struct{}{} // Adquirir semáforo - - go func(p int) { - defer wg.Done() - defer func() { <-sem }() // Liberar semáforo - - address := fmt.Sprintf("%s:%d", target, p) - conn, err := net.DialTimeout("tcp", address, 500*time.Millisecond) - - if err == nil { - mutex.Lock() - openPorts = append(openPorts, p) - mutex.Unlock() - conn.Close() - } - }(port) - } - - // Esperar a que terminen todas las goroutines - wg.Wait() - - // Ordenar e imprimir los resultados - fmt.Printf("\nEscaneo completado en %s\n", time.Since(startTime)) - fmt.Printf("Puertos abiertos en %s:\n", target) - - if len(openPorts) == 0 { - fmt.Println("No se encontraron puertos abiertos") - } else { - for _, port := range openPorts { - fmt.Printf("- Puerto %d: abierto\n", port) - } - fmt.Printf("\nTotal de puertos abiertos: %d\n", len(openPorts)) - } -} From bbea864ee52a907bd1e2ab3cbd795a4ae8aa505e Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 10:42:51 +0200 Subject: [PATCH 07/12] Streamv2beta --- hola_mundo.py | 1 + src/cai/cli.py | 2 +- src/cai/repl/commands/agent.py | 2 +- src/cai/repl/commands/model.py | 2 +- src/cai/repl/ui/banner.py | 4 +- src/cai/repl/ui/prompt.py | 2 +- src/cai/tools/common.py | 112 +++++++++++------ src/cai/tools/reconnaissance/exec_code.py | 71 ++++++++--- src/cai/util.py | 139 ++++++++++++++++++++-- 9 files changed, 266 insertions(+), 69 deletions(-) create mode 100644 hola_mundo.py diff --git a/hola_mundo.py b/hola_mundo.py new file mode 100644 index 00000000..f1409692 --- /dev/null +++ b/hola_mundo.py @@ -0,0 +1 @@ +print('Hola mundo') diff --git a/src/cai/cli.py b/src/cai/cli.py index 0a980e27..b691a493 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -355,7 +355,7 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= title="[bold]Session Summary[/bold]", title_align="left" ) - console.print(time_panel) + console.print(time_panel, end="") print_session_summary(console, metrics, logging_path) diff --git a/src/cai/repl/commands/agent.py b/src/cai/repl/commands/agent.py index dec5b0ee..6be0d2b8 100644 --- a/src/cai/repl/commands/agent.py +++ b/src/cai/repl/commands/agent.py @@ -238,7 +238,7 @@ class AgentCommand(Command): os.environ["CAI_AGENT_TYPE"] = selected_agent_key console.print( - f"[green]Switched to agent: {agent_name}[/green]") + f"[green]Switched to agent: {agent_name}[/green]", end="") visualize_agent_graph(agent) return True diff --git a/src/cai/repl/commands/model.py b/src/cai/repl/commands/model.py index 04d2ad52..07cb0480 100644 --- a/src/cai/repl/commands/model.py +++ b/src/cai/repl/commands/model.py @@ -421,7 +421,7 @@ class ModelCommand(Command): change_message, border_style="green", title="Model Changed" - ) + ), end="" ) return True diff --git a/src/cai/repl/ui/banner.py b/src/cai/repl/ui/banner.py index b1acdb2e..1c52de87 100644 --- a/src/cai/repl/ui/banner.py +++ b/src/cai/repl/ui/banner.py @@ -177,7 +177,7 @@ def display_banner(console: Console): [white] Bug bounty-ready AI[/white] """ - console.print(banner) + console.print(banner, end="") # # Create a table showcasing CAI framework capabilities # # @@ -361,4 +361,4 @@ def display_quick_guide(console: Console): border_style="blue", padding=(1, 2), title_align="center" - )) + ), end="") diff --git a/src/cai/repl/ui/prompt.py b/src/cai/repl/ui/prompt.py index 35de00e0..60358333 100644 --- a/src/cai/repl/ui/prompt.py +++ b/src/cai/repl/ui/prompt.py @@ -96,7 +96,7 @@ def get_user_input( # Get user input with all features return prompt( - [('class:prompt', '\nCAI> ')], + [('class:prompt', 'CAI> ')], completer=command_completer, style=create_prompt_style(), history=FileHistory(str(history_file)), diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index a0159c15..2650a827 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -434,7 +434,7 @@ def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None): return error_msg -def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None): +def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None, custom_args=None): """Runs command locally in the specified workspace_dir.""" # Make sure we're in active time mode for tool execution stop_idle_timer() @@ -453,17 +453,17 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t # Parse command into parts for display parts = command.strip().split(' ', 1) - cmd = parts[0] if parts else "" + cmd_var = parts[0] if parts else "" args_param_val = parts[1] if len(parts) > 1 else "" # Renamed to avoid conflict with tool_args dict key # For generic Linux commands, standardize the tool_name format if not tool_name: - tool_name = f"{cmd}_command" if cmd else "command" + tool_name = f"{cmd_var}_command" if cmd_var else "command" # Create args dictionary with non-empty values only tool_args = {} - if cmd: - tool_args["command"] = cmd + if cmd_var: + tool_args["command"] = cmd_var if args_param_val and args_param_val.strip(): tool_args["args"] = args_param_val @@ -471,9 +471,16 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t tool_args["workspace"] = os.path.basename(target_dir) tool_args["full_command"] = command + # If custom args were provided, merge them with the default args + if custom_args is not None: + if isinstance(custom_args, dict): + # Merge the dictionaries, with custom args taking precedence + for key, value in custom_args.items(): + tool_args[key] = value + # For generic commands, ensure we have a unique call_id if not call_id: - call_id = f"cmd_{cmd}_{str(uuid.uuid4())[:8]}" + call_id = f"cmd_{cmd_var}_{str(uuid.uuid4())[:8]}" # Initialize/use the call_id for this streaming session call_id = start_tool_streaming(tool_name, tool_args, call_id) @@ -561,11 +568,11 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t # 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 "" + cmd_var = parts[0] if parts else "" args_param_val = parts[1] if len(parts) > 1 else "" # Generate a standard tool name for consistency with streaming - standard_tool_name = tool_name or (f"{cmd}_command" if cmd else "command") + standard_tool_name = tool_name or (f"{cmd_var}_command" if cmd_var else "command") # Calculate a consistent command key - must match the format used in cli_print_tool_output command_key = f"{standard_tool_name}:{args_param_val}" @@ -575,7 +582,37 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t return output.strip() if stdout: - print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") + # Create a tool_args dictionary for non-streaming display + # that matches the format used in streaming + tool_display_args = { + "command": cmd_var, + "args": args_param_val, + "full_command": command, + "workspace": os.path.basename(target_dir) + } + + # If custom args were provided, merge them with the default args + if custom_args is not None and isinstance(custom_args, dict): + for key, value in custom_args.items(): + tool_display_args[key] = value + + # Calculate execution info + tool_execution_time = time.time() - process_start_time + exec_info = { + "status": "completed" if result.returncode == 0 else "error", + "return_code": result.returncode, + "environment": "Local", + "host": os.path.basename(target_dir), + "tool_time": tool_execution_time + } + + # Display the command output with rich formatting + cli_print_tool_output( + tool_name=standard_tool_name, + args=tool_display_args, + output=output, + execution_info=exec_info + ) return output.strip() except subprocess.TimeoutExpired as e: @@ -587,13 +624,13 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t 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 "" + cmd_var = parts[0] if parts else "" + args_var = parts[1] if len(parts) > 1 else "" # Ensure tool_args has complete information tool_args = { - "command": cmd, - "args": args if args.strip() else "", + "command": cmd_var, + "args": args_var if args_var.strip() else "", "full_command": command, "environment": "Local", "workspace": os.path.basename(target_dir) @@ -604,7 +641,7 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t "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) + finish_tool_streaming(tool_name or f"{cmd_var}_command", tool_args, error_msg, call_id, execution_info) if stdout: print("\033[32m" + error_msg + "\033[0m") @@ -618,13 +655,13 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t 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 "" + cmd_var = parts[0] if parts else "" + args_var = parts[1] if len(parts) > 1 else "" # Ensure tool_args has complete information tool_args = { - "command": cmd, - "args": args if args.strip() else "", + "command": cmd_var, + "args": args_var if args_var.strip() else "", "full_command": command, "environment": "Local", "workspace": os.path.basename(target_dir) @@ -635,7 +672,7 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t "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) + finish_tool_streaming(tool_name or f"{cmd_var}_command", tool_args, error_msg, call_id, execution_info) print(color(error_msg, fg="red")) return error_msg @@ -647,7 +684,7 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arguments # noqa: E501 async_mode=False, session_id=None, - timeout=100, stream=False, call_id=None, tool_name=None): + timeout=100, stream=False, call_id=None, tool_name=None, args=None): """ Run command in the appropriate environment (Docker, CTF, SSH, Local) and workspace. @@ -663,6 +700,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg call_id: Unique ID for the command execution (for streaming) tool_name: Name of the tool being executed (for display in streaming output). If None, the tool name will be derived from the command. + args: Additional arguments for the tool (for display and context). Returns: str: Command output, status message, or session ID. @@ -673,17 +711,17 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # 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 "" + cmd_name = parts[0] if parts else "" + cmd_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]}" + call_id = f"cmd_{cmd_name}_{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" + tool_name = f"{cmd_name}_command" if cmd_name else "command" try: # If session_id is provided, send command to that session @@ -750,8 +788,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # Create args dictionary with standardized format tool_args = { - "command": cmd, - "args": args if args.strip() else "", + "command": cmd_name, + "args": cmd_args if cmd_args.strip() else "", "full_command": command, "container": container_id[:12], "environment": "Container", @@ -887,7 +925,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg 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()) + return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir(), args) except Exception as e: # Handle other errors @@ -908,7 +946,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg 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()) + return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir(), args) # Handle Synchronous Execution in Container try: @@ -945,7 +983,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_active_timer() start_idle_timer() # Fallback to local execution, preserving workspace context - return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501 # Switch back to idle mode after command completes stop_active_timer() @@ -961,7 +999,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_active_timer() start_idle_timer() # Fallback to local execution on timeout - return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501 except Exception as e: # pylint: disable=broad-except error_msg = f"Error executing command in container: {str(e)}" print(color(f"{context_msg} {error_msg}", fg="red")) @@ -970,7 +1008,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stop_active_timer() start_idle_timer() # Fallback to local execution on other errors - return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501 # --- CTF Execution --- if ctf: @@ -981,8 +1019,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # Create args dictionary with standardized format tool_args = { - "command": cmd, - "args": args if args.strip() else "", + "command": cmd_name, + "args": cmd_args if cmd_args.strip() else "", "full_command": command, "environment": "CTF", "workspace": os.path.basename(_get_workspace_dir()) @@ -1066,8 +1104,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # Create args dictionary with standardized format tool_args = { - "command": cmd, - "args": args if args.strip() else "", + "command": cmd_name, + "args": cmd_args if cmd_args.strip() else "", "full_command": command, "ssh_host": ssh_connection, "environment": "SSH" @@ -1213,7 +1251,9 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg timeout, stream=stream, call_id=call_id, - tool_name=tool_name + tool_name=tool_name, + workspace_dir=_get_workspace_dir(), + custom_args=args ) # Switch back to idle mode after local command completes diff --git a/src/cai/tools/reconnaissance/exec_code.py b/src/cai/tools/reconnaissance/exec_code.py index 794ac990..d532719a 100644 --- a/src/cai/tools/reconnaissance/exec_code.py +++ b/src/cai/tools/reconnaissance/exec_code.py @@ -37,62 +37,101 @@ def execute_code(code: str = "", language: str = "python", "python": "py", "php": "php", "bash": "sh", + "shell": "sh", # Add shell as alias for bash "ruby": "rb", "perl": "pl", "golang": "go", + "go": "go", # Add go as alias for golang "javascript": "js", + "js": "js", # Add js as alias for javascript "typescript": "ts", + "ts": "ts", # Add ts as alias for typescript "rust": "rs", "csharp": "cs", + "cs": "cs", # Add cs as alias for csharp "java": "java", - "kotlin": "kt" + "kotlin": "kt", + "c": "c", # Add C language + "cpp": "cpp", # Add C++ language + "c++": "cpp" # Add C++ language alias } - ext = extensions.get(language.lower(), "txt") + # Normalize language to lowercase + language = language.lower() + ext = extensions.get(language, "txt") full_filename = f"{filename}.{ext}" + # Create code file with content create_cmd = f"cat << 'EOF' > {full_filename}\n{code}\nEOF" - result = run_command(create_cmd, ctf=ctf) + result = run_command(create_cmd, ctf=ctf, stream=True, tool_name="execute_code") if "error" in result.lower(): return f"Failed to create code file: {result}" - if language.lower() == "python": + + # Prepare execution command based on language + if language in ["python", "py"]: exec_cmd = f"python3 {full_filename}" - elif language.lower() == "php": + elif language in ["php"]: exec_cmd = f"php {full_filename}" - elif language.lower() in ["bash", "sh"]: + elif language in ["bash", "sh", "shell"]: exec_cmd = f"bash {full_filename}" - elif language.lower() == "ruby": + elif language in ["ruby", "rb"]: exec_cmd = f"ruby {full_filename}" - elif language.lower() == "perl": + elif language in ["perl", "pl"]: exec_cmd = f"perl {full_filename}" - elif language.lower() == "golang" or language.lower() == "go": + elif language in ["golang", "go"]: temp_dir = f"/tmp/go_exec_{filename}" run_command(f"mkdir -p {temp_dir}", ctf=ctf) run_command(f"cp {full_filename} {temp_dir}/main.go", ctf=ctf) run_command(f"cd {temp_dir} && go mod init temp", ctf=ctf) exec_cmd = f"cd {temp_dir} && go run main.go" - elif language.lower() == "javascript": + elif language in ["javascript", "js"]: exec_cmd = f"node {full_filename}" - elif language.lower() == "typescript": + elif language in ["typescript", "ts"]: exec_cmd = f"ts-node {full_filename}" - elif language.lower() == "rust": + elif language in ["rust", "rs"]: # For Rust, we need to compile first run_command(f"rustc {full_filename} -o {filename}", ctf=ctf) exec_cmd = f"./{filename}" - elif language.lower() == "csharp": + elif language in ["csharp", "cs"]: # For C#, compile with dotnet run_command(f"dotnet build {full_filename}", ctf=ctf) exec_cmd = f"dotnet run {full_filename}" - elif language.lower() == "java": + elif language in ["java"]: # For Java, compile first run_command(f"javac {full_filename}", ctf=ctf) exec_cmd = f"java {filename}" - elif language.lower() == "kotlin": + elif language in ["kotlin", "kt"]: # For Kotlin, compile first run_command(f"kotlinc {full_filename} -include-runtime -d {filename}.jar", ctf=ctf) exec_cmd = f"java -jar {filename}.jar" + elif language in ["c"]: + # For C, compile with gcc + run_command(f"gcc {full_filename} -o {filename}", ctf=ctf) + exec_cmd = f"./{filename}" + elif language in ["cpp", "c++"]: + # For C++, compile with g++ + run_command(f"g++ {full_filename} -o {filename}", ctf=ctf) + exec_cmd = f"./{filename}" else: return f"Unsupported language: {language}" - output = run_command(exec_cmd, ctf=ctf, timeout=timeout) + # Execute the code with syntax-highlighted output + # Create a custom tool args dictionary to send language and code info to the tool output function + tool_args = { + "command": "execute", + "language": language, + "filename": filename, + "code": code, # Include the code for syntax highlighting + "timeout": timeout + } + + # Run the command with streaming to get syntax highlighting + output = run_command( + exec_cmd, + ctf=ctf, + timeout=timeout, + stream=True, + tool_name="execute_code", + args=tool_args + ) return output diff --git a/src/cai/util.py b/src/cai/util.py index ed2d7a22..8941c661 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -22,6 +22,10 @@ from dataclasses import dataclass, field from typing import Dict, Optional import time import threading +from rich.syntax import Syntax # Import Syntax for highlighting +from rich.panel import Panel +from rich.console import Group +from rich.box import ROUNDED # Global timing variables for tracking active and idle time _active_timer_start = None @@ -1183,7 +1187,7 @@ def create_agent_streaming_context(agent_name, counter, model): Text.assemble(header, content, footer), border_style="blue", box=ROUNDED, - padding=(1, 2), + padding=(0, 1), title="[bold]Agent Streaming Response[/bold]", title_align="left", width=panel_width, @@ -1244,7 +1248,7 @@ def update_agent_streaming_content(context, text_delta): Text.assemble(context["header"], context["content"], context["footer"]), border_style="blue", box=ROUNDED, - padding=(1, 2), + padding=(0, 1), title="[bold]Agent Streaming Response[/bold]", title_align="left", width=context.get("panel_width", 100), @@ -1349,13 +1353,12 @@ def finish_agent_streaming(context, final_stats=None): Text.assemble( context["header"], context["content"], - Text("\n\n"), - tokens_text if tokens_text else Text(""), + tokens_text if tokens_text else Text(""), context["footer"] ), border_style="blue", box=ROUNDED, - padding=(1, 2), + padding=(0, 1), title="[bold]Agent Streaming Response[/bold]", title_align="left", width=context.get("panel_width", 100), @@ -1522,7 +1525,7 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut content, title=title, border_style=border_style, - padding=(1, 2), + padding=(0, 1), box=ROUNDED, title_align="left" ) @@ -1586,18 +1589,34 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Create a console for output console = Console(theme=theme) - # Get the panel content + # Get the panel content - with syntax highlighting header, content = _create_tool_panel_content(tool_name, args, output, execution_info, token_info) - # Determine border style + # Determine border style based on status border_style = "blue" # Default for non-streaming + if execution_info: + status = execution_info.get('status', 'completed') + if status == "completed": + border_style = "green" + title = "[bold green]Tool Output [Completed][/bold green]" + elif status == "error": + border_style = "red" + title = "[bold red]Tool Output [Error][/bold red]" + elif status == "timeout": + border_style = "red" + title = "[bold red]Tool Output [Timeout][/bold red]" + else: + title = "[bold blue]Tool Output[/bold blue]" + else: + title = "[bold blue]Tool Output[/bold blue]" + # Create the panel panel = Panel( content, - title="[bold blue]Tool Output[/bold blue]", + title=title, border_style=border_style, - padding=(1, 2), + padding=(0, 1), box=ROUNDED, title_align="left" ) @@ -1613,6 +1632,10 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut 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 + from rich.syntax import Syntax # Import Syntax for highlighting + from rich.panel import Panel + from rich.console import Group + from rich.box import ROUNDED # Format arguments for display args_str = _format_tool_args(args) @@ -1655,7 +1678,101 @@ def _create_tool_panel_content(tool_name, args, output, execution_info=None, tok # Create token information if available token_content = _create_token_info_display(token_info) - # Assemble the full content + # Determine if we need specialized content formatting + group_content = [header] + + # Special handling for execute_code tool + if tool_name == "execute_code": + try: + # Parse args to get language and code + language = "python" # Default + code = "" + filename = "exploit" + + if isinstance(args, dict): + language = args.get("language", "python") + code = args.get("code", "") + filename = args.get("filename", "exploit") + elif isinstance(args, str): + # Try to extract language and code from args string + import re + lang_match = re.search(r"language\s*=\s*['\"]?([a-zA-Z0-9+]+)['\"]?", args) + if lang_match: + language = lang_match.group(1) + + # Create syntax highlighted code panel if we have code + if code: + syntax = Syntax(code, language, theme="monokai", line_numbers=True, + background_color="#272822", indent_guides=True) + code_panel = Panel( + syntax, + title=f"Code ({language})", + border_style="cyan", + title_align="left", + box=ROUNDED, + padding=(0, 1) + ) + + # Create output panel with proper highlighting + output_syntax = Syntax(output, "text", theme="monokai", + background_color="#272822") + output_panel = Panel( + output_syntax, + title="Output", + border_style="green", + title_align="left", + box=ROUNDED, + padding=(0, 1) + ) + + # Add code and output panels to content + group_content.append("\n") + group_content.append(code_panel) + group_content.append("\n") + group_content.append(output_panel) + + # Add token info if available + if token_content: + group_content.append("\n") + group_content.append(token_content) + + return header, Group(*group_content) + except Exception: + # Fallback if syntax highlighting fails + pass + + # Special handling for generic_linux_command + elif tool_name == "generic_linux_command" or "command" in tool_name: + try: + # Highlight the output as bash + output_syntax = Syntax(output, "bash", theme="monokai", + background_color="#272822") + + # Create a panel for the formatted output + output_panel = Panel( + output_syntax, + title="Command Output", + border_style="green", + title_align="left", + box=ROUNDED, + padding=(0, 1) + ) + + # Assemble content with highlighted output + group_content.append("\n") + group_content.append(output_panel) + + # Add token info if available + if token_content: + group_content.append("\n") + group_content.append(token_content) + + return header, Group(*group_content) + except Exception: + # Fallback if syntax highlighting fails + pass + + # Default content assembly for other tools content = Text() content.append(header) content.append("\n\n") From fa59140542f79d12cf5118b3afbc54c53d068ee4 Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 10:42:57 +0200 Subject: [PATCH 08/12] Streamv2beta --- hola_mundo.py | 1 - 1 file changed, 1 deletion(-) delete mode 100644 hola_mundo.py diff --git a/hola_mundo.py b/hola_mundo.py deleted file mode 100644 index f1409692..00000000 --- a/hola_mundo.py +++ /dev/null @@ -1 +0,0 @@ -print('Hola mundo') From c96b3d44d4b1071b9a2adefbe12062e6d77118cf Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 10:56:29 +0200 Subject: [PATCH 09/12] Fix message lists --- src/cai/cli.py | 73 ++++++- .../agents/models/openai_chatcompletions.py | 82 +++++++- src/cai/util.py | 183 +++++++++++++++++- 3 files changed, 333 insertions(+), 5 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index b691a493..7027da7a 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -490,14 +490,85 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= else: # Use non-streamed response response = asyncio.run(Runner.run(agent, conversation_input)) + + # Process the response items for item in response.new_items: + # Handle tool call output items (tool results) if isinstance(item, ToolCallOutputItem): + # First, ensure there's a corresponding assistant message with tool_calls + # before adding the tool response to prevent the OpenAI error + assistant_with_tool_call_exists = False + tool_call_id = item.raw_item["call_id"] + + for msg in message_history: + if (msg.get("role") == "assistant" and + msg.get("tool_calls") and + any(tc.get("id") == tool_call_id for tc in msg.get("tool_calls", []))): + assistant_with_tool_call_exists = True + break + + # If no matching assistant message exists, create one first + if not assistant_with_tool_call_exists: + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": tool_call_id, + "type": "function", + "function": { + "name": "unknown_function", + "arguments": "{}" + } + }] + } + add_to_message_history(tool_call_msg) + + # Now add the tool response tool_msg = { "role": "tool", - "tool_call_id": item.raw_item["call_id"], # Use consistent format with streaming + "tool_call_id": tool_call_id, "content": item.output, } add_to_message_history(tool_msg) + + # Make sure that assistant messages with tool calls are also added to message_history + # This is especially important for non-streaming mode + if hasattr(agent, 'model'): + # Access the _Converter directly from the OpenAIChatCompletionsModel implementation + from cai.sdk.agents.models.openai_chatcompletions import _Converter + + # Check if recent_tool_calls exists and process them + if hasattr(_Converter, 'recent_tool_calls'): + for call_id, call_info in _Converter.recent_tool_calls.items(): + # Only process new tool calls that haven't been added to message history yet + tool_call_found = False + for msg in message_history: + if (msg.get("role") == "assistant" and + msg.get("tool_calls") and + any(tc.get("id") == call_id for tc in msg.get("tool_calls", []))): + tool_call_found = True + break + + if not tool_call_found: + # Add the assistant message with the tool call + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": call_id, + "type": "function", + "function": { + "name": call_info.get('name', ''), + "arguments": call_info.get('arguments', '{}') + } + }] + } + add_to_message_history(tool_call_msg) + + # Final validation to ensure message history follows OpenAI's requirements + # Ensure every tool message has a preceding assistant message with matching tool_call_id + from cai.util import fix_message_list + message_history[:] = fix_message_list(message_history) turn_count += 1 # Stop measuring active time and start measuring idle time again diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index cd921ac9..14d51b8c 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -1433,10 +1433,17 @@ class OpenAIChatCompletionsModel(Model): if tracing.include_data(): span.span_data.input = converted_messages - # Ensure message list has correct structure regardless of error condition + # IMPORTANT: Always sanitize the message list to prevent tool call errors + # This is critical to fix common errors with tool/assistant sequences try: from cai.util import fix_message_list + prev_length = len(converted_messages) converted_messages = fix_message_list(converted_messages) + new_length = len(converted_messages) + + # Log if the message list was changed significantly + if new_length != prev_length: + logger.debug(f"Message list was fixed: {prev_length} -> {new_length} messages") except Exception as e: logger.warning(f"Failed to fix message list: {e}") @@ -1623,15 +1630,49 @@ class OpenAIChatCompletionsModel(Model): elif ("An assistant message with 'tool_calls'" in str(e) or "`tool_use` blocks must be followed by a user message with `tool_result`" in str(e) or # noqa: E501 # pylint: disable=C0301 - "An assistant message with 'tool_calls' must be followed by tool messages" in str(e)): # Añadir esta condición + "An assistant message with 'tool_calls' must be followed by tool messages" in str(e) or + "messages with role 'tool' must be a response to a preceeding message with 'tool_calls'" in str(e)): print(f"Error: {str(e)}") + + # Use the pretty message history printer instead of the simple loop + try: + from cai.util import print_message_history + print("\nCurrent message sequence causing the error:") + print_message_history(kwargs["messages"], title="Message Sequence Error") + except ImportError: + # Fall back to simple printing if the function isn't available + print("\nCurrent message sequence causing the error:") + for i, msg in enumerate(kwargs["messages"]): + role = msg.get("role", "unknown") + content_type = ( + "text" if isinstance(msg.get("content"), str) else + "list" if isinstance(msg.get("content"), list) else + "None" if msg.get("content") is None else + type(msg.get("content")).__name__ + ) + tool_calls = "with tool_calls" if msg.get("tool_calls") else "" + tool_call_id = f", tool_call_id: {msg.get('tool_call_id')}" if msg.get("tool_call_id") else "" + + print(f" [{i}] {role}{tool_call_id} (content: {content_type}) {tool_calls}") + # NOTE: EDGE CASE: Report Agent CTRL C error # # This fix CTRL-C error when message list is incomplete # When a tool is not finished but the LLM generates a tool call try: from cai.util import fix_message_list - kwargs["messages"] = fix_message_list(kwargs["messages"]) + print("Attempting to fix message sequence...") + fixed_messages = fix_message_list(kwargs["messages"]) + + # Show the fixed messages if they're different + if fixed_messages != kwargs["messages"]: + try: + from cai.util import print_message_history + print_message_history(fixed_messages, title="Fixed Message Sequence") + except ImportError: + print("Messages fixed successfully.") + + kwargs["messages"] = fixed_messages except Exception as fix_error: print(f"Failed to fix message sequence: {fix_error}") @@ -2446,6 +2487,41 @@ class _Converter: # Continue with normal processing flush_assistant_message() + + # CRITICAL: Verify this tool message has a matching assistant message before adding it + # Find if there's any assistant message with a tool call matching this ID + has_matching_assistant_message = False + for msg in result: + if ( + msg.get("role") == "assistant" and + msg.get("tool_calls") and + any(tc.get("id") == call_id for tc in msg.get("tool_calls", [])) + ): + has_matching_assistant_message = True + break + + # If no matching assistant message, create one + if not has_matching_assistant_message: + # Create a synthetic assistant message with this tool call + asst_msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": "{}" # Use empty object as default arguments + } + } + ] + } + # Add to result list + result.append(asst_msg) + logger.debug(f"Created synthetic assistant message for tool call {call_id}") + + # Now add the tool message msg: ChatCompletionToolMessageParam = { "role": "tool", "tool_call_id": func_output["call_id"], diff --git a/src/cai/util.py b/src/cai/util.py index 8941c661..1595b211 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -26,6 +26,7 @@ from rich.syntax import Syntax # Import Syntax for highlighting from rich.panel import Panel from rich.console import Group from rich.box import ROUNDED +from rich.table import Table # Global timing variables for tracking active and idle time _active_timer_start = None @@ -549,6 +550,8 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 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). + 6. Each 'tool' message must be immediately preceded by an 'assistant' message + with matching tool_call_id in its tool_calls. Args: messages (List[dict]): List of message dictionaries containing @@ -611,6 +614,101 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 # Update mapping tool_call_map[tool_id] = {"assistant_idx": len(sanitized_messages) - 2, "tool_idx": len(sanitized_messages) - 1} + # Second pass - ensure correct sequence (tool messages must directly follow their assistant messages) + # This fixes the error "messages with role 'tool' must be a response to a preceeding message with 'tool_calls'" + i = 0 + while i < len(sanitized_messages): + msg = sanitized_messages[i] + + # Check if this is a tool message that might be out of sequence + if msg.get("role") == "tool" and msg.get("tool_call_id"): + tool_id = msg.get("tool_call_id") + + # If this isn't the first message, check if the previous message is a matching assistant message + if i > 0: + prev_msg = sanitized_messages[i-1] + + # Check if the previous message is an assistant message with matching tool_call_id + is_valid_sequence = ( + prev_msg.get("role") == "assistant" and + prev_msg.get("tool_calls") and + any(tc.get("id") == tool_id for tc in prev_msg.get("tool_calls", [])) + ) + + if not is_valid_sequence: + # Find the assistant message with this tool_call_id + assistant_idx = None + for j, assistant_msg in enumerate(sanitized_messages): + if (assistant_msg.get("role") == "assistant" and + assistant_msg.get("tool_calls") and + any(tc.get("id") == tool_id for tc in assistant_msg.get("tool_calls", []))): + assistant_idx = j + break + + # If we found a matching assistant message, move this tool message right after it + if assistant_idx is not None: + # Remember to save the tool message + tool_msg = sanitized_messages.pop(i) + + # Insert right after the assistant message + sanitized_messages.insert(assistant_idx + 1, tool_msg) + + # Adjust i to account for the move + if assistant_idx < i: + # We moved the message backward, so i should point to the next message + # which is now at position i (since we removed a message before it) + continue + else: + # We moved the message forward, so i should now point to the message + # that is now at position i + continue + else: + # No matching assistant message found - create one + assistant_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": tool_id, + "type": "function", + "function": { + "name": "unknown_function", + "arguments": "{}" + } + }] + } + + # Insert the assistant message before the tool message + sanitized_messages.insert(i, assistant_msg) + + # Skip past both messages + i += 2 + continue + else: + # This tool message is at index 0, which means there's no preceding assistant message + # Create a dummy assistant message + assistant_msg = { + "role": "assistant", + "content": None, + "tool_calls": [{ + "id": tool_id, + "type": "function", + "function": { + "name": "unknown_function", + "arguments": "{}" + } + }] + } + + # Insert the assistant message before the tool message + sanitized_messages.insert(0, assistant_msg) + + # Skip past both messages + i += 2 + continue + + # Move to the next message + i += 1 + # Final validation - ensure all tool calls have responses for tool_id, indices in list(tool_call_map.items()): if indices["tool_idx"] is None: @@ -2098,4 +2196,87 @@ def finish_tool_streaming(tool_name, args, output, call_id, execution_info=None, # 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 + cli_print_tool_output._streaming_sessions[call_id]['is_complete'] = True + +def print_message_history(messages, title="Message History"): + """ + Pretty-print a sequence of messages with enhanced debug information. + + Args: + messages (List[dict]): List of message dictionaries to display + title (str, optional): Title to display above the message history + """ + from rich.console import Console + from rich.panel import Panel + from rich.text import Text + from rich.table import Table + + console = Console() + + # Create a table for displaying messages + table = Table(show_header=True, header_style="bold magenta", expand=True) + table.add_column("#", style="dim", width=3) + table.add_column("Role", style="cyan", width=10) + table.add_column("Content", width=40) + table.add_column("Metadata", width=30) + + # Process each message + for i, msg in enumerate(messages): + # Get role with color based on type + role = msg.get("role", "unknown") + role_style = { + "user": "green", + "assistant": "blue", + "system": "yellow", + "tool": "magenta" + }.get(role, "white") + + # Get content preview + content = msg.get("content") + content_preview = "" + if content is None: + content_preview = "[dim]None[/dim]" + elif isinstance(content, str): + # Truncate and escape long content + content_preview = (content[:37] + "...") if len(content) > 40 else content + content_preview = content_preview.replace("\n", "\\n") + elif isinstance(content, list): + content_preview = f"[list with {len(content)} items]" + else: + content_preview = f"[{type(content).__name__}]" + + # Gather metadata + metadata = [] + if msg.get("tool_calls"): + tc_count = len(msg["tool_calls"]) + tc_info = [] + for tc in msg["tool_calls"]: + tc_id = tc.get("id", "unknown") + tc_name = tc.get("function", {}).get("name", "unknown") if "function" in tc else "unknown" + tc_info.append(f"{tc_name}({tc_id})") + metadata.append(f"tool_calls[{tc_count}]: {', '.join(tc_info)}") + + if msg.get("tool_call_id"): + metadata.append(f"tool_call_id: {msg['tool_call_id']}") + + metadata_str = ", ".join(metadata) + + # Add row to table + table.add_row( + str(i), + f"[{role_style}]{role}[/{role_style}]", + content_preview, + metadata_str + ) + + # Create the panel with the table + panel = Panel( + table, + title=f"[bold]{title}[/bold]", + expand=False + ) + + # Display the panel + console.print(panel) + + return len(messages) # Return message count for convenience \ No newline at end of file From f63d7068474fd09cf14a2ceb54d24d8c13edc2ac Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 12:54:18 +0200 Subject: [PATCH 10/12] Streaming finished --- pyproject.toml | 3 +- .../agents/models/openai_chatcompletions.py | 7 + src/cai/util.py | 414 +++++++++++++----- 3 files changed, 311 insertions(+), 113 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 650ff948..385cbd03 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,6 @@ license = {text = "Dual-licensed MIT and Proprietary"} authors = [{ name = "Alias Robotics", email = "research@aliasrobotics.com" }] dependencies = [ # logging and visualization - "flask>=2.0, <3", "folium>=0.15.0, <1", "matplotlib>=3.0, <4", "numpy>=1.21, <3", @@ -26,7 +25,7 @@ dependencies = [ "prompt_toolkit>=3.0.39", "dotenv>=0.9.9", "litellm>=1.63.7", - "mako>=1.3.9", + "mako>=1.3.8", "mcp; python_version >= '3.10'", "mkdocs>=1.6.0", "mkdocs-material>=9.6.0", diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 14d51b8c..0e75b9de 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -420,6 +420,10 @@ class OpenAIChatCompletionsModel(Model): # Only display the agent message if we haven't already shown the tool output if should_display_message: + # Ensure we're in non-streaming mode for proper markdown parsing + previous_stream_setting = os.environ.get('CAI_STREAM', 'false') + os.environ['CAI_STREAM'] = 'false' # Force non-streaming mode for markdown parsing + # Print the agent message for CLI display cli_print_agent_messages( agent_name=getattr(self, 'agent_name', 'Agent'), @@ -444,6 +448,9 @@ class OpenAIChatCompletionsModel(Model): tool_output=None, # Don't pass tool output here, we're using direct display suppress_empty=True # Suppress empty panels ) + + # Restore previous streaming setting + os.environ['CAI_STREAM'] = previous_stream_setting # --- Add assistant tool call to message_history if present --- # If the response contains tool_calls, add them to message_history as assistant messages diff --git a/src/cai/util.py b/src/cai/util.py index 1595b211..b130f367 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -970,28 +970,116 @@ def parse_message_content(message): """ Parse a message object to extract its textual content. Only processes messages that don't have tool calls. + Detects markdown code blocks and applies syntax highlighting in non-streaming mode. + Also formats other markdown elements like headers, lists, and text formatting. Args: message: Can be a string or a Message object with content attribute Returns: - str: The extracted content as a string + str or rich.console.Group: The extracted content as a string or as a rich Group with Syntax highlighting """ - # Check if this is a duplicate print from OpenAIChatCompletionsModel - # If message is already a string, return it + from rich.console import Group + from rich.syntax import Syntax + from rich.text import Text + from rich.markdown import Markdown + import re + + # Extract the raw content + raw_content = "" + + # If message is already a string, use it if isinstance(message, str): - return message - + raw_content = message # If message is a Message object with content attribute - if hasattr(message, 'content') and message.content is not None: - return message.content - + elif hasattr(message, 'content') and message.content is not None: + raw_content = message.content # If message is a dict with content key - if isinstance(message, dict) and 'content' in message: - return message['content'] - + elif isinstance(message, dict) and 'content' in message: + raw_content = message['content'] # If we can't extract content, convert to string - return str(message) + else: + raw_content = str(message) + + # Check if streaming is enabled + streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' + + # Only apply markdown formatting in non-streaming mode + if not streaming_enabled and raw_content: + # Check if content contains markdown code blocks with improved regex + code_block_pattern = r'```(\w*)\s*([\s\S]*?)\s*```' + matches = re.findall(code_block_pattern, raw_content, re.DOTALL) + + if matches: + # Prepare to process markdown with code blocks highlighted + elements = [] + last_end = 0 + + # Find all code blocks with improved regex pattern + for match in re.finditer(r'```(\w*)\s*([\s\S]*?)\s*```', raw_content, re.DOTALL): + # Get text before the code block + start = match.start() + if start > last_end: + text_before = raw_content[last_end:start] + + # Process markdown in the text before the code block + if text_before.strip(): + md = Markdown(text_before) + elements.append(md) + + # Process the code block + lang = match.group(1) or "text" + code = match.group(2) + + # Use the language mapping helper to get proper syntax highlighting + syntax_lang = get_language_from_code_block(lang) + + # Create syntax highlighted code + syntax = Syntax( + code, + syntax_lang, + theme="monokai", + line_numbers=True, + word_wrap=True, + background_color="#272822" + ) + elements.append(syntax) + + last_end = match.end() + + # Add any remaining text after the last code block + if last_end < len(raw_content): + text_after = raw_content[last_end:] + + # Process markdown in the text after the code block + if text_after.strip(): + md = Markdown(text_after) + elements.append(md) + + return Group(*elements) + else: + # If no code blocks, but still contains markdown, use Rich's markdown renderer + # Check for markdown elements (headers, lists, formatting) + has_markdown = any([ + # Headers + re.search(r'^#{1,6}\s+\w+', raw_content, re.MULTILINE), + # Lists + re.search(r'^\s*[-*+]\s+\w+', raw_content, re.MULTILINE), + re.search(r'^\s*\d+\.\s+\w+', raw_content, re.MULTILINE), + # Bold/Italic + '**' in raw_content, + '*' in raw_content and not '**' in raw_content, + '__' in raw_content, + '_' in raw_content and not '__' in raw_content, + # Links + re.search(r'\[.+?\]\(.+?\)', raw_content) + ]) + + if has_markdown: + return Group(Markdown(raw_content)) + + # For streaming mode or no markdown, return the raw content + return raw_content def parse_message_tool_call(message, tool_output=None): """ @@ -1152,23 +1240,33 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli if suppress_empty and not parsed_message and not tool_panels: return + # Check if parsed_message is empty or "null" + is_empty_message = (parsed_message == "null" or parsed_message == "" or + (isinstance(parsed_message, str) and not parsed_message.strip())) + # Also skip if the only message is "null" or empty - if parsed_message == "null" or parsed_message == "": + if is_empty_message: if suppress_empty and not tool_panels: return + # Check if we have Group content from markdown parsing + is_rich_content = False + from rich.console import Group + if isinstance(parsed_message, Group): + is_rich_content = True + # Special handling for Reasoner Agent if agent_name == "Reasoner Agent": text.append(f"[{counter}] ", style="bold red") text.append(f"Agent: {agent_name} ", style="bold yellow") - if parsed_message: + if parsed_message and not is_rich_content: text.append(f">> {parsed_message} ", style="green") text.append(f"[{timestamp}", style="dim") if model: text.append(f" ({os.getenv('CAI_SUPPORT_MODEL')})", style="bold blue") text.append("]", style="dim") - elif not parsed_message: + elif is_empty_message: # When parsed_message is empty, only include timestamp and model info text.append(f"Agent: {agent_name} ", style="bold green") text.append(f"[{timestamp}", style="dim") @@ -1178,7 +1276,7 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli else: text.append(f"[{counter}] ", style="bold cyan") text.append(f"Agent: {agent_name} ", style="bold green") - if parsed_message: + if parsed_message and not is_rich_content: text.append(f">> {parsed_message} ", style="yellow") text.append(f"[{timestamp}", style="dim") if model: @@ -1206,19 +1304,52 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli total_cost ) # Only append token information if there is a parsed message - if parsed_message: + if parsed_message and not is_rich_content: text.append(tokens_text) - panel = Panel( - text, - border_style="red" if agent_name == "Reasoner Agent" else "blue", - box=ROUNDED, - padding=(0, 1), - title=("[bold]Reasoning Analysis[/bold]" - if agent_name == "Reasoner Agent" - else "[bold]Agent Interaction[/bold]"), - title_align="left" - ) + # Create the panel content based on whether we have rich content or not + from rich.panel import Panel + from rich.console import Group + + if is_rich_content: + # For rich content, create a Group with the header, content, and tokens + panel_content = [] + panel_content.append(text) + + # Add spacing between header and content for better readability + panel_content.append(Text("\n")) + + # Add the Group with highlighted content + panel_content.append(parsed_message) + + # Add token information at the bottom with proper spacing + if tokens_text: + panel_content.append(Text("\n")) + panel_content.append(tokens_text) + + panel = Panel( + Group(*panel_content), + border_style="red" if agent_name == "Reasoner Agent" else "blue", + box=ROUNDED, + padding=(1, 1), # Increased padding for better appearance + title=("[bold]Reasoning Analysis[/bold]" + if agent_name == "Reasoner Agent" + else "[bold]Agent Interaction[/bold]"), + title_align="left" + ) + else: + # For regular text content, use the original panel format + panel = Panel( + text, + border_style="red" if agent_name == "Reasoner Agent" else "blue", + box=ROUNDED, + padding=(0, 1), + title=("[bold]Reasoning Analysis[/bold]" + if agent_name == "Reasoner Agent" + else "[bold]Agent Interaction[/bold]"), + title_align="left" + ) + #console.print("\n") console.print(panel) @@ -1490,6 +1621,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, streaming=False): """ Print a tool call output to the command line. @@ -1512,6 +1644,10 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut if not output and not call_id and not streaming: return + # Skip early for execute_code tool in non-streaming mode + if tool_name == "execute_code" 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 = {} @@ -1723,20 +1859,20 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut console.print(panel) except ImportError: - # Fall back to simple formatting if Rich is not available - _print_simple_tool_output(tool_name, args, output, execution_info, token_info) + pass + # 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 - from rich.syntax import Syntax # Import Syntax for highlighting + from rich.syntax import Syntax # Import Syntax for highlighting from rich.panel import Panel from rich.console import Group from rich.box import ROUNDED - # Format arguments for display - args_str = _format_tool_args(args) + # Format arguments for display, passing tool_name for specific formatting + args_str = _format_tool_args(args, tool_name=tool_name) # Get timing information timing_info, tool_time = _get_timing_info(execution_info) @@ -1778,73 +1914,72 @@ def _create_tool_panel_content(tool_name, args, output, execution_info=None, tok # Determine if we need specialized content formatting group_content = [header] - + # Special handling for execute_code tool - if tool_name == "execute_code": - try: - # Parse args to get language and code - language = "python" # Default - code = "" - filename = "exploit" + if tool_name == "execute_code" and isinstance(args, dict): + command_type = args.get("command") + actual_code = args.get("code") + language = args.get("language", "python") # Default to python + + # For execute command, show code and output panels + if command_type == "execute" and actual_code: + # Ensure language is a string for Syntax + language_str = get_language_from_code_block(str(language)) - if isinstance(args, dict): - language = args.get("language", "python") - code = args.get("code", "") - filename = args.get("filename", "exploit") - elif isinstance(args, str): - # Try to extract language and code from args string - import re - lang_match = re.search(r"language\s*=\s*['\"]?([a-zA-Z0-9+]+)['\"]?", args) - if lang_match: - language = lang_match.group(1) - - # Create syntax highlighted code panel if we have code - if code: - syntax = Syntax(code, language, theme="monokai", line_numbers=True, - background_color="#272822", indent_guides=True) - code_panel = Panel( - syntax, - title=f"Code ({language})", - border_style="cyan", - title_align="left", - box=ROUNDED, - padding=(0, 1) - ) + code_syntax = Syntax(actual_code, language_str, theme="monokai", + line_numbers=True, background_color="#272822", + indent_guides=True, word_wrap=True) + code_panel = Panel( + code_syntax, + title=f"Code ({language_str})", + border_style="cyan", + title_align="left", + box=ROUNDED, + padding=(0,1) + ) + group_content.extend([Text("\n"), code_panel]) + + # Panel for the output of the executed code + if output: + # Try to highlight output as text, or specific language if known (e.g. json) + output_lang = "text" + try: + json.loads(output) # Check if output is JSON + output_lang = "json" + except json.JSONDecodeError: + pass # Not JSON, keep as text - # Create output panel with proper highlighting - output_syntax = Syntax(output, "text", theme="monokai", - background_color="#272822") + output_syntax = Syntax(output, output_lang, theme="monokai", + background_color="#272822", word_wrap=True) output_panel = Panel( output_syntax, title="Output", border_style="green", title_align="left", box=ROUNDED, - padding=(0, 1) + padding=(0,1) ) - - # Add code and output panels to content - group_content.append("\n") - group_content.append(code_panel) - group_content.append("\n") - group_content.append(output_panel) - - # Add token info if available - if token_content: - group_content.append("\n") - group_content.append(token_content) - - return header, Group(*group_content) - except Exception: - # Fallback if syntax highlighting fails - pass - - # Special handling for generic_linux_command - elif tool_name == "generic_linux_command" or "command" in tool_name: + group_content.extend([Text("\n"), output_panel]) + # For other commands (like cat) or if no code, just show output if any + elif output: + output_syntax = Syntax(output, "text", theme="monokai", + background_color="#272822", word_wrap=True) + output_panel = Panel( + output_syntax, + title="Output", + border_style="green", + title_align="left", + box=ROUNDED, + padding=(0,1) + ) + group_content.extend([Text("\n"), output_panel]) + + # Special handling for generic_linux_command or any command containing 'command' + elif "command" in tool_name.lower() or "shell" in tool_name.lower(): try: # Highlight the output as bash output_syntax = Syntax(output, "bash", theme="monokai", - background_color="#272822") + background_color="#272822", word_wrap=True) # Create a panel for the formatted output output_panel = Panel( @@ -1857,34 +1992,20 @@ def _create_tool_panel_content(tool_name, args, output, execution_info=None, tok ) # Assemble content with highlighted output - group_content.append("\n") - group_content.append(output_panel) + group_content.extend([Text("\n"), output_panel]) - # Add token info if available - if token_content: - group_content.append("\n") - group_content.append(token_content) - - return header, Group(*group_content) except Exception: - # Fallback if syntax highlighting fails - pass - - # Default content assembly for other tools - content = Text() - content.append(header) - content.append("\n\n") - content.append(output) + # Fallback if syntax highlighting fails, just add raw output + group_content.extend([Text("\n"), Text(output)]) # Add token info if available if token_content: - content.append("\n\n") - content.append(token_content) + group_content.extend([Text("\n"), token_content]) - return header, content + return header, Group(*group_content) # Helper function to format tool arguments -def _format_tool_args(args): +def _format_tool_args(args, tool_name=None): """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): @@ -1894,7 +2015,7 @@ def _format_tool_args(args): parsed_dict = json.loads(args) # Recursively call with the parsed dict for consistent formatting - return _format_tool_args(parsed_dict) + return _format_tool_args(parsed_dict, tool_name=tool_name) except json.JSONDecodeError: # Not valid JSON, or not a dict; return as is return args @@ -1913,11 +2034,17 @@ def _format_tool_args(args): # Skip special flags if key in ["async_mode", "streaming"] and not value: continue + + value_str = str(value) + # Format the value if isinstance(value, str): - arg_parts.append(f"{key}={value}") + # Truncate long string values that are not specifically handled above + if len(value_str) > 70 and key not in ["code", "args"]: + value_str = value_str[:67] + "..." + arg_parts.append(f"{key}={value_str}") else: - arg_parts.append(f"{key}={value}") + arg_parts.append(f"{key}={value_str}") return ", ".join(arg_parts) else: return str(args) @@ -1947,7 +2074,6 @@ def _get_timing_info(execution_info=None): 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.""" @@ -2279,4 +2405,70 @@ def print_message_history(messages, title="Message History"): # Display the panel console.print(panel) - return len(messages) # Return message count for convenience \ No newline at end of file + return len(messages) # Return message count for convenience + +def get_language_from_code_block(lang_identifier): + """ + Maps a language identifier from a markdown code block to a proper syntax + highlighting language name. Handles common aliases and defaults. + + Args: + lang_identifier (str): Language identifier from markdown code block + + Returns: + str: Proper language name for syntax highlighting + """ + # Convert to lowercase and strip whitespace + lang = lang_identifier.lower().strip() if lang_identifier else "" + + # Map common language aliases to their proper names + lang_map = { + # Empty strings or unknown + "": "text", + # Python variants + "py": "python", + "python3": "python", + # JavaScript variants + "js": "javascript", + "jsx": "jsx", + "ts": "typescript", + "tsx": "tsx", + "typescript": "typescript", + # Shell variants + "sh": "bash", + "shell": "bash", + "console": "bash", + "terminal": "bash", + # Web languages + "html": "html", + "css": "css", + "json": "json", + "xml": "xml", + "yml": "yaml", + "yaml": "yaml", + # C family + "c": "c", + "cpp": "cpp", + "c++": "cpp", + "csharp": "csharp", + "cs": "csharp", + "java": "java", + # Other common languages + "go": "go", + "golang": "go", + "ruby": "ruby", + "rb": "ruby", + "rust": "rust", + "php": "php", + "sql": "sql", + "diff": "diff", + "markdown": "markdown", + "md": "markdown", + # Default fallback + "text": "text", + "plaintext": "text", + "txt": "text", + } + + # Return mapped language or default to the original if not in map + return lang_map.get(lang, lang or "text") From 81a431ca9e08fd30253600d1e414f1ab9109236c Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 13:02:09 +0200 Subject: [PATCH 11/12] Final touches --- src/cai/util.py | 33 ++++++++++++++------------------- 1 file changed, 14 insertions(+), 19 deletions(-) diff --git a/src/cai/util.py b/src/cai/util.py index b130f367..771a7853 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -1332,9 +1332,7 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli border_style="red" if agent_name == "Reasoner Agent" else "blue", box=ROUNDED, padding=(1, 1), # Increased padding for better appearance - title=("[bold]Reasoning Analysis[/bold]" - if agent_name == "Reasoner Agent" - else "[bold]Agent Interaction[/bold]"), + title="", title_align="left" ) else: @@ -1344,12 +1342,9 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli border_style="red" if agent_name == "Reasoner Agent" else "blue", box=ROUNDED, padding=(0, 1), - title=("[bold]Reasoning Analysis[/bold]" - if agent_name == "Reasoner Agent" - else "[bold]Agent Interaction[/bold]"), + title="", title_align="left" ) - #console.print("\n") console.print(panel) @@ -1417,7 +1412,7 @@ def create_agent_streaming_context(agent_name, counter, model): border_style="blue", box=ROUNDED, padding=(0, 1), - title="[bold]Agent Streaming Response[/bold]", + title="Stream", title_align="left", width=panel_width, expand=True @@ -1478,7 +1473,7 @@ def update_agent_streaming_content(context, text_delta): border_style="blue", box=ROUNDED, padding=(0, 1), - title="[bold]Agent Streaming Response[/bold]", + title="Stream", title_align="left", width=context.get("panel_width", 100), expand=True @@ -1588,7 +1583,7 @@ def finish_agent_streaming(context, final_stats=None): border_style="blue", box=ROUNDED, padding=(0, 1), - title="[bold]Agent Streaming Response[/bold]", + title="Stream", title_align="left", width=context.get("panel_width", 100), expand=True @@ -1744,13 +1739,13 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Create panel title based on status if status == "running": - title = "[bold yellow]Tool Execution [Running][/bold yellow]" + title = "[bold yellow]Running[/bold yellow]" elif status == "completed": - title = "[bold green]Tool Execution [Completed][/bold green]" + title = "[bold green]Completed[/bold green]" elif status == "error": - title = "[bold red]Tool Execution [Error][/bold red]" + title = "[bold red]Error[/bold red]" elif status == "timeout": - title = "[bold red]Tool Execution [Timeout][/bold red]" + title = "[bold red]Timeout[/bold red]" else: title = "[bold blue]Tool Execution[/bold blue]" @@ -1833,17 +1828,17 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut status = execution_info.get('status', 'completed') if status == "completed": border_style = "green" - title = "[bold green]Tool Output [Completed][/bold green]" + title = f"[bold green]{tool_name}({args_str}) [Completed][/bold green]" elif status == "error": border_style = "red" - title = "[bold red]Tool Output [Error][/bold red]" + title = f"[bold red]{tool_name}({args_str}) [Error][/bold red]" elif status == "timeout": border_style = "red" - title = "[bold red]Tool Output [Timeout][/bold red]" + title = f"[bold red]{tool_name}({args_str}) [Timeout][/bold red]" else: - title = "[bold blue]Tool Output[/bold blue]" + title = f"[bold blue]{tool_name}({args_str})[/bold blue]" else: - title = "[bold blue]Tool Output[/bold blue]" + title = f"[bold blue]{tool_name}({args_str})[/bold blue]" # Create the panel panel = Panel( From 5d745455a2a9cb650634f5fa390b1acce581d6ea Mon Sep 17 00:00:00 2001 From: luijait Date: Sun, 11 May 2025 13:35:08 +0200 Subject: [PATCH 12/12] Fix test --- .../agents/models/openai_chatcompletions.py | 69 ++++++------------- src/cai/util.py | 3 + 2 files changed, 23 insertions(+), 49 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 0e75b9de..afd1ffdf 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -2407,22 +2407,22 @@ class _Converter: # Update execution timing if we have the start time if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: - tool_call = cls.recent_tool_calls[call_id] - if 'start_time' in tool_call: + tool_call_details = cls.recent_tool_calls[call_id] # Renamed for clarity + if 'start_time' in tool_call_details: end_time = time.time() - tool_execution_time = end_time - tool_call['start_time'] + tool_execution_time = end_time - tool_call_details['start_time'] # Update the execution info - if 'execution_info' in tool_call: - tool_call['execution_info']['end_time'] = end_time - tool_call['execution_info']['tool_time'] = tool_execution_time + if 'execution_info' in tool_call_details: + tool_call_details['execution_info']['end_time'] = end_time + tool_call_details['execution_info']['tool_time'] = tool_execution_time # If this is the first tool being executed, record the total time from conversation start if not hasattr(cls, 'conversation_start_time'): - cls.conversation_start_time = tool_call['start_time'] + cls.conversation_start_time = tool_call_details['start_time'] - total_time = end_time - getattr(cls, 'conversation_start_time', tool_call['start_time']) - tool_call['execution_info']['total_time'] = total_time + total_time = end_time - getattr(cls, 'conversation_start_time', tool_call_details['start_time']) + tool_call_details['execution_info']['total_time'] = total_time # Store the output so it can be accessed later if not hasattr(cls, 'tool_outputs'): @@ -2439,10 +2439,10 @@ class _Converter: 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', {}) + tool_call_details = cls.recent_tool_calls[call_id] # Renamed for clarity + tool_name = tool_call_details.get('name', 'Unknown Tool') + tool_args = tool_call_details.get('arguments', {}) + execution_info = tool_call_details.get('execution_info', {}) # Get token counts from the OpenAIChatCompletionsModel if available model_instance = None @@ -2467,14 +2467,14 @@ class _Converter: # 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) + model_name_str = str(model_instance.model) # Ensure model name is string token_info['interaction_cost'] = calculate_model_cost( - model_name, + model_name_str, token_info['interaction_input_tokens'], token_info['interaction_output_tokens'] ) token_info['total_cost'] = calculate_model_cost( - model_name, + model_name_str, token_info['total_input_tokens'], token_info['total_output_tokens'] ) @@ -2495,39 +2495,10 @@ class _Converter: # Continue with normal processing flush_assistant_message() - # CRITICAL: Verify this tool message has a matching assistant message before adding it - # Find if there's any assistant message with a tool call matching this ID - has_matching_assistant_message = False - for msg in result: - if ( - msg.get("role") == "assistant" and - msg.get("tool_calls") and - any(tc.get("id") == call_id for tc in msg.get("tool_calls", [])) - ): - has_matching_assistant_message = True - break - - # If no matching assistant message, create one - if not has_matching_assistant_message: - # Create a synthetic assistant message with this tool call - asst_msg = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": call_id, - "type": "function", - "function": { - "name": tool_name, - "arguments": "{}" # Use empty object as default arguments - } - } - ] - } - # Add to result list - result.append(asst_msg) - logger.debug(f"Created synthetic assistant message for tool call {call_id}") - + # REMOVED THE BLOCK THAT CREATED A SYNTHETIC ASSISTANT MESSAGE HERE + # The responsibility for ensuring a preceding assistant message + # is now fully deferred to fix_message_list, called later. + # Now add the tool message msg: ChatCompletionToolMessageParam = { "role": "tool", diff --git a/src/cai/util.py b/src/cai/util.py index 771a7853..f786d7a3 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -1821,6 +1821,9 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Get the panel content - with syntax highlighting header, content = _create_tool_panel_content(tool_name, args, output, execution_info, token_info) + # Format args for the title display + args_str = _format_tool_args(args, tool_name=tool_name) + # Determine border style based on status border_style = "blue" # Default for non-streaming