hermes-agent/tui_gateway/mcp_oauth_sessions.py

"""Session-backed MCP OAuth flows for the gateway (mcp.servers.oauth.*).

This mirrors the *provider* OAuth model used by the dashboard
(``/api/providers/oauth/{id}/start`` + ``/poll/{session_id}``) rather than the
FastAPI-request-coupled MCP dashboard flow: a ``start`` primitive kicks off a
background worker and returns ``{session_id, auth_url, flow}``; a ``poll``
primitive reports ``{status: pending|approved|error}`` until the tokens land on
disk for that server in that profile.

The underlying token machinery is the *same* one the CLI ``hermes mcp login``
uses — ``hermes_cli.mcp_config._probe_single_server`` under
``tools.mcp_oauth.force_interactive_oauth`` — so no OAuth logic is reimplemented
here. The only new piece is decoupling the two browser callbacks (authorization
URL out, ``code``/``state`` back in) from a FastAPI ``Request``:

* ``tools.mcp_dashboard_oauth.DashboardOAuthFlow`` already provides the two
  thread-safe rendezvous points (``publish_authorization_url`` /
  ``deliver_callback``). We reuse it verbatim as the bridge object.
* Instead of routing the browser redirect through a FastAPI callback route, we
  run a tiny loopback HTTP listener on ``127.0.0.1:<port>/callback`` and set the
  flow's ``redirect_uri`` to it. When the provider redirects the user's browser
  there, the listener calls ``flow.deliver_callback(...)``. This is the same
  loopback strategy the CLI uses by default, just wired to the shared bridge.

Client contract (what the desktop plugin does):
  1. call ``mcp.servers.oauth.start(profile, name)`` → ``{session_id, auth_url}``
  2. open ``auth_url`` in the native browser (``openExternal``)
  3. poll ``mcp.servers.oauth.poll(profile, name, session_id)`` until
     ``status == "approved"`` (tokens persisted) or ``"error"``.
"""

from __future__ import annotations

import http.server
import secrets
import threading
import time
from pathlib import Path
from typing import Any, Dict, Optional
from urllib.parse import parse_qs, urlparse

# Session registry: session_id -> record. A record wraps the shared
# DashboardOAuthFlow bridge plus a bit of gateway bookkeeping.
_sessions: Dict[str, Dict[str, Any]] = {}
_sessions_lock = threading.Lock()

# How long a completed/abandoned session lingers before GC (seconds).
_SESSION_TTL_SECONDS = 900
# Cap concurrent in-flight flows so a runaway client can't exhaust ports/threads.
_MAX_PENDING = 12


def _gc_sessions() -> None:
    """Drop expired sessions. Called opportunistically on start."""
    cutoff = time.time() - _SESSION_TTL_SECONDS
    with _sessions_lock:
        stale = [sid for sid, rec in _sessions.items() if rec["created_at"] < cutoff]
        for sid in stale:
            rec = _sessions.pop(sid, None)
            if rec is not None:
                _shutdown_listener(rec)


def _shutdown_listener(rec: Dict[str, Any]) -> None:
    server = rec.get("httpd")
    if server is not None:
        try:
            server.shutdown()
        except Exception:
            pass
        try:
            server.server_close()
        except Exception:
            pass
        rec["httpd"] = None


