From 79ee16c9b4521a5fae98f5eae2e2e3b6e2975ac9 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Thu, 10 Sep 2026 23:44:23 +0000 Subject: [PATCH 1/4] mcp-proxy: drive per-connection login on a 401 (generic, all agents) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Give every coding agent a working login flow for connection-backed AI Gateway mcp-services endpoints, done in the stdio proxy every agent already spawns — no new library (cf. mcp-remote) and no per-agent OAuth app. When AI Gateway has no per-user connection credential it answers with HTTP 401 (RFC 9728). The proxy's httpx auth hook already sees every response, so on a 401 for a connection-backed URL it runs the Databricks CLI U2M login with an RFC 8707 resource indicator (`databricks auth login --resource `, using the CLI's own registered redirect — no --client-id), then retries with a fresh token. A resource-aware /oidc drives the connection's SaaS login before minting the token, so the retry succeeds — transparently to the agent, which just sees the request authenticate rather than a failed tools/list. A later credential revoke re-triggers the login on the next 401. New module mcp_connection_login holds connection_from_url + run_connection_login; mcp_proxy._build_token_auth gains the login-on-401 retry. Unit-tested. Depends on the CLI --resource flag (databricks/cli#6621) and /oidc resource handling (login). Co-authored-by: Isaac --- src/ucode/mcp_connection_login.py | 97 ++++++++++++++++++++++++++++++ src/ucode/mcp_proxy.py | 54 ++++++++++++----- tests/test_mcp_connection_login.py | 76 +++++++++++++++++++++++ tests/test_mcp_proxy.py | 91 ++++++++++++++++++++++++++-- 4 files changed, 296 insertions(+), 22 deletions(-) create mode 100644 src/ucode/mcp_connection_login.py create mode 100644 tests/test_mcp_connection_login.py diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py new file mode 100644 index 00000000..4c994034 --- /dev/null +++ b/src/ucode/mcp_connection_login.py @@ -0,0 +1,97 @@ +"""Per-connection login for AI Gateway MCP services, driven by the proxy on a 401. + +A connection-backed AI Gateway MCP service (e.g. ``system.ai.github``) needs a +per-user connection credential before its tools can be used. Until the user has +logged in to the underlying SaaS, AI Gateway answers requests with an HTTP 401 +(RFC 9728 ``WWW-Authenticate``). + +The ``ug mcp-proxy`` bridge (see ``mcp_proxy``) sees that 401 in its httpx auth +flow and, for a connection-backed service, runs the Databricks CLI U2M login +with an RFC 8707 ``resource`` indicator naming the service, then retries the +request. A resource-aware ``/oidc`` drives the connection's own SaaS login +before minting the token, so the credential exists on retry — transparently to +the coding agent, which just sees the connection authenticate and succeed. This +is the behaviour of a generic OAuth MCP bridge (e.g. ``mcp-remote``), done in +ucode with the Databricks CLI so no extra library or per-agent OAuth app is +needed. Requires the CLI ``--resource`` flag (databricks/cli#6621). +""" + +from __future__ import annotations + +import subprocess + +# AI Gateway MCP service endpoints look like +# ``https:///ai-gateway/mcp-services/..``. +AIGW_MCP_SERVICES_SEGMENT = "/ai-gateway/mcp-services/" + +# Login can pop a browser and wait for the user to complete the SaaS login, so +# allow generously more than a token refresh would take. +_LOGIN_TIMEOUT_SECONDS = 300 + + +def connection_from_url(url: str) -> str | None: + """Return the connection FQN of an AI Gateway MCP service URL, or ``None``. + + ``https://ws/ai-gateway/mcp-services/system.ai.github`` -> ``system.ai.github``. + A URL that is not an mcp-services endpoint (or names no service) returns + ``None`` — only connection-backed services get the login-on-401 treatment. + """ + marker = url.find(AIGW_MCP_SERVICES_SEGMENT) + if marker == -1: + return None + tail = url[marker + len(AIGW_MCP_SERVICES_SEGMENT) :] + # Strip any trailing path (``/tools/list``), query, or fragment. + connection = tail.split("/")[0].split("?")[0].split("#")[0] + return connection or None + + +def run_connection_login( + resource_url: str, + workspace: str, + *, + profile: str | None = None, + login_binary: str = "databricks", +) -> tuple[bool, str]: + """Run the CLI U2M login with an RFC 8707 resource indicator for this service. + + ``resource_url`` is the MCP service endpoint (also the proxy's upstream URL); + it is sent as ``--resource`` so a resource-aware ``/oidc`` drives the + connection's SaaS login before issuing the token. Uses the Databricks CLI's + own default client, whose loopback redirect is already registered — no + ``--client-id`` needed. Returns ``(ok, message)``; ``message`` is the CLI's + own output on failure so the caller can surface it. + """ + argv = [ + login_binary, + "auth", + "login", + "--host", + workspace.rstrip("/"), + "--resource", + resource_url, + ] + if profile: + argv += ["--profile", profile] + try: + result = subprocess.run( + argv, + check=False, + capture_output=True, + text=True, + timeout=_LOGIN_TIMEOUT_SECONDS, + ) + except OSError as exc: + return False, f"could not run '{login_binary} auth login': {exc}" + except subprocess.TimeoutExpired: + return False, "login timed out waiting for the browser flow to complete" + if result.returncode == 0: + return True, "signed in" + detail = (result.stderr or result.stdout or "").strip() + return False, detail or f"login exited with code {result.returncode}" + + +__all__ = [ + "AIGW_MCP_SERVICES_SEGMENT", + "connection_from_url", + "run_connection_login", +] diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 0d2a68cb..f352f76e 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -43,6 +43,7 @@ from mcp.server.stdio import stdio_server from ucode.databricks import ensure_pat_bearer, get_databricks_token +from ucode.mcp_connection_login import connection_from_url, run_connection_login # 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 +112,50 @@ 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 _build_token_auth(workspace: str, profile: str | None, url: str): + """Build an httpx ``Auth`` that injects a fresh bearer and logs in on a 401. 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.""" + contract in httpx and httpx2. + + For a connection-backed AI Gateway mcp-services endpoint, an HTTP 401 means + the per-user connection credential is missing. We drive the connection login + once (browser, via ``run_connection_login`` -> ``databricks auth login + --resource``) and retry with a fresh token, so the coding agent just sees the + request authenticate and succeed rather than a failed ``tools/list``.""" httpx = _httpx() + connection = connection_from_url(url) + + def _mint_bearer(request): + # get_databricks_token honors the DATABRICKS_BEARER short-circuit and PAT + # profiles internally; --use-pat is surfaced via the env ucode set. A + # RuntimeError means auth is dead (expired refresh token, logged-out + # profile). Raising from inside auth_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) + except RuntimeError as exc: + raise ProxyAuthError(str(exc)) from exc + request.headers["Authorization"] = f"Bearer {token}" class _DatabricksTokenAuth(httpx.Auth): def auth_flow(self, request): - # 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 - # 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) - except RuntimeError as exc: - raise ProxyAuthError(str(exc)) from exc - request.headers["Authorization"] = f"Bearer {token}" + _mint_bearer(request) + response = yield request + # Only connection-backed services have a per-user login to drive; a 401 + # from anything else is a real auth failure, left to surface as-is. + if connection is None or response.status_code != 401: + return + # Blocks this proxy while the browser login runs — acceptable, since the + # agent is only waiting on this one connect; the CLI runs its own + # callback listener out of process. + ok, detail = run_connection_login(url, workspace, profile=profile) + if not ok: + raise ProxyAuthError(f"connection login for '{connection}' failed: {detail}") + _mint_bearer(request) yield request return _DatabricksTokenAuth() @@ -168,7 +190,7 @@ async def _pump_upstream[T]( async def _run(url: str, workspace: str, profile: str | None) -> None: httpx = _httpx() - auth = _build_token_auth(workspace, profile) + auth = _build_token_auth(workspace, profile, 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/tests/test_mcp_connection_login.py b/tests/test_mcp_connection_login.py new file mode 100644 index 00000000..f1b3b798 --- /dev/null +++ b/tests/test_mcp_connection_login.py @@ -0,0 +1,76 @@ +"""Tests for the per-connection MCP login helpers (mcp_connection_login). + +Network-free: URL/connection parsing and the login runner with the subprocess +monkeypatched. +""" + +from __future__ import annotations + +import subprocess + +from ucode import mcp_connection_login as mcl + +WS = "https://ws.staging.cloud.databricks.com" +AIGW_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github" + + +class TestConnectionFromUrl: + def test_plain_endpoint(self): + assert mcl.connection_from_url(AIGW_URL) == "system.ai.github" + + def test_with_trailing_path_and_query(self): + assert mcl.connection_from_url(f"{AIGW_URL}/tools/list?x=1") == "system.ai.github" + + def test_non_aigw_url_is_none(self): + assert mcl.connection_from_url(f"{WS}/api/2.0/mcp/functions/system/ai") is None + + def test_missing_service_is_none(self): + assert mcl.connection_from_url(f"{WS}/ai-gateway/mcp-services/") is None + + +class TestRunConnectionLogin: + def _fake_run(self, captured, *, returncode, stderr=""): + def _run(argv, **kwargs): + captured.append(argv) + return subprocess.CompletedProcess(argv, returncode, stdout="", stderr=stderr) + + return _run + + def test_success_sends_resource_and_host_without_client_id(self, monkeypatch): + captured: list[list[str]] = [] + monkeypatch.setattr(mcl.subprocess, "run", self._fake_run(captured, returncode=0)) + + ok, message = mcl.run_connection_login(AIGW_URL, WS, profile="p") + + assert ok and message == "signed in" + argv = captured[0] + assert argv[:3] == ["databricks", "auth", "login"] + assert "--resource" in argv and AIGW_URL in argv + assert "--host" in argv and WS in argv + assert "--profile" in argv and "p" in argv + # Uses the CLI's default client (its own registered redirect), so no --client-id. + assert "--client-id" not in argv + + def test_failure_returns_cli_detail(self, monkeypatch): + captured: list[list[str]] = [] + monkeypatch.setattr( + mcl.subprocess, "run", self._fake_run(captured, returncode=1, stderr="nope") + ) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok and message == "nope" + + def test_timeout_is_reported(self, monkeypatch): + def _run(argv, **kwargs): + raise subprocess.TimeoutExpired(argv, 1) + + monkeypatch.setattr(mcl.subprocess, "run", _run) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok and "timed out" in message + + def test_binary_missing_is_reported(self, monkeypatch): + def _run(argv, **kwargs): + raise OSError("not found") + + monkeypatch.setattr(mcl.subprocess, "run", _run) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok and "could not run" in message diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index bf7a3536..d6c06c1c 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -61,7 +61,7 @@ def test_proxy_imports_the_streamable_http_client_shared_by_both_majors(): class TestDatabricksTokenAuth: def test_injects_bearer_from_minted_token(self, monkeypatch): monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok-123") - auth = mcp_proxy._build_token_auth(WS, "uc-dogfood") + auth = mcp_proxy._build_token_auth(WS, "uc-dogfood", URL) request = httpx.Request("POST", URL) # auth_flow is a generator that yields the (mutated) request. @@ -73,7 +73,7 @@ def test_auth_is_an_instance_of_the_selected_httpx_auth(self, monkeypatch): # The auth must subclass the *same* httpx flavor's Auth as the transport, # or the SDK's AsyncClient won't accept it. monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "t") - auth = mcp_proxy._build_token_auth(WS, None) + auth = mcp_proxy._build_token_auth(WS, None, URL) assert isinstance(auth, mcp_proxy._httpx().Auth) @@ -84,7 +84,7 @@ def test_calls_get_token_with_workspace_and_profile(self, monkeypatch): "get_databricks_token", lambda ws, profile: calls.append((ws, profile)) or "t", ) - auth = mcp_proxy._build_token_auth(WS, "myprofile") + auth = mcp_proxy._build_token_auth(WS, "myprofile", URL) list(auth.auth_flow(httpx.Request("POST", URL))) @@ -95,7 +95,7 @@ def test_mints_a_fresh_token_per_request(self, monkeypatch): # picked up mid-session without the proxy tracking expiry itself. tokens = iter(["first", "second"]) monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: next(tokens)) - auth = mcp_proxy._build_token_auth(WS, None) + auth = mcp_proxy._build_token_auth(WS, None, URL) r1 = httpx.Request("POST", URL) r2 = httpx.Request("POST", URL) @@ -107,7 +107,7 @@ def test_mints_a_fresh_token_per_request(self, monkeypatch): def test_auth_flow_yields_the_same_request(self, monkeypatch): monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "t") - auth = mcp_proxy._build_token_auth(WS, None) + auth = mcp_proxy._build_token_auth(WS, None, URL) request = httpx.Request("POST", URL) yielded = list(auth.auth_flow(request)) @@ -122,12 +122,91 @@ def boom(ws, profile): raise RuntimeError("no access token; run `databricks auth login`") monkeypatch.setattr(mcp_proxy, "get_databricks_token", boom) - auth = mcp_proxy._build_token_auth(WS, "p") + auth = mcp_proxy._build_token_auth(WS, "p", URL) with pytest.raises(mcp_proxy.ProxyAuthError, match="databricks auth login"): list(auth.auth_flow(httpx.Request("POST", URL))) +CONN_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github" + + +def _drive_auth_flow(auth, request, responses): + """Drive an httpx auth_flow generator, feeding ``responses`` back per yield. + + Returns the list of requests the flow yielded (one per attempt).""" + gen = auth.auth_flow(request) + yielded = [next(gen)] + for response in responses: + try: + yielded.append(gen.send(response)) + except StopIteration: + break + return yielded + + +def _response(status): + return httpx.Response(status, request=httpx.Request("POST", CONN_URL)) + + +class TestConnectionLoginOn401: + def test_no_401_does_not_trigger_login(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + logins: list = [] + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") + ) + auth = mcp_proxy._build_token_auth(WS, "p", CONN_URL) + + yielded = _drive_auth_flow(auth, httpx.Request("POST", CONN_URL), [_response(200)]) + + assert len(yielded) == 1 # no retry + assert logins == [] + + def test_401_drives_login_and_retries_with_fresh_token(self, monkeypatch): + tokens = iter(["stale", "fresh"]) + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: next(tokens)) + calls: list[tuple] = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda url, ws, **k: calls.append((url, ws, k.get("profile"))) or (True, "signed in"), + ) + auth = mcp_proxy._build_token_auth(WS, "p", CONN_URL) + + request = httpx.Request("POST", CONN_URL) + yielded = _drive_auth_flow(auth, request, [_response(401), _response(200)]) + + # Logged in for this connection's URL, then retried with the fresh token. + assert calls == [(CONN_URL, WS, "p")] + assert len(yielded) == 2 + assert yielded[1].headers["Authorization"] == "Bearer fresh" + + def test_login_failure_is_terminal(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: (False, "user cancelled") + ) + auth = mcp_proxy._build_token_auth(WS, "p", CONN_URL) + + with pytest.raises(mcp_proxy.ProxyAuthError, match="user cancelled"): + _drive_auth_flow(auth, httpx.Request("POST", CONN_URL), [_response(401)]) + + def test_non_connection_url_401_does_not_trigger_login(self, monkeypatch): + # A 401 from a non-mcp-services endpoint is a real auth failure, left as-is. + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + logins: list = [] + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") + ) + auth = mcp_proxy._build_token_auth(WS, "p", URL) # URL is not connection-backed + + yielded = _drive_auth_flow(auth, httpx.Request("POST", URL), [_response(401)]) + + assert len(yielded) == 1 # no retry + assert logins == [] + + class TestPump: def test_forwards_all_messages_in_order(self): async def scenario() -> list[str]: From 0cff035837736d5d0115d02d8fdaa0d93107a479 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 01:43:29 +0000 Subject: [PATCH 2/4] mcp-proxy: make the connection login non-blocking and surface its URL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes the two defects that made the proxy hang "connecting…" on a 401 instead of behaving like a generic OAuth MCP bridge (mcp-remote): 1. Non-blocking: the browser login ran via a synchronous subprocess inside the sync httpx auth_flow, which the async client executes on the event-loop thread — freezing the transport (stdio pumps included) for the whole login. Add async_auth_flow that offloads run_connection_login to a worker thread (anyio.to_thread.run_sync), so the loop stays responsive and cancellable while the user completes the browser flow. sync auth_flow kept for parity; both share the decision + login logic. 2. Visible URL: run_connection_login captured the CLI's output, hiding the authorize URL. Route the CLI's stdout+stderr to the proxy's stderr (fd 2, the agent's MCP log) — never fd 1 (the JSON-RPC wire) — and let the CLI open the browser, so the login is discoverable exactly like mcp-remote's. Co-authored-by: Isaac --- src/ucode/mcp_connection_login.py | 33 ++++++++++++++++++----- src/ucode/mcp_proxy.py | 42 +++++++++++++++++++++++------- tests/test_mcp_connection_login.py | 24 ++++++++++++++--- tests/test_mcp_proxy.py | 31 ++++++++++++++++++++++ 4 files changed, 109 insertions(+), 21 deletions(-) diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py index 4c994034..bb3c8557 100644 --- a/src/ucode/mcp_connection_login.py +++ b/src/ucode/mcp_connection_login.py @@ -19,6 +19,7 @@ from __future__ import annotations import subprocess +import sys # AI Gateway MCP service endpoints look like # ``https:///ai-gateway/mcp-services/..``. @@ -58,8 +59,13 @@ def run_connection_login( it is sent as ``--resource`` so a resource-aware ``/oidc`` drives the connection's SaaS login before issuing the token. Uses the Databricks CLI's own default client, whose loopback redirect is already registered — no - ``--client-id`` needed. Returns ``(ok, message)``; ``message`` is the CLI's - own output on failure so the caller can surface it. + ``--client-id`` needed. + + The CLI opens the browser to complete the login and prints the authorize URL. + We route its output to **stderr** (never stdout — that is the proxy's MCP + JSON-RPC wire), so a coding agent surfaces it in the server's log and the URL + stays visible when the browser can't open (e.g. a headless remote). Returns + ``(ok, message)``; on failure ``message`` points at that log. """ argv = [ login_binary, @@ -72,22 +78,35 @@ def run_connection_login( ] if profile: argv += ["--profile", profile] + connection = connection_from_url(resource_url) or resource_url + print( + f"ucode mcp-proxy: '{connection}' needs a one-time connection sign-in. Opening your " + "browser to complete it — if it doesn't open, use the authorization URL printed below.", + file=sys.stderr, + flush=True, + ) try: + # stdout -> stderr: the CLI's prompts and authorize URL reach the agent's + # MCP log (fd 2) without corrupting this process's stdout (fd 1, the MCP + # JSON-RPC stream). stdin is closed since the flow is browser-driven. result = subprocess.run( argv, check=False, - capture_output=True, - text=True, timeout=_LOGIN_TIMEOUT_SECONDS, + stdin=subprocess.DEVNULL, + stdout=sys.stderr, + stderr=sys.stderr, ) except OSError as exc: return False, f"could not run '{login_binary} auth login': {exc}" except subprocess.TimeoutExpired: - return False, "login timed out waiting for the browser flow to complete" + return False, "connection sign-in timed out waiting for the browser flow to complete" if result.returncode == 0: return True, "signed in" - detail = (result.stderr or result.stdout or "").strip() - return False, detail or f"login exited with code {result.returncode}" + return ( + False, + f"connection sign-in did not complete (CLI exited {result.returncode}; see the log above)", + ) __all__ = [ diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index f352f76e..5647e0b2 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -39,6 +39,7 @@ from typing import Protocol, Self import anyio +from anyio.to_thread import run_sync from mcp.client.streamable_http import streamable_http_client from mcp.server.stdio import stdio_server @@ -124,7 +125,12 @@ def _build_token_auth(workspace: str, profile: str | None, url: str): the per-user connection credential is missing. We drive the connection login once (browser, via ``run_connection_login`` -> ``databricks auth login --resource``) and retry with a fresh token, so the coding agent just sees the - request authenticate and succeed rather than a failed ``tools/list``.""" + request authenticate and succeed rather than a failed ``tools/list``. + + The proxy runs on an async event loop, so the login (a blocking subprocess + that waits on the browser) is offloaded to a worker thread in + ``async_auth_flow`` — the loop keeps servicing the stdio pumps and stays + cancellable while the user completes the browser flow, instead of freezing.""" httpx = _httpx() connection = connection_from_url(url) @@ -141,21 +147,37 @@ def _mint_bearer(request): raise ProxyAuthError(str(exc)) from exc request.headers["Authorization"] = f"Bearer {token}" + def _login_and_remint(request): + # Blocking: drives the browser login, then re-mints the now-valid bearer. + ok, detail = run_connection_login(url, workspace, profile=profile) + if not ok: + raise ProxyAuthError(f"connection login for '{connection}' failed: {detail}") + _mint_bearer(request) + + def _needs_login(response) -> bool: + # Only connection-backed services have a per-user login to drive; a 401 from + # anything else is a real auth failure, left to surface as-is. + return connection is not None and response.status_code == 401 + class _DatabricksTokenAuth(httpx.Auth): + # Async is the real path (the proxy uses an AsyncClient); the sync flow is + # kept for completeness/parity. Both share the same decision + login logic. def auth_flow(self, request): _mint_bearer(request) response = yield request - # Only connection-backed services have a per-user login to drive; a 401 - # from anything else is a real auth failure, left to surface as-is. - if connection is None or response.status_code != 401: + if not _needs_login(response): return - # Blocks this proxy while the browser login runs — acceptable, since the - # agent is only waiting on this one connect; the CLI runs its own - # callback listener out of process. - ok, detail = run_connection_login(url, workspace, profile=profile) - if not ok: - raise ProxyAuthError(f"connection login for '{connection}' failed: {detail}") + _login_and_remint(request) + yield request + + async def async_auth_flow(self, request): _mint_bearer(request) + response = yield request + if not _needs_login(response): + return + # Offload the blocking browser login so the event loop keeps running + # (mcp-remote-style: the transport stays responsive, not frozen). + await run_sync(lambda: _login_and_remint(request)) yield request return _DatabricksTokenAuth() diff --git a/tests/test_mcp_connection_login.py b/tests/test_mcp_connection_login.py index f1b3b798..4a9a9bfb 100644 --- a/tests/test_mcp_connection_login.py +++ b/tests/test_mcp_connection_login.py @@ -51,13 +51,29 @@ def test_success_sends_resource_and_host_without_client_id(self, monkeypatch): # Uses the CLI's default client (its own registered redirect), so no --client-id. assert "--client-id" not in argv - def test_failure_returns_cli_detail(self, monkeypatch): + def test_nonzero_exit_reports_failure(self, monkeypatch): + # The CLI's own output streams live to stderr (the agent's MCP log), so on + # failure we return a pointer to that log rather than captured text. captured: list[list[str]] = [] + monkeypatch.setattr(mcl.subprocess, "run", self._fake_run(captured, returncode=1)) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok + assert "did not complete" in message and "1" in message + + def test_output_is_routed_to_stderr_not_stdout(self, monkeypatch): + # stdout must never be captured to the proxy's stdout (the MCP wire); the + # CLI's URL/prompts go to this process's stderr. + seen: dict = {} monkeypatch.setattr( - mcl.subprocess, "run", self._fake_run(captured, returncode=1, stderr="nope") + mcl.subprocess, + "run", + lambda argv, **kw: seen.update(kw) or subprocess.CompletedProcess(argv, 0), ) - ok, message = mcl.run_connection_login(AIGW_URL, WS) - assert not ok and message == "nope" + ok, _ = mcl.run_connection_login(AIGW_URL, WS) + assert ok + assert seen.get("stdout") is mcl.sys.stderr + assert seen.get("stderr") is mcl.sys.stderr + assert "capture_output" not in seen def test_timeout_is_reported(self, monkeypatch): def _run(argv, **kwargs): diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index d6c06c1c..e2d810b8 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -206,6 +206,37 @@ def test_non_connection_url_401_does_not_trigger_login(self, monkeypatch): assert len(yielded) == 1 # no retry assert logins == [] + def test_async_flow_offloads_login_and_retries(self, monkeypatch): + # The real path is async (AsyncClient). async_auth_flow must offload the + # blocking login to a worker thread and retry with a fresh token. + tokens = iter(["stale", "fresh"]) + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: next(tokens)) + calls: list[tuple] = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda url, ws, **k: calls.append((url, k.get("profile"))) or (True, "signed in"), + ) + auth = mcp_proxy._build_token_auth(WS, "p", CONN_URL) + + async def scenario(): + # The flow mutates one request object in place, so capture the header + # value at each yield rather than comparing object references. + req = httpx.Request("POST", CONN_URL) + gen = auth.async_auth_flow(req) + await gen.__anext__() + first_auth = req.headers["Authorization"] + await gen.asend(_response(401)) + retry_auth = req.headers["Authorization"] + with pytest.raises(StopAsyncIteration): + await gen.asend(_response(200)) + return first_auth, retry_auth + + first_auth, retry_auth = anyio.run(scenario) + assert calls == [(CONN_URL, "p")] # login fired (offloaded), once + assert first_auth == "Bearer stale" + assert retry_auth == "Bearer fresh" + class TestPump: def test_forwards_all_messages_in_order(self): From 81cf9d70fb6354cc4962984182e0c6321bdb966a Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 02:23:34 +0000 Subject: [PATCH 3/4] mcp-proxy: drive the connection login at connect-time (mcp-remote style) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Makes the generic proxy behave like a generic OAuth MCP bridge (mcp-remote): the connection login happens while the agent shows "connecting…", and the browser opens on its own — instead of racing the agent's tools/list timeout or burying the URL. - Connect-time login: before opening the bridge, serve() probes the connection (_connection_login_required: a lightweight initialize + tools/list to AI Gateway); on a 401 it drives run_connection_login *then*, so the agent's session comes up already authenticated. The on-401 retry in the auth hook stays as a mid-session fallback (credential revoked while connected). PAT profiles skip it (no connection OAuth). - Browser auto-open: the login inherits the environment (incl. $BROWSER), so databricks-cli opens the browser on the user's machine; the authorize URL is the printed fallback. Co-authored-by: Isaac --- src/ucode/mcp_proxy.py | 65 +++++++++++++++++++++++++++++++++++++++++ tests/test_mcp_proxy.py | 56 +++++++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+) diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 5647e0b2..4c2c1ab4 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -249,6 +249,57 @@ def _preflight_token(workspace: str, profile: str | None) -> None: get_databricks_token(workspace, profile) +# Timeout for the startup probe MCP round-trips (initialize + tools/list). Short: +# it's a liveness check, not the login (which has its own generous timeout). +_PROBE_TIMEOUT_SECONDS = 15 + + +def _connection_login_required(url: str, workspace: str, profile: str | None) -> bool: + """Whether the connection-backed MCP service answers ``tools/list`` with a 401. + + A lightweight probe run at startup (before the bridge) so the connection login + happens during the agent's "connecting…" phase — like a generic OAuth MCP + bridge — instead of on the agent's first ``tools/list``, where it would race + the agent's tool-fetch timeout. Any non-401 outcome (authenticated, or a + network/transport hiccup) returns ``False`` so startup is never blocked on a + false alarm; a genuine missing credential surfaces again on the live request.""" + httpx = _httpx() + try: + token = get_databricks_token(workspace, profile) + except RuntimeError: + return False # dead databricks auth; let _preflight_token/the bridge report it + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + initialize = { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "ucode-mcp-proxy", "version": "0"}, + }, + } + try: + with httpx.Client(timeout=_PROBE_TIMEOUT_SECONDS) as client: + init_response = client.post(url, headers=headers, json=initialize) + session_id = init_response.headers.get("mcp-session-id") + if session_id: + headers["mcp-session-id"] = session_id + client.post( + url, headers=headers, json={"jsonrpc": "2.0", "method": "notifications/initialized"} + ) + tools = client.post( + url, headers=headers, json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"} + ) + return tools.status_code == 401 + except httpx.HTTPError: + return False + + def _unwrap_proxy_error(exc: BaseException) -> ProxyAuthError | ProxyTransportError | None: """Find a known proxy error in an exception (or ExceptionGroup) tree. @@ -289,6 +340,20 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool except RuntimeError as exc: _fail_fast(str(exc)) + # Connect-time connection login (generic OAuth-bridge behaviour): before the + # bridge starts, if a connection-backed service is unauthenticated, drive the + # login now — while the agent shows "connecting…" — so the session comes up + # already connected instead of failing the agent's first tools/list. The + # browser opens (databricks-cli honours $BROWSER, inherited here) or the + # authorize URL is printed to this stderr. PAT profiles have no connection + # OAuth to drive. The on-401 retry in the auth hook remains as a mid-session + # fallback (e.g. the credential is revoked while connected). + connection = None if use_pat else connection_from_url(url) + if connection is not None and _connection_login_required(url, workspace, profile): + ok, detail = run_connection_login(url, workspace, profile=profile) + if not ok: + _fail_fast(f"connection login for '{connection}' failed: {detail}") + try: anyio.run(_run, url, workspace, profile) except BaseException as exc: # noqa: BLE001 - re-raised unless it's a known proxy failure diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index e2d810b8..1effe359 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -406,6 +406,62 @@ def test_use_pat_without_a_resolvable_pat_exits_before_serving(self, monkeypatch assert started == [] # never opened the bridge assert "no personal access token" in capsys.readouterr().err + def test_connect_time_login_runs_before_the_bridge_when_required(self, monkeypatch): + order: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr(mcp_proxy, "_connection_login_required", lambda url, ws, profile: True) + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda url, ws, **k: order.append(("login", url)) or (True, "signed in"), + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: order.append(("bridge",))) + + mcp_proxy.serve(CONN_URL, WS, "p") + + # Login (during "connecting…") happens before the bridge opens. + assert order == [("login", CONN_URL), ("bridge",)] + + def test_authenticated_connection_skips_connect_time_login(self, monkeypatch): + logins: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr(mcp_proxy, "_connection_login_required", lambda url, ws, profile: False) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) + + mcp_proxy.serve(CONN_URL, WS, "p") + + assert logins == [] # already authenticated -> no login + + def test_connect_time_login_failure_is_terminal(self, monkeypatch): + started: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr(mcp_proxy, "_connection_login_required", lambda url, ws, profile: True) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: (False, "user cancelled") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: started.append("bridge")) + + with pytest.raises(SystemExit): + mcp_proxy.serve(CONN_URL, WS, "p") + + assert started == [] # never opened the bridge + + def test_use_pat_skips_the_connect_time_probe(self, monkeypatch): + probed: list = [] + monkeypatch.setattr(mcp_proxy, "ensure_pat_bearer", lambda profile: True) + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "_connection_login_required", lambda *a: probed.append(1) or True + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) + + mcp_proxy.serve(CONN_URL, WS, "p", use_pat=True) + + assert probed == [] # PAT has no connection OAuth to probe/drive + def test_oauth_path_never_touches_pat(self, monkeypatch): # Without use_pat, ensure_pat_bearer must not be consulted at all. called: list[str] = [] From 35c1a1dcd27b27a98d476ddbbb490f2990d07f7f Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 03:06:05 +0000 Subject: [PATCH 4/4] mcp-proxy: connection login at connect (the mcp-remote pattern), drop the probe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Simplify to what a generic OAuth MCP bridge (mcp-remote) does: authenticate at connect, then serve. Before opening the bridge, serve() runs a blocking `databricks auth login --resource ` for connection-backed services — the databricks-cli equivalent of mcp-remote's in-process OAuth, where --resource also routes /oidc through the connection sign-in (/mcp-service-login). The agent blocks on "connecting…" while it runs (browser opens via $BROWSER, or the URL is printed), then the session comes up authenticated, so AI Gateway is never asked to elicit a login. Idempotent: once signed in it returns immediately. Removes the redundant startup probe (Claude already fires initialize+tools/list; the proxy shouldn't duplicate that) and the on-401 retry inside the auth hook (the connect-time login makes it unnecessary). _build_token_auth is back to a plain per-request bearer read of the session that login established. Co-authored-by: Isaac --- src/ucode/mcp_proxy.py | 158 +++++++-------------------- tests/test_mcp_proxy.py | 233 +++++++++++----------------------------- 2 files changed, 100 insertions(+), 291 deletions(-) diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 4c2c1ab4..0a803061 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -39,7 +39,6 @@ from typing import Protocol, Self import anyio -from anyio.to_thread import run_sync from mcp.client.streamable_http import streamable_http_client from mcp.server.stdio import stdio_server @@ -113,71 +112,35 @@ def _fail_fast(message: str) -> None: raise SystemExit(AUTH_FAILURE_EXIT_CODE) -def _build_token_auth(workspace: str, profile: str | None, url: str): - """Build an httpx ``Auth`` that injects a fresh bearer and logs in on a 401. +def _build_token_auth(workspace: str, profile: str | None): + """Build an httpx ``Auth`` that injects a fresh bearer on every request. 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. - For a connection-backed AI Gateway mcp-services endpoint, an HTTP 401 means - the per-user connection credential is missing. We drive the connection login - once (browser, via ``run_connection_login`` -> ``databricks auth login - --resource``) and retry with a fresh token, so the coding agent just sees the - request authenticate and succeed rather than a failed ``tools/list``. - - The proxy runs on an async event loop, so the login (a blocking subprocess - that waits on the browser) is offloaded to a worker thread in - ``async_auth_flow`` — the loop keeps servicing the stdio pumps and stays - cancellable while the user completes the browser flow, instead of freezing.""" + The bearer is the Databricks *workspace* token, read from the session that + ``serve`` already established via ``databricks auth login --resource`` before + the bridge opened (so a connection-backed service is signed in, credential + and all — see ``serve``). This only *reads* that session's token (refreshing + it as it nears expiry); it never authenticates on its own, so it can't hand + the gateway a token that skips the connection login.""" httpx = _httpx() - connection = connection_from_url(url) - - def _mint_bearer(request): - # get_databricks_token honors the DATABRICKS_BEARER short-circuit and PAT - # profiles internally; --use-pat is surfaced via the env ucode set. A - # RuntimeError means auth is dead (expired refresh token, logged-out - # profile). Raising from inside auth_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) - except RuntimeError as exc: - raise ProxyAuthError(str(exc)) from exc - request.headers["Authorization"] = f"Bearer {token}" - - def _login_and_remint(request): - # Blocking: drives the browser login, then re-mints the now-valid bearer. - ok, detail = run_connection_login(url, workspace, profile=profile) - if not ok: - raise ProxyAuthError(f"connection login for '{connection}' failed: {detail}") - _mint_bearer(request) - - def _needs_login(response) -> bool: - # Only connection-backed services have a per-user login to drive; a 401 from - # anything else is a real auth failure, left to surface as-is. - return connection is not None and response.status_code == 401 class _DatabricksTokenAuth(httpx.Auth): - # Async is the real path (the proxy uses an AsyncClient); the sync flow is - # kept for completeness/parity. Both share the same decision + login logic. def auth_flow(self, request): - _mint_bearer(request) - response = yield request - if not _needs_login(response): - return - _login_and_remint(request) - yield request - - async def async_auth_flow(self, request): - _mint_bearer(request) - response = yield request - if not _needs_login(response): - return - # Offload the blocking browser login so the event loop keeps running - # (mcp-remote-style: the transport stays responsive, not frozen). - await run_sync(lambda: _login_and_remint(request)) + # get_databricks_token honors the DATABRICKS_BEARER short-circuit and + # PAT profiles internally; --use-pat is surfaced via the env ucode set. + # A RuntimeError means the session is dead (expired refresh token, + # logged-out profile). Raising from inside auth_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. + try: + token = get_databricks_token(workspace, profile) + except RuntimeError as exc: + raise ProxyAuthError(str(exc)) from exc + request.headers["Authorization"] = f"Bearer {token}" yield request return _DatabricksTokenAuth() @@ -212,7 +175,7 @@ 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) + auth = _build_token_auth(workspace, profile) # 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. @@ -249,57 +212,6 @@ def _preflight_token(workspace: str, profile: str | None) -> None: get_databricks_token(workspace, profile) -# Timeout for the startup probe MCP round-trips (initialize + tools/list). Short: -# it's a liveness check, not the login (which has its own generous timeout). -_PROBE_TIMEOUT_SECONDS = 15 - - -def _connection_login_required(url: str, workspace: str, profile: str | None) -> bool: - """Whether the connection-backed MCP service answers ``tools/list`` with a 401. - - A lightweight probe run at startup (before the bridge) so the connection login - happens during the agent's "connecting…" phase — like a generic OAuth MCP - bridge — instead of on the agent's first ``tools/list``, where it would race - the agent's tool-fetch timeout. Any non-401 outcome (authenticated, or a - network/transport hiccup) returns ``False`` so startup is never blocked on a - false alarm; a genuine missing credential surfaces again on the live request.""" - httpx = _httpx() - try: - token = get_databricks_token(workspace, profile) - except RuntimeError: - return False # dead databricks auth; let _preflight_token/the bridge report it - headers = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - "Accept": "application/json, text/event-stream", - } - initialize = { - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": { - "protocolVersion": "2024-11-05", - "capabilities": {}, - "clientInfo": {"name": "ucode-mcp-proxy", "version": "0"}, - }, - } - try: - with httpx.Client(timeout=_PROBE_TIMEOUT_SECONDS) as client: - init_response = client.post(url, headers=headers, json=initialize) - session_id = init_response.headers.get("mcp-session-id") - if session_id: - headers["mcp-session-id"] = session_id - client.post( - url, headers=headers, json={"jsonrpc": "2.0", "method": "notifications/initialized"} - ) - tools = client.post( - url, headers=headers, json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"} - ) - return tools.status_code == 401 - except httpx.HTTPError: - return False - - def _unwrap_proxy_error(exc: BaseException) -> ProxyAuthError | ProxyTransportError | None: """Find a known proxy error in an exception (or ExceptionGroup) tree. @@ -332,6 +244,22 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool "Set DATABRICKS_BEARER, or reconfigure the profile." ) + # Connection-backed AI Gateway services need a per-user connection credential + # (e.g. a SaaS login) before their tools can be used. Drive that login *here*, + # before the bridge opens — a blocking `databricks auth login --resource + # `, which a resource-aware /oidc routes through the connection's own + # sign-in (/mcp-service-login) before minting the token. The agent blocks on + # "connecting…" while it runs (the browser opens, or the URL is printed to this + # stderr), then the session comes up already authenticated — so AI Gateway is + # only ever called with a valid credential and never has to elicit a login. It + # is idempotent: once signed in, the login returns immediately with no prompt. + # PAT profiles have no connection OAuth to drive, so they skip it. + connection = None if use_pat else connection_from_url(url) + if connection is not None: + ok, detail = run_connection_login(url, workspace, profile=profile) + if not ok: + _fail_fast(f"connection login for '{connection}' failed: {detail}") + # Pre-flight the token before opening the bridge. Without this, the first # token failure surfaces from inside the transport's task group, where it can # stall the process instead of erroring out. @@ -340,20 +268,6 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool except RuntimeError as exc: _fail_fast(str(exc)) - # Connect-time connection login (generic OAuth-bridge behaviour): before the - # bridge starts, if a connection-backed service is unauthenticated, drive the - # login now — while the agent shows "connecting…" — so the session comes up - # already connected instead of failing the agent's first tools/list. The - # browser opens (databricks-cli honours $BROWSER, inherited here) or the - # authorize URL is printed to this stderr. PAT profiles have no connection - # OAuth to drive. The on-401 retry in the auth hook remains as a mid-session - # fallback (e.g. the credential is revoked while connected). - connection = None if use_pat else connection_from_url(url) - if connection is not None and _connection_login_required(url, workspace, profile): - ok, detail = run_connection_login(url, workspace, profile=profile) - if not ok: - _fail_fast(f"connection login for '{connection}' failed: {detail}") - try: anyio.run(_run, url, workspace, profile) except BaseException as exc: # noqa: BLE001 - re-raised unless it's a known proxy failure diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index 1effe359..9b9e2721 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -61,7 +61,7 @@ def test_proxy_imports_the_streamable_http_client_shared_by_both_majors(): class TestDatabricksTokenAuth: def test_injects_bearer_from_minted_token(self, monkeypatch): monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok-123") - auth = mcp_proxy._build_token_auth(WS, "uc-dogfood", URL) + auth = mcp_proxy._build_token_auth(WS, "uc-dogfood") request = httpx.Request("POST", URL) # auth_flow is a generator that yields the (mutated) request. @@ -73,7 +73,7 @@ def test_auth_is_an_instance_of_the_selected_httpx_auth(self, monkeypatch): # The auth must subclass the *same* httpx flavor's Auth as the transport, # or the SDK's AsyncClient won't accept it. monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "t") - auth = mcp_proxy._build_token_auth(WS, None, URL) + auth = mcp_proxy._build_token_auth(WS, None) assert isinstance(auth, mcp_proxy._httpx().Auth) @@ -84,7 +84,7 @@ def test_calls_get_token_with_workspace_and_profile(self, monkeypatch): "get_databricks_token", lambda ws, profile: calls.append((ws, profile)) or "t", ) - auth = mcp_proxy._build_token_auth(WS, "myprofile", URL) + auth = mcp_proxy._build_token_auth(WS, "myprofile") list(auth.auth_flow(httpx.Request("POST", URL))) @@ -95,7 +95,7 @@ def test_mints_a_fresh_token_per_request(self, monkeypatch): # picked up mid-session without the proxy tracking expiry itself. tokens = iter(["first", "second"]) monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: next(tokens)) - auth = mcp_proxy._build_token_auth(WS, None, URL) + auth = mcp_proxy._build_token_auth(WS, None) r1 = httpx.Request("POST", URL) r2 = httpx.Request("POST", URL) @@ -107,7 +107,7 @@ def test_mints_a_fresh_token_per_request(self, monkeypatch): def test_auth_flow_yields_the_same_request(self, monkeypatch): monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "t") - auth = mcp_proxy._build_token_auth(WS, None, URL) + auth = mcp_proxy._build_token_auth(WS, None) request = httpx.Request("POST", URL) yielded = list(auth.auth_flow(request)) @@ -122,7 +122,7 @@ def boom(ws, profile): raise RuntimeError("no access token; run `databricks auth login`") monkeypatch.setattr(mcp_proxy, "get_databricks_token", boom) - auth = mcp_proxy._build_token_auth(WS, "p", URL) + auth = mcp_proxy._build_token_auth(WS, "p") with pytest.raises(mcp_proxy.ProxyAuthError, match="databricks auth login"): list(auth.auth_flow(httpx.Request("POST", URL))) @@ -131,113 +131,6 @@ def boom(ws, profile): CONN_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github" -def _drive_auth_flow(auth, request, responses): - """Drive an httpx auth_flow generator, feeding ``responses`` back per yield. - - Returns the list of requests the flow yielded (one per attempt).""" - gen = auth.auth_flow(request) - yielded = [next(gen)] - for response in responses: - try: - yielded.append(gen.send(response)) - except StopIteration: - break - return yielded - - -def _response(status): - return httpx.Response(status, request=httpx.Request("POST", CONN_URL)) - - -class TestConnectionLoginOn401: - def test_no_401_does_not_trigger_login(self, monkeypatch): - monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") - logins: list = [] - monkeypatch.setattr( - mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") - ) - auth = mcp_proxy._build_token_auth(WS, "p", CONN_URL) - - yielded = _drive_auth_flow(auth, httpx.Request("POST", CONN_URL), [_response(200)]) - - assert len(yielded) == 1 # no retry - assert logins == [] - - def test_401_drives_login_and_retries_with_fresh_token(self, monkeypatch): - tokens = iter(["stale", "fresh"]) - monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: next(tokens)) - calls: list[tuple] = [] - monkeypatch.setattr( - mcp_proxy, - "run_connection_login", - lambda url, ws, **k: calls.append((url, ws, k.get("profile"))) or (True, "signed in"), - ) - auth = mcp_proxy._build_token_auth(WS, "p", CONN_URL) - - request = httpx.Request("POST", CONN_URL) - yielded = _drive_auth_flow(auth, request, [_response(401), _response(200)]) - - # Logged in for this connection's URL, then retried with the fresh token. - assert calls == [(CONN_URL, WS, "p")] - assert len(yielded) == 2 - assert yielded[1].headers["Authorization"] == "Bearer fresh" - - def test_login_failure_is_terminal(self, monkeypatch): - monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") - monkeypatch.setattr( - mcp_proxy, "run_connection_login", lambda *a, **k: (False, "user cancelled") - ) - auth = mcp_proxy._build_token_auth(WS, "p", CONN_URL) - - with pytest.raises(mcp_proxy.ProxyAuthError, match="user cancelled"): - _drive_auth_flow(auth, httpx.Request("POST", CONN_URL), [_response(401)]) - - def test_non_connection_url_401_does_not_trigger_login(self, monkeypatch): - # A 401 from a non-mcp-services endpoint is a real auth failure, left as-is. - monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") - logins: list = [] - monkeypatch.setattr( - mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") - ) - auth = mcp_proxy._build_token_auth(WS, "p", URL) # URL is not connection-backed - - yielded = _drive_auth_flow(auth, httpx.Request("POST", URL), [_response(401)]) - - assert len(yielded) == 1 # no retry - assert logins == [] - - def test_async_flow_offloads_login_and_retries(self, monkeypatch): - # The real path is async (AsyncClient). async_auth_flow must offload the - # blocking login to a worker thread and retry with a fresh token. - tokens = iter(["stale", "fresh"]) - monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: next(tokens)) - calls: list[tuple] = [] - monkeypatch.setattr( - mcp_proxy, - "run_connection_login", - lambda url, ws, **k: calls.append((url, k.get("profile"))) or (True, "signed in"), - ) - auth = mcp_proxy._build_token_auth(WS, "p", CONN_URL) - - async def scenario(): - # The flow mutates one request object in place, so capture the header - # value at each yield rather than comparing object references. - req = httpx.Request("POST", CONN_URL) - gen = auth.async_auth_flow(req) - await gen.__anext__() - first_auth = req.headers["Authorization"] - await gen.asend(_response(401)) - retry_auth = req.headers["Authorization"] - with pytest.raises(StopAsyncIteration): - await gen.asend(_response(200)) - return first_auth, retry_auth - - first_auth, retry_auth = anyio.run(scenario) - assert calls == [(CONN_URL, "p")] # login fired (offloaded), once - assert first_auth == "Bearer stale" - assert retry_auth == "Bearer fresh" - - class TestPump: def test_forwards_all_messages_in_order(self): async def scenario() -> list[str]: @@ -406,62 +299,6 @@ def test_use_pat_without_a_resolvable_pat_exits_before_serving(self, monkeypatch assert started == [] # never opened the bridge assert "no personal access token" in capsys.readouterr().err - def test_connect_time_login_runs_before_the_bridge_when_required(self, monkeypatch): - order: list = [] - monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) - monkeypatch.setattr(mcp_proxy, "_connection_login_required", lambda url, ws, profile: True) - monkeypatch.setattr( - mcp_proxy, - "run_connection_login", - lambda url, ws, **k: order.append(("login", url)) or (True, "signed in"), - ) - monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: order.append(("bridge",))) - - mcp_proxy.serve(CONN_URL, WS, "p") - - # Login (during "connecting…") happens before the bridge opens. - assert order == [("login", CONN_URL), ("bridge",)] - - def test_authenticated_connection_skips_connect_time_login(self, monkeypatch): - logins: list = [] - monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) - monkeypatch.setattr(mcp_proxy, "_connection_login_required", lambda url, ws, profile: False) - monkeypatch.setattr( - mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") - ) - monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) - - mcp_proxy.serve(CONN_URL, WS, "p") - - assert logins == [] # already authenticated -> no login - - def test_connect_time_login_failure_is_terminal(self, monkeypatch): - started: list = [] - monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) - monkeypatch.setattr(mcp_proxy, "_connection_login_required", lambda url, ws, profile: True) - monkeypatch.setattr( - mcp_proxy, "run_connection_login", lambda *a, **k: (False, "user cancelled") - ) - monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: started.append("bridge")) - - with pytest.raises(SystemExit): - mcp_proxy.serve(CONN_URL, WS, "p") - - assert started == [] # never opened the bridge - - def test_use_pat_skips_the_connect_time_probe(self, monkeypatch): - probed: list = [] - monkeypatch.setattr(mcp_proxy, "ensure_pat_bearer", lambda profile: True) - monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) - monkeypatch.setattr( - mcp_proxy, "_connection_login_required", lambda *a: probed.append(1) or True - ) - monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) - - mcp_proxy.serve(CONN_URL, WS, "p", use_pat=True) - - assert probed == [] # PAT has no connection OAuth to probe/drive - def test_oauth_path_never_touches_pat(self, monkeypatch): # Without use_pat, ensure_pat_bearer must not be consulted at all. called: list[str] = [] @@ -558,6 +395,64 @@ def raise_other(func, *args): with pytest.raises(ValueError, match="some transport bug"): mcp_proxy.serve(URL, WS, "p") + def test_connection_backed_url_logs_in_before_the_bridge(self, monkeypatch): + # A connection-backed mcp-services URL drives `databricks auth login + # --resource` up front (blocking), then opens the bridge — so the session + # is authenticated before AI Gateway is ever called. + order: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda url, ws, **k: ( + order.append(("login", url, k.get("profile"))) or (True, "signed in") + ), + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: order.append(("bridge",))) + + mcp_proxy.serve(CONN_URL, WS, "p") + + assert order == [("login", CONN_URL, "p"), ("bridge",)] + + def test_non_connection_url_skips_the_login(self, monkeypatch): + logins: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) + + mcp_proxy.serve(URL, WS, "p") # URL is not an mcp-services endpoint + + assert logins == [] + + def test_connection_login_failure_exits_before_the_bridge(self, monkeypatch): + started: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: (False, "user cancelled") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: started.append("bridge")) + + with pytest.raises(SystemExit) as excinfo: + mcp_proxy.serve(CONN_URL, WS, "p") + + assert excinfo.value.code == mcp_proxy.AUTH_FAILURE_EXIT_CODE + assert started == [] # never opened the bridge + + def test_use_pat_skips_the_connection_login(self, monkeypatch): + logins: list = [] + monkeypatch.setattr(mcp_proxy, "ensure_pat_bearer", lambda profile: True) + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) + + mcp_proxy.serve(CONN_URL, WS, "p", use_pat=True) + + assert logins == [] # PAT has no connection OAuth to drive + class TestPreflightToken: def test_passes_through_when_a_token_is_available(self, monkeypatch):