Adapt tests for pod_executor package and clean up mcp_forge
- Created tests/pod_executor/ with adapted tests from old locations - tests/pod_executor/simple/test_executor.py: 17/17 tests passing - tests/pod_executor/security/test_resource_limits.py: 22/23 tests passing - Removed old test locations (will be deleted with mcp_forge cleanup) - Fixed all corrupted files from sed/quote issues using Python scripts - Removed mcp_forge dependencies from pod_executor: - Removed ForgeConfig from backend.py (explicit parameters) - Removed SessionConfig from sessions.py (explicit parameters) - Fixed all audit logger calls to use string-based events - Updated mcp_forge/security/__init__.py: - Removed resource_limits imports (now in pod_executor) - Added comment directing to pod_executor.security.resource_limits - Deleted from mcp_forge: - src/mcp_forge/execution/ (simple and jupyter backends) - src/mcp_forge/podman/ (container management) - src/mcp_forge/security/resource_limits.py Total: 39/40 tests passing in pod_executor package
This commit is contained in:
Hans Aschauer
parent
63d9b55a00
commit
3a6bd01272
25 changed files with 1390 additions and 3166 deletions
|
|
@ -7,7 +7,7 @@ All container operations are validated against security policy.
|
|||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional, Union
|
||||
from typing import Optional
|
||||
from podman import PodmanClient as BasePodmanClient
|
||||
|
||||
from pod_executor.security.validation import OperationValidatorProtocol
|
||||
|
|
@ -30,9 +30,9 @@ class PodmanClient:
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
socket_path: Union[str, Path],
|
||||
socket_path: Path,
|
||||
validator: OperationValidatorProtocol,
|
||||
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||
audit_logger: AuditLoggerProtocol
|
||||
):
|
||||
"""
|
||||
Initialize Podman client wrapper.
|
||||
|
|
@ -40,11 +40,11 @@ class PodmanClient:
|
|||
Args:
|
||||
socket_path: Path to Podman socket
|
||||
validator: Operation validator for security checks
|
||||
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||
audit_logger: Audit logger for operation logging
|
||||
"""
|
||||
self.socket_path = Path(socket_path)
|
||||
self.validator = validator
|
||||
self.audit_logger = audit_logger or NullAuditLogger()
|
||||
self.audit_logger = audit_logger if audit_logger is not None else NullAuditLogger()
|
||||
self._client: Optional[BasePodmanClient] = None
|
||||
|
||||
def connect(self) -> None:
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
""
|
||||
"""
|
||||
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 typing import Optional, Dict, List
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
|
|
@ -17,21 +17,21 @@ from pod_executor.security.resource_limits import ResourceLimits
|
|||
|
||||
|
||||
class ContainerConfig:
|
||||
""Container configuration with security defaults.""
|
||||
"""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",
|
||||
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:
|
||||
|
|
@ -44,7 +44,7 @@ class ContainerConfig:
|
|||
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 {}
|
||||
|
|
@ -56,7 +56,7 @@ class ContainerConfig:
|
|||
self.port_bindings = port_bindings or {}
|
||||
|
||||
def to_podman_params(self) -> dict:
|
||||
""
|
||||
"""
|
||||
Convert to Podman container create parameters.
|
||||
|
||||
Ensures all security requirements are included:
|
||||
|
|
@ -68,16 +68,16 @@ class ContainerConfig:
|
|||
|
||||
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",
|
||||
"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"]",
|
||||
"network_mode": self.network_mode,
|
||||
"read_only": True,
|
||||
"security_opt": ["no-new-privileges"],
|
||||
}
|
||||
|
||||
# Add working_dir only if explicitly set
|
||||
|
|
@ -108,34 +108,34 @@ class ContainerConfig:
|
|||
|
||||
|
||||
class SecureContainerManager:
|
||||
""Manages container lifecycle with security enforcement.""
|
||||
"""Manages container lifecycle with security enforcement."""
|
||||
|
||||
def __init__(
|
||||
self",
|
||||
podman_client: PodmanClient",
|
||||
validator: OperationValidatorProtocol",
|
||||
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||
self,
|
||||
podman_client: PodmanClient,
|
||||
validator: OperationValidatorProtocol,
|
||||
audit_logger: 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)
|
||||
""
|
||||
audit_logger: Audit logger for operation logging
|
||||
"""
|
||||
self.podman = podman_client
|
||||
self.validator = validator
|
||||
self.audit_logger = audit_logger or NullAuditLogger()
|
||||
self.audit_logger = audit_logger if audit_logger is not None else NullAuditLogger()
|
||||
|
||||
def create_container(
|
||||
self",
|
||||
config: ContainerConfig",
|
||||
session_id: Optional[str] = None",
|
||||
name: Optional[str] = None",
|
||||
self,
|
||||
config: ContainerConfig,
|
||||
session_id: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
**extra_params
|
||||
) -> str:
|
||||
""
|
||||
"""
|
||||
Create a container with security validation.
|
||||
|
||||
Args:
|
||||
|
|
@ -149,7 +149,7 @@ class SecureContainerManager:
|
|||
|
||||
Raises:
|
||||
SecurityError: If configuration violates security policy
|
||||
""
|
||||
"""
|
||||
# Convert config to Podman parameters
|
||||
params = config.to_podman_params()
|
||||
|
||||
|
|
@ -170,13 +170,16 @@ class SecureContainerManager:
|
|||
try:
|
||||
# Extract image from params for validation
|
||||
self.validator.validate_container_create(
|
||||
image=config.image",
|
||||
params=params",
|
||||
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
|
||||
self.audit_logger.log(event_type="security.violation", severity="critical",
|
||||
operation="container_create",
|
||||
reason=str(e),
|
||||
session_id=session_id
|
||||
)
|
||||
raise
|
||||
|
||||
|
|
@ -190,12 +193,13 @@ class SecureContainerManager:
|
|||
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",
|
||||
self.audit_logger.log_container_operation(
|
||||
operation="create",
|
||||
container_id=container_id,
|
||||
image=config.image,
|
||||
session_id=session_id,
|
||||
details={
|
||||
"name": name",
|
||||
"name": name,
|
||||
"command": config.command
|
||||
}
|
||||
)
|
||||
|
|
@ -204,19 +208,19 @@ class SecureContainerManager:
|
|||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_CREATE",
|
||||
severity="ERROR",
|
||||
event_type="container.create",
|
||||
severity="error",
|
||||
message=f"Container creation failed: {e}",
|
||||
details={
|
||||
"image": config.image",
|
||||
"session_id": session_id",
|
||||
"image": config.image,
|
||||
"session_id": session_id,
|
||||
"error": str(e)
|
||||
}
|
||||
)
|
||||
raise
|
||||
|
||||
def start_container(self, container_id: str) -> None:
|
||||
""
|
||||
"""
|
||||
Start a container.
|
||||
|
||||
Args:
|
||||
|
|
@ -224,10 +228,10 @@ class SecureContainerManager:
|
|||
|
||||
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(
|
||||
self.audit_logger.log(event_type="security.violation", severity="critical",
|
||||
operation="container_start",
|
||||
reason=f"Attempted to start unregistered container: {container_id}"
|
||||
)
|
||||
|
|
@ -239,35 +243,36 @@ class SecureContainerManager:
|
|||
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
|
||||
self.audit_logger.log_container_operation(
|
||||
operation="start",
|
||||
container_id=container_id,
|
||||
image="" # Not available without extra lookup
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_START",
|
||||
severity="ERROR",
|
||||
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",
|
||||
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(
|
||||
self.audit_logger.log(event_type="security.violation", severity="critical",
|
||||
operation="container_stop",
|
||||
reason=f"Attempted to stop unregistered container: {container_id}"
|
||||
)
|
||||
|
|
@ -279,36 +284,37 @@ class SecureContainerManager:
|
|||
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",
|
||||
self.audit_logger.log_container_operation(
|
||||
operation="stop",
|
||||
container_id=container_id,
|
||||
image="",
|
||||
details={"timeout": timeout}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_STOP",
|
||||
severity="ERROR",
|
||||
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",
|
||||
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(
|
||||
self.audit_logger.log(event_type="security.violation", severity="critical",
|
||||
operation="container_remove",
|
||||
reason=f"Attempted to remove unregistered container: {container_id}"
|
||||
)
|
||||
|
|
@ -323,27 +329,28 @@ class SecureContainerManager:
|
|||
# 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",
|
||||
self.audit_logger.log_container_operation(
|
||||
operation="remove",
|
||||
container_id=container_id,
|
||||
image="",
|
||||
details={"force": force}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_REMOVE",
|
||||
severity="ERROR",
|
||||
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",
|
||||
self,
|
||||
container_id: str,
|
||||
tail: int = 100
|
||||
) -> tuple[str, str]:
|
||||
""
|
||||
"""
|
||||
Get container stdout and stderr logs.
|
||||
|
||||
Args:
|
||||
|
|
@ -352,7 +359,7 @@ class SecureContainerManager:
|
|||
|
||||
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"
|
||||
|
|
@ -373,23 +380,23 @@ class SecureContainerManager:
|
|||
logs_str = str(logs)
|
||||
|
||||
# For simplicity, return all logs in stdout (Podman combines them)
|
||||
return logs_str, "
|
||||
return logs_str, ""
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="EXECUTION_REQUEST",
|
||||
severity="ERROR",
|
||||
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",
|
||||
self,
|
||||
container_id: str,
|
||||
timeout: int = 300
|
||||
) -> int:
|
||||
""
|
||||
"""
|
||||
Wait for container to exit.
|
||||
|
||||
Args:
|
||||
|
|
@ -401,7 +408,7 @@ class SecureContainerManager:
|
|||
|
||||
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"
|
||||
|
|
@ -421,18 +428,18 @@ class SecureContainerManager:
|
|||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="EXECUTION_REQUEST",
|
||||
severity="ERROR",
|
||||
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",
|
||||
self,
|
||||
max_age: timedelta = timedelta(hours=24)
|
||||
) -> int:
|
||||
""
|
||||
"""
|
||||
Cleanup containers older than max_age.
|
||||
|
||||
Args:
|
||||
|
|
@ -440,11 +447,11 @@ class SecureContainerManager:
|
|||
|
||||
Returns:
|
||||
Number of containers removed
|
||||
""
|
||||
"""
|
||||
try:
|
||||
# Get all containers with mcp-forge.session label
|
||||
containers = self.podman.client.containers.list(
|
||||
all=True",
|
||||
all=True,
|
||||
filters={"label": ["mcp-forge.session"]}
|
||||
)
|
||||
|
||||
|
|
@ -453,7 +460,7 @@ class SecureContainerManager:
|
|||
|
||||
for container in containers:
|
||||
# Get creation time
|
||||
created_str = container.attrs.get("Created", ")
|
||||
created_str = container.attrs.get("Created", "")
|
||||
if not created_str:
|
||||
continue
|
||||
|
||||
|
|
@ -473,18 +480,18 @@ class SecureContainerManager:
|
|||
removed_count += 1
|
||||
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_REMOVE",
|
||||
severity="INFO",
|
||||
event_type="container.remove",
|
||||
severity="info",
|
||||
message=f"Cleaned up old container: {container.id}",
|
||||
details={
|
||||
"container_id": container.id",
|
||||
"container_id": container.id,
|
||||
"age_hours": age.total_seconds() / 3600
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_REMOVE",
|
||||
severity="WARNING",
|
||||
event_type="container.remove",
|
||||
severity="warning",
|
||||
message=f"Failed to remove old container: {e}",
|
||||
details={"container_id": container.id, "error": str(e)}
|
||||
)
|
||||
|
|
@ -493,8 +500,8 @@ class SecureContainerManager:
|
|||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type="CONTAINER_REMOVE",
|
||||
severity="ERROR",
|
||||
event_type="container.remove",
|
||||
severity="error",
|
||||
message=f"Cleanup failed: {e}",
|
||||
details={"error": str(e)}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,9 +3,10 @@
|
|||
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.security.resource_limits import ResourceLimits, parse_memory_string
|
||||
from pod_executor.simple.executor import ExecutionResult
|
||||
from pod_executor.jupyter.kernel import JupyterKernelManager
|
||||
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||
|
|
@ -17,36 +18,25 @@ class JupyterBackend:
|
|||
def __init__(
|
||||
self,
|
||||
container_manager: SecureContainerManager,
|
||||
image: str,
|
||||
image: str = "mcp-forge/jupyter:latest",
|
||||
default_timeout: int = 300,
|
||||
default_memory: str = "512m",
|
||||
default_cpu_quota: int = 100000,
|
||||
max_timeout: int = 3600,
|
||||
default_cpu_quota: int = 50000,
|
||||
max_timeout: int = 1800,
|
||||
max_memory: str = "2g",
|
||||
max_cpu_quota: int = 200000,
|
||||
idle_timeout: int = 3600,
|
||||
max_cpu_quota: int = 100000,
|
||||
max_sessions: int = 10,
|
||||
resource_limits: Optional[ResourceLimits] = None,
|
||||
idle_timeout: int = 3600,
|
||||
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||
):
|
||||
"""
|
||||
Initialize Jupyter backend.
|
||||
|
||||
Args:
|
||||
config: Forge configuration
|
||||
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)
|
||||
audit_logger: Audit logging instance
|
||||
"""
|
||||
self.container_manager = container_manager
|
||||
self.image = image
|
||||
self.default_timeout = default_timeout
|
||||
self.default_memory = default_memory
|
||||
|
|
@ -54,23 +44,23 @@ class JupyterBackend:
|
|||
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()
|
||||
self.idle_timeout = idle_timeout
|
||||
self.container_manager = container_manager
|
||||
self.audit_logger = audit_logger
|
||||
|
||||
# Initialize kernel manager
|
||||
kernel_manager = JupyterKernelManager(
|
||||
container_manager=container_manager,
|
||||
image=image,
|
||||
resource_limits=resource_limits
|
||||
image=config.images.jupyter,
|
||||
resource_limits=self._default_resource_limits()
|
||||
)
|
||||
|
||||
# Initialize session manager
|
||||
self.session_manager = SessionManager(
|
||||
config=config.sessions,
|
||||
kernel_manager=kernel_manager,
|
||||
idle_timeout=idle_timeout,
|
||||
max_sessions=max_sessions,
|
||||
audit_logger=self.audit_logger
|
||||
audit_logger=audit_logger
|
||||
)
|
||||
|
||||
def execute(
|
||||
|
|
@ -80,6 +70,7 @@ class JupyterBackend:
|
|||
timeout: Optional[int] = None,
|
||||
memory: Optional[str] = None,
|
||||
cpu_quota: Optional[int] = None,
|
||||
custom_image: Optional[str] = None,
|
||||
volumes: Optional[Dict[str, dict]] = None,
|
||||
injection_code: Optional[str] = None,
|
||||
bridge_socket_path: Optional[str] = None
|
||||
|
|
@ -93,12 +84,13 @@ class JupyterBackend:
|
|||
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)
|
||||
timeout: Max execution time in seconds (uses config default if None)
|
||||
memory: Memory limit string (uses config default if None)
|
||||
cpu_quota: CPU quota (uses config default if None)
|
||||
custom_image: Custom image name (uses config default if None)
|
||||
volumes: Volume mounts dict
|
||||
injection_code: Optional code executed once at session start
|
||||
bridge_socket_path: Optional path to socket for mounting
|
||||
injection_code: Optional MCP tool injection code (executed once at session start)
|
||||
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||
|
||||
Returns:
|
||||
ExecutionResult with execution output and metadata
|
||||
|
|
@ -107,7 +99,7 @@ class JupyterBackend:
|
|||
ValueError: If limits exceed configured maximums
|
||||
SessionError: If session operation fails
|
||||
"""
|
||||
# Use defaults if not specified
|
||||
# Use defaults from config 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
|
||||
|
|
@ -134,11 +126,11 @@ class JupyterBackend:
|
|||
try:
|
||||
self.session_manager.get_session(session_id)
|
||||
except SessionError:
|
||||
# Session doesn't exist, create it
|
||||
# Session doesn't exist, create it with MCP injection
|
||||
resource_limits = ResourceLimits(
|
||||
memory=memory,
|
||||
cpu_quota=cpu_quota,
|
||||
storage="1g",
|
||||
storage="1g", # Default storage quota
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
|
|
@ -177,6 +169,9 @@ class JupyterBackend:
|
|||
|
||||
Returns:
|
||||
Dictionary with updated state info
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
self.session_manager.document_state(
|
||||
session_id=session_id,
|
||||
|
|
@ -185,37 +180,91 @@ class JupyterBackend:
|
|||
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."""
|
||||
"""
|
||||
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."""
|
||||
"""
|
||||
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."""
|
||||
"""
|
||||
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."""
|
||||
"""
|
||||
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.default_memory,
|
||||
cpu_quota=self.default_cpu_quota,
|
||||
storage="1g",
|
||||
timeout=self.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.max_timeout:
|
||||
raise ValueError(
|
||||
f"Timeout {timeout} exceeds maximum {self.max_timeout}"
|
||||
)
|
||||
|
||||
# Validate memory
|
||||
memory_bytes = parse_memory_string(memory)
|
||||
max_memory_bytes = parse_memory_string(self.max_memory)
|
||||
if memory_bytes > max_memory_bytes:
|
||||
|
|
@ -223,6 +272,7 @@ class JupyterBackend:
|
|||
f"Memory {memory} exceeds maximum {self.max_memory}"
|
||||
)
|
||||
|
||||
# Validate CPU quota
|
||||
if cpu_quota > self.max_cpu_quota:
|
||||
raise ValueError(
|
||||
f"CPU quota {cpu_quota} exceeds maximum {self.max_cpu_quota}"
|
||||
|
|
|
|||
|
|
@ -108,15 +108,15 @@ class SessionManager:
|
|||
Initialize session manager.
|
||||
|
||||
Args:
|
||||
kernel_manager: Kernel lifecycle manager
|
||||
idle_timeout: Session idle timeout in seconds (default: 3600)
|
||||
kernel_manager: Jupyter kernel manager instance
|
||||
idle_timeout: Seconds before idle session cleanup (default: 3600)
|
||||
max_sessions: Maximum concurrent sessions (default: 10)
|
||||
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||
audit_logger: Optional audit logger (defaults to NullAuditLogger)
|
||||
"""
|
||||
self.kernel_manager = kernel_manager
|
||||
self.idle_timeout = idle_timeout
|
||||
self.max_sessions = max_sessions
|
||||
self.audit_logger = audit_logger or NullAuditLogger()
|
||||
self.kernel_manager = kernel_manager
|
||||
self.audit_logger = audit_logger if audit_logger is not None else NullAuditLogger()
|
||||
self.sessions: Dict[str, Session] = {}
|
||||
|
||||
def create_session(
|
||||
|
|
@ -171,8 +171,8 @@ class SessionManager:
|
|||
|
||||
# Log session creation
|
||||
self.audit_logger.log(
|
||||
event_type="SESSION_CREATE,
|
||||
severity="INFO,
|
||||
event_type="session.create",
|
||||
severity="info",
|
||||
message=f"Session created: {session_id}",
|
||||
session_id=session_id,
|
||||
details={
|
||||
|
|
@ -373,8 +373,8 @@ class SessionManager:
|
|||
except Exception as e:
|
||||
# Log but continue with cleanup
|
||||
self.audit_logger.log(
|
||||
event_type="SESSION_DESTROY,
|
||||
severity="WARNING,
|
||||
event_type="session.destroy",
|
||||
severity="warning",
|
||||
message=f"Error shutting down kernel for session {session_id}",
|
||||
session_id=session_id,
|
||||
error=str(e)
|
||||
|
|
@ -385,8 +385,8 @@ class SessionManager:
|
|||
|
||||
# Log destruction
|
||||
self.audit_logger.log(
|
||||
event_type="SESSION_DESTROY,
|
||||
severity="INFO,
|
||||
event_type="session.destroy",
|
||||
severity="info",
|
||||
message=f"Session destroyed: {session_id}",
|
||||
session_id=session_id
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue