Merge pull request #128 from 73ai/sandbox-commands-cli

sandboxing for shell command execution

Author: priyanshujain

cli/CLAUDE.md

@@ -0,0 +1,2 @@
+- use uv for running python scripts
+- Write minimal comments without extra explanations

cli/pyproject.toml

@@ -27,6 +27,7 @@ dependencies = [
     "pyyaml>=6.0.2",
     "openai>=1.98.0",
     "anthropic>=0.60.0",
+    "docker>=7.1.0",
 ]
 
 [project.urls]
@@ -37,6 +38,7 @@ infragpt = "infragpt.main:cli"
 
 
 
+
 [dependency-groups]
 dev = [
     "pytest>=8.4.1",

cli/src/infragpt/agent.py

@@ -12,13 +12,11 @@
 from .history import log_interaction
 from .tools import ToolExecutionCancelled
 
-# Import prompt_toolkit for better input handling
 import pathlib
 from prompt_toolkit import PromptSession
 from prompt_toolkit.history import FileHistory
 
 
-# Create a simple system prompt instead of using the old prompts module
 def get_system_prompt():
     return """You are an intelligent shell operations assistant. You help users with:
 
@@ -36,7 +34,6 @@ def get_system_prompt():
 Be concise but thorough in your responses."""
 
 
-# Initialize console for rich output
 console = Console()
 
 
@@ -68,25 +65,18 @@ def add_message(
                 "execute_shell_command"  # Add tool name for compatibility
             )
 
-        # Keep system message separate
         if role == "system":
             self.system_message = message_dict
         else:
             self.messages.append(message_dict)
-
-            # Maintain context window
             if len(self.messages) > self.max_messages:
                 self.messages = self.messages[-self.max_messages :]
 
     def get_context_messages(self) -> List[Dict[str, Any]]:
         """Get messages formatted for LLM API."""
         context = []
-
-        # Add system message first
         if self.system_message:
             context.append(self.system_message)
-
-        # Add conversation messages
         context.extend(self.messages)
 
         return context
@@ -106,35 +96,22 @@ def __init__(self, model_string: str, api_key: str, verbose: bool = False):
         self.verbose = verbose
         self.context = ConversationContext()
 
-        # Create LLM adapter
         self.llm_adapter = get_llm_adapter(
             model_string=model_string, api_key=api_key, verbose=verbose
         )
 
-        # Initialize command history session
         self._setup_command_history()
-
-        # Don't set up custom signal handlers - let prompt_toolkit handle them
-
-        # Initialize system message
         self._initialize_system_message()
 
     def _setup_command_history(self):
         """Setup command history with persistent storage."""
         try:
-            # Create history directory following InfraGPT conventions
             history_dir = pathlib.Path.home() / ".infragpt"
             history_dir.mkdir(exist_ok=True)
-
-            # History file for command-line input
             history_file = history_dir / "history"
 
-            # WARNING: FileHistory writes clear-text inputs to disk.
-            # This is a security risk if users enter secrets at the prompt.
-            # Create PromptSession with FileHistory
+            # WARNING: FileHistory writes clear-text inputs to disk - security risk if users enter secrets
             self.prompt_session = PromptSession(history=FileHistory(str(history_file)))
-
-            # Best-effort: restrict file permissions to user-only
             try:
                 history_file.chmod(0o600)
             except Exception:
@@ -144,7 +121,6 @@ def _setup_command_history(self):
                 console.print(f"[dim]Command history: {history_file}[/dim]")
 
         except Exception as e:
-            # Fallback to no history if setup fails
             self.prompt_session = PromptSession()
             if self.verbose:
                 console.print(
@@ -167,8 +143,6 @@ def run_interactive_session(self):
         )
 
         console.print(f"[yellow]Model:[/yellow] [bold]{self.model_string}[/bold]")
-
-        # Validate API key
         console.print("[dim]Validating API key...[/dim]")
         try:
             if self.llm_adapter.validate_api_key():
@@ -180,34 +154,26 @@ def run_interactive_session(self):
             console.print(f"[red]✗ API key validation failed: {e}[/red]")
             return
 
-        # Show initial prompt
         console.print("[bold cyan]What would you like me to help with?[/bold cyan]")
         console.print(
             "[dim]Press Ctrl+D to exit, Ctrl+C to interrupt operations[/dim]\n"
         )
 
         while True:
             try:
-                # Get user input
                 user_input = self._get_user_input()
                 if not user_input:
-                    continue  # Go back to prompt for empty input
+                    continue
 
-                # Check for exit commands
                 if user_input.lower() in ["exit", "quit", "bye"]:
                     break
 
-                # Add user message to context
                 self.context.add_message("user", user_input)
-
-                # Process with LLM
                 self._process_user_input(user_input)
 
             except KeyboardInterrupt:
-                # This shouldn't happen with prompt_toolkit, but handle just in case
                 continue
             except EOFError:
-                # Ctrl+D - exit the application
                 console.print("\n[dim]EOF received (Ctrl+D). Exiting...[/dim]")
                 break
 
@@ -218,19 +184,15 @@ def _get_user_input(self) -> str:
         try:
             return self.prompt_session.prompt("> ")
         except KeyboardInterrupt:
-            # Ctrl+C should just return empty string to continue
             return ""
         except EOFError:
-            # Ctrl+D - let this propagate for proper exit handling
             raise
 
     def _process_user_input(self, user_input: str):
         """Process user input with direct SDK streaming and tool execution."""
         try:
-            # Get context messages
             messages = self.context.get_context_messages()
 
-            # Debug: Show message structure if verbose
             if self.verbose:
                 console.print(f"[dim]Context has {len(messages)} messages[/dim]")
                 for i, msg in enumerate(messages):
@@ -242,21 +204,15 @@ def _process_user_input(self, user_input: str):
                         f"[dim]  {i}: {role} (content: {content_len} chars, tools: {has_tools}, tool_id: {has_tool_id})[/dim]"
                     )
 
