diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 4285cc78..d65b58c6 100644 --- a/src/ucode/mcp.py +++ b/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__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. diff --git a/src/ucode/mcp_oauth.py b/src/ucode/mcp_oauth.py new file mode 100644 index 00000000..9b8b9765 --- /dev/null +++ b/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", +] diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 3a31c71b..d82aa311 100644 --- a/tests/test_mcp.py +++ b/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): diff --git a/tests/test_mcp_oauth.py b/tests/test_mcp_oauth.py new file mode 100644 index 00000000..d0186002 --- /dev/null +++ b/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