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

tinyagent/code_agent/tiny_code_agent.py

@@ -1,14 +1,26 @@
 import traceback
+import os
+import json
 from textwrap import dedent
 from typing import Optional, List, Dict, Any
 from pathlib import Path
 from tinyagent import TinyAgent, tool
 from tinyagent.hooks.logging_manager import LoggingManager
+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 .helper import translate_tool_for_code_agent, load_template, render_system_prompt, prompt_code_example, prompt_qwen_helper
 
 
+DEFAULT_SUMMARY_SYSTEM_PROMPT = (
+    "You are an expert coding assistant. Your goal is to generate a concise, structured summary "
+    "of the conversation below that captures all essential information needed to continue "
+    "development after context replacement. Include tasks performed, code areas modified or "
+    "reviewed, key decisions or assumptions, test results or errors, and outstanding tasks or next steps."
+    
+)
+
 class TinyCodeAgent:
     """
     A TinyAgent specialized for code execution tasks.
@@ -27,10 +39,15 @@ class TinyCodeAgent:
         code_tools: Optional[List[Any]] = None,
         authorized_imports: Optional[List[str]] = None,
         system_prompt_template: Optional[str] = None,
+        system_prompt: Optional[str] = None,
         provider_config: Optional[Dict[str, Any]] = None,
         user_variables: Optional[Dict[str, Any]] = None,
         pip_packages: Optional[List[str]] = None,
         local_execution: bool = False,
+        check_string_obfuscation: bool = True,
+        default_workdir: Optional[str] = None,
+        summary_config: Optional[Dict[str, Any]] = None,
+        ui: Optional[str] = None,
         **agent_kwargs
     ):
         """
@@ -50,6 +67,11 @@ class TinyCodeAgent:
             pip_packages: List of additional Python packages to install in Modal environment
             local_execution: If True, uses Modal's .local() method for local execution. 
                                 If False, uses Modal's .remote() method for cloud execution (default: False)
+            check_string_obfuscation: If True (default), check for string obfuscation techniques. Set to False to allow 
+                                legitimate use of base64 encoding and other string manipulations.
+            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).
             **agent_kwargs: Additional arguments passed to TinyAgent
         """
         self.model = model
@@ -63,6 +85,8 @@ class TinyCodeAgent:
         self.pip_packages = pip_packages or []
         self.local_execution = local_execution
         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
         
         # Create the code execution provider
         self.code_provider = self._create_provider(provider, self.provider_config)
@@ -72,23 +96,32 @@ class TinyCodeAgent:
             self.code_provider.set_user_variables(self.user_variables)
         
         # Build system prompt
-        self.system_prompt = self._build_system_prompt(system_prompt_template)
+        self.static_system_prompt= system_prompt
+        self.system_prompt =  self._build_system_prompt(system_prompt_template)
+        
         
-        # Create the underlying TinyAgent
+        self.summary_config = summary_config or {}
+
+        # Create the underlying TinyAgent with summary configuration
         self.agent = TinyAgent(
             model=model,
             api_key=api_key,
             system_prompt=self.system_prompt,
             logger=log_manager.get_logger('tinyagent.tiny_agent') if log_manager else None,
+            summary_config=summary_config,
             **agent_kwargs
         )
         
-        # Add the code execution tool
-        self._setup_code_execution_tool()
+        # Add the code execution tools
+        self._setup_code_execution_tools()
         
         # Add LLM tools (not code tools - those go to the provider)
         if self.tools:
             self.agent.add_tools(self.tools)
+
+        # Add the selected UI callback
+        if ui:
+            self.add_ui_callback(ui)
     
     def _create_provider(self, provider_type: str, config: Dict[str, Any]) -> CodeExecutionProvider:
         """Create a code execution provider based on the specified type."""
@@ -104,6 +137,7 @@ class TinyCodeAgent:
             final_config = config.copy()
             final_config["pip_packages"] = final_pip_packages
             final_config["authorized_imports"] = final_authorized_imports
+            final_config["check_string_obfuscation"] = self.check_string_obfuscation
             
             return ModalProvider(
                 log_manager=self.log_manager,
@@ -117,7 +151,9 @@ class TinyCodeAgent:
     def _build_system_prompt(self, template_path: Optional[str] = None) -> str:
         """Build the system prompt for the code agent."""
         # Use default template if none provided
-        if template_path is None:
+        if self.static_system_prompt is not None:
+            return self.static_system_prompt
+        elif template_path is None :
             template_path = str(Path(__file__).parent.parent / "prompts" / "code_agent.yaml")
         
         # Translate code tools to code agent format
@@ -242,8 +278,8 @@ class TinyCodeAgent:
         
         return "\n".join(code_tools_lines)
     
-    def _setup_code_execution_tool(self):
-        """Set up the run_python tool using the code provider."""
+    def _setup_code_execution_tools(self):
+        """Set up the code execution tools using the code provider."""
         @tool(name="run_python", description=dedent("""
         This tool receives Python code and executes it in a sandboxed environment.
         During each intermediate step, you can use 'print()' to save important information.
