From 23a5c33be9fb6aae528a3843ca8ff21cb037eb9a Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Thu, 10 Sep 2026 20:43:04 +0000 Subject: [PATCH 1/2] ug mcp add: register connection-backed MCP services as direct HTTP for OAuth-client-capable agents MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Connection-backed AI Gateway MCP services (…/ai-gateway/mcp-services/…) need a per-user connection login before their tools can be called. Registered as the stdio proxy, an agent only ever sees "connected" (the proxy injects a workspace token) and the login can't be triggered from /mcp. Registered as a direct-HTTP server, the agent does the OAuth itself: /mcp shows "needs authentication" → Authenticate → /oidc → /mcp-service-login → login → connected. /oidc has no dynamic client registration, so the agent must present a pre-registered public client. Only Claude Code's `mcp add` supports pinning one (--client-id); its published claude-code app has the loopback /callback redirect Claude uses. 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. ug mcp add now registers these services as direct HTTP with the agent's OAuth client (AGENT_OAUTH_CLIENT — today just claude → claude-code) when that client is registered on the workspace, probed via a back-channel token-endpoint check (401 invalid_client = absent, cached per workspace+client). Everything else keeps the stdio proxy: non-connection MCPs, the skills registry, PAT auth, agents with no mapped OAuth client, and workspaces where the client isn't published. Co-authored-by: Isaac --- src/ucode/mcp.py | 85 ++++++++++++++++++++++++++++- src/ucode/mcp_oauth.py | 118 ++++++++++++++++++++++++++++++++++++++++ tests/test_mcp.py | 59 ++++++++++++++++++++ tests/test_mcp_oauth.py | 79 +++++++++++++++++++++++++++ 4 files changed, 339 insertions(+), 2 deletions(-) create mode 100644 src/ucode/mcp_oauth.py create mode 100644 tests/test_mcp_oauth.py diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index b40a8ed6..2e8484e6 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -43,6 +43,11 @@ list_mcp_services, 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, @@ -58,6 +63,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 @@ -151,6 +171,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 ( @@ -293,8 +353,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 73b28f39..f12b2a0f 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"] @@ -291,6 +294,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 From 693a49b8866c94d86abdde5fce28248cdc2be23f Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 11 Sep 2026 04:31:03 +0000 Subject: [PATCH 2/2] Extend direct-HTTP MCP OAuth to Cursor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cursor's mcp.json accepts a pre-registered OAuth client via an `auth.CLIENT_ID` object (no dynamic client registration, which /oidc lacks), and drives the login to its fixed loopback `/callback` redirect — the same precondition claude-code satisfies. So for a connection-backed mcp-services endpoint, when the workspace has the `cursor-desktop` OAuth app (probed + cached like claude-code), register a url+auth server in ~/.cursor/mcp.json instead of the stdio proxy, giving Cursor its native connection login. Everything else (no published app, other agents, PAT, non-connection MCPs) still falls back to the generic proxy. Co-authored-by: Isaac --- src/ucode/agents/cursor.py | 31 +++++++++++++++++++++++ src/ucode/mcp.py | 52 +++++++++++++++++++++++--------------- src/ucode/mcp_oauth.py | 9 +++++++ tests/test_agent_cursor.py | 32 +++++++++++++++++++++++ tests/test_mcp.py | 52 ++++++++++++++++++++++++++++++++++---- 5 files changed, 151 insertions(+), 25 deletions(-) diff --git a/src/ucode/agents/cursor.py b/src/ucode/agents/cursor.py index db567684..6ca87ce1 100644 --- a/src/ucode/agents/cursor.py +++ b/src/ucode/agents/cursor.py @@ -53,6 +53,37 @@ def write_mcp_server_config(name: str, argv: list[str]) -> bool: return removed +def build_http_mcp_server_entry(url: str, client_id: str) -> dict: + # Cursor's remote-MCP-with-OAuth schema: a `url` server plus an `auth` object + # naming a pre-registered OAuth client. Cursor drives the OAuth itself (to its + # fixed `http://localhost:8787/callback` redirect) instead of the stdio proxy, + # so the user gets Cursor's native connection login. Scopes are omitted — the + # MCP protected-resource metadata advertises them, and Cursor discovers them + # (the same way Claude Code's `--transport http --client-id` flow does). + return { + "url": url, + "auth": {"CLIENT_ID": client_id}, + } + + +def write_http_mcp_server_config(name: str, url: str, client_id: str) -> bool: + """Add (or replace) one **OAuth HTTP** MCP server entry in ~/.cursor/mcp.json. + + Used for connection-backed AI Gateway services when the workspace has Cursor's + OAuth client published: Cursor authenticates directly rather than going through + the token-injecting stdio proxy. Merges into `mcpServers` like the stdio path; + returns True when an entry with this name was already present.""" + existing = read_json_safe(CURSOR_MCP_CONFIG_PATH) + mcp_servers = existing.get("mcpServers") + if not isinstance(mcp_servers, dict): + mcp_servers = {} + removed = name in mcp_servers + mcp_servers[name] = build_http_mcp_server_entry(url, client_id) + existing["mcpServers"] = mcp_servers + write_json_file(CURSOR_MCP_CONFIG_PATH, existing) + return removed + + def remove_mcp_server_config(name: str) -> bool: """Surgically remove one MCP server entry from ~/.cursor/mcp.json. diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 2e8484e6..718ec165 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -45,6 +45,7 @@ ) from ucode.mcp_oauth import ( CLAUDE_CODE_OAUTH_CLIENT_ID, + CURSOR_OAUTH_CLIENT_ID, MCP_OAUTH_CALLBACK_PORT, oauth_client_available, ) @@ -68,15 +69,22 @@ 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} +# These agents can pin a pre-registered OAuth client and drive the `/oidc` login +# themselves (so `/mcp` shows "needs authentication" / Cursor shows a login), which +# is a much better experience than the stdio proxy for a connection-backed service: +# - Claude Code: `claude mcp add --transport http --client-id `. +# - Cursor: a `url` server with an `auth.CLIENT_ID` in ~/.cursor/mcp.json. +# Both need the app *published on the workspace* (checked per-workspace via +# `oauth_client_available`) and its loopback `/callback` redirect registered on +# `/oidc` — which lacks dynamic client registration, so a pre-registered client is +# required. Agents whose `mcp add` accept only a static bearer, not an OAuth client +# — codex (`--bearer-token-env-var`), gemini (`--header`) — stay on the stdio proxy +# even where their apps exist; add one here (with its registration branch below) +# once its CLI can pin a client. +AGENT_OAUTH_CLIENT = { + "claude": CLAUDE_CODE_OAUTH_CLIENT_ID, + "cursor": CURSOR_OAUTH_CLIENT_ID, +} class _Back: @@ -354,12 +362,12 @@ def configure_client_mcp_server( always_load: bool = False, ) -> list[str]: # 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. + # the agent drives the connection login natively — but only for an agent that can + # pin an OAuth client (AGENT_OAUTH_CLIENT: Claude Code, Cursor) 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 @@ -368,11 +376,15 @@ def configure_client_mcp_server( 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 + if client == "claude": + 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 + if client == "cursor": + removed = cursor.write_http_mcp_server_config(name, url, client_id=oauth_client) + return [MCP_USER_SCOPE] if removed else [] # Every other case registers the `ucode mcp-proxy ...` stdio command; the proxy # forwards to `url` and refreshes the Databricks token itself. Only the diff --git a/src/ucode/mcp_oauth.py b/src/ucode/mcp_oauth.py index 9b8b9765..17d3bbfc 100644 --- a/src/ucode/mcp_oauth.py +++ b/src/ucode/mcp_oauth.py @@ -31,6 +31,14 @@ CLAUDE_CODE_OAUTH_CLIENT_ID = "claude-code" MCP_OAUTH_CALLBACK_PORT = 3118 +# Published public OAuth client Cursor authenticates with for OAuth MCP servers. +# Cursor's `mcp.json` accepts a pre-registered `auth.CLIENT_ID` (no dynamic client +# registration, which `/oidc` lacks) and drives the login to Cursor's fixed loopback +# redirect `http://localhost:8787/callback`. `/oidc` matches loopback redirects by +# path (RFC 8252 §8.4), so the registered `/callback` path is what matters — the +# same requirement `claude-code` satisfies. +CURSOR_OAUTH_CLIENT_ID = "cursor-desktop" + # 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" @@ -113,6 +121,7 @@ def oauth_client_available(workspace: str, client_id: str) -> bool: __all__ = [ "CLAUDE_CODE_OAUTH_CLIENT_ID", + "CURSOR_OAUTH_CLIENT_ID", "MCP_OAUTH_CALLBACK_PORT", "oauth_client_available", ] diff --git a/tests/test_agent_cursor.py b/tests/test_agent_cursor.py index dcdf287c..865ae844 100644 --- a/tests/test_agent_cursor.py +++ b/tests/test_agent_cursor.py @@ -68,6 +68,38 @@ def test_reports_replaced_entry(self, tmp_path, monkeypatch): ) +AIGW_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github" + + +class TestHttpMcpServerEntry: + def test_builds_url_plus_auth_entry_with_pinned_client(self): + entry = cursor.build_http_mcp_server_entry(AIGW_URL, "cursor-desktop") + # A direct OAuth HTTP server: url + auth.CLIENT_ID, no stdio command and no + # static bearer — Cursor drives the OAuth itself. Scopes come from the MCP + # protected-resource metadata, so none are hard-coded here. + assert entry == {"url": AIGW_URL, "auth": {"CLIENT_ID": "cursor-desktop"}} + assert "command" not in entry + assert "headers" not in entry + + def test_write_merges_and_reports_replacement(self, tmp_path, monkeypatch): + config_file = tmp_path / "mcp.json" + monkeypatch.setattr(cursor, "CURSOR_MCP_CONFIG_PATH", config_file) + config_file.write_text( + json.dumps({"mcpServers": {"proxy": {"command": "keep"}}}), + encoding="utf-8", + ) + + removed = cursor.write_http_mcp_server_config("github", AIGW_URL, "cursor-desktop") + + written = json.loads(config_file.read_text()) + assert removed is False + assert written["mcpServers"]["proxy"] == {"command": "keep"} + assert written["mcpServers"]["github"] == { + "url": AIGW_URL, + "auth": {"CLIENT_ID": "cursor-desktop"}, + } + + class TestRemoveMcpServerConfig: def test_removes_without_clobbering_others(self, tmp_path, monkeypatch): config_file = tmp_path / "mcp.json" diff --git a/tests/test_mcp.py b/tests/test_mcp.py index f12b2a0f..dd0832e4 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -339,16 +339,58 @@ def test_claude_aigw_service_with_pat_keeps_proxy(self, monkeypatch): 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]]] = [] + def _capture_cursor(self, monkeypatch, *, cursor_client_available: bool): + http_calls: list[tuple[str, str, str]] = [] + proxy_calls: list[tuple[str, list[str]]] = [] + monkeypatch.setattr( + mcp, "oauth_client_available", lambda ws, client_id: cursor_client_available + ) + monkeypatch.setattr( + mcp.cursor, + "write_http_mcp_server_config", + lambda name, url, client_id: http_calls.append((name, url, client_id)) or False, + ) monkeypatch.setattr( mcp.cursor, "write_mcp_server_config", - lambda name, argv: calls.append((name, argv)) or False, + lambda name, argv: proxy_calls.append((name, argv)) or False, ) + return http_calls, proxy_calls + + def test_cursor_aigw_service_registers_http_when_client_available(self, monkeypatch): + http_calls, proxy_calls = self._capture_cursor(monkeypatch, cursor_client_available=True) + mcp.configure_client_mcp_server("cursor", "github", AIGW_MCP_URL, WS, "p") + assert http_calls == [("github", AIGW_MCP_URL, mcp.CURSOR_OAUTH_CLIENT_ID)] + assert proxy_calls == [] + + def test_cursor_aigw_service_falls_back_to_proxy_without_client(self, monkeypatch): + http_calls, proxy_calls = self._capture_cursor(monkeypatch, cursor_client_available=False) mcp.configure_client_mcp_server("cursor", "github", AIGW_MCP_URL, WS, "p") - assert len(calls) == 1 + assert http_calls == [] + assert len(proxy_calls) == 1 # workspaces without cursor-desktop keep the stdio proxy + + def test_cursor_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_cursor(monkeypatch, cursor_client_available=True) + mcp.configure_client_mcp_server("cursor", "github", AIGW_MCP_URL, WS, "p", use_pat=True) + assert http_calls == [] + assert len(proxy_calls) == 1 + + def test_oauthless_client_keeps_proxy_and_skips_probe(self, monkeypatch): + # An agent with no mapped OAuth client (codex) always proxies, and must not + # even probe /oidc — there's nothing it could pin. + probed: list[str] = [] + proxy_calls: list[tuple[str, list[str]]] = [] + monkeypatch.setattr( + mcp, "oauth_client_available", lambda ws, client_id: probed.append(client_id) or True + ) + monkeypatch.setattr( + mcp, "add_codex_mcp_server", lambda name, argv: proxy_calls.append((name, argv)) + ) + monkeypatch.setattr(mcp, "remove_codex_mcp_server", lambda name: False) + mcp.configure_client_mcp_server("codex", "github", AIGW_MCP_URL, WS, "p") + assert len(proxy_calls) == 1 + assert probed == [] # AGENT_OAUTH_CLIENT has no entry for codex → no probe class TestMcpPicker: