"""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: 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 (defaults to NullAuditLogger) """ self.idle_timeout = idle_timeout self.max_sessions = max_sessions 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( 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" )