@@ -261,8 +297,17 @@ class TinyCodeAgent:
         async def run_python(code_lines: List[str], timeout: int = 120) -> str:
             """Execute Python code using the configured provider."""
             try:
+                # Before execution, ensure provider has the latest user variables
+                if self.user_variables:
+                    self.code_provider.set_user_variables(self.user_variables)
+                    
                 result = await self.code_provider.execute_python(code_lines, timeout)
-                return str(result)
+                
+                # After execution, update TinyCodeAgent's user_variables from the provider
+                # This ensures they stay in sync
+                self.user_variables = self.code_provider.get_user_variables()
+                
+                return json.dumps(result)
             except Exception as e:
                 print("!"*100)
                 COLOR = {
@@ -272,9 +317,133 @@ class TinyCodeAgent:
                 print(f"{COLOR['RED']}{str(e)}{COLOR['ENDC']}")
                 print(f"{COLOR['RED']}{traceback.format_exc()}{COLOR['ENDC']}")
                 print("!"*100)
-                return f"Error executing code: {str(e)}"
+                
+                # Even after an exception, update user_variables from the provider
+                # This ensures any variables that were successfully created/modified are preserved
+                self.user_variables = self.code_provider.get_user_variables()
+                
+                return json.dumps({"error": f"Error executing code: {str(e)}"})
+        
+        @tool(name="bash", description=dedent("""
+        This tool executes shell commands securely in a sandboxed environment.
+        Only a limited set of safe commands are allowed for security reasons.
+        Before executing the command, please follow these steps:
+          
+          1. Directory Verification:
+             - If the command will create new directories or files, first use ls to verify the parent directory exists and is the correct location
+             - For example, before running "mkdir foo/bar", first use ls to check that "foo" exists and is the intended parent directory
+          
+          2. Command Execution:
+             - Always quote file paths that contain spaces with double quotes (e.g., cd "path with spaces/file.txt")
+             - Examples of proper quoting:
+               - cd "/Users/name/My Documents" (correct)
+               - cd /Users/name/My Documents (incorrect - will fail)
+               - python "/path/with spaces/script.py" (correct)
+               - python /path/with spaces/script.py (incorrect - will fail)
+             - After ensuring proper quoting, execute the command.
+             - Capture the output of the command.
+          
+          Usage notes:
+            - The command argument is required.
+            - 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 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" +
+            - Try to maintain your current working directory throughout the session by using absolute paths and avoiding usage of `cd`. You may use `cd` if the User explicitly requests it.
+              <good-example>
+              pytest /foo/bar/tests
+              </good-example>
+              <bad-example>
+              cd /foo/bar && pytest tests
+              </bad-example>
+
+        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).
+        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:
+            """Execute shell commands securely using the configured provider."""
+            try:
+                # Use the default working directory if none is specified
+                effective_workdir = absolute_workdir or self.default_workdir
+                print(f" {command} to {description}")
+                # Verify that the working directory exists
+                if effective_workdir and not os.path.exists(effective_workdir):
+                    return json.dumps({
+                        "stdout": "",
+                        "stderr": f"Working directory does not exist: {effective_workdir}",
+                        "exit_code": 1
+                    })
+                
+                if effective_workdir and not os.path.isdir(effective_workdir):
+                    return json.dumps({
+                        "stdout": "",
+                        "stderr": f"Path is not a directory: {effective_workdir}",
+                        "exit_code": 1
+                    })
+                
+                result = await self.code_provider.execute_shell(command, timeout, effective_workdir)
+                return json.dumps(result)
+            except Exception as e:
+                COLOR = {
+                    "RED": "\033[91m",
+                    "ENDC": "\033[0m",
+                }
+                print(f"{COLOR['RED']}{str(e)}{COLOR['ENDC']}")
+                print(f"{COLOR['RED']}{traceback.format_exc()}{COLOR['ENDC']}")
+                
+                return json.dumps({"error": f"Error executing shell command: {str(e)}"})
         
         self.agent.add_tool(run_python)
+        self.agent.add_tool(run_shell)
+    
+    def set_default_workdir(self, workdir: str, create_if_not_exists: bool = False):
+        """
+        Set the default working directory for shell commands.
+        
+        Args:
+            workdir: The path to use as the default working directory
+            create_if_not_exists: If True, create the directory if it doesn't exist
+        
+        Raises:
+            ValueError: If the directory doesn't exist and create_if_not_exists is False
+            OSError: If there's an error creating the directory
+        """
+        workdir = os.path.expanduser(workdir)  # Expand user directory if needed
+        
+        if not os.path.exists(workdir):
+            if create_if_not_exists:
+                try:
+                    os.makedirs(workdir, exist_ok=True)
+                    print(f"Created directory: {workdir}")
+                except OSError as e:
+                    raise OSError(f"Failed to create directory {workdir}: {str(e)}")
+            else:
+                raise ValueError(f"Directory does not exist: {workdir}")
+        
+        if not os.path.isdir(workdir):
+            raise ValueError(f"Path is not a directory: {workdir}")
+            
+        self.default_workdir = workdir
+    
+    def get_default_workdir(self) -> str:
+        """
+        Get the current default working directory for shell commands.
+        
+        Returns:
+            The current default working directory path
+        """
+        return self.default_workdir
     
     async def run(self, user_input: str, max_turns: int = 10) -> str:
         """
@@ -289,6 +458,22 @@ class TinyCodeAgent:
         """
         return await self.agent.run(user_input, max_turns)
     
+    async def resume(self, max_turns: int = 10) -> str:
+        """
+        Resume the conversation without adding a new user message.
+        
+        This method continues the conversation from the current state,
+        allowing the agent to process the existing conversation history
+        and potentially take additional actions.
+        
+        Args:
+            max_turns: Maximum number of conversation turns
+            
+        Returns:
+            The agent's response
+        """
+        return await self.agent.resume(max_turns)
+    
     async def connect_to_server(self, command: str, args: List[str], **kwargs):
         """Connect to an MCP server."""
         return await self.agent.connect_to_server(command, args, **kwargs)
@@ -537,6 +722,78 @@ class TinyCodeAgent:
         """Get the session ID."""
         return self.agent.session_id 
 
+    def set_check_string_obfuscation(self, enabled: bool):
+        """
+        Enable or disable string obfuscation detection.
+        
+        Args:
+            enabled: If True, check for string obfuscation techniques. If False, allow
+                    legitimate use of base64 encoding and other string manipulations.
+        """
+        self.check_string_obfuscation = enabled
+        
+        # Update the provider with the new setting
+        if hasattr(self.code_provider, 'check_string_obfuscation'):
+            self.code_provider.check_string_obfuscation = enabled
+
+    async def summarize(self) -> str:
+        """
+        Generate a summary of the current conversation history.
+        
+        Args:
+        Returns:
+            A string containing the conversation summary
+        """
+        # Use the underlying TinyAgent's summarize_conversation method
+        return await self.agent.summarize()
+        
+    async def compact(self) -> bool:
+        """
+        Compact the conversation history by replacing it with a summary.
+        
+        This method delegates to the underlying TinyAgent's compact method,
+        which:
+        1. Generates a summary of the current conversation
+        2. If successful, replaces the conversation with just [system, user] messages
+           where the user message contains the summary
+        3. Returns True if compaction was successful, False otherwise
+        
+        Returns:
+            Boolean indicating whether the compaction was successful
+        """
+        return await self.agent.compact()
+
+    def add_ui_callback(self, ui_type: str, optimized: bool = True):
+        """
+        Adds a UI callback to the agent based on the type.
+        
+        Args:
+            ui_type: The type of UI callback ('rich' or 'jupyter')
+            optimized: Whether to use the optimized version (default: True for better performance)
+        """
+        if ui_type == 'rich':
+            ui_callback = RichCodeUICallback(
+                logger=self.log_manager.get_logger('tinyagent.hooks.rich_code_ui_callback') if self.log_manager else None
+            )
+            self.add_callback(ui_callback)
+        elif ui_type == 'jupyter':
+            if optimized:
+                from tinyagent.hooks.jupyter_notebook_callback import OptimizedJupyterNotebookCallback
+                ui_callback = OptimizedJupyterNotebookCallback(
+                    logger=self.log_manager.get_logger('tinyagent.hooks.jupyter_notebook_callback') if self.log_manager else None,
+                    max_visible_turns=20,    # Limit visible turns for performance
+                    max_content_length=100000,  # Limit total content
+                    enable_markdown=True,    # Keep markdown but optimized
+                    show_raw_responses=False # Show formatted responses
+                )
+            else:
+                ui_callback = JupyterNotebookCallback(
+                    logger=self.log_manager.get_logger('tinyagent.hooks.jupyter_notebook_callback') if self.log_manager else None
+                )
+            self.add_callback(ui_callback)
+        else:
+            self.log_manager.get_logger(__name__).warning(f"Unknown UI type: {ui_type}. No UI callback will be added.")
+
 
 # Example usage demonstrating both LLM tools and code tools
 async def run_example():
@@ -548,6 +805,7 @@ async def run_example():
     Code tools: Available in the Python execution environment
     """
     from tinyagent import tool
+    import os
     
     # Example LLM tool - available to the LLM for direct calling
     @tool(name="search_web", description="Search the web for information")
@@ -576,7 +834,9 @@ async def run_example():
             "sample_data": [1, 2, 3, 4, 5, 10, 15, 20]
         },
         authorized_imports=["tinyagent", "gradio", "requests", "numpy", "pandas"],  # Explicitly specify authorized imports
-        local_execution=False  # Remote execution via Modal (default)
+        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
     )
     
     # Connect to MCP servers
@@ -593,6 +853,13 @@ async def run_example():
     print(response_remote)
     print("\n" + "="*80 + "\n")
     
+    # Test the resume functionality
+    print("🔄 Testing resume functionality (continuing without new user input)")
+    resume_response = await agent_remote.resume(max_turns=3)
+    print("Resume Response:")
+    print(resume_response)
+    print("\n" + "="*80 + "\n")
+    
     # Now test with local execution
     print("🏠 Testing TinyCodeAgent with LOCAL execution")
     agent_local = TinyCodeAgent(
@@ -603,7 +870,8 @@ async def run_example():
             "sample_data": [1, 2, 3, 4, 5, 10, 15, 20]
         },
         authorized_imports=["tinyagent", "gradio", "requests"],  # More restricted imports for local execution
-        local_execution=True  # Local execution
+        local_execution=True,  # Local execution
+        check_string_obfuscation=True
     )
     
     # Connect to MCP servers
