From 237d409e599968dc2c2a7214594c6b39edc4976d Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Thu, 10 Sep 2026 00:03:22 +0000 Subject: [PATCH] Drive MCP connection login (RFC 8707 resource) from the mcp-proxy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When an AI Gateway MCP service is backed by a per-user connection (e.g. system.ai.github) with no stored credential, the gateway answers a tools call with an RFC 9728 challenge (HTTP 401 + WWW-Authenticate resource_metadata). The proxy currently forwards that to the coding agent as a plain "tools fetch failed" — the connection login never happens. Teach the proxy's httpx auth to detect that challenge and drive an interactive resource-scoped OAuth login: authorization-code + PKCE against /oidc using the published databricks-cli app, with the MCP endpoint URL sent as the RFC 8707 resource indicator so /oidc routes through the connection's /mcp-service-login page. After login the connection credential exists, so the same request is retried and succeeds. - resource_login.py: standalone stdlib OAuth-with-resource flow + loopback callback. - mcp_proxy.py: async_auth_flow detects the challenge, runs the login off the event loop, and retries once. Skipped under DATABRICKS_BEARER (headless/CI). - Tests for authorize-URL/PKCE construction and challenge detection. Co-authored-by: Isaac --- src/ucode/mcp_proxy.py | 83 ++++++++++-- src/ucode/resource_login.py | 254 +++++++++++++++++++++++++++++++++++ tests/test_mcp_proxy.py | 2 +- tests/test_resource_login.py | 89 ++++++++++++ 4 files changed, 419 insertions(+), 9 deletions(-) create mode 100644 src/ucode/resource_login.py create mode 100644 tests/test_resource_login.py diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 0d2a68cb..0089ee5d 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -33,16 +33,20 @@ from __future__ import annotations +import functools +import os import sys from collections.abc import AsyncIterator from types import ModuleType, TracebackType from typing import Protocol, Self import anyio +from anyio import to_thread from mcp.client.streamable_http import streamable_http_client from mcp.server.stdio import stdio_server from ucode.databricks import ensure_pat_bearer, get_databricks_token +from ucode.resource_login import ConnectionLoginError, login_for_connection # Exit code used when the proxy cannot continue. MCP clients surface a non-zero # exit far more usefully than a timeout, so bail out instead of hanging. @@ -111,29 +115,90 @@ def _fail_fast(message: str) -> None: raise SystemExit(AUTH_FAILURE_EXIT_CODE) -def _build_token_auth(workspace: str, profile: str | None): - """Build an httpx ``Auth`` that injects a fresh bearer on every request. +def _is_connection_login_challenge(response) -> bool: + """True for the AI Gateway's RFC 9728 connection-login challenge. + + The gateway answers a tools call against a connection with no per-user + credential with HTTP 401 + ``WWW-Authenticate: Bearer resource_metadata="…"``. + That is distinct from a dead *workspace* token (which fails earlier, when the + token is minted): here the workspace token is valid, but the connection needs + its own login.""" + if response.status_code != 401: + return False + return "resource_metadata" in response.headers.get("www-authenticate", "").lower() + + +def _build_token_auth(workspace: str, profile: str | None, resource: str | None = None): + """Build an httpx ``Auth`` that injects a fresh bearer on every request and, + when the gateway asks for a connection login, drives it and retries. The base class comes from whichever httpx the SDK uses (see ``_httpx``), so the returned auth is accepted by that SDK's ``AsyncClient``. Behaviour is - identical across flavours — ``Auth.auth_flow`` has the same generator - contract in httpx and httpx2.""" + identical across flavours — the ``auth_flow``/``async_auth_flow`` generator + contract is the same in httpx and httpx2. + + ``resource`` is the MCP endpoint URL being bridged; it is sent as the RFC 8707 + resource indicator on the connection-login OAuth so ``/oidc`` routes through + the connection's ``/mcp-service-login`` page.""" httpx = _httpx() class _DatabricksTokenAuth(httpx.Auth): - def auth_flow(self, request): + def __init__(self) -> None: + # One connection login per proxy lifetime — guards against a retry loop + # if login "succeeds" but the credential still isn't usable. + self._attempted_connection_login = False + + def _apply_token(self, request, *, force_refresh: bool = False) -> None: # get_databricks_token honors the DATABRICKS_BEARER short-circuit and # PAT profiles internally; --use-pat is surfaced via the env ucode set. # A RuntimeError here means auth is dead (expired refresh token, - # logged-out profile). Raising it from inside auth_flow would tear + # logged-out profile). Raising it from inside the flow would tear # through the transport's task group and stall the process until the # client times out, so translate it into a terminal ProxyAuthError the # caller reports cleanly. try: - token = get_databricks_token(workspace, profile) + # Only pass force_refresh when set — get_databricks_token defaults to + # False, so the bare call is equivalent and keeps the common path simple. + token = ( + get_databricks_token(workspace, profile, force_refresh=True) + if force_refresh + else get_databricks_token(workspace, profile) + ) except RuntimeError as exc: raise ProxyAuthError(str(exc)) from exc request.headers["Authorization"] = f"Bearer {token}" + + def auth_flow(self, request): + # Sync path: unused by the async transport, kept for the Auth contract. + self._apply_token(request) + yield request + + async def async_auth_flow(self, request): + self._apply_token(request) + response = yield request + if ( + resource is None + or self._attempted_connection_login + or not _is_connection_login_challenge(response) + # DATABRICKS_BEARER is the headless/CI path — no browser to drive an + # interactive login, so surface the challenge instead of hanging. + or os.environ.get("DATABRICKS_BEARER", "").strip() + ): + return + self._attempted_connection_login = True + await response.aread() + try: + # Blocking (browser + loopback server): run off the event loop. + await to_thread.run_sync( + functools.partial(login_for_connection, workspace, resource) + ) + except ConnectionLoginError as exc: + # Login didn't complete — leave the original 401 for the client to + # surface rather than crashing the proxy. + print(f"ucode mcp-proxy: {exc}", file=sys.stderr, flush=True) + return + # Connection credential now exists; retry the same request. + self._apply_token(request, force_refresh=True) yield request return _DatabricksTokenAuth() @@ -168,7 +233,9 @@ async def _pump_upstream[T]( async def _run(url: str, workspace: str, profile: str | None) -> None: httpx = _httpx() - auth = _build_token_auth(workspace, profile) + # `url` is the MCP endpoint we bridge to; pass it as the RFC 8707 resource so a + # connection-login challenge can be answered for this exact service. + auth = _build_token_auth(workspace, profile, resource=url) # 2.x-native shape: hand the transport a pre-built AsyncClient carrying our # per-request auth. Works on mcp 1.28+ and 2.x; `streamable_http_client` # yields a (read, write) pair in both. diff --git a/src/ucode/resource_login.py b/src/ucode/resource_login.py new file mode 100644 index 00000000..94112a87 --- /dev/null +++ b/src/ucode/resource_login.py @@ -0,0 +1,254 @@ +"""Interactive OAuth login for a specific MCP connection (RFC 8707 resource). + +When an AI Gateway MCP service is backed by a per-user connection (e.g. +``system.ai.github``) and the user hasn't logged in to that connection yet, the +gateway answers a tools call with an RFC 9728 challenge: HTTP 401 + +``WWW-Authenticate: Bearer resource_metadata="…"``. A plain Databricks workspace +token (what ``databricks auth token`` mints) is not enough — the *connection* +still needs a login. + +This module drives that login. It runs a standard OAuth authorization-code + PKCE +flow against the workspace ``/oidc`` server using the published ``databricks-cli`` +app, but adds the **RFC 8707 ``resource`` indicator** naming the MCP service. That +indicator is what makes ``/oidc`` route the browser through the connection's +``/mcp-service-login`` page: the user signs in to the backing SaaS (GitHub, …), +the connection credential is stored server-side, and the flow returns here with an +authorization code. After that the caller's normal Databricks token works, because +the connection credential now exists. + +stdlib only (no new dependency): ``http.server`` for the loopback callback, +``urllib`` for the token exchange, ``secrets``/``hashlib`` for PKCE. +""" + +from __future__ import annotations + +import base64 +import hashlib +import http.server +import json +import secrets +import sys +import threading +import webbrowser +from urllib import request as urllib_request +from urllib.parse import parse_qs, urlencode, urlparse + +# The published "Databricks CLI" OAuth app. It is a public (PKCE) client with +# loopback redirect URIs registered, so no client secret is needed. ``databricks`` +# itself authenticates with this same client id. +DATABRICKS_CLI_CLIENT_ID = "databricks-cli" +# One of the app's registered redirect URIs. Must match exactly what OIDC has on +# file for the client, so this is not freely configurable. +DEFAULT_CALLBACK_PORT = 8020 +# `all-apis` for the workspace token; `offline_access` for a refresh token. +DEFAULT_SCOPES = ("all-apis", "offline_access") +# How long to wait for the user to finish the browser login before giving up. +LOGIN_TIMEOUT_SECONDS = 300 + + +class ConnectionLoginError(RuntimeError): + """The interactive connection login could not be completed.""" + + +def _b64url(raw: bytes) -> str: + """base64url without padding, per RFC 7636.""" + return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii") + + +def _pkce_pair() -> tuple[str, str]: + """Return (code_verifier, code_challenge) for PKCE S256.""" + verifier = _b64url(secrets.token_bytes(64)) + challenge = _b64url(hashlib.sha256(verifier.encode("ascii")).digest()) + return verifier, challenge + + +def build_authorize_url( + workspace: str, + resource: str, + *, + client_id: str = DATABRICKS_CLI_CLIENT_ID, + callback_port: int = DEFAULT_CALLBACK_PORT, + scopes: tuple[str, ...] = DEFAULT_SCOPES, + code_challenge: str, + state: str, +) -> str: + """Build the ``/oidc/v1/authorize`` URL carrying the RFC 8707 resource. + + Split out from the flow so it can be unit-tested without any network or + browser. ``resource`` is the MCP endpoint URL the proxy is bridging to.""" + query = urlencode( + { + "response_type": "code", + "client_id": client_id, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + "redirect_uri": f"http://localhost:{callback_port}", + "scope": " ".join(scopes), + "state": state, + # RFC 8707: names the connection-backed MCP service so /oidc redirects + # through its /mcp-service-login page instead of issuing a bare token. + "resource": resource, + "prompt": "consent", + } + ) + return f"{workspace.rstrip('/')}/oidc/v1/authorize?{query}" + + +class _CallbackHandler(http.server.BaseHTTPRequestHandler): + """Captures the ``code``/``state`` the OIDC server redirects back with.""" + + # Set by the server instance before handling. + expected_state: str = "" + result: dict[str, str] = {} + + def do_GET(self) -> None: # noqa: N802 - BaseHTTPRequestHandler contract + params = parse_qs(urlparse(self.path).query) + code = params.get("code", [""])[0] + state = params.get("state", [""])[0] + server_result = type(self).result + if code and state == type(self).expected_state: + server_result["code"] = code + body = b"Login complete. You can close this tab." + else: + server_result["error"] = params.get("error", ["state_mismatch_or_no_code"])[0] + body = b"Login failed. You can close this tab." + self.send_response(200) + self.send_header("Content-Type", "text/html") + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 - match base signature + # Silence the default stderr access log — stderr is surfaced to the MCP client. + return + + +def _await_authorization_code(callback_port: int, state: str) -> str: + """Serve the loopback redirect until the auth code arrives (or timeout).""" + _CallbackHandler.expected_state = state + _CallbackHandler.result = {} + server = http.server.HTTPServer(("127.0.0.1", callback_port), _CallbackHandler) + server.timeout = LOGIN_TIMEOUT_SECONDS + + done = threading.Event() + + def _serve() -> None: + # One real redirect ends the flow; loop so a stray request (favicon, etc.) + # doesn't consume the single handle_request budget. + while not done.is_set(): + server.handle_request() + if _CallbackHandler.result: + done.set() + + thread = threading.Thread(target=_serve, daemon=True) + thread.start() + if not done.wait(timeout=LOGIN_TIMEOUT_SECONDS): + server.server_close() + raise ConnectionLoginError( + f"timed out after {LOGIN_TIMEOUT_SECONDS}s waiting for the connection login" + ) + server.server_close() + + result = _CallbackHandler.result + if "code" not in result: + raise ConnectionLoginError(f"connection login failed: {result.get('error', 'unknown')}") + return result["code"] + + +def _exchange_code( + workspace: str, + code: str, + *, + client_id: str, + callback_port: int, + code_verifier: str, + resource: str, +) -> str: + """Exchange the authorization code for an access token at ``/oidc/v1/token``.""" + body = urlencode( + { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": f"http://localhost:{callback_port}", + "client_id": client_id, + "code_verifier": code_verifier, + "resource": resource, + } + ).encode("ascii") + request = urllib_request.Request( + f"{workspace.rstrip('/')}/oidc/v1/token", + data=body, + method="POST", + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + try: + with urllib_request.urlopen(request, timeout=30) as response: + payload = json.loads(response.read().decode("utf-8")) + except Exception as exc: # noqa: BLE001 - surface any exchange failure uniformly + raise ConnectionLoginError(f"token exchange failed: {exc}") from exc + token = payload.get("access_token") + if not token: + raise ConnectionLoginError("token exchange returned no access_token") + return token + + +def login_for_connection( + workspace: str, + resource: str, + *, + client_id: str = DATABRICKS_CLI_CLIENT_ID, + callback_port: int = DEFAULT_CALLBACK_PORT, + scopes: tuple[str, ...] = DEFAULT_SCOPES, +) -> str: + """Run the resource-scoped OAuth login for an MCP connection. + + Opens the browser to the workspace ``/oidc`` authorize endpoint with the RFC + 8707 ``resource`` indicator, serves the loopback redirect, and exchanges the + returned code for a token. On success the connection credential is stored + server-side, so the caller's subsequent MCP requests succeed. + + Blocking (browser + local HTTP server); call it from a worker thread when on + an event loop. Returns the freshly minted access token. + """ + verifier, challenge = _pkce_pair() + state = secrets.token_urlsafe(24) + authorize_url = build_authorize_url( + workspace, + resource, + client_id=client_id, + callback_port=callback_port, + scopes=scopes, + code_challenge=challenge, + state=state, + ) + + # stdout is the MCP wire — every human-facing hint goes to stderr. Print the URL + # too: on a remote/SSH host the auto-open may not reach the user's browser. + print( + f"ucode mcp-proxy: connection login required — opening browser to sign in.\n" + f" If it doesn't open, visit:\n {authorize_url}", + file=sys.stderr, + flush=True, + ) + try: + webbrowser.open(authorize_url) + except Exception: # noqa: BLE001 - a headless open failure is non-fatal; URL was printed + pass + + code = _await_authorization_code(callback_port, state) + return _exchange_code( + workspace, + code, + client_id=client_id, + callback_port=callback_port, + code_verifier=verifier, + resource=resource, + ) + + +__all__ = [ + "ConnectionLoginError", + "DATABRICKS_CLI_CLIENT_ID", + "DEFAULT_CALLBACK_PORT", + "build_authorize_url", + "login_for_connection", +] diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index bf7a3536..544e400f 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -230,7 +230,7 @@ async def stop_bridge(*args, **kwargs): yield monkeypatch.setattr(httpx_module, "AsyncClient", CapturingClient) - monkeypatch.setattr(mcp_proxy, "_build_token_auth", lambda *args: object()) + monkeypatch.setattr(mcp_proxy, "_build_token_auth", lambda *args, **kwargs: object()) monkeypatch.setattr(mcp_proxy, "streamable_http_client", stop_bridge) with pytest.raises(StopBridge): diff --git a/tests/test_resource_login.py b/tests/test_resource_login.py new file mode 100644 index 00000000..75a56d02 --- /dev/null +++ b/tests/test_resource_login.py @@ -0,0 +1,89 @@ +"""Tests for resource-scoped MCP connection login (RFC 8707) and the proxy's +detection of the AI Gateway connection-login challenge. + +Network-free: the browser + loopback + token-exchange steps are not exercised +here; these cover the pure URL / PKCE construction and challenge classification. +""" + +from __future__ import annotations + +import base64 +import hashlib +from urllib.parse import parse_qs, urlparse + +from ucode import mcp_proxy, resource_login + +WS = "https://example.staging.cloud.databricks.com" +RESOURCE = f"{WS}/ai-gateway/mcp-services/system.ai.github" + + +def test_build_authorize_url_carries_resource_and_pkce(): + url = resource_login.build_authorize_url( + WS, + RESOURCE, + code_challenge="test-challenge", + state="test-state", + ) + parsed = urlparse(url) + assert parsed.path == "/oidc/v1/authorize" + q = parse_qs(parsed.query) + assert q["response_type"] == ["code"] + assert q["client_id"] == [resource_login.DATABRICKS_CLI_CLIENT_ID] + assert q["code_challenge"] == ["test-challenge"] + assert q["code_challenge_method"] == ["S256"] + assert q["redirect_uri"] == [f"http://localhost:{resource_login.DEFAULT_CALLBACK_PORT}"] + assert q["state"] == ["test-state"] + # The RFC 8707 resource indicator is what makes /oidc route through the + # connection's /mcp-service-login page. + assert q["resource"] == [RESOURCE] + assert "all-apis" in q["scope"][0] + assert "offline_access" in q["scope"][0] + + +def test_build_authorize_url_trims_trailing_slash(): + url = resource_login.build_authorize_url(WS + "/", RESOURCE, code_challenge="c", state="s") + assert url.startswith(f"{WS}/oidc/v1/authorize?") + + +def test_pkce_pair_is_valid_s256(): + verifier, challenge = resource_login._pkce_pair() + expected = ( + base64.urlsafe_b64encode(hashlib.sha256(verifier.encode("ascii")).digest()) + .rstrip(b"=") + .decode("ascii") + ) + assert challenge == expected + # base64url, no padding. + assert "=" not in verifier and "=" not in challenge + + +class _FakeResponse: + def __init__(self, status_code: int, headers: dict[str, str]) -> None: + self.status_code = status_code + self.headers = headers + + +def test_is_connection_login_challenge_true_on_401_with_resource_metadata(): + resp = _FakeResponse( + 401, + { + "www-authenticate": 'Bearer resource_metadata="https://ws/.well-known/oauth-protected-resource/x"' + }, + ) + assert mcp_proxy._is_connection_login_challenge(resp) is True + + +def test_is_connection_login_challenge_false_without_resource_metadata(): + # A generic 401 (e.g. workspace auth) is not a connection-login challenge. + assert mcp_proxy._is_connection_login_challenge(_FakeResponse(401, {})) is False + assert ( + mcp_proxy._is_connection_login_challenge( + _FakeResponse(401, {"www-authenticate": 'Bearer error="invalid_token"'}) + ) + is False + ) + + +def test_is_connection_login_challenge_false_on_non_401(): + resp = _FakeResponse(200, {"www-authenticate": 'Bearer resource_metadata="https://ws/x"'}) + assert mcp_proxy._is_connection_login_challenge(resp) is False