From 264d2969e2910b65b7a512cc052c7ff9345e1b40 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Thu, 10 Sep 2026 23:44:23 +0000 Subject: [PATCH 01/10] 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 6cfbb49810a8127d2bf32ce2889437dbcef35228 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 01:43:29 +0000 Subject: [PATCH 02/10] 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 e2de85c215c8889bfb45290429491f7844cbca3f Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 02:23:34 +0000 Subject: [PATCH 03/10] 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 12c9061b36a92ef54c9b8b107ed6cb248fd3db6a Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 03:06:05 +0000 Subject: [PATCH 04/10] 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 2dffbf141d75cee1e6b0efb40b69f9544bdc820d Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Tue, 15 Sep 2026 20:03:37 +0000 Subject: [PATCH 05/10] 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 c1e3f6dc4cab0d212ec364c7abf3e94051102bf8 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Tue, 22 Sep 2026 00:20:30 +0000 Subject: [PATCH 06/10] 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") ) From e000015daa09e5201f873a69673ca2551890d8df Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Tue, 22 Sep 2026 21:23:52 +0000 Subject: [PATCH 07/10] mcp-proxy: drive the connection login lazily on a 401, not eagerly at startup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The connect-time login blocked the bridge from opening until it finished. With N configured connection-backed servers all missing their credential, an agent spawns N proxies at once, each launches a blocking `databricks auth login`, and they stall on browser flows / racing loopback ports until every one hits the agent's MCP startup timeout — all N servers fail to start (observed with 5 on Codex, all timing out at 120s). Move the login into the httpx auth flow, triggered by a 401: the bridge opens immediately (tools/list needs no credential and returns the full list), and `databricks auth login --resource` runs only when a request actually gets a 401 — i.e. when a tool you invoke needs the per-user connection credential. So a browser opens only for a service you use, one at a time; an already-signed-in service never prompts; and startup never blocks. The login runs at most once per session, off the event loop (async_auth_flow + anyio.to_thread) so the bridge's pumps aren't starved. PAT profiles have no connection OAuth, so they skip it. This drops the eager credential-state gate added in the previous commit (the lazy 401 signal supersedes it) along with its UC-REST probe. Co-authored-by: Isaac --- src/ucode/mcp_connection_login.py | 96 ------------ src/ucode/mcp_proxy.py | 130 +++++++++-------- tests/test_mcp_connection_login.py | 76 ---------- tests/test_mcp_proxy.py | 227 ++++++++++++++--------------- 4 files changed, 185 insertions(+), 344 deletions(-) diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py index 1d61a7657..37b2efa1c 100644 --- a/src/ucode/mcp_connection_login.py +++ b/src/ucode/mcp_connection_login.py @@ -20,14 +20,6 @@ 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/..``. @@ -37,25 +29,6 @@ # 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``. @@ -73,70 +46,6 @@ 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. @@ -229,11 +138,6 @@ 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 da54c3933..2dcde9428 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -39,16 +39,12 @@ from typing import Protocol, Self import anyio +from anyio import to_thread from mcp.client.streamable_http import streamable_http_client from mcp.server.stdio import stdio_server from ucode.databricks import ensure_pat_bearer, get_databricks_token -from ucode.mcp_connection_login import ( - CREDENTIAL_MISSING, - connection_credential_state, - connection_from_url, - run_connection_login, -) +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. @@ -117,35 +113,72 @@ 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. - - 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. - - 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.""" +def _build_token_auth(url: str, workspace: str, profile: str | None, *, use_pat: bool = False): + """Build an httpx ``Auth`` that injects a fresh bearer and drives the + per-connection login **lazily, on a 401** — not eagerly at startup. + + The bearer is the Databricks *workspace* token (read from the CLI session, + refreshed as it nears expiry). For a connection-backed AI Gateway service, + the gateway answers with an RFC 9728 ``401`` until the user holds the per-user + connection credential. Rather than block startup logging in every configured + server (which pops N browsers and stalls the agent's MCP startup), we let the + bridge come up immediately — ``tools/list`` needs no credential — and only + run ``databricks auth login --resource`` when a request actually gets a 401. + So a browser opens only for a service you actually *use*, one at a time, and + an already-signed-in service never prompts. The login runs at most once per + session; PAT profiles have no connection OAuth to drive, so they never do it. + + The base class comes from whichever httpx the SDK uses (see ``_httpx``); the + sync and async flavours share the same generator contract. The async client + uses ``async_auth_flow``, so the blocking login runs off the event loop.""" httpx = _httpx() + # PAT auth has no interactive OAuth to drive, so never treat it as connection-backed. + connection = None if use_pat else connection_from_url(url) + login = {"attempted": False} + + def _mint(request): + # get_databricks_token honors the DATABRICKS_BEARER short-circuit and PAT + # profiles internally. 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, 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}" + + def _needs_connection_login(response) -> bool: + # A 401 from a connection-backed service means the per-user credential is + # missing. Only drive the login once per session — if it doesn't resolve + # the 401, re-running it on every subsequent request would loop browsers. + return connection is not None and response.status_code == 401 and not login["attempted"] + + def _login_or_fail() -> None: + login["attempted"] = True + ok, detail = run_connection_login(url, workspace, profile=profile) + if not ok: + raise ProxyAuthError(f"connection login for '{connection}' failed: {detail}") 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 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}" + _mint(request) + response = yield request + if not _needs_connection_login(response): + return + _login_or_fail() + _mint(request) + yield request + + async def async_auth_flow(self, request): + _mint(request) + response = yield request + if not _needs_connection_login(response): + return + # Run the blocking browser login off the event loop so the bridge's + # other pumps aren't starved while the user completes the sign-in. + await to_thread.run_sync(_login_or_fail) + _mint(request) yield request return _DatabricksTokenAuth() @@ -178,9 +211,9 @@ async def _pump_upstream[T]( raise ProxyTransportError("upstream MCP transport closed unexpectedly") -async def _run(url: str, workspace: str, profile: str | None) -> None: +async def _run(url: str, workspace: str, profile: str | None, use_pat: bool = False) -> None: httpx = _httpx() - auth = _build_token_auth(workspace, profile) + auth = _build_token_auth(url, workspace, profile, use_pat=use_pat) # 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,29 +282,12 @@ 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. - # - # 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: - 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}") + # The per-connection sign-in is NOT driven here. Doing it eagerly at startup + # blocks the agent's MCP startup and, with N configured connection-backed + # servers, pops N browsers at once and stalls until each times out. Instead + # the bridge opens immediately (tools/list needs no credential) and the login + # is driven lazily, on a 401, from the httpx auth flow (see _build_token_auth) + # — so a browser opens only for a service whose tool you actually call. # 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 @@ -282,7 +298,7 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool _fail_fast(str(exc)) try: - anyio.run(_run, url, workspace, profile) + anyio.run(_run, url, workspace, profile, use_pat) except BaseException as exc: # noqa: BLE001 - re-raised unless it's a known proxy failure # Errors raised inside the transport arrive wrapped by its task group. # Report expected auth/transport failures without hiding programming bugs. diff --git a/tests/test_mcp_connection_login.py b/tests/test_mcp_connection_login.py index afe1de91c..bfb662fd5 100644 --- a/tests/test_mcp_connection_login.py +++ b/tests/test_mcp_connection_login.py @@ -28,82 +28,6 @@ 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 ec8831156..3b2e0b97b 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -10,7 +10,6 @@ import httpx import pytest -from ucode import mcp_connection_login as mcl from ucode import mcp_proxy WS = "https://example.databricks.com" @@ -62,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(URL, WS, "uc-dogfood") request = httpx.Request("POST", URL) # auth_flow is a generator that yields the (mutated) request. @@ -74,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(URL, WS, None) assert isinstance(auth, mcp_proxy._httpx().Auth) @@ -85,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(URL, WS, "myprofile") list(auth.auth_flow(httpx.Request("POST", URL))) @@ -96,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(URL, WS, None) r1 = httpx.Request("POST", URL) r2 = httpx.Request("POST", URL) @@ -108,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(URL, WS, None) request = httpx.Request("POST", URL) yielded = list(auth.auth_flow(request)) @@ -123,11 +122,114 @@ 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(URL, WS, "p") with pytest.raises(mcp_proxy.ProxyAuthError, match="databricks auth login"): list(auth.auth_flow(httpx.Request("POST", URL))) + # --- lazy on-401 connection login --------------------------------------- + + @staticmethod + def _run_flow(auth, request, response): + """Drive the sync auth_flow, feeding `response` after the first yield. + Returns the yielded requests (1 = no retry, 2 = logged in and retried).""" + gen = auth.auth_flow(request) + yielded = [next(gen)] + try: + yielded.append(gen.send(response)) + except StopIteration: + pass + return yielded + + def test_on_401_from_connection_service_runs_login_then_retries(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + logins: list = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda url, ws, **k: logins.append((url, k.get("profile"))) or (True, "signed in"), + ) + auth = mcp_proxy._build_token_auth(CONN_URL, WS, "p") + yielded = self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + assert logins == [(CONN_URL, "p")] # login driven once, only on the 401 + assert len(yielded) == 2 # retried after signing in + assert yielded[1].headers["Authorization"] == "Bearer tok" + + def test_non_connection_401_is_not_retried(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(URL, WS, "p") # URL is not an mcp-services endpoint + yielded = self._run_flow(auth, httpx.Request("POST", URL), httpx.Response(401)) + assert logins == [] and len(yielded) == 1 + + def test_success_response_never_logs_in(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(CONN_URL, WS, "p") + yielded = self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(200)) + assert logins == [] and len(yielded) == 1 # tools/list etc. never trigger a browser + + def test_login_runs_at_most_once_per_session(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(CONN_URL, WS, "p") + # Two 401s in the same session — the browser login must fire only once. + self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + assert logins == [1] + + def test_use_pat_never_drives_connection_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(CONN_URL, WS, "p", use_pat=True) + yielded = self._run_flow(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + assert logins == [] and len(yielded) == 1 # PAT has no connection OAuth to drive + + def test_login_failure_becomes_a_proxy_auth_error(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(CONN_URL, WS, "p") + gen = auth.auth_flow(httpx.Request("POST", CONN_URL)) + next(gen) + with pytest.raises(mcp_proxy.ProxyAuthError, match="user cancelled"): + gen.send(httpx.Response(401)) + + def test_async_auth_flow_drives_login_on_401(self, monkeypatch): + # The proxy uses the async client, so async_auth_flow is the real path; the + # blocking login runs off the event loop via anyio.to_thread. + 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(CONN_URL, WS, "p") + + async def scenario(): + gen = auth.async_auth_flow(httpx.Request("POST", CONN_URL)) + await gen.__anext__() + try: + await gen.asend(httpx.Response(401)) + except StopAsyncIteration: + pass + await gen.aclose() + return logins + + assert anyio.run(scenario) == [1] + CONN_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github" @@ -234,7 +336,7 @@ async def stop_bridge(*args, **kwargs): yield monkeypatch.setattr(httpx_module, "AsyncClient", CapturingClient) - monkeypatch.setattr(mcp_proxy, "_build_token_auth", lambda *args: object()) + monkeypatch.setattr(mcp_proxy, "_build_token_auth", lambda *args, **kwargs: object()) monkeypatch.setattr(mcp_proxy, "streamable_http_client", stop_bridge) with pytest.raises(StopBridge): @@ -259,7 +361,7 @@ def fake_run(func, *args): assert captured["func"] is mcp_proxy._run # PAT is resolved inside the token mint, so _run takes no use_pat arg. - assert captured["args"] == (URL, WS, "uc-dogfood") + assert captured["args"] == (URL, WS, "uc-dogfood", False) def test_defaults_profile_none(self, monkeypatch): captured: dict = {} @@ -268,7 +370,7 @@ def test_defaults_profile_none(self, monkeypatch): mcp_proxy.serve(URL, WS) - assert captured["args"] == (URL, WS, None) + assert captured["args"] == (URL, WS, None, False) def test_use_pat_exports_the_bearer_before_serving(self, monkeypatch): # PAT auth: the profile's static token must be exported (ensure_pat_bearer) @@ -396,111 +498,6 @@ 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_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", - 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_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) - 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, "connection_credential_state", lambda *a, **k: mcp_proxy.CREDENTIAL_MISSING - ) - 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 acb8ddc9e4e0617a9cf6c08e3d80074702d370d8 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 23 Sep 2026 04:46:41 +0000 Subject: [PATCH 08/10] mcp_connection_login: make the sign-in message reusable by `ug mcp login` Drop the proxy-specific `ucode mcp-proxy:` prefix from the connect-time sign-in message so `run_connection_login` / `connection_from_url` can be reused verbatim by `ug mcp login` (PR #654), which stacks on this one. Behavior is unchanged for the proxy. Co-authored-by: Isaac --- src/ucode/mcp_connection_login.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py index 37b2efa1c..2da6a5d82 100644 --- a/src/ucode/mcp_connection_login.py +++ b/src/ucode/mcp_connection_login.py @@ -107,8 +107,8 @@ def run_connection_login( if profile: argv += ["--profile", profile] 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.", + f"Signing in to '{connection}' — opening your browser to complete the connection login; " + "if it doesn't open, use the authorization URL printed below.", file=sys.stderr, flush=True, ) From 3403e58e0e884eb153ab7f73183f9dc43765b19f Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 23 Sep 2026 15:51:34 +0000 Subject: [PATCH 09/10] mcp: one source of truth for the /ai-gateway/mcp-services/ path + FQN parse Consolidate the AI Gateway mcp-services URL knowledge that had drifted into three copies: the path-segment literal lived in databricks.py (URL building), mcp.py (AIGW_MCP_SERVICES_PATH, URL-shape classification), and mcp_connection_login.py (AIGW_MCP_SERVICES_SEGMENT, connection-backed detection). Define AIGW_MCP_SERVICES_SEGMENT once in databricks.py (where the URL is built) and import it into mcp.py and mcp_connection_login.py. Also reuse connection_from_url in mcp._mcp_server_location instead of a parallel url.split(), so the URL -> FQN parse lives in one place too. No behavior change (the parser additionally strips a trailing slash/query, which is strictly more correct for the label). Co-authored-by: Isaac --- src/ucode/databricks.py | 9 ++++++++- src/ucode/mcp.py | 20 +++++++++----------- src/ucode/mcp_connection_login.py | 5 +---- tests/test_mcp.py | 6 +++--- 4 files changed, 21 insertions(+), 19 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a4c27e772..fd73981e9 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1975,6 +1975,13 @@ def delete_coding_agent_config(workspace: str, token: str, name: str) -> str | N # --- MCP services (parallel to model services) ----------------------------- +# Canonical path segment of an AI Gateway MCP-services endpoint +# (``https:///ai-gateway/mcp-services/..``). This is +# the single source of truth: URL building (below), connection-backed detection +# (`mcp_connection_login.connection_from_url`), and URL-shape classification +# (`mcp.py`) all reference this one constant. +AIGW_MCP_SERVICES_SEGMENT = "/ai-gateway/mcp-services/" + _MCP_SERVICE_NAME_PREFIX = "mcp-services/" @@ -2042,7 +2049,7 @@ def list_mcp_services( def build_mcp_service_url(workspace: str, full_name: str) -> str: - return f"{workspace}/ai-gateway/mcp-services/{full_name}" + return f"{workspace}{AIGW_MCP_SERVICES_SEGMENT}{full_name}" def build_skills_mcp_url(workspace: str, locations: list[str]) -> str: diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index ec2eed073..88046969c 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -19,6 +19,7 @@ from ucode.config_io import restore_file from ucode.constants import MCP_CLEANUP_SCOPES, MCP_USER_SCOPE from ucode.databricks import ( + AIGW_MCP_SERVICES_SEGMENT, PermissionDeniedError, apply_pat_environment, build_mcp_proxy_argv, @@ -31,6 +32,7 @@ list_mcp_services, workspace_hostname, ) +from ucode.mcp_connection_login import connection_from_url from ucode.mcp_oauth import ( CLAUDE_CODE_OAUTH_CLIENT_ID, CURSOR_OAUTH_CLIENT_ID, @@ -57,10 +59,6 @@ spinner, ) -# AI Gateway MCP-services endpoints carry this path segment. These are the -# connection-backed services that need a per-user connection login. -AIGW_MCP_SERVICES_PATH = "/ai-gateway/mcp-services/" - # Workspace-relative path fragments for the V2 AI Gateway MCP endpoints, shared by the URL-shape # checks (`_is_app_mcp_server`, `_mcp_server_location`) so the set stays in one place. MCP_EXTERNAL_PATH = "/api/2.0/mcp/external/" @@ -259,7 +257,7 @@ def _oauth_http_client(client: str, workspace: str, *, use_pat: bool) -> str | N This is the single source of truth for that choice, shared by the per-server (:func:`configure_client_mcp_server`) and batched (:func:`_managed_mcp_entry`) paths. Whether it applies to a *specific* server additionally requires a connection-backed mcp-services URL - (``AIGW_MCP_SERVICES_PATH``), which the caller checks per server. It is URL-independent, so + (``AIGW_MCP_SERVICES_SEGMENT``), which the caller checks per server. It is URL-independent, so callers can compute it once per (client, workspace) rather than once per server.""" oauth_client = AGENT_OAUTH_CLIENT.get(client) if oauth_client is not None and not use_pat and oauth_client_available(workspace, oauth_client): @@ -282,7 +280,7 @@ def configure_client_mcp_server( # proxy: non-connection MCPs, the skills registry, PAT auth, agents without a mapped OAuth # client, and workspaces where the mapped client isn't published. http_client = _oauth_http_client(client, workspace, use_pat=use_pat) - if http_client is not None and AIGW_MCP_SERVICES_PATH in url: + if http_client is not None and AIGW_MCP_SERVICES_SEGMENT in url: if client == "claude": removed_scopes = [ scope @@ -753,7 +751,7 @@ def _is_app_mcp_server(server: dict) -> bool: return False stripped = url.rstrip("/") known = ( - AIGW_MCP_SERVICES_PATH, + AIGW_MCP_SERVICES_SEGMENT, MCP_EXTERNAL_PATH, MCP_GENIE_PATH, MCP_VECTOR_SEARCH_PATH, @@ -867,7 +865,7 @@ def _agent_managed_file_entries( if agent == "claude": # A native HTTP+OAuth entry is only valid for a connection-backed mcp-services URL; # anything else stays on the stdio proxy so both delivery paths resolve identically. - if AIGW_MCP_SERVICES_PATH not in url: + if AIGW_MCP_SERVICES_SEGMENT not in url: continue entries[name] = claude.managed_mcp_entry(url) elif agent == "codex": @@ -1199,7 +1197,7 @@ def _managed_mcp_entry( workspace) (or ``None``); the caller computes it once per client so the batch loop doesn't re-probe ``oauth_client_available`` per server. HTTP+OAuth applies here only when that client is set AND this server's URL is a connection-backed mcp-services URL; otherwise the stdio proxy.""" - use_http = http_client is not None and AIGW_MCP_SERVICES_PATH in url + use_http = http_client is not None and AIGW_MCP_SERVICES_SEGMENT in url if client == "claude": if use_http: return claude.managed_mcp_entry(url) @@ -2147,8 +2145,8 @@ def _mcp_server_location(server: dict) -> str: return "skills" url = str(server.get("url") or "") stripped = url.rstrip("/") - if AIGW_MCP_SERVICES_PATH in url: - return url.split(AIGW_MCP_SERVICES_PATH, 1)[1] or "mcp-service" + if AIGW_MCP_SERVICES_SEGMENT in url: + return connection_from_url(url) or "mcp-service" if MCP_EXTERNAL_PATH in url: return f"connection:{stripped.rsplit('/', 1)[-1]}" if MCP_GENIE_PATH in url: diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py index 2da6a5d82..d1ae3c2a7 100644 --- a/src/ucode/mcp_connection_login.py +++ b/src/ucode/mcp_connection_login.py @@ -21,9 +21,7 @@ import subprocess import sys -# AI Gateway MCP service endpoints look like -# ``https:///ai-gateway/mcp-services/..``. -AIGW_MCP_SERVICES_SEGMENT = "/ai-gateway/mcp-services/" +from ucode.databricks import AIGW_MCP_SERVICES_SEGMENT # 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. @@ -137,7 +135,6 @@ def run_connection_login( __all__ = [ - "AIGW_MCP_SERVICES_SEGMENT", "connection_from_url", "run_connection_login", ] diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 4007199d6..b66f8c984 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -3233,7 +3233,7 @@ def test_eligible_agents_use_managed_file_not_user_scope(self, monkeypatch): resolved=[ { "name": "system-ai-github", - "url": f"https://host{mcp.AIGW_MCP_SERVICES_PATH}github", + "url": f"https://host{mcp.AIGW_MCP_SERVICES_SEGMENT}github", } ], applied=applied, @@ -3257,7 +3257,7 @@ def test_claude_non_mcp_services_url_falls_back_while_mcp_services_goes_native( ): # Claude's managed file only takes native mcp-services entries; any other URL falls back to # the user-scope proxy (matching configure_client_mcp_server). Codex proxies both. - github_url = f"https://host{mcp.AIGW_MCP_SERVICES_PATH}github" + github_url = f"https://host{mcp.AIGW_MCP_SERVICES_SEGMENT}github" custom_url = "https://apps.example/custom/mcp" state = {"workspace": WS, "profile": None, "managed_mcp_servers": []} applied: dict = {} @@ -3353,7 +3353,7 @@ def test_failed_managed_write_falls_back_to_user_scope(self, monkeypatch): def test_migration_removes_prior_user_scope_for_managed_file_agents(self, monkeypatch): # Servers a prior configure registered at user scope for an agent now on the managed file are # unregistered from that agent. - sg_url = f"https://host{mcp.AIGW_MCP_SERVICES_PATH}sg" + sg_url = f"https://host{mcp.AIGW_MCP_SERVICES_SEGMENT}sg" previous = [{"name": "sg", "url": sg_url, "clients": ["claude", "codex"]}] state = {"workspace": WS, "profile": None, "managed_mcp_servers": list(previous)} applied: dict = {} From c28a77b3d2d19cc12b2bda7867ba5793782a964a Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Wed, 23 Sep 2026 20:42:08 +0000 Subject: [PATCH 10/10] mcp-proxy: rename _login_or_fail -> _connection_login_or_fail (review nit) Clearer name for the closure that drives the per-connection sign-in on a 401. Co-authored-by: Isaac --- src/ucode/mcp_proxy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 2dcde9428..95ad38d8f 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -154,7 +154,7 @@ def _needs_connection_login(response) -> bool: # the 401, re-running it on every subsequent request would loop browsers. return connection is not None and response.status_code == 401 and not login["attempted"] - def _login_or_fail() -> None: + def _connection_login_or_fail() -> None: login["attempted"] = True ok, detail = run_connection_login(url, workspace, profile=profile) if not ok: @@ -166,7 +166,7 @@ def auth_flow(self, request): response = yield request if not _needs_connection_login(response): return - _login_or_fail() + _connection_login_or_fail() _mint(request) yield request @@ -177,7 +177,7 @@ async def async_auth_flow(self, request): return # Run the blocking browser login off the event loop so the bridge's # other pumps aren't starved while the user completes the sign-in. - await to_thread.run_sync(_login_or_fail) + await to_thread.run_sync(_connection_login_or_fail) _mint(request) yield request