@@ -655,6 +923,44 @@ async def run_example():
     print("Local Agent Validation Response:")
     print(response2_local)
     
+    # Test shell execution
+    print("\n" + "="*80)
+    print("🐚 Testing shell execution")
+    
+    shell_prompt = "Run 'ls -la' to list files in the current directory."
+    
+    response_shell = await agent_remote.run(shell_prompt)
+    print("Shell Execution Response:")
+    print(response_shell)
+    
+    # Test default working directory functionality
+    print("\n" + "="*80)
+    print("🏠 Testing default working directory functionality")
+    
+    # Set a custom default working directory
+    custom_dir = os.path.expanduser("~")  # Use home directory as an example
+    agent_remote.set_default_workdir(custom_dir)
+    print(f"Set default working directory to: {custom_dir}")
+    
+    # Create a new directory for testing
+    test_dir = os.path.join(os.getcwd(), "test_workdir")
+    print(f"Setting default working directory with auto-creation: {test_dir}")
+    agent_remote.set_default_workdir(test_dir, create_if_not_exists=True)
+    
+    # Run shell command without specifying workdir - should use the default
+    shell_prompt_default_dir = "Run 'pwd' to show the current working directory."
+    
+    response_shell_default = await agent_remote.run(shell_prompt_default_dir)
+    print("Shell Execution with Default Working Directory:")
+    print(response_shell_default)
+    
+    # Run shell command with explicit workdir - should override the default
+    shell_prompt_explicit_dir = "Run 'pwd' in the /tmp directory."
+    
+    response_shell_explicit = await agent_remote.run(shell_prompt_explicit_dir)
+    print("Shell Execution with Explicit Working Directory:")
+    print(response_shell_explicit)
+    
     await agent_remote.close()
     await agent_local.close()
 

