"""Session management for Python REPL sessions."""
import sys
import os
import uuid
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional, List
from dotenv import load_dotenv
import logging
from .models import SessionInfo, ExecutionResult
from .code_executor import CodeExecutor
logger = logging.getLogger(__name__)
def setup_project_venv(project_dir: Path):
"""Set up virtual environment for project directory."""
import subprocess
venv_path = project_dir / ".venv"
if not venv_path.exists():
logger.info(f"No .venv found in {project_dir}, creating one...")
# Initialize project if no pyproject.toml
if not (project_dir / "pyproject.toml").exists():
logger.info(f"No pyproject.toml found, running uv init...")
subprocess.run(["uv", "init"], cwd=project_dir, check=True)
# Create virtual environment
logger.info(f"Creating virtual environment...")
subprocess.run(["uv", "venv"], cwd=project_dir, check=True)
# Install dependencies if they exist
if (project_dir / "pyproject.toml").exists():
logger.info(f"Installing dependencies...")
subprocess.run(["uv", "sync"], cwd=project_dir, check=False) # Don't fail if no deps
return venv_path
def get_venv_python_path(project_dir: Path):
"""Get the Python executable path for the project's venv."""
venv_path = project_dir / ".venv"
if venv_path.exists():
python_exe = venv_path / "bin" / "python"
if python_exe.exists():
return python_exe
return None
def get_venv_site_packages(project_dir: Path):
"""Get the site-packages path for the project's venv."""
venv_path = project_dir / ".venv"
if venv_path.exists():
site_packages = venv_path / "lib" / f"python{sys.version_info.major}.{sys.version_info.minor}" / "site-packages"
if site_packages.exists():
return site_packages
return None
class Session:
"""Individual REPL session with isolated virtual environment."""
def __init__(self, session_id: str, name: str, working_directory: str):
self.id = session_id
self.name = name
self.working_directory = Path(working_directory).resolve()
self.created_at = datetime.utcnow()
self.active = True
self.executor = CodeExecutor()
self.env_vars: Dict[str, str] = {}
# Venv-related state
self.venv_path = None
self.venv_python = None
self.venv_site_packages = None
self.project_python_path = []
def initialize(self):
"""Initialize session environment with proper venv setup."""
logger.info(f"Initializing session {self.id} in {self.working_directory}")
# Ensure directory exists
self.working_directory.mkdir(parents=True, exist_ok=True)
# Set up or find virtual environment
self.venv_path = setup_project_venv(self.working_directory)
self.venv_python = get_venv_python_path(self.working_directory)
self.venv_site_packages = get_venv_site_packages(self.working_directory)
# Build Python path for this session
self.project_python_path = []
# Add project root for local imports
self.project_python_path.append(str(self.working_directory))
# Add venv site-packages if available
if self.venv_site_packages:
self.project_python_path.append(str(self.venv_site_packages))
logger.info(f"Using virtual environment: {self.venv_path}")
else:
logger.warning(f"No virtual environment site-packages found for {self.working_directory}")
# Load .env file if it exists
env_file = self.working_directory / ".env"
if env_file.exists():
# Change to project directory temporarily for .env loading
original_cwd = os.getcwd()
try:
os.chdir(self.working_directory)
load_dotenv(env_file)
logger.info(f"Loaded environment variables from {env_file}")
finally:
os.chdir(original_cwd)
logger.info(f"Session {self.id} ready with Python path: {self.project_python_path}")
async def execute_code(self, code: str) -> ExecutionResult:
"""Execute code in this session's venv context."""
if not self.active:
raise ValueError(f"Session {self.id} is not active")
# Store original state
original_cwd = os.getcwd()
original_sys_path = sys.path.copy()
try:
# Change to project directory
os.chdir(self.working_directory)
# Prepend session's Python path to sys.path
sys.path = self.project_python_path + sys.path
# Execute the code with session's context
result = await self.executor.execute(code)
return result
finally:
# Always restore original state
os.chdir(original_cwd)
sys.path[:] = original_sys_path
def cleanup(self):
"""Clean up session environment."""
self.active = False
logger.info(f"Cleaned up session {self.id}")
def set_env_var(self, name: str, value: str):
"""Set environment variable for this session."""
os.environ[name] = value
self.env_vars[name] = value
logger.debug(f"Set environment variable {name} in session {self.id}")
def get_env_vars(self) -> Dict[str, str]:
"""Get all environment variables."""
return dict(os.environ)
def to_session_info(self) -> SessionInfo:
"""Convert to SessionInfo model."""
return SessionInfo(
id=self.id,
name=self.name,
working_directory=str(self.working_directory),
created_at=self.created_at,
active=self.active
)
class SessionManager:
"""Manages multiple REPL sessions."""
def __init__(self):
self.sessions: Dict[str, Session] = {}
self.active_session_id: Optional[str] = None
async def create_session(self, name: Optional[str] = None, working_directory: Optional[str] = None) -> str:
"""Create a new REPL session."""
session_id = str(uuid.uuid4())
if name is None:
name = f"session-{len(self.sessions) + 1}"
if working_directory is None:
working_directory = os.getcwd()
# Validate that working directory exists
working_dir_path = Path(working_directory).resolve()
if not working_dir_path.exists():
raise ValueError(f"Working directory does not exist: {working_directory}")
# Create session
session = Session(session_id, name, str(working_dir_path))
session.initialize()
self.sessions[session_id] = session
# Make this the active session if it's the first one
if self.active_session_id is None:
self.active_session_id = session_id
logger.info(f"Created session {session_id} ({name}) in {working_directory}")
return session_id
def get_session(self, session_id: Optional[str] = None) -> Session:
"""Get session by ID, or active session if no ID provided."""
if session_id is None:
session_id = self.active_session_id
if session_id is None:
raise ValueError("No active session")
if session_id not in self.sessions:
raise ValueError(f"Session {session_id} not found")
return self.sessions[session_id]
def list_sessions(self) -> List[SessionInfo]:
"""List all sessions."""
return [session.to_session_info() for session in self.sessions.values()]
def switch_session(self, session_id: str) -> bool:
"""Switch to a different session."""
if session_id not in self.sessions:
raise ValueError(f"Session {session_id} not found")
session = self.sessions[session_id]
if not session.active:
raise ValueError(f"Session {session_id} is not active")
# Clean up current session if different
if self.active_session_id and self.active_session_id != session_id:
current_session = self.sessions[self.active_session_id]
current_session.cleanup()
# Initialize the new session
session.initialize()
self.active_session_id = session_id
logger.info(f"Switched to session {session_id}")
return True
def close_session(self, session_id: str) -> bool:
"""Close and remove a session."""
if session_id not in self.sessions:
raise ValueError(f"Session {session_id} not found")
session = self.sessions[session_id]
session.cleanup()
# Remove from sessions
del self.sessions[session_id]
# If this was the active session, clear it
if self.active_session_id == session_id:
self.active_session_id = None
# If there are other sessions, make the first one active
if self.sessions:
next_session_id = next(iter(self.sessions.keys()))
self.switch_session(next_session_id)
logger.info(f"Closed session {session_id}")
return True
async def execute_code(self, code: str, session_id: Optional[str] = None) -> ExecutionResult:
"""Execute code in specified session or active session."""
session = self.get_session(session_id)
return await session.execute_code(code)
def set_env_var(self, name: str, value: str, session_id: Optional[str] = None):
"""Set environment variable in specified session."""
session = self.get_session(session_id)
session.set_env_var(name, value)
def get_env_vars(self, session_id: Optional[str] = None) -> Dict[str, str]:
"""Get environment variables from specified session."""
session = self.get_session(session_id)
return session.get_env_vars()
def cleanup_all_sessions(self):
"""Clean up all sessions."""
for session in self.sessions.values():
if session.active:
session.cleanup()
self.sessions.clear()
self.active_session_id = None
logger.info("Cleaned up all sessions")
