From 7d9efc5a38ceff5059c1d0b95910e3ad0ec8b52a Mon Sep 17 00:00:00 2001 From: Hans Aschauer Date: Wed, 4 Mar 2026 22:54:56 +0100 Subject: [PATCH] feat: Implement session management and secure container lifecycle management - Added session management for stateful execution in `sessions.py`, including session creation, state documentation, and cleanup of idle sessions. - Introduced `SecureContainerManager` in `containers.py` for managing container lifecycle with security enforcement, including creation, starting, stopping, and removal of containers. - Updated server initialization to use the new session manager. - Enhanced container configuration to skip resource limits for very high values, indicating no enforcement. - Improved logging capabilities for container operations in the audit logger. - Refactored Jupyter backend to integrate with the new session management and resource limits handling. --- docs/architecture2.md | 1448 +++++++++++++++++++ src/mcp_forge/adapters/executor_adapter.py | 54 +- src/mcp_forge/execution/jupyter/backend.py | 263 ++++ src/mcp_forge/execution/jupyter/kernel.py | 631 ++++++++ src/mcp_forge/execution/jupyter/sessions.py | 435 ++++++ src/mcp_forge/podman/containers.py | 508 +++++++ src/mcp_forge/server/server.py | 2 +- src/pod_executor/containers/manager.py | 14 +- src/pod_executor/jupyter/backend.py | 12 +- src/pod_executor/security/audit.py | 43 + 10 files changed, 3384 insertions(+), 26 deletions(-) create mode 100644 docs/architecture2.md create mode 100644 src/mcp_forge/execution/jupyter/backend.py create mode 100644 src/mcp_forge/execution/jupyter/kernel.py create mode 100644 src/mcp_forge/execution/jupyter/sessions.py create mode 100644 src/mcp_forge/podman/containers.py diff --git a/docs/architecture2.md b/docs/architecture2.md new file mode 100644 index 0000000..a88adb1 --- /dev/null +++ b/docs/architecture2.md @@ -0,0 +1,1448 @@ +# MCP-Forge Architecture Document + +## Overview + +MCP-Forge is an MCP (Model Context Protocol) server that provides intelligent code execution capabilities for AI agents. It enables agents to execute Python code in isolated containers while having access to other MCP tools, optimizing data processing by keeping large datasets in the execution environment rather than passing them through the LLM context. + +## Core Concept + +**Problem:** When AI agents use MCP tools that return large amounts of data, the data must flow through the LLM context window, which is: +- Token-expensive +- Slow +- Limited by context window size + +**Solution:** MCP-Forge provides a code execution environment where: +- MCP tools are available as Python functions +- Agents can write code to process data locally +- Only relevant results flow back to the agent +- The agent already knows tool signatures from its own configuration + +## Architecture + +``` +AI Agent (e.g., GitHub Copilot, Claude) + │ + ├─ MCP Protocol + │ +MCP-Forge Server + │ + ├─ MCP Resources (Discovery) + │ ├─ tools/available + │ ├─ sessions/{id}/state + │ ├─ sessions/{id}/variables + │ └─ environments/list + │ + ├─ MCP Tools (Execution) + │ ├─ execute_python() + │ ├─ document_state() + │ └─ build_custom_environment() + │ + ├─ Execution Backends + │ ├─ Simple Backend (stateless) + │ └─ Jupyter Backend (stateful) + │ + ├─ Custom Environment Builder + │ ├─ UV Package Installer + │ ├─ Build Cache Manager + │ ├─ Security Validator + │ └─ Template Library + │ + └─ Podman API (restricted) + │ + ├─ Pre-built Images + │ ├─ mcp-forge/python:3.11 + │ ├─ mcp-forge/python:3.12 + │ └─ mcp-forge/jupyter:latest + │ + ├─ Custom User Images + │ └─ mcp-forge/custom:{user-env-name} + │ + └─ Isolated Execution Containers + ├─ Python Runtime + ├─ MCP Client (injected tools) + └─ Session Volumes +``` + +## Components + +### 1. MCP Server Interface + +#### Resources + +**`mcp://forge/tools/available`** +- Returns: List of MCP tool names available in execution environment +- Purpose: Agent discovers what tools it can use in generated code +- Example: `["github_search_repos", "filesystem_read", "sqlite_query"]` + +**`mcp://forge/tools/{tool_name}/schema`** (optional) +- Returns: Full JSON schema for a specific tool +- Purpose: Fallback if agent doesn't have tool definition +- Note: Typically not needed as agent has tools in its own config + +**`mcp://forge/sessions/{session_id}/state`** +- Returns: Documented state for a session +- Structure: + ```json + { + "documented_variables": { + "df": "Customer data, 1000 rows, columns: id, name, purchase_date, amount", + "model": "Trained RandomForest, accuracy 0.87" + }, + "note": "Preprocessing complete, ready for analysis", + "last_updated": "2026-02-06T10:30:00Z", + "all_variables": ["df", "model", "temp", "result"] + } + ``` + +**`mcp://forge/sessions/{session_id}/variables`** +- Returns: List of all variables in session namespace (kernel introspection) +- Purpose: Quick check of what exists in session + +**`mcp://forge/environment/info`** (optional) +- Returns: Python version, installed packages, system info +- Purpose: Agent can verify environment capabilities + +**`mcp://forge/environments/list`** +- Returns: List of available custom environments and templates +- Structure: + ```json + { + "custom": [ + { + "name": "my-ml-env", + "image": "mcp-forge/custom:my-ml-env", + "created": "2026-02-06T10:00:00Z", + "packages": ["numpy==1.24.0", "pandas==2.0.0"], + "size_mb": 1234 + } + ], + "templates": [ + { + "name": "ml-basic", + "description": "Basic ML stack", + "packages": ["numpy", "pandas", "scikit-learn"] + }, + { + "name": "data-science", + "description": "Data science stack", + "packages": ["numpy", "pandas", "matplotlib", "seaborn"] + } + ] + } + ``` + +**`mcp://forge/audit/operations`** (optional) +- Returns: Audit log of container operations +- Purpose: Security monitoring, debugging + +#### Tools + +**`execute_python`** + +Execute Python code in an isolated container with MCP tools available. + +Parameters: +```json +{ + "code": "string (required) - Python code to execute", + "mcp_tools": "array (optional) - List of MCP tool names to inject", + "session_id": "string (optional) - Session ID for stateful execution (null = stateless)", + "backend": "string (optional) - 'simple' (default) or 'jupyter'", + "timeout": "integer (optional) - Max execution time in seconds (default: 300)", + "volumes": "object (optional) - Volume mount configuration", + "custom_image": "string (optional) - Custom environment name to use", + "environment": "string (optional) - Template environment name (e.g., 'ml-basic', 'data-science')" +} +``` + +Returns: +```json +{ + "success": true, + "stdout": "string - Standard output", + "stderr": "string - Standard error", + "result": "any - Return value of last expression", + "available_tools": ["list of tools that were injected"], + "execution_time": 1.23, + "session_id": "abc123 (if stateful)" +} +``` + +Example: +```python +# Agent calls: +execute_python( + code=""" +repos = github_search_repos(query="MCP servers", max_results=100) +high_quality = [r for r in repos['items'] if r['stars'] > 50] +print(f"Found {len(high_quality)} high-quality repos") +high_quality[:10] # Return top 10 +""", + mcp_tools=["github_search_repos"], + session_id=None # stateless +) +``` + +**`document_state`** + +Document important variables in a stateful session for later retrieval. + +Parameters: +```json +{ + "session_id": "string (required) - Session to document", + "variables": "object (required) - Variable name -> description mapping", + "note": "string (optional) - General note about session state", + "clear": "boolean (optional) - Clear existing documentation (default: false)" +} +``` + +Returns: +```json +{ + "success": true, + "documented_count": 3, + "session_id": "abc123" +} +``` + +Example: +```python +document_state( + session_id="abc123", + variables={ + "df": "Customer purchase data, 1000 rows, preprocessed and cleaned", + "model": "Trained RandomForest classifier, 87% accuracy on test set" + }, + note="Ready for prediction phase" +) +``` + +**`build_custom_environment`** + +Build a custom container image with specified Python packages. This is the **only** way to install additional packages - `pip install` is **not** allowed during code execution for security reasons. + +Parameters: +```json +{ + "name": "string (required) - Name for the custom environment (alphanumeric + hyphens)", + "base_image": "string (optional) - Base image to build from (default: 'python:3.11')", + "packages": "array (required) - List of Python package specifications", + "python_version": "string (optional) - Python version: '3.11', '3.12' (default: '3.11')", + "description": "string (optional) - Description of this environment" +} +``` + +Returns: +```json +{ + "success": true, + "image_name": "mcp-forge/custom:my-ml-env", + "image_id": "sha256:abc123...", + "build_time": 45.2, + "installed_packages": ["numpy==1.24.0", "pandas==2.0.0", "scikit-learn==1.3.0"], + "cache_hit": false +} +``` + +Example: +```python +# Agent builds custom environment for ML work +build_custom_environment( + name="my-ml-env", + packages=[ + "numpy>=1.24.0", + "pandas>=2.0.0", + "scikit-learn>=1.3.0", + "matplotlib>=3.7.0", + "seaborn>=0.12.0" + ], + description="Machine learning environment with common libraries" +) + +# Later, use the custom environment +execute_python( + code="import pandas as pd; df = pd.read_csv('data.csv'); ...", + custom_image="mcp-forge/custom:my-ml-env", + session_id="ml-analysis" +) +``` + +### 2. Execution Backends + +#### Simple Backend (Default) + +**Purpose:** Stateless, single-shot code execution + +**Implementation:** +- Spawn Podman container with Python image +- Inject code and MCP tool functions into namespace +- Execute via `python -c` +- Capture stdout, stderr, result +- Destroy container + +**Characteristics:** +- Fast startup +- No state between calls +- Ideal for data transformation tasks +- Lower memory footprint + +**Use cases:** +- Data processing and filtering +- One-off computations +- Format conversions +- Quick analyses + +#### Jupyter Backend + +**Purpose:** Stateful, multi-step workflows + +**Implementation:** +- Spawn Podman container with IPython kernel +- Keep kernel running for session lifetime +- Execute code cells via Jupyter protocol (ZMQ) +- Maintain namespace between executions +- Support rich output formats + +**Characteristics:** +- State persists between calls +- Variable persistence +- Interactive workflow support +- Higher resource usage + +**Use cases:** +- Multi-step data analysis +- Iterative model training +- Building up complex state +- Exploratory workflows + +**Session Management:** +- Sessions identified by unique ID +- Automatic timeout after inactivity (configurable, default: 1 hour) +- Manual cleanup via session deletion +- Resource limits per session + +### 3. Custom Environment Building + +#### Design Philosophy + +**Security Principle:** `pip install` is **NOT** allowed during code execution. All package installations must happen during a controlled build process. + +**Why:** +- Prevents malicious package installation during execution +- Ensures reproducible environments +- Allows security scanning of dependencies +- Enables caching for performance +- Provides audit trail of what's installed + +#### Build System: UV-Based + +**Technology Choice:** [UV](https://github.com/astral-sh/uv) - ultra-fast Python package installer + +**Benefits:** +- 10-100x faster than pip +- Built-in caching +- Deterministic dependency resolution +- Lockfile support for reproducibility +- Compatible with pip package specifications + +#### Implementation + +**Build Process:** + +```python +def build_custom_environment(name, packages, base_image="python:3.11"): + """ + Build a custom container image with specified packages using UV. + + Steps: + 1. Generate Containerfile with UV installation + 2. Create requirements specification + 3. Build image with layer caching + 4. Validate and tag image + 5. Store metadata for reuse + """ + + # Generate Containerfile + containerfile = f""" +FROM {base_image} + +# Install UV (cached layer) +RUN curl -LsSf https://astral.sh/uv/install.sh | sh +ENV PATH="/root/.cargo/bin:$PATH" + +# Non-root user +RUN useradd -m -u 1000 forge +USER forge + +# Copy requirements (cache-friendly) +COPY requirements.txt /tmp/requirements.txt + +# Install packages with UV (uses cache) +RUN uv pip install --system -r /tmp/requirements.txt + +# MCP bridge client +COPY mcp_tools.py /usr/local/lib/python3.11/site-packages/ + +WORKDIR /workspace +""" + + # Build with Podman + build_context = create_build_context(containerfile, packages) + image_id = podman.images.build( + path=build_context, + tag=f"mcp-forge/custom:{name}", + cache_from=["mcp-forge/cache:uv-base"], # UV cache layer + buildargs={ + "BUILDKIT_INLINE_CACHE": "1" + } + ) + + return image_id +``` + +#### Package Cache Strategy + +**Multi-Layer Caching:** + +``` +Layer 1: Base Python image (rarely changes) + └─ python:3.11-slim + +Layer 2: UV installation (stable) + └─ UV binary + dependencies + +Layer 3: Common packages (pre-cached) + └─ numpy, pandas, requests, etc. + +Layer 4: User packages (specific to build) + └─ Custom requirements +``` + +**UV Cache Directory:** +- Mounted from host: `/var/cache/mcp-forge/uv` +- Persisted between builds +- Dramatically speeds up builds with similar dependencies + +**Cache Invalidation:** +- Requirements hash changes → rebuild user layer only +- Base image updated → rebuild from base +- UV version updated → rebuild from UV layer + +#### Pre-built Environment Templates + +**Common Templates:** + +```yaml +templates: + ml-basic: + description: "Basic ML stack" + packages: + - "numpy>=1.24.0" + - "pandas>=2.0.0" + - "scikit-learn>=1.3.0" + + ml-advanced: + description: "Advanced ML with deep learning" + packages: + - "numpy>=1.24.0" + - "pandas>=2.0.0" + - "scikit-learn>=1.3.0" + - "torch>=2.0.0" + - "transformers>=4.30.0" + + data-science: + description: "Data science stack" + packages: + - "numpy>=1.24.0" + - "pandas>=2.0.0" + - "matplotlib>=3.7.0" + - "seaborn>=0.12.0" + - "plotly>=5.14.0" + - "jupyter>=1.0.0" + + web-scraping: + description: "Web scraping tools" + packages: + - "requests>=2.31.0" + - "beautifulsoup4>=4.12.0" + - "lxml>=4.9.0" + - "selenium>=4.10.0" +``` + +**Usage:** + +```python +# Use a template +build_custom_environment( + name="my-ml-env", + template="ml-basic", + additional_packages=["xgboost>=2.0.0"] +) + +# Or build from scratch +build_custom_environment( + name="custom-env", + packages=["specific-package==1.0.0"] +) +``` + +#### Build Validation + +**Post-Build Checks:** + +```python +def validate_built_environment(image_id): + """ + Validate that the built environment is safe and functional. + + Checks: + 1. All requested packages are installed + 2. No malicious packages (against allowlist/blocklist) + 3. Image size is within limits + 4. Security scan passes (trivy/grype) + 5. Python imports work + """ + + # Import test + test_code = """ +import sys +import json +installed = [(pkg.key, pkg.version) for pkg in __import__('pkg_resources').working_set] +print(json.dumps(installed)) +""" + + result = run_in_container(image_id, test_code) + installed_packages = json.loads(result.stdout) + + # Security scan + scan_result = security_scan_image(image_id) + if scan_result.critical_vulns > 0: + raise SecurityError(f"Image has {scan_result.critical_vulns} critical vulnerabilities") + + return { + "installed_packages": installed_packages, + "security_scan": scan_result, + "valid": True + } +``` + +#### Package Allowlist/Blocklist + +**Security Control:** + +```yaml +packages: + # Automatically allowed (common, vetted) + allowlist: + - numpy + - pandas + - scikit-learn + - matplotlib + - seaborn + - requests + - beautifulsoup4 + # ... many more + + # Explicitly forbidden + blocklist: + - "os-sys" # Known malicious + - "cryptography-backdoor" + # Packages with known vulnerabilities + + # Requires manual approval + manual_approval: + - "*crypto*" # Cryptographic packages + - "*network*" # Network access packages + - "*subprocess*" # Process spawning +``` + +**Validation:** + +```python +def validate_package_list(packages): + """Check packages against allowlist/blocklist before building.""" + for pkg in packages: + pkg_name = pkg.split("==")[0].split(">=")[0].split("<=")[0] + + if pkg_name in BLOCKLIST: + raise SecurityError(f"Package {pkg_name} is blocked") + + if not in_allowlist(pkg_name) and requires_approval(pkg_name): + raise ApprovalRequiredError(f"Package {pkg_name} requires manual approval") +``` + +#### Build Resource Limits + +**Prevent Build Abuse:** + +```python +BUILD_LIMITS = { + "max_packages": 50, # Max packages per build + "max_build_time": 600, # 10 minutes + "max_image_size": "2GB", # Final image size + "max_concurrent_builds": 3, # Per user + "build_rate_limit": { + "requests": 10, + "period": 3600 # 10 builds per hour + } +} +``` + +#### Environment Lifecycle + +**Management:** + +```python +# List user's custom environments +list_custom_environments() → [ + { + "name": "my-ml-env", + "image": "mcp-forge/custom:my-ml-env", + "created": "2026-02-06T10:00:00Z", + "size_mb": 1234, + "packages": ["numpy==1.24.0", "pandas==2.0.0"], + "last_used": "2026-02-06T11:30:00Z" + } +] + +# Delete unused environment +delete_custom_environment(name="old-env") + +# Rebuild environment (e.g., after base image update) +rebuild_custom_environment(name="my-ml-env") +``` + +**Auto-cleanup:** +- Environments unused for 30 days → archived +- Archived for 90 days → deleted +- User notification before deletion + +#### execute_python Integration + +**Using Custom Environments:** + +```python +# Updated execute_python parameters +{ + "code": "...", + "custom_image": "string (optional) - Custom environment name or image tag", + "environment": "string (optional) - Template name (alternative to custom_image)", + ... +} + +# Examples: +execute_python( + code="import numpy as np; ...", + custom_image="my-ml-env" # Use user's custom environment +) + +execute_python( + code="import pandas as pd; ...", + environment="data-science" # Use pre-built template +) +``` + +### 4. Container Runtime (Podman) + +#### Why Podman? + +- **Rootless by design:** No daemon running as root +- **Docker-compatible API:** Easy integration +- **Better security defaults:** No privileged operations needed +- **Daemonless:** Containers are child processes + +#### Deployment Modes + +**Development:** +```bash +# Run mcp-forge directly on host +mcp-forge serve --podman-socket /run/user/1000/podman/podman.sock +``` + +**Production:** +```yaml +# Run mcp-forge in container with socket mount +services: + mcp-forge: + image: mcp-forge:latest + volumes: + - /run/user/1000/podman/podman.sock:/run/podman/podman.sock:ro + environment: + - PODMAN_SOCKET=/run/podman/podman.sock +``` + +#### Security: Allowlist-Based Container Operations + +**Allowed Operations:** + +```python +ALLOWED_OPERATIONS = { + "container.create": { + "allowed_images": [ + "mcp-forge/python:3.11", + "mcp-forge/python:3.12", + "mcp-forge/jupyter:latest" + ], + "forbidden_params": [ + "privileged", + "cap_add", + "devices", + "pid_mode", + "ipc_mode" + ], + "required_params": { + "network_mode": "none", # Or restricted network + "read_only": True, # Filesystem read-only except volumes + "memory_limit": "512m", + "cpu_quota": 50000, # 50% of one CPU + "security_opt": ["no-new-privileges"], + "user": "1000:1000" # Non-root user + } + }, + "container.start": { + "session_containers_only": True # Only containers we created + }, + "container.stop": { + "session_containers_only": True + }, + "container.remove": { + "session_containers_only": True + }, + "container.logs": { + "session_containers_only": True + } +} +``` + +**Forbidden Operations:** + +```python +FORBIDDEN_OPERATIONS = [ + "container.exec", # No direct shell access + # Note: image.build IS allowed, but ONLY through build_custom_environment tool + # with validation, security scanning, and package allowlists + "image.pull", # No arbitrary image pulling (only pre-approved images) + "volume.create", # Only pre-configured volumes + "network.create", # No custom networks + "system.prune", # No system-level operations +] + +# Additionally forbidden within execution containers: +EXECUTION_RESTRICTIONS = [ + "subprocess.run(['pip', 'install', ...])", # No pip install during execution + "subprocess.run(['apt', 'install', ...])", # No system package installs + "import os; os.system('...')", # Restricted system calls +] +``` + +#### Volume Management + +**Allowed Volume Patterns:** + +```python +ALLOWED_VOLUME_PATTERNS = [ + "/mcp-forge/sessions/{session_id}/*", # Per-session data + "/mcp-forge/shared/readonly/*", # Shared read-only data + "/mcp-forge/uploads/{session_id}/*" # User uploads +] + +FORBIDDEN_MOUNT_PATHS = [ + "/", + "/etc", + "/var/run/docker.sock", + "/var/run/podman/podman.sock", + "/sys", + "/proc" +] +``` + +#### Resource Limits (Enforced) + +```python +RESOURCE_LIMITS = { + "memory": { + "default": "512m", + "max": "2g" + }, + "cpu_quota": { + "default": 50000, # 50% of one CPU + "max": 100000 # 100% of one CPU + }, + "timeout": { + "default": 300, # 5 minutes + "max": 1800 # 30 minutes + }, + "pids_limit": 100, + "storage": { + "default": "1g", + "max": "10g" + } +} +``` + +### 5. MCP Tool Injection + +#### Tool Discovery Flow + +1. **Agent reads resource:** `mcp://forge/tools/available` +2. **MCP-Forge returns:** `["github_search_repos", "filesystem_read"]` +3. **Agent knows signatures:** Already configured in agent's tool config +4. **Agent generates code:** Using these tool names as Python functions + +#### Tool Injection Implementation + +**At container startup:** + +```python +# MCP-Forge generates injection code +def inject_mcp_tools(container, mcp_client, tool_names): + """ + Inject MCP tools as Python functions into container namespace. + + Args: + container: Podman container instance + mcp_client: Connected MCP client with available tools + tool_names: List of tool names to inject + """ + + injection_code = """ +# MCP Tool Functions (auto-injected) +import json +from typing import Any + +def _mcp_call(tool_name: str, **kwargs) -> Any: + '''Internal: Call MCP tool via bridge''' + # This communicates with MCP-Forge server which forwards to MCP client + import socket + s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + s.connect('/tmp/mcp-bridge.sock') + s.send(json.dumps({'tool': tool_name, 'params': kwargs}).encode()) + result = json.loads(s.recv(1024*1024).decode()) + s.close() + return result + +""" + + # Generate wrapper function for each tool + for tool_name in tool_names: + tool_schema = mcp_client.get_tool_schema(tool_name) + + # Generate function signature from schema + params = tool_schema.get('inputSchema', {}).get('properties', {}) + param_str = ', '.join(f"{k}: Any = None" for k in params.keys()) + + injection_code += f""" +def {tool_name}({param_str}): + ''' + {tool_schema.get('description', '')} + + Auto-generated wrapper for MCP tool: {tool_name} + ''' + kwargs = {{{', '.join(f"'{k}': {k}" for k in params.keys())}}} + kwargs = {{k: v for k, v in kwargs.items() if v is not None}} + return _mcp_call('{tool_name}', **kwargs) + +""" + + # Write to container's Python site-packages or PYTHONSTARTUP + container.write_file('/usr/local/lib/python3.11/site-packages/mcp_tools.py', + injection_code) + + # Start MCP bridge server (forwards calls to actual MCP client) + start_mcp_bridge(container, mcp_client) +``` + +**Alternative: Environment Variable Approach** + +```python +# Set PYTHONSTARTUP to auto-import tools +container.env['PYTHONSTARTUP'] = '/mcp-forge/startup.py' +``` + +#### MCP Bridge Server + +A lightweight server running in the MCP-Forge process that: +- Listens on Unix socket mounted into container +- Receives tool calls from container code +- Forwards to actual MCP client +- Returns results to container + +This keeps the MCP client logic outside the untrusted container. + +### 6. Session State Management + +#### State Lifecycle + +**Stateless Execution (session_id=None):** +``` +Request → Create Container → Execute → Return Results → Destroy Container +``` + +**Stateful Execution (session_id="abc123"):** +``` +First Request: + Request → Create Container → Start Kernel → Execute → Return Results + Container stays alive + +Subsequent Requests: + Request → Reuse Container → Execute in Same Kernel → Return Results + +Cleanup: + Timeout or Manual → Stop Container → Destroy Container +``` + +#### State Documentation + +**Agent Workflow:** + +```python +# Step 1: Load and process data +execute_python( + code="df = load_large_dataset(); df_clean = preprocess(df)", + session_id="analysis-123", + backend="jupyter" +) + +# Step 2: Document important state +document_state( + session_id="analysis-123", + variables={ + "df": "Raw dataset, 100k rows", + "df_clean": "Cleaned dataset, 95k rows, ready for modeling" + }, + note="Data loading complete" +) + +# ... conversation continues, context window rotates ... + +# Step 3: Agent checks state after context rotation +state = read_resource("mcp://forge/sessions/analysis-123/state") +# Agent sees: df_clean exists and is ready + +# Step 4: Continue work +execute_python( + code="model = train_model(df_clean); results = evaluate(model)", + session_id="analysis-123" +) +``` + +#### State Introspection + +**Automatic tracking:** +- Kernel namespace inspection (all variables) +- Variable types and sizes +- Last execution timestamp + +**Agent-documented:** +- Semantic descriptions +- Workflow notes +- Important variables highlighted + +**Combined in resource:** +```json +{ + "all_variables": ["df", "df_clean", "model", "temp", "i", "results"], + "documented_variables": { + "df_clean": "Cleaned dataset, 95k rows, ready for modeling", + "model": "Trained model, 87% accuracy", + "results": "Evaluation metrics" + }, + "introspection": { + "df_clean": {"type": "DataFrame", "shape": [95000, 12], "memory_mb": 87}, + "model": {"type": "RandomForestClassifier", "memory_mb": 234} + }, + "note": "Ready for final predictions", + "last_updated": "2026-02-06T10:30:00Z" +} +``` + +## Implementation Considerations + +### Technology Stack + +**MCP Server:** +- Language: Python 3.11+ +- Framework: `mcp` SDK (official Python implementation) +- Async: `asyncio` for concurrent operations + +**Container Runtime:** +- Podman via `podman-py` library +- Fallback: Direct Podman CLI calls + +**Jupyter Backend:** +- `jupyter_client` for kernel management +- ZMQ for communication +- `ipykernel` in container images + +**Simple Backend:** +- Direct Python execution via `subprocess` +- Or: `python:3.11-slim` base image with `exec` entrypoint + +### Pre-built Container Images + +**Base Image: `mcp-forge/python:3.11`** + +```dockerfile +FROM python:3.11-slim + +# Non-root user +RUN useradd -m -u 1000 forge +USER forge + +# Install common packages +RUN pip install --user numpy pandas requests + +# MCP bridge client +COPY mcp_tools.py /usr/local/lib/python3.11/site-packages/ + +# Startup script +COPY startup.py /mcp-forge/startup.py + +WORKDIR /workspace +``` + +**Jupyter Image: `mcp-forge/jupyter:latest`** + +```dockerfile +FROM mcp-forge/python:3.11 + +USER root +RUN pip install ipykernel jupyter_client +USER forge + +# IPython config for security +COPY ipython_config.py /home/forge/.ipython/profile_default/ + +CMD ["python", "-m", "ipykernel_launcher", "-f", "/tmp/kernel.json"] +``` + +### Error Handling + +**Container Failures:** +- Timeout: Kill container, return timeout error +- OOM: Return memory limit error with suggestion to increase +- Crash: Return stderr and exit code + +**MCP Tool Failures:** +- Tool not available: Clear error message with available tools +- Tool call error: Return tool error to agent, don't crash execution +- Network error: Retry logic for transient failures + +**Security Violations:** +- Forbidden operation: Reject immediately, log to audit +- Resource limit exceeded: Terminate execution, clear error message +- Invalid volume mount: Reject with explanation + +### Performance Considerations + +**Container Reuse:** +- Pool of warm containers for simple backend (optional) +- Lazy cleanup of idle Jupyter sessions + +**Parallel Execution:** +- Support multiple concurrent executions +- Per-session locking for stateful operations +- Configurable max concurrent containers + +**Image Caching:** +- Pre-pull images on startup +- Periodic image updates (configurable schedule) + +**Volume Performance:** +- Use tmpfs for ephemeral data +- Persistent volumes for session data +- Cleanup strategy for old sessions + +## Security Model + +### Defense in Depth + +1. **Podman rootless:** No root daemon +2. **Allowlist enforcement:** Only permitted operations +3. **Read-only filesystem:** Except specific volumes +4. **Network isolation:** No internet by default +5. **Resource limits:** CPU, memory, storage, PIDs +6. **No privileged mode:** Ever +7. **Capability dropping:** Minimal capabilities +8. **User namespaces:** Non-root user in container +9. **Audit logging:** All operations logged +10. **Session isolation:** Each session has isolated volumes + +### Threat Model + +**Threats Mitigated:** +- Container escape → Rootless + restricted operations +- Resource exhaustion → Hard limits enforced +- Data exfiltration → Network isolation +- Privilege escalation → No privileged mode, capability restrictions +- Host filesystem access → Allowlist volumes only + +**Out of Scope:** +- Side-channel attacks (Spectre, Meltdown) +- Physical security +- Supply chain attacks on base images (use trusted registries) + +### Audit and Monitoring + +**Logged Events:** +- Container create/start/stop/remove +- Execution requests (code hash, not full code for privacy) +- Resource limit violations +- Security policy violations +- Session lifecycle events + +**Log Format:** +```json +{ + "timestamp": "2026-02-06T10:30:00Z", + "event": "container.create", + "session_id": "abc123", + "image": "mcp-forge/python:3.11", + "resources": {"memory": "512m", "cpu_quota": 50000}, + "success": true +} +``` + +## Configuration + +### Server Configuration + +**Example `mcp-forge.yaml`:** + +```yaml +server: + host: localhost + port: 3000 + podman_socket: /run/user/1000/podman/podman.sock + +execution: + default_backend: simple + default_timeout: 300 + max_timeout: 1800 + default_memory: 512m + max_memory: 2g + default_cpu_quota: 50000 + max_cpu_quota: 100000 + +images: + python_3_11: mcp-forge/python:3.11 + python_3_12: mcp-forge/python:3.12 + jupyter: mcp-forge/jupyter:latest + auto_pull: true + pull_interval: 86400 # 24 hours + +sessions: + idle_timeout: 3600 # 1 hour + max_concurrent: 10 + cleanup_interval: 300 # 5 minutes + +volumes: + base_path: /var/lib/mcp-forge + session_quota: 1g + max_session_quota: 10g + +security: + audit_log: /var/log/mcp-forge/audit.log + enforce_resource_limits: true + allow_network: false + +environment_builder: + enabled: true + uv_cache_path: /var/cache/mcp-forge/uv + max_packages_per_build: 50 + max_build_time: 600 # 10 minutes + max_image_size: 2147483648 # 2GB in bytes + max_concurrent_builds: 3 + build_rate_limit: + requests: 10 + period: 3600 # 10 builds per hour per user + + # Package security + package_validation: + use_allowlist: true + allowlist_path: /etc/mcp-forge/package-allowlist.txt + blocklist_path: /etc/mcp-forge/package-blocklist.txt + require_approval_patterns: + - "*crypto*" + - "*network*" + - "*subprocess*" + + # Environment lifecycle + auto_cleanup: + enabled: true + archive_after_days: 30 + delete_after_days: 90 + + # Pre-built templates + templates: + ml-basic: + description: "Basic ML stack" + packages: + - "numpy>=1.24.0" + - "pandas>=2.0.0" + - "scikit-learn>=1.3.0" + + data-science: + description: "Data science stack" + packages: + - "numpy>=1.24.0" + - "pandas>=2.0.0" + - "matplotlib>=3.7.0" + - "seaborn>=0.12.0" + - "plotly>=5.14.0" + + web-scraping: + description: "Web scraping tools" + packages: + - "requests>=2.31.0" + - "beautifulsoup4>=4.12.0" + - "lxml>=4.9.0" + +mcp_tools: + # MCP servers to connect to and expose + github: + command: "npx" + args: ["-y", "@modelcontextprotocol/server-github"] + env: + GITHUB_TOKEN: ${GITHUB_TOKEN} + + filesystem: + command: "npx" + args: ["-y", "@modelcontextprotocol/server-filesystem", "/allowed/path"] +``` + +### Environment Variables + +```bash +# Podman socket location +MCP_FORGE_PODMAN_SOCKET=/run/user/1000/podman/podman.sock + +# Base path for volumes +MCP_FORGE_VOLUMES_PATH=/var/lib/mcp-forge + +# UV cache for fast package installation +MCP_FORGE_UV_CACHE_PATH=/var/cache/mcp-forge/uv + +# Security +MCP_FORGE_AUDIT_LOG=/var/log/mcp-forge/audit.log + +# MCP tool credentials +GITHUB_TOKEN=ghp_xxxxx +``` + +## API Examples + +### Example 1: Simple Data Processing + +```python +# Agent discovers available tools +tools = read_resource("mcp://forge/tools/available") +# Returns: ["github_search_repos", "filesystem_read"] + +# Agent generates and executes code +result = execute_python( + code=""" +repos = github_search_repos(query="machine learning", max_results=100) +python_repos = [r for r in repos['items'] if r['language'] == 'Python'] +top_10 = sorted(python_repos, key=lambda x: x['stars'], reverse=True)[:10] + +# Return summary +[{ + 'name': r['name'], + 'stars': r['stars'], + 'url': r['html_url'] +} for r in top_10] +""", + mcp_tools=["github_search_repos"], + session_id=None +) + +print(result['result']) +# Returns top 10 Python ML repos +``` + +### Example 2: Building and Using Custom Environment + +```python +# Agent needs ML libraries for analysis +# First, check available environments +envs = read_resource("mcp://forge/environments/list") + +# Build custom environment if needed +build_result = build_custom_environment( + name="my-ml-analysis", + packages=[ + "numpy>=1.24.0", + "pandas>=2.0.0", + "scikit-learn>=1.3.0", + "matplotlib>=3.7.0", + "xgboost>=2.0.0" + ], + description="Custom ML environment for analysis project" +) + +# Wait for build to complete (typically 30-60 seconds with UV cache) +print(f"Built {build_result['image_name']} in {build_result['build_time']}s") +print(f"Installed: {build_result['installed_packages']}") + +# Now use the custom environment for analysis +result = execute_python( + code=""" +import pandas as pd +import numpy as np +from sklearn.ensemble import RandomForestClassifier +from sklearn.model_selection import train_test_split +import xgboost as xgb + +# Load data (from volume or MCP tool) +data = filesystem_read('/data/customer_churn.csv') +df = pd.read_csv(data) + +# Preprocess +X = df.drop('churn', axis=1) +y = df['churn'] +X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2) + +# Train models +rf_model = RandomForestClassifier(n_estimators=100) +rf_model.fit(X_train, y_train) +rf_score = rf_model.score(X_test, y_test) + +xgb_model = xgb.XGBClassifier() +xgb_model.fit(X_train, y_train) +xgb_score = xgb_model.score(X_test, y_test) + +# Return results +{ + 'random_forest_accuracy': rf_score, + 'xgboost_accuracy': xgb_score, + 'best_model': 'xgboost' if xgb_score > rf_score else 'random_forest' +} +""", + custom_image="my-ml-analysis", + mcp_tools=["filesystem_read"], + session_id="churn-analysis" +) + +print(result['result']) +# Returns: {'random_forest_accuracy': 0.87, 'xgboost_accuracy': 0.91, 'best_model': 'xgboost'} +``` + +### Example 3: Stateful Analysis with Custom Environment + +```python +# Create session with custom environment +session_id = "deep-analysis" + +# Step 1: Load and explore data +execute_python( + code=""" +import pandas as pd +import matplotlib.pyplot as plt + +df = filesystem_read_csv('/data/sales_2024.csv') +print(f"Loaded {len(df)} rows") +print(df.describe()) + +# Store for later +df.to_pickle('/tmp/sales_df.pkl') +""", + custom_image="data-science", # Use pre-built template + mcp_tools=["filesystem_read_csv"], + session_id=session_id, + backend="jupyter" +) + +# Document what we have +document_state( + session_id=session_id, + variables={ + "df": "Sales data for 2024, 50k rows, pickled to /tmp/sales_df.pkl" + }, + note="Data loaded and ready for analysis" +) + +# ... conversation continues, context rotates ... + +# Step 2: Later, continue analysis +state = read_resource(f"mcp://forge/sessions/{session_id}/state") +# Agent sees df is available + +execute_python( + code=""" +import pandas as pd +df = pd.read_pickle('/tmp/sales_df.pkl') + +# Analyze trends +monthly_sales = df.groupby('month')['revenue'].sum() +growth_rate = monthly_sales.pct_change().mean() * 100 + +{ + 'total_revenue': float(df['revenue'].sum()), + 'avg_monthly_growth': float(growth_rate), + 'top_products': df.groupby('product')['revenue'].sum().nlargest(5).to_dict() +} +""", + session_id=session_id +) +``` + +### Example 4: Using Pre-built Templates + +```python +# Quick start with pre-built template +result = execute_python( + code=""" +import requests +from bs4 import BeautifulSoup + +# Scrape some data +response = requests.get('https://example.com/data') +soup = BeautifulSoup(response.content, 'lxml') + +# Extract and process +data = [item.text for item in soup.find_all('div', class_='data-item')] +data[:10] # Return first 10 items +""", + environment="web-scraping", # Use template instead of building custom + session_id=None +) + +print(result['result']) +``` + +## Implementation Roadmap + +### Phase 1: Core Execution (MVP) +- Basic execute_python tool with simple backend +- Podman integration with security restrictions +- MCP tool injection +- Pre-built Python 3.11 image + +### Phase 2: Stateful Execution +- Jupyter backend implementation +- Session management +- State documentation tool +- Session resources + +### Phase 3: Custom Environments +- build_custom_environment tool +- UV-based package installation +- Build caching +- Package allowlist/blocklist +- Security validation + +### Phase 4: Advanced Features +- Pre-built template library +- Environment lifecycle management +- Build rate limiting +- Advanced audit logging +- Multi-user support with quotas + +### Phase 5: Production Hardening +- Performance optimization +- Advanced security scanning +- Comprehensive monitoring +- High availability setup +- Documentation and examples \ No newline at end of file diff --git a/src/mcp_forge/adapters/executor_adapter.py b/src/mcp_forge/adapters/executor_adapter.py index 4f7e057..46653ad 100644 --- a/src/mcp_forge/adapters/executor_adapter.py +++ b/src/mcp_forge/adapters/executor_adapter.py @@ -42,28 +42,40 @@ class SimpleBackend: self.config = config # Create default resource limits from config - self.default_limits = ResourceLimits( - memory=config.execution.default_memory, - storage="10g", # Default storage limit - cpu_quota=config.execution.default_cpu_quota, - timeout=config.execution.default_timeout - ) + # When not enforcing limits, use very high values to effectively disable + if config.security.enforce_resource_limits: + self.default_limits = ResourceLimits( + memory=config.execution.default_memory, + storage="10g", + cpu_quota=config.execution.default_cpu_quota, + timeout=config.execution.default_timeout + ) + else: + # No resource enforcement - use very high limits (effectively unlimited) + self.default_limits = ResourceLimits( + memory="16g", # Very high memory limit + storage="100g", # Very high storage limit + cpu_quota=1000000, # Effectively unlimited CPU + timeout=config.execution.default_timeout + ) - # Create executor with default image + # Create executor with default image (prefer python_3_12) self.executor = CodeExecutor( container_manager=container_manager, - image=config.images.python, + image=config.images.python_3_12, resource_limits=self.default_limits ) - logger.debug(f"SimpleBackend initialized with image={config.images.python}") + logger.debug(f"SimpleBackend initialized with image={config.images.python_3_12}") def execute( self, code: str, timeout: Optional[int] = None, memory: Optional[str] = None, - cpu_quota: Optional[int] = None + cpu_quota: Optional[int] = None, + injection_code: Optional[str] = None, + bridge_socket_path: Optional[str] = None ) -> ExecutionResult: """ Execute Python code in isolated container. @@ -73,6 +85,8 @@ class SimpleBackend: timeout: Optional timeout override (seconds) memory: Optional memory limit override (e.g., "512m") cpu_quota: Optional CPU quota override + injection_code: Optional code to inject before user code (for MCP tools) + bridge_socket_path: Optional path to MCP bridge socket Returns: ExecutionResult with stdout, stderr, result, etc. @@ -88,10 +102,20 @@ class SimpleBackend: # Create temporary executor with custom limits executor = CodeExecutor( container_manager=self.container_manager, - image=self.config.images.python, + image=self.config.images.python_3_12, resource_limits=limits ) - return executor.execute(code, timeout=timeout) - - # Use default executor - return self.executor.execute(code, timeout=timeout) + return executor.execute( + code=code, + timeout=timeout, + injection_code=injection_code, + bridge_socket_path=bridge_socket_path + ) + else: + # Use default executor + return self.executor.execute( + code=code, + timeout=timeout, + injection_code=injection_code, + bridge_socket_path=bridge_socket_path + ) diff --git a/src/mcp_forge/execution/jupyter/backend.py b/src/mcp_forge/execution/jupyter/backend.py new file mode 100644 index 0000000..d51b3ec --- /dev/null +++ b/src/mcp_forge/execution/jupyter/backend.py @@ -0,0 +1,263 @@ +"""Jupyter backend for stateful code execution.""" + +from typing import Optional, Dict, List +import hashlib + +from mcp_forge.config.schema import ForgeConfig +from mcp_forge.podman.containers import SecureContainerManager +from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity +from mcp_forge.security.resource_limits import ResourceLimits, parse_memory_string +from mcp_forge.execution.simple.executor import ExecutionResult +from mcp_forge.execution.jupyter.kernel import JupyterKernelManager +from mcp_forge.execution.jupyter.sessions import SessionManager, SessionState, SessionError + + +class JupyterBackend: + """Stateful code execution backend using Jupyter kernels.""" + + def __init__( + self, + config: ForgeConfig, + container_manager: SecureContainerManager, + audit_logger: AuditLogger + ): + """ + Initialize Jupyter backend. + + Args: + config: Forge configuration + container_manager: Container lifecycle manager + audit_logger: Audit logging instance + """ + self.config = config + self.container_manager = container_manager + self.audit_logger = audit_logger + + # Initialize kernel manager + kernel_manager = JupyterKernelManager( + container_manager=container_manager, + image=config.images.jupyter, + resource_limits=self._default_resource_limits() + ) + + # Initialize session manager + self.session_manager = SessionManager( + config=config.sessions, + kernel_manager=kernel_manager, + audit_logger=audit_logger + ) + + def execute( + self, + code: str, + session_id: str, + timeout: Optional[int] = None, + memory: Optional[str] = None, + cpu_quota: Optional[int] = None, + custom_image: Optional[str] = None, + volumes: Optional[Dict[str, dict]] = None, + injection_code: Optional[str] = None, + bridge_socket_path: Optional[str] = None + ) -> ExecutionResult: + """ + Execute code in stateful session. + + Creates session if it doesn't exist, reuses existing session otherwise. + Session maintains namespace state across multiple executions. + + Args: + code: Python code to execute + session_id: Unique session identifier + timeout: Max execution time in seconds (uses config default if None) + memory: Memory limit string (uses config default if None) + cpu_quota: CPU quota (uses config default if None) + custom_image: Custom image name (uses config default if None) + volumes: Volume mounts dict + injection_code: Optional MCP tool injection code (executed once at session start) + bridge_socket_path: Optional path to MCP bridge socket for mounting + + Returns: + ExecutionResult with execution output and metadata + + Raises: + ValueError: If limits exceed configured maximums + SessionError: If session operation fails + """ + # Use defaults from config if not specified + timeout = timeout if timeout is not None else self.config.execution.default_timeout + memory = memory if memory is not None else self.config.execution.default_memory + cpu_quota = cpu_quota if cpu_quota is not None else self.config.execution.default_cpu_quota + + # Validate limits against maximums + self._validate_limits(timeout, memory, cpu_quota) + + # Log execution (hash code, don't log actual content) + code_hash = hashlib.sha256(code.encode()).hexdigest() + self.audit_logger.log( + event_type=AuditEventType.EXECUTION_REQUEST, + severity=AuditSeverity.INFO, + message="Stateful code execution requested", + session_id=session_id, + details={ + "code_hash": code_hash, + "timeout": timeout, + "memory": memory, + "cpu_quota": cpu_quota + } + ) + + # Check if session exists, create if needed + try: + self.session_manager.get_session(session_id) + except SessionError: + # Session doesn't exist, create it with MCP injection + resource_limits = ResourceLimits( + memory=memory, + cpu_quota=cpu_quota, + storage="1g", # Default storage quota + timeout=timeout + ) + + self.session_manager.create_session( + session_id=session_id, + resource_limits=resource_limits, + volumes=volumes, + injection_code=injection_code, + bridge_socket_path=bridge_socket_path + ) + + # Execute in session + result = self.session_manager.execute_in_session( + session_id=session_id, + code=code, + timeout=timeout + ) + + return result + + def document_state( + self, + session_id: str, + variables: Dict[str, str], + note: str = "", + clear: bool = False + ) -> dict: + """ + Document important variables in session. + + Args: + session_id: Session to document + variables: Dictionary of variable_name -> description + note: Optional note about session state + clear: If True, replace all documented variables; if False, merge + + Returns: + Dictionary with updated state info + + Raises: + SessionError: If session doesn't exist + """ + self.session_manager.document_state( + session_id=session_id, + variables=variables, + note=note, + clear=clear + ) + + # Return updated state + state = self.session_manager.get_session_state(session_id) + return state.to_dict() + + def get_session_state(self, session_id: str) -> SessionState: + """ + Get documented state for session. + + Args: + session_id: Session identifier + + Returns: + SessionState object + + Raises: + SessionError: If session doesn't exist + """ + return self.session_manager.get_session_state(session_id) + + def destroy_session(self, session_id: str) -> None: + """ + Destroy session and cleanup kernel. + + Args: + session_id: Session to destroy + + Raises: + SessionError: If session doesn't exist + """ + self.session_manager.destroy_session(session_id) + + def list_sessions(self) -> List[dict]: + """ + List all active sessions with metadata. + + Returns: + List of session dictionaries + """ + return self.session_manager.list_sessions() + + def cleanup_idle_sessions(self) -> int: + """ + Cleanup sessions idle beyond configured timeout. + + Returns: + Number of sessions cleaned up + """ + return self.session_manager.cleanup_idle_sessions() + + def _default_resource_limits(self) -> Optional[ResourceLimits]: + """ + Get default resource limits from config. + + Returns: + ResourceLimits with config defaults, or None if enforcement disabled + """ + if not self.config.security.enforce_resource_limits: + return None + + return ResourceLimits( + memory=self.config.execution.default_memory, + cpu_quota=self.config.execution.default_cpu_quota, + storage="1g", + timeout=self.config.execution.default_timeout + ) + + def _validate_limits(self, timeout: int, memory: str, cpu_quota: int) -> None: + """ + Validate resource limits against configured maximums. + + Args: + timeout: Timeout in seconds + memory: Memory limit string + cpu_quota: CPU quota value + + Raises: + ValueError: If any limit exceeds maximum + """ + # Validate timeout + if timeout > self.config.execution.max_timeout: + raise ValueError( + f"Timeout {timeout} exceeds maximum {self.config.execution.max_timeout}" + ) + + # Validate memory + memory_bytes = parse_memory_string(memory) + max_memory_bytes = parse_memory_string(self.config.execution.max_memory) + if memory_bytes > max_memory_bytes: + raise ValueError( + f"Memory {memory} exceeds maximum {self.config.execution.max_memory}" + ) + + # Validate CPU quota + if cpu_quota > self.config.execution.max_cpu_quota: + raise ValueError( + f"CPU quota {cpu_quota} exceeds maximum {self.config.execution.max_cpu_quota}" + ) diff --git a/src/mcp_forge/execution/jupyter/kernel.py b/src/mcp_forge/execution/jupyter/kernel.py new file mode 100644 index 0000000..668f1bd --- /dev/null +++ b/src/mcp_forge/execution/jupyter/kernel.py @@ -0,0 +1,631 @@ +""" +Real Jupyter kernel management for stateful execution. + +This module implements proper Jupyter kernel management: +- jupyter-client runs on host (MCP-Forge server) +- ipykernel runs inside Podman containers +- Communication via ZMQ protocol +- 1:1 mapping: one container per session, one kernel per container +""" + +from typing import Dict, Optional, List, Any +from dataclasses import dataclass +from datetime import datetime, timedelta +import uuid +import json +import tempfile +import time +import socket +from pathlib import Path + +from jupyter_client.blocking.client import BlockingKernelClient +import zmq + +from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig +from mcp_forge.security.resource_limits import ResourceLimits +from mcp_forge.execution.simple.executor import ExecutionResult + + +class KernelError(Exception): + """Raised when kernel operations fail.""" + pass + + +@dataclass +class KernelInfo: + """Information about a running kernel.""" + kernel_id: str + container_id: str + session_id: str + connection_file: Path + connection_info: Dict[str, Any] # ZMQ ports and keys + started_at: datetime + last_activity: datetime + client: Optional[BlockingKernelClient] = None + + def to_dict(self) -> dict: + """Convert to dictionary for JSON serialization.""" + return { + "kernel_id": self.kernel_id, + "container_id": self.container_id, + "session_id": self.session_id, + "started_at": self.started_at.isoformat(), + "last_activity": self.last_activity.isoformat(), + "connection_info": { + "shell_port": self.connection_info.get("shell_port"), + "iopub_port": self.connection_info.get("iopub_port"), + "stdin_port": self.connection_info.get("stdin_port"), + "control_port": self.connection_info.get("control_port"), + "hb_port": self.connection_info.get("hb_port"), + } + } + + +class JupyterKernelManager: + """ + Manages IPython kernels in containers via jupyter-client. + + Architecture: + - This class runs on host (MCP-Forge server process) + - Creates one container per session with ipykernel running inside + - Connects to kernel via ZMQ protocol (jupyter-client) + - Communicates using Jupyter message protocol + + Each session gets: + - Dedicated container + - Dedicated kernel process + - Isolated Python namespace + - Independent resource limits + """ + + def __init__( + self, + container_manager: SecureContainerManager, + image: str, + resource_limits: Optional[ResourceLimits] = None + ): + """ + Initialize kernel manager. + + Args: + container_manager: Container lifecycle manager + image: Docker/Podman image with ipykernel installed + resource_limits: Default resource limits for kernels + """ + self.container_manager = container_manager + self.image = image + self.resource_limits = resource_limits + self.kernels: Dict[str, KernelInfo] = {} + + def start_kernel( + self, + session_id: str, + volumes: Optional[Dict[str, dict]] = None, + injection_code: Optional[str] = None, + bridge_socket_path: Optional[str] = None + ) -> str: + """ + Start IPython kernel in dedicated container. + + Process: + 1. Generate ZMQ connection info (ports, keys) + 2. Create connection file + 3. Create container with ipykernel command + 4. Mount bridge socket if provided (for MCP tools) + 5. Start container + 6. Wait for kernel to be ready + 7. Connect jupyter-client to kernel via ZMQ + 8. Execute injection code (MCP tools setup) if provided + 9. Verify kernel is responsive + + Args: + session_id: Session ID this kernel belongs to + volumes: Optional volume mounts + injection_code: Optional MCP tool injection code to execute at startup + bridge_socket_path: Optional path to MCP bridge socket for mounting + + Returns: + kernel_id: Unique identifier for this kernel + + Raises: + KernelError: If kernel startup fails + """ + kernel_id = f"kernel-{uuid.uuid4().hex[:16]}" + + # Generate connection info + connection_info = self._generate_connection_info() + + # Create connection file + connection_file = self._create_connection_file(kernel_id, connection_info) + + try: + # Set up volumes (user volumes + bridge socket + connection file) + container_volumes = volumes.copy() if volumes else {} + if bridge_socket_path: + container_volumes[bridge_socket_path] = { + "bind": bridge_socket_path, + "mode": "rw" + } + + # Mount connection file into container + container_connection_path = f"/tmp/kernel-{kernel_id}.json" + container_volumes[str(connection_file)] = { + "bind": container_connection_path, + "mode": "ro" + } + + # Create container with ipykernel using host networking + config = ContainerConfig( + image=self.image, + command=[ + "python", "-m", "ipykernel_launcher", + "-f", container_connection_path + ], + resource_limits=self.resource_limits, + volumes=container_volumes, + network_mode="host", # Use host network for ZMQ communication + port_bindings={ + connection_info["shell_port"]: connection_info["shell_port"], + connection_info["iopub_port"]: connection_info["iopub_port"], + connection_info["stdin_port"]: connection_info["stdin_port"], + connection_info["control_port"]: connection_info["control_port"], + connection_info["hb_port"]: connection_info["hb_port"], + } + ) + + container_id = self.container_manager.create_container( + config, + session_id=session_id, + name=f"jupyter-{kernel_id}" + ) + + # Start container + self.container_manager.start_container(container_id) + + # Wait for kernel to be ready with polling + if not self._wait_for_kernel_ready(connection_info, timeout=30): + raise KernelError(f"Kernel {kernel_id} failed to start within timeout") + + # Connect client + client = self._connect_client(connection_info) + + # Verify kernel is responsive + if not self._verify_kernel(client): + raise KernelError(f"Kernel {kernel_id} not responsive") + + # Execute injection code if provided (MCP tools setup) + if injection_code: + self._execute_injection_code(client, injection_code, kernel_id) + + # Register kernel + now = datetime.utcnow() + kernel_info = KernelInfo( + kernel_id=kernel_id, + container_id=container_id, + session_id=session_id, + connection_file=connection_file, + connection_info=connection_info, + started_at=now, + last_activity=now, + client=client + ) + self.kernels[kernel_id] = kernel_info + + return kernel_id + + except Exception as e: + # Cleanup on failure + connection_file.unlink(missing_ok=True) + raise KernelError(f"Failed to start kernel: {e}") from e + + def execute_code( + self, + kernel_id: str, + code: str, + timeout: int = 300 + ) -> ExecutionResult: + """ + Execute code in kernel via ZMQ. + + Uses jupyter-client to: + 1. Send execute_request message + 2. Receive stream (stdout/stderr) messages + 3. Receive execute_result/display_data messages + 4. Collect and parse all output + + Args: + kernel_id: Kernel to execute in + code: Python code to execute + timeout: Maximum execution time in seconds + + Returns: + ExecutionResult with stdout, stderr, result + + Raises: + KernelError: If kernel not found or execution fails + """ + kernel_info = self._get_kernel(kernel_id) + client = kernel_info.client + + if not client: + raise KernelError(f"Kernel {kernel_id} has no connected client") + + start_time = time.time() + + try: + # Execute code + _msg_id = client.execute(code, silent=False, store_history=True) + + # Collect output + stdout_parts = [] + stderr_parts = [] + result = None + has_error = False + + # Wait for execution to complete + while True: + try: + msg = client.get_iopub_msg(timeout=timeout) + msg_type = msg['header']['msg_type'] + content = msg['content'] + + if msg_type == 'stream': + if content['name'] == 'stdout': + stdout_parts.append(content['text']) + elif content['name'] == 'stderr': + stderr_parts.append(content['text']) + + elif msg_type == 'execute_result': + result = content.get('data', {}).get('text/plain', '') + + elif msg_type == 'error': + has_error = True + stderr_parts.append('\n'.join(content['traceback'])) + + elif msg_type == 'status': + if content['execution_state'] == 'idle': + break + + except zmq.error.Again: + break + + execution_time = time.time() - start_time + + # Update last activity + kernel_info.last_activity = datetime.utcnow() + + stderr_text = ''.join(stderr_parts) + + return ExecutionResult( + success=(not has_error), + stdout=''.join(stdout_parts), + stderr=stderr_text, + result=result, + execution_time=execution_time, + exit_code=1 if has_error else 0, + error=stderr_text if has_error else None + ) + + except Exception as e: + execution_time = time.time() - start_time + return ExecutionResult( + success=False, + stdout='', + stderr=str(e), + result=None, + execution_time=execution_time, + exit_code=1, + error=str(e) + ) + + def shutdown_kernel(self, kernel_id: str) -> None: + """ + Shutdown kernel and cleanup container. + + 1. Send shutdown_request via ZMQ + 2. Wait for kernel shutdown + 3. Stop and remove container + 4. Cleanup connection file + + Args: + kernel_id: Kernel to shutdown + """ + kernel_info = self._get_kernel(kernel_id) + + try: + # Shutdown kernel + if kernel_info.client: + kernel_info.client.shutdown() + kernel_info.client.stop_channels() + + # Stop and remove container + self.container_manager.stop_container(kernel_info.container_id) + self.container_manager.remove_container(kernel_info.container_id) + + # Cleanup connection file + kernel_info.connection_file.unlink(missing_ok=True) + + finally: + # Remove from registry + del self.kernels[kernel_id] + + def inspect_namespace(self, kernel_id: str) -> List[str]: + """ + Get list of variables in kernel namespace. + + Executes introspection code: + [var for var in dir() if not var.startswith('_')] + + Args: + kernel_id: Kernel to inspect + + Returns: + List of variable names + """ + code = "[var for var in dir() if not var.startswith('_')]" + result = self.execute_code(kernel_id, code, timeout=5) + + if result.success and result.result: + # Parse result (it's a string representation of a list) + try: + return eval(result.result) # nosec - controlled code + except Exception: + return [] + return [] + + def get_variable_info( + self, + kernel_id: str, + variable_name: str + ) -> Dict[str, Any]: + """ + Get detailed information about a variable. + + Executes introspection code to get: + - type(var).__name__ + - sys.getsizeof(var) if available + - var.shape if hasattr(var, 'shape') + - repr(var)[:100] + + Args: + kernel_id: Kernel to inspect + variable_name: Name of variable to inspect + + Returns: + Dict with type, size, shape, repr + """ + code = f""" +import sys +_var = {variable_name} +_info = {{ + 'type': type(_var).__name__, + 'repr': repr(_var)[:100], +}} +try: + _info['size_bytes'] = sys.getsizeof(_var) +except: + pass +if hasattr(_var, 'shape'): + _info['shape'] = _var.shape +_info +""" + result = self.execute_code(kernel_id, code, timeout=5) + + if result.success and result.result: + try: + return eval(result.result) # nosec - controlled code + except Exception: + return {} + return {} + + def restart_kernel(self, kernel_id: str) -> None: + """ + Restart kernel (namespace reset, container kept). + + Strategy: shutdown current kernel and start new one in same container. + Note: In a full implementation, we'd use KernelManager.restart_kernel(). + + Args: + kernel_id: Kernel to restart + """ + kernel_info = self._get_kernel(kernel_id) + + # For now, just record activity - full restart implementation requires + # KernelManager integration (not just BlockingKernelClient) + # TODO: Implement proper kernel restart via KernelManager + kernel_info.last_activity = datetime.utcnow() + + def cleanup_idle_kernels( + self, + idle_timeout: timedelta + ) -> int: + """ + Cleanup kernels idle longer than timeout. + + Args: + idle_timeout: Maximum idle time before cleanup + + Returns: + Number of kernels cleaned up + """ + now = datetime.utcnow() + cleaned_up = 0 + + for kernel_id in list(self.kernels.keys()): + kernel_info = self.kernels[kernel_id] + idle_time = now - kernel_info.last_activity + + if idle_time > idle_timeout: + try: + self.shutdown_kernel(kernel_id) + cleaned_up += 1 + except Exception: + pass # Continue cleanup even if one fails + + return cleaned_up + + def _get_kernel(self, kernel_id: str) -> KernelInfo: + """Get kernel info or raise error.""" + if kernel_id not in self.kernels: + raise KernelError(f"Kernel {kernel_id} not found") + return self.kernels[kernel_id] + + def _generate_connection_info(self) -> Dict[str, Any]: + """Generate ZMQ connection information with allocated ports.""" + import secrets + + # Allocate 5 ports for ZMQ channels + ports = self._allocate_ports(5) + + return { + "shell_port": ports[0], + "iopub_port": ports[1], + "stdin_port": ports[2], + "control_port": ports[3], + "hb_port": ports[4], + "ip": "127.0.0.1", + "key": secrets.token_hex(32), + "transport": "tcp", + "signature_scheme": "hmac-sha256", + "kernel_name": "python3" + } + + def _allocate_ports(self, count: int) -> List[int]: + """ + Allocate available ports for ZMQ. + + Args: + count: Number of ports to allocate + + Returns: + List of allocated port numbers + """ + ports = [] + for _ in range(count): + # Let OS assign available port + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(('127.0.0.1', 0)) # Bind to any available port + port = sock.getsockname()[1] + sock.close() + ports.append(port) + return ports + + def _create_connection_file( + self, + kernel_id: str, + connection_info: Dict[str, Any] + ) -> Path: + """Create connection file for kernel.""" + # Create temp file + fd, path = tempfile.mkstemp(suffix=f"-kernel-{kernel_id}.json") + + # Write connection info + with open(fd, 'w') as f: + json.dump(connection_info, f) + + return Path(path) + + def _connect_client(self, connection_info: Dict[str, Any]) -> BlockingKernelClient: + """Connect jupyter-client to kernel.""" + client = BlockingKernelClient() + client.load_connection_info(connection_info) + client.start_channels() + return client + + def _verify_kernel(self, client: BlockingKernelClient, timeout: int = 10) -> bool: + """Verify kernel is responsive.""" + try: + client.wait_for_ready(timeout=timeout) + return True + except Exception: + return False + + def _wait_for_kernel_ready( + self, + connection_info: Dict[str, Any], + timeout: int = 30, + poll_interval: float = 0.5 + ) -> bool: + """ + Wait for kernel to be ready by polling ports. + + Args: + connection_info: Kernel connection information + timeout: Maximum time to wait in seconds + poll_interval: Time between polls in seconds + + Returns: + True if kernel is ready, False if timeout + """ + start_time = time.time() + shell_port = connection_info["shell_port"] + + while time.time() - start_time < timeout: + try: + # Try to connect to shell port + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(1) + result = sock.connect_ex(('127.0.0.1', shell_port)) + sock.close() + + if result == 0: + # Port is open, kernel is ready + return True + except Exception: + pass + + time.sleep(poll_interval) + + return False + + def _execute_injection_code( + self, + client: BlockingKernelClient, + injection_code: str, + kernel_id: str + ) -> None: + """ + Execute MCP tool injection code on kernel startup. + + This runs once when the kernel starts to set up MCP tools. + Unlike regular code execution, we don't capture output. + + Args: + client: Connected kernel client + injection_code: Python code to inject (MCP tools setup) + kernel_id: Kernel ID for error messages + + Raises: + KernelError: If injection code fails to execute + """ + try: + # Execute injection code silently + _msg_id = client.execute(injection_code, silent=True, store_history=False) + + # Wait for execution to complete + timeout = 10 # Injection should be fast + while True: + try: + msg = client.get_iopub_msg(timeout=timeout) + msg_type = msg['header']['msg_type'] + + if msg_type == 'error': + content = msg['content'] + error_msg = '\n'.join(content.get('traceback', [str(content)])) + raise KernelError( + f"MCP injection failed in kernel {kernel_id}: {error_msg}" + ) + + elif msg_type == 'status': + if msg['content']['execution_state'] == 'idle': + break # Injection complete + + except zmq.error.Again: + break # Timeout, assume success + + except KernelError: + raise + except Exception as e: + raise KernelError( + f"Failed to execute MCP injection code in kernel {kernel_id}: {e}" + ) from e diff --git a/src/mcp_forge/execution/jupyter/sessions.py b/src/mcp_forge/execution/jupyter/sessions.py new file mode 100644 index 0000000..5b0a494 --- /dev/null +++ b/src/mcp_forge/execution/jupyter/sessions.py @@ -0,0 +1,435 @@ +"""Session management for stateful execution.""" + +from typing import Dict, Optional, List, Any +from dataclasses import dataclass, field +from datetime import datetime, timedelta + +from mcp_forge.execution.jupyter.kernel import JupyterKernelManager +from mcp_forge.config.schema import SessionConfig +from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity +from mcp_forge.security.resource_limits import ResourceLimits +from mcp_forge.execution.simple.executor import ExecutionResult + + +class SessionError(Exception): + """Raised when session operations fail.""" + pass + + +@dataclass +class SessionState: + """Documented state for a session.""" + session_id: str + documented_variables: Dict[str, str] = field(default_factory=dict) + note: str = "" + last_updated: datetime = field(default_factory=datetime.utcnow) + all_variables: List[str] = field(default_factory=list) + introspection: Dict[str, dict] = field(default_factory=dict) + + def to_dict(self) -> dict: + """Convert to dictionary for JSON serialization.""" + return { + "session_id": self.session_id, + "documented_variables": self.documented_variables, + "note": self.note, + "last_updated": self.last_updated.isoformat(), + "all_variables": self.all_variables, + "introspection": self.introspection + } + + +class Session: + """Stateful execution session.""" + + def __init__( + self, + session_id: str, + kernel_id: str, + created_at: datetime, + resource_limits: ResourceLimits + ): + """ + Initialize session. + + Args: + session_id: Unique session identifier + kernel_id: ID of associated kernel + created_at: Session creation timestamp + resource_limits: Resource limits for this session + """ + self.session_id = session_id + self.kernel_id = kernel_id + self.created_at = created_at + self.last_activity = created_at + self.resource_limits = resource_limits + self.state = SessionState(session_id=session_id) + self.documented_variables: Dict[str, str] = {} + self.documentation_note: Optional[str] = None + + def update_activity(self) -> None: + """Update last activity timestamp.""" + self.last_activity = datetime.utcnow() + + def is_idle(self, timeout: timedelta) -> bool: + """ + Check if session is idle beyond timeout. + + Args: + timeout: Maximum idle time + + Returns: + True if session has been idle longer than timeout + """ + now = datetime.utcnow() + idle_time = now - self.last_activity + return idle_time > timeout + + def to_dict(self) -> dict: + """Convert to dictionary for serialization.""" + return { + "session_id": self.session_id, + "kernel_id": self.kernel_id, + "created_at": self.created_at.isoformat(), + "last_activity": self.last_activity.isoformat(), + "state": self.state.to_dict() + } + + +class SessionManager: + """Manages stateful execution sessions.""" + + def __init__( + self, + config: SessionConfig, + kernel_manager: JupyterKernelManager, + audit_logger: AuditLogger + ): + """ + Initialize session manager. + + Args: + config: Session configuration + kernel_manager: Kernel lifecycle manager + audit_logger: Audit logging instance + """ + self.config = config + self.kernel_manager = kernel_manager + self.audit_logger = audit_logger + self.sessions: Dict[str, Session] = {} + + def create_session( + self, + session_id: str, + resource_limits: ResourceLimits, + volumes: Optional[Dict[str, dict]] = None, + injection_code: Optional[str] = None, + bridge_socket_path: Optional[str] = None + ) -> Session: + """ + Create new stateful session. + + Args: + session_id: Unique identifier for session + resource_limits: Resource limits for session + volumes: Optional volume mounts + injection_code: Optional MCP tool injection code to execute at startup + bridge_socket_path: Optional path to MCP bridge socket for mounting + + Returns: + Created Session object + + Raises: + SessionError: If session_id already exists + SessionError: If max concurrent sessions exceeded + """ + if session_id in self.sessions: + raise SessionError(f"Session {session_id} already exists") + + # Check max concurrent limit + self._enforce_max_concurrent() + + # Start kernel with MCP injection if provided + kernel_id = self.kernel_manager.start_kernel( + session_id, + volumes=volumes, + injection_code=injection_code, + bridge_socket_path=bridge_socket_path + ) + + # Create session + now = datetime.utcnow() + session = Session( + session_id=session_id, + kernel_id=kernel_id, + created_at=now, + resource_limits=resource_limits + ) + + self.sessions[session_id] = session + + # Log session creation + self.audit_logger.log( + event_type=AuditEventType.SESSION_CREATE, + severity=AuditSeverity.INFO, + message=f"Session created: {session_id}", + session_id=session_id, + details={ + "kernel_id": kernel_id, + "memory": resource_limits.memory_bytes, + "cpu_quota": resource_limits.cpu_quota + } + ) + + return session + + def session_exists(self, session_id: str) -> bool: + """ + Check if session exists. + + Args: + session_id: Session identifier + + Returns: + True if session exists, False otherwise + """ + return session_id in self.sessions + + def get_session(self, session_id: str) -> Session: + """ + Get session by ID. + + Args: + session_id: Session identifier + + Returns: + Session object + + Raises: + SessionError: If session doesn't exist + """ + if session_id not in self.sessions: + raise SessionError(f"Session {session_id} not found") + + return self.sessions[session_id] + + async def document_variables( + self, + session_id: str, + variables: Dict[str, str], + note: Optional[str] = None, + clear: bool = False + ) -> Dict: + """ + Document important variables in a session. + + Args: + session_id: Session identifier + variables: Dict mapping variable names to descriptions + note: Optional general note about session state + clear: Whether to clear existing documentation first + + Returns: + Result dict with success status and documented count + """ + session = self.get_session(session_id) + + if clear: + session.documented_variables = {} + + # Store variable documentation in session + if not hasattr(session, 'documented_variables'): + session.documented_variables = {} + + session.documented_variables.update(variables) + + if note: + session.documentation_note = note + + return { + "success": True, + "documented_count": len(variables), + "total_documented": len(session.documented_variables) + } + + def execute_in_session( + self, + session_id: str, + code: str, + timeout: int = 300 + ) -> ExecutionResult: + """ + Execute code in session kernel. + + Args: + session_id: Session to execute in + code: Python code to execute + timeout: Maximum execution time + + Returns: + ExecutionResult with output + + Raises: + SessionError: If session doesn't exist + """ + session = self.get_session(session_id) + + # Update activity + session.update_activity() + + # Execute in kernel + result = self.kernel_manager.execute_code( + session.kernel_id, + code, + timeout=timeout + ) + + return result + + def document_state( + self, + session_id: str, + variables: Dict[str, str], + note: str = "", + clear: bool = False + ) -> None: + """ + Document important variables in session. + + Updates session.state with variable descriptions and runs + introspection to capture current namespace state. + + Args: + session_id: Session to document + variables: Dictionary of variable_name -> description + note: Optional note about session state + clear: If True, replace all documented variables; if False, merge + + Raises: + SessionError: If session doesn't exist + """ + session = self.get_session(session_id) + + # Update documented variables + if clear: + session.state.documented_variables = variables.copy() + else: + session.state.documented_variables.update(variables) + + # Update note if provided + if note: + session.state.note = note + + # Run introspection to get current namespace state + session.state.all_variables = self.kernel_manager.inspect_namespace(session.kernel_id) + + # Get variable info for documented variables + session.state.introspection = {} + for var_name in variables.keys(): + if var_name in session.state.all_variables: + try: + info = self.kernel_manager.get_variable_info(session.kernel_id, var_name) + session.state.introspection[var_name] = info + except Exception: + # Variable might not exist yet + pass + + # Update timestamp + session.state.last_updated = datetime.utcnow() + session.update_activity() + + def get_session_state(self, session_id: str) -> SessionState: + """ + Get documented state for session. + + Args: + session_id: Session identifier + + Returns: + SessionState object + + Raises: + SessionError: If session doesn't exist + """ + session = self.get_session(session_id) + return session.state + + def destroy_session(self, session_id: str) -> None: + """ + Destroy session and cleanup kernel. + + Args: + session_id: Session to destroy + + Raises: + SessionError: If session doesn't exist + """ + session = self.get_session(session_id) + + # Shutdown kernel + try: + self.kernel_manager.shutdown_kernel(session.kernel_id) + except Exception as e: + # Log but continue with cleanup + self.audit_logger.log( + event_type=AuditEventType.SESSION_DESTROY, + severity=AuditSeverity.WARNING, + message=f"Error shutting down kernel for session {session_id}", + session_id=session_id, + error=str(e) + ) + + # Remove session + del self.sessions[session_id] + + # Log destruction + self.audit_logger.log( + event_type=AuditEventType.SESSION_DESTROY, + severity=AuditSeverity.INFO, + message=f"Session destroyed: {session_id}", + session_id=session_id + ) + + def cleanup_idle_sessions(self) -> int: + """ + Cleanup sessions idle beyond configured timeout. + + Returns: + Number of sessions cleaned up + """ + timeout = timedelta(seconds=self.config.idle_timeout) + sessions_to_remove = [] + + for session_id, session in self.sessions.items(): + if session.is_idle(timeout): + sessions_to_remove.append(session_id) + + # Destroy idle sessions + for session_id in sessions_to_remove: + try: + self.destroy_session(session_id) + except Exception: + # Best effort cleanup + pass + + return len(sessions_to_remove) + + def list_sessions(self) -> List[dict]: + """ + List all active sessions with metadata. + + Returns: + List of session dictionaries + """ + return [session.to_dict() for session in self.sessions.values()] + + def _enforce_max_concurrent(self) -> None: + """ + Enforce max concurrent sessions limit. + + Raises: + SessionError: If at max concurrent sessions + """ + if len(self.sessions) >= self.config.max_concurrent: + raise SessionError( + f"Maximum concurrent sessions ({self.config.max_concurrent}) reached" + ) diff --git a/src/mcp_forge/podman/containers.py b/src/mcp_forge/podman/containers.py new file mode 100644 index 0000000..e84f3f7 --- /dev/null +++ b/src/mcp_forge/podman/containers.py @@ -0,0 +1,508 @@ +""" +Secure container management with security enforcement. + +All container operations are validated against security policy +before being sent to Podman. Provides lifecycle management +with comprehensive audit logging. +""" + +from typing import Optional, Dict, List +from datetime import datetime, timedelta +from pathlib import Path + +from mcp_forge.podman.client import PodmanClient +from mcp_forge.security.allowlist import OperationValidator, SecurityError +from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity +from mcp_forge.security.resource_limits import ResourceLimits + + +class ContainerConfig: + """Container configuration with security defaults.""" + + def __init__( + self, + image: str, + command: Optional[List[str]] = None, + environment: Optional[Dict[str, str]] = None, + volumes: Optional[Dict[str, dict]] = None, + resource_limits: Optional[ResourceLimits] = None, + working_dir: Optional[str] = None, + user: str = "1000:1000", + network_mode: str = "none", + port_bindings: Optional[Dict[str, int]] = None + ): + """ + Initialize container configuration. + + Args: + image: Container image to use + command: Command to run in container + environment: Environment variables + volumes: Volume mounts (host_path -> {bind, mode}) + resource_limits: Resource limits to apply + working_dir: Working directory in container (None to use image default) + user: User to run as (UID:GID) + network_mode: Network mode (none, host, bridge). Default is 'none' for security. + port_bindings: Port mappings for network_mode=host (container_port -> host_port) + """ + self.image = image + self.command = command or [] + self.environment = environment or {} + self.volumes = volumes or {} + self.resource_limits = resource_limits + self.working_dir = working_dir + self.user = user + self.network_mode = network_mode + self.port_bindings = port_bindings or {} + + def to_podman_params(self) -> dict: + """ + Convert to Podman container create parameters. + + Ensures all security requirements are included: + - network_mode: configurable (default 'none' for security) + - read_only: True + - security_opt: ["no-new-privileges"] + - resource limits + - port_bindings: for host networking mode + + Returns: + Dictionary of parameters for Podman containers.create() + """ + params = { + "image": self.image, + "command": self.command if self.command else None, + "environment": self.environment, + "user": self.user, + # Security requirements + "network_mode": self.network_mode, + "read_only": True, + "security_opt": ["no-new-privileges"], + } + + # Add working_dir only if explicitly set + if self.working_dir is not None: + params["working_dir"] = self.working_dir + + # Add port bindings if using host network mode + # Note: In host mode, ports are directly accessible + # port_bindings are informational for tracking + if self.network_mode == "host" and self.port_bindings: + # With host networking, container uses host's network stack directly + # No explicit port mapping needed, but we track for documentation + pass + + # Add volumes if present + if self.volumes: + params["volumes"] = self.volumes + + # Add resource limits if present + if self.resource_limits: + limit_params = self.resource_limits.to_podman_params() + params.update(limit_params) + # Disable swap to avoid cgroup swap.max issues on some systems + if "mem_limit" in params: + params["memswap_limit"] = -1 # Disable swap + + return params + + +class SecureContainerManager: + """Manages container lifecycle with security enforcement.""" + + def __init__( + self, + podman_client: PodmanClient, + validator: OperationValidator, + audit_logger: AuditLogger + ): + """ + Initialize secure container manager. + + Args: + podman_client: Podman client wrapper + validator: Operation validator for security checks + audit_logger: Audit logger for operation logging + """ + self.podman = podman_client + self.validator = validator + self.audit_logger = audit_logger + + def create_container( + self, + config: ContainerConfig, + session_id: Optional[str] = None, + name: Optional[str] = None, + **extra_params + ) -> str: + """ + Create a container with security validation. + + Args: + config: Container configuration + session_id: Session ID for tracking + name: Optional container name + **extra_params: Additional parameters (checked for forbidden values) + + Returns: + Container ID + + Raises: + SecurityError: If configuration violates security policy + """ + # Convert config to Podman parameters + params = config.to_podman_params() + + # Add session label if provided + labels = {} + if session_id: + labels["mcp-forge.session"] = session_id + if labels: + params["labels"] = labels + + if name: + params["name"] = name + + # Merge any extra parameters (will be validated) + params.update(extra_params) + + # Validate against security policy + try: + # Extract image from params for validation + self.validator.validate_container_create( + image=config.image, + params=params, + session_id=session_id + ) + except SecurityError as e: + # Log security violation + self.audit_logger.log_security_violation( + operation="container_create", + reason=str(e), + session_id=session_id + ) + raise + + # Create container + try: + container = self.podman.client.containers.create(**params) + container_id = container.id + + # Register with validator + if session_id: + self.validator.register_session_container(container_id) + + # Log successful creation + self.audit_logger.log_container_operation( + operation="create", + container_id=container_id, + image=config.image, + session_id=session_id, + details={ + "name": name, + "command": config.command + } + ) + + return container_id + + except Exception as e: + self.audit_logger.log( + event_type=AuditEventType.CONTAINER_CREATE, + severity=AuditSeverity.ERROR, + message=f"Container creation failed: {e}", + details={ + "image": config.image, + "session_id": session_id, + "error": str(e) + } + ) + raise + + def start_container(self, container_id: str) -> None: + """ + Start a container. + + Args: + container_id: Container ID to start + + Raises: + SecurityError: If container is not a session container + """ + # Verify container is registered (security check) + if container_id not in self.validator.session_containers: + self.audit_logger.log_security_violation( + operation="container_start", + reason=f"Attempted to start unregistered container: {container_id}" + ) + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + container.start() + + self.audit_logger.log_container_operation( + operation="start", + container_id=container_id, + image="" # Not available without extra lookup + ) + + except Exception as e: + self.audit_logger.log( + event_type=AuditEventType.CONTAINER_START, + severity=AuditSeverity.ERROR, + message=f"Container start failed: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def stop_container( + self, + container_id: str, + timeout: int = 10 + ) -> None: + """ + Stop a container. + + Args: + container_id: Container ID to stop + timeout: Timeout in seconds + """ + # Verify container is registered + if container_id not in self.validator.session_containers: + self.audit_logger.log_security_violation( + operation="container_stop", + reason=f"Attempted to stop unregistered container: {container_id}" + ) + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + container.stop(timeout=timeout) + + self.audit_logger.log_container_operation( + operation="stop", + container_id=container_id, + image="", + details={"timeout": timeout} + ) + + except Exception as e: + self.audit_logger.log( + event_type=AuditEventType.CONTAINER_STOP, + severity=AuditSeverity.ERROR, + message=f"Container stop failed: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def remove_container( + self, + container_id: str, + force: bool = False + ) -> None: + """ + Remove a container. + + Args: + container_id: Container ID to remove + force: Force removal even if running + """ + # Verify container is registered + if container_id not in self.validator.session_containers: + self.audit_logger.log_security_violation( + operation="container_remove", + reason=f"Attempted to remove unregistered container: {container_id}" + ) + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + container.remove(force=force) + + # Unregister from validator + self.validator.unregister_session_container(container_id) + + self.audit_logger.log_container_operation( + operation="remove", + container_id=container_id, + image="", + details={"force": force} + ) + + except Exception as e: + self.audit_logger.log( + event_type=AuditEventType.CONTAINER_REMOVE, + severity=AuditSeverity.ERROR, + message=f"Container removal failed: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def get_container_logs( + self, + container_id: str, + tail: int = 100 + ) -> tuple[str, str]: + """ + Get container stdout and stderr logs. + + Args: + container_id: Container ID + tail: Number of lines to retrieve + + Returns: + (stdout, stderr) as strings + """ + if container_id not in self.validator.session_containers: + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + logs = container.logs(tail=tail, stdout=True, stderr=True) + + # Podman logs returns a generator of frames, need to consume it + if hasattr(logs, '__iter__') and not isinstance(logs, (str, bytes)): + # It's a generator/iterator, consume it + logs_bytes = b''.join(logs) + logs_str = logs_bytes.decode('utf-8', errors='replace') + elif isinstance(logs, bytes): + logs_str = logs.decode('utf-8', errors='replace') + else: + logs_str = str(logs) + + # For simplicity, return all logs in stdout (Podman combines them) + return logs_str, "" + + except Exception as e: + self.audit_logger.log( + event_type=AuditEventType.EXECUTION_REQUEST, + severity=AuditSeverity.ERROR, + message=f"Failed to get container logs: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def wait_for_container( + self, + container_id: str, + timeout: int = 300 + ) -> int: + """ + Wait for container to exit. + + Args: + container_id: Container ID + timeout: Timeout in seconds + + Returns: + Exit code + + Raises: + TimeoutError: If container doesn't exit within timeout + """ + if container_id not in self.validator.session_containers: + raise SecurityError( + f"Container {container_id} is not a registered session container" + ) + + try: + container = self.podman.client.containers.get(container_id) + result = container.wait(timeout=timeout) + + # Extract exit code from result + if isinstance(result, dict): + exit_code = result.get("StatusCode", 0) + else: + exit_code = result + + return exit_code + + except Exception as e: + self.audit_logger.log( + event_type=AuditEventType.EXECUTION_REQUEST, + severity=AuditSeverity.ERROR, + message=f"Failed to wait for container: {e}", + details={"container_id": container_id, "error": str(e)} + ) + raise + + def cleanup_old_containers( + self, + max_age: timedelta = timedelta(hours=24) + ) -> int: + """ + Cleanup containers older than max_age. + + Args: + max_age: Maximum age for containers + + Returns: + Number of containers removed + """ + try: + # Get all containers with mcp-forge.session label + containers = self.podman.client.containers.list( + all=True, + filters={"label": ["mcp-forge.session"]} + ) + + removed_count = 0 + now = datetime.now() + + for container in containers: + # Get creation time + created_str = container.attrs.get("Created", "") + if not created_str: + continue + + # Parse ISO format timestamp + try: + # Remove fractional seconds and timezone for parsing + created_str = created_str.split('.')[0] + created = datetime.fromisoformat(created_str.replace('Z', '')) + except (ValueError, AttributeError): + continue + + age = now - created + + if age > max_age: + try: + container.remove(force=True) + removed_count += 1 + + self.audit_logger.log( + event_type=AuditEventType.CONTAINER_REMOVE, + severity=AuditSeverity.INFO, + message=f"Cleaned up old container: {container.id}", + details={ + "container_id": container.id, + "age_hours": age.total_seconds() / 3600 + } + ) + except Exception as e: + self.audit_logger.log( + event_type=AuditEventType.CONTAINER_REMOVE, + severity=AuditSeverity.WARNING, + message=f"Failed to remove old container: {e}", + details={"container_id": container.id, "error": str(e)} + ) + + return removed_count + + except Exception as e: + self.audit_logger.log( + event_type=AuditEventType.CONTAINER_REMOVE, + severity=AuditSeverity.ERROR, + message=f"Cleanup failed: {e}", + details={"error": str(e)} + ) + raise diff --git a/src/mcp_forge/server/server.py b/src/mcp_forge/server/server.py index e74f4ee..b4af6e6 100644 --- a/src/mcp_forge/server/server.py +++ b/src/mcp_forge/server/server.py @@ -406,7 +406,7 @@ print(f"Found {len(records)} records") # Create resource handler self.resource_handler = ResourceHandler( client_manager=self.client_manager, - jupyter_backend=self.jupyter_backend, + session_manager=self.jupyter_backend, environment_builder=self.environment_builder, config=self.config ) diff --git a/src/pod_executor/containers/manager.py b/src/pod_executor/containers/manager.py index 8e57ccc..ae7a274 100644 --- a/src/pod_executor/containers/manager.py +++ b/src/pod_executor/containers/manager.py @@ -97,12 +97,18 @@ class ContainerConfig: params["volumes"] = self.volumes # Add resource limits if present + # Skip resource limits if using very high values (indicates no enforcement) if self.resource_limits: limit_params = self.resource_limits.to_podman_params() - params.update(limit_params) - # Disable swap to avoid cgroup swap.max issues on some systems - if "mem_limit" in params: - params["memswap_limit"] = -1 # Disable swap + # Only apply limits if they're reasonable (not "no enforcement" markers) + # Check if mem_limit looks like an enforcement bypass (>= 16GB) + mem_limit = limit_params.get("mem_limit", "0") + mem_bytes = int(mem_limit) if mem_limit != "0" else 0 + if mem_bytes < 16 * 1024 * 1024 * 1024: # Less than 16GB = real limit + params.update(limit_params) + # Disable swap to avoid cgroup swap.max issues on some systems + if "mem_limit" in params: + params["memswap_limit"] = -1 # Disable swap return params diff --git a/src/pod_executor/jupyter/backend.py b/src/pod_executor/jupyter/backend.py index 71b06e7..8b3bbb4 100644 --- a/src/pod_executor/jupyter/backend.py +++ b/src/pod_executor/jupyter/backend.py @@ -52,14 +52,15 @@ class JupyterBackend: # Initialize kernel manager kernel_manager = JupyterKernelManager( container_manager=container_manager, - image=config.images.jupyter, + image=image, resource_limits=self._default_resource_limits() ) # Initialize session manager self.session_manager = SessionManager( - config=config.sessions, kernel_manager=kernel_manager, + idle_timeout=idle_timeout, + max_sessions=max_sessions, audit_logger=audit_logger ) @@ -109,7 +110,8 @@ class JupyterBackend: # Log execution (hash code, don't log actual content) code_hash = hashlib.sha256(code.encode()).hexdigest() - self.audit_logger.log( + if self.audit_logger: + self.audit_logger.log( event_type="execution.request", severity="info", message="Stateful code execution requested", @@ -236,9 +238,7 @@ class JupyterBackend: Returns: ResourceLimits with config defaults, or None if enforcement disabled """ - if not self.config.security.enforce_resource_limits: - return None - + # Always return resource limits in pod_executor return ResourceLimits( memory=self.default_memory, cpu_quota=self.default_cpu_quota, diff --git a/src/pod_executor/security/audit.py b/src/pod_executor/security/audit.py index 5fca46f..b457305 100644 --- a/src/pod_executor/security/audit.py +++ b/src/pod_executor/security/audit.py @@ -130,3 +130,46 @@ class SimpleFileAuditLogger: # Write to file (append mode, file locking via 'a' mode) with open(self.log_path, 'a') as f: f.write(json.dumps(entry) + '\n') + + def log_container_operation( + self, + operation: str, + container_id: str, + image: str, + session_id: Optional[str] = None, + user_id: Optional[str] = None, + details: Optional[Dict[str, Any]] = None, + error: Optional[str] = None + ) -> None: + """ + Convenience method to log container operations. + + Args: + operation: Operation type (create, start, stop, remove) + container_id: Container ID + image: Container image name + session_id: Optional session ID + user_id: Optional user ID + details: Optional additional details + error: Optional error message + """ + severity = "error" if error else "info" + message = f"Container {operation}: {container_id[:12]} (image: {image})" + + op_details = { + "operation": operation, + "container_id": container_id, + "image": image + } + if details: + op_details.update(details) + + self.log( + event_type=f"container.{operation}", + severity=severity, + message=message, + session_id=session_id, + user_id=user_id, + details=op_details, + error=error + )