-            # Show thinking and stream response
             console.print("\n[dim]Thinking...[/dim]")
 
             response_content = ""
             first_content = True
 
-            # Stream response using new adapter with interrupt checking
-            # Note: stream_with_tools already handles the complete tool execution cycle
-            # including getting the final response after tool execution
             try:
                 for chunk in self.llm_adapter.stream_with_tools(messages):
-                    # Handle content streaming
                     if chunk.content:
                         if first_content:
-                            # Clear thinking message and show A: prefix
                             console.print(
                                 "\033[1A\033[K", end=""
                             )  # Move up and clear line
@@ -266,7 +222,6 @@ def _process_user_input(self, user_input: str):
                         response_content += chunk.content
                         console.print(chunk.content, end="")
 
-                    # Handle finish reason
                     if chunk.finish_reason:
                         if self.verbose:
                             console.print(
@@ -276,23 +231,15 @@ def _process_user_input(self, user_input: str):
                 console.print("\n[yellow]Operation cancelled by user.[/yellow]")
                 return
 
-            # Add newline after streaming
             if response_content:
                 console.print()
-
-            # Add assistant message to context if we have content
-            if response_content:
                 self.context.add_message("assistant", response_content)
 
-            # Log interaction
             self._log_interaction(user_input, response_content)
 
         except ToolExecutionCancelled:
-            # User cancelled tool execution - just return to prompt
-            # No need to print anything - the tool already printed a message
             return
         except KeyboardInterrupt:
-            # Ctrl+C during streaming - just return to prompt (message already printed by LLM adapter)
             return
         except Exception as e:
             console.print(f"[bold red]Error processing input:[/bold red] {e}")
@@ -304,19 +251,15 @@ def _process_user_input(self, user_input: str):
     def _log_interaction(self, user_input: str, response: str):
         """Log the interaction for history. Sensitive fields are excluded explicitly."""
         try:
-            # Only safe fields, do not include api_key or any secrets
             interaction_data = {
                 "user_input": user_input,
                 "assistant_response": response,
                 "model": self.model_string,
                 "timestamp": datetime.utcnow().isoformat() + "Z",
-                # Not persisted: allow-list will drop this; used only for error reporting in history.py
                 "verbose": self.verbose,
             }
-            # Note: self.api_key is NOT logged ever
             log_interaction("agent_conversation_v2", interaction_data)
         except Exception as e:
-            # Don't let logging failures interrupt the session
             if self.verbose:
                 console.print(f"[dim]Warning: Could not log interaction: {e}[/dim]")