diff --git a/Dockerfile b/Dockerfile index 5310652..038c713 100644 --- a/Dockerfile +++ b/Dockerfile @@ -97,5 +97,11 @@ RUN npm install playwright@1.53.0 -g RUN npx playwright@1.53.0 install +# Create a non-root user +RUN useradd -ms /bin/bash coderunner + +# Switch to the non-root user +USER coderunner + # Use the entrypoint script ENTRYPOINT ["/entrypoint.sh"] diff --git a/README.md b/README.md index 61f1845..6cf6956 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ # CodeRunner: Run AI Generated Code Locally -CodeRunner is an MCP (Model Context Protocol) server that executes AI-generated code in a sandboxed environment on your Mac using Apple's native [containers](https://github.com/apple/container). +CodeRunner is a cross-platform MCP (Model Context Protocol) server that executes AI-generated code in a secure, sandboxed environment. It supports macOS via Apple's native [containers](https://github.com/apple/container) and Linux/Windows via Docker. **Key use case:** Process your local files (videos, images, documents, data) with remote LLMs like Claude or ChatGPT without uploading your files to the cloud. The LLM generates code that runs locally on your machine to analyze, transform, or process your files. @@ -24,18 +24,23 @@ CodeRunner is an MCP (Model Context Protocol) server that executes AI-generated ## Quick Start -**Prerequisites:** Mac with macOS and Apple Silicon (M1/M2/M3/M4), Python 3.10+ +**Prerequisites:** +* **For macOS:** Apple Silicon (M1/M2/M3/M4) and the [Apple Container](https://github.com/apple/container/releases) tool installed. +* **For Linux/Windows:** [Docker](https://docs.docker.com/get-docker/) installed and running. ```bash git clone https://github.com/instavm/coderunner.git cd coderunner chmod +x install.sh -sudo ./install.sh +./install.sh ``` -MCP server will be available at: http://coderunner.local:8222/mcp +The script will detect your operating system and set up CodeRunner accordingly. -**Install required packages** (use virtualenv and note the python path): +* **On macOS:** The MCP server will be available at `http://coderunner.local:8222/mcp` +* **On Linux/Windows:** The MCP server will be available at `http://localhost:8222/mcp` + +**Install required packages for examples:** ```bash pip install -r examples/requirements.txt ``` @@ -169,11 +174,13 @@ Code runs in an isolated container with VM-level isolation. Your host system and From [@apple/container](https://github.com/apple/container/blob/main/docs/technical-overview.md): >Each container has the isolation properties of a full VM, using a minimal set of core utilities and dynamic libraries to reduce resource utilization and attack surface. +On Linux and Windows, CodeRunner uses Docker for similar containerization and security benefits. + ## Architecture CodeRunner consists of: -- **Sandbox Container:** Isolated execution environment with Jupyter kernel -- **MCP Server:** Handles communication between AI models and the sandbox +- **Sandbox Container:** Isolated execution environment (Apple Container or Docker) with a Jupyter kernel. +- **MCP Server:** Handles communication between AI models and the sandbox. ## Examples diff --git a/install.sh b/install.sh index cd3d795..49151d9 100755 --- a/install.sh +++ b/install.sh @@ -1,77 +1,89 @@ #!/bin/bash +# --- Helper Functions --- + +# Function to check if a command exists +command_exists() { + command -v "$1" &> /dev/null +} + # Function to get current macOS version get_macos_version() { sw_vers -productVersion | awk -F. '{print $1 "." $2}' } -# Check the system type -if [[ "$OSTYPE" != "darwin"* ]]; then - echo "❌ This script is intended for macOS systems only. Exiting." - exit 1 -fi +# --- Main Installation Logic --- -# Check macOS version -macos_version=$(get_macos_version) -if (( $(echo "$macos_version < 26.0" | bc -l) )); then - echo "Warning: Your macOS version is $macos_version. Version 26.0 or later is recommended. Some features of 'container' might not work properly." -else +echo "Starting CodeRunner Setup..." + +# --- macOS Specific Setup --- +if [[ "$OSTYPE" == "darwin"* ]]; then echo "✅ macOS system detected." -fi -download_url="https://github.com/apple/container/releases/download/0.3.0/container-0.3.0-installer-signed.pkg" + # Check macOS version + macos_version=$(get_macos_version) + if (( $(echo "$macos_version < 26.0" | bc -l) )); then + echo "⚠️ Warning: Your macOS version is $macos_version. Version 26.0 or later is recommended for Apple Container." + fi -# Check if container is installed and display its version -if command -v container &> /dev/null -then - echo "Apple 'container' tool detected. Current version:" + # Check for Apple Container tool + if command_exists container; then + echo "✅ Apple 'container' tool detected." container --version - current_version=$(container --version | awk '{print $4}') - echo $current_version - target_version=$(echo $download_url | awk -F'/' '{print $8}') + else + echo "❌ Apple 'container' tool not found." + echo "Please install it from: https://github.com/apple/container/releases" + exit 1 + fi + echo "Starting Apple Container services..." + container system start + sudo container system dns create local + container system dns default set local - if [ "$current_version" != "$target_version" ]; then - echo "Consider updating to version $target_version. Download it here: $download_url" - fi + echo "Pulling the latest image for Apple Container..." + container image pull instavm/coderunner - echo "Stopping any running Apple 'container' processes..." -else - echo "Apple 'container' tool not detected. Proceeding with installation..." - - # Download and install the Apple 'container' tool - echo "Downloading Apple 'container' tool..." - curl -Lo container-installer.pkg "$download_url" - - echo "Installing Apple 'container' tool..." - sudo installer -pkg container-installer.pkg -target / -fi + echo "→ Ensuring coderunner assets directory exists..." + ASSETS_SRC="$HOME/.coderunner/assets" + mkdir -p "$ASSETS_SRC" -echo "Starting the Sandbox Container..." -container system start + echo "🚀 Starting CodeRunner container..." + container run --volume "$ASSETS_SRC:/app/uploads" --name coderunner --detach --rm --cpus 8 --memory 4g instavm/coderunner -echo "Setting up local network domain..." + echo "✅ Setup complete! MCP server is available at http://coderunner.local:8222/mcp" -# Run the commands for setting up the local network -echo "Running: sudo container system dns create local" -sudo container system dns create local +# --- Docker-based Setup for Linux/Other --- +else + echo "✅ Non-macOS system detected. Setting up with Docker." -echo "Running: container system dns default set local" -container system dns default set local + # Check for Docker + if ! command_exists docker; then + echo "❌ Docker is not installed. Please install Docker to continue." + echo "Visit: https://docs.docker.com/get-docker/" + exit 1 + fi -echo "Starting the Sandbox Container..." -container system start + echo "✅ Docker is installed." + # Check if Docker daemon is running + if ! docker info &> /dev/null; then + echo "❌ Docker daemon is not running. Please start Docker and re-run this script." + exit 1 + fi -echo "Pulling the latest image: instavm/coderunner" -container image pull instavm/coderunner + echo "Pulling the latest image from Docker Hub..." + docker pull instavm/coderunner -echo "→ Ensuring coderunner assets directory…" -ASSETS_SRC="$HOME/.coderunner/assets" -mkdir -p "$ASSETS_SRC" + echo "→ Ensuring coderunner assets directory exists..." + ASSETS_SRC="$HOME/.coderunner/assets" + mkdir -p "$ASSETS_SRC" -# Run the command to start the sandbox container -echo "Running: container run --name coderunner --detach --rm --cpus 8 --memory 4g instavm/coderunner" -container run --volume "$ASSETS_SRC:/app/uploads" --name coderunner --detach --rm --cpus 8 --memory 4g instavm/coderunner + echo "🚀 Starting CodeRunner container using Docker..." + docker run -d --rm --name coderunner \ + -p 8222:8222 \ + -v "$ASSETS_SRC:/app/uploads" \ + instavm/coderunner -echo "✅ Setup complete. MCP server is available at http://coderunner.local:8222/mcp" \ No newline at end of file + echo "✅ Setup complete! MCP server is available at http://localhost:8222/mcp" +fi diff --git a/kernel_manager.py b/kernel_manager.py new file mode 100644 index 0000000..bbd8fb6 --- /dev/null +++ b/kernel_manager.py @@ -0,0 +1,278 @@ + +import asyncio +import json +import logging +import time +import uuid +from typing import Dict, Optional, Set +from dataclasses import dataclass, field +from enum import Enum +from datetime import datetime, timedelta + +import aiofiles +import websockets +import httpx + +from utils import create_jupyter_request + +logger = logging.getLogger(__name__) + +# Kernel pool configuration +MAX_KERNELS = 5 +MIN_KERNELS = 2 +KERNEL_TIMEOUT = 300 # 5 minutes +KERNEL_HEALTH_CHECK_INTERVAL = 30 # 30 seconds +MAX_RETRY_ATTEMPTS = 3 +RETRY_BACKOFF_BASE = 2 # exponential backoff base + +# Jupyter connection settings +JUPYTER_WS_URL = "ws://127.0.0.1:8888" +JUPYTER_HTTP_URL = "http://127.0.0.1:8888" + +# Enhanced WebSocket settings +WEBSOCKET_TIMEOUT = 600 # 10 minutes for long operations +WEBSOCKET_PING_INTERVAL = 30 +WEBSOCKET_PING_TIMEOUT = 10 + +# --- CUSTOM EXCEPTIONS --- + +class KernelError(Exception): + """Base exception for kernel-related errors""" + pass + +class NoKernelAvailableError(KernelError): + """Raised when no kernels are available in the pool""" + pass + +class KernelExecutionError(KernelError): + """Raised when kernel execution fails""" + pass + +class KernelTimeoutError(KernelError): + """Raised when kernel operation times out""" + pass + +# --- KERNEL MANAGEMENT CLASSES --- + +class KernelState(Enum): + HEALTHY = "healthy" + BUSY = "busy" + UNRESPONSIVE = "unresponsive" + FAILED = "failed" + +@dataclass +class KernelInfo: + kernel_id: str + state: KernelState = KernelState.HEALTHY + last_used: datetime = field(default_factory=datetime.now) + last_health_check: datetime = field(default_factory=datetime.now) + current_operation: Optional[str] = None + failure_count: int = 0 + + def is_available(self) -> bool: + return self.state == KernelState.HEALTHY + + def needs_health_check(self) -> bool: + return datetime.now() - self.last_health_check > timedelta(seconds=KERNEL_HEALTH_CHECK_INTERVAL) + +class KernelPool: + def __init__(self, kernel_id_file_path): + self.kernels: Dict[str, KernelInfo] = {} + self.lock = asyncio.Lock() + self.busy_kernels: Set[str] = set() + self._initialized = False + self._health_check_task: Optional[asyncio.Task] = None + self.kernel_id_file_path = kernel_id_file_path + + async def initialize(self): + """Initialize the kernel pool with minimum number of kernels""" + if self._initialized: + return + + async with self.lock: + logger.info("Initializing kernel pool...") + + # Try to use existing kernel first + existing_kernel = await self._get_existing_kernel() + if existing_kernel: + self.kernels[existing_kernel] = KernelInfo(kernel_id=existing_kernel) + logger.info(f"Added existing kernel to pool: {existing_kernel}") + + # Create additional kernels to reach minimum + while len(self.kernels) < MIN_KERNELS: + kernel_id = await self._create_new_kernel() + if kernel_id: + self.kernels[kernel_id] = KernelInfo(kernel_id=kernel_id) + logger.info(f"Created new kernel: {kernel_id}") + else: + logger.warning("Failed to create minimum number of kernels") + break + + self._initialized = True + # Start health check background task + self._health_check_task = asyncio.create_task(self._health_check_loop()) + logger.info(f"Kernel pool initialized with {len(self.kernels)} kernels") + + async def get_available_kernel(self) -> Optional[str]: + """Get an available kernel from the pool""" + if not self._initialized: + await self.initialize() + + async with self.lock: + # Find healthy, available kernel + for kernel_id, kernel_info in self.kernels.items(): + if kernel_info.is_available() and kernel_id not in self.busy_kernels: + self.busy_kernels.add(kernel_id) + kernel_info.state = KernelState.BUSY + kernel_info.last_used = datetime.now() + logger.info(f"Assigned kernel {kernel_id} to operation") + return kernel_id + + # No available kernels, try to create a new one if under limit + if len(self.kernels) < MAX_KERNELS: + kernel_id = await self._create_new_kernel() + if kernel_id: + kernel_info = KernelInfo(kernel_id=kernel_id, state=KernelState.BUSY) + self.kernels[kernel_id] = kernel_info + self.busy_kernels.add(kernel_id) + logger.info(f"Created and assigned new kernel: {kernel_id}") + return kernel_id + + logger.warning("No available kernels in pool") + return None + + async def release_kernel(self, kernel_id: str, failed: bool = False): + """Release a kernel back to the pool""" + async with self.lock: + if kernel_id in self.busy_kernels: + self.busy_kernels.remove(kernel_id) + + if kernel_id in self.kernels: + kernel_info = self.kernels[kernel_id] + if failed: + kernel_info.failure_count += 1 + kernel_info.state = KernelState.FAILED + logger.warning(f"Kernel {kernel_id} marked as failed (failures: {kernel_info.failure_count})") + + # Remove failed kernel if it has too many failures + if kernel_info.failure_count >= MAX_RETRY_ATTEMPTS: + await self._remove_kernel(kernel_id) + # Create replacement kernel + new_kernel_id = await self._create_new_kernel() + if new_kernel_id: + self.kernels[new_kernel_id] = KernelInfo(kernel_id=new_kernel_id) + else: + kernel_info.state = KernelState.HEALTHY + kernel_info.current_operation = None + logger.info(f"Released kernel {kernel_id} back to pool") + + async def _get_existing_kernel(self) -> Optional[str]: + """Try to get kernel ID from existing file""" + try: + async with aiofiles.open(self.kernel_id_file_path, mode='r') as f: + kernel_id = (await f.read()).strip() + if kernel_id and await self._check_kernel_health(kernel_id): + return kernel_id + except FileNotFoundError: + # This is a normal case if the server is starting for the first time. + pass + except Exception as e: + logger.warning(f"Could not read or validate existing kernel from {self.kernel_id_file_path}: {e}") + return None + + async def _create_new_kernel(self) -> Optional[str]: + """Create a new Jupyter kernel""" + try: + async with httpx.AsyncClient() as client: + response = await client.post( + f"{JUPYTER_HTTP_URL}/api/kernels", + json={"name": "python3"}, + timeout=30.0 + ) + if response.status_code == 201: + kernel_data = response.json() + kernel_id = kernel_data["id"] + logger.info(f"Created new kernel: {kernel_id}") + return kernel_id + else: + logger.error(f"Failed to create kernel: {response.status_code}") + except Exception as e: + logger.error(f"Error creating kernel: {e}") + return None + + async def _remove_kernel(self, kernel_id: str): + """Remove and shutdown a kernel""" + try: + async with httpx.AsyncClient() as client: + await client.delete( + f"{JUPYTER_HTTP_URL}/api/kernels/{kernel_id}", + timeout=10.0 + ) + logger.info(f"Removed kernel: {kernel_id}") + except Exception as e: + logger.warning(f"Error removing kernel {kernel_id}: {e}") + + if kernel_id in self.kernels: + del self.kernels[kernel_id] + if kernel_id in self.busy_kernels: + self.busy_kernels.remove(kernel_id) + + async def _check_kernel_health(self, kernel_id: str) -> bool: + """Check if a kernel is healthy by sending a simple command""" + try: + jupyter_ws_url = f"{JUPYTER_WS_URL}/api/kernels/{kernel_id}/channels" + async with websockets.connect( + jupyter_ws_url, + ping_interval=WEBSOCKET_PING_INTERVAL, + ping_timeout=WEBSOCKET_PING_TIMEOUT + ) as ws: + # Send simple health check command + msg_id, request_json = create_jupyter_request("1+1") + await ws.send(request_json) + + # Wait for response with timeout + start_time = time.time() + while time.time() - start_time < 10: # 10 second timeout for health check + try: + message_str = await asyncio.wait_for(ws.recv(), timeout=2.0) + message_data = json.loads(message_str) + parent_msg_id = message_data.get("parent_header", {}).get("msg_id") + + if parent_msg_id == msg_id: + msg_type = message_data.get("header", {}).get("msg_type") + if msg_type == "status" and message_data.get("content", {}).get("execution_state") == "idle": + return True + except asyncio.TimeoutError: + continue + return False + except Exception as e: + logger.warning(f"Health check failed for kernel {kernel_id}: {e}") + return False + + async def _health_check_loop(self): + """Background task to monitor kernel health""" + while True: + try: + await asyncio.sleep(KERNEL_HEALTH_CHECK_INTERVAL) + async with self.lock: + unhealthy_kernels = [] + for kernel_id, kernel_info in self.kernels.items(): + if kernel_info.needs_health_check() and kernel_id not in self.busy_kernels: + if await self._check_kernel_health(kernel_id): + kernel_info.last_health_check = datetime.now() + kernel_info.state = KernelState.HEALTHY + else: + kernel_info.state = KernelState.UNRESPONSIVE + unhealthy_kernels.append(kernel_id) + + # Remove unhealthy kernels and create replacements + for kernel_id in unhealthy_kernels: + logger.warning(f"Removing unhealthy kernel: {kernel_id}") + await self._remove_kernel(kernel_id) + # Create replacement if below minimum + if len(self.kernels) < MIN_KERNELS: + new_kernel_id = await self._create_new_kernel() + if new_kernel_id: + self.kernels[new_kernel_id] = KernelInfo(kernel_id=new_kernel_id) + except Exception as e: + logger.error(f"Error in health check loop: {e}") diff --git a/server.py b/server.py index d2b4e10..938721b 100644 --- a/server.py +++ b/server.py @@ -1,26 +1,20 @@ -# --- IMPORTS --- + + import asyncio -import base64 -import binascii -import json import logging import os import pathlib import time -import uuid -from typing import Dict, Optional, Set -from dataclasses import dataclass, field -from enum import Enum -from datetime import datetime, timedelta - -import aiofiles +import json import websockets -import httpx -# Import Context for progress reporting + from mcp.server.fastmcp import FastMCP, Context from playwright.async_api import async_playwright from bs4 import BeautifulSoup -import socket + +from kernel_manager import KernelPool, NoKernelAvailableError, KernelExecutionError, KernelTimeoutError, JUPYTER_WS_URL, WEBSOCKET_PING_INTERVAL, WEBSOCKET_PING_TIMEOUT, WEBSOCKET_TIMEOUT, MAX_RETRY_ATTEMPTS, RETRY_BACKOFF_BASE, KernelError +from utils import create_jupyter_request + # --- CONFIGURATION & SETUP --- logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" @@ -30,315 +24,15 @@ # Initialize the MCP server with a descriptive name for the toolset mcp = FastMCP("CodeRunner") -# Kernel pool configuration -MAX_KERNELS = 5 -MIN_KERNELS = 2 -KERNEL_TIMEOUT = 300 # 5 minutes -KERNEL_HEALTH_CHECK_INTERVAL = 30 # 30 seconds -MAX_RETRY_ATTEMPTS = 3 -RETRY_BACKOFF_BASE = 2 # exponential backoff base - -# Jupyter connection settings -JUPYTER_WS_URL = "ws://127.0.0.1:8888" -JUPYTER_HTTP_URL = "http://127.0.0.1:8888" - -# Enhanced WebSocket settings -WEBSOCKET_TIMEOUT = 600 # 10 minutes for long operations -WEBSOCKET_PING_INTERVAL = 30 -WEBSOCKET_PING_TIMEOUT = 10 - -# Directory configuration (ensure this matches your Jupyter/Docker setup) -# This directory must be accessible by both this script and the Jupyter kernel. +# Directory configuration SHARED_DIR = pathlib.Path("/app/uploads") SHARED_DIR.mkdir(exist_ok=True) KERNEL_ID_FILE_PATH = SHARED_DIR / "python_kernel_id.txt" -def resolve_with_system_dns(hostname): - try: - return socket.gethostbyname(hostname) - except socket.gaierror as e: - print(f"Error resolving {hostname}: {e}") - return None - PLAYWRIGHT_WS_URL =f"ws://127.0.0.1:3000/" -# --- CUSTOM EXCEPTIONS --- - -class KernelError(Exception): - """Base exception for kernel-related errors""" - pass - -class NoKernelAvailableError(KernelError): - """Raised when no kernels are available in the pool""" - pass - -class KernelExecutionError(KernelError): - """Raised when kernel execution fails""" - pass - -class KernelTimeoutError(KernelError): - """Raised when kernel operation times out""" - pass - -# --- KERNEL MANAGEMENT CLASSES --- - -class KernelState(Enum): - HEALTHY = "healthy" - BUSY = "busy" - UNRESPONSIVE = "unresponsive" - FAILED = "failed" - -@dataclass -class KernelInfo: - kernel_id: str - state: KernelState = KernelState.HEALTHY - last_used: datetime = field(default_factory=datetime.now) - last_health_check: datetime = field(default_factory=datetime.now) - current_operation: Optional[str] = None - failure_count: int = 0 - - def is_available(self) -> bool: - return self.state == KernelState.HEALTHY - - def needs_health_check(self) -> bool: - return datetime.now() - self.last_health_check > timedelta(seconds=KERNEL_HEALTH_CHECK_INTERVAL) - -class KernelPool: - def __init__(self): - self.kernels: Dict[str, KernelInfo] = {} - self.lock = asyncio.Lock() - self.busy_kernels: Set[str] = set() - self._initialized = False - self._health_check_task: Optional[asyncio.Task] = None - - async def initialize(self): - """Initialize the kernel pool with minimum number of kernels""" - if self._initialized: - return - - async with self.lock: - logger.info("Initializing kernel pool...") - - # Try to use existing kernel first - existing_kernel = await self._get_existing_kernel() - if existing_kernel: - self.kernels[existing_kernel] = KernelInfo(kernel_id=existing_kernel) - logger.info(f"Added existing kernel to pool: {existing_kernel}") - - # Create additional kernels to reach minimum - while len(self.kernels) < MIN_KERNELS: - kernel_id = await self._create_new_kernel() - if kernel_id: - self.kernels[kernel_id] = KernelInfo(kernel_id=kernel_id) - logger.info(f"Created new kernel: {kernel_id}") - else: - logger.warning("Failed to create minimum number of kernels") - break - - self._initialized = True - # Start health check background task - self._health_check_task = asyncio.create_task(self._health_check_loop()) - logger.info(f"Kernel pool initialized with {len(self.kernels)} kernels") - - async def get_available_kernel(self) -> Optional[str]: - """Get an available kernel from the pool""" - if not self._initialized: - await self.initialize() - - async with self.lock: - # Find healthy, available kernel - for kernel_id, kernel_info in self.kernels.items(): - if kernel_info.is_available() and kernel_id not in self.busy_kernels: - self.busy_kernels.add(kernel_id) - kernel_info.state = KernelState.BUSY - kernel_info.last_used = datetime.now() - logger.info(f"Assigned kernel {kernel_id} to operation") - return kernel_id - - # No available kernels, try to create a new one if under limit - if len(self.kernels) < MAX_KERNELS: - kernel_id = await self._create_new_kernel() - if kernel_id: - kernel_info = KernelInfo(kernel_id=kernel_id, state=KernelState.BUSY) - self.kernels[kernel_id] = kernel_info - self.busy_kernels.add(kernel_id) - logger.info(f"Created and assigned new kernel: {kernel_id}") - return kernel_id - - logger.warning("No available kernels in pool") - return None - - async def release_kernel(self, kernel_id: str, failed: bool = False): - """Release a kernel back to the pool""" - async with self.lock: - if kernel_id in self.busy_kernels: - self.busy_kernels.remove(kernel_id) - - if kernel_id in self.kernels: - kernel_info = self.kernels[kernel_id] - if failed: - kernel_info.failure_count += 1 - kernel_info.state = KernelState.FAILED - logger.warning(f"Kernel {kernel_id} marked as failed (failures: {kernel_info.failure_count})") - - # Remove failed kernel if it has too many failures - if kernel_info.failure_count >= MAX_RETRY_ATTEMPTS: - await self._remove_kernel(kernel_id) - # Create replacement kernel - new_kernel_id = await self._create_new_kernel() - if new_kernel_id: - self.kernels[new_kernel_id] = KernelInfo(kernel_id=new_kernel_id) - else: - kernel_info.state = KernelState.HEALTHY - kernel_info.current_operation = None - logger.info(f"Released kernel {kernel_id} back to pool") - - async def _get_existing_kernel(self) -> Optional[str]: - """Try to get kernel ID from existing file""" - try: - async with aiofiles.open(KERNEL_ID_FILE_PATH, mode='r') as f: - kernel_id = (await f.read()).strip() - if kernel_id and await self._check_kernel_health(kernel_id): - return kernel_id - except FileNotFoundError: - # This is a normal case if the server is starting for the first time. - pass - except Exception as e: - logger.warning(f"Could not read or validate existing kernel from {KERNEL_ID_FILE_PATH}: {e}") - return None - - async def _create_new_kernel(self) -> Optional[str]: - """Create a new Jupyter kernel""" - try: - async with httpx.AsyncClient() as client: - response = await client.post( - f"{JUPYTER_HTTP_URL}/api/kernels", - json={"name": "python3"}, - timeout=30.0 - ) - if response.status_code == 201: - kernel_data = response.json() - kernel_id = kernel_data["id"] - logger.info(f"Created new kernel: {kernel_id}") - return kernel_id - else: - logger.error(f"Failed to create kernel: {response.status_code}") - except Exception as e: - logger.error(f"Error creating kernel: {e}") - return None - - async def _remove_kernel(self, kernel_id: str): - """Remove and shutdown a kernel""" - try: - async with httpx.AsyncClient() as client: - await client.delete( - f"{JUPYTER_HTTP_URL}/api/kernels/{kernel_id}", - timeout=10.0 - ) - logger.info(f"Removed kernel: {kernel_id}") - except Exception as e: - logger.warning(f"Error removing kernel {kernel_id}: {e}") - - if kernel_id in self.kernels: - del self.kernels[kernel_id] - if kernel_id in self.busy_kernels: - self.busy_kernels.remove(kernel_id) - - async def _check_kernel_health(self, kernel_id: str) -> bool: - """Check if a kernel is healthy by sending a simple command""" - try: - jupyter_ws_url = f"{JUPYTER_WS_URL}/api/kernels/{kernel_id}/channels" - async with websockets.connect( - jupyter_ws_url, - ping_interval=WEBSOCKET_PING_INTERVAL, - ping_timeout=WEBSOCKET_PING_TIMEOUT - ) as ws: - # Send simple health check command - msg_id, request_json = create_jupyter_request("1+1") - await ws.send(request_json) - - # Wait for response with timeout - start_time = time.time() - while time.time() - start_time < 10: # 10 second timeout for health check - try: - message_str = await asyncio.wait_for(ws.recv(), timeout=2.0) - message_data = json.loads(message_str) - parent_msg_id = message_data.get("parent_header", {}).get("msg_id") - - if parent_msg_id == msg_id: - msg_type = message_data.get("header", {}).get("msg_type") - if msg_type == "status" and message_data.get("content", {}).get("execution_state") == "idle": - return True - except asyncio.TimeoutError: - continue - return False - except Exception as e: - logger.warning(f"Health check failed for kernel {kernel_id}: {e}") - return False - - async def _health_check_loop(self): - """Background task to monitor kernel health""" - while True: - try: - await asyncio.sleep(KERNEL_HEALTH_CHECK_INTERVAL) - async with self.lock: - unhealthy_kernels = [] - for kernel_id, kernel_info in self.kernels.items(): - if kernel_info.needs_health_check() and kernel_id not in self.busy_kernels: - if await self._check_kernel_health(kernel_id): - kernel_info.last_health_check = datetime.now() - kernel_info.state = KernelState.HEALTHY - else: - kernel_info.state = KernelState.UNRESPONSIVE - unhealthy_kernels.append(kernel_id) - - # Remove unhealthy kernels and create replacements - for kernel_id in unhealthy_kernels: - logger.warning(f"Removing unhealthy kernel: {kernel_id}") - await self._remove_kernel(kernel_id) - # Create replacement if below minimum - if len(self.kernels) < MIN_KERNELS: - new_kernel_id = await self._create_new_kernel() - if new_kernel_id: - self.kernels[new_kernel_id] = KernelInfo(kernel_id=new_kernel_id) - except Exception as e: - logger.error(f"Error in health check loop: {e}") - # Global kernel pool instance -kernel_pool = KernelPool() - - - -# --- HELPER FUNCTION --- -def create_jupyter_request(code: str) -> tuple[str, str]: - """ - Creates a Jupyter execute_request message. - Returns a tuple: (msg_id, json_payload_string) - """ - msg_id = uuid.uuid4().hex - session_id = uuid.uuid4().hex - - request = { - "header": { - "msg_id": msg_id, - "username": "mcp_client", - "session": session_id, - "msg_type": "execute_request", - "version": "5.3", - }, - "parent_header": {}, - "metadata": {}, - "content": { - "code": code, - "silent": False, - "store_history": False, - "user_expressions": {}, - "allow_stdin": False, - "stop_on_error": True, - }, - "buffers": [], - } - return msg_id, json.dumps(request) +kernel_pool = KernelPool(KERNEL_ID_FILE_PATH) # --- ENHANCED EXECUTION WITH RETRY LOGIC --- @@ -531,4 +225,19 @@ async def navigate_and_get_all_visible_text(url: str) -> str: # Use the streamable_http_app as it's the modern standard -app = mcp.streamable_http_app() \ No newline at end of file +app = mcp.streamable_http_app() + +# --- MAIN ENTRY POINT --- +if __name__ == "__main__": + import uvicorn + + # Start the kernel pool initialization in the background + asyncio.create_task(kernel_pool.initialize()) + + # Start the FastAPI server + uvicorn.run( + app, + host=os.getenv("FASTMCP_HOST", "0.0.0.0"), + port=int(os.getenv("FASTMCP_PORT", "8222")), + log_level="info", + ) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..40c2435 --- /dev/null +++ b/utils.py @@ -0,0 +1,33 @@ + +import uuid +import json + +def create_jupyter_request(code: str) -> tuple[str, str]: + """ + Creates a Jupyter execute_request message. + Returns a tuple: (msg_id, json_payload_string) + """ + msg_id = uuid.uuid4().hex + session_id = uuid.uuid4().hex + + request = { + "header": { + "msg_id": msg_id, + "username": "mcp_client", + "session": session_id, + "msg_type": "execute_request", + "version": "5.3", + }, + "parent_header": {}, + "metadata": {}, + "content": { + "code": code, + "silent": False, + "store_history": False, + "user_expressions": {}, + "allow_stdin": False, + "stop_on_error": True, + }, + "buffers": [], + } + return msg_id, json.dumps(request)