Session Management

┌─────────────────────────────────────────────────────────────┐
│                      SESSION LIFECYCLE                       │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  1. Client sends initialize (no session ID yet)             │
│     ──► POST /mcp                                           │
│         Content-Type: application/json                      │
│         {"method": "initialize", ...}                       │
│                                                             │
│  2. Server responds with new session ID                     │
│     ◄── 200 OK                                              │
│         Mcp-Session-Id: abc123-def456                       │
│         {"result": {"capabilities": ...}}                   │
│                                                             │
│  3. Client includes session ID in all future requests       │
│     ──► POST /mcp                                           │
│         Mcp-Session-Id: abc123-def456                       │
│         {"method": "tools/list", ...}                       │
│                                                             │
│  4. Server validates session and processes request          │
│     ◄── 200 OK                                              │
│         Mcp-Session-Id: abc123-def456                       │
│         {"result": {"tools": [...]}}                        │
│                                                             │
└─────────────────────────────────────────────────────────────┘

"""
session_store.py - MCP session management.
"""

import uuid
import time
from dataclasses import dataclass, field
from typing import Dict, Optional, Any


@dataclass
class MCPSession:
    """Represents an MCP client session."""
    session_id: str
    created_at: float = field(default_factory=time.time)
    last_active: float = field(default_factory=time.time)
    initialized: bool = False
    client_info: Dict[str, Any] = field(default_factory=dict)
    capabilities: Dict[str, Any] = field(default_factory=dict)

    def touch(self):
        """Update last active time."""
        self.last_active = time.time()

    def is_expired(self, timeout_seconds: int = 3600) -> bool:
        """Check if session has expired."""
        return time.time() - self.last_active > timeout_seconds


class SessionStore:
    """Store and manage MCP sessions."""

    def __init__(self, session_timeout: int = 3600):
        """
        Initialize session store.

        Args:
            session_timeout: Session expiry in seconds (default: 1 hour)
        """
        self.sessions: Dict[str, MCPSession] = {}
        self.session_timeout = session_timeout

    def create_session(self) -> MCPSession:
        """Create a new session."""
        session_id = str(uuid.uuid4())
        session = MCPSession(session_id=session_id)
        self.sessions[session_id] = session
        return session

    def get_session(self, session_id: str) -> Optional[MCPSession]:
        """
        Get a session by ID.

        Returns None if not found or expired.
        """
        session = self.sessions.get(session_id)
        if session is None:
            return None
        if session.is_expired(self.session_timeout):
            self.delete_session(session_id)
            return None
        session.touch()
        return session

    def delete_session(self, session_id: str) -> bool:
        """Delete a session. Returns True if deleted."""
        if session_id in self.sessions:
            del self.sessions[session_id]
            return True
        return False

    def cleanup_expired(self):
        """Remove all expired sessions."""
        expired = [
            sid for sid, session in self.sessions.items()
            if session.is_expired(self.session_timeout)
        ]
        for sid in expired:
            del self.sessions[sid]
        return len(expired)


# Global session store (or inject via dependency)
sessions = SessionStore()

from session_store import sessions, MCPSession


class MCPHTTPHandler(BaseHTTPRequestHandler):
    """MCP handler with session support."""

    def get_session(self) -> Optional[MCPSession]:
        """Get session from request header."""
        session_id = self.headers.get("Mcp-Session-Id")
        if session_id:
            return sessions.get_session(session_id)
        return None

    def add_session_header(self, session: MCPSession):
        """Add session ID to response headers."""
        self.send_header("Mcp-Session-Id", session.session_id)

    def do_POST(self):
        """Handle POST with session management."""
        if self.path != "/mcp":
            self.send_error(404, "Not Found")
            return

        # Parse request
        body = self.read_request_body()
        request, error = self.parse_json_rpc(body)
        if error:
            self.send_json_response(error, status=400)
            return

        method = request.get("method")

        # Handle initialize (creates new session)
        if method == "initialize":
            session = sessions.create_session()
            response = self.handle_initialize(request, session)
            self.send_json_response_with_session(response, session)
            return

        # All other methods require existing session
        session = self.get_session()
        if session is None:
            self.send_json_response({
                "jsonrpc": "2.0",
                "id": request.get("id"),
                "error": {
                    "code": -32600,
                    "message": "No valid session. Call 'initialize' first."
                }
            }, status=400)
            return

        # Check if session is initialized
        if not session.initialized and method != "notifications/initialized":
            self.send_json_response({
                "jsonrpc": "2.0",
                "id": request.get("id"),
                "error": {
                    "code": -32600,
                    "message": "Session not initialized"
                }
            }, status=400)
            return

        # Handle the request
        response = self.handle_mcp_request(request, session)
        self.send_json_response_with_session(response, session)

    def send_json_response_with_session(self, data, session):
        """Send JSON response with session header."""
        body = json.dumps(data).encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.add_session_header(session)
        self.end_headers()
        self.wfile.write(body)

def handle_initialize(self, request, session: MCPSession):
    """Handle initialize and set up session."""
    params = request.get("params", {})
    request_id = request.get("id")

    # Store client info in session
    session.client_info = params.get("clientInfo", {})
    session.capabilities = params.get("capabilities", {})

    # Server capabilities
    server_capabilities = {
        "tools": {},
        "resources": {},
        "prompts": {}
    }

    # Build response
    response = {
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {
            "protocolVersion": "2024-11-05",
            "capabilities": server_capabilities,
            "serverInfo": {
                "name": "my-mcp-server",
                "version": "1.0.0"
            }
        }
    }

    return response

def handle_initialized_notification(self, request, session: MCPSession):
    """Handle notifications/initialized."""
    # Mark session as fully initialized
    session.initialized = True
    # Notifications don't return a response
    return None