From 8567f43544ea543c2ae7326a2c26b36fc615e795 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Mon, 14 Sep 2026 23:06:51 +0000 Subject: [PATCH 1/2] cursor: dedup mcp.json writers into a shared _upsert_mcp_server helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit write_mcp_server_config (stdio) and write_http_mcp_server_config (OAuth HTTP) were identical except the entry builder — both now delegate to _upsert_mcp_server, so the read/merge/write is defined once. Behavior-identical. Co-authored-by: Isaac --- src/ucode/agents/cursor.py | 36 ++++++++++++++---------------------- 1 file changed, 14 insertions(+), 22 deletions(-) diff --git a/src/ucode/agents/cursor.py b/src/ucode/agents/cursor.py index 6ca87ce17..f5e6739b0 100644 --- a/src/ucode/agents/cursor.py +++ b/src/ucode/agents/cursor.py @@ -36,23 +36,26 @@ def build_mcp_server_entry(argv: list[str]) -> dict: } -def write_mcp_server_config(name: str, argv: list[str]) -> bool: - """Add (or replace) one MCP server entry in ~/.cursor/mcp.json. - - Merges into the existing `mcpServers` map so unrelated entries the user - already configured survive. Returns True when an entry with this name was - already present (i.e. this was a replacement).""" +def _upsert_mcp_server(name: str, entry: dict) -> bool: + """Add (or replace) one entry in ~/.cursor/mcp.json's `mcpServers`, merging so + unrelated entries the user already configured survive. Returns True when an + entry with this name was already present (i.e. this was a replacement).""" 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_mcp_server_entry(argv) + mcp_servers[name] = entry existing["mcpServers"] = mcp_servers write_json_file(CURSOR_MCP_CONFIG_PATH, existing) return removed +def write_mcp_server_config(name: str, argv: list[str]) -> bool: + """Add (or replace) a stdio (`ucode mcp-proxy`) MCP server in ~/.cursor/mcp.json.""" + return _upsert_mcp_server(name, build_mcp_server_entry(argv)) + + 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 @@ -67,21 +70,10 @@ def build_http_mcp_server_entry(url: str, client_id: str) -> dict: 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 + """Add (or replace) an **OAuth HTTP** MCP server (url + pre-registered client) in + ~/.cursor/mcp.json, so Cursor drives the connection login itself instead of the + token-injecting stdio proxy.""" + return _upsert_mcp_server(name, build_http_mcp_server_entry(url, client_id)) def remove_mcp_server_config(name: str) -> bool: From 5e4d5358463e7f1ab2c1f394cc56b99a364c429e Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Mon, 14 Sep 2026 23:42:28 +0000 Subject: [PATCH 2/2] mcp: move Claude MCP-registration helpers into agents/claude.py The claude-specific `claude mcp add/add-json/remove` wrappers (add_claude_mcp_server, add_claude_http_mcp_server, remove_claude_mcp_server) lived in ucode.mcp even though every other agent's registration lives in its own agents/.py module. Move them next to the Claude agent so mcp.py dispatches uniformly (claude.X, like cursor.X / opencode.X). - Shared MCP scope constants (MCP_USER_SCOPE, MCP_CLEANUP_SCOPES) move to the leaf ucode.constants so both mcp and agents.claude import them with no cycle; claude.py's default args can then reference them at module load. - _is_missing_mcp_server_output stays in mcp.py (shared by the codex/gemini removers); claude.remove_claude_mcp_server imports it lazily to avoid the mcp -> agents.claude load cycle (the same pattern claude.py already used). - claude.py's web_search register/unregister now call the local helpers directly instead of lazily importing them from mcp. - Relocate the corresponding unit tests to tests/test_agent_claude.py. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 112 +++++++++++++++++++++-- src/ucode/constants.py | 6 ++ src/ucode/mcp.py | 116 +++--------------------- tests/test_agent_claude.py | 178 +++++++++++++++++++++++++++++++------ tests/test_mcp.py | 137 ++-------------------------- 5 files changed, 277 insertions(+), 272 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index b5cc72071..175594c60 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -26,6 +26,8 @@ ) from ucode.constants import ( LOOPBACK_HOST, + MCP_CLEANUP_SCOPES, + MCP_USER_SCOPE, MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER, ) @@ -49,6 +51,7 @@ reconcile_managed_file, revert_managed_file, ) +from ucode.mcp_oauth import CLAUDE_CODE_OAUTH_CLIENT_ID, MCP_OAUTH_CALLBACK_PORT from ucode.smart_routing import v2 as smart_routing_v2 from ucode.smart_routing.claude_hooks import ( remove_smart_routing_hooks, @@ -502,6 +505,106 @@ def _enforce_model_default_hierarchy( return selected_default_model +def add_claude_mcp_server( + name: str, + server: list[str] | dict, + scope: str = MCP_USER_SCOPE, + *, + always_load: bool = False, +) -> None: + # Three registration shapes share this helper. The plain proxy path passes an + # argv list (`ucode mcp-proxy ...`), registered via `claude mcp add ... -- ` + # where `--` fences the proxy's own flags off from claude's parser. The + # web_search server passes a full stdio entry dict with its own env, which only + # `add-json` can express — so a dict routes there. Finally, `always_load` (the + # skills registry) needs `alwaysLoad: true`, which plain `mcp add` can't set, so + # build a stdio entry dict and route it to add-json too. + if isinstance(server, dict): + cmd = ["claude", "mcp", "add-json", name, json.dumps(server), "-s", scope] + elif always_load: + entry = { + "type": "stdio", + "command": server[0], + "args": list(server[1:]), + "alwaysLoad": True, + } + cmd = ["claude", "mcp", "add-json", name, json.dumps(entry), "-s", scope] + else: + cmd = ["claude", "mcp", "add", name, "-s", scope, "--", *server] + try: + subprocess.run( + cmd, + check=True, + capture_output=True, + text=True, + timeout=30, + ) + except subprocess.CalledProcessError as exc: + 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 remove_claude_mcp_server(name: str, scope: str) -> bool: + # Imported lazily: `_is_missing_mcp_server_output` is a shared CLI-output matcher + # in ucode.mcp (used by the codex/gemini removers too), and ucode.mcp imports + # this module at load time — a function-level import avoids that cycle. + from ucode.mcp import _is_missing_mcp_server_output + + try: + subprocess.run( + ["claude", "mcp", "remove", name, "-s", scope], + check=True, + capture_output=True, + text=True, + timeout=30, + ) + return True + except subprocess.CalledProcessError as exc: + output = f"{exc.stderr or ''}\n{exc.stdout or ''}" + if _is_missing_mcp_server_output(output): + return False + raise RuntimeError(f"Failed to remove MCP server '{name}' via claude CLI.") from exc + + def _register_web_search_mcp(workspace: str, search_model: str, profile: str | None = None) -> bool: """Register (or replace) the web_search MCP server in Claude Code's user scope via `claude mcp add-json`. Removes any prior entry first so re-runs @@ -510,13 +613,6 @@ def _register_web_search_mcp(workspace: str, search_model: str, profile: str | N Returns True if registration succeeded. Failures are non-blocking: we warn and return False so the rest of `ucode claude` setup can complete. """ - # Imported lazily to avoid a circular import via ucode.mcp -> ucode.agents. - from ucode.mcp import ( - MCP_CLEANUP_SCOPES, - add_claude_mcp_server, - remove_claude_mcp_server, - ) - for scope in MCP_CLEANUP_SCOPES: try: remove_claude_mcp_server(WEB_SEARCH_MCP_NAME, scope) @@ -548,8 +644,6 @@ def _web_search_mcp_is_current(state: dict, entry: dict) -> bool: def _unregister_web_search_mcp() -> None: """Remove the web_search MCP server from all scopes. Used by revert.""" - from ucode.mcp import MCP_CLEANUP_SCOPES, remove_claude_mcp_server - for scope in MCP_CLEANUP_SCOPES: try: remove_claude_mcp_server(WEB_SEARCH_MCP_NAME, scope) diff --git a/src/ucode/constants.py b/src/ucode/constants.py index 8c08838a0..c129da93e 100644 --- a/src/ucode/constants.py +++ b/src/ucode/constants.py @@ -5,3 +5,9 @@ MODEL_PROVIDER_SERVICE_HEADER = "Databricks-Model-Provider-Service" MODEL_SERVICE_PARENT_SCHEMA_HEADER = "Databricks-Model-Service-Parent-Schema" + +# MCP server registration scopes. Claude Code supports local/project/user; the +# other CLIs only take the user-scope name. Kept here (a leaf module) so both +# `ucode.mcp` and `ucode.agents.claude` can import them without an import cycle. +MCP_USER_SCOPE = "user" +MCP_CLEANUP_SCOPES = ("local", "project", MCP_USER_SCOPE) diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 9a69b5182..ae0256794 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -2,7 +2,6 @@ from __future__ import annotations -import json import os import shutil import string @@ -28,8 +27,9 @@ from questionary.question import Question from questionary.styles import merge_styles_default -from ucode.agents import copilot, cursor, gemini, opencode +from ucode.agents import claude, copilot, cursor, gemini, opencode from ucode.config_io import restore_file +from ucode.constants import MCP_CLEANUP_SCOPES, MCP_USER_SCOPE from ucode.databricks import ( PermissionDeniedError, apply_pat_environment, @@ -46,7 +46,6 @@ from ucode.mcp_oauth import ( CLAUDE_CODE_OAUTH_CLIENT_ID, CURSOR_OAUTH_CLIENT_ID, - MCP_OAUTH_CALLBACK_PORT, oauth_client_available, ) from ucode.state import load_full_state, load_state, save_state @@ -60,8 +59,6 @@ spinner, ) -MCP_USER_SCOPE = "user" -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 @@ -142,84 +139,6 @@ class _Back: MCP_ADD_PREFIX = "add:" -def add_claude_mcp_server( - name: str, - server: list[str] | dict, - scope: str = MCP_USER_SCOPE, - *, - always_load: bool = False, -) -> None: - # Three registration shapes share this helper. The plain proxy path passes an - # argv list (`ucode mcp-proxy ...`), registered via `claude mcp add ... -- ` - # where `--` fences the proxy's own flags off from claude's parser. The - # web_search server (agents/claude.py) passes a full stdio entry dict with its - # own env, which only `add-json` can express — so a dict routes there. Finally, - # `always_load` (the skills registry) needs `alwaysLoad: true`, which plain - # `mcp add` can't set, so build a stdio entry dict and route it to add-json too. - if isinstance(server, dict): - cmd = ["claude", "mcp", "add-json", name, json.dumps(server), "-s", scope] - elif always_load: - entry = { - "type": "stdio", - "command": server[0], - "args": list(server[1:]), - "alwaysLoad": True, - } - cmd = ["claude", "mcp", "add-json", name, json.dumps(entry), "-s", scope] - else: - cmd = ["claude", "mcp", "add", name, "-s", scope, "--", *server] - try: - subprocess.run( - cmd, - check=True, - capture_output=True, - text=True, - timeout=30, - ) - except subprocess.CalledProcessError as exc: - 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 ( @@ -230,23 +149,6 @@ def _is_missing_mcp_server_output(output: str) -> bool: ) -def remove_claude_mcp_server(name: str, scope: str) -> bool: - try: - subprocess.run( - ["claude", "mcp", "remove", name, "-s", scope], - check=True, - capture_output=True, - text=True, - timeout=30, - ) - return True - except subprocess.CalledProcessError as exc: - output = f"{exc.stderr or ''}\n{exc.stdout or ''}" - if _is_missing_mcp_server_output(output): - return False - raise RuntimeError(f"Failed to remove MCP server '{name}' via claude CLI.") from exc - - def add_codex_mcp_server(name: str, argv: list[str]) -> None: # `--` fences the proxy argv off from codex's own flag parser, registering # it as a stdio server (codex spawns the command and speaks MCP over it). @@ -378,9 +280,11 @@ def configure_client_mcp_server( ): if client == "claude": removed_scopes = [ - scope for scope in MCP_CLEANUP_SCOPES if remove_claude_mcp_server(name, scope) + scope + for scope in MCP_CLEANUP_SCOPES + if claude.remove_claude_mcp_server(name, scope) ] - add_claude_http_mcp_server(name, url, client_id=oauth_client) + claude.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) @@ -394,9 +298,9 @@ def configure_client_mcp_server( argv = build_mcp_proxy_argv(url, workspace, profile, use_pat=use_pat) if client == "claude": removed_scopes = [ - scope for scope in MCP_CLEANUP_SCOPES if remove_claude_mcp_server(name, scope) + scope for scope in MCP_CLEANUP_SCOPES if claude.remove_claude_mcp_server(name, scope) ] - add_claude_mcp_server(name, argv, MCP_USER_SCOPE, always_load=always_load) + claude.add_claude_mcp_server(name, argv, MCP_USER_SCOPE, always_load=always_load) return removed_scopes if client == "codex": removed = remove_codex_mcp_server(name) @@ -420,7 +324,9 @@ def configure_client_mcp_server( def remove_client_mcp_server(client: str, name: str) -> list[str]: if client == "claude": - return [scope for scope in MCP_CLEANUP_SCOPES if remove_claude_mcp_server(name, scope)] + return [ + scope for scope in MCP_CLEANUP_SCOPES if claude.remove_claude_mcp_server(name, scope) + ] if client == "codex": return [MCP_USER_SCOPE] if remove_codex_mcp_server(name) else [] if client == "gemini": diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 4a8a0ac14..201d903ef 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -5,8 +5,9 @@ import json import os import shlex +import subprocess from pathlib import Path -from unittest.mock import Mock +from unittest.mock import MagicMock, Mock import pytest @@ -16,6 +17,15 @@ from ucode.state import MANAGED_OVERLAY_KEY WS = "https://example.databricks.com" +# A connection MCP proxy argv, used by the Claude MCP-registration helper tests. +# The leading element is the resolved `ucode` binary path, so tests assert the tail. +GH_URL = f"{WS}/api/2.0/mcp/external/github" + + +def _proxy_argv() -> list[str]: + from ucode.databricks import build_mcp_proxy_argv + + return build_mcp_proxy_argv(GH_URL, WS, "p") @pytest.fixture(autouse=True) @@ -1067,6 +1077,130 @@ def deny_managed_write(*args, **kwargs): ) +class TestAddClaudeMcpServer: + def test_registers_stdio_proxy_command(self, monkeypatch): + calls: list[dict] = [] + + def fake_run(args, **kwargs): + calls.append({"args": args, "kwargs": kwargs}) + return MagicMock(returncode=0) + + monkeypatch.setattr(claude.subprocess, "run", fake_run) + + claude.add_claude_mcp_server("github", _proxy_argv()) + + args = calls[0]["args"] + assert args[:4] == ["claude", "mcp", "add", "github"] + assert args[4:6] == ["-s", "user"] + # `--` fences the proxy argv; everything after it is the stdio command. + assert args[6] == "--" + assert args[7:] == _proxy_argv() + + def test_always_load_routes_through_add_json_stdio_entry(self, monkeypatch): + # The skills registry needs `alwaysLoad: true`, which plain `mcp add` + # can't set — so the proxy argv is wrapped in a stdio entry dict and + # registered via add-json instead. + calls: list[dict] = [] + + def fake_run(args, **kwargs): + calls.append({"args": args, "kwargs": kwargs}) + return MagicMock(returncode=0) + + monkeypatch.setattr(claude.subprocess, "run", fake_run) + + claude.add_claude_mcp_server("skills", _proxy_argv(), always_load=True) + + args = calls[0]["args"] + assert args[:4] == ["claude", "mcp", "add-json", "skills"] + entry = json.loads(args[4]) + assert entry == { + "type": "stdio", + "command": _proxy_argv()[0], + "args": _proxy_argv()[1:], + "alwaysLoad": True, + } + assert args[5:] == ["-s", "user"] + + def test_dict_entry_routes_through_add_json(self, monkeypatch): + # The web_search server registers a full stdio entry dict with its own + # env, which only `add-json` can express — a dict must route there rather + # than through the proxy `mcp add -- ` path. + calls: list[dict] = [] + + def fake_run(args, **kwargs): + calls.append({"args": args, "kwargs": kwargs}) + return MagicMock(returncode=0) + + monkeypatch.setattr(claude.subprocess, "run", fake_run) + + entry = {"type": "stdio", "command": "ucode", "args": ["mcp", "web-search"]} + claude.add_claude_mcp_server("web_search", entry) + + args = calls[0]["args"] + assert args[:4] == ["claude", "mcp", "add-json", "web_search"] + assert json.loads(args[4]) == entry + assert args[5:] == ["-s", "user"] + + +class TestRemoveClaudeMcpServer: + def test_returns_true_when_server_removed(self, monkeypatch): + calls: list[list[str]] = [] + + def fake_run(args, **kwargs): + calls.append(args) + return MagicMock(returncode=0) + + monkeypatch.setattr(claude.subprocess, "run", fake_run) + + assert claude.remove_claude_mcp_server("github", "user") is True + assert calls == [["claude", "mcp", "remove", "github", "-s", "user"]] + + def test_returns_false_when_server_missing(self, monkeypatch): + def fake_run(args, **kwargs): + raise subprocess.CalledProcessError(1, args, stderr="No MCP server named github found") + + monkeypatch.setattr(claude.subprocess, "run", fake_run) + + assert claude.remove_claude_mcp_server("github", "user") is False + + def test_returns_false_when_project_local_server_missing(self, monkeypatch): + def fake_run(args, **kwargs): + raise subprocess.CalledProcessError( + 1, + args, + stderr="No project-local MCP server found with name: github", + ) + + monkeypatch.setattr(claude.subprocess, "run", fake_run) + + assert claude.remove_claude_mcp_server("github", "project") is False + + def test_returns_false_when_user_scoped_server_missing(self, monkeypatch): + def fake_run(args, **kwargs): + raise subprocess.CalledProcessError( + 1, + args, + stderr="No user-scoped MCP server found with name: github", + ) + + monkeypatch.setattr(claude.subprocess, "run", fake_run) + + assert claude.remove_claude_mcp_server("github", "user") is False + + def test_unexpected_failure_raises(self, monkeypatch): + def fake_run(args, **kwargs): + raise subprocess.CalledProcessError(1, args, stderr="permission denied") + + monkeypatch.setattr(claude.subprocess, "run", fake_run) + + try: + claude.remove_claude_mcp_server("github", "user") + except RuntimeError as exc: + assert "Failed to remove MCP server 'github'" in str(exc) + else: + raise AssertionError("expected RuntimeError") + + class TestRegisterWebSearchMcp: def test_skips_registration_when_entry_is_current(self, monkeypatch): entry = claude._web_search_mcp_entry(WS, "m", "profile") @@ -1085,37 +1219,33 @@ def test_detects_registration_drift(self, monkeypatch): assert claude._web_search_mcp_is_current(state, entry) is False def test_clears_existing_then_adds(self, monkeypatch): - import ucode.mcp as mcp_mod - removed: list[str] = [] added: list = [] monkeypatch.setattr( - mcp_mod, "remove_claude_mcp_server", lambda name, scope: removed.append(scope) or True + claude, "remove_claude_mcp_server", lambda name, scope: removed.append(scope) or True ) monkeypatch.setattr( - mcp_mod, + claude, "add_claude_mcp_server", - lambda name, entry, scope=mcp_mod.MCP_USER_SCOPE: added.append((name, entry, scope)), + lambda name, entry, scope=claude.MCP_USER_SCOPE: added.append((name, entry, scope)), ) claude._register_web_search_mcp(WS, "databricks-gpt-5") - assert removed == list(mcp_mod.MCP_CLEANUP_SCOPES) + assert removed == list(claude.MCP_CLEANUP_SCOPES) assert len(added) == 1 name, entry, _ = added[0] assert name == "web_search" assert entry["env"]["UCODE_WEB_SEARCH_MODEL"] == "databricks-gpt-5" def test_remove_failures_are_swallowed(self, monkeypatch): - import ucode.mcp as mcp_mod - def boom(name, scope): raise RuntimeError("nope") added: list = [] - monkeypatch.setattr(mcp_mod, "remove_claude_mcp_server", boom) + monkeypatch.setattr(claude, "remove_claude_mcp_server", boom) monkeypatch.setattr( - mcp_mod, + claude, "add_claude_mcp_server", - lambda name, entry, scope=mcp_mod.MCP_USER_SCOPE: added.append(name), + lambda name, entry, scope=claude.MCP_USER_SCOPE: added.append(name), ) claude._register_web_search_mcp(WS, "m") assert added == ["web_search"] @@ -1123,27 +1253,23 @@ def boom(name, scope): def test_add_failure_is_non_blocking_and_warns(self, monkeypatch, capsys): # Regression: a failing `claude mcp add-json` used to abort the whole # `ucode claude` setup. It must now warn and return False instead. - import ucode.mcp as mcp_mod - - monkeypatch.setattr(mcp_mod, "remove_claude_mcp_server", lambda name, scope: False) + monkeypatch.setattr(claude, "remove_claude_mcp_server", lambda name, scope: False) - def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE): + def boom(name, entry, scope=claude.MCP_USER_SCOPE): raise RuntimeError("Failed to add MCP server 'web_search' via claude CLI.") - monkeypatch.setattr(mcp_mod, "add_claude_mcp_server", boom) + monkeypatch.setattr(claude, "add_claude_mcp_server", boom) result = claude._register_web_search_mcp(WS, "m") assert result is False captured = capsys.readouterr() assert "web_search" in captured.out.lower() or "web search" in captured.out.lower() def test_add_success_returns_true(self, monkeypatch): - import ucode.mcp as mcp_mod - - monkeypatch.setattr(mcp_mod, "remove_claude_mcp_server", lambda name, scope: False) + monkeypatch.setattr(claude, "remove_claude_mcp_server", lambda name, scope: False) monkeypatch.setattr( - mcp_mod, + claude, "add_claude_mcp_server", - lambda name, entry, scope=mcp_mod.MCP_USER_SCOPE: None, + lambda name, entry, scope=claude.MCP_USER_SCOPE: None, ) assert claude._register_web_search_mcp(WS, "m") is True @@ -1151,19 +1277,17 @@ def test_write_tool_config_completes_when_mcp_registration_fails(self, monkeypat # Regression for issue #100: a `claude mcp add-json` failure must not # block the rest of `ucode claude` setup (state save, managed-key # marking, etc.) from completing. - import ucode.mcp as mcp_mod - monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) monkeypatch.setattr(claude, "read_json_safe", lambda path: {}) monkeypatch.setattr(claude, "write_json_file", lambda path, payload: None) saved: list[dict] = [] monkeypatch.setattr(claude, "save_state", lambda state: saved.append(state)) - monkeypatch.setattr(mcp_mod, "remove_claude_mcp_server", lambda name, scope: False) + monkeypatch.setattr(claude, "remove_claude_mcp_server", lambda name, scope: False) - def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE): + def boom(name, entry, scope=claude.MCP_USER_SCOPE): raise RuntimeError("Failed to add MCP server 'web_search' via claude CLI.") - monkeypatch.setattr(mcp_mod, "add_claude_mcp_server", boom) + monkeypatch.setattr(claude, "add_claude_mcp_server", boom) state = {"workspace": WS, "codex_models": ["databricks-gpt-5"]} result = claude.write_tool_config(state, "databricks-claude-sonnet-4") diff --git a/tests/test_mcp.py b/tests/test_mcp.py index b977b787e..e41a0ca34 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2,14 +2,13 @@ from __future__ import annotations -import json -import subprocess import threading from unittest.mock import MagicMock import pytest from ucode import mcp +from ucode.agents import claude WS = "https://example.databricks.com" CLAUDE_STATE = {"workspace": WS, "available_tools": ["claude"]} @@ -69,71 +68,6 @@ def test_use_pat_appends_flag_and_profile_optional(self): assert "--profile" not in no_profile -class TestAddClaudeMcpServer: - def test_registers_stdio_proxy_command(self, monkeypatch): - calls: list[dict] = [] - - def fake_run(args, **kwargs): - calls.append({"args": args, "kwargs": kwargs}) - return MagicMock(returncode=0) - - monkeypatch.setattr(mcp.subprocess, "run", fake_run) - - mcp.add_claude_mcp_server("github", _proxy_argv()) - - args = calls[0]["args"] - assert args[:4] == ["claude", "mcp", "add", "github"] - assert args[4:6] == ["-s", "user"] - # `--` fences the proxy argv; everything after it is the stdio command. - assert args[6] == "--" - assert args[7:] == _proxy_argv() - - def test_always_load_routes_through_add_json_stdio_entry(self, monkeypatch): - # The skills registry needs `alwaysLoad: true`, which plain `mcp add` - # can't set — so the proxy argv is wrapped in a stdio entry dict and - # registered via add-json instead. - calls: list[dict] = [] - - def fake_run(args, **kwargs): - calls.append({"args": args, "kwargs": kwargs}) - return MagicMock(returncode=0) - - monkeypatch.setattr(mcp.subprocess, "run", fake_run) - - mcp.add_claude_mcp_server("skills", _proxy_argv(), always_load=True) - - args = calls[0]["args"] - assert args[:4] == ["claude", "mcp", "add-json", "skills"] - entry = json.loads(args[4]) - assert entry == { - "type": "stdio", - "command": _proxy_argv()[0], - "args": _proxy_argv()[1:], - "alwaysLoad": True, - } - assert args[5:] == ["-s", "user"] - - def test_dict_entry_routes_through_add_json(self, monkeypatch): - # The web_search server (agents/claude.py) registers a full stdio entry - # dict with its own env, which only `add-json` can express — a dict must - # route there rather than through the proxy `mcp add -- ` path. - calls: list[dict] = [] - - def fake_run(args, **kwargs): - calls.append({"args": args, "kwargs": kwargs}) - return MagicMock(returncode=0) - - monkeypatch.setattr(mcp.subprocess, "run", fake_run) - - entry = {"type": "stdio", "command": "ucode", "args": ["mcp", "web-search"]} - mcp.add_claude_mcp_server("web_search", entry) - - args = calls[0]["args"] - assert args[:4] == ["claude", "mcp", "add-json", "web_search"] - assert json.loads(args[4]) == entry - assert args[5:] == ["-s", "user"] - - class TestAddCodexMcpServer: def test_registers_stdio_proxy_command(self, monkeypatch): calls: list[dict] = [] @@ -175,65 +109,6 @@ def fake_run(args, **kwargs): assert call["kwargs"]["env"]["GEMINI_CLI_HOME"] == str(mcp.gemini.GEMINI_HOME_DIR) -class TestRemoveClaudeMcpServer: - def test_returns_true_when_server_removed(self, monkeypatch): - calls: list[list[str]] = [] - - def fake_run(args, **kwargs): - calls.append(args) - return MagicMock(returncode=0) - - monkeypatch.setattr(mcp.subprocess, "run", fake_run) - - assert mcp.remove_claude_mcp_server("github", "user") is True - assert calls == [["claude", "mcp", "remove", "github", "-s", "user"]] - - def test_returns_false_when_server_missing(self, monkeypatch): - def fake_run(args, **kwargs): - raise subprocess.CalledProcessError(1, args, stderr="No MCP server named github found") - - monkeypatch.setattr(mcp.subprocess, "run", fake_run) - - assert mcp.remove_claude_mcp_server("github", "user") is False - - def test_returns_false_when_project_local_server_missing(self, monkeypatch): - def fake_run(args, **kwargs): - raise subprocess.CalledProcessError( - 1, - args, - stderr="No project-local MCP server found with name: github", - ) - - monkeypatch.setattr(mcp.subprocess, "run", fake_run) - - assert mcp.remove_claude_mcp_server("github", "project") is False - - def test_returns_false_when_user_scoped_server_missing(self, monkeypatch): - def fake_run(args, **kwargs): - raise subprocess.CalledProcessError( - 1, - args, - stderr="No user-scoped MCP server found with name: github", - ) - - monkeypatch.setattr(mcp.subprocess, "run", fake_run) - - assert mcp.remove_claude_mcp_server("github", "user") is False - - def test_unexpected_failure_raises(self, monkeypatch): - def fake_run(args, **kwargs): - raise subprocess.CalledProcessError(1, args, stderr="permission denied") - - monkeypatch.setattr(mcp.subprocess, "run", fake_run) - - try: - mcp.remove_claude_mcp_server("github", "user") - except RuntimeError as exc: - assert "Failed to remove MCP server 'github'" in str(exc) - else: - raise AssertionError("expected RuntimeError") - - class TestCursorMcpClient: def test_cursor_registered_as_mcp_only_client(self): assert "cursor" in mcp.MCP_CLIENTS @@ -300,14 +175,14 @@ def _capture_claude(self, monkeypatch, *, claude_code_available: bool): 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(claude, "remove_claude_mcp_server", lambda name, scope: False) monkeypatch.setattr( - mcp, + claude, "add_claude_http_mcp_server", lambda name, url, **kw: http_calls.append((name, url)), ) monkeypatch.setattr( - mcp, + claude, "add_claude_mcp_server", lambda name, argv, scope=mcp.MCP_USER_SCOPE, **kw: proxy_calls.append((name, argv)), ) @@ -717,8 +592,8 @@ def test_skips_existing_server_state_by_name(self, monkeypatch): monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) monkeypatch.setattr(mcp, "discover_app_mcp_servers", lambda workspace, profile=None: []) _patch_mcp_choices(monkeypatch, "github") - monkeypatch.setattr(mcp, "remove_claude_mcp_server", lambda name, scope: False) - monkeypatch.setattr(mcp, "add_claude_mcp_server", lambda name, entry, scope: None) + monkeypatch.setattr(claude, "remove_claude_mcp_server", lambda name, scope: False) + monkeypatch.setattr(claude, "add_claude_mcp_server", lambda name, entry, scope: None) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) assert mcp.configure_mcp_command() == 0