tinyagent-py 0.0.15__py3-none-any.whl → 0.0.16__py3-none-any.whl

tinyagent/code_agent/tiny_code_agent.py

@@ -10,7 +10,10 @@ from tinyagent.hooks.rich_code_ui_callback import RichCodeUICallback
 from tinyagent.hooks.jupyter_notebook_callback import JupyterNotebookCallback
 from .providers.base import CodeExecutionProvider
 from .providers.modal_provider import ModalProvider
+from .providers.seatbelt_provider import SeatbeltProvider
 from .helper import translate_tool_for_code_agent, load_template, render_system_prompt, prompt_code_example, prompt_qwen_helper
+from .utils import truncate_output, format_truncation_message
+import datetime
 
 
 DEFAULT_SUMMARY_SYSTEM_PROMPT = (
@@ -26,7 +29,16 @@ class TinyCodeAgent:
     A TinyAgent specialized for code execution tasks.
     
     This class provides a high-level interface for creating agents that can execute
-    Python code using various providers (Modal, Docker, local execution, etc.).
+    Python code using various providers (Modal, SeatbeltProvider for macOS sandboxing, etc.).
+    
+    Features include:
+    - Code execution in sandboxed environments
+    - Shell command execution with safety checks
+    - Environment variable management (SeatbeltProvider)
+    - File system access controls
+    - Memory management and conversation summarization
+    - Git checkpoint automation
+    - Output truncation controls
     """
     
     def __init__(
@@ -48,6 +60,8 @@ class TinyCodeAgent:
         default_workdir: Optional[str] = None,
         summary_config: Optional[Dict[str, Any]] = None,
         ui: Optional[str] = None,
+        truncation_config: Optional[Dict[str, Any]] = None,
+        auto_git_checkpoint: bool = False,
         **agent_kwargs
     ):
         """
@@ -72,7 +86,33 @@ class TinyCodeAgent:
             default_workdir: Default working directory for shell commands. If None, the current working directory is used.
             summary_config: Optional configuration for generating conversation summaries
             ui: The user interface callback to use ('rich', 'jupyter', or None).
+            truncation_config: Configuration for output truncation (max_tokens, max_lines)
+            auto_git_checkpoint: If True, automatically create git checkpoints after each successful shell command
             **agent_kwargs: Additional arguments passed to TinyAgent
+            
+        Provider Config Options:
+            For SeatbeltProvider:
+                - seatbelt_profile: String containing seatbelt profile rules
+                - seatbelt_profile_path: Path to a file containing seatbelt profile rules
+                - python_env_path: Path to the Python environment to use
+                - bypass_shell_safety: If True, bypass shell command safety checks (default: True for seatbelt)
+                - additional_safe_shell_commands: Additional shell commands to consider safe
+                - additional_safe_control_operators: Additional shell control operators to consider safe
+                - additional_read_dirs: List of additional directories to allow read access to
+                - additional_write_dirs: List of additional directories to allow write access to
+                - environment_variables: Dictionary of environment variables to make available in the sandbox
+            
+            For ModalProvider:
+                - pip_packages: List of additional Python packages to install
+                - authorized_imports: List of authorized Python imports
+                - bypass_shell_safety: If True, bypass shell command safety checks (default: False for modal)
+                - additional_safe_shell_commands: Additional shell commands to consider safe
+                - additional_safe_control_operators: Additional shell control operators to consider safe
+                
+        Truncation Config Options:
+            - max_tokens: Maximum number of tokens to keep in output (default: 3000)
+            - max_lines: Maximum number of lines to keep in output (default: 250)
+            - enabled: Whether truncation is enabled (default: True)
         """
         self.model = model
         self.api_key = api_key
@@ -87,6 +127,15 @@ class TinyCodeAgent:
         self.provider = provider  # Store provider type for reuse
         self.check_string_obfuscation = check_string_obfuscation
         self.default_workdir = default_workdir or os.getcwd()  # Default to current working directory if not specified
+        self.auto_git_checkpoint = auto_git_checkpoint  # Enable/disable automatic git checkpoints
+        
+        # Set up truncation configuration with defaults
+        default_truncation = {
+            "max_tokens": 3000,
+            "max_lines": 250,
+            "enabled": True
+        }
+        self.truncation_config = {**default_truncation, **(truncation_config or {})}
         
         # Create the code execution provider
         self.code_provider = self._create_provider(provider, self.provider_config)
@@ -139,12 +188,70 @@ class TinyCodeAgent:
             final_config["authorized_imports"] = final_authorized_imports
             final_config["check_string_obfuscation"] = self.check_string_obfuscation
             
+            # Shell safety configuration (default to False for Modal)
+            bypass_shell_safety = config.get("bypass_shell_safety", False)
+            additional_safe_shell_commands = config.get("additional_safe_shell_commands", None)
+            additional_safe_control_operators = config.get("additional_safe_control_operators", None)
+            
             return ModalProvider(
                 log_manager=self.log_manager,
                 code_tools=self.code_tools,
                 local_execution=self.local_execution,
+                bypass_shell_safety=bypass_shell_safety,
+                additional_safe_shell_commands=additional_safe_shell_commands,
+                additional_safe_control_operators=additional_safe_control_operators,
                 **final_config
             )
+        elif provider_type.lower() == "seatbelt":
+            # Check if seatbelt is supported on this system
+            if not SeatbeltProvider.is_supported():
+                raise ValueError("Seatbelt provider is not supported on this system. It requires macOS with sandbox-exec.")
+            
+            # Seatbelt only works with local execution
+            if not self.local_execution:
+                raise ValueError("Seatbelt provider requires local execution mode. Please set local_execution=True.")
+            
+            # Create a copy of the config without the parameters we'll pass directly
+            filtered_config = config.copy()
+            for key in ['seatbelt_profile', 'seatbelt_profile_path', 'python_env_path', 
+                        'bypass_shell_safety', 'additional_safe_shell_commands', 
+                        'additional_safe_control_operators', 'additional_read_dirs',
+                        'additional_write_dirs', 'environment_variables']:
+                if key in filtered_config:
+                    filtered_config.pop(key)
+            
+            # Get seatbelt profile configuration
+            seatbelt_profile = config.get("seatbelt_profile", None)
+            seatbelt_profile_path = config.get("seatbelt_profile_path", None)
+            python_env_path = config.get("python_env_path", None)
+            
+            # Shell safety configuration (default to True for Seatbelt)
+            bypass_shell_safety = config.get("bypass_shell_safety", True)
+            additional_safe_shell_commands = config.get("additional_safe_shell_commands", None)
+            additional_safe_control_operators = config.get("additional_safe_control_operators", None)
+            
+            # Additional directory access configuration
+            additional_read_dirs = config.get("additional_read_dirs", None)
+            additional_write_dirs = config.get("additional_write_dirs", None)
+            
+            # Environment variables to make available in the sandbox
+            environment_variables = config.get("environment_variables", {})
+            
+            # Create the seatbelt provider
+            return SeatbeltProvider(
+                log_manager=self.log_manager,
+                code_tools=self.code_tools,
+                seatbelt_profile=seatbelt_profile,
+                seatbelt_profile_path=seatbelt_profile_path,
+                python_env_path=python_env_path,
+                bypass_shell_safety=bypass_shell_safety,
+                additional_safe_shell_commands=additional_safe_shell_commands,
+                additional_safe_control_operators=additional_safe_control_operators,
+                additional_read_dirs=additional_read_dirs,
+                additional_write_dirs=additional_write_dirs,
+                environment_variables=environment_variables,
+                **filtered_config
+            )
         else:
             raise ValueError(f"Unsupported provider type: {provider_type}")
     
@@ -307,6 +414,24 @@ class TinyCodeAgent:
                 # This ensures they stay in sync
                 self.user_variables = self.code_provider.get_user_variables()
                 
+                # Apply truncation if enabled
+                if self.truncation_config["enabled"] and "printed_output" in result:
+                    truncated_output, is_truncated, original_tokens, original_lines = truncate_output(
+                        result["printed_output"],
+                        max_tokens=self.truncation_config["max_tokens"],
+                        max_lines=self.truncation_config["max_lines"]
+                    )
+                    
+                    if is_truncated:
+                        result["printed_output"] = format_truncation_message(
+                            truncated_output,
+                            is_truncated,
+                            original_tokens,
+                            original_lines,
+                            self.truncation_config["max_lines"],
+                            "python_output"
+                        )
+                
                 return json.dumps(result)
             except Exception as e:
                 print("!"*100)
@@ -348,7 +473,7 @@ class TinyCodeAgent:
             - You have the capability to call multiple tools in a single response. When multiple independent pieces of information are requested, batch your tool calls together for optimal performance.
             - You can specify an optional timeout in milliseconds (up to 600000ms / 10 minutes). If not specified, commands will timeout after 120000ms (2 minutes).
             - It is very helpful if you write a clear, concise description of what this command does in 5-10 words.
-            - If the output exceeds 30000 characters, output will be truncated before being returned to you.
+            - If the output is too large, it will be truncated before being returned to you.
             
            - If you _still_ need to run `grep`, STOP. ALWAYS USE ripgrep at `rg` first, which all ${PRODUCT_NAME} users have pre-installed.
           "  - When issuing multiple commands, use the ; or && operator to separate them. DO NOT use newlines (newlines are ok in quoted strings).\n" +
@@ -359,19 +484,66 @@ class TinyCodeAgent:
               <bad-example>
               cd /foo/bar && pytest tests
               </bad-example>
+        
+        ## IMPORTANT: Bash Tool Usage
+        
+        When using the bash tool, you MUST provide all required parameters:
+        
+        **Correct Usage:**
+        ```
+        bash(
+            command=["ls", "-la"],
+            absolute_workdir="/path/to/directory", 
+            description="List files in directory"
+        )
+        ```
+        
+        **For creating files with content, use these safe patterns:**
+        
+        1. **Simple file creation:**
+        ```
+        bash(
+            command=["touch", "filename.txt"],
+            absolute_workdir="/working/directory",
+            description="Create empty file"
+        )
+        ```
+        
+        2. **Write content using cat and heredoc:**
+        ```
+        bash(
+            command=["sh", "-c", "cat > filename.txt << 'EOF'\nYour content here\nEOF"],
+            absolute_workdir="/working/directory", 
+            description="Create file with content"
+        )
+        ```
+        
+        3. **Write content using echo:**
+        ```
+        bash(
+            command=["sh", "-c", "echo 'Your content' > filename.txt"],
+            absolute_workdir="/working/directory",
+            description="Write content to file"
+        )
+        ```
+        
+        **Never:**
+        - Call bash() without all required parameters
+        - Use complex nested quotes without testing
+        - Try to create large files in a single command (break into parts)
 
         Args:
             command: list[str]: The shell command to execute as a list of strings.  Example: ["ls", "-la"] or ["cat", "file.txt"]
                 
             absolute_workdir: str: could be presented workdir in the system prompt or one of the subdirectories of the workdir. This is the only allowed path, and accessing else will result in an error.
             description: str: A clear, concise description of what this command does in 5-10 words.
-            timeout: int: Maximum execution time in seconds (default: 30).
+            timeout: int: Maximum execution time in seconds (default: 60).
         Returns:
             Dictionary with stdout, stderr, and exit_code from the command execution.
             If the command is rejected for security reasons, stderr will contain the reason.
             The stdout will include information about which working directory was used.
         """))
