headroom/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, Protocol

logger = logging.getLogger(__name__)

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


TerminationCause = str  # "client_disconnect" | "upstream_disconnect" | "upstream_error"
# | "client_error" | "response_completed" | "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[Any] = 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": [
                getattr(t, "get_name", lambda: "")() 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.
        """
        handle = self._sessions.pop(session_id, None)
        if handle is None:
            return None
        handle.termination_cause = cause
        self._active_relay_tasks = max(
            0, self._active_relay_tasks - len(handle.relay_tasks)
        )
        handle.relay_tasks.clear()
        return handle

    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)