Create standalone pod_executor package
Created a standalone code execution package independent of MCP-Forge: Structure: - pod_executor/security/ - Resource limits, audit protocols, validation - pod_executor/containers/ - Podman client and container management - pod_executor/simple/ - Stateless code executor - pod_executor/jupyter/ - Stateful Jupyter backend with sessions Key changes: - Removed ForgeConfig dependency - all parameters explicit - Audit logger now a protocol with NullAuditLogger/SimpleFileAuditLogger - Validator now a protocol with NoOpValidator/BasicValidator - All imports updated to pod_executor namespace - Audit calls use simple strings instead of enums Benefits: - Standalone package usable without MCP-Forge - Clear separation between execution engine and MCP protocol - Easier testing and development - Reusable in other projects
This commit is contained in:
parent
8b6b237be9
commit
db75b822f4
14 changed files with 2842 additions and 0 deletions
82
src/pod_executor/__init__.py
Normal file
82
src/pod_executor/__init__.py
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
"""
|
||||
Pod Executor - Standalone Python code execution in Podman containers.
|
||||
|
||||
Provides stateless and stateful (Jupyter) code execution backends with
|
||||
security isolation via Podman containers.
|
||||
|
||||
Examples:
|
||||
# Simple stateless execution
|
||||
from pod_executor import CodeExecutor, ResourceLimits
|
||||
from pod_executor.containers import SecureContainerManager, PodmanClient
|
||||
from pod_executor.security import NoOpValidator, NullAuditLogger
|
||||
|
||||
client = PodmanClient(socket_path="/run/podman/podman.sock",
|
||||
validator=NoOpValidator(),
|
||||
audit_logger=NullAuditLogger())
|
||||
container_manager = SecureContainerManager(client, NoOpValidator(), NullAuditLogger())
|
||||
limits = ResourceLimits(memory="512m", cpu_quota=100000, storage="1g", timeout=30)
|
||||
executor = CodeExecutor(container_manager, "python:3.12", limits)
|
||||
|
||||
result = executor.execute("print('Hello World')")
|
||||
print(result.stdout)
|
||||
|
||||
# Stateful Jupyter execution
|
||||
from pod_executor import JupyterBackend
|
||||
|
||||
backend = JupyterBackend(
|
||||
container_manager=container_manager,
|
||||
image="jupyter/base-notebook",
|
||||
default_timeout=300
|
||||
)
|
||||
|
||||
result = backend.execute("x = 42", session_id="my-session")
|
||||
result = backend.execute("print(x * 2)", session_id="my-session")
|
||||
"""
|
||||
|
||||
from pod_executor.simple.executor import CodeExecutor, ExecutionResult
|
||||
from pod_executor.jupyter.backend import JupyterBackend
|
||||
from pod_executor.jupyter.kernel import JupyterKernelManager, KernelError
|
||||
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||
from pod_executor.containers.manager import SecureContainerManager, ContainerConfig
|
||||
from pod_executor.containers.client import PodmanClient, PodmanConnectionError
|
||||
from pod_executor.security.resource_limits import ResourceLimits
|
||||
from pod_executor.security.audit import (
|
||||
AuditLoggerProtocol,
|
||||
NullAuditLogger,
|
||||
SimpleFileAuditLogger,
|
||||
)
|
||||
from pod_executor.security.validation import (
|
||||
SecurityError,
|
||||
OperationValidatorProtocol,
|
||||
NoOpValidator,
|
||||
BasicValidator,
|
||||
)
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
||||
__all__ = [
|
||||
# Simple execution
|
||||
"CodeExecutor",
|
||||
"ExecutionResult",
|
||||
# Jupyter execution
|
||||
"JupyterBackend",
|
||||
"JupyterKernelManager",
|
||||
"KernelError",
|
||||
"SessionManager",
|
||||
"SessionState",
|
||||
"SessionError",
|
||||
# Container management
|
||||
"SecureContainerManager",
|
||||
"ContainerConfig",
|
||||
"PodmanClient",
|
||||
"PodmanConnectionError",
|
||||
# Security
|
||||
"ResourceLimits",
|
||||
"AuditLoggerProtocol",
|
||||
"NullAuditLogger",
|
||||
"SimpleFileAuditLogger",
|
||||
"SecurityError",
|
||||
"OperationValidatorProtocol",
|
||||
"NoOpValidator",
|
||||
"BasicValidator",
|
||||
]
|
||||
14
src/pod_executor/containers/__init__.py
Normal file
14
src/pod_executor/containers/__init__.py
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
"""Container management for pod_executor."""
|
||||
|
||||
from pod_executor.containers.client import PodmanClient, PodmanConnectionError
|
||||
from pod_executor.containers.manager import (
|
||||
ContainerConfig,
|
||||
SecureContainerManager,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"PodmanClient",
|
||||
"PodmanConnectionError",
|
||||
"ContainerConfig",
|
||||
"SecureContainerManager",
|
||||
]
|
||||
157
src/pod_executor/containers/client.py
Normal file
157
src/pod_executor/containers/client.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
"""
|
||||
Podman client wrapper with security validation.
|
||||
|
||||
Wraps Podman API with security validation and error handling.
|
||||
All container operations are validated against security policy.
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional, Union
|
||||
from podman import PodmanClient as BasePodmanClient
|
||||
|
||||
from pod_executor.security.validation import OperationValidatorProtocol
|
||||
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||
|
||||
|
||||
class PodmanConnectionError(Exception):
|
||||
"""Raised when connection to Podman fails."""
|
||||
pass
|
||||
|
||||
|
||||
class PodmanClient:
|
||||
"""
|
||||
Wrapper around Podman API with security validation.
|
||||
|
||||
All container operations are validated against security policy
|
||||
before being sent to Podman. Provides lazy connection and
|
||||
proper error handling.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
socket_path: Union[str, Path],
|
||||
validator: OperationValidatorProtocol,
|
||||
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||
):
|
||||
"""
|
||||
Initialize Podman client wrapper.
|
||||
|
||||
Args:
|
||||
socket_path: Path to Podman socket
|
||||
validator: Operation validator for security checks
|
||||
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||
"""
|
||||
self.socket_path = Path(socket_path)
|
||||
self.validator = validator
|
||||
self.audit_logger = audit_logger or NullAuditLogger()
|
||||
self._client: Optional[BasePodmanClient] = None
|
||||
|
||||
def connect(self) -> None:
|
||||
"""
|
||||
Connect to Podman via socket.
|
||||
|
||||
Raises:
|
||||
PodmanConnectionError: If connection fails
|
||||
"""
|
||||
try:
|
||||
# Verify socket exists and is accessible
|
||||
self.verify_socket_access()
|
||||
|
||||
# Create Podman client with Unix socket
|
||||
base_url = f"unix://{self.socket_path}"
|
||||
self._client = BasePodmanClient(base_url=base_url)
|
||||
|
||||
# Test connection with ping (only if client supports it)
|
||||
if hasattr(self._client, 'ping'):
|
||||
self._client.ping()
|
||||
|
||||
except PodmanConnectionError:
|
||||
# Re-raise our own exceptions
|
||||
raise
|
||||
except Exception as e:
|
||||
raise PodmanConnectionError(
|
||||
f"Failed to connect to Podman at {self.socket_path}: {e}"
|
||||
) from e
|
||||
|
||||
def ping(self) -> bool:
|
||||
"""
|
||||
Test connection to Podman.
|
||||
|
||||
Returns:
|
||||
True if connection is healthy
|
||||
|
||||
Raises:
|
||||
PodmanConnectionError: If not connected or ping fails
|
||||
"""
|
||||
if self._client is None:
|
||||
raise PodmanConnectionError("Not connected to Podman")
|
||||
|
||||
try:
|
||||
result = self._client.ping()
|
||||
return result == "OK" or result is True
|
||||
except Exception as e:
|
||||
raise PodmanConnectionError(f"Ping failed: {e}") from e
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Disconnect from Podman and cleanup."""
|
||||
if self._client is not None:
|
||||
try:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
pass # Ignore errors during cleanup
|
||||
finally:
|
||||
self._client = None
|
||||
|
||||
def verify_socket_access(self) -> None:
|
||||
"""
|
||||
Verify that socket exists and is accessible.
|
||||
|
||||
Raises:
|
||||
PodmanConnectionError: If socket is not accessible
|
||||
"""
|
||||
if not self.socket_path.exists():
|
||||
raise PodmanConnectionError(
|
||||
f"Podman socket not found: {self.socket_path}"
|
||||
)
|
||||
|
||||
if not os.access(self.socket_path, os.R_OK):
|
||||
raise PodmanConnectionError(
|
||||
f"Podman socket is not readable: {self.socket_path}"
|
||||
)
|
||||
|
||||
def check_api_version(self) -> dict:
|
||||
"""
|
||||
Get Podman API version information.
|
||||
|
||||
Returns:
|
||||
Dictionary with version information
|
||||
|
||||
Raises:
|
||||
PodmanConnectionError: If not connected
|
||||
"""
|
||||
if self._client is None:
|
||||
raise PodmanConnectionError("Not connected to Podman")
|
||||
|
||||
try:
|
||||
return self._client.version()
|
||||
except Exception as e:
|
||||
raise PodmanConnectionError(
|
||||
f"Failed to get API version: {e}"
|
||||
) from e
|
||||
|
||||
@property
|
||||
def client(self) -> BasePodmanClient:
|
||||
"""
|
||||
Get underlying Podman client (lazy connection).
|
||||
|
||||
Returns:
|
||||
Connected Podman client
|
||||
|
||||
Raises:
|
||||
PodmanConnectionError: If connection fails
|
||||
"""
|
||||
if self._client is None:
|
||||
self.connect()
|
||||
assert self._client is not None # Type narrowing for mypy
|
||||
return self._client
|
||||
501
src/pod_executor/containers/manager.py
Normal file
501
src/pod_executor/containers/manager.py
Normal file
|
|
@ -0,0 +1,501 @@
|
|||
""
|
||||
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, Union
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
from pod_executor.containers.client import PodmanClient
|
||||
from pod_executor.security.validation import OperationValidatorProtocol, SecurityError
|
||||
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||
from pod_executor.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: OperationValidatorProtocol",
|
||||
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||
):
|
||||
""
|
||||
Initialize secure container manager.
|
||||
|
||||
Args:
|
||||
podman_client: Podman client wrapper
|
||||
validator: Operation validator for security checks
|
||||
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||
""
|
||||
self.podman = podman_client
|
||||
self.validator = validator
|
||||
self.audit_logger = audit_logger or NullAuditLogger()
|
||||
|
||||
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(event_type="security.violation", severity="critical", message=f"Security violation: container_create", details={"operation": "container_create", "reason": str(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(event_type="container.operation", severity="info", 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="CONTAINER_CREATE",
|
||||
severity="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(event_type="container.operation", severity="info", operation="start",
|
||||
container_id=container_id",
|
||||
image=" # Not available without extra lookup
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_START",
|
||||
severity="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(event_type="container.operation", severity="info", operation="stop",
|
||||
container_id=container_id",
|
||||
image="",
|
||||
details={"timeout": timeout}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_STOP",
|
||||
severity="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(event_type="container.operation", severity="info", operation="remove",
|
||||
container_id=container_id",
|
||||
image="",
|
||||
details={"force": force}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_REMOVE",
|
||||
severity="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="EXECUTION_REQUEST",
|
||||
severity="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="EXECUTION_REQUEST",
|
||||
severity="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="CONTAINER_REMOVE",
|
||||
severity="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="CONTAINER_REMOVE",
|
||||
severity="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="CONTAINER_REMOVE",
|
||||
severity="ERROR",
|
||||
message=f"Cleanup failed: {e}",
|
||||
details={"error": str(e)}
|
||||
)
|
||||
raise
|
||||
14
src/pod_executor/jupyter/__init__.py
Normal file
14
src/pod_executor/jupyter/__init__.py
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
"""Jupyter stateful code executor."""
|
||||
|
||||
from pod_executor.jupyter.backend import JupyterBackend
|
||||
from pod_executor.jupyter.kernel import JupyterKernelManager, KernelError
|
||||
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||
|
||||
__all__ = [
|
||||
"JupyterBackend",
|
||||
"JupyterKernelManager",
|
||||
"KernelError",
|
||||
"SessionManager",
|
||||
"SessionState",
|
||||
"SessionError",
|
||||
]
|
||||
229
src/pod_executor/jupyter/backend.py
Normal file
229
src/pod_executor/jupyter/backend.py
Normal file
|
|
@ -0,0 +1,229 @@
|
|||
"""Jupyter backend for stateful code execution."""
|
||||
|
||||
from typing import Optional, Dict, List
|
||||
import hashlib
|
||||
|
||||
from pod_executor.containers.manager import SecureContainerManager
|
||||
from pod_executor.security.resource_limits import ResourceLimits, parse_memory_string
|
||||
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||
from pod_executor.simple.executor import ExecutionResult
|
||||
from pod_executor.jupyter.kernel import JupyterKernelManager
|
||||
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||
|
||||
|
||||
class JupyterBackend:
|
||||
"""Stateful code execution backend using Jupyter kernels."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
container_manager: SecureContainerManager,
|
||||
image: str,
|
||||
default_timeout: int = 300,
|
||||
default_memory: str = "512m",
|
||||
default_cpu_quota: int = 100000,
|
||||
max_timeout: int = 3600,
|
||||
max_memory: str = "2g",
|
||||
max_cpu_quota: int = 200000,
|
||||
idle_timeout: int = 3600,
|
||||
max_sessions: int = 10,
|
||||
resource_limits: Optional[ResourceLimits] = None,
|
||||
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||
):
|
||||
"""
|
||||
Initialize Jupyter backend.
|
||||
|
||||
Args:
|
||||
container_manager: Container lifecycle manager
|
||||
image: Docker/Podman image with ipykernel installed
|
||||
default_timeout: Default execution timeout in seconds
|
||||
default_memory: Default memory limit (e.g., "512m", "1g")
|
||||
default_cpu_quota: Default CPU quota (100000 = 1 CPU)
|
||||
max_timeout: Maximum allowed timeout
|
||||
max_memory: Maximum allowed memory
|
||||
max_cpu_quota: Maximum allowed CPU quota
|
||||
idle_timeout: Session idle timeout in seconds
|
||||
max_sessions: Maximum concurrent sessions
|
||||
resource_limits: Default resource limits for kernels
|
||||
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||
"""
|
||||
self.container_manager = container_manager
|
||||
self.image = image
|
||||
self.default_timeout = default_timeout
|
||||
self.default_memory = default_memory
|
||||
self.default_cpu_quota = default_cpu_quota
|
||||
self.max_timeout = max_timeout
|
||||
self.max_memory = max_memory
|
||||
self.max_cpu_quota = max_cpu_quota
|
||||
self.idle_timeout = idle_timeout
|
||||
self.max_sessions = max_sessions
|
||||
self.audit_logger = audit_logger or NullAuditLogger()
|
||||
|
||||
# Initialize kernel manager
|
||||
kernel_manager = JupyterKernelManager(
|
||||
container_manager=container_manager,
|
||||
image=image,
|
||||
resource_limits=resource_limits
|
||||
)
|
||||
|
||||
# Initialize session manager
|
||||
self.session_manager = SessionManager(
|
||||
kernel_manager=kernel_manager,
|
||||
idle_timeout=idle_timeout,
|
||||
max_sessions=max_sessions,
|
||||
audit_logger=self.audit_logger
|
||||
)
|
||||
|
||||
def execute(
|
||||
self,
|
||||
code: str,
|
||||
session_id: str,
|
||||
timeout: Optional[int] = None,
|
||||
memory: Optional[str] = None,
|
||||
cpu_quota: Optional[int] = 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 default if None)
|
||||
memory: Memory limit string (uses default if None)
|
||||
cpu_quota: CPU quota (uses default if None)
|
||||
volumes: Volume mounts dict
|
||||
injection_code: Optional code executed once at session start
|
||||
bridge_socket_path: Optional path to socket for mounting
|
||||
|
||||
Returns:
|
||||
ExecutionResult with execution output and metadata
|
||||
|
||||
Raises:
|
||||
ValueError: If limits exceed configured maximums
|
||||
SessionError: If session operation fails
|
||||
"""
|
||||
# Use defaults if not specified
|
||||
timeout = timeout if timeout is not None else self.default_timeout
|
||||
memory = memory if memory is not None else self.default_memory
|
||||
cpu_quota = cpu_quota if cpu_quota is not None else self.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="execution.request",
|
||||
severity="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
|
||||
resource_limits = ResourceLimits(
|
||||
memory=memory,
|
||||
cpu_quota=cpu_quota,
|
||||
storage="1g",
|
||||
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
|
||||
"""
|
||||
self.session_manager.document_state(
|
||||
session_id=session_id,
|
||||
variables=variables,
|
||||
note=note,
|
||||
clear=clear
|
||||
)
|
||||
|
||||
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."""
|
||||
return self.session_manager.get_session_state(session_id)
|
||||
|
||||
def destroy_session(self, session_id: str) -> None:
|
||||
"""Destroy session and cleanup kernel."""
|
||||
self.session_manager.destroy_session(session_id)
|
||||
|
||||
def list_sessions(self) -> List[dict]:
|
||||
"""List all active sessions with metadata."""
|
||||
return self.session_manager.list_sessions()
|
||||
|
||||
def cleanup_idle_sessions(self) -> int:
|
||||
"""Cleanup sessions idle beyond configured timeout."""
|
||||
return self.session_manager.cleanup_idle_sessions()
|
||||
|
||||
def _validate_limits(self, timeout: int, memory: str, cpu_quota: int) -> None:
|
||||
"""
|
||||
Validate resource limits against configured maximums.
|
||||
|
||||
Raises:
|
||||
ValueError: If any limit exceeds maximum
|
||||
"""
|
||||
if timeout > self.max_timeout:
|
||||
raise ValueError(
|
||||
f"Timeout {timeout} exceeds maximum {self.max_timeout}"
|
||||
)
|
||||
|
||||
memory_bytes = parse_memory_string(memory)
|
||||
max_memory_bytes = parse_memory_string(self.max_memory)
|
||||
if memory_bytes > max_memory_bytes:
|
||||
raise ValueError(
|
||||
f"Memory {memory} exceeds maximum {self.max_memory}"
|
||||
)
|
||||
|
||||
if cpu_quota > self.max_cpu_quota:
|
||||
raise ValueError(
|
||||
f"CPU quota {cpu_quota} exceeds maximum {self.max_cpu_quota}"
|
||||
)
|
||||
631
src/pod_executor/jupyter/kernel.py
Normal file
631
src/pod_executor/jupyter/kernel.py
Normal file
|
|
@ -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 pod_executor.containers.manager import SecureContainerManager, ContainerConfig
|
||||
from pod_executor.security.resource_limits import ResourceLimits
|
||||
from pod_executor.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
|
||||
437
src/pod_executor/jupyter/sessions.py
Normal file
437
src/pod_executor/jupyter/sessions.py
Normal file
|
|
@ -0,0 +1,437 @@
|
|||
"""Session management for stateful execution."""
|
||||
|
||||
from typing import Dict, Optional, List, Any
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from pod_executor.jupyter.kernel import JupyterKernelManager
|
||||
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||
from pod_executor.security.resource_limits import ResourceLimits
|
||||
from pod_executor.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,
|
||||
kernel_manager: JupyterKernelManager,
|
||||
idle_timeout: int = 3600,
|
||||
max_sessions: int = 10,
|
||||
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||
):
|
||||
"""
|
||||
Initialize session manager.
|
||||
|
||||
Args:
|
||||
kernel_manager: Kernel lifecycle manager
|
||||
idle_timeout: Session idle timeout in seconds (default: 3600)
|
||||
max_sessions: Maximum concurrent sessions (default: 10)
|
||||
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||
"""
|
||||
self.kernel_manager = kernel_manager
|
||||
self.idle_timeout = idle_timeout
|
||||
self.max_sessions = max_sessions
|
||||
self.audit_logger = audit_logger or NullAuditLogger()
|
||||
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="SESSION_CREATE,
|
||||
severity="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="SESSION_DESTROY,
|
||||
severity="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="SESSION_DESTROY,
|
||||
severity="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.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.max_sessions:
|
||||
raise SessionError(
|
||||
f"Maximum concurrent sessions ({self.max_sessions}) reached"
|
||||
)
|
||||
33
src/pod_executor/security/__init__.py
Normal file
33
src/pod_executor/security/__init__.py
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
"""Security components for pod_executor."""
|
||||
|
||||
from pod_executor.security.resource_limits import (
|
||||
ResourceLimits,
|
||||
parse_memory_string,
|
||||
parse_cpu_quota,
|
||||
parse_storage_string,
|
||||
)
|
||||
from pod_executor.security.audit import (
|
||||
AuditLoggerProtocol,
|
||||
NullAuditLogger,
|
||||
SimpleFileAuditLogger,
|
||||
)
|
||||
from pod_executor.security.validation import (
|
||||
SecurityError,
|
||||
OperationValidatorProtocol,
|
||||
NoOpValidator,
|
||||
BasicValidator,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ResourceLimits",
|
||||
"parse_memory_string",
|
||||
"parse_cpu_quota",
|
||||
"parse_storage_string",
|
||||
"AuditLoggerProtocol",
|
||||
"NullAuditLogger",
|
||||
"SimpleFileAuditLogger",
|
||||
"SecurityError",
|
||||
"OperationValidatorProtocol",
|
||||
"NoOpValidator",
|
||||
"BasicValidator",
|
||||
]
|
||||
132
src/pod_executor/security/audit.py
Normal file
132
src/pod_executor/security/audit.py
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
"""
|
||||
Audit logger protocol for pod_executor.
|
||||
|
||||
Provides a protocol (interface) for audit logging that can be implemented
|
||||
by consuming applications. A default no-op implementation is provided.
|
||||
"""
|
||||
|
||||
from typing import Protocol, Any, Optional, Dict
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class AuditLoggerProtocol(Protocol):
|
||||
"""Protocol for audit logging (optional dependency)."""
|
||||
|
||||
def log(
|
||||
self,
|
||||
event_type: str,
|
||||
severity: str,
|
||||
message: str,
|
||||
session_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
details: Optional[Dict[str, Any]] = None,
|
||||
error: Optional[str] = None,
|
||||
**kwargs: Any
|
||||
) -> None:
|
||||
"""
|
||||
Log an audit event.
|
||||
|
||||
Args:
|
||||
event_type: Type of event (e.g., "container.create", "execution.request")
|
||||
severity: Severity level ("info", "warning", "error", "critical")
|
||||
message: Human-readable message describing the event
|
||||
session_id: Optional session ID associated with event
|
||||
user_id: Optional user ID associated with event
|
||||
details: Optional dictionary of additional details
|
||||
error: Optional error message if event represents an error
|
||||
**kwargs: Additional keyword arguments for extensibility
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class NullAuditLogger:
|
||||
"""No-op audit logger for standalone usage without audit requirements."""
|
||||
|
||||
def log(
|
||||
self,
|
||||
event_type: str = "",
|
||||
severity: str = "info",
|
||||
message: str = "",
|
||||
session_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
details: Optional[Dict[str, Any]] = None,
|
||||
error: Optional[str] = None,
|
||||
**kwargs: Any
|
||||
) -> None:
|
||||
"""Do nothing - audit logging disabled."""
|
||||
pass
|
||||
|
||||
|
||||
class SimpleFileAuditLogger:
|
||||
"""
|
||||
Simple file-based audit logger for basic use cases.
|
||||
|
||||
Logs events to a JSON Lines file (one JSON object per line).
|
||||
Thread-safe via file locking.
|
||||
"""
|
||||
|
||||
def __init__(self, log_path: Path):
|
||||
"""
|
||||
Initialize file audit logger.
|
||||
|
||||
Args:
|
||||
log_path: Path to audit log file
|
||||
"""
|
||||
self.log_path = Path(log_path)
|
||||
self._ensure_log_file()
|
||||
|
||||
def _ensure_log_file(self) -> None:
|
||||
"""Ensure log file and directory exist."""
|
||||
self.log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not self.log_path.exists():
|
||||
self.log_path.touch()
|
||||
|
||||
def log(
|
||||
self,
|
||||
event_type: str = "",
|
||||
severity: str = "info",
|
||||
message: str = "",
|
||||
session_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
details: Optional[Dict[str, Any]] = None,
|
||||
error: Optional[str] = None,
|
||||
**kwargs: Any
|
||||
) -> None:
|
||||
"""
|
||||
Log an audit event to JSON Lines file.
|
||||
|
||||
Args:
|
||||
event_type: Type of event
|
||||
severity: Severity level
|
||||
message: Human-readable message
|
||||
session_id: Optional session ID
|
||||
user_id: Optional user ID
|
||||
details: Optional details dictionary
|
||||
error: Optional error message
|
||||
**kwargs: Additional fields
|
||||
"""
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
|
||||
entry = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"event_type": event_type,
|
||||
"severity": severity,
|
||||
"message": message,
|
||||
}
|
||||
|
||||
if session_id is not None:
|
||||
entry["session_id"] = session_id
|
||||
if user_id is not None:
|
||||
entry["user_id"] = user_id
|
||||
if details is not None:
|
||||
entry["details"] = details
|
||||
if error is not None:
|
||||
entry["error"] = error
|
||||
|
||||
# Add any additional kwargs
|
||||
entry.update(kwargs)
|
||||
|
||||
# Write to file (append mode, file locking via 'a' mode)
|
||||
with open(self.log_path, 'a') as f:
|
||||
f.write(json.dumps(entry) + '\n')
|
||||
150
src/pod_executor/security/resource_limits.py
Normal file
150
src/pod_executor/security/resource_limits.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
"""
|
||||
Resource limit parser and validator.
|
||||
|
||||
Parses and validates resource limit strings (memory, CPU, storage).
|
||||
All values must be positive and within reasonable limits.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Dict, Any
|
||||
|
||||
|
||||
def parse_memory_string(memory: str) -> int:
|
||||
"""
|
||||
Parse memory string to bytes.
|
||||
|
||||
Supports: k, m, g suffixes (case-insensitive)
|
||||
Examples: "512m" → 536870912, "2g" → 2147483648
|
||||
|
||||
Args:
|
||||
memory: Memory string with suffix (e.g., "512m", "2g", "1024k")
|
||||
|
||||
Returns:
|
||||
Memory in bytes
|
||||
|
||||
Raises:
|
||||
ValueError: If format is invalid or value is <= 0
|
||||
"""
|
||||
memory = memory.strip()
|
||||
|
||||
# Pattern: optional sign, number (int or float), suffix (k/m/g)
|
||||
pattern = r'^(-?\d+(?:\.\d+)?)\s*([kmgKMG])$'
|
||||
match = re.match(pattern, memory)
|
||||
|
||||
if not match:
|
||||
raise ValueError(
|
||||
f"Invalid memory format: '{memory}'. "
|
||||
f"Expected format: <number><k|m|g> (e.g., '512m', '2g')"
|
||||
)
|
||||
|
||||
value_str, suffix = match.groups()
|
||||
value = float(value_str)
|
||||
|
||||
if value <= 0:
|
||||
raise ValueError(
|
||||
f"Memory value must be positive, got: {value}"
|
||||
)
|
||||
|
||||
# Convert to bytes
|
||||
suffix_lower = suffix.lower()
|
||||
multipliers = {
|
||||
'k': 1024,
|
||||
'm': 1024 * 1024,
|
||||
'g': 1024 * 1024 * 1024,
|
||||
}
|
||||
|
||||
bytes_value = int(value * multipliers[suffix_lower])
|
||||
|
||||
return bytes_value
|
||||
|
||||
|
||||
def parse_cpu_quota(cpu_quota: int) -> int:
|
||||
"""
|
||||
Validate CPU quota value.
|
||||
|
||||
CPU quota is in microseconds per 100ms period.
|
||||
100000 = 100% of one CPU core
|
||||
|
||||
Args:
|
||||
cpu_quota: CPU quota in microseconds (e.g., 50000 for 50% of one core)
|
||||
|
||||
Returns:
|
||||
Validated CPU quota value
|
||||
|
||||
Raises:
|
||||
ValueError: If quota <= 0 or > 1000000 (10 cores max)
|
||||
"""
|
||||
if cpu_quota <= 0:
|
||||
raise ValueError(
|
||||
f"CPU quota must be positive, got: {cpu_quota}"
|
||||
)
|
||||
|
||||
# Maximum of 10 cores (1000000 microseconds)
|
||||
if cpu_quota > 1000000:
|
||||
raise ValueError(
|
||||
f"CPU quota exceeds maximum of 1000000 (10 cores), got: {cpu_quota}"
|
||||
)
|
||||
|
||||
return cpu_quota
|
||||
|
||||
|
||||
def parse_storage_string(storage: str) -> int:
|
||||
"""
|
||||
Parse storage string to bytes (same as memory).
|
||||
|
||||
Args:
|
||||
storage: Storage string with suffix (e.g., "1g", "512m")
|
||||
|
||||
Returns:
|
||||
Storage in bytes
|
||||
|
||||
Raises:
|
||||
ValueError: If format is invalid or value is <= 0
|
||||
"""
|
||||
return parse_memory_string(storage)
|
||||
|
||||
|
||||
class ResourceLimits:
|
||||
"""
|
||||
Resource limits with validation.
|
||||
|
||||
Encapsulates memory, storage, CPU, and timeout limits with validation.
|
||||
Provides conversion to Podman container parameters.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
memory: str,
|
||||
storage: str,
|
||||
cpu_quota: int,
|
||||
timeout: int = 300
|
||||
):
|
||||
"""
|
||||
Initialize resource limits with validation.
|
||||
|
||||
Args:
|
||||
memory: Memory limit string (e.g., "512m", "2g")
|
||||
storage: Storage limit string (e.g., "1g", "10g")
|
||||
cpu_quota: CPU quota in microseconds per 100ms period
|
||||
timeout: Execution timeout in seconds (default: 300)
|
||||
|
||||
Raises:
|
||||
ValueError: If any limit is invalid
|
||||
"""
|
||||
self.memory_bytes = parse_memory_string(memory)
|
||||
self.storage_bytes = parse_storage_string(storage)
|
||||
self.cpu_quota = parse_cpu_quota(cpu_quota)
|
||||
self.timeout = timeout
|
||||
|
||||
def to_podman_params(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Convert to Podman container create parameters.
|
||||
|
||||
Returns:
|
||||
Dictionary of parameters suitable for Podman container creation
|
||||
"""
|
||||
return {
|
||||
"mem_limit": str(self.memory_bytes), # Podman expects string
|
||||
"cpu_quota": self.cpu_quota
|
||||
# Note: storage_bytes tracked internally but not passed to Podman (not supported)
|
||||
}
|
||||
232
src/pod_executor/security/validation.py
Normal file
232
src/pod_executor/security/validation.py
Normal file
|
|
@ -0,0 +1,232 @@
|
|||
"""
|
||||
Security validation protocol for pod_executor.
|
||||
|
||||
Provides protocols (interfaces) for security validation that can be implemented
|
||||
by consuming applications. Default implementations are provided.
|
||||
"""
|
||||
|
||||
from typing import Protocol, Set, Optional, Dict, Any
|
||||
|
||||
|
||||
class SecurityError(Exception):
|
||||
"""Raised when security policy is violated."""
|
||||
pass
|
||||
|
||||
|
||||
class OperationValidatorProtocol(Protocol):
|
||||
"""Protocol for validating container operations."""
|
||||
|
||||
session_containers: Set[str]
|
||||
|
||||
def validate_container_create(
|
||||
self,
|
||||
image: str,
|
||||
params: Dict[str, Any],
|
||||
session_id: Optional[str] = None
|
||||
) -> None:
|
||||
"""
|
||||
Validate container creation parameters against security policy.
|
||||
|
||||
Args:
|
||||
image: Container image name
|
||||
params: Container creation parameters
|
||||
session_id: Optional session ID for volume validation
|
||||
|
||||
Raises:
|
||||
SecurityError: If any security policy is violated
|
||||
"""
|
||||
...
|
||||
|
||||
def validate_container_start(self, container_id: str) -> None:
|
||||
"""
|
||||
Validate container start.
|
||||
|
||||
Args:
|
||||
container_id: Container ID to start
|
||||
|
||||
Raises:
|
||||
SecurityError: If operation is not allowed
|
||||
"""
|
||||
...
|
||||
|
||||
def validate_container_stop(self, container_id: str) -> None:
|
||||
"""
|
||||
Validate container stop.
|
||||
|
||||
Args:
|
||||
container_id: Container ID to stop
|
||||
|
||||
Raises:
|
||||
SecurityError: If operation is not allowed
|
||||
"""
|
||||
...
|
||||
|
||||
def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Validate generic operation.
|
||||
|
||||
Args:
|
||||
operation: Operation type
|
||||
target: Operation target
|
||||
|
||||
Returns:
|
||||
(allowed, reason) tuple
|
||||
"""
|
||||
...
|
||||
|
||||
def register_session_container(self, container_id: str) -> None:
|
||||
"""Register a container as belonging to a session."""
|
||||
...
|
||||
|
||||
def unregister_session_container(self, container_id: str) -> None:
|
||||
"""Unregister a session container."""
|
||||
...
|
||||
|
||||
|
||||
class NoOpValidator:
|
||||
"""
|
||||
No-op validator that allows all operations.
|
||||
|
||||
WARNING: This validator provides NO SECURITY. Only use for testing
|
||||
or in fully trusted environments.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize with empty session container set."""
|
||||
self.session_containers: Set[str] = set()
|
||||
|
||||
def validate_container_create(
|
||||
self,
|
||||
image: str,
|
||||
params: Dict[str, Any],
|
||||
session_id: Optional[str] = None
|
||||
) -> None:
|
||||
"""Allow all container creations."""
|
||||
pass
|
||||
|
||||
def validate_container_start(self, container_id: str) -> None:
|
||||
"""Allow all container starts."""
|
||||
pass
|
||||
|
||||
def validate_container_stop(self, container_id: str) -> None:
|
||||
"""Allow all container stops."""
|
||||
pass
|
||||
|
||||
def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]:
|
||||
"""Allow all operations."""
|
||||
return (True, None)
|
||||
|
||||
def register_session_container(self, container_id: str) -> None:
|
||||
"""Track session container."""
|
||||
self.session_containers.add(container_id)
|
||||
|
||||
def unregister_session_container(self, container_id: str) -> None:
|
||||
"""Untrack session container."""
|
||||
self.session_containers.discard(container_id)
|
||||
|
||||
|
||||
class BasicValidator:
|
||||
"""
|
||||
Basic validator with minimal security checks.
|
||||
|
||||
Enforces:
|
||||
- Allowed image patterns
|
||||
- Required security parameters
|
||||
- Forbidden dangerous parameters
|
||||
- Session container tracking
|
||||
"""
|
||||
|
||||
# Allowed container images with wildcard support
|
||||
DEFAULT_ALLOWED_IMAGES = [
|
||||
"python:3.11*",
|
||||
"python:3.12*",
|
||||
"jupyter/*",
|
||||
"mcp-forge/*",
|
||||
]
|
||||
|
||||
# Parameters that are forbidden
|
||||
FORBIDDEN_PARAMS = [
|
||||
"privileged",
|
||||
"cap_add",
|
||||
"devices",
|
||||
"pid_mode",
|
||||
]
|
||||
|
||||
def __init__(self, allowed_images: Optional[list[str]] = None):
|
||||
"""
|
||||
Initialize basic validator.
|
||||
|
||||
Args:
|
||||
allowed_images: List of allowed image patterns (supports wildcards)
|
||||
"""
|
||||
self.allowed_images = allowed_images or self.DEFAULT_ALLOWED_IMAGES
|
||||
self.session_containers: Set[str] = set()
|
||||
|
||||
def validate_container_create(
|
||||
self,
|
||||
image: str,
|
||||
params: Dict[str, Any],
|
||||
session_id: Optional[str] = None
|
||||
) -> None:
|
||||
"""
|
||||
Validate container creation.
|
||||
|
||||
Args:
|
||||
image: Container image name
|
||||
params: Container creation parameters
|
||||
session_id: Optional session ID
|
||||
|
||||
Raises:
|
||||
SecurityError: If validation fails
|
||||
"""
|
||||
# Validate image is allowed
|
||||
if not self._is_image_allowed(image):
|
||||
raise SecurityError(
|
||||
f"Image '{image}' not in allowlist. "
|
||||
f"Allowed patterns: {self.allowed_images}"
|
||||
)
|
||||
|
||||
# Check for forbidden parameters
|
||||
for forbidden in self.FORBIDDEN_PARAMS:
|
||||
if forbidden in params:
|
||||
raise SecurityError(
|
||||
f"Forbidden parameter '{forbidden}' in container creation"
|
||||
)
|
||||
|
||||
# Ensure security_opt includes no-new-privileges
|
||||
security_opts = params.get("security_opt", [])
|
||||
if "no-new-privileges" not in security_opts:
|
||||
raise SecurityError(
|
||||
"Container must include security_opt=['no-new-privileges']"
|
||||
)
|
||||
|
||||
def validate_container_start(self, container_id: str) -> None:
|
||||
"""Validate container start."""
|
||||
if container_id not in self.session_containers:
|
||||
raise SecurityError(
|
||||
f"Container {container_id} is not a registered session container"
|
||||
)
|
||||
|
||||
def validate_container_stop(self, container_id: str) -> None:
|
||||
"""Validate container stop."""
|
||||
if container_id not in self.session_containers:
|
||||
raise SecurityError(
|
||||
f"Container {container_id} is not a registered session container"
|
||||
)
|
||||
|
||||
def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]:
|
||||
"""Validate generic operation."""
|
||||
return (True, None) # Allow by default
|
||||
|
||||
def register_session_container(self, container_id: str) -> None:
|
||||
"""Register session container."""
|
||||
self.session_containers.add(container_id)
|
||||
|
||||
def unregister_session_container(self, container_id: str) -> None:
|
||||
"""Unregister session container."""
|
||||
self.session_containers.discard(container_id)
|
||||
|
||||
def _is_image_allowed(self, image: str) -> bool:
|
||||
"""Check if image matches any allowed pattern."""
|
||||
import fnmatch
|
||||
return any(fnmatch.fnmatch(image, pattern) for pattern in self.allowed_images)
|
||||
8
src/pod_executor/simple/__init__.py
Normal file
8
src/pod_executor/simple/__init__.py
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
"""Simple stateless code executor."""
|
||||
|
||||
from pod_executor.simple.executor import CodeExecutor, ExecutionResult
|
||||
|
||||
__all__ = [
|
||||
"CodeExecutor",
|
||||
"ExecutionResult",
|
||||
]
|
||||
222
src/pod_executor/simple/executor.py
Normal file
222
src/pod_executor/simple/executor.py
Normal file
|
|
@ -0,0 +1,222 @@
|
|||
"""Code execution in isolated containers."""
|
||||
|
||||
from typing import Any, Optional
|
||||
from dataclasses import dataclass, asdict
|
||||
import json
|
||||
import time
|
||||
import textwrap
|
||||
import uuid
|
||||
|
||||
from pod_executor.containers.manager import SecureContainerManager, ContainerConfig
|
||||
from pod_executor.security.resource_limits import ResourceLimits
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExecutionResult:
|
||||
"""Result of code execution."""
|
||||
success: bool
|
||||
stdout: str
|
||||
stderr: str
|
||||
result: Optional[Any]
|
||||
execution_time: float
|
||||
exit_code: int
|
||||
error: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
class CodeExecutor:
|
||||
"""Executes Python code in isolated containers."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
container_manager: SecureContainerManager,
|
||||
image: str,
|
||||
resource_limits: ResourceLimits
|
||||
):
|
||||
self.container_manager = container_manager
|
||||
self.image = image
|
||||
self.resource_limits = resource_limits
|
||||
|
||||
def execute(
|
||||
self,
|
||||
code: str,
|
||||
timeout: Optional[int] = None,
|
||||
injection_code: Optional[str] = None,
|
||||
bridge_socket_path: Optional[str] = None
|
||||
) -> ExecutionResult:
|
||||
"""
|
||||
Execute Python code in a fresh container.
|
||||
|
||||
Process:
|
||||
1. Create container with code
|
||||
2. Start container
|
||||
3. Wait for completion (with timeout)
|
||||
4. Capture stdout/stderr
|
||||
5. Extract result from last expression
|
||||
6. Cleanup container
|
||||
|
||||
Args:
|
||||
code: Python code to execute
|
||||
timeout: Maximum execution time in seconds
|
||||
injection_code: Optional MCP tool injection code to prepend
|
||||
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||
|
||||
Returns:
|
||||
ExecutionResult with stdout, stderr, result, and timing
|
||||
"""
|
||||
start_time = time.time()
|
||||
container_id = None
|
||||
|
||||
# Use provided timeout or default from resource limits
|
||||
exec_timeout = timeout if timeout is not None else self.resource_limits.timeout
|
||||
|
||||
try:
|
||||
# Prepare code wrapper (with injection if provided)
|
||||
wrapped_code = self._prepare_code(code, injection_code=injection_code)
|
||||
|
||||
# Generate a session ID for this execution to register the container
|
||||
import uuid
|
||||
session_id = f"simple-exec-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
# Set up volumes for MCP bridge socket if provided
|
||||
volumes = {}
|
||||
if bridge_socket_path:
|
||||
volumes[bridge_socket_path] = {
|
||||
"bind": bridge_socket_path,
|
||||
"mode": "rw"
|
||||
}
|
||||
|
||||
# Create container configuration
|
||||
config = ContainerConfig(
|
||||
image=self.image,
|
||||
command=["python", "-c", wrapped_code],
|
||||
resource_limits=self.resource_limits,
|
||||
volumes=volumes if volumes else None
|
||||
)
|
||||
|
||||
# Create and start container with session_id for proper registration
|
||||
container_id = self.container_manager.create_container(config, session_id=session_id)
|
||||
self.container_manager.start_container(container_id)
|
||||
|
||||
# Wait for completion
|
||||
exit_code = self.container_manager.wait_for_container(
|
||||
container_id,
|
||||
timeout=exec_timeout
|
||||
)
|
||||
|
||||
# Get logs
|
||||
stdout, stderr = self.container_manager.get_container_logs(container_id)
|
||||
|
||||
# Parse output to extract result
|
||||
result, error = self._parse_output(stdout)
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return ExecutionResult(
|
||||
success=(exit_code == 0 and error is None),
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
result=result,
|
||||
execution_time=execution_time,
|
||||
exit_code=exit_code,
|
||||
error=error
|
||||
)
|
||||
|
||||
except TimeoutError as e:
|
||||
execution_time = time.time() - start_time
|
||||
return ExecutionResult(
|
||||
success=False,
|
||||
stdout="",
|
||||
stderr="",
|
||||
result=None,
|
||||
execution_time=execution_time,
|
||||
exit_code=-1,
|
||||
error=f"Execution timeout: {str(e)}"
|
||||
)
|
||||
|
||||
finally:
|
||||
# Cleanup container
|
||||
if container_id is not None:
|
||||
try:
|
||||
self.container_manager.remove_container(container_id)
|
||||
except Exception:
|
||||
pass # Best effort cleanup
|
||||
|
||||
def _prepare_code(self, code: str, injection_code: Optional[str] = None) -> str:
|
||||
"""
|
||||
Wrap code to capture result and handle errors.
|
||||
|
||||
Wraps code in try/except and captures:
|
||||
- Last expression result
|
||||
- Exceptions with traceback
|
||||
- Execution metadata
|
||||
|
||||
Args:
|
||||
code: User code to execute
|
||||
injection_code: Optional MCP tool injection code to prepend
|
||||
|
||||
Returns wrapped code that outputs JSON to stdout.
|
||||
"""
|
||||
# Prepend injection code if provided
|
||||
if injection_code:
|
||||
full_code = injection_code + "\n\n" + code
|
||||
else:
|
||||
full_code = code
|
||||
|
||||
# Escape the code for embedding in exec string
|
||||
escaped_code = full_code.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')
|
||||
|
||||
wrapper_template = '''
|
||||
import sys
|
||||
import json
|
||||
import traceback
|
||||
|
||||
def __mcp_execute():
|
||||
result = None
|
||||
error = None
|
||||
|
||||
try:
|
||||
# Execute user code
|
||||
exec_globals = {}
|
||||
exec("""%s""", exec_globals)
|
||||
|
||||
# Try to get result from last expression
|
||||
result = exec_globals.get('_', None)
|
||||
|
||||
except SyntaxError as e:
|
||||
error = f"SyntaxError: {e.msg} (line {e.lineno})"
|
||||
except Exception as e:
|
||||
error = f"{type(e).__name__}: {str(e)}"
|
||||
|
||||
# Output result as JSON
|
||||
print(json.dumps({"result": result, "error": error}))
|
||||
|
||||
__mcp_execute()
|
||||
'''
|
||||
|
||||
return wrapper_template % escaped_code
|
||||
|
||||
def _parse_output(self, stdout: str) -> tuple[Optional[Any], Optional[str]]:
|
||||
"""
|
||||
Parse execution output to extract result and error.
|
||||
|
||||
Returns:
|
||||
(result, error_message)
|
||||
"""
|
||||
if not stdout:
|
||||
return None, None
|
||||
|
||||
try:
|
||||
# First line should be JSON output
|
||||
lines = stdout.split('\n', 1)
|
||||
json_line = lines[0]
|
||||
|
||||
data = json.loads(json_line)
|
||||
return data.get("result"), data.get("error")
|
||||
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
# If can't parse JSON, treat entire output as result
|
||||
return None, None
|
||||
Loading…
Add table
Add a link
Reference in a new issue