From 88f2574ac8048171acf3e77cbb038328f730c9f3 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Thu, 10 Sep 2026 23:44:23 +0000 Subject: [PATCH 1/6] 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 000000000..4c994034f --- /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 c01741698..21edec794 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 000000000..f1b3b7988 --- /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 0ded63d9a..527d61333 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 bf13b70368340739c42fbb142853127c9c9ff0be Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 01:43:29 +0000 Subject: [PATCH 2/6] 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 4c994034f..bb3c8557b 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 21edec794..2486ff177 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 f1b3b7988..4a9a9bfb4 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 527d61333..b7d7426df 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 a526acb20243a9546da5fd4e3c96418fbd975700 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 02:23:34 +0000 Subject: [PATCH 3/6] 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 2486ff177..6ed26008b 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 b7d7426df..ab3c321fd 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 a0dcd848c4fcc9ade4ef0c75b2d0bcd7c59f8891 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 03:06:05 +0000 Subject: [PATCH 4/6] 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 6ed26008b..486e705df 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 ab3c321fd..47e59f706 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): From 608e8d75febd15651580e98e087b64f42c6789e1 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Tue, 15 Sep 2026 20:03:37 +0000 Subject: [PATCH 5/6] mcp-proxy: clear error when the Databricks CLI lacks --resource The connect-time connection login runs `databricks auth login --resource`; an older CLI without databricks/cli#6621 rejects the flag and the login exits with a cryptic parse error (seen as a bare "connection closed" in the agent). Add a one-time `auth login --help` pre-check and, when `--resource` is absent, return a clear "upgrade your Databricks CLI" message instead of attempting the doomed login. Fail-open if --help can't run. Co-authored-by: Isaac --- src/ucode/mcp_connection_login.py | 29 +++++++++++++++++++++++++- tests/test_mcp_connection_login.py | 33 +++++++++++++++++++++++++----- 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py index bb3c8557b..37b2efa1c 100644 --- a/src/ucode/mcp_connection_login.py +++ b/src/ucode/mcp_connection_login.py @@ -46,6 +46,27 @@ def connection_from_url(url: str) -> str | None: return connection or None +def _cli_supports_resource_flag(login_binary: str) -> bool: + """Whether `` auth login`` advertises the ``--resource`` flag. + + The connection sign-in needs a Databricks CLI with ``--resource`` + (databricks/cli#6621). An older CLI rejects the flag and the login exits with + a cryptic parse error, so we check ``--help`` up front to give a clear message + instead. Fail-open (assume supported) if ``--help`` can't be run — the real + login attempt will surface any genuine failure.""" + try: + result = subprocess.run( + [login_binary, "auth", "login", "--help"], + check=False, + timeout=20, + capture_output=True, + text=True, + ) + except (OSError, subprocess.TimeoutExpired): + return True + return "--resource" in f"{result.stdout or ''}{result.stderr or ''}" + + def run_connection_login( resource_url: str, workspace: str, @@ -67,6 +88,13 @@ def run_connection_login( stays visible when the browser can't open (e.g. a headless remote). Returns ``(ok, message)``; on failure ``message`` points at that log. """ + connection = connection_from_url(resource_url) or resource_url + if not _cli_supports_resource_flag(login_binary): + return False, ( + f"the Databricks CLI ('{login_binary}') has no `--resource` flag, so the " + f"'{connection}' connection sign-in can't run. Upgrade the CLI " + "(databricks/cli#6621) and retry." + ) argv = [ login_binary, "auth", @@ -78,7 +106,6 @@ 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.", diff --git a/tests/test_mcp_connection_login.py b/tests/test_mcp_connection_login.py index 4a9a9bfb4..bfb662fd5 100644 --- a/tests/test_mcp_connection_login.py +++ b/tests/test_mcp_connection_login.py @@ -31,6 +31,11 @@ def test_missing_service_is_none(self): class TestRunConnectionLogin: def _fake_run(self, captured, *, returncode, stderr=""): def _run(argv, **kwargs): + # The `--resource`-support pre-check runs `auth login --help` first. + if "--help" in argv: + return subprocess.CompletedProcess( + argv, 0, stdout="--resource stringArray", stderr="" + ) captured.append(argv) return subprocess.CompletedProcess(argv, returncode, stdout="", stderr=stderr) @@ -64,11 +69,14 @@ 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", - lambda argv, **kw: seen.update(kw) or subprocess.CompletedProcess(argv, 0), - ) + + def _run(argv, **kw): + if "--help" in argv: # the --resource pre-check; not the login call under test + return subprocess.CompletedProcess(argv, 0, stdout="--resource", stderr="") + seen.update(kw) + return subprocess.CompletedProcess(argv, 0) + + monkeypatch.setattr(mcl.subprocess, "run", _run) ok, _ = mcl.run_connection_login(AIGW_URL, WS) assert ok assert seen.get("stdout") is mcl.sys.stderr @@ -90,3 +98,18 @@ def _run(argv, **kwargs): monkeypatch.setattr(mcl.subprocess, "run", _run) ok, message = mcl.run_connection_login(AIGW_URL, WS) assert not ok and "could not run" in message + + def test_old_cli_without_resource_flag_reports_clearly(self, monkeypatch): + # `auth login --help` lacking `--resource` => an old CLI (no databricks/cli#6621). + # We must report that clearly and never attempt the login (the flag would error). + def _run(argv, **kwargs): + if "--help" in argv: + return subprocess.CompletedProcess( + argv, 0, stdout="usage: login [--host]", stderr="" + ) + raise AssertionError("login must not run when --resource is unsupported") + + monkeypatch.setattr(mcl.subprocess, "run", _run) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok + assert "--resource" in message and "Upgrade" in message From 15c80f014e5bb51e7f6df323be702f60919efd07 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Tue, 22 Sep 2026 00:20:30 +0000 Subject: [PATCH 6/6] mcp-proxy: only run the connection login when the credential is actually missing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The connect-time login ran `databricks auth login --resource` unconditionally on every proxy startup for every connection-backed server. But `databricks auth login` always re-runs the browser OAuth (it does not short-circuit on a cached token), so N configured servers meant N browser windows on every session — and concurrent `auth login` runs also race on the shared token cache. Gate the login on the actual per-user credential state (new `connection_credential_state`, using the same Unity Catalog REST APIs as `ug mcp login`): only run the login when the credential is confirmed MISSING. When it is already PRESENT — or the state can't be determined (UNKNOWN) or the connection isn't OAuth-U2M (NO_LOGIN) — skip the browser entirely. `tools/list` still works without the credential, so UNKNOWN fails safe rather than risking a spurious browser. This makes repeat sessions open zero browsers and first sessions prompt only for the servers genuinely needing a sign-in. Co-authored-by: Isaac --- src/ucode/mcp_connection_login.py | 96 ++++++++++++++++++++++++++++++ src/ucode/mcp_proxy.py | 29 ++++++--- tests/test_mcp_connection_login.py | 76 +++++++++++++++++++++++ tests/test_mcp_proxy.py | 56 +++++++++++++++-- 4 files changed, 245 insertions(+), 12 deletions(-) diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py index 37b2efa1c..1d61a7657 100644 --- a/src/ucode/mcp_connection_login.py +++ b/src/ucode/mcp_connection_login.py @@ -20,6 +20,14 @@ import subprocess import sys +from urllib.parse import quote + +from ucode.databricks import ( + _http_get_json, + _scim_me, + get_databricks_token, + workspace_hostname, +) # AI Gateway MCP service endpoints look like # ``https:///ai-gateway/mcp-services/..``. @@ -29,6 +37,25 @@ # allow generously more than a token refresh would take. _LOGIN_TIMEOUT_SECONDS = 300 +# Connection securable kinds that use a per-user OAuth (U2M) credential — i.e. the +# service needs a connection sign-in. Anything else (PAT/basic/service-managed) +# has no per-user login to drive. Mirrors the webapp's known-OAuth-kinds set. +_OAUTH_U2M_CONNECTION_KINDS = frozenset( + { + "CONNECTION_HTTP_OAUTH_U2M_MAPPING", + "CONNECTION_HTTP_DCR", + "CONNECTION_SLACK_OAUTH_U2M_MAPPING", + } +) + +# Credential-state outcomes for :func:`connection_credential_state`. +CREDENTIAL_PRESENT = "present" # signed in — skip the login +CREDENTIAL_MISSING = "missing" # confirmed no credential yet — run the login +CREDENTIAL_NO_LOGIN = "no_login_needed" # not an OAuth-U2M connection — skip +CREDENTIAL_UNKNOWN = ( + "unknown" # couldn't determine — skip (don't open a browser we're unsure about) +) + def connection_from_url(url: str) -> str | None: """Return the connection FQN of an AI Gateway MCP service URL, or ``None``. @@ -46,6 +73,70 @@ def connection_from_url(url: str) -> str | None: return connection or None +def _strip_connections_prefix(name: str | None) -> str | None: + """UC returns a connection reference as ``connections/``; the REST path + wants the bare name.""" + if not name: + return None + prefix = "connections/" + return name[len(prefix) :] if name.startswith(prefix) else name + + +def connection_credential_state( + resource_url: str, workspace: str, *, profile: str | None = None +) -> str: + """Best-effort per-user credential state for a connection-backed MCP service. + + Gates the connect-time login: the proxy should only run ``databricks auth + login`` (which opens a browser) when the credential is genuinely **missing**, + not on every session for a service the user already signed in to — otherwise + N configured servers would each pop a browser every session. + + Uses the same Unity Catalog REST APIs as ``ug mcp login``: resolve the + service's backing connection, then read the current user's credential + ``provisioning_info.state``. Returns one of ``CREDENTIAL_{PRESENT,MISSING, + NO_LOGIN,UNKNOWN}``. Anything but ``MISSING`` means "don't open a browser": + ``PRESENT``/``NO_LOGIN`` are definitive, and ``UNKNOWN`` (any probe failure) + deliberately fails safe — ``tools/list`` still works, and a tool call can + surface the login later — rather than risk a spurious browser. + """ + full = connection_from_url(resource_url) + if not full: + return CREDENTIAL_NO_LOGIN + try: + token = get_databricks_token(workspace, profile) + except Exception: # noqa: BLE001 - a dead token is reported by _preflight_token, not here + return CREDENTIAL_UNKNOWN + user = (_scim_me(workspace, token) or {}).get("userName") + if not user: + return CREDENTIAL_UNKNOWN + prefix = f"https://{workspace_hostname(workspace)}/api/2.1/unity-catalog" + details, err = _http_get_json(f"{prefix}/mcp-services/{quote(full, safe='')}", token) + if err is not None or not isinstance(details, dict): + return CREDENTIAL_UNKNOWN + source = (details.get("config") or {}).get("source_connection") or {} + if source.get("securable_kind") not in _OAUTH_U2M_CONNECTION_KINDS: + return CREDENTIAL_NO_LOGIN + conn = _strip_connections_prefix(source.get("name")) + service_id = details.get("id") + if not conn or not service_id: + return CREDENTIAL_UNKNOWN + cred_url = ( + f"{prefix}/connections/{quote(conn, safe='')}/user-credentials/" + f"{quote(user, safe='')}?dependent.mcp_service.id={quote(str(service_id), safe='')}" + ) + cred, cred_err = _http_get_json(cred_url, token) + if cred_err is not None: + # 404 is the authoritative "no credential yet"; any other error is inconclusive. + return CREDENTIAL_MISSING if cred_err.startswith("HTTP 404") else CREDENTIAL_UNKNOWN + if not isinstance(cred, dict): + return CREDENTIAL_UNKNOWN + state = ((cred.get("connection_user_credential") or {}).get("provisioning_info") or {}).get( + "state" + ) + return CREDENTIAL_PRESENT if state == "ACTIVE" else CREDENTIAL_MISSING + + def _cli_supports_resource_flag(login_binary: str) -> bool: """Whether `` auth login`` advertises the ``--resource`` flag. @@ -138,6 +229,11 @@ def run_connection_login( __all__ = [ "AIGW_MCP_SERVICES_SEGMENT", + "CREDENTIAL_MISSING", + "CREDENTIAL_NO_LOGIN", + "CREDENTIAL_PRESENT", + "CREDENTIAL_UNKNOWN", + "connection_credential_state", "connection_from_url", "run_connection_login", ] diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 486e705df..da54c3933 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -43,7 +43,12 @@ 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 +from ucode.mcp_connection_login import ( + CREDENTIAL_MISSING, + connection_credential_state, + 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. @@ -250,15 +255,23 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool # `, 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. + # stderr), then the session comes up already authenticated. + # + # But `databricks auth login` always re-runs the browser OAuth — it does NOT + # short-circuit on a cached token — so doing it unconditionally would pop a + # browser for EVERY connection-backed server on EVERY session. So gate it on + # the actual credential state: only log in when the credential is confirmed + # missing. When it's already present (or the state can't be determined) we + # skip the browser entirely — this is what makes N configured servers not + # equal N browser windows per session. PAT profiles have no connection OAuth + # to drive, so they skip it too. 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}") + state = connection_credential_state(url, workspace, profile=profile) + if state == CREDENTIAL_MISSING: + 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 diff --git a/tests/test_mcp_connection_login.py b/tests/test_mcp_connection_login.py index bfb662fd5..afe1de91c 100644 --- a/tests/test_mcp_connection_login.py +++ b/tests/test_mcp_connection_login.py @@ -28,6 +28,82 @@ def test_missing_service_is_none(self): assert mcl.connection_from_url(f"{WS}/ai-gateway/mcp-services/") is None +class TestConnectionCredentialState: + """The gate that stops the proxy re-running the browser login every session.""" + + _DETAILS = { + "id": "svc-1", + "config": { + "source_connection": { + "name": "connections/github", + "securable_kind": "CONNECTION_HTTP_OAUTH_U2M_MAPPING", + } + }, + } + + def _wire(self, monkeypatch, *, details=_DETAILS, details_err=None, cred=None, cred_err=None): + monkeypatch.setattr(mcl, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mcl, "_scim_me", lambda ws, token: {"userName": "u@databricks.com"}) + + def _http(url, token, *a, **k): + if "/mcp-services/" in url: + return details, details_err + if "/user-credentials/" in url: + return cred, cred_err + return None, "unexpected url" + + monkeypatch.setattr(mcl, "_http_get_json", _http) + + def test_present_when_credential_active(self, monkeypatch): + self._wire( + monkeypatch, + cred={"connection_user_credential": {"provisioning_info": {"state": "ACTIVE"}}}, + ) + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_PRESENT + + def test_missing_on_404(self, monkeypatch): + self._wire(monkeypatch, cred_err="HTTP 404 Not Found") + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_MISSING + + def test_missing_when_state_not_active(self, monkeypatch): + self._wire( + monkeypatch, + cred={"connection_user_credential": {"provisioning_info": {"state": "PROVISIONING"}}}, + ) + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_MISSING + + def test_no_login_for_non_oauth_connection(self, monkeypatch): + details = { + "id": "x", + "config": { + "source_connection": {"name": "connections/c", "securable_kind": "CONNECTION_MYSQL"} + }, + } + self._wire(monkeypatch, details=details) + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_NO_LOGIN + + def test_unknown_on_service_lookup_error(self, monkeypatch): + self._wire(monkeypatch, details=None, details_err="HTTP 500") + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_UNKNOWN + + def test_unknown_on_non_404_credential_error(self, monkeypatch): + self._wire(monkeypatch, cred_err="HTTP 403 Forbidden") + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_UNKNOWN + + def test_unknown_when_token_unavailable(self, monkeypatch): + def boom(ws, profile): + raise RuntimeError("no token") + + monkeypatch.setattr(mcl, "get_databricks_token", boom) + assert mcl.connection_credential_state(AIGW_URL, WS) == mcl.CREDENTIAL_UNKNOWN + + def test_no_login_for_non_aigw_url(self, monkeypatch): + assert ( + mcl.connection_credential_state(f"{WS}/api/2.0/mcp/functions/x", WS) + == mcl.CREDENTIAL_NO_LOGIN + ) + + class TestRunConnectionLogin: def _fake_run(self, captured, *, returncode, stderr=""): def _run(argv, **kwargs): diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index 47e59f706..ec8831156 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -10,6 +10,7 @@ import httpx import pytest +from ucode import mcp_connection_login as mcl from ucode import mcp_proxy WS = "https://example.databricks.com" @@ -395,12 +396,17 @@ 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. + def test_connection_backed_url_logs_in_before_the_bridge_when_credential_missing( + self, monkeypatch + ): + # A connection-backed mcp-services URL whose credential is MISSING drives + # `databricks auth login --resource` up front (blocking), then opens the + # bridge — so the session is authenticated before AI Gateway is called. order: list = [] monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "connection_credential_state", lambda *a, **k: mcp_proxy.CREDENTIAL_MISSING + ) monkeypatch.setattr( mcp_proxy, "run_connection_login", @@ -414,6 +420,45 @@ def test_connection_backed_url_logs_in_before_the_bridge(self, monkeypatch): assert order == [("login", CONN_URL, "p"), ("bridge",)] + def test_present_credential_skips_the_login_and_opens_the_bridge(self, monkeypatch): + # The key scaling fix: an already-signed-in connection must NOT re-run the + # browser login — otherwise N configured servers pop N browsers per session. + logins: list = [] + opened: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "connection_credential_state", lambda *a, **k: mcl.CREDENTIAL_PRESENT + ) + monkeypatch.setattr( + mcp_proxy, "run_connection_login", lambda *a, **k: logins.append(1) or (True, "") + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: opened.append("bridge")) + + mcp_proxy.serve(CONN_URL, WS, "p") + + assert logins == [] # no browser + assert opened == ["bridge"] # session still comes up + + def test_unknown_or_no_login_state_skips_the_login(self, monkeypatch): + # UNKNOWN (probe failed) and NO_LOGIN (non-OAuth connection) both fail safe: + # no eager browser — tools/list still works and a tool call can surface it later. + for state in (mcl.CREDENTIAL_UNKNOWN, mcl.CREDENTIAL_NO_LOGIN): + logins: list = [] + monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) + monkeypatch.setattr( + mcp_proxy, "connection_credential_state", lambda *a, s=state, **k: s + ) + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, _l=logins, **k: _l.append(1) or (True, ""), + ) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: None) + + mcp_proxy.serve(CONN_URL, WS, "p") + + assert logins == [], f"state={state} must not open a browser" + def test_non_connection_url_skips_the_login(self, monkeypatch): logins: list = [] monkeypatch.setattr(mcp_proxy, "_preflight_token", lambda ws, profile: None) @@ -429,6 +474,9 @@ def test_non_connection_url_skips_the_login(self, monkeypatch): 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, "connection_credential_state", lambda *a, **k: mcp_proxy.CREDENTIAL_MISSING + ) monkeypatch.setattr( mcp_proxy, "run_connection_login", lambda *a, **k: (False, "user cancelled") )