tinyagent/code_agent/utils.py

@@ -1,7 +1,10 @@
 import sys
 import cloudpickle
+import subprocess
+import os
 from typing import Dict, Any, List
 from .safety import validate_code_safety, function_safety_context
+import shlex
 
 
 def clean_response(resp: Dict[str, Any]) -> Dict[str, Any]:
@@ -41,6 +44,96 @@ def make_session_blob(ns: dict) -> bytes:
     return cloudpickle.dumps(clean)
 
 
+def _run_shell(
+    command: List[str],
+    timeout: int = 10,
+    workdir: str = None
+) -> Dict[str, Any]:
+    """
+    Execute a shell command securely with proper timeout and error handling.
+    
+    Args:
+        command: List of command parts to execute
+        timeout: Maximum execution time in seconds
+        workdir: Working directory for command execution
+        
+    Returns:
+        Dictionary containing execution results with keys:
+        - stdout: stdout from the execution
+        - stderr: stderr from the execution
+        - exit_code: exit code from the command
+    """
+    try:
+        # Set working directory if provided
+        cwd = os.path.expanduser(workdir) if workdir else None
+        
+        # Check if this is a command that needs bash -c wrapping
+        if len(command) > 0:
+            # If the command already uses bash -c, use it directly
+            if command[0] == "bash" and len(command) >= 3 and command[1] in ["-c", "-lc"]:
+                process = subprocess.run(
+                    command,
+                    shell=False,  # No need for shell=True as we're explicitly using bash -c
+                    capture_output=True,
+                    text=True,
+                    timeout=timeout,
+                    cwd=cwd,
+                    check=False
+                )
+            else:
+                # For all other commands, wrap in bash -c to handle shell operators
+                # and properly quote arguments that need quoting
+                
+                # Shell operators that should not be quoted
+                shell_operators = ['|', '&&', '||', '>', '<', '>>', '<<', ';']
+                
+                # Quote each part that needs quoting
+                quoted_parts = []
+                for part in command:
+                    if part in shell_operators:
+                        # Don't quote shell operators
+                        quoted_parts.append(part)
+                    else:
+                        # Use shlex.quote to properly escape special characters
+                        quoted_parts.append(shlex.quote(part))
+                
+                shell_command = " ".join(quoted_parts)
+                process = subprocess.run(
+                    ["bash", "-c", shell_command],
+                    shell=False,  # Using explicit bash -c instead of shell=True
+                    capture_output=True,
+                    text=True,
+                    timeout=timeout,
+                    cwd=cwd,
+                    check=False
+                )
+        else:
+            # Empty command
+            return {
+                "stdout": "",
+                "stderr": "Empty command",
+                "exit_code": 1
+            }
+        
+        return {
+            "stdout": process.stdout,
+            "stderr": process.stderr,
+            "exit_code": process.returncode
+        }
+    except subprocess.TimeoutExpired:
+        return {
+            "stdout": "",
+            "stderr": f"Command timed out after {timeout} seconds",
+            "exit_code": 124  # Standard timeout exit code
+        }
+    except Exception as e:
+        return {
+            "stdout": "",
+            "stderr": f"Error executing command: {str(e)}",
+            "exit_code": 1
+        }
+
+
 def _run_python(
     code: str,
     globals_dict: Dict[str, Any] | None = None,
@@ -48,6 +141,7 @@ def _run_python(
     authorized_imports: List[str] | None = None,
     authorized_functions: List[str] | None = None,
     trusted_code: bool = False,
+    check_string_obfuscation: bool = True,
 ):
     """
     Execute Python code in a controlled environment with proper error handling.
@@ -59,6 +153,7 @@ def _run_python(
         authorized_imports: List of authorized imports that user code may access. Wildcards (e.g. "numpy.*") are supported. A value of None disables the allow-list and only blocks dangerous modules.
         authorized_functions: List of authorized dangerous functions that user code may access. A value of None disables the allow-list and blocks all dangerous functions.
         trusted_code: If True, skip security checks. Should only be used for framework code, tools, or default executed code.
+        check_string_obfuscation: If True (default), check for string obfuscation techniques. Set to False to allow legitimate use of base64 encoding and other string manipulations.
         
     Returns:
         Dictionary containing execution results
@@ -74,7 +169,8 @@ def _run_python(
     # 1. Static safety analysis – refuse code containing dangerous imports or functions
     # ------------------------------------------------------------------
     validate_code_safety(code, authorized_imports=authorized_imports, 
-                        authorized_functions=authorized_functions, trusted_code=trusted_code)
+                        authorized_functions=authorized_functions, trusted_code=trusted_code,
+                        check_string_obfuscation=check_string_obfuscation)
 
     # Make copies to avoid mutating the original parameters
     globals_dict = globals_dict or {}
@@ -116,22 +212,34 @@ def _run_python(
     #updated_globals['print'] = custom_print
     
     # Parse the code
-    tree = ast.parse(code, mode="exec")
-    compiled = compile(tree, filename="<ast>", mode="exec")
+    try:
+        tree = ast.parse(code, mode="exec")
+        compiled = compile(tree, filename="<ast>", mode="exec")
+    except SyntaxError as e:
+        # Return syntax error without executing
+        return {
+            "printed_output": "", 
+            "return_value": None, 
+            "stderr": "", 
+            "error_traceback": f"Syntax error: {str(e)}",
+            "updated_globals": updated_globals,
+            "updated_locals": updated_locals
+        }
+    
     stdout_buf = io.StringIO()
     stderr_buf = io.StringIO()   
     # Execute with exception handling
     error_traceback = None
     output = None
 
+    # Merge all variables into globals to avoid scoping issues with generator expressions
+    # When exec() is called with both globals and locals, generator expressions can't
+    # access local variables. By using only globals, everything runs in global scope.
+    merged_globals = updated_globals.copy()
+    merged_globals.update(updated_locals)
+
     with contextlib.redirect_stdout(stdout_buf), contextlib.redirect_stderr(stderr_buf):
         try:
-            # Merge all variables into globals to avoid scoping issues with generator expressions
-            # When exec() is called with both globals and locals, generator expressions can't
-            # access local variables. By using only globals, everything runs in global scope.
-            merged_globals = updated_globals.copy()
-            merged_globals.update(updated_locals)
-            
             # Add 'exec' to authorized_functions for internal use
             internal_authorized_functions = ['exec','eval']
             if authorized_functions is not None and not isinstance(authorized_functions, bool):
@@ -152,6 +260,18 @@ def _run_python(
         except Exception:
             # Capture the full traceback as a string
             error_traceback = traceback.format_exc()
+            
+            # CRITICAL FIX: Even when an exception occurs, we need to update the globals and locals
+            # with any variables that were successfully created/modified before the exception
+            for key, value in merged_globals.items():
+                # Skip special variables and modules
+                if key.startswith('__') or key in ['builtins', 'traceback', 'contextlib', 'io', 'ast', 'sys']:
+                    continue
+                    
+                # Update both dictionaries with the current state
+                if key in updated_locals or key not in updated_globals:
+                    updated_locals[key] = value
+                updated_globals[key] = value
 
     # Join all captured output
     #printed_output = ''.join(output_buffer)  

tinyagent/hooks/__init__.py

@@ -2,4 +2,6 @@
 from .rich_ui_callback import RichUICallback
 from .rich_code_ui_callback import RichCodeUICallback
 from .logging_manager import LoggingManager
-__all__ = ["RichUICallback", "RichCodeUICallback", "LoggingManager"]
+from .token_tracker import TokenTracker, UsageStats, create_token_tracker
+
+__all__ = ["RichUICallback", "RichCodeUICallback", "LoggingManager", "TokenTracker", "UsageStats", "create_token_tracker"]