mirror of https://github.com/aliasrobotics/cai.git
Add a lot of stream
This commit is contained in:
luijait
parent
355b82b1f5
commit
618d982b2d
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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")
|
||||
|
|
@ -0,0 +1 @@
|
|||
print('Hello, World!')
|
||||
|
|
@ -0,0 +1 @@
|
|||
print('Hello, World!')
|
||||
|
|
@ -0,0 +1 @@
|
|||
print("¡Hola Mundo!")
|
||||
|
|
@ -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")
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
998
src/cai/util.py
998
src/cai/util.py
File diff suppressed because it is too large
Load Diff
|
|
@ -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")
|
||||
|
|
@ -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)
|
||||
Loading…
Reference in New Issue