diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py new file mode 100644 index 00000000..bb3c8557 --- /dev/null +++ b/src/ucode/mcp_connection_login.py @@ -0,0 +1,116 @@ +"""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 +import sys + +# 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. + + 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, + "auth", + "login", + "--host", + workspace.rstrip("/"), + "--resource", + resource_url, + ] + 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, + 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, "connection sign-in timed out waiting for the browser flow to complete" + if result.returncode == 0: + return True, "signed in" + return ( + False, + f"connection sign-in did not complete (CLI exited {result.returncode}; see the log above)", + ) + + +__all__ = [ + "AIGW_MCP_SERVICES_SEGMENT", + "connection_from_url", + "run_connection_login", +] diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 0d2a68cb..0a803061 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. @@ -117,18 +118,24 @@ def _build_token_auth(workspace: str, profile: str | None): 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. + + 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() 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. + # 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: @@ -237,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. diff --git a/tests/test_mcp_connection_login.py b/tests/test_mcp_connection_login.py new file mode 100644 index 00000000..4a9a9bfb --- /dev/null +++ b/tests/test_mcp_connection_login.py @@ -0,0 +1,92 @@ +"""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_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", + lambda argv, **kw: seen.update(kw) or subprocess.CompletedProcess(argv, 0), + ) + 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): + raise subprocess.TimeoutExpired(argv, 1) + + monkeypatch.setattr(mcl.subprocess, "run", _run) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok and "timed out" in message + + def test_binary_missing_is_reported(self, monkeypatch): + def _run(argv, **kwargs): + raise OSError("not found") + + monkeypatch.setattr(mcl.subprocess, "run", _run) + ok, message = mcl.run_connection_login(AIGW_URL, WS) + assert not ok and "could not run" in message diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index bf7a3536..9b9e2721 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -128,6 +128,9 @@ def boom(ws, profile): list(auth.auth_flow(httpx.Request("POST", URL))) +CONN_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github" + + class TestPump: def test_forwards_all_messages_in_order(self): async def scenario() -> list[str]: @@ -392,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):