Building an AI Agent's Code Execution Environment

AI Code Execution Environment

Want to give your AI agent the power to write and execute code safely? Here's how to build a secure Python execution environment in under 30 minutes using Jupyter Kernel Gateway.

Why This Matters

As AI agents become more capable, they need secure environments to execute code. Traditional approaches like subprocess calls or eval() are dangerous and limited. Jupyter Kernel Gateway offers a production-ready solution that:

  • Provides isolated Python environments for code execution
  • Supports real-time output streaming
  • Handles proper cleanup of resources
  • Enables controlled access to system resources
  • Works with any AI agent or application through a simple HTTP API

The Solution: A Three-Part System

  1. Kernel Gateway Server: Manages Python kernels and handles code execution requests
  2. Client Library: Provides a clean interface for your AI agent to send code
  3. Security Layer: Controls what the agent can and cannot do

Let's build this step by step.

Step 1: Setting Up the Kernel Gateway

First, we'll create a Docker container running Jupyter Kernel Gateway. This isolates the execution environment and makes deployment easy:

FROM jupyter/base-notebook:python-3.10
 
USER root
# Install system dependencies if needed
RUN apt-get update && apt-get install -y --no-install-recommends     build-essential     && rm -rf /var/lib/apt/lists/*
 
USER ${NB_UID}
# Install Jupyter Kernel Gateway and other dependencies
RUN pip install --no-cache-dir jupyter_kernel_gateway==2.5.2 numpy pandas matplotlib seaborn scikit-learn
 
# Add security configurations
COPY kernel_security.py /opt/conda/lib/python3.10/site-packages/
ENV JUPYTER_KERNEL_GATEWAY_KERNEL_SECURITY_CLASS=kernel_security.KernelSecurity
 
EXPOSE 8888
CMD ["jupyter", "kernelgateway", \
     "--KernelGatewayApp.ip=0.0.0.0", \
     "--KernelGatewayApp.allow_origin=*", \
     "--KernelGatewayApp.allow_credentials=true", \
     "--KernelGatewayApp.allow_headers=*"]
docker build -t ai-kernel-gateway .
docker run -p 8888:8888 ai-kernel-gateway

Step 2: Creating the Client Interface

We'll create a Python client that makes it easy for your AI agent to:

  • Start new kernel sessions
  • Execute code and get results
  • Handle streaming output
  • Manage kernel lifecycle

The client handles all the WebSocket communication and message formatting, providing a clean interface:

from kernel_client import KernelClient
 
client = KernelClient("http://localhost:8888")
result = await client.execute_code("print('Hello from AI!')")
print(result.output)  # "Hello from AI!"

Step 3: Adding Security Controls

The security layer lets you:

  • Whitelist allowed Python modules
  • Set memory and CPU limits
  • Control filesystem access
  • Monitor execution time
from jupyter_kernel_gateway.gateway.handlers import IPythonHandler
import sys
from importlib.util import find_spec
 
class KernelSecurity:
    """Security layer for controlling kernel execution environment"""
    
    def __init__(self):
        self.allowed_modules = {
            'numpy', 'pandas', 'matplotlib', 'seaborn', 
            'sklearn', 'math', 'statistics', 'datetime'
        }
        self.blocked_modules = {'os', 'subprocess', 'sys', 'socket'}
    
    def check_code(self, code):
        """Analyze code for potential security issues"""
        # Basic static analysis
        if any(blocked in code for blocked in ['eval', 'exec', '__import__']):
            raise SecurityError("Prohibited functions detected")
        
        # Check imports
        import_lines = [line for line in code.split('\n') 
                       if line.strip().startswith(('import ', 'from '))]
        
        for line in import_lines:
            module = line.split()[1].split('.')[0]
            if module in self.blocked_modules:
                raise SecurityError(f"Module {module} is not allowed")
            if module not in self.allowed_modules:
                if not self._is_stdlib_module(module):
                    raise SecurityError(f"Module {module} is not whitelisted")
    
    def _is_stdlib_module(self, module_name):
        """Check if a module is part of Python's standard library"""
        spec = find_spec(module_name)
        if spec is None:
            return False
        return 'site-packages' not in str(spec.origin)
 
class SecurityError(Exception):
    pass

Complete Example

Here's a real-world example of an AI agent using this system to analyze data:

# AI agent code
async def analyze_data(data_path):
    code = """
    import pandas as pd
    df = pd.read_csv('data.csv')
    mean_value = df['value'].mean()
    print(f"Average value: {mean_value}")
    """
    result = await kernel_client.execute_code(code)
    return result.output

Getting Started

  1. Clone the repository with the complete implementation
  2. Follow the setup instructions in README.md
  3. Run the example code to test your environment

The full code includes:

  • Dockerfile and setup scripts
  • Python client library
  • Security configuration
  • Example AI agent integration
  • Comprehensive tests

Next Steps

Once you have the basic environment running, you can:

  1. Add custom security rules
  2. Integrate with your AI agent
  3. Deploy to production
  4. Add support for additional languages

Code Repository

https://github.com/ColeMurray/ai-code-agent-environment/tree/main

I’d love to hear from you if you found this post helpful or have any questions. Reach out