438 lines
13 KiB
Python
438 lines
13 KiB
Python
|
"""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"
|
||
|
|
)
|