teng-lin--notebooklm-py/src/notebooklm/_runtime/auth.py

"""Auth refresh coordinator helper for the client runtime.

Owns the auth refresh state machine and snapshot serialization:

* ``_refresh_lock`` — single-flight lock guarding refresh-task creation. Lazy
  because ``asyncio.Lock()`` needs a running loop in some Python versions and
  the coordinator can be constructed outside one.
* ``_refresh_task`` — the shared in-flight refresh task. Slot is intentionally
  preserved across waiter cancellation so siblings can still join, and is
  replaced only on the next refresh wave once the existing task hits
  ``done()`` (see :meth:`await_refresh` docstring).
* ``_refresh_callback`` — the user-supplied async callable that performs the
  actual refresh. ``None`` disables refresh-on-401.
* ``_auth_snapshot_lock`` — serializes the four-scalar reads in
  :meth:`snapshot` with the two-scalar writes in :meth:`update_auth_tokens`
  so RPC snapshots cannot observe a torn ``(csrf, session_id)`` pair while
  refresh is in flight. Intentionally distinct from ``_refresh_lock``:
  mixing them would re-introduce the reentrancy ambiguity that
  snapshot-side serialization was added to avoid.

Design constraints (load-bearing — see tests/unit/test_refresh_*.py and
tests/integration/concurrency/test_refresh_cancellation_propagation.py):

* ``__init__`` MUST be event-loop-agnostic — it stores only a plain callable
  and ``None`` placeholders. Never call ``asyncio.get_running_loop()`` or
  instantiate ``asyncio.*`` primitives at construction time.
* :meth:`await_refresh` MUST hold no lock across ``await self._refresh_callback()``.
  The refresh lock gates *task creation* only; the await on the task itself
  happens outside the lock so other waiters can join. Mixing this contract
  would silently deadlock waiters on a slow callback.
* :meth:`update_auth_tokens` writes ONLY ``auth.csrf_token`` and
  ``auth.session_id`` under the snapshot lock. It does NOT touch the
  http client. The cookie-jar sync is a separate concern handled by
  :meth:`update_auth_headers` (sync, no await — it runs the
  ``kernel.get_http_client().cookies`` read outside any auth lock).
* The ``_refresh_task`` slot is intentionally NOT cleared when a waiter is
  cancelled mid-shield — concurrency tests assert task identity across
  cancellation so siblings joined to the same single-flight refresh see the
  same completion.

Collaborator surface:

The coordinator depends on explicit per-method collaborators:
:meth:`snapshot` and :meth:`update_auth_tokens` take an
``auth: AuthTokens`` kwarg, :meth:`update_auth_headers` takes
``auth: AuthTokens`` plus ``kernel: Kernel``, and the lock-wait metric is
recorded through ``self._metrics`` (already supplied at construction).
This keeps every dependency on the coordinator's surface concrete and
narrow.
"""

from __future__ import annotations

import asyncio
import logging
import time
from collections.abc import Awaitable, Callable, Coroutine
from typing import TYPE_CHECKING, Any, cast

from .._loop_affinity import assert_bound_loop
from .._loop_bound import LoopBoundPrimitive
from .._request_types import AuthSnapshot
from ..auth import AuthTokens
from .config import CORE_LOGGER_NAME

if TYPE_CHECKING:
    from .._client_metrics import ClientMetrics
    from .._kernel import Kernel

# Logger name pinned via :data:`CORE_LOGGER_NAME` so log filters in
# tests — e.g. ``caplog.at_level("DEBUG", logger=CORE_LOGGER_NAME)`` —
# keep matching after the extraction.
logger = logging.getLogger(CORE_LOGGER_NAME)


