Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license
This commit is contained in:
@@ -0,0 +1,429 @@
|
||||
"""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"``.
|
||||
|
||||
Remote-backend variant (client-side callback): when the desktop app runs on a
|
||||
DIFFERENT machine than the gateway (SSH/Tailscale remote backend), the
|
||||
gateway-side ``127.0.0.1`` listener is unreachable from the user's browser —
|
||||
the redirect lands on the user's machine where nothing is listening, and the
|
||||
flow times out. For that topology the client binds its OWN loopback listener
|
||||
(same pattern as the desktop's native gateway login), passes its
|
||||
``redirect_uri`` to ``start`` (``client_redirect_uri``), and relays the
|
||||
provider redirect back via ``deliver_callback_flow`` /
|
||||
``mcp.servers.oauth.callback``. State verification stays server-side in
|
||||
``DashboardOAuthFlow.deliver_callback`` — a relayed code with the wrong
|
||||
``state`` is rejected exactly like a forged loopback hit.
|
||||
"""
|
||||
|
||||
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 _validate_client_redirect_uri(uri: str) -> str:
|
||||
"""Validate a client-supplied loopback redirect URI.
|
||||
|
||||
Only plain-http loopback URLs are accepted (``http://127.0.0.1:<port>/...``
|
||||
or ``http://localhost:<port>/...``), mirroring RFC 8252 native-app rules —
|
||||
the client hosts a one-shot listener on ITS machine, so anything else
|
||||
(public hosts, https proxies, schemes) is rejected to keep the gateway from
|
||||
pinning an attacker-controlled redirect into a DCR registration.
|
||||
"""
|
||||
parsed = urlparse(str(uri or "").strip())
|
||||
host = (parsed.hostname or "").lower()
|
||||
if (
|
||||
parsed.scheme != "http"
|
||||
or host not in ("127.0.0.1", "localhost", "::1")
|
||||
or not parsed.port
|
||||
or parsed.username is not None
|
||||
or parsed.password is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"client_redirect_uri must be a loopback http URL like "
|
||||
"http://127.0.0.1:<port>/callback"
|
||||
)
|
||||
return f"http://{'[' + host + ']' if ':' in host else host}:{parsed.port}{parsed.path or '/callback'}"
|
||||
|
||||
|
||||
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,
|
||||
client_redirect_uri: Optional[str] = None,
|
||||
) -> 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.
|
||||
|
||||
``client_redirect_uri`` (remote-backend variant): a loopback callback URL
|
||||
the CLIENT hosts on its own machine. When set (and valid), no gateway-side
|
||||
listener is bound — the OAuth ``redirect_uri`` is pinned to the client's
|
||||
listener, and the client relays the redirect's ``code``/``state`` back via
|
||||
``deliver_callback_flow``. Invalid values raise ``ValueError``.
|
||||
"""
|
||||
from tools.mcp_dashboard_oauth import DashboardOAuthFlow
|
||||
|
||||
if client_redirect_uri is not None:
|
||||
client_redirect_uri = _validate_client_redirect_uri(client_redirect_uri)
|
||||
|
||||
_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,
|
||||
)
|
||||
if client_redirect_uri:
|
||||
# Remote-backend variant: the CLIENT hosts the callback listener on its
|
||||
# own machine and relays the code via deliver_callback_flow(). No
|
||||
# gateway-side listener is bound — a 127.0.0.1 port here would be
|
||||
# unreachable from the user's browser anyway.
|
||||
httpd = None
|
||||
flow.redirect_uri = client_redirect_uri
|
||||
else:
|
||||
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
|
||||
|
||||
|
||||
def deliver_callback_flow(
|
||||
session_id: str,
|
||||
server_name: str,
|
||||
*,
|
||||
code: Optional[str],
|
||||
state: Optional[str],
|
||||
error: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Relay a client-captured OAuth redirect into a session's flow.
|
||||
|
||||
Remote-backend companion to ``start_flow(client_redirect_uri=...)``: the
|
||||
desktop app's loopback listener caught the provider redirect on the USER'S
|
||||
machine and forwards ``code``/``state`` (or ``error``) here. Security
|
||||
properties are unchanged from the gateway-listener path — the underlying
|
||||
``DashboardOAuthFlow.deliver_callback`` verifies ``state`` against the
|
||||
pinned authorization request (constant-time compare) and rejects replays,
|
||||
so a forged or replayed relay fails identically to a forged loopback hit.
|
||||
|
||||
Returns ``{ok: true}`` on acceptance or ``{ok: false, error_message}``.
|
||||
"""
|
||||
with _sessions_lock:
|
||||
rec = _sessions.get(session_id)
|
||||
if rec is None:
|
||||
return {"ok": False, "error_message": "OAuth session not found or expired"}
|
||||
if rec["server_name"] != server_name:
|
||||
return {"ok": False, "error_message": "server name mismatch for session"}
|
||||
|
||||
flow = rec["flow"]
|
||||
try:
|
||||
flow.deliver_callback(code=code, state=state, error=error)
|
||||
except ValueError as exc:
|
||||
return {"ok": False, "error_message": str(exc)}
|
||||
return {"ok": True, "session_id": session_id}
|
||||
Reference in New Issue
Block a user