mcp-proxy: drive per-connection login on a 401 (generic, all agents)

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://<ws>/ai-gateway/mcp-services/<catalog>.<schema>.<service>``.
+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",
+]

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
+    # <mcp-url>`, 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.

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

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):