def _start_loopback_listener(flow) -> "http.server.HTTPServer":
    """Bind a loopback callback listener that feeds the flow's deliver_callback.

    Returns the running HTTPServer (already serving on a daemon thread). The
    bound port is read back off ``server.server_address`` so the caller can set
    ``flow.redirect_uri`` to the matching ``/callback`` URL BEFORE the worker
    starts the OAuth flow (the redirect URI must be pinned at authorization).
    """

    class _Handler(http.server.BaseHTTPRequestHandler):
        def do_GET(self):  # noqa: N802 — stdlib naming
            parsed = urlparse(self.path)
            if parsed.path.rstrip("/") not in ("/callback", ""):
                self.send_response(404)
                self.end_headers()
                return
            qs = parse_qs(parsed.query)
            code = (qs.get("code") or [None])[0]
            state = (qs.get("state") or [None])[0]
            error = (qs.get("error") or [None])[0]
            body = b"<h1>Authorization received</h1><p>You can close this tab and return to Hermes.</p>"
            status = 200
            try:
                flow.deliver_callback(code=code, state=state, error=error)
            except Exception:
                body = b"<h1>OAuth callback rejected</h1><p>The callback was invalid or already used.</p>"
                status = 400
            self.send_response(status)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.end_headers()
            try:
                self.wfile.write(body)
            except Exception:
                pass

        def log_message(self, *_a):  # silence stdlib request logging
            return

    httpd = http.server.HTTPServer(("127.0.0.1", 0), _Handler)
    threading.Thread(
        target=httpd.serve_forever,
        kwargs={"poll_interval": 0.5},
        daemon=True,
        name=f"mcp-oauth-cb-{flow.server_name}",
    ).start()
    return httpd


def _worker(session_id: str, hermes_home: str, server_name: str, cfg: dict, reconnect_live: bool) -> None:
    """Drive the interactive MCP OAuth probe under the shared dashboard bridge.

    Structurally identical to ``web_server._run_dashboard_mcp_oauth`` — the same
    HERMES_HOME override + secret-scope + force_interactive_oauth +
    dashboard_oauth_flow wrapping around ``_probe_single_server`` — but keyed to
    our session record instead of a FastAPI request. On success the token file
    exists on disk (verified via ``_oauth_tokens_present``) and the server config
    is (re)saved into the profile's config.yaml.
    """
    from hermes_cli.mcp_config import (
        _oauth_tokens_present,
        _probe_single_server,
        _save_mcp_server,
    )
    from hermes_constants import reset_hermes_home_override, set_hermes_home_override

    rec = _sessions.get(session_id)
    flow = rec["flow"] if rec else None
    try:
        from agent.secret_scope import (
            build_profile_secret_scope,
            reset_secret_scope,
            set_secret_scope,
        )
        from tools.mcp_dashboard_oauth import dashboard_oauth_flow
        from tools.mcp_oauth import force_interactive_oauth
        from tools.mcp_oauth_manager import get_manager

        home_token = set_hermes_home_override(hermes_home)
        secret_token = set_secret_scope(build_profile_secret_scope(Path(hermes_home)))
        try:
            with force_interactive_oauth(), dashboard_oauth_flow(flow):
                from tools.mcp_oauth import HermesTokenStorage

                manager = get_manager()
                storage = HermesTokenStorage(server_name)
                backup = storage.snapshot()
                previous_entry = None
                try:
                    previous_entry = manager.remove(server_name, hermes_home=hermes_home)
                    tools = _probe_single_server(
                        server_name,
                        cfg,
                        connect_timeout=max(float(cfg.get("connect_timeout", 0) or 0), 315),
                    )
                    if not _oauth_tokens_present(server_name):
                        raise RuntimeError(
                            "The server responded, but no OAuth token was obtained — "
                            "this provider may require a manually-registered OAuth client."
                        )
                    _save_mcp_server(server_name, cfg)
                    if flow is not None:
                        flow.tools = [{"name": t, "description": d} for t, d in tools]
                        flow.mark_approved()
                    if reconnect_live:
                        from tools.mcp_tool import reconnect_mcp_server

                        reconnect_mcp_server(server_name)
                except Exception:
                    storage.restore(backup, only_if_absent=True)
                    manager.restore_entry(server_name, previous_entry, hermes_home=hermes_home)
                    raise
        finally:
            reset_secret_scope(secret_token)
            reset_hermes_home_override(home_token)
    except Exception as exc:
        msg = str(exc)
        try:
            from tools.mcp_oauth import humanize_oauth_registration_error

            humanized = humanize_oauth_registration_error(
                server_name, exc, server_url=cfg.get("url") if isinstance(cfg, dict) else None
            )
            if humanized:
                msg = humanized
        except Exception:
            pass
        if flow is not None:
            flow.mark_error(msg)
    finally:
        if flow is not None:
            flow.mark_worker_done()
        if rec is not None:
            _shutdown_listener(rec)


