"""Dashboard-mediated callback bridge for MCP OAuth: the SDK still does discovery, DCR, PKCE, state validation and token exchange; this only moves the two human/browser callbacks from a loopback listener into the already-authenticated dashboard session.""" from __future__ import annotations import asyncio import contextvars import logging import secrets import threading import time from contextlib import contextmanager from dataclasses import dataclass, field from typing import Iterator from urllib.parse import parse_qs, urlparse logger = logging.getLogger(__name__) @contextmanager def contextvar_set(var: contextvars.ContextVar, value) -> Iterator[None]: """Set *var* to *value* for the block, restoring the previous value after.""" token = var.set(value) try: yield finally: var.reset(token) def _event_field(): return field(default_factory=threading.Event, init=False, repr=False) def exception_message(exc: BaseException) -> str: """``str(exc)``, or the type name when it is empty (bare ``TimeoutError()``/``RuntimeError()``) so ``mark_error`` never records a blank cause.""" return str(exc) or type(exc).__name__ @dataclass class DashboardOAuthFlow: flow_id: str server_name: str profile: str | None hermes_home: str redirect_uri: str reconnect_live: bool = False created_at: float = field(default_factory=time.time) status: str = "starting" authorization_url: str | None = None error: str | None = None tools: list[dict] = field(default_factory=list) discovery_error: str = "" # The user abandoned this flow: terminal for good, never re-minted (see publish_authorization_url). cancelled: bool = field(default=False, init=False) expected_state: str | None = field(default=None, init=False) _callback: tuple[str, str | None, str | None] | None = field(default=None, init=False, repr=False) _callback_error: str | None = field(default=None, init=False, repr=False) _authorization_ready: threading.Event = _event_field() _callback_ready: threading.Event = _event_field() _worker_done: threading.Event = _event_field() _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False) async def publish_authorization_url(self, url: str) -> None: """Record the SDK's authorization URL (with its ``state``) for the dashboard to show. A flow whose EARLIER attempt already ended is re-minted here instead of raising: the MCP server task inherits this handle for the whole life of the task, so ``RuntimeError: OAuth flow already ended`` escaped the SDK's auth flow on every retry and parked the server ("failed initial connection after 3 attempts") with no way back short of restarting Hermes (#114739). A flow the user CANCELLED stays terminal — the retrying worker must not reopen what they abandoned. """ state = parse_qs(urlparse(url).query).get("state", [None])[0] if not state: raise ValueError("OAuth authorization URL did not include state") with self._lock: if self.cancelled: raise RuntimeError("OAuth flow already ended: cancelled by user") if self.status in {"approved", "error"}: self._reopen_ended_attempt() self.expected_state = state self.authorization_url = url self.status = "authorization_required" self._authorization_ready.set() def _reopen_ended_attempt(self) -> None: """Reset an ended attempt so the next authorization URL mints a fresh state (caller holds ``self._lock``). The ended attempt's URL, state and spent authorization code are dropped — replaying them would hand the SDK a code the provider already burned.""" logger.info("MCP OAuth: dashboard flow %s for '%s' ended (%s); re-minting for the next attempt", self.flow_id, self.server_name, self.status) self.authorization_url = None self.expected_state = None self._callback = None self._callback_error = None self._callback_ready.clear() self.error = None self.status = "starting" async def wait_for_authorization_url(self, timeout: float = 30.0) -> str: if not await asyncio.to_thread(self._authorization_ready.wait, timeout): raise TimeoutError("Timed out waiting for MCP authorization URL") if not self.authorization_url: raise RuntimeError(self.error or "MCP OAuth flow ended before authorization") return self.authorization_url def deliver_callback( self, *, code: str | None, state: str | None, error: str | None, iss: str | None = None ) -> None: """Hand the browser redirect to the waiting flow; ``state`` must match exactly. ``iss`` (RFC 9207) is carried through — see ``tools.mcp_oauth._parse_redirect_query``. """ with self._lock: if self._callback_ready.is_set(): raise ValueError("OAuth callback already received") if self.expected_state is None or state is None or not secrets.compare_digest(self.expected_state, state): raise ValueError("OAuth callback state mismatch") if error: self._callback_error = error elif code: self._callback = (code, state, iss) else: self._callback_error = "OAuth callback did not include code or error" self._callback_ready.set() async def wait_for_callback(self, timeout: float = 300.0) -> tuple[str, str | None, str | None]: if not await asyncio.to_thread(self._callback_ready.wait, timeout): raise TimeoutError("Timed out waiting for MCP OAuth callback") if self._callback_error: raise RuntimeError(f"OAuth authorization failed: {self._callback_error}") if self._callback is None: raise RuntimeError( f"MCP OAuth flow for '{self.server_name}' ended without an authorization code " f"(status={self.status}, error={self.error!r})") return self._callback def mark_approved(self) -> None: with self._lock: if self.status == "error": raise RuntimeError("OAuth flow already ended") self.status = "approved" self.error = None def mark_error(self, error: str, *, cancelled: bool = False) -> None: """Fail the flow with *error*; the first reason wins. Waking the callback waiter makes the worker fail too, and its follow-on ``mark_error`` must not clobber the cause the user needs. ``cancelled`` marks a user cancellation, which ``publish_authorization_url`` refuses to re-mint (unlike an ordinary failure, which a retry may reopen). It is recorded even when the flow has already ended: the user's "no" outranks the recorded cause and keeps the handle terminal. """ with self._lock: self.cancelled = self.cancelled or cancelled if self.status in {"approved", "error"}: return self.status = "error" self.error = error or "MCP OAuth flow failed before the callback was received (empty error message)" # The SDK's callback waiter reads _callback_error, not error — without # this copy, a failure marked before any browser redirect (worker # crash, authorization-URL timeout, user cancel) wakes the waiter # with no callback and no error, surfacing as the generic # "did not include an authorization code" RuntimeError while the # real cause stays unread. Guarded so a late mark_error cannot # override an already delivered callback. if self._callback is None and self._callback_error is None: self._callback_error = self.error self._authorization_ready.set() self._callback_ready.set() def snapshot(self) -> dict: with self._lock: return {"flow_id": self.flow_id, "server_name": self.server_name, "status": self.status, "authorization_url": self.authorization_url, "error": self.error} def mark_worker_done(self) -> None: self._worker_done.set() @property def worker_done(self) -> bool: return self._worker_done.is_set() _current_dashboard_flow: contextvars.ContextVar[DashboardOAuthFlow | None] = contextvars.ContextVar("mcp_dashboard_oauth_flow", default=None) def dashboard_oauth_flow(flow: DashboardOAuthFlow): """Make *flow* the active dashboard OAuth flow for the block (ContextVar-scoped).""" return contextvar_set(_current_dashboard_flow, flow) def get_dashboard_oauth_flow() -> DashboardOAuthFlow | None: return _current_dashboard_flow.get()