notebook-ai-extension/opencode_bridge/session_manager.py

"""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, []))