mcp-forge/src/pod_executor/jupyter/sessions.py

"""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"
            )