""" Session Management Service for Lovdata Chat This service manages the lifecycle of OpenCode containers for individual user sessions. Each session gets its own isolated container with a dedicated working directory. """ import os import uuid import json import asyncio from datetime import datetime, timedelta from pathlib import Path from typing import Dict, Optional, List from contextlib import asynccontextmanager import docker from docker.errors import DockerException, NotFound from fastapi import FastAPI, HTTPException, BackgroundTasks, Request, Response from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel import uvicorn import httpx # Configuration SESSIONS_DIR = Path("/app/sessions") SESSIONS_FILE = Path("/app/sessions/sessions.json") CONTAINER_IMAGE = "lovdata-opencode:latest" MAX_CONCURRENT_SESSIONS = 3 # Workstation limit SESSION_TIMEOUT_MINUTES = 60 # Auto-cleanup after 1 hour CONTAINER_MEMORY_LIMIT = "4g" CONTAINER_CPU_QUOTA = 100000 # 1 CPU core class SessionData(BaseModel): session_id: str container_name: str container_id: Optional[str] = None host_dir: str port: Optional[int] = None created_at: datetime last_accessed: datetime status: str = "creating" # creating, running, stopped, error class SessionManager: def __init__(self): # Use Docker library 7.1.0 with Unix socket support import docker self.docker_client = docker.from_env() # Test the connection self.docker_client.ping() print("Docker library client initialized successfully") self.sessions: Dict[str, SessionData] = {} self._load_sessions() def _load_sessions(self): """Load session data from persistent storage""" if SESSIONS_FILE.exists(): try: with open(SESSIONS_FILE, "r") as f: data = json.load(f) for session_id, session_dict in data.items(): # Convert datetime strings back to datetime objects session_dict["created_at"] = datetime.fromisoformat( session_dict["created_at"] ) session_dict["last_accessed"] = datetime.fromisoformat( session_dict["last_accessed"] ) self.sessions[session_id] = SessionData(**session_dict) except (json.JSONDecodeError, KeyError) as e: print(f"Warning: Could not load sessions file: {e}") self.sessions = {} def _save_sessions(self): """Save session data to persistent storage""" SESSIONS_DIR.mkdir(exist_ok=True) data = {} for session_id, session in self.sessions.items(): data[session_id] = session.dict() with open(SESSIONS_FILE, "w") as f: json.dump(data, f, indent=2, default=str) def _generate_session_id(self) -> str: """Generate a unique session ID""" return str(uuid.uuid4()).replace("-", "")[:16] def _get_available_port(self) -> int: """Find an available port for the container""" used_ports = {s.port for s in self.sessions.values() if s.port} port = 8081 # Start from 8081 to avoid conflicts while port in used_ports: port += 1 return port def _check_container_limits(self) -> bool: """Check if we're within concurrent session limits""" if not self.docker_client: return False active_sessions = sum( 1 for s in self.sessions.values() if s.status in ["creating", "running"] ) return active_sessions < MAX_CONCURRENT_SESSIONS async def create_session(self) -> SessionData: """Create a new OpenCode session with dedicated container""" if not self._check_container_limits(): raise HTTPException( status_code=429, detail=f"Maximum concurrent sessions ({MAX_CONCURRENT_SESSIONS}) reached", ) session_id = self._generate_session_id() container_name = f"opencode-{session_id}" host_dir = str(SESSIONS_DIR / session_id) port = self._get_available_port() # Create host directory Path(host_dir).mkdir(parents=True, exist_ok=True) session = SessionData( session_id=session_id, container_name=container_name, host_dir=host_dir, port=port, created_at=datetime.now(), last_accessed=datetime.now(), status="creating", ) self.sessions[session_id] = session self._save_sessions() # Start container in background asyncio.create_task(self._start_container(session)) return session async def _start_container(self, session: SessionData): """Start the OpenCode container for a session""" try: # Create and start the OpenCode container container = self.docker_client.containers.run( "lovdata-opencode:latest", # Will be built from the Dockerfile name=session.container_name, volumes={session.host_dir: {"bind": "/app/somedir", "mode": "rw"}}, ports={f"8080/tcp": session.port}, detach=True, environment={ "MCP_SERVER": os.getenv("MCP_SERVER", ""), "OPENAI_API_KEY": os.getenv("OPENAI_API_KEY", ""), "ANTHROPIC_API_KEY": os.getenv("ANTHROPIC_API_KEY", ""), "GOOGLE_API_KEY": os.getenv("GOOGLE_API_KEY", ""), }, network_mode="bridge", ) session.container_id = container.id session.status = "running" self._save_sessions() print(f"Container {session.container_name} started on port {session.port}") except Exception as e: session.status = "error" self._save_sessions() print(f"Failed to start container {session.container_name}: {e}") async def get_session(self, session_id: str) -> Optional[SessionData]: """Get session information""" session = self.sessions.get(session_id) if session: session.last_accessed = datetime.now() self._save_sessions() return session async def list_sessions(self) -> List[SessionData]: """List all sessions""" return list(self.sessions.values()) async def cleanup_expired_sessions(self): """Clean up expired sessions and their containers""" now = datetime.now() expired_sessions = [] for session_id, session in self.sessions.items(): # Check if session has expired if now - session.last_accessed > timedelta(minutes=SESSION_TIMEOUT_MINUTES): expired_sessions.append(session_id) # Stop and remove container try: container = self.docker_client.containers.get( session.container_name ) container.stop(timeout=10) container.remove() print(f"Cleaned up container {session.container_name}") except Exception as e: print(f"Error cleaning up container {session.container_name}: {e}") # Remove session directory try: import shutil shutil.rmtree(session.host_dir) print(f"Removed session directory {session.host_dir}") except OSError as e: print(f"Error removing session directory {session.host_dir}: {e}") for session_id in expired_sessions: del self.sessions[session_id] if expired_sessions: self._save_sessions() print(f"Cleaned up {len(expired_sessions)} expired sessions") # Global session manager instance session_manager = SessionManager() @asynccontextmanager async def lifespan(app: FastAPI): """Application lifespan manager""" # Startup print("Starting Session Management Service") # Start cleanup task async def cleanup_task(): while True: await session_manager.cleanup_expired_sessions() await asyncio.sleep(300) # Run every 5 minutes cleanup_coro = asyncio.create_task(cleanup_task()) yield # Shutdown print("Shutting down Session Management Service") cleanup_coro.cancel() app = FastAPI( title="Lovdata Chat Session Manager", description="Manages isolated OpenCode containers for Norwegian legal research sessions", version="1.0.0", lifespan=lifespan, ) @app.post("/sessions", response_model=SessionData) async def create_session(): """Create a new session with dedicated container""" try: session = await session_manager.create_session() return session except HTTPException: raise except Exception as e: raise HTTPException( status_code=500, detail=f"Failed to create session: {str(e)}" ) @app.get("/sessions/{session_id}", response_model=SessionData) async def get_session(session_id: str): """Get session information""" session = await session_manager.get_session(session_id) if not session: raise HTTPException(status_code=404, detail="Session not found") return session @app.get("/sessions", response_model=List[SessionData]) async def list_sessions(): """List all active sessions""" return await session_manager.list_sessions() @app.delete("/sessions/{session_id}") async def delete_session(session_id: str, background_tasks: BackgroundTasks): """Delete a session and its container""" session = await session_manager.get_session(session_id) if not session: raise HTTPException(status_code=404, detail="Session not found") # Schedule cleanup background_tasks.add_task(session_manager.cleanup_expired_sessions) # Remove from sessions immediately del session_manager.sessions[session_id] session_manager._save_sessions() return {"message": f"Session {session_id} scheduled for deletion"} @app.post("/cleanup") async def trigger_cleanup(): """Manually trigger cleanup of expired sessions""" await session_manager.cleanup_expired_sessions() return {"message": "Cleanup completed"} @app.api_route( "/session/{session_id}/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], ) async def proxy_to_session(request: Request, session_id: str, path: str): """Proxy requests to session containers based on session ID in URL""" session = await session_manager.get_session(session_id) if not session or session.status != "running": raise HTTPException(status_code=404, detail="Session not found or not running") # Proxy the request to the container (use host.docker.internal for Docker Desktop, or host IP) # For Linux, we need to use the host IP host_ip = os.getenv("HOST_IP", "172.17.0.1") # Default Docker bridge IP container_url = f"http://{host_ip}:{session.port}" # Prepare the request URL url = f"{container_url}/{path}" if request.url.query: url += f"?{request.url.query}" # Get request body body = await request.body() # Prepare headers (exclude host header to avoid conflicts) headers = dict(request.headers) headers.pop("host", None) # Make the proxy request async with httpx.AsyncClient(timeout=30.0) as client: try: response = await client.request( method=request.method, url=url, headers=headers, content=body, follow_redirects=False, # Let the client handle redirects ) # Return the response return Response( content=response.content, status_code=response.status_code, headers=dict(response.headers), ) except httpx.TimeoutException: raise HTTPException( status_code=504, detail="Request to session container timed out" ) except httpx.RequestError as e: raise HTTPException( status_code=502, detail=f"Failed to connect to session container: {str(e)}", ) @app.get("/health") async def health_check(): """Health check endpoint""" docker_ok = False try: # Check Docker connectivity session_manager.docker_client.ping() docker_ok = True except: docker_ok = False return { "status": "healthy" if docker_ok else "unhealthy", "docker": docker_ok, "active_sessions": len( [s for s in session_manager.sessions.values() if s.status == "running"] ), "timestamp": datetime.now().isoformat(), } if __name__ == "__main__": uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)