-        async def run_shell(command: List[str], absolute_workdir: str,  description: str, timeout: int = 30) -> str:
+        async def run_shell(command: List[str], absolute_workdir: str,  description: str, timeout: int = 60) -> str:
             """Execute shell commands securely using the configured provider."""
             try:
                 # Use the default working directory if none is specified
@@ -393,6 +565,30 @@ class TinyCodeAgent:
                     })
                 
                 result = await self.code_provider.execute_shell(command, timeout, effective_workdir)
+                
+                # Apply truncation if enabled
+                if self.truncation_config["enabled"] and "stdout" in result and result["stdout"]:
+                    truncated_output, is_truncated, original_tokens, original_lines = truncate_output(
+                        result["stdout"],
+                        max_tokens=self.truncation_config["max_tokens"],
+                        max_lines=self.truncation_config["max_lines"]
+                    )
+                    
+                    if is_truncated:
+                        result["stdout"] = format_truncation_message(
+                            truncated_output,
+                            is_truncated,
+                            original_tokens,
+                            original_lines,
+                            self.truncation_config["max_lines"],
+                            "bash_output"
+                        )
+                
+                # Create a git checkpoint if auto_git_checkpoint is enabled
+                if self.auto_git_checkpoint and result.get("exit_code", 1) == 0:
+                    checkpoint_result = await self._create_git_checkpoint(command, description, effective_workdir)
+                    self.log_manager.get_logger(__name__).info(f"Git checkpoint {effective_workdir} result: {checkpoint_result}")
+                
                 return json.dumps(result)
             except Exception as e:
                 COLOR = {
@@ -407,6 +603,64 @@ class TinyCodeAgent:
         self.agent.add_tool(run_python)
         self.agent.add_tool(run_shell)
     
+    async def _create_git_checkpoint(self, command: List[str], description: str, workdir: str) -> Dict[str, Any]:
+        """
+        Create a git checkpoint after command execution.
+        
+        Args:
+            command: The command that was executed
+            description: Description of the command
+            workdir: Working directory where the command was executed
+            
+        Returns:
+            Dictionary with stdout and stderr from the git operations
+        """
+        try:
+            # Format the command for the commit message
+            cmd_str = " ".join(command)
+            
+            # Check if there are changes to commit
+            git_check_cmd = ["bash", "-c", "if ! git diff-index --quiet HEAD --; then echo 'changes_exist'; else echo 'no_changes'; fi"]
+            check_result = await self.code_provider.execute_shell(git_check_cmd, 10, workdir)
+            
+            # If no changes or check failed, return early
+            if check_result.get("exit_code", 1) != 0 or "no_changes" in check_result.get("stdout", ""):
+                return {"stdout": "No changes detected, skipping git checkpoint", "stderr": ""}
+            
+            # Stage all changes
+            git_add_cmd = ["git", "add", "-A"]
+            add_result = await self.code_provider.execute_shell(git_add_cmd, 30, workdir)
+            
+            if add_result.get("exit_code", 1) != 0:
+                return {
+                    "stdout": "",
+                    "stderr": f"Failed to stage changes: {add_result.get('stderr', '')}"
+                }
+            
+            # Create commit with command description and timestamp
+            timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+            commit_msg = f"Checkpoint: {description} @ {timestamp}\n\nCommand: {cmd_str}"
+            git_commit_cmd = ["git", "commit", "-m", commit_msg, "--no-gpg-sign"]
+            commit_result = await self.code_provider.execute_shell(git_commit_cmd, 30, workdir)
+            
+            if commit_result.get("exit_code", 1) != 0:
+                return {
+                    "stdout": "",
+                    "stderr": f"Failed to create commit: {commit_result.get('stderr', '')}"
+                }
+            
+            # Get the first line of the commit message without using split with \n in f-string
+            first_line = commit_msg.split("\n")[0]
+            return {
+                "stdout": f"✓ Git checkpoint created: {first_line}",
+                "stderr": ""
+            }
+        except Exception as e:
+            return {
+                "stdout": "",
+                "stderr": f"Error creating git checkpoint: {str(e)}"
+            }
+    
     def set_default_workdir(self, workdir: str, create_if_not_exists: bool = False):
         """
         Set the default working directory for shell commands.
@@ -676,6 +930,17 @@ class TinyCodeAgent:
         """
         return self.authorized_imports.copy()
     
+    @classmethod
+    def is_seatbelt_supported(cls) -> bool:
+        """
+        Check if the seatbelt provider is supported on this system.
+        
+        Returns:
+            True if seatbelt is supported (macOS with sandbox-exec), False otherwise
+        """
+        from .providers.seatbelt_provider import SeatbeltProvider
+        return SeatbeltProvider.is_supported()
+    
     def remove_authorized_import(self, import_name: str):
         """
         Remove an authorized import.
@@ -794,6 +1059,119 @@ class TinyCodeAgent:
         else:
             self.log_manager.get_logger(__name__).warning(f"Unknown UI type: {ui_type}. No UI callback will be added.")
 
+    def set_truncation_config(self, config: Dict[str, Any]):
+        """
+        Set the truncation configuration.
+        
+        Args:
+            config: Dictionary containing truncation configuration options:
+                - max_tokens: Maximum number of tokens to keep in output
+                - max_lines: Maximum number of lines to keep in output
+                - enabled: Whether truncation is enabled
+        """
+        self.truncation_config.update(config)
+    
+    def get_truncation_config(self) -> Dict[str, Any]:
+        """
+        Get the current truncation configuration.
+        
+        Returns:
+            Dictionary containing truncation configuration
+        """
+        return self.truncation_config.copy()
+    
+    def enable_truncation(self, enabled: bool = True):
+        """
+        Enable or disable output truncation.
+        
+        Args:
+            enabled: Whether to enable truncation
+        """
+        self.truncation_config["enabled"] = enabled
+
+    def enable_auto_git_checkpoint(self, enabled: bool = True):
+        """
+        Enable or disable automatic git checkpoint creation after successful shell commands.
+        
+        Args:
+            enabled: If True, automatically create git checkpoints. If False, do not create them.
+        """
+        self.auto_git_checkpoint = enabled
+
+    def get_auto_git_checkpoint_status(self) -> bool:
+        """
+        Get the current status of auto_git_checkpoint.
+        
+        Returns:
+            True if auto_git_checkpoint is enabled, False otherwise.
+        """
+        return self.auto_git_checkpoint
+    
+    def set_environment_variables(self, env_vars: Dict[str, str]):
+        """
+        Set environment variables for the code execution provider.
+        Currently only supported for SeatbeltProvider.
+        
+        Args:
+            env_vars: Dictionary of environment variable name -> value pairs
+            
+        Raises:
+            AttributeError: If the provider doesn't support environment variables
+        """
+        if hasattr(self.code_provider, 'set_environment_variables'):
+            self.code_provider.set_environment_variables(env_vars)
+        else:
+            raise AttributeError(f"Provider {self.provider} does not support environment variables")
+    
+    def add_environment_variable(self, name: str, value: str):
+        """
+        Add a single environment variable for the code execution provider.
+        Currently only supported for SeatbeltProvider.
+        
+        Args:
+            name: Environment variable name
+            value: Environment variable value
+            
+        Raises:
+            AttributeError: If the provider doesn't support environment variables
+        """
+        if hasattr(self.code_provider, 'add_environment_variable'):
+            self.code_provider.add_environment_variable(name, value)
+        else:
+            raise AttributeError(f"Provider {self.provider} does not support environment variables")
+    
+    def remove_environment_variable(self, name: str):
+        """
+        Remove an environment variable from the code execution provider.
+        Currently only supported for SeatbeltProvider.
+        
+        Args:
+            name: Environment variable name to remove
+            
+        Raises:
+            AttributeError: If the provider doesn't support environment variables
+        """
+        if hasattr(self.code_provider, 'remove_environment_variable'):
+            self.code_provider.remove_environment_variable(name)
+        else:
+            raise AttributeError(f"Provider {self.provider} does not support environment variables")
+    
+    def get_environment_variables(self) -> Dict[str, str]:
+        """
+        Get a copy of current environment variables from the code execution provider.
+        Currently only supported for SeatbeltProvider.
+        
+        Returns:
+            Dictionary of current environment variables
+            
+        Raises:
+            AttributeError: If the provider doesn't support environment variables
+        """
+        if hasattr(self.code_provider, 'get_environment_variables'):
+            return self.code_provider.get_environment_variables()
+        else:
+            raise AttributeError(f"Provider {self.provider} does not support environment variables")
+
 
 # Example usage demonstrating both LLM tools and code tools
 async def run_example():
@@ -836,7 +1214,12 @@ async def run_example():
         authorized_imports=["tinyagent", "gradio", "requests", "numpy", "pandas"],  # Explicitly specify authorized imports
         local_execution=False,  # Remote execution via Modal (default)
         check_string_obfuscation=True,
-        default_workdir=os.path.join(os.getcwd(), "examples")  # Set a default working directory for shell commands
+        default_workdir=os.path.join(os.getcwd(), "examples"),  # Set a default working directory for shell commands
+        truncation_config={
+            "max_tokens": 3000,
+            "max_lines": 250,
+            "enabled": True
+        }
     )
     
     # Connect to MCP servers
@@ -961,6 +1344,292 @@ async def run_example():
     print("Shell Execution with Explicit Working Directory:")
     print(response_shell_explicit)
     
+    # Test truncation functionality
+    print("\n" + "="*80)
+    print("✂️ Testing output truncation")
+    
+    # Configure truncation with smaller limits for testing
+    agent_remote.set_truncation_config({
+        "max_tokens": 100,  # Very small limit for testing
+        "max_lines": 5      # Very small limit for testing
+    })
+    
+    # Generate a large output to test truncation
+    large_output_prompt = """
+    Generate a large output by printing a lot of text. Create a Python script that:
+    1. Prints numbers from 1 to 1000
+    2. For each number, also print its square and cube
+    3. Add random text for each line to make it longer
+    """
+    
+    response_truncated = await agent_remote.run(large_output_prompt)
+    print("Truncated Output Response:")
+    print(response_truncated)
+    
+    # Test disabling truncation
+    print("\n" + "="*80)
+    print("🔄 Testing with truncation disabled")
+    
+    agent_remote.enable_truncation(False)
+    response_untruncated = await agent_remote.run("Run the same script again but limit to 20 numbers")
+    print("Untruncated Output Response:")
+    print(response_untruncated)
+    
+    # Test git checkpoint functionality
+    print("\n" + "="*80)
+    print("🔄 Testing git checkpoint functionality")
+    
+    # Enable git checkpoints
+    agent_remote.enable_auto_git_checkpoint(True)
+    print(f"Auto Git Checkpoint enabled: {agent_remote.get_auto_git_checkpoint_status()}")
+    
+    # Create a test file to demonstrate git checkpoint
+    git_test_prompt = """
+    Create a new file called test_file.txt with some content, then modify it, and observe
+    that git checkpoints are created automatically after each change.
+    """
+    
+    git_response = await agent_remote.run(git_test_prompt)
+    print("Git Checkpoint Response:")
+    print(git_response)
+    
+    # Disable git checkpoints
+    agent_remote.enable_auto_git_checkpoint(False)
+    print(f"Auto Git Checkpoint disabled: {agent_remote.get_auto_git_checkpoint_status()}")
+    
+    # Test seatbelt provider if supported
+    if TinyCodeAgent.is_seatbelt_supported():
+        print("\n" + "="*80)
+        print("🔒 Testing TinyCodeAgent with SEATBELT provider (sandboxed execution)")
+        
+        # Create a test directory for read/write access
+        test_read_dir = os.path.join(os.getcwd(), "test_read_dir")
+        test_write_dir = os.path.join(os.getcwd(), "test_write_dir")
+        
+        # Create directories if they don't exist
+        os.makedirs(test_read_dir, exist_ok=True)
+        os.makedirs(test_write_dir, exist_ok=True)
+        
+        # Create a test file in the read directory
+        with open(os.path.join(test_read_dir, "test.txt"), "w") as f:
+            f.write("This is a test file for reading")
+        
+        # Create a simple seatbelt profile
+        seatbelt_profile = """(version 1)
+        
+        ; Default to deny everything
+        (deny default)
+        
+        ; Allow network connections with proper DNS resolution
+        (allow network*)
+        (allow network-outbound)
+        (allow mach-lookup)
+        
+        ; Allow process execution
+        (allow process-exec)
+        (allow process-fork)
+        (allow signal (target self))
+        
+        ; Restrict file read to current path and system files
+        (deny file-read* (subpath "/Users"))
+        (allow file-read*
+          (subpath "{os.getcwd()}")
+          (subpath "/usr")
+          (subpath "/System")
+          (subpath "/Library")
+          (subpath "/bin")
+          (subpath "/sbin")
+          (subpath "/opt")
+          (subpath "/private/tmp")
+          (subpath "/private/var/tmp")
+          (subpath "/dev")
+          (subpath "/etc")
+          (literal "/")
+          (literal "/."))
+        
+        ; Allow write access to specified folder and temp directories
+        (deny file-write* (subpath "/"))
+        (allow file-write*
+          (subpath "{os.getcwd()}")
+          (subpath "/private/tmp")
+          (subpath "/private/var/tmp")
+          (subpath "/dev"))
+        
+        ; Allow standard device operations
+        (allow file-write-data
+          (literal "/dev/null")
+          (literal "/dev/dtracehelper")
+          (literal "/dev/tty")
+          (literal "/dev/stdout")
+          (literal "/dev/stderr"))
+        
+        ; Allow iokit operations needed for system functions
+        (allow iokit-open)
+        
+        ; Allow shared memory operations
+        (allow ipc-posix-shm)
+        
+        ; Allow basic system operations
+        (allow file-read-metadata)
+        (allow process-info-pidinfo)
+        (allow process-info-setcontrol)
+        """
+        
+        # Create TinyCodeAgent with seatbelt provider
+        agent_seatbelt = TinyCodeAgent(
+            model="gpt-4.1-mini",
+            tools=[search_web],  # LLM tools
+            code_tools=[data_processor],  # Code tools
+            user_variables={
+                "sample_data": [1, 2, 3, 4, 5, 10, 15, 20]
+            },
+            provider="seatbelt",  # Use seatbelt provider
+            provider_config={
+                "seatbelt_profile": seatbelt_profile,
+                # Alternatively, you can specify a path to a seatbelt profile file:
+                # "seatbelt_profile_path": "/path/to/seatbelt.sb",
+                # "python_env_path": "/path/to/python/env",  # Optional path to Python environment
+                
+                # Specify additional directories for read/write access
+                "additional_read_dirs": [test_read_dir],
+                "additional_write_dirs": [test_write_dir],
+                
+                # Allow git commands
+                "bypass_shell_safety": True,
+                "additional_safe_shell_commands": ["git"],
+                
+                # Environment variables to make available in the sandbox
+                "environment_variables": {
+                    "TEST_READ_DIR": test_read_dir,
+                    "TEST_WRITE_DIR": test_write_dir,
+                    "PROJECT_NAME": "TinyAgent Seatbelt Demo",
+                    "BUILD_VERSION": "1.0.0"
+                }
+            },
+            local_execution=True,  # Required for seatbelt
+            check_string_obfuscation=True,
+            truncation_config={
+                "max_tokens": 500,
+                "max_lines": 20,
+                "enabled": True
+            }
+        )
+        
+        # Connect to MCP servers
+        await agent_seatbelt.connect_to_server("npx", ["-y", "@openbnb/mcp-server-airbnb", "--ignore-robots-txt"])
+        await agent_seatbelt.connect_to_server("npx", ["-y", "@modelcontextprotocol/server-sequential-thinking"])
+        
+        # Test the seatbelt agent
+        response_seatbelt = await agent_seatbelt.run("""
+        I have some sample data. Please use the data_processor tool in Python to analyze my sample_data
+        and show me the results.
+        """)
+        
+        print("Seatbelt Agent Response:")
+        print(response_seatbelt)
+        
+        # Test shell execution in sandbox
+        shell_prompt_sandbox = "Run 'ls -la' to list files in the current directory."
+        
+        response_shell_sandbox = await agent_seatbelt.run(shell_prompt_sandbox)
+        print("Shell Execution in Sandbox:")
+        print(response_shell_sandbox)
+        
+        # Test reading from the additional read directory
+        read_prompt = f"Read the contents of the file in the test_read_dir directory."
+        
+        response_read = await agent_seatbelt.run(read_prompt)
+        print("Reading from Additional Read Directory:")
+        print(response_read)
+        
+        # Test writing to the additional write directory
+        write_prompt = f"Write a file called 'output.txt' with the text 'Hello from sandbox!' in the test_write_dir directory."
+        
+        response_write = await agent_seatbelt.run(write_prompt)
+        print("Writing to Additional Write Directory:")
+        print(response_write)
+        
+        # Test environment variables
+        print("\n" + "="*80)
+        print("🔧 Testing environment variables functionality")
+        
+        # Add additional environment variables dynamically
+        agent_seatbelt.add_environment_variable("CUSTOM_VAR", "custom_value")
+        agent_seatbelt.add_environment_variable("DEBUG_MODE", "true")
+        
+        # Get and display current environment variables
+        current_env_vars = agent_seatbelt.get_environment_variables()
+        print(f"Current environment variables: {list(current_env_vars.keys())}")
+        
+        # Test accessing environment variables in Python and shell
+        env_test_prompt = """
+        Test the environment variables we set:
+        1. In Python, use os.environ to check for CUSTOM_VAR and DEBUG_MODE
+        2. In a shell command, use 'echo $CUSTOM_VAR' and 'echo $DEBUG_MODE'
+        3. Also check the TEST_READ_DIR and TEST_WRITE_DIR variables that were set during initialization
+        4. Show all environment variables that start with 'TEST_' or 'CUSTOM_' or 'DEBUG_'
+        """
+        
+        response_env_test = await agent_seatbelt.run(env_test_prompt)
+        print("Environment Variables Test:")
+        print(response_env_test)
+        
+        # Update environment variables
+        agent_seatbelt.set_environment_variables({
+            "CUSTOM_VAR": "updated_value",
+            "NEW_VAR": "new_value",
+            "API_KEY": "test_api_key_123"
+        })
+        
+        # Test updated environment variables
+        updated_env_test_prompt = """
+        Test the updated environment variables:
+        1. Check that CUSTOM_VAR now has the value 'updated_value'
+        2. Check that NEW_VAR is available with value 'new_value'
+        3. Check that API_KEY is available with value 'test_api_key_123'
+        4. Verify that DEBUG_MODE is no longer available (should have been removed by set operation)
+        """
+        
+        response_updated_env = await agent_seatbelt.run(updated_env_test_prompt)
+        print("Updated Environment Variables Test:")
+        print(response_updated_env)
+        
+        # Remove a specific environment variable
+        agent_seatbelt.remove_environment_variable("API_KEY")
+        
+        # Test that the removed variable is no longer available
+        removed_env_test_prompt = """
+        Test that API_KEY environment variable has been removed:
+        1. Try to access API_KEY in Python - it should not be available
+        2. Use shell command 'echo $API_KEY' - it should be empty
+        3. List all current environment variables that start with 'CUSTOM_' or 'NEW_'
+        """
+        
+        response_removed_env = await agent_seatbelt.run(removed_env_test_prompt)
+        print("Removed Environment Variable Test:")
+        print(response_removed_env)
+        
+        # Test git commands with the custom configuration
+        git_prompt = "Run 'git status' to show the current git status."
+        
+        response_git = await agent_seatbelt.run(git_prompt)
+        print("Git Command Execution:")
+        print(response_git)
+        
+        # Clean up test directories
+        import shutil
+        try:
+            shutil.rmtree(test_read_dir)
+            shutil.rmtree(test_write_dir)
+            print("Cleaned up test directories")
+        except Exception as e:
+            print(f"Error cleaning up test directories: {str(e)}")
+        
+        await agent_seatbelt.close()
+    else:
+        print("\n" + "="*80)
+        print("⚠️  Seatbelt provider is not supported on this system. Skipping seatbelt tests.")
+    
     await agent_remote.close()
     await agent_local.close()