class AuthRefreshCoordinator(LoopBoundPrimitive):
    """Owns refresh single-flight, snapshot serialization, and auth-header sync.

    Field names (``_refresh_lock``, ``_refresh_task``, ``_refresh_callback``,
    ``_auth_snapshot_lock``) are kept stable so coordinator-specific
    tests remain easy to audit.
    """

    def __init__(
        self,
        *,
        refresh_callback: Callable[[], Awaitable[AuthTokens]] | None = None,
        metrics: ClientMetrics | None = None,
    ) -> None:
        # Lazily-created — ``asyncio.Lock()`` needs a running loop in some
        # Python versions, and this object can be constructed outside one.
        self._refresh_lock: asyncio.Lock | None = None
        self._refresh_task: asyncio.Task[AuthTokens] | None = None
        self._refresh_callback: Callable[[], Awaitable[AuthTokens]] | None = refresh_callback
        # ``await_refresh`` records lock-wait latency via this metrics dep.
        # The same ``self._metrics`` slot is read by :meth:`snapshot` and
        # :meth:`update_auth_tokens` too. ``None`` is a safe
        # fallback for tests that construct the coordinator standalone
        # without a metrics collaborator; the lock-wait latency is simply not
        # recorded in that case.
        self._metrics: ClientMetrics | None = metrics
        # Distinct from ``_refresh_lock`` — see module docstring.
        self._auth_snapshot_lock: asyncio.Lock | None = None
        # ``_bound_loop`` (the loop-affinity guard consulted by
        # :meth:`await_refresh` before touching the lazy ``_refresh_lock``)
        # and ``set_bound_loop`` are provided by the
        # :class:`~notebooklm._loop_bound.LoopBoundPrimitive` base. This
        # coordinator only stores the binding, so it uses the default no-op
        # ``_on_loop_rebind`` (the lazy locks are never held across
        # ``open()`` and are rebuilt implicitly per ``open()``).

    @property
    def has_refresh_callback(self) -> bool:
        """``True`` iff a refresh callback was wired at construction.

        Used by :class:`notebooklm._middleware.auth_refresh.AuthRefreshMiddleware`
        to gate the refresh-and-retry branch: a client constructed without
        a ``refresh_callback`` should propagate auth errors directly.
        Exposing this as a property avoids reaching into the private
        ``_refresh_callback`` attribute from outside the coordinator.
        """
        return self._refresh_callback is not None

    # ------------------------------------------------------------------
    # Lazy lock accessors. Both follow the same race-free check-then-assign
    # pattern as ``_reqid_lock``: asyncio is single-threaded, so no other
    # coroutine can execute between the ``is None`` check and the
    # assignment unless we ``await`` — and we don't.
    # ------------------------------------------------------------------

    def get_refresh_lock(self) -> asyncio.Lock:
        """Return the lazily-initialised refresh lock.

        Concurrent callers resolve to the *same* instance because allocation
        is synchronous and asyncio is single-threaded; this preserves the
        single-flight refresh-task creation invariant in :meth:`await_refresh`.
        """
        if self._refresh_lock is None:
            self._refresh_lock = asyncio.Lock()
        return self._refresh_lock

    def get_auth_snapshot_lock(self) -> asyncio.Lock:
        """Return the lazily-initialised auth-snapshot lock.

        Held only across the four scalar reads in :meth:`snapshot` and the
        two scalar writes in :meth:`update_auth_tokens` — never across an
        ``await`` — so RPC throughput is not serialized to refresh latency.
        """
        if self._auth_snapshot_lock is None:
            self._auth_snapshot_lock = asyncio.Lock()
        return self._auth_snapshot_lock

    # ------------------------------------------------------------------
    # Auth snapshot + token write — the load-bearing AST-guarded pair.
    # These two methods are the canonical implementations of the
    # concurrency invariants: ``snapshot`` holds the ``_auth_snapshot_lock``
    # across four synchronous scalar reads, and ``update_auth_tokens``
    # forbids any ``await`` inside the csrf/session_id mutation try-block.
    # The AST guards in ``tests/unit/test_concurrency_refresh_race.py``
    # (``test_snapshot_acquires_auth_snapshot_lock`` and
    # ``test_update_auth_tokens_has_no_await_inside_mutation_block``)
    # inspect THIS module's source via ``inspect.getsource(...)`` + AST
    # parsing — any structural change to either method body (e.g.
    # extracting a helper, refactoring the lock dance, adding an
    # ``await`` mid-mutation) will trip those guards. Previously a facade
    # method mirrored each body; now ``snapshot`` and ``update_auth_tokens``
    # here are the only implementations, so there is no second body to keep
    # in sync.
    # ------------------------------------------------------------------

    async def snapshot(self, *, auth: AuthTokens) -> AuthSnapshot:
        """Capture the current auth scalars as a frozen snapshot.

        Acquires :attr:`_auth_snapshot_lock` for the four scalar reads so a
        concurrent :meth:`update_auth_tokens` cannot interleave between
        ``csrf_token`` / ``session_id`` / ``authuser`` / ``account_email``.
        The critical section is purely synchronous attribute reads — no
        ``await`` — so the lock is uncontested in steady state and refresh's
        tiny write block cannot block RPC throughput.

        The whole-attempt atomicity for ``(csrf, sid, cookies)`` on the wire
        is completed at the transport terminal:
        :meth:`RuntimeTransport.refresh_request_for_current_auth` captures a
        fresh snapshot, rebuilds the envelope, and
        :meth:`RuntimeTransport.terminal` calls ``Kernel.post`` with no await
        between materialization and the POST (see the AST guards in
        ``tests/unit/test_concurrency_refresh_race.py``). This lock guarantees
        the four scalars in the returned snapshot are coherent with each other;
        the terminal no-await rule keeps the cookie axis aligned with the
        materialized envelope.

        ``auth`` is passed explicitly per call; the lock-wait metric is
        recorded through ``self._metrics`` (supplied at construction).
        """
        wait_start = time.perf_counter()
        async with self.get_auth_snapshot_lock():
            if self._metrics is not None:
                self._metrics.record_lock_wait(time.perf_counter() - wait_start)
            return AuthSnapshot(
                csrf_token=auth.csrf_token,
                session_id=auth.session_id,
                authuser=auth.authuser,
                account_email=auth.account_email,
            )

    async def update_auth_tokens(
        self,
        *,
        auth: AuthTokens,
        csrf: str,
        session_id: str,
    ) -> None:
        """Atomically update ``auth.csrf_token`` + ``auth.session_id`` only.

        Does NOT touch the http client — the cookie-jar sync is the separate
        :meth:`update_auth_headers` concern. Conflating the two would let a
        snapshot acquired between this method and the header sync observe a
        new token pair against stale cookies, which is exactly the torn-state
        scenario the snapshot lock exists to prevent.

        ``auth`` is passed explicitly (no ``_AuthRefreshHost`` shape).
        """
        lock = self.get_auth_snapshot_lock()
        wait_start = time.perf_counter()
        await lock.acquire()
        try:
            # ``record_lock_wait`` lives INSIDE the ``try`` so a metric-side
            # exception (e.g. a misconfigured spy in tests, or a runtime bug
            # in :class:`ClientMetrics`) cannot leave the snapshot lock held
            # — the ``finally`` releases unconditionally. The call is
            # synchronous so the no-await guard pinned by
            # ``test_update_auth_tokens_has_no_await_inside_mutation_block``
            # still holds.
            if self._metrics is not None:
                self._metrics.record_lock_wait(time.perf_counter() - wait_start)
            auth.csrf_token = csrf
            auth.session_id = session_id
        finally:
            lock.release()

    def update_auth_headers(self, *, auth: AuthTokens, kernel: Kernel) -> None:
        """Sync ``auth.cookie_jar`` with the live HTTP client's jar.

        Synchronous on purpose — no await — so callers can run this without
        any auth lock held. The httpx client's cookie jar is authoritative
        once the session is open; re-injecting startup cookies here would
        overwrite cookies refreshed during redirects to
        ``accounts.google.com``.

        ``auth`` and ``kernel`` are passed explicitly per call so the
        coordinator does not need an owner-shaped host.

        Raises:
            RuntimeError: If the kernel's HTTP client is not initialised (the
                error originates from :meth:`Kernel.get_http_client`).
        """
        auth.cookie_jar = kernel.get_http_client().cookies

    # ------------------------------------------------------------------
    # Single-flight refresh task.
    # ------------------------------------------------------------------

    async def await_refresh(self) -> None:
        """Run / join the shared refresh task.

        Concurrent callers share one refresh task so a thundering herd of
        401s on the same client triggers exactly one token refresh. The lock
        protects task-creation only; the await on the task itself happens
        outside the lock so other callers can join.

        The join is wrapped in :func:`asyncio.shield` so that a caller
        cancelled while waiting — e.g. via ``asyncio.wait_for(..., timeout=...)``
        — unwinds locally without propagating the ``CancelledError`` into the
        *shared* refresh task. Without the shield, one cancelled waiter would
        cancel the underlying task, taking down every sibling joined to the
        same single-flight refresh. The slot at :attr:`_refresh_task` is left
        intact across the cancellation and is replaced only on the next
        refresh wave once the current task transitions to ``done()``.

        This method takes no host parameter — the metrics dependency it needs
        is supplied via the ``metrics`` kwarg on :meth:`__init__`. The other
        coordinator methods (``snapshot``, ``update_auth_tokens``,
        ``update_auth_headers``) likewise take explicit per-method
        collaborators rather than an owner facade.
        """
        # Catch cross-loop refresh before touching ``_refresh_lock``.
        # The lock is lazily bound to the loop that first awaited
        # ``get_refresh_lock`` — a cross-loop call would hang on the
        # ``await lock.acquire()`` if we let it through.
        assert_bound_loop(self._bound_loop)
        if self._refresh_callback is None:
            raise RuntimeError(
                "AuthRefreshCoordinator.await_refresh called without a "
                "refresh_callback configured — wire one via "
                "AuthRefreshCoordinator(refresh_callback=...) (or by "
                "constructing NotebookLMClient with refresh_callback=...) before "
                "triggering an auth refresh."
            )

        # Lazy-init the lock on first refresh attempt. Every concurrent
        # caller resolves to the same instance because ``get_refresh_lock``
        # runs synchronously in a single-threaded asyncio loop, so the
        # single-flight task creation below is preserved.
        lock = self.get_refresh_lock()
        wait_start = time.perf_counter()
        await lock.acquire()
        try:
            # ``record_lock_wait`` lives INSIDE the ``try`` so a metric-side
            # exception (e.g. a misconfigured spy in tests, or a runtime bug
            # in :class:`ClientMetrics`) cannot leave the refresh lock held —
            # the ``finally`` releases unconditionally. Mirrors the same
            # hardening on :meth:`update_auth_tokens`; the call is synchronous
            # so no await runs between acquiring and releasing the lock.
            if self._metrics is not None:
                self._metrics.record_lock_wait(time.perf_counter() - wait_start)
            if self._refresh_task is not None and not self._refresh_task.done():
                refresh_task = self._refresh_task
                logger.debug("Joining existing refresh task")
            else:
                coro = cast(Coroutine[Any, Any, AuthTokens], self._refresh_callback())
                self._refresh_task = asyncio.create_task(coro)
                refresh_task = self._refresh_task
        finally:
            lock.release()

        await asyncio.shield(refresh_task)

    async def cancel_inflight_refresh(self) -> None:
        """Cancel any in-flight refresh task during ``ClientLifecycle.close``.

        Mirrors the legacy close block previously inlined in
        :meth:`ClientLifecycle.close` so the lifecycle never touches the
        private ``_refresh_task`` slot on this coordinator:

        - **No-op** when ``_refresh_task is None`` — a freshly-opened client
          that never triggered an auth refresh has no task to cancel.
        - **No-op** when ``_refresh_task.done()`` — a refresh wave that
          already finished must not be re-cancelled (it would be harmless
          but ``gather(return_exceptions=True)`` would still log noise).
        - **Cancel** an unfinished task and ``await`` it via
          ``asyncio.gather(..., return_exceptions=True)`` so the resulting
          :class:`asyncio.CancelledError` is absorbed and ``close()`` itself
          stays non-raising in the normal racing case.

        Slot-preservation invariant (CRITICAL — load-bearing): the
        ``self._refresh_task`` slot is INTENTIONALLY left intact after a
        cancel. Sibling waiters joined to the same single-flight refresh
        (see :meth:`await_refresh` and the ``asyncio.shield`` it wraps
        around the join) read the slot to identify the shared task; clearing
        it here would break the concurrency invariant pinned by
        ``tests/unit/test_runtime_auth.py::test_await_refresh_cancellation_preserves_task_slot``.
        The slot is replaced only on the NEXT refresh wave once the current
        task transitions to ``done()`` — never here, never in close.

        Behavior is equivalent to:

            refresh_task = host._auth_coord._refresh_task
            if refresh_task is not None and not refresh_task.done():
                refresh_task.cancel()
                await asyncio.gather(refresh_task, return_exceptions=True)

        Regression coverage:
        ``tests/unit/concurrency/test_session_close_refresh_race.py`` and
        the three focused unit tests added with this method in
        ``tests/unit/test_runtime_auth.py`` (the two companion
        ``reset_after_open`` tests for :class:`TransportDrainTracker` live
        in ``tests/unit/test_runtime_lifecycle.py``).
        """
        refresh_task = self._refresh_task
        if refresh_task is not None and not refresh_task.done():
            refresh_task.cancel()
            await asyncio.gather(refresh_task, return_exceptions=True)


__all__ = ["AuthRefreshCoordinator"]