ug mcp add: register connection-backed MCP services as direct HTTP for OAuth-client-capable agents

src/ucode/mcp.py

@@ -47,6 +47,11 @@
     list_vector_search_catalog_schemas,
     workspace_hostname,
 )
+from ucode.mcp_oauth import (
+    CLAUDE_CODE_OAUTH_CLIENT_ID,
+    MCP_OAUTH_CALLBACK_PORT,
+    oauth_client_available,
+)
 from ucode.state import load_full_state, load_state, save_state
 from ucode.ui import (
     console,
@@ -62,6 +67,21 @@
 MCP_CLEANUP_SCOPES = ("local", "project", MCP_USER_SCOPE)
 MCP_PICKER_VISIBLE_ROWS = 10
 
+# 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/"
+
+# Per-agent published OAuth app used for the direct-HTTP MCP connection login.
+# ONLY Claude Code's `mcp add` supports pinning an OAuth client (`--client-id`),
+# so only it can drive the `/oidc` login for a connection-backed service and show
+# `/mcp` "needs authentication". The other agents' `mcp add` accept only a static
+# bearer, not an OAuth client — codex (`--bearer-token-env-var`), gemini
+# (`--header`), cursor (config) — so despite their published apps (codex-cli,
+# cursor-desktop, …) they can't drive this flow and stay on the stdio proxy. Add
+# an agent here — with its own `add_<agent>_http_mcp_server` — once its CLI gains
+# `--client-id` support.
+AGENT_OAUTH_CLIENT = {"claude": CLAUDE_CODE_OAUTH_CLIENT_ID}
+
 
 class _Back:
     """Sentinel type: a wizard step returns the `_BACK` instance when the user
@@ -162,6 +182,46 @@ def add_claude_mcp_server(
         raise RuntimeError(f"Failed to add MCP server '{name}' via claude CLI.") from exc
 
 
+def add_claude_http_mcp_server(
+    name: str,
+    url: str,
+    scope: str = MCP_USER_SCOPE,
+    *,
+    client_id: str = CLAUDE_CODE_OAUTH_CLIENT_ID,
+    callback_port: int = MCP_OAUTH_CALLBACK_PORT,
+) -> None:
+    """Register a Databricks MCP endpoint as a **direct HTTP** server so Claude
+    Code is the OAuth client and drives the RFC 8707 connection login itself.
+
+    Unlike the stdio proxy (which injects a plain workspace token and hides the
+    per-user connection state), a direct HTTP server lets Claude Code do MCP OAuth
+    against ``/oidc`` with the ``resource`` indicator: on a missing/expired
+    connection credential, ``/mcp`` shows "needs authentication" and Authenticate
+    runs the login (``/oidc`` -> ``/mcp-service-login``). ``client_id`` is the
+    published ``claude-code`` app (it has the loopback ``/callback`` redirect
+    registered); the callback port is arbitrary because ``/oidc`` ignores the port
+    for loopback redirects."""
+    cmd = [
+        "claude",
+        "mcp",
+        "add",
+        "--transport",
+        "http",
+        "-s",
+        scope,
+        "--client-id",
+        client_id,
+        "--callback-port",
+        str(callback_port),
+        name,
+        url,
+    ]
+    try:
+        subprocess.run(cmd, check=True, capture_output=True, text=True, timeout=30)
+    except subprocess.CalledProcessError as exc:
+        raise RuntimeError(f"Failed to add HTTP MCP server '{name}' via claude CLI.") from exc
+
+
 def _is_missing_mcp_server_output(output: str) -> bool:
     normalized = output.lower()
     return (
@@ -304,8 +364,29 @@ def configure_client_mcp_server(
     use_pat: bool = False,
     always_load: bool = False,
 ) -> list[str]:
-    # Every client registers the same `ucode mcp-proxy ...` stdio command; the
-    # proxy forwards to `url` and refreshes the Databricks token itself. Only the
+    # Connection-backed AI Gateway MCP services register as a direct HTTP server so
+    # the agent drives the connection login natively — but only for an agent whose
+    # CLI can pin an OAuth client (AGENT_OAUTH_CLIENT: today just Claude) and only
+    # when that client is registered on the workspace. Everything else keeps the
+    # stdio proxy: non-connection MCPs, the skills registry (`always_load`), PAT
+    # auth (no interactive OAuth), agents without a mapped OAuth client, and
+    # workspaces where the mapped client isn't published.
+    oauth_client = AGENT_OAUTH_CLIENT.get(client)
+    if (
+        oauth_client is not None
+        and AIGW_MCP_SERVICES_PATH in url
+        and not always_load
+        and not use_pat
+        and oauth_client_available(workspace, oauth_client)
+    ):
+        removed_scopes = [
+            scope for scope in MCP_CLEANUP_SCOPES if remove_claude_mcp_server(name, scope)
+        ]
+        add_claude_http_mcp_server(name, url, client_id=oauth_client)
+        return removed_scopes
+
+    # Every other case registers the `ucode mcp-proxy ...` stdio command; the proxy
+    # forwards to `url` and refreshes the Databricks token itself. Only the
     # per-client registration syntax differs. `always_load` (skills registry) is
     # a Claude-only hint to load the server's tools at session start; other
     # clients don't support it and ignore it.

src/ucode/mcp_oauth.py

@@ -0,0 +1,118 @@
+"""Discovery of the OAuth client used for direct-HTTP MCP registration.
+
+A connection-backed AI Gateway MCP service (e.g. ``system.ai.github``) needs a
+per-user connection login before its tools can be called. When a coding agent
+registers the service as a **direct HTTP** MCP server, the agent itself drives
+that login via OAuth against the workspace ``/oidc`` — and `/oidc` has no dynamic
+client registration, so the agent must present a **pre-registered public client**.
+
+``claude-code`` is that published client for Claude Code: it has the loopback
+``/callback`` redirect Claude Code uses registered (``databricks-cli`` does not,
+so Claude's direct-HTTP OAuth is rejected against it). Where a workspace has
+``claude-code``, ucode registers these services as direct HTTP so ``/mcp`` shows
+"needs authentication" and Authenticate drives the login natively. Not every
+workspace has it yet, so we probe — and cache the answer per workspace.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import urllib.error
+import urllib.request
+from urllib.parse import urlencode
+
+from ucode.config_io import APP_DIR
+
+# Published public OAuth client Claude Code authenticates with. Any loopback
+# callback port works — `/oidc` strips the port when matching loopback redirects
+# (RFC 8252 §8.4) — but the redirect *path* (`/callback`) must be registered,
+# which this client has and `databricks-cli` does not.
+CLAUDE_CODE_OAUTH_CLIENT_ID = "claude-code"
+MCP_OAUTH_CALLBACK_PORT = 3118
+
+# Published apps rarely appear/disappear, so a per-workspace probe result is good
+# for a while; delete the cache file to force a re-probe.
+_CACHE_PATH = APP_DIR / "oauth_client_cache.json"
+_CACHE_TTL_SECONDS = 7 * 24 * 3600
+
+
+def _probe_oauth_client(workspace: str, client_id: str) -> bool:
+    """True if ``client_id`` is a registered OAuth app on the workspace's ``/oidc``.
+
+    Back-channel and unauthenticated: POST a throwaway ``authorization_code`` grant
+    to ``/oidc/v1/token``. An **unknown** client fails client authentication (HTTP
+    401 ``invalid_client``); a **known** client gets past that to a grant error
+    (HTTP 400 ``invalid_request`` — "Invalid authorization code"). We only read
+    which of the two it is; the dummy code always fails, harmlessly. This is the
+    only user-level check available — the authorize endpoint redirects to SSO
+    before validating the client, and the published-app API needs account-admin."""
+    body = urlencode(
+        {
+            "grant_type": "authorization_code",
+            "code": "ucode-probe-not-a-real-code",
+            "redirect_uri": f"http://localhost:{MCP_OAUTH_CALLBACK_PORT}/callback",
+            "client_id": client_id,
+            # `/oidc` rejects a code_verifier shorter than 43 chars *before* it
+            # validates the client, so pad past that to reach the client check.
+            "code_verifier": "u" * 43,
+        }
+    ).encode("ascii")
+    request = urllib.request.Request(
+        f"{workspace.rstrip('/')}/oidc/v1/token",
+        data=body,
+        method="POST",
+        headers={"Content-Type": "application/x-www-form-urlencoded"},
+    )
+    try:
+        urllib.request.urlopen(request, timeout=10)  # noqa: S310 - fixed https workspace URL
+        return True  # a 2xx for a dummy code is unexpected, but means the client is valid
+    except urllib.error.HTTPError as exc:
+        # 401 invalid_client => not registered; any other error (400 for the bad
+        # code) => the client IS registered.
+        return exc.code != 401
+    except OSError:
+        # Network failure: don't claim availability — the caller falls back to the
+        # stdio proxy, which is always safe.
+        return False
+
+
+def _read_cache() -> dict:
+    try:
+        return json.loads(_CACHE_PATH.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return {}
+
+
+def _write_cache(cache: dict) -> None:
+    try:
+        _CACHE_PATH.write_text(json.dumps(cache, indent=2), encoding="utf-8")
+    except OSError:
+        pass  # best-effort: a write failure just means we re-probe next time
+
+
+def oauth_client_available(workspace: str, client_id: str) -> bool:
+    """Whether the workspace has ``client_id`` as a registered OAuth app, cached
+    per (workspace, client_id).
+
+    Cached in ``APP_DIR`` with a weekly TTL so ``ug mcp add`` doesn't probe every
+    run. Negative results are cached too (workspaces that don't have it yet)."""
+    ws = workspace.rstrip("/")
+    cache = _read_cache()
+    entry = cache.get(ws, {}).get(client_id)
+    if entry and (time.time() - entry.get("checked_at", 0)) < _CACHE_TTL_SECONDS:
+        return bool(entry.get("available"))
+    available = _probe_oauth_client(ws, client_id)
+    cache.setdefault(ws, {})[client_id] = {
+        "available": available,
+        "checked_at": time.time(),
+    }
+    _write_cache(cache)
+    return available
+
+
+__all__ = [
+    "CLAUDE_CODE_OAUTH_CLIENT_ID",
+    "MCP_OAUTH_CALLBACK_PORT",
+    "oauth_client_available",
+]

tests/test_mcp.py

@@ -36,6 +36,9 @@ def test_no_changes_falls_back_to_saved(self):
 # The proxy argv every client registers as a stdio command. The leading element
 # is the resolved `ucode` binary path, so tests assert the tail (the stable part).
 GH_URL = f"{WS}/api/2.0/mcp/external/github"
+# A connection-backed AI Gateway MCP service (3-part FQN) — the URL form that
+# registers as direct HTTP for Claude when the claude-code client is available.
+AIGW_MCP_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github"
 PROXY_TAIL = ["mcp-proxy", "--url", GH_URL, "--host", WS, "--profile", "p"]
 
 
@@ -315,6 +318,62 @@ def test_configures_copilot_with_proxy_argv(self, monkeypatch):
         # Copilot receives the proxy argv, not a URL/bearer entry.
         assert calls == [("github", _proxy_argv())]
 
+    def _capture_claude(self, monkeypatch, *, claude_code_available: bool):
+        http_calls: list[tuple[str, str]] = []
+        proxy_calls: list[tuple[str, list[str]]] = []
+        monkeypatch.setattr(
+            mcp, "oauth_client_available", lambda ws, client_id: claude_code_available
+        )
+        monkeypatch.setattr(mcp, "remove_claude_mcp_server", lambda name, scope: False)
+        monkeypatch.setattr(
+            mcp,
+            "add_claude_http_mcp_server",
+            lambda name, url, **kw: http_calls.append((name, url)),
+        )
+        monkeypatch.setattr(
+            mcp,
+            "add_claude_mcp_server",
+            lambda name, argv, scope=mcp.MCP_USER_SCOPE, **kw: proxy_calls.append((name, argv)),
+        )
+        return http_calls, proxy_calls
+
+    def test_claude_aigw_service_registers_http_when_client_available(self, monkeypatch):
+        http_calls, proxy_calls = self._capture_claude(monkeypatch, claude_code_available=True)
+        mcp.configure_client_mcp_server("claude", "github", AIGW_MCP_URL, WS, "p")
+        assert http_calls == [("github", AIGW_MCP_URL)]
+        assert proxy_calls == []
+
+    def test_claude_aigw_service_falls_back_to_proxy_without_client(self, monkeypatch):
+        http_calls, proxy_calls = self._capture_claude(monkeypatch, claude_code_available=False)
+        mcp.configure_client_mcp_server("claude", "github", AIGW_MCP_URL, WS, "p")
+        assert http_calls == []
+        assert len(proxy_calls) == 1  # workspaces without claude-code keep the stdio proxy
+
+    def test_claude_non_aigw_url_keeps_proxy(self, monkeypatch):
+        # External/genie/vector-search/functions MCPs have no per-user connection login.
+        http_calls, proxy_calls = self._capture_claude(monkeypatch, claude_code_available=True)
+        mcp.configure_client_mcp_server("claude", "github", GH_URL, WS, "p")
+        assert http_calls == []
+        assert len(proxy_calls) == 1
+
+    def test_claude_aigw_service_with_pat_keeps_proxy(self, monkeypatch):
+        # PAT auth has no interactive OAuth, so it can't use the HTTP login path.
+        http_calls, proxy_calls = self._capture_claude(monkeypatch, claude_code_available=True)
+        mcp.configure_client_mcp_server("claude", "github", AIGW_MCP_URL, WS, "p", use_pat=True)
+        assert http_calls == []
+        assert len(proxy_calls) == 1
+
+    def test_non_claude_client_keeps_proxy_for_aigw_service(self, monkeypatch):
+        # Only Claude's HTTP+OAuth registration is wired; other clients proxy.
+        calls: list[tuple[str, list[str]]] = []
+        monkeypatch.setattr(
+            mcp.cursor,
+            "write_mcp_server_config",
+            lambda name, argv: calls.append((name, argv)) or False,
+        )
+        mcp.configure_client_mcp_server("cursor", "github", AIGW_MCP_URL, WS, "p")
+        assert len(calls) == 1
+
 
 class TestMcpPicker:
     def test_prompt_uses_scrolling_checkbox_selector(self, monkeypatch):

tests/test_mcp_oauth.py

@@ -0,0 +1,79 @@
+"""Tests for OAuth-client discovery + per-workspace caching (mcp_oauth).
+
+Network-free: the token-endpoint probe is monkeypatched; these cover the
+known/unknown/error classification and the cache behaviour.
+"""
+
+from __future__ import annotations
+
+import urllib.error
+
+from ucode import mcp_oauth
+
+WS = "https://ws.staging.cloud.databricks.com"
+
+
+def _http_error(code: int) -> urllib.error.HTTPError:
+    return urllib.error.HTTPError(url="x", code=code, msg="m", hdrs=None, fp=None)
+
+
+class TestProbeOauthClient:
+    def test_known_client_400_is_available(self, monkeypatch):
+        # 400 invalid_request ("Invalid authorization code") => client is registered.
+        monkeypatch.setattr(
+            mcp_oauth.urllib.request,
+            "urlopen",
+            lambda req, timeout=0: (_ for _ in ()).throw(_http_error(400)),
+        )
+        assert mcp_oauth._probe_oauth_client(WS, "claude-code") is True
+
+    def test_unknown_client_401_is_absent(self, monkeypatch):
+        # 401 invalid_client => client not registered on this workspace.
+        monkeypatch.setattr(
+            mcp_oauth.urllib.request,
+            "urlopen",
+            lambda req, timeout=0: (_ for _ in ()).throw(_http_error(401)),
+        )
+        assert mcp_oauth._probe_oauth_client(WS, "claude-code") is False
+
+    def test_network_error_is_absent(self, monkeypatch):
+        # A network failure must not claim availability — caller falls back to the proxy.
+        monkeypatch.setattr(
+            mcp_oauth.urllib.request,
+            "urlopen",
+            lambda req, timeout=0: (_ for _ in ()).throw(OSError("down")),
+        )
+        assert mcp_oauth._probe_oauth_client(WS, "claude-code") is False
+
+
+class TestOauthClientAvailableCache:
+    def test_probes_once_then_serves_from_cache(self, monkeypatch, tmp_path):
+        monkeypatch.setattr(mcp_oauth, "_CACHE_PATH", tmp_path / "cache.json")
+        calls: list[tuple[str, str]] = []
+        monkeypatch.setattr(
+            mcp_oauth, "_probe_oauth_client", lambda ws, cid: calls.append((ws, cid)) or True
+        )
+        assert mcp_oauth.oauth_client_available(WS, "claude-code") is True
+        assert mcp_oauth.oauth_client_available(WS, "claude-code") is True
+        assert len(calls) == 1  # second call served from cache, no re-probe
+
+    def test_negative_result_is_cached_too(self, monkeypatch, tmp_path):
+        monkeypatch.setattr(mcp_oauth, "_CACHE_PATH", tmp_path / "cache.json")
+        calls: list[int] = []
+        monkeypatch.setattr(
+            mcp_oauth, "_probe_oauth_client", lambda ws, cid: calls.append(1) or False
+        )
+        assert mcp_oauth.oauth_client_available(WS, "claude-code") is False
+        assert mcp_oauth.oauth_client_available(WS, "claude-code") is False
+        assert len(calls) == 1
+
+    def test_expired_entry_reprobes(self, monkeypatch, tmp_path):
+        monkeypatch.setattr(mcp_oauth, "_CACHE_PATH", tmp_path / "cache.json")
+        monkeypatch.setattr(mcp_oauth, "_CACHE_TTL_SECONDS", -1)  # every entry immediately stale
+        calls: list[int] = []
+        monkeypatch.setattr(
+            mcp_oauth, "_probe_oauth_client", lambda ws, cid: calls.append(1) or True
+        )
+        mcp_oauth.oauth_client_available(WS, "claude-code")
+        mcp_oauth.oauth_client_available(WS, "claude-code")
+        assert len(calls) == 2  # re-probed because the cached entry expired