"""Per-notebook session manager for OpenCode. Maps each notebookPath to one OR MORE OpenCode session IDs. One of the bound sessions is the "active" one (the one used for the next ``/edit``). The others are historical — they were created in this notebook at some point but the user switched away (e.g. started a new conversation topic). Lifecycle: - ``get_or_create`` returns the active session, lazily creating one on first use. This is the path used by ``EditHandler`` on every prompt. - ``create_new`` always creates a brand-new session on OpenCode Serve, binds it, sets it as active. Used by the "New session" UI action. - ``bind_existing`` / ``set_active`` / ``unbind`` / ``delete_session`` are the management operations for the per-notebook session list. Async-safety: a per-notebook ``asyncio.Lock`` serializes ``get_or_create`` so concurrent prompts on the same notebook don't double-create. The lock survives across calls (held in ``_locks``) so subsequent calls can reuse it. """ from __future__ import annotations import asyncio import logging from typing import Callable, Optional from .opencode_client import OpenCodeClient log = logging.getLogger("opencode_bridge.session_manager") ClientFactory = Callable[[], OpenCodeClient] class SessionManager: def __init__(self, client_factory: ClientFactory) -> None: self._client_factory = client_factory # notebookPath -> active sessionId self._active: dict[str, str] = {} # notebookPath -> ordered list of sessionIds bound to this # notebook (the first is the oldest; the active one may be # anywhere in the list). Used by the UI to render the session # picker. self._bound: dict[str, list[str]] = {} # Per-notebook lock for serializing lazy-create. self._locks: dict[str, asyncio.Lock] = {} # ---------- active-session lookup / create ---------- async def get_or_create(self, notebook_path: str) -> str: """Return the active session id for this notebook, creating one if no active session is bound yet. This is what ``EditHandler`` calls on every prompt. """ active = self._active.get(notebook_path) if active is not None: return active lock = self._locks.setdefault(notebook_path, asyncio.Lock()) async with lock: active = self._active.get(notebook_path) if active is not None: return active return await self._create_and_bind(notebook_path, title=None) async def _create_and_bind( self, notebook_path: str, title: Optional[str] ) -> str: client = self._client_factory() session = await client.create_session( title=title or "jupyter:%s" % notebook_path ) sid = session["id"] self._bind(notebook_path, sid) log.info("created opencode session %s for %s", sid, notebook_path) return sid def get_active(self, notebook_path: str) -> Optional[str]: """Return the active sessionId for the notebook, or None if no active session has been bound/created yet. Does NOT create one. """ return self._active.get(notebook_path) def peek(self, notebook_path: str) -> Optional[str]: """Alias for :meth:`get_active` (legacy name).""" return self.get_active(notebook_path) def has_session(self, notebook_path: str) -> bool: """True iff at least one session has ever been bound to this notebook (active or not).""" return notebook_path in self._bound # ---------- mutation: bind / unbind / set_active / create_new ---------- def _bind(self, notebook_path: str, session_id: str) -> None: """Add session_id to the bound list (if not already) and set it as active. Idempotent on re-bind. """ self._active[notebook_path] = session_id bound = self._bound.setdefault(notebook_path, []) if session_id not in bound: bound.append(session_id) async def create_new( self, notebook_path: str, title: Optional[str] = None ) -> str: """Create a brand new session on OpenCode Serve, bind it to this notebook, set as active, return its id. """ return await self._create_and_bind(notebook_path, title) def bind_existing(self, notebook_path: str, session_id: str) -> None: """Bind an existing sessionId to this notebook and set as active. Does NOT call OpenCode — the caller must verify the session exists (e.g. via ``GET /session``). """ self._bind(notebook_path, session_id) def set_active(self, notebook_path: str, session_id: str) -> bool: """Switch the active session. Returns False if ``session_id`` is not currently bound to this notebook (callers should bind first if uncertain). """ bound = self._bound.get(notebook_path, []) if session_id not in bound: return False self._active[notebook_path] = session_id return True def unbind(self, notebook_path: str, session_id: str) -> bool: """Remove ``session_id`` from this notebook's bindings. If it was the active one, also clear active (next prompt will lazy-create a new session). Returns True if it was bound. """ bound = self._bound.get(notebook_path) if not bound or session_id not in bound: return False bound.remove(session_id) if self._active.get(notebook_path) == session_id: self._active.pop(notebook_path, None) if not bound: self._bound.pop(notebook_path, None) return True def unbind_active(self, notebook_path: str) -> bool: """Remove the active session binding only (leaves the rest of the notebook's bound sessions intact). Next prompt will lazy-create a new active session. Returns True if there was an active binding to remove. """ active = self._active.pop(notebook_path, None) if active is None: return False bound = self._bound.get(notebook_path) if bound is not None: bound.remove(active) if not bound: self._bound.pop(notebook_path, None) return True async def delete_session( self, notebook_path: str, session_id: str ) -> bool: """Delete the session on OpenCode Serve AND remove its binding. Order matters: we ask OpenCode to delete FIRST, and only remove the local binding if OpenCode confirmed the delete. That way a failed delete leaves the binding intact (the user can retry); a deleted session never leaves a phantom binding behind. Returns True if the session was actually deleted (i.e. it was bound AND OpenCode acknowledged the delete). If the session is not bound to this notebook, do NOT call OpenCode (we don't want to silently delete sessions owned by other tools). """ if not self._is_bound(notebook_path, session_id): return False try: client = self._client_factory() deleted = await client.delete_session(session_id) except Exception: log.warning("failed to delete opencode session %s", session_id) return False # Whether the delete succeeded or just returned False (404), the # session is gone on the OpenCode side — drop the binding to # keep the UI consistent. We return whether OpenCode said # "yes" so the caller can surface a warning on a 404. self.unbind(notebook_path, session_id) return deleted def _is_bound(self, notebook_path: str, session_id: str) -> bool: bound = self._bound.get(notebook_path) return bool(bound) and session_id in bound # ---------- legacy / compat ---------- async def release(self, notebook_path: str) -> bool: """Tear down ALL state for this notebook: delete every bound session on OpenCode Serve AND clear every local binding. Returns True iff at least one session was actually deleted on the OpenCode side. (For deleting a SPECIFIC session use :meth:`delete_session`.) """ bound = list(self._bound.get(notebook_path, [])) deleted_any = False for sid in bound: if await self.delete_session(notebook_path, sid): deleted_any = True # `delete_session` clears the binding. Belt-and-suspenders: drop # any straggler state. self._bound.pop(notebook_path, None) self._active.pop(notebook_path, None) self._locks.pop(notebook_path, None) return deleted_any def invalidate(self, notebook_path: str) -> bool: """Drop the cached active sessionId without calling OpenCode. Returns True if a binding was removed. Use this when an upstream error indicates the session is dead (e.g. 404). """ was_active = self._active.pop(notebook_path, None) if was_active is None: return False bound = self._bound.get(notebook_path) if bound is not None: try: bound.remove(was_active) except ValueError: pass if not bound: self._bound.pop(notebook_path, None) return True def list_sessions(self) -> list[dict]: """Debug: list all (notebookPath, activeSessionId) pairs across every notebook. One row per notebook — does NOT enumerate the full bound-list per notebook. Use :meth:`list_sessions_for_notebook` for that. """ return [ {"notebookPath": path, "sessionId": sid} for path, sid in sorted(self._active.items()) ] def list_sessions_for_notebook(self, notebook_path: str) -> list[str]: """All sessionIds bound to this notebook (in bind order: oldest first, newest last). Empty list if none. """ return list(self._bound.get(notebook_path, []))