"""Session registry: tracks active project sessions by conversation ID.""" from __future__ import annotations import asyncio import json import logging from dataclasses import dataclass, field, asdict from pathlib import Path from typing import Optional from agent.sdk_session import ( SDKSession, SessionProgress, DEFAULT_PERMISSION_MODE, VALID_PERMISSION_MODES, ) logger = logging.getLogger(__name__) DEFAULT_IDLE_TIMEOUT = 30 * 60 PERSISTENCE_FILE = Path(__file__).parent.parent / "data" / "sessions.json" @dataclass class Session: conv_id: str cwd: str owner_id: str = "" last_activity: float = 0.0 idle_timeout: int = DEFAULT_IDLE_TIMEOUT permission_mode: str = field(default_factory=lambda: DEFAULT_PERMISSION_MODE) chat_id: str | None = None # Runtime only — not serialized sdk_session: SDKSession | None = field(default=None, repr=False) def touch(self) -> None: self.last_activity = asyncio.get_event_loop().time() def to_dict(self) -> dict: d = asdict(self) d.pop("sdk_session", None) return d @classmethod def from_dict(cls, data: dict) -> "Session": data.pop("sdk_session", None) # Migration: remove old cc_runner fields if present in persisted data for old_key in ("cc_session_id", "started", "cc_timeout"): data.pop(old_key, None) return cls(**data) class SessionManager: """Registry of active Claude Code project sessions with persistence and user isolation.""" def __init__(self) -> None: self._sessions: dict[str, Session] = {} self._lock = asyncio.Lock() self._reaper_task: Optional[asyncio.Task] = None async def start(self) -> None: self._load() loop = asyncio.get_event_loop() self._reaper_task = loop.create_task(self._reaper_loop()) async def stop(self) -> None: if self._reaper_task: self._reaper_task.cancel() async with self._lock: for session in self._sessions.values(): if session.sdk_session: await session.sdk_session.close() self._sessions.clear() if PERSISTENCE_FILE.exists(): PERSISTENCE_FILE.unlink() async def create( self, conv_id: str, working_dir: str, owner_id: str = "", idle_timeout: int = DEFAULT_IDLE_TIMEOUT, permission_mode: str = DEFAULT_PERMISSION_MODE, chat_id: str | None = None, ) -> Session: async with self._lock: session = Session( conv_id=conv_id, cwd=working_dir, owner_id=owner_id, idle_timeout=idle_timeout, permission_mode=permission_mode, chat_id=chat_id, ) self._sessions[conv_id] = session self._save() logger.info( "Created session %s (owner=...%s) in %s (idle=%ds, perm=%s)", conv_id, owner_id[-8:] if owner_id else "-", working_dir, idle_timeout, permission_mode, ) return session # --- Secretary model: async send (returns immediately) --- async def send_message( self, conv_id: str, message: str, user_id: Optional[str] = None, chat_id: Optional[str] = None, ) -> str: """Send a message to the session. Returns immediately; result pushed to Feishu on completion.""" session = self._get_session(conv_id, user_id) session.touch() self._ensure_sdk_session(session, chat_id) return await session.sdk_session.send(message, chat_id) async def send_and_wait( self, conv_id: str, message: str, user_id: Optional[str] = None, chat_id: Optional[str] = None, ) -> str: """Send and wait for completion. For LLM agent tool calls that need the result.""" session = self._get_session(conv_id, user_id) session.touch() self._ensure_sdk_session(session, chat_id) return await session.sdk_session.send_and_wait(message, chat_id) # --- Kept for backward compatibility (used by bot/commands.py _cmd_new) --- async def send( self, conv_id: str, message: str, user_id: Optional[str] = None, direct: bool = False, ) -> str: """Backward-compatible send. Maps to send_message (async, secretary model).""" from orchestrator.tools import get_current_chat chat_id = get_current_chat() return await self.send_message(conv_id, message, user_id=user_id, chat_id=chat_id) # --- Progress, interrupt, approve --- def get_progress(self, conv_id: str, user_id: Optional[str] = None) -> SessionProgress | None: """Query session progress. Primary interface for the secretary AI.""" session = self._sessions.get(conv_id) if not session: return None if session.owner_id and user_id and session.owner_id != user_id: return None if session.sdk_session: return session.sdk_session.get_progress() return SessionProgress() async def interrupt(self, conv_id: str, user_id: Optional[str] = None) -> bool: """Interrupt the currently running task in a session.""" session = self._get_session(conv_id, user_id) if session.sdk_session: await session.sdk_session.interrupt() return True return False async def approve(self, conv_id: str, approved: bool) -> None: """Resolve a pending tool approval for a session.""" session = self._sessions.get(conv_id) if session and session.sdk_session: await session.sdk_session.approve(approved) async def answer_question(self, conv_id: str, answers: dict[str, str]) -> None: """Resolve a pending AskUserQuestion with user's answers.""" session = self._sessions.get(conv_id) if session and session.sdk_session: await session.sdk_session.answer_question(answers) # --- Close, list, permission --- async def close(self, conv_id: str, user_id: Optional[str] = None) -> bool: async with self._lock: session = self._sessions.get(conv_id) if session is None: return False if session.owner_id and user_id and session.owner_id != user_id: raise PermissionError(f"Session {conv_id} belongs to another user") if session.sdk_session: await session.sdk_session.close() del self._sessions[conv_id] self._save() logger.info("Closed session %s", conv_id) return True def list_sessions(self, user_id: Optional[str] = None) -> list[dict]: sessions = self._sessions.values() if user_id: sessions = [s for s in sessions if not s.owner_id or s.owner_id == user_id] return [ { "conv_id": s.conv_id, "cwd": s.cwd, "owner_id": s.owner_id[-8:] if s.owner_id else None, "busy": s.sdk_session._busy if s.sdk_session else False, "idle_timeout": s.idle_timeout, "permission_mode": s.permission_mode, } for s in sessions ] def set_permission_mode(self, conv_id: str, mode: str, user_id: Optional[str] = None) -> None: """Change the permission mode for an existing session.""" session = self._sessions.get(conv_id) if session is None: raise KeyError(f"No session for conv_id={conv_id!r}") if session.owner_id and user_id and session.owner_id != user_id: raise PermissionError(f"Session {conv_id} belongs to another user") if mode not in VALID_PERMISSION_MODES: raise ValueError(f"Invalid permission mode {mode!r}. Valid: {VALID_PERMISSION_MODES}") session.permission_mode = mode if session.sdk_session: asyncio.create_task(session.sdk_session.set_permission_mode(mode)) self._save() logger.info("Set permission_mode=%s for session %s", mode, conv_id) # --- Internal --- def _get_session(self, conv_id: str, user_id: Optional[str] = None) -> Session: session = self._sessions.get(conv_id) if session is None: raise KeyError(f"No session for conv_id={conv_id!r}") if session.owner_id and user_id and session.owner_id != user_id: raise PermissionError(f"Session {conv_id} belongs to another user") return session def _ensure_sdk_session(self, session: Session, chat_id: str | None = None) -> None: if session.sdk_session is None: session.sdk_session = SDKSession( conv_id=session.conv_id, cwd=session.cwd, owner_id=session.owner_id, permission_mode=session.permission_mode, chat_id=chat_id or session.chat_id, ) def _save(self) -> None: try: data = {cid: s.to_dict() for cid, s in self._sessions.items()} PERSISTENCE_FILE.parent.mkdir(parents=True, exist_ok=True) with open(PERSISTENCE_FILE, "w", encoding="utf-8") as f: json.dump(data, f, indent=2) logger.debug("Saved %d sessions to %s", len(data), PERSISTENCE_FILE) except Exception: logger.exception("Failed to save sessions") def _load(self) -> None: if not PERSISTENCE_FILE.exists(): return try: with open(PERSISTENCE_FILE, "r", encoding="utf-8") as f: data = json.load(f) for cid, sdata in data.items(): self._sessions[cid] = Session.from_dict(sdata) logger.info("Loaded %d sessions from %s", len(self._sessions), PERSISTENCE_FILE) except Exception: logger.exception("Failed to load sessions") async def _reaper_loop(self) -> None: while True: await asyncio.sleep(60) await self._reap_idle() async def _reap_idle(self) -> None: now = asyncio.get_event_loop().time() async with self._lock: to_close = [] for cid, s in self._sessions.items(): if s.last_activity > 0 and (now - s.last_activity) > s.idle_timeout: to_close.append(cid) for cid in to_close: session = self._sessions[cid] if session.sdk_session: await session.sdk_session.close() del self._sessions[cid] logger.info("Reaped idle session %s", cid) if to_close: self._save() manager = SessionManager()