def start_flow(
    hermes_home: str,
    server_name: str,
    cfg: dict,
    *,
    reconnect_live: bool = False,
    url_timeout: float = 30.0,
) -> Dict[str, Any]:
    """Begin an MCP OAuth flow and return ``{session_id, auth_url, flow}``.

    ``cfg`` is the server's resolved config dict (must have ``url`` and be
    OAuth-capable). ``hermes_home`` is the already-resolved profile home dir
    string. Blocks up to ``url_timeout`` for the worker to publish the browser
    authorization URL, then returns it.
    """
    from tools.mcp_dashboard_oauth import DashboardOAuthFlow

    _gc_sessions()

    with _sessions_lock:
        pending = sum(
            1
            for r in _sessions.values()
            if not r["flow"].worker_done
        )
        if pending >= _MAX_PENDING:
            raise RuntimeError("Too many MCP OAuth flows are already in progress")
        if any(
            r["server_name"] == server_name
            and r["hermes_home"] == hermes_home
            and not r["flow"].worker_done
            for r in _sessions.values()
        ):
            raise RuntimeError(f"MCP OAuth for '{server_name}' is already in progress")

    session_id = secrets.token_urlsafe(24)
    flow = DashboardOAuthFlow(
        flow_id=session_id,
        server_name=server_name,
        profile=None,
        hermes_home=hermes_home,
        redirect_uri="",  # set below once the loopback port is known
        reconnect_live=reconnect_live,
    )
    httpd = _start_loopback_listener(flow)
    port = httpd.server_address[1]
    flow.redirect_uri = f"http://127.0.0.1:{port}/callback"

    rec = {
        "session_id": session_id,
        "server_name": server_name,
        "hermes_home": hermes_home,
        "flow": flow,
        "httpd": httpd,
        "created_at": time.time(),
    }
    with _sessions_lock:
        _sessions[session_id] = rec

    threading.Thread(
        target=_worker,
        args=(session_id, hermes_home, server_name, dict(cfg), reconnect_live),
        daemon=True,
        name=f"mcp-oauth-{server_name}",
    ).start()

    try:
        auth_url = None
        # wait_for_authorization_url is async; run its wait synchronously.
        deadline = time.time() + url_timeout
        while time.time() < deadline:
            snap = flow.snapshot()
            if snap.get("authorization_url"):
                auth_url = snap["authorization_url"]
                break
            if snap.get("status") == "error":
                raise RuntimeError(snap.get("error") or "MCP OAuth flow failed before authorization")
            time.sleep(0.1)
        if not auth_url:
            raise TimeoutError("Timed out waiting for MCP authorization URL")
    except Exception:
        flow.mark_error("Timed out waiting for MCP authorization URL")
        _shutdown_listener(rec)
        raise

    return {
        "session_id": session_id,
        "auth_url": auth_url,
        # "pkce" mirrors the provider-OAuth ``flow`` discriminator: the client
        # opens a URL then polls (no user_code to type, unlike device_code).
        "flow": "pkce",
    }


def poll_flow(session_id: str, server_name: str) -> Dict[str, Any]:
    """Poll a session's status → ``{status, error_message?, auth_url?, tools?}``.

    ``status`` is one of ``pending`` | ``approved`` | ``error`` — the same
    vocabulary as the provider poll endpoint (``authorization_required`` from
    the underlying bridge maps to ``pending`` since the client only needs to
    know whether to keep waiting).
    """
    with _sessions_lock:
        rec = _sessions.get(session_id)
    if rec is None:
        return {"status": "error", "error_message": "OAuth session not found or expired"}
    if rec["server_name"] != server_name:
        return {"status": "error", "error_message": "server name mismatch for session"}

    flow = rec["flow"]
    snap = flow.snapshot()
    raw = snap.get("status")
    if raw == "approved":
        status = "approved"
    elif raw == "error":
        status = "error"
    else:
        status = "pending"
    out: Dict[str, Any] = {
        "session_id": session_id,
        "status": status,
        "error_message": snap.get("error"),
        "auth_url": snap.get("authorization_url"),
    }
    if status == "approved":
        out["tools"] = list(getattr(flow, "tools", []) or [])
    return out