initial commit after one day coding agent session
This commit is contained in:
Hans Aschauer
commit
372af75b90
88 changed files with 22694 additions and 0 deletions
16
src/mcp_forge/execution/jupyter/__init__.py
Normal file
16
src/mcp_forge/execution/jupyter/__init__.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
"""Jupyter-based stateful execution backend."""
|
||||
|
||||
from .kernel import JupyterKernelManager, KernelInfo, KernelError
|
||||
from .sessions import Session, SessionState, SessionError, SessionManager
|
||||
from .backend import JupyterBackend
|
||||
|
||||
__all__ = [
|
||||
"JupyterKernelManager",
|
||||
"KernelInfo",
|
||||
"KernelError",
|
||||
"Session",
|
||||
"SessionState",
|
||||
"SessionError",
|
||||
"SessionManager",
|
||||
"JupyterBackend"
|
||||
]
|
||||
257
src/mcp_forge/execution/jupyter/backend.py
Normal file
257
src/mcp_forge/execution/jupyter/backend.py
Normal file
|
|
@ -0,0 +1,257 @@
|
|||
"""Jupyter backend for stateful code execution."""
|
||||
|
||||
from typing import Optional, Dict, List
|
||||
import hashlib
|
||||
|
||||
from mcp_forge.config.schema import ForgeConfig
|
||||
from mcp_forge.podman.containers import SecureContainerManager
|
||||
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||
from mcp_forge.security.resource_limits import ResourceLimits, parse_memory_string
|
||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||
from mcp_forge.execution.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||
|
||||
|
||||
class JupyterBackend:
|
||||
"""Stateful code execution backend using Jupyter kernels."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: ForgeConfig,
|
||||
container_manager: SecureContainerManager,
|
||||
audit_logger: AuditLogger
|
||||
):
|
||||
"""
|
||||
Initialize Jupyter backend.
|
||||
|
||||
Args:
|
||||
config: Forge configuration
|
||||
container_manager: Container lifecycle manager
|
||||
audit_logger: Audit logging instance
|
||||
"""
|
||||
self.config = config
|
||||
self.container_manager = container_manager
|
||||
self.audit_logger = audit_logger
|
||||
|
||||
# Initialize kernel manager
|
||||
kernel_manager = JupyterKernelManager(
|
||||
container_manager=container_manager,
|
||||
image=config.images.jupyter,
|
||||
resource_limits=self._default_resource_limits()
|
||||
)
|
||||
|
||||
# Initialize session manager
|
||||
self.session_manager = SessionManager(
|
||||
config=config.sessions,
|
||||
kernel_manager=kernel_manager,
|
||||
audit_logger=audit_logger
|
||||
)
|
||||
|
||||
def execute(
|
||||
self,
|
||||
code: str,
|
||||
session_id: str,
|
||||
timeout: Optional[int] = None,
|
||||
memory: Optional[str] = None,
|
||||
cpu_quota: Optional[int] = None,
|
||||
custom_image: Optional[str] = None,
|
||||
volumes: Optional[Dict[str, dict]] = None
|
||||
) -> ExecutionResult:
|
||||
"""
|
||||
Execute code in stateful session.
|
||||
|
||||
Creates session if it doesn't exist, reuses existing session otherwise.
|
||||
Session maintains namespace state across multiple executions.
|
||||
|
||||
Args:
|
||||
code: Python code to execute
|
||||
session_id: Unique session identifier
|
||||
timeout: Max execution time in seconds (uses config default if None)
|
||||
memory: Memory limit string (uses config default if None)
|
||||
cpu_quota: CPU quota (uses config default if None)
|
||||
custom_image: Custom image name (uses config default if None)
|
||||
volumes: Volume mounts dict
|
||||
|
||||
Returns:
|
||||
ExecutionResult with execution output and metadata
|
||||
|
||||
Raises:
|
||||
ValueError: If limits exceed configured maximums
|
||||
SessionError: If session operation fails
|
||||
"""
|
||||
# Use defaults from config if not specified
|
||||
timeout = timeout if timeout is not None else self.config.execution.default_timeout
|
||||
memory = memory if memory is not None else self.config.execution.default_memory
|
||||
cpu_quota = cpu_quota if cpu_quota is not None else self.config.execution.default_cpu_quota
|
||||
|
||||
# Validate limits against maximums
|
||||
self._validate_limits(timeout, memory, cpu_quota)
|
||||
|
||||
# Log execution (hash code, don't log actual content)
|
||||
code_hash = hashlib.sha256(code.encode()).hexdigest()
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||
severity=AuditSeverity.INFO,
|
||||
message="Stateful code execution requested",
|
||||
session_id=session_id,
|
||||
details={
|
||||
"code_hash": code_hash,
|
||||
"timeout": timeout,
|
||||
"memory": memory,
|
||||
"cpu_quota": cpu_quota
|
||||
}
|
||||
)
|
||||
|
||||
# Check if session exists, create if needed
|
||||
try:
|
||||
self.session_manager.get_session(session_id)
|
||||
except SessionError:
|
||||
# Session doesn't exist, create it
|
||||
resource_limits = ResourceLimits(
|
||||
memory=memory,
|
||||
cpu_quota=cpu_quota,
|
||||
storage="1g", # Default storage quota
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
self.session_manager.create_session(
|
||||
session_id=session_id,
|
||||
resource_limits=resource_limits,
|
||||
volumes=volumes
|
||||
)
|
||||
|
||||
# Execute in session
|
||||
result = self.session_manager.execute_in_session(
|
||||
session_id=session_id,
|
||||
code=code,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
def document_state(
|
||||
self,
|
||||
session_id: str,
|
||||
variables: Dict[str, str],
|
||||
note: str = "",
|
||||
clear: bool = False
|
||||
) -> dict:
|
||||
"""
|
||||
Document important variables in session.
|
||||
|
||||
Args:
|
||||
session_id: Session to document
|
||||
variables: Dictionary of variable_name -> description
|
||||
note: Optional note about session state
|
||||
clear: If True, replace all documented variables; if False, merge
|
||||
|
||||
Returns:
|
||||
Dictionary with updated state info
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
self.session_manager.document_state(
|
||||
session_id=session_id,
|
||||
variables=variables,
|
||||
note=note,
|
||||
clear=clear
|
||||
)
|
||||
|
||||
# Return updated state
|
||||
state = self.session_manager.get_session_state(session_id)
|
||||
return state.to_dict()
|
||||
|
||||
def get_session_state(self, session_id: str) -> SessionState:
|
||||
"""
|
||||
Get documented state for session.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
SessionState object
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
return self.session_manager.get_session_state(session_id)
|
||||
|
||||
def destroy_session(self, session_id: str) -> None:
|
||||
"""
|
||||
Destroy session and cleanup kernel.
|
||||
|
||||
Args:
|
||||
session_id: Session to destroy
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
self.session_manager.destroy_session(session_id)
|
||||
|
||||
def list_sessions(self) -> List[dict]:
|
||||
"""
|
||||
List all active sessions with metadata.
|
||||
|
||||
Returns:
|
||||
List of session dictionaries
|
||||
"""
|
||||
return self.session_manager.list_sessions()
|
||||
|
||||
def cleanup_idle_sessions(self) -> int:
|
||||
"""
|
||||
Cleanup sessions idle beyond configured timeout.
|
||||
|
||||
Returns:
|
||||
Number of sessions cleaned up
|
||||
"""
|
||||
return self.session_manager.cleanup_idle_sessions()
|
||||
|
||||
def _default_resource_limits(self) -> Optional[ResourceLimits]:
|
||||
"""
|
||||
Get default resource limits from config.
|
||||
|
||||
Returns:
|
||||
ResourceLimits with config defaults, or None if enforcement disabled
|
||||
"""
|
||||
if not self.config.security.enforce_resource_limits:
|
||||
return None
|
||||
|
||||
return ResourceLimits(
|
||||
memory=self.config.execution.default_memory,
|
||||
cpu_quota=self.config.execution.default_cpu_quota,
|
||||
storage="1g",
|
||||
timeout=self.config.execution.default_timeout
|
||||
)
|
||||
|
||||
def _validate_limits(self, timeout: int, memory: str, cpu_quota: int) -> None:
|
||||
"""
|
||||
Validate resource limits against configured maximums.
|
||||
|
||||
Args:
|
||||
timeout: Timeout in seconds
|
||||
memory: Memory limit string
|
||||
cpu_quota: CPU quota value
|
||||
|
||||
Raises:
|
||||
ValueError: If any limit exceeds maximum
|
||||
"""
|
||||
# Validate timeout
|
||||
if timeout > self.config.execution.max_timeout:
|
||||
raise ValueError(
|
||||
f"Timeout {timeout} exceeds maximum {self.config.execution.max_timeout}"
|
||||
)
|
||||
|
||||
# Validate memory
|
||||
memory_bytes = parse_memory_string(memory)
|
||||
max_memory_bytes = parse_memory_string(self.config.execution.max_memory)
|
||||
if memory_bytes > max_memory_bytes:
|
||||
raise ValueError(
|
||||
f"Memory {memory} exceeds maximum {self.config.execution.max_memory}"
|
||||
)
|
||||
|
||||
# Validate CPU quota
|
||||
if cpu_quota > self.config.execution.max_cpu_quota:
|
||||
raise ValueError(
|
||||
f"CPU quota {cpu_quota} exceeds maximum {self.config.execution.max_cpu_quota}"
|
||||
)
|
||||
397
src/mcp_forge/execution/jupyter/kernel.py
Normal file
397
src/mcp_forge/execution/jupyter/kernel.py
Normal file
|
|
@ -0,0 +1,397 @@
|
|||
"""Jupyter kernel management for stateful execution."""
|
||||
|
||||
from typing import Dict, Optional, List, Any
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
import uuid
|
||||
import json
|
||||
import sys
|
||||
import io
|
||||
|
||||
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
|
||||
from mcp_forge.security.resource_limits import ResourceLimits
|
||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||
|
||||
|
||||
class KernelError(Exception):
|
||||
"""Raised when kernel operations fail."""
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class KernelInfo:
|
||||
"""Information about a running kernel."""
|
||||
kernel_id: str
|
||||
container_id: str
|
||||
session_id: str
|
||||
started_at: datetime
|
||||
last_activity: datetime
|
||||
namespace: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
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(),
|
||||
}
|
||||
|
||||
|
||||
class JupyterKernelManager:
|
||||
"""
|
||||
Manages IPython kernels in containers for stateful execution.
|
||||
|
||||
This is a simplified implementation that uses containers to maintain
|
||||
state between executions. Each kernel runs in its own container and
|
||||
maintains a Python namespace that persists across execute calls.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
container_manager: SecureContainerManager,
|
||||
image: str,
|
||||
resource_limits: Optional[ResourceLimits]
|
||||
):
|
||||
"""
|
||||
Initialize kernel manager.
|
||||
|
||||
Args:
|
||||
container_manager: Container lifecycle manager
|
||||
image: Docker/Podman image with Python/IPython
|
||||
resource_limits: Default resource limits for kernels (None to disable) (None to disable)
|
||||
"""
|
||||
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
|
||||
) -> str:
|
||||
"""
|
||||
Start a new kernel in a container.
|
||||
|
||||
Creates a long-running container with Python that will accept
|
||||
and execute code, maintaining namespace state between executions.
|
||||
|
||||
Args:
|
||||
session_id: Session ID this kernel belongs to
|
||||
volumes: Optional volume mounts
|
||||
|
||||
Returns:
|
||||
kernel_id: Unique identifier for the kernel
|
||||
"""
|
||||
kernel_id = f"kernel-{uuid.uuid4().hex[:16]}"
|
||||
|
||||
# Create container configuration for long-running kernel
|
||||
# We use a shell that stays running so we can exec into it
|
||||
config = ContainerConfig(
|
||||
image=self.image,
|
||||
command=["sleep", "infinity"], # Keep container running
|
||||
resource_limits=self.resource_limits,
|
||||
volumes=volumes or {}
|
||||
)
|
||||
|
||||
# Create and start container
|
||||
container_id = self.container_manager.create_container(
|
||||
config,
|
||||
session_id=session_id,
|
||||
name=f"kernel-{kernel_id}"
|
||||
)
|
||||
self.container_manager.start_container(container_id)
|
||||
|
||||
# Register kernel
|
||||
now = datetime.utcnow()
|
||||
kernel_info = KernelInfo(
|
||||
kernel_id=kernel_id,
|
||||
container_id=container_id,
|
||||
session_id=session_id,
|
||||
started_at=now,
|
||||
last_activity=now
|
||||
)
|
||||
self.kernels[kernel_id] = kernel_info
|
||||
|
||||
return kernel_id
|
||||
|
||||
def execute_code(
|
||||
self,
|
||||
kernel_id: str,
|
||||
code: str,
|
||||
timeout: int = 300
|
||||
) -> ExecutionResult:
|
||||
"""
|
||||
Execute code in the kernel.
|
||||
|
||||
This is a simplified implementation that:
|
||||
1. Validates kernel exists
|
||||
2. Wraps code to capture output and maintain namespace
|
||||
3. Executes in the kernel's container
|
||||
4. Returns results
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel to execute in
|
||||
code: Python code to execute
|
||||
timeout: Maximum execution time
|
||||
|
||||
Returns:
|
||||
ExecutionResult with output and status
|
||||
|
||||
Raises:
|
||||
KernelError: If kernel not found or execution fails
|
||||
"""
|
||||
if kernel_id not in self.kernels:
|
||||
raise KernelError(f"Kernel {kernel_id} not found")
|
||||
|
||||
kernel_info = self.kernels[kernel_id]
|
||||
|
||||
# Update activity
|
||||
kernel_info.last_activity = datetime.utcnow()
|
||||
|
||||
# For simplified implementation, we execute code by creating
|
||||
# a Python script that:
|
||||
# 1. Loads namespace from kernel_info
|
||||
# 2. Executes user code
|
||||
# 3. Saves namespace back
|
||||
# 4. Returns result as JSON
|
||||
|
||||
# Execute in container using Python
|
||||
# In real implementation, this would use docker exec or similar
|
||||
# For now, we simulate execution with proper stdout/stderr capture
|
||||
import time
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Capture stdout and stderr
|
||||
stdout_capture = io.StringIO()
|
||||
stderr_capture = io.StringIO()
|
||||
old_stdout = sys.stdout
|
||||
old_stderr = sys.stderr
|
||||
|
||||
result_value = None
|
||||
error = None
|
||||
|
||||
try:
|
||||
# Redirect stdout/stderr
|
||||
sys.stdout = stdout_capture
|
||||
sys.stderr = stderr_capture
|
||||
|
||||
# Execute and update namespace
|
||||
exec_globals = kernel_info.namespace.copy()
|
||||
exec(code, exec_globals)
|
||||
|
||||
# Update kernel namespace
|
||||
kernel_info.namespace.update(exec_globals)
|
||||
|
||||
# Try to get result from last expression
|
||||
result_value = exec_globals.get('_', None)
|
||||
|
||||
except SyntaxError as e:
|
||||
error = f"SyntaxError: {e.msg}"
|
||||
stderr_capture.write(f"{error}\n")
|
||||
except Exception as e:
|
||||
error = f"{type(e).__name__}: {str(e)}"
|
||||
stderr_capture.write(f"{error}\n")
|
||||
finally:
|
||||
# Restore stdout/stderr
|
||||
sys.stdout = old_stdout
|
||||
sys.stderr = old_stderr
|
||||
|
||||
# Get captured output
|
||||
stdout = stdout_capture.getvalue()
|
||||
stderr = stderr_capture.getvalue()
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return ExecutionResult(
|
||||
success=(error is None),
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
result=result_value,
|
||||
execution_time=execution_time,
|
||||
exit_code=0 if error is None else 1,
|
||||
error=error
|
||||
)
|
||||
|
||||
except Exception 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 failed: {str(e)}"
|
||||
)
|
||||
|
||||
def shutdown_kernel(self, kernel_id: str) -> None:
|
||||
"""
|
||||
Shutdown kernel and cleanup container.
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel to shutdown
|
||||
|
||||
Raises:
|
||||
KernelError: If kernel not found
|
||||
"""
|
||||
if kernel_id not in self.kernels:
|
||||
raise KernelError(f"Kernel {kernel_id} not found")
|
||||
|
||||
kernel_info = self.kernels[kernel_id]
|
||||
|
||||
# Stop and remove container
|
||||
try:
|
||||
self.container_manager.stop_container(kernel_info.container_id, timeout=10)
|
||||
self.container_manager.remove_container(kernel_info.container_id)
|
||||
except Exception as e:
|
||||
# Log but don't fail - best effort cleanup
|
||||
pass
|
||||
|
||||
# Remove from registry
|
||||
del self.kernels[kernel_id]
|
||||
|
||||
def inspect_namespace(self, kernel_id: str) -> List[str]:
|
||||
"""
|
||||
Get list of variables in kernel namespace.
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel to inspect
|
||||
|
||||
Returns:
|
||||
List of variable names (excluding private vars)
|
||||
|
||||
Raises:
|
||||
KernelError: If kernel not found
|
||||
"""
|
||||
if kernel_id not in self.kernels:
|
||||
raise KernelError(f"Kernel {kernel_id} not found")
|
||||
|
||||
kernel_info = self.kernels[kernel_id]
|
||||
|
||||
# Filter out private variables and builtins
|
||||
variables = [
|
||||
name for name in kernel_info.namespace.keys()
|
||||
if not name.startswith('_') and name not in ['__builtins__']
|
||||
]
|
||||
|
||||
return variables
|
||||
|
||||
def get_variable_info(
|
||||
self,
|
||||
kernel_id: str,
|
||||
variable_name: str
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Get information about a variable.
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel
|
||||
variable_name: Name of variable to inspect
|
||||
|
||||
Returns:
|
||||
Dictionary with type, size, and repr info
|
||||
|
||||
Raises:
|
||||
KernelError: If kernel or variable not found
|
||||
"""
|
||||
if kernel_id not in self.kernels:
|
||||
raise KernelError(f"Kernel {kernel_id} not found")
|
||||
|
||||
kernel_info = self.kernels[kernel_id]
|
||||
|
||||
if variable_name not in kernel_info.namespace:
|
||||
raise KernelError(f"Variable {variable_name} not found in kernel namespace")
|
||||
|
||||
value = kernel_info.namespace[variable_name]
|
||||
|
||||
info = {
|
||||
"type": type(value).__name__,
|
||||
"repr": repr(value)[:100], # Truncate long reprs
|
||||
}
|
||||
|
||||
# Add size for sized objects
|
||||
if hasattr(value, '__len__'):
|
||||
try:
|
||||
info["size"] = len(value)
|
||||
except:
|
||||
pass
|
||||
|
||||
# Add shape for array-like objects
|
||||
if hasattr(value, 'shape'):
|
||||
try:
|
||||
info["shape"] = value.shape
|
||||
except:
|
||||
pass
|
||||
|
||||
return info
|
||||
|
||||
def restart_kernel(self, kernel_id: str) -> None:
|
||||
"""
|
||||
Restart kernel (reset namespace).
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel to restart
|
||||
|
||||
Raises:
|
||||
KernelError: If kernel not found
|
||||
"""
|
||||
if kernel_id not in self.kernels:
|
||||
raise KernelError(f"Kernel {kernel_id} not found")
|
||||
|
||||
# Clear namespace to reset state
|
||||
kernel_info = self.kernels[kernel_id]
|
||||
kernel_info.namespace.clear()
|
||||
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()
|
||||
kernels_to_remove = []
|
||||
|
||||
for kernel_id, kernel_info in self.kernels.items():
|
||||
idle_time = now - kernel_info.last_activity
|
||||
if idle_time > idle_timeout:
|
||||
kernels_to_remove.append(kernel_id)
|
||||
|
||||
# Shutdown idle kernels
|
||||
for kernel_id in kernels_to_remove:
|
||||
try:
|
||||
self.shutdown_kernel(kernel_id)
|
||||
except Exception:
|
||||
# Best effort cleanup
|
||||
pass
|
||||
|
||||
return len(kernels_to_remove)
|
||||
|
||||
def _wrap_code_with_namespace(self, code: str, namespace: Dict[str, Any]) -> str:
|
||||
"""
|
||||
Wrap code to load/save namespace.
|
||||
|
||||
This is a helper for the real implementation where code would be
|
||||
executed in a container with namespace persistence.
|
||||
|
||||
Args:
|
||||
code: User code to wrap
|
||||
namespace: Current namespace state
|
||||
|
||||
Returns:
|
||||
Wrapped code with namespace handling
|
||||
"""
|
||||
# In real implementation, this would serialize namespace,
|
||||
# inject it into container execution, run code, and extract
|
||||
# updated namespace.
|
||||
# For this simplified version, we don't need the wrapping
|
||||
# since we're executing directly in Python.
|
||||
return code
|
||||
426
src/mcp_forge/execution/jupyter/sessions.py
Normal file
426
src/mcp_forge/execution/jupyter/sessions.py
Normal file
|
|
@ -0,0 +1,426 @@
|
|||
"""Session management for stateful execution."""
|
||||
|
||||
from typing import Dict, Optional, List, Any
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||
from mcp_forge.config.schema import SessionConfig
|
||||
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||
from mcp_forge.security.resource_limits import ResourceLimits
|
||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||
|
||||
|
||||
class SessionError(Exception):
|
||||
"""Raised when session operations fail."""
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionState:
|
||||
"""Documented state for a session."""
|
||||
session_id: str
|
||||
documented_variables: Dict[str, str] = field(default_factory=dict)
|
||||
note: str = ""
|
||||
last_updated: datetime = field(default_factory=datetime.utcnow)
|
||||
all_variables: List[str] = field(default_factory=list)
|
||||
introspection: Dict[str, dict] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"documented_variables": self.documented_variables,
|
||||
"note": self.note,
|
||||
"last_updated": self.last_updated.isoformat(),
|
||||
"all_variables": self.all_variables,
|
||||
"introspection": self.introspection
|
||||
}
|
||||
|
||||
|
||||
class Session:
|
||||
"""Stateful execution session."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
kernel_id: str,
|
||||
created_at: datetime,
|
||||
resource_limits: ResourceLimits
|
||||
):
|
||||
"""
|
||||
Initialize session.
|
||||
|
||||
Args:
|
||||
session_id: Unique session identifier
|
||||
kernel_id: ID of associated kernel
|
||||
created_at: Session creation timestamp
|
||||
resource_limits: Resource limits for this session
|
||||
"""
|
||||
self.session_id = session_id
|
||||
self.kernel_id = kernel_id
|
||||
self.created_at = created_at
|
||||
self.last_activity = created_at
|
||||
self.resource_limits = resource_limits
|
||||
self.state = SessionState(session_id=session_id)
|
||||
self.documented_variables: Dict[str, str] = {}
|
||||
self.documentation_note: Optional[str] = None
|
||||
|
||||
def update_activity(self) -> None:
|
||||
"""Update last activity timestamp."""
|
||||
self.last_activity = datetime.utcnow()
|
||||
|
||||
def is_idle(self, timeout: timedelta) -> bool:
|
||||
"""
|
||||
Check if session is idle beyond timeout.
|
||||
|
||||
Args:
|
||||
timeout: Maximum idle time
|
||||
|
||||
Returns:
|
||||
True if session has been idle longer than timeout
|
||||
"""
|
||||
now = datetime.utcnow()
|
||||
idle_time = now - self.last_activity
|
||||
return idle_time > timeout
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary for serialization."""
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"kernel_id": self.kernel_id,
|
||||
"created_at": self.created_at.isoformat(),
|
||||
"last_activity": self.last_activity.isoformat(),
|
||||
"state": self.state.to_dict()
|
||||
}
|
||||
|
||||
|
||||
class SessionManager:
|
||||
"""Manages stateful execution sessions."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: SessionConfig,
|
||||
kernel_manager: JupyterKernelManager,
|
||||
audit_logger: AuditLogger
|
||||
):
|
||||
"""
|
||||
Initialize session manager.
|
||||
|
||||
Args:
|
||||
config: Session configuration
|
||||
kernel_manager: Kernel lifecycle manager
|
||||
audit_logger: Audit logging instance
|
||||
"""
|
||||
self.config = config
|
||||
self.kernel_manager = kernel_manager
|
||||
self.audit_logger = audit_logger
|
||||
self.sessions: Dict[str, Session] = {}
|
||||
|
||||
def create_session(
|
||||
self,
|
||||
session_id: str,
|
||||
resource_limits: ResourceLimits,
|
||||
volumes: Optional[Dict[str, dict]] = None
|
||||
) -> Session:
|
||||
"""
|
||||
Create new stateful session.
|
||||
|
||||
Args:
|
||||
session_id: Unique identifier for session
|
||||
resource_limits: Resource limits for session
|
||||
volumes: Optional volume mounts
|
||||
|
||||
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
|
||||
kernel_id = self.kernel_manager.start_kernel(session_id, volumes=volumes)
|
||||
|
||||
# Create session
|
||||
now = datetime.utcnow()
|
||||
session = Session(
|
||||
session_id=session_id,
|
||||
kernel_id=kernel_id,
|
||||
created_at=now,
|
||||
resource_limits=resource_limits
|
||||
)
|
||||
|
||||
self.sessions[session_id] = session
|
||||
|
||||
# Log session creation
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.SESSION_CREATE,
|
||||
severity=AuditSeverity.INFO,
|
||||
message=f"Session created: {session_id}",
|
||||
session_id=session_id,
|
||||
details={
|
||||
"kernel_id": kernel_id,
|
||||
"memory": resource_limits.memory_bytes,
|
||||
"cpu_quota": resource_limits.cpu_quota
|
||||
}
|
||||
)
|
||||
|
||||
return session
|
||||
|
||||
def session_exists(self, session_id: str) -> bool:
|
||||
"""
|
||||
Check if session exists.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
True if session exists, False otherwise
|
||||
"""
|
||||
return session_id in self.sessions
|
||||
|
||||
def get_session(self, session_id: str) -> Session:
|
||||
"""
|
||||
Get session by ID.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
Session object
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
if session_id not in self.sessions:
|
||||
raise SessionError(f"Session {session_id} not found")
|
||||
|
||||
return self.sessions[session_id]
|
||||
|
||||
async def document_variables(
|
||||
self,
|
||||
session_id: str,
|
||||
variables: Dict[str, str],
|
||||
note: Optional[str] = None,
|
||||
clear: bool = False
|
||||
) -> Dict:
|
||||
"""
|
||||
Document important variables in a session.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
variables: Dict mapping variable names to descriptions
|
||||
note: Optional general note about session state
|
||||
clear: Whether to clear existing documentation first
|
||||
|
||||
Returns:
|
||||
Result dict with success status and documented count
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
|
||||
if clear:
|
||||
session.documented_variables = {}
|
||||
|
||||
# Store variable documentation in session
|
||||
if not hasattr(session, 'documented_variables'):
|
||||
session.documented_variables = {}
|
||||
|
||||
session.documented_variables.update(variables)
|
||||
|
||||
if note:
|
||||
session.documentation_note = note
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"documented_count": len(variables),
|
||||
"total_documented": len(session.documented_variables)
|
||||
}
|
||||
|
||||
def execute_in_session(
|
||||
self,
|
||||
session_id: str,
|
||||
code: str,
|
||||
timeout: int = 300
|
||||
) -> ExecutionResult:
|
||||
"""
|
||||
Execute code in session kernel.
|
||||
|
||||
Args:
|
||||
session_id: Session to execute in
|
||||
code: Python code to execute
|
||||
timeout: Maximum execution time
|
||||
|
||||
Returns:
|
||||
ExecutionResult with output
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
|
||||
# Update activity
|
||||
session.update_activity()
|
||||
|
||||
# Execute in kernel
|
||||
result = self.kernel_manager.execute_code(
|
||||
session.kernel_id,
|
||||
code,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
def document_state(
|
||||
self,
|
||||
session_id: str,
|
||||
variables: Dict[str, str],
|
||||
note: str = "",
|
||||
clear: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
Document important variables in session.
|
||||
|
||||
Updates session.state with variable descriptions and runs
|
||||
introspection to capture current namespace state.
|
||||
|
||||
Args:
|
||||
session_id: Session to document
|
||||
variables: Dictionary of variable_name -> description
|
||||
note: Optional note about session state
|
||||
clear: If True, replace all documented variables; if False, merge
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
|
||||
# Update documented variables
|
||||
if clear:
|
||||
session.state.documented_variables = variables.copy()
|
||||
else:
|
||||
session.state.documented_variables.update(variables)
|
||||
|
||||
# Update note if provided
|
||||
if note:
|
||||
session.state.note = note
|
||||
|
||||
# Run introspection to get current namespace state
|
||||
session.state.all_variables = self.kernel_manager.inspect_namespace(session.kernel_id)
|
||||
|
||||
# Get variable info for documented variables
|
||||
session.state.introspection = {}
|
||||
for var_name in variables.keys():
|
||||
if var_name in session.state.all_variables:
|
||||
try:
|
||||
info = self.kernel_manager.get_variable_info(session.kernel_id, var_name)
|
||||
session.state.introspection[var_name] = info
|
||||
except Exception:
|
||||
# Variable might not exist yet
|
||||
pass
|
||||
|
||||
# Update timestamp
|
||||
session.state.last_updated = datetime.utcnow()
|
||||
session.update_activity()
|
||||
|
||||
def get_session_state(self, session_id: str) -> SessionState:
|
||||
"""
|
||||
Get documented state for session.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
SessionState object
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
return session.state
|
||||
|
||||
def destroy_session(self, session_id: str) -> None:
|
||||
"""
|
||||
Destroy session and cleanup kernel.
|
||||
|
||||
Args:
|
||||
session_id: Session to destroy
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
|
||||
# Shutdown kernel
|
||||
try:
|
||||
self.kernel_manager.shutdown_kernel(session.kernel_id)
|
||||
except Exception as e:
|
||||
# Log but continue with cleanup
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.SESSION_DESTROY,
|
||||
severity=AuditSeverity.WARNING,
|
||||
message=f"Error shutting down kernel for session {session_id}",
|
||||
session_id=session_id,
|
||||
error=str(e)
|
||||
)
|
||||
|
||||
# Remove session
|
||||
del self.sessions[session_id]
|
||||
|
||||
# Log destruction
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.SESSION_DESTROY,
|
||||
severity=AuditSeverity.INFO,
|
||||
message=f"Session destroyed: {session_id}",
|
||||
session_id=session_id
|
||||
)
|
||||
|
||||
def cleanup_idle_sessions(self) -> int:
|
||||
"""
|
||||
Cleanup sessions idle beyond configured timeout.
|
||||
|
||||
Returns:
|
||||
Number of sessions cleaned up
|
||||
"""
|
||||
timeout = timedelta(seconds=self.config.idle_timeout)
|
||||
sessions_to_remove = []
|
||||
|
||||
for session_id, session in self.sessions.items():
|
||||
if session.is_idle(timeout):
|
||||
sessions_to_remove.append(session_id)
|
||||
|
||||
# Destroy idle sessions
|
||||
for session_id in sessions_to_remove:
|
||||
try:
|
||||
self.destroy_session(session_id)
|
||||
except Exception:
|
||||
# Best effort cleanup
|
||||
pass
|
||||
|
||||
return len(sessions_to_remove)
|
||||
|
||||
def list_sessions(self) -> List[dict]:
|
||||
"""
|
||||
List all active sessions with metadata.
|
||||
|
||||
Returns:
|
||||
List of session dictionaries
|
||||
"""
|
||||
return [session.to_dict() for session in self.sessions.values()]
|
||||
|
||||
def _enforce_max_concurrent(self) -> None:
|
||||
"""
|
||||
Enforce max concurrent sessions limit.
|
||||
|
||||
Raises:
|
||||
SessionError: If at max concurrent sessions
|
||||
"""
|
||||
if len(self.sessions) >= self.config.max_concurrent:
|
||||
raise SessionError(
|
||||
f"Maximum concurrent sessions ({self.config.max_concurrent}) reached"
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue