/headroom/proxy/ws_session_registry.py

"""WebSocket session registry for Codex relay lifecycle tracking.

Unit 3 of the Codex-proxy resilience plan. Every ``/v1/responses`` WS
session is explicitly registered on accept and deregistered in the
outermost ``finally`` of the handler. The registry provides:

* First-class visibility of active sessions (``/debug/ws-sessions``
  from Unit 5 consumes :meth:`WebSocketSessionRegistry.snapshot`).
* Gauges for Prometheus (``active_ws_sessions``, ``active_relay_tasks``).
* A home for relay-task references so the handler's orchestrator can
  attach them at creation time for introspection, without the registry
  itself owning or cancelling them — cancellation is the handler's job.

Design notes
------------
* Single-event-loop usage. All mutations happen on the proxy's event
  loop from within the handler coroutine, so no asyncio.Lock / asyncio
  primitives are needed. Python dict mutations are atomic under the
  GIL; snapshot copies happen under the event loop's cooperative
  scheduling, so iteration + mutation cannot interleave across an
  ``await`` point.
* ``deregister`` must be idempotent. The handler calls it from the
  outermost ``finally`` — if upstream never connected, the session may
  have been registered or may not have been, depending on how far
  handshake got. Either way ``deregister`` is safe.
* ``deregister`` clears the handle's ``relay_tasks`` list so the
  registry does not retain references to task coroutine frames after a
  session ends.
"""

from __future__ import annotations

import logging
import time
from collections.abc import Iterable
from dataclasses import dataclass, field
from typing import Any, Literal, Protocol

logger = logging.getLogger(__name__)

__all__ = [
    "WebSocketSessionRegistry",
    "WSSessionHandle",
    "TerminationCause",
]


TerminationCause = Literal[
    "client_disconnect",
    "upstream_disconnect",
    "upstream_error",
    "client_error",
    "client_cancel",
    "response_completed",
    "client_timeout",
    "unknown",
]


class _TaskLike(Protocol):
    def done(self) -> bool: ...
    def cancel(self) -> bool: ...
    def get_name(self) -> str: ...


@dataclass
class WSSessionHandle:
    """Per-session state entry held in :class:`WebSocketSessionRegistry`.

    Timestamps use :func:`time.perf_counter` so age computations are
    monotonic and independent of wall-clock adjustments.
    """

    session_id: str
    request_id: str
    client_addr: str | None = None
    upstream_url: str | None = None
    started_at: float = field(default_factory=time.perf_counter)
    last_activity_at: float = field(default_factory=time.perf_counter)
    relay_tasks: list[_TaskLike] = field(default_factory=list)
    termination_cause: TerminationCause | None = None

    def mark_activity(self) -> None:
        self.last_activity_at = time.perf_counter()

    def age_seconds(self) -> float:
        return max(0.0, time.perf_counter() - self.started_at)

    def to_snapshot_dict(self) -> dict[str, Any]:
        return {
            "session_id": self.session_id,
            "request_id": self.request_id,
            "client_addr": self.client_addr,
            "upstream_url": self.upstream_url,
            "age_seconds": self.age_seconds(),
            "idle_seconds": max(0.0, time.perf_counter() - self.last_activity_at),
            "relay_task_count": len(self.relay_tasks),
            "relay_task_names": [t.get_name() for t in self.relay_tasks],
            "termination_cause": self.termination_cause,
        }


class WebSocketSessionRegistry:
    """In-memory registry of active Codex WS sessions.

    Methods are safe to call from the event loop. Unit 5's
    ``/debug/ws-sessions`` endpoint consumes :meth:`snapshot`; the
    Prometheus exporter reads :meth:`active_count` and
    :meth:`active_relay_task_count`.
    """

    def __init__(self) -> None:
        self._sessions: dict[str, WSSessionHandle] = {}
        # Tracked separately from ``sum(len(h.relay_tasks))`` so that
        # repeated snapshotting is O(1) rather than O(N * tasks).
        self._active_relay_tasks = 0

    # ------------------------------------------------------------------
    # Core lifecycle
    # ------------------------------------------------------------------

    def register(self, handle: WSSessionHandle) -> None:
        """Register a session. Idempotent by ``session_id``.

        If the session id is already present, the existing entry is
        replaced and the active count stays the same. This matches the
        invariant that one ``session_id`` corresponds to at most one
        active session at a time.
        """
        existing = self._sessions.get(handle.session_id)
        if existing is not None:
            # Re-registration: release any tasks we were tracking on the
            # old handle before replacing it.
            self._active_relay_tasks -= len(existing.relay_tasks)
            existing.relay_tasks.clear()
        self._sessions[handle.session_id] = handle
        self._active_relay_tasks += len(handle.relay_tasks)

    def deregister(
        self, session_id: str, cause: TerminationCause = "unknown"
    ) -> WSSessionHandle | None:
        """Remove a session. Idempotent: returns ``None`` if unknown.

        Also clears the handle's ``relay_tasks`` list so the registry
        stops holding references to coroutine frames after the session
        ends.

        The handle is returned with ``relay_tasks`` already cleared, but
        the caller can still observe the number of tasks the registry
        released via the ``released_tasks`` attribute set on the handle
        before the clear. Prefer :meth:`deregister_and_count` for new
        callers that need that number for Prometheus dec calls — it
        couples the handle pop and the count read so they cannot drift.
        """
        handle, _count = self._deregister_internal(session_id, cause)
        return handle

    def deregister_and_count(
        self, session_id: str, cause: TerminationCause = "unknown"
    ) -> tuple[WSSessionHandle | None, int]:
        """Deregister and return (handle, released_task_count).

        This is the preferred API for handlers that need to decrement a
        Prometheus gauge by the number of relay tasks that were attached
        to the session: capturing the count separately (via
        ``len(handle.relay_tasks)`` before :meth:`deregister`) and then
        calling :meth:`deregister` risks drift if the registry's
        bookkeeping is ever changed. ``deregister_and_count`` returns
        both atomically.

        Returns ``(None, 0)`` when the session was not registered.
        """
        return self._deregister_internal(session_id, cause)

    def _deregister_internal(
        self, session_id: str, cause: TerminationCause
    ) -> tuple[WSSessionHandle | None, int]:
        handle = self._sessions.pop(session_id, None)
        if handle is None:
            return None, 0
        handle.termination_cause = cause
        released = len(handle.relay_tasks)
        self._active_relay_tasks = max(0, self._active_relay_tasks - released)
        handle.relay_tasks.clear()
        return handle, released

    def attach_tasks(self, session_id: str, tasks: Iterable[_TaskLike]) -> None:
        """Attach relay tasks to an existing session (merge, not replace).

        If ``session_id`` is not registered, this is a no-op (handler
        should have registered before spawning tasks; defensive so a
        race during deregister doesn't crash the handler).
        """
        handle = self._sessions.get(session_id)
        if handle is None:
            return
        task_list = list(tasks)
        handle.relay_tasks.extend(task_list)
        self._active_relay_tasks += len(task_list)
        handle.mark_activity()

    # ------------------------------------------------------------------
    # Introspection
    # ------------------------------------------------------------------

    def get(self, session_id: str) -> WSSessionHandle | None:
        return self._sessions.get(session_id)

    def active_count(self) -> int:
        return len(self._sessions)

    def active_relay_task_count(self) -> int:
        return self._active_relay_tasks

    def snapshot(self) -> list[dict[str, Any]]:
        """JSON-serializable view of the registry (for ``/debug/ws-sessions``)."""
        return [handle.to_snapshot_dict() for handle in self._sessions.values()]

    # ------------------------------------------------------------------
    # Convenience
    # ------------------------------------------------------------------

    def __contains__(self, session_id: str) -> bool:
        return session_id in self._sessions

    def __len__(self) -> int:
        return len(self._sessions)