diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 251d92fb7..c81e1fa5d 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1074,6 +1074,17 @@ def read_managed_mcp_urls() -> dict[str, str]: } +def read_user_mcp_names() -> set[str] | None: + """Names in Claude's user-scope ``mcpServers``; ``None`` when the config can't be parsed. + + Read-only, so ug can tell a name is free before picking an entry name.""" + config = _read_claude_config_for_rewrite(claude_mcp_config_path()) + if config is None: + return None + servers = config.get("mcpServers") + return set(servers) if isinstance(servers, dict) else set() + + def _register_web_search_mcp( workspace: str, search_model: str, diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 3408c9365..8531a6897 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -866,6 +866,17 @@ def _read_user_config_for_rewrite(path: Path) -> tomlkit.TOMLDocument | None: return None +def read_user_mcp_names() -> set[str] | None: + """Names in Codex's user-scope ``[mcp_servers]``; ``None`` when the config can't be parsed. + + Read-only, so ug can tell a name is free before picking an entry name.""" + doc = _read_user_config_for_rewrite(user_mcp_config_path()) + if doc is None: + return None + servers = doc.get(MANAGED_MCP_CONFIG_KEY) + return set(servers) if isinstance(servers, dict) else set() + + def write_user_mcp_servers(add: dict[str, dict], remove: set[str]) -> set[str]: """Apply ``add``/``remove`` to Codex's user-scope ``[mcp_servers]`` (``~/.codex/config.toml``, or under ``$CODEX_HOME``) in a single read-modify-write, instead of one ``codex mcp`` subprocess diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 6d8155763..ac00a1fde 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -7,7 +7,8 @@ import shutil import subprocess import threading -from collections.abc import Callable +from collections import Counter +from collections.abc import Callable, Iterable from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any from urllib.parse import urlparse @@ -644,15 +645,87 @@ def _server_choice(name: str, checked: bool, title: str | None = None) -> questi ) +def _native_mcp_names() -> set[str] | None: + """Names in the developer's own Claude and Codex MCP config; ``None`` if either can't be parsed.""" + names: set[str] = set() + for reader in (claude.read_user_mcp_names, codex.read_user_mcp_names): + found = reader() + if found is None: + return None + names |= found + return names + + +def mcp_service_entry_names( + full_names: Iterable[str], workspace: str, existing_servers: Iterable[dict] = () +) -> dict[str, str]: + """Map each MCP service full name (`..`) to its agent entry name. + + A service already registered in ``existing_servers`` (ug state, managed) keeps that entry's + name, so nothing is ever renamed. A new service gets its leaf (`exa` for `main.default.exa`) + only when the name is completely free: no existing entry, no entry in the developer's native + Claude/Codex config (any URL), no other service in ``full_names`` with the same leaf, and the + native config was readable. Otherwise it gets the dashed full name (`main-default-exa`; agent + CLIs reject dots in names).""" + unique = sorted(set(full_names)) + leaf_counts = Counter(full_name.rsplit(".", 1)[-1] for full_name in unique) + name_by_url: dict[str, str] = {} + taken: set[str] = set() + for server in existing_servers: + if name := _server_name(server): + taken.add(name) + if isinstance(server.get("url"), str): + name_by_url.setdefault(server["url"], name) + native = ( + _native_mcp_names() + if any(build_mcp_service_url(workspace, n) not in name_by_url for n in unique) + else set() + ) + names: dict[str, str] = {} + for full_name in unique: + leaf = full_name.rsplit(".", 1)[-1] + if (url := build_mcp_service_url(workspace, full_name)) in name_by_url: + names[full_name] = name_by_url[url] + elif native is not None and leaf_counts[leaf] == 1 and leaf not in taken | native: + names[full_name] = leaf + else: + names[full_name] = full_name.replace(".", "-") + return names + + +def _mcp_service_selection_names( + selections: Iterable[str], workspace: str, existing_servers: Iterable[dict] +) -> dict[str, str]: + """Entry names for the MCP-service selections in ``selections``, decided as one set.""" + prefix = MCP_SERVICE_SELECTION_PREFIX + full_names = [s.removeprefix(prefix) for s in selections if s.startswith(prefix)] + return mcp_service_entry_names([n for n in full_names if n], workspace, existing_servers) + + def _add_choice(selection: str, title: str) -> questionary.Choice: return questionary.Choice(title=title, value=f"{MCP_ADD_PREFIX}{selection}") +def _mcp_service_holder( + name: str, names: set[str], entry_urls: dict[str, str | None] +) -> str | None: + """Entry name in ``names`` that registers MCP service ``name`` (leaf or dashed), else None. + + An entry counts only if its recorded URL is this service's, so an unrelated `exa` that merely + shares the leaf does not (an entry with no recorded URL is trusted).""" + for candidate in (name.replace(".", "-"), name.rsplit(".", 1)[-1]): + url = entry_urls.get(candidate) + if candidate in names and (url is None or url.endswith(f"/{name}")): + return candidate + return None + + def _mcp_service_choice( name: str, known_names: set[str], additive: bool, managed_names: set[str] | None = None, + entry_urls: dict[str, str | None] | None = None, ) -> questionary.Choice | None: """Picker choice for one MCP-service full name (`..`), or None to hide it. @@ -661,11 +734,11 @@ def _mcp_service_choice( shows rows you can act on, so two kinds are hidden: a managed service (never ug's to add or remove) and, under `mcp add`, an already-configured one (nothing to add). Replace mode still lists ug's own as a removable toggle; an unregistered service is an add-choice.""" - registered_as = name.replace(".", "-") + urls = entry_urls or {} display_title = f"MCP: {name}" - if registered_as in (managed_names or set()): + if _mcp_service_holder(name, managed_names or set(), urls) is not None: return None - if registered_as in known_names: + if (registered_as := _mcp_service_holder(name, known_names, urls)) is not None: return None if additive else _server_choice(registered_as, True, display_title) return _add_choice(f"{MCP_SERVICE_SELECTION_PREFIX}{name}", display_title) @@ -680,6 +753,7 @@ def build_mcp_picker_choices( available_uc_functions_servers: list[dict] | None = None, additive: bool = False, managed_names: set[str] | None = None, + entry_urls: dict[str, str | None] | None = None, ) -> list[questionary.Choice | questionary.Separator]: original_by_name = _servers_by_name(original_servers) managed = managed_names or set() @@ -711,8 +785,11 @@ def push(choice: questionary.Choice | None) -> None: # Picker shows the dotted UC name; state/agents store the dashed form # (see resolver). The shared helper is also used by the background walk that # streams more services in, so up-front and streamed rows match exactly. - push(_mcp_service_choice(name, known_names, additive, managed)) - displayed_names.add(name.replace(".", "-")) + push(_mcp_service_choice(name, known_names, additive, managed, entry_urls)) + # Only the entry this service holds is represented by its row. + held = _mcp_service_holder(name, known_names, entry_urls or {}) + if held is not None: + displayed_names.add(held) for name in available_external_names: display_title = f"Connection: {name}" @@ -810,6 +887,7 @@ def prompt_for_mcp_server_choices( managed_names: set[str] | None = None, background_loader: Callable[[Callable[[list[questionary.Choice]], None], threading.Event], None] | None = None, + entry_urls: dict[str, str | None] | None = None, ) -> list[str] | None | _Back: """Show the MCP server picker. Returns the list of selected values, `None` if cancelled (Ctrl-C), or `_BACK` if `allow_back` and the user pressed Left @@ -839,6 +917,7 @@ def prompt_for_mcp_server_choices( available_uc_functions_servers, additive=additive, managed_names=managed_names, + entry_urls=entry_urls, ), style=picker_style(), instruction=instruction, @@ -881,7 +960,11 @@ def _is_app_mcp_server(server: dict) -> bool: def _resolve_managed_mcp_servers( - selector: dict, workspace: str, profile: str | None, clients: list[str] + selector: dict, + workspace: str, + profile: str | None, + clients: list[str], + existing_servers: Iterable[dict] = (), ) -> list[dict]: """Resolve a managed ``mcp_servers`` selector into MCP server entries for ``clients``. @@ -900,7 +983,9 @@ def _resolve_managed_mcp_servers( """ location = selector.get("unity_catalog_location") if isinstance(location, str) and location: - return _resolve_location_mcp_servers(workspace, profile, clients, location, []) + return _resolve_location_mcp_servers( + workspace, profile, clients, location, [], existing_servers=existing_servers + ) names = [n for n in (selector.get("names") or []) if isinstance(n, str) and n] if not names: return [] @@ -923,8 +1008,9 @@ def _resolve_managed_mcp_servers( # still discovers, since a whole-schema pointer has no explicit set to register. working: list[dict] = [] seen: set[str] = set() + entry_names = mcp_service_entry_names(names, workspace, existing_servers) for full_name in sorted(set(names)): - entry_name = full_name.replace(".", "-") + entry_name = entry_names[full_name] if entry_name in seen: continue seen.add(entry_name) @@ -1085,7 +1171,14 @@ def reconcile_managed_mcp_servers(managed: dict, agents: set[str]) -> list[dict] workspace, profile, clients = setup_mcp_clients( state, "Managed MCP Servers", agents=scope, action_note="Registering for" ) - working = _resolve_managed_mcp_servers(selector, workspace, profile, clients) + working = _resolve_managed_mcp_servers( + selector, + workspace, + profile, + clients, + # Managed entries first (file + fallback), so an entry already there keeps its name. + [*managed_mcp_servers(state), *(state.get("mcp_servers") or [])], + ) else: # No MCP client configured: nothing to register, but still clear what a prior configure left. workspace = state.get("workspace") @@ -1196,6 +1289,7 @@ def _resolve_mcp_selection( available_genie_servers: list[dict] | None = None, available_vector_search_servers: list[dict] | None = None, available_uc_functions_servers: list[dict] | None = None, + service_names: dict[str, str] | None = None, ) -> tuple[str, str]: if selection.startswith(APP_MCP_SELECTION_PREFIX): app_name = selection.removeprefix(APP_MCP_SELECTION_PREFIX) @@ -1233,8 +1327,10 @@ def _resolve_mcp_selection( if not full_name: raise RuntimeError("missing MCP service name") # Agent CLIs (claude/codex/gemini) reject dots in registered names. - # URL keeps the UC `..` form; entry name uses dashes. - return full_name.replace(".", "-"), build_mcp_service_url(workspace, full_name) + # URL keeps the UC `..` form; the entry name is the dashed full name unless + # `service_names` (see `mcp_service_entry_names`) picked another for this run. + entry_name = (service_names or {}).get(full_name) or full_name.replace(".", "-") + return entry_name, build_mcp_service_url(workspace, full_name) if selection.startswith(VECTOR_SEARCH_SELECTION_PREFIX): return _resolve_catalog_schema_selection( @@ -1317,6 +1413,7 @@ def _mcp_services_background_loader( known_names: set[str], additive: bool, managed_names: set[str] | None = None, + entry_urls: dict[str, str | None] | None = None, ) -> Callable[[Callable[[list[questionary.Choice]], None], threading.Event], None]: """Return a picker `background_loader` that runs the workspace-wide MCP-services walk and streams each schema's newly-found services into the open picker as choices, so the walk @@ -1332,7 +1429,11 @@ def on_services(new_names: list[str]) -> None: new_choices = [ choice for name in new_names - if (choice := _mcp_service_choice(name, known_names, additive, managed_names)) + if ( + choice := _mcp_service_choice( + name, known_names, additive, managed_names, entry_urls + ) + ) is not None ] if new_choices: @@ -1660,6 +1761,7 @@ def _resolve_location_mcp_servers( location: str, original_servers: list[dict], services: set[str] | None = None, + existing_servers: Iterable[dict] = (), ) -> list[dict]: """Build the desired MCP server list for ``--location .``. @@ -1716,9 +1818,10 @@ def _resolve_location_mcp_servers( ] original_by_name = _servers_by_name(original_servers) + entry_names = mcp_service_entry_names(names, workspace, [*original_servers, *existing_servers]) working_servers: list[dict] = [] for full_name in names: - entry_name = full_name.replace(".", "-") + entry_name = entry_names[full_name] original = original_by_name.get(entry_name) original_clients = list((original or {}).get("clients") or []) merged_clients = original_clients + [c for c in clients if c not in original_clients] @@ -1906,12 +2009,18 @@ def _configure_v2_mcp_selectors( picker_servers = [s for s in original_mcp_servers if s.get("kind") != SKILLS_MCP_KIND] original_by_name = _servers_by_name(picker_servers) - managed_names = managed_mcp_server_names(state, set(clients)) + managed_servers = managed_mcp_servers(state, set(clients)) + managed_names = {name for server in managed_servers if (name := _server_name(server))} + service_names = _mcp_service_selection_names( + selectors, workspace, [*original_mcp_servers, *managed_servers] + ) working_mcp_servers: list[dict] = list(skills_servers) working_names: set[str] = set() skipped_managed: set[str] = set() for selection in selectors: - entry_name, url = _resolve_mcp_selection(selection, workspace, available_app_servers) + entry_name, url = _resolve_mcp_selection( + selection, workspace, available_app_servers, service_names=service_names + ) if entry_name in working_names: continue if entry_name in managed_names: @@ -2000,12 +2109,19 @@ def configure_mcp_command( state, "Add MCP Servers" if append else "MCP Servers", agents=agents ) - managed_names = managed_mcp_server_names(state, set(clients)) + managed_servers = managed_mcp_servers(state, set(clients)) + managed_names = {name for server in managed_servers if (name := _server_name(server))} original_mcp_servers_for_location: list[dict] = list(state.get("mcp_servers") or []) if location is not None: working_mcp_servers = _resolve_location_mcp_servers( - workspace, profile, clients, location, original_mcp_servers_for_location, services + workspace, + profile, + clients, + location, + original_mcp_servers_for_location, + services, + existing_servers=managed_servers, ) working_mcp_servers = _drop_managed_servers(working_mcp_servers, managed_names) if append: @@ -2037,6 +2153,9 @@ def configure_mcp_command( if "apps" in excluded_sources: picker_servers = [s for s in picker_servers if not _is_app_mcp_server(s)] original_by_name = _servers_by_name(picker_servers) + entry_urls = { + n: s.get("url") for s in [*picker_servers, *managed_servers] if (n := _server_name(s)) + } # Single source (MCP services), so there's no "choose sources" step — discover the fast # `system.ai` list, show the picker immediately, and let the workspace-wide walk stream in @@ -2048,6 +2167,7 @@ def configure_mcp_command( set(original_by_name) | managed_names, additive=append, managed_names=managed_names, + entry_urls=entry_urls, ) selections = prompt_for_mcp_server_choices( discovered["external"], @@ -2060,6 +2180,7 @@ def configure_mcp_command( additive=append, managed_names=managed_names, background_loader=services_loader, + entry_urls=entry_urls, ) if selections is None or isinstance(selections, _Back): # No back-navigation without the source step; `_Back` can't occur, but keep the guard @@ -2071,6 +2192,33 @@ def configure_mcp_command( available_vector_search_servers = discovered["vector_search"] available_uc_functions_servers = discovered["uc_functions"] + # Names other selections in this run will hold (kept entries, and new non-service adds), so a + # service's leaf never collides with one of them. + run_names: list[dict] = [] + for selection in selections: + if not selection.startswith(MCP_ADD_PREFIX): + run_names.append({"name": selection}) + elif not selection.removeprefix(MCP_ADD_PREFIX).startswith(MCP_SERVICE_SELECTION_PREFIX): + try: + run_names.append( + { + "name": _resolve_mcp_selection( + selection.removeprefix(MCP_ADD_PREFIX), + workspace, + discovered["apps"], + discovered["genie"], + discovered["vector_search"], + discovered["uc_functions"], + )[0] + } + ) + except RuntimeError: + pass # reported when the selection is resolved below + service_names = _mcp_service_selection_names( + [s.removeprefix(MCP_ADD_PREFIX) for s in selections if s.startswith(MCP_ADD_PREFIX)], + workspace, + [*original_mcp_servers, *managed_servers, *run_names], + ) working_mcp_servers: list[dict] = list(skills_servers) working_names: set[str] = set() add_selections: list[str] = [] @@ -2092,6 +2240,7 @@ def configure_mcp_command( available_genie_mcp_servers, available_vector_search_servers, available_uc_functions_servers, + service_names=service_names, ) except RuntimeError as exc: print_warning(f"Skipped MCP selection `{selection}`: {exc}.") diff --git a/tests/conftest.py b/tests/conftest.py index 3b8ff90f2..c318b1f8a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -35,6 +35,7 @@ def _isolate_ucode_state(tmp_path, monkeypatch): import ucode.managed_files as managed_files_mod import ucode.os_compatibility.subprocess_cross_os as subprocess_cross_os_mod import ucode.state as state_mod + from ucode.agents import claude as claude_mod from ucode.agents import codex as codex_mod state_dir = tmp_path / ".ucode" @@ -57,6 +58,11 @@ def _isolate_ucode_state(tmp_path, monkeypatch): codex_mod, "CODEX_MODEL_CATALOG_PATH", state_dir / "codex-model-catalog.json" ) monkeypatch.setattr(codex_mod, "CODEX_CONFIG_PATH", tmp_path / ".codex" / "ucode.config.toml") + # `mcp` reads the developer's native Claude/Codex MCP config to avoid name clashes; never the real one. + monkeypatch.setattr(claude_mod, "CLAUDE_MCP_CONFIG_PATH", tmp_path / ".claude.json") + monkeypatch.setattr(codex_mod, "LEGACY_CODEX_CONFIG_PATH", tmp_path / ".codex" / "config.toml") + monkeypatch.delenv("CLAUDE_CONFIG_DIR", raising=False) + monkeypatch.delenv("CODEX_HOME", raising=False) def reject_privileged_write(path, _desired_text): pytest.fail( diff --git a/tests/test_mcp.py b/tests/test_mcp.py index f57147a3e..b5ae9a6d3 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2,13 +2,14 @@ from __future__ import annotations +import json import threading from unittest.mock import MagicMock import pytest from ucode import mcp -from ucode.agents import claude +from ucode.agents import claude, codex WS = "https://example.databricks.com" CLAUDE_STATE = {"workspace": WS, "available_tools": ["claude"]} @@ -1661,17 +1662,17 @@ def fake_list(workspace, token, parent): assert seen == {"parent": "system.ai"} assert picker_called == [] - assert [c[1] for c in configured] == ["system-ai-github", "system-ai-slack"] + assert [c[1] for c in configured] == ["github", "slack"] assert configured[0][2] == f"{WS}/ai-gateway/mcp-services/system.ai.github" assert saved_states[-1]["mcp_servers"] == [ { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "proxy", "clients": ["claude"], }, { - "name": "system-ai-slack", + "name": "slack", "url": f"{WS}/ai-gateway/mcp-services/system.ai.slack", "auth": "proxy", "clients": ["claude"], @@ -1712,10 +1713,10 @@ def test_replaces_servers_outside_location(self, monkeypatch): assert mcp.configure_mcp_command(location="system.ai") == 0 assert removed == [("claude", "databricks-sql")] - assert [c[1] for c in configured] == ["system-ai-github"] + assert [c[1] for c in configured] == ["github"] assert saved_states[-1]["mcp_servers"] == [ { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "proxy", "clients": ["claude"], @@ -1842,10 +1843,10 @@ def test_keeps_servers_outside_location(self, monkeypatch): # Nothing is removed; the new service is added and the outside one kept. assert removed == [] - assert [c[1] for c in configured] == ["system-ai-github"] + assert [c[1] for c in configured] == ["github"] assert saved_states[-1]["mcp_servers"] == [ { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "proxy", "clients": ["claude"], @@ -1885,7 +1886,7 @@ def test_services_subset_keeps_others_in_location(self, monkeypatch): assert removed == [] names = [s["name"] for s in saved_states[-1]["mcp_servers"]] - assert names == ["system-ai-github", "system-ai-slack"] + assert names == ["github", "system-ai-slack"] def test_empty_services_is_a_noop(self, monkeypatch): """`mcp add --names ""` has nothing to add, so it's a no-op that never @@ -1918,7 +1919,7 @@ def test_agents_scopes_registration_to_named_agent(self, monkeypatch): assert mcp.add_mcp_command(location="system.ai", agents={"claude"}) == 0 - assert configured == [("claude", "system-ai-github")] + assert configured == [("claude", "github")] assert saved_states[-1]["mcp_servers"][0]["clients"] == ["claude"] def test_agents_not_configured_raises(self, monkeypatch): @@ -2099,10 +2100,10 @@ def test_configures_only_the_requested_subset(self, monkeypatch): ) # slack is dropped; only the two requested services are configured. - assert sorted(c[1] for c in configured) == ["system-ai-github", "system-ai-gmail"] + assert sorted(c[1] for c in configured) == ["github", "gmail"] assert sorted(s["name"] for s in saved_states[-1]["mcp_servers"]) == [ - "system-ai-github", - "system-ai-gmail", + "github", + "gmail", ] def test_matches_bare_short_names(self, monkeypatch): @@ -2122,7 +2123,7 @@ def test_matches_bare_short_names(self, monkeypatch): assert mcp.configure_mcp_command(location="system.ai", services={"github"}) == 0 - assert [c[1] for c in configured] == ["system-ai-github"] + assert [c[1] for c in configured] == ["github"] def test_unknown_requested_service_warns_and_skips(self, monkeypatch): configured: list[tuple[str, str, str, dict]] = [] @@ -2149,7 +2150,7 @@ def test_unknown_requested_service_warns_and_skips(self, monkeypatch): ) # The known service is still configured; the unknown one is reported, not fatal. - assert [c[1] for c in configured] == ["system-ai-github"] + assert [c[1] for c in configured] == ["github"] assert any("system.ai.ghost" in w for w in warnings) def test_empty_services_removes_everything(self, monkeypatch): @@ -2234,10 +2235,10 @@ def test_adds_and_removes_to_match_new_selection(self, monkeypatch): # slack removed, gmail added, github untouched (entry unchanged). assert removed == [("claude", "system-ai-slack")] - assert [c[1] for c in configured] == ["system-ai-gmail"] + assert [c[1] for c in configured] == ["gmail"] assert sorted(s["name"] for s in saved_states[-1]["mcp_servers"]) == [ + "gmail", "system-ai-github", - "system-ai-gmail", ] def test_full_names_without_location_derive_the_schema(self, monkeypatch): @@ -2261,7 +2262,7 @@ def fake_list(workspace, token, parent): assert mcp.configure_mcp_command(services={"system.ai.github", "system.ai.slack"}) == 0 assert seen == {"parent": "system.ai"} - assert sorted(c[1] for c in configured) == ["system-ai-github", "system-ai-slack"] + assert sorted(c[1] for c in configured) == ["github", "slack"] def test_short_name_without_location_raises(self): try: @@ -3234,7 +3235,7 @@ class TestResolveManagedMcpServers: def test_location_registers_the_whole_schema(self, monkeypatch): captured: dict = {} - def fake_loc(ws, profile, clients, location, original, services=None): + def fake_loc(ws, profile, clients, location, original, services=None, existing_servers=()): captured.update(location=location, services=services, original=original) return [{"name": "system-ai-a", "url": "u", "auth": "proxy", "clients": clients}] @@ -3264,8 +3265,8 @@ def test_names_build_entries_directly_without_discovery(self, monkeypatch): ["claude", "codex"], ) by_name = {s["name"]: s for s in out} - assert set(by_name) == {"system-ai-github", "system-ai-jira", "main-tools-foo"} - github = by_name["system-ai-github"] + assert set(by_name) == {"github", "jira", "foo"} + github = by_name["github"] assert github["url"] == mcp.build_mcp_service_url(WS, "system.ai.github") assert github["auth"] == "proxy" assert github["clients"] == ["claude", "codex"] @@ -3279,7 +3280,7 @@ def test_duplicate_names_collapse_to_one_entry(self, monkeypatch): out = mcp._resolve_managed_mcp_servers( {"names": ["system.ai.github", "system.ai.github"]}, WS, None, ["claude"] ) - assert [s["name"] for s in out] == ["system-ai-github"] + assert [s["name"] for s in out] == ["github"] def test_malformed_names_are_skipped_with_a_warning(self, monkeypatch): # A bare (`github`) or over-qualified (`a.b.c.d`) name isn't a full FQN and can't be @@ -3298,7 +3299,7 @@ def test_malformed_names_are_skipped_with_a_warning(self, monkeypatch): out = mcp._resolve_managed_mcp_servers( {"names": ["github", "a.b.c.d", "system.ai.slack"]}, WS, None, ["claude"] ) - assert {s["name"] for s in out} == {"system-ai-slack"} + assert {s["name"] for s in out} == {"slack"} assert warned and "github" in warned[0] and "a.b.c.d" in warned[0] def test_all_malformed_names_returns_empty_without_discovery(self, monkeypatch): @@ -3332,7 +3333,7 @@ def fake_setup(s, section, **kw): monkeypatch.setattr( mcp, "_resolve_managed_mcp_servers", - lambda selector, ws, profile, clients: [dict(s, clients=clients) for s in resolved], + lambda selector, ws, profile, clients, *_: [dict(s, clients=clients) for s in resolved], ) monkeypatch.setattr( mcp, @@ -3385,7 +3386,7 @@ def test_location_config_registers(self, monkeypatch): monkeypatch.setattr( mcp, "_resolve_managed_mcp_servers", - lambda selector, ws, profile, clients: ( + lambda selector, ws, profile, clients, *_: ( seen_selector.update(selector) or [dict(s, clients=clients) for s in [{"name": "system-ai-a", "url": "u"}]] ), @@ -3396,6 +3397,32 @@ def test_location_config_registers(self, monkeypatch): assert seen_selector == {"unity_catalog_location": "system.ai"} assert [s["name"] for s in out] == ["system-ai-a"] + def test_existing_managed_file_entry_keeps_its_name_on_upgrade(self, monkeypatch): + url = mcp.build_mcp_service_url(WS, "system.ai.github") + state = {"workspace": WS, "profile": None, "managed_mcp_servers": []} + applied: dict = {} + real_resolve = mcp._resolve_managed_mcp_servers + self._wire( + monkeypatch, + state=state, + configured=["claude"], + resolved=[], + applied=applied, + saved={}, + setup_calls=[], + ) + monkeypatch.setattr(mcp, "_resolve_managed_mcp_servers", real_resolve) + monkeypatch.setattr(mcp, "_native_mcp_names", lambda: set()) + monkeypatch.setattr(mcp.claude, "read_managed_mcp_urls", lambda: {"system-ai-github": url}) + monkeypatch.setattr(mcp.codex, "read_managed_mcp_urls", lambda: {}) + monkeypatch.setattr(mcp.claude, "managed_mcp_uses_managed_file", lambda *a, **k: False) + monkeypatch.setattr(mcp.claude, "reconcile_managed_mcp", lambda *a, **k: True) + monkeypatch.setattr(mcp.codex, "reconcile_managed_mcp", lambda *a, **k: True) + out = mcp.reconcile_managed_mcp_servers( + {"mcp_servers": {"names": ["system.ai.github"]}}, {"claude"} + ) + assert [s["name"] for s in out] == ["system-ai-github"] + def test_reconcile_removes_servers_the_admin_dropped(self, monkeypatch): # A config that no longer lists any MCP server clears the previously-registered managed set. previous = [{"name": "old", "url": "u", "auth": "proxy", "clients": ["claude"]}] @@ -3930,7 +3957,10 @@ def test_location_add_skips_an_already_managed_service(self, monkeypatch, capsys configured: list[str] = [] saved_states: list[dict] = [] _stub_location_base(monkeypatch, {**CLAUDE_STATE}) - self._stub_managed_files(monkeypatch, claude_urls={"system-ai-github": "managed-url"}) + self._stub_managed_files( + monkeypatch, + claude_urls={"system-ai-github": f"{WS}/ai-gateway/mcp-services/system.ai.github"}, + ) monkeypatch.setattr( mcp, "list_mcp_services", @@ -3946,8 +3976,8 @@ def test_location_add_skips_an_already_managed_service(self, monkeypatch, capsys assert mcp.add_mcp_command(location="system.ai") == 0 # Only the unmanaged service is registered, and the skip is reported rather than silent. - assert configured == ["system-ai-slack"] - assert [s["name"] for s in saved_states[-1]["mcp_servers"]] == ["system-ai-slack"] + assert configured == ["slack"] + assert [s["name"] for s in saved_states[-1]["mcp_servers"]] == ["slack"] out = capsys.readouterr().out assert "system-ai-github" in out assert "managed configuration" in out @@ -4362,3 +4392,156 @@ def test_setup_mcp_clients_user_messages_name_ug_not_ucode(monkeypatch, capsys): mcp.setup_mcp_clients(state, "Add MCP Servers", require_auth=False, agents={"codex"}) with pytest.raises(RuntimeError, match="Run `ug configure` first"): mcp.setup_mcp_clients({}, "Add MCP Servers") + + +class TestMcpServiceEntryNames: + """A service gets its short name only when it is completely free; nothing is ever renamed.""" + + @staticmethod + def _url(full_name): + return mcp.build_mcp_service_url(WS, full_name) + + def _run(self, monkeypatch, services, state_extra=None, location="main.default"): + """Run `ug mcp add --location` and return ({name: url} configured, saved servers).""" + state = {**CLAUDE_STATE, **(state_extra or {})} + saved: list[dict] = [] + configured: dict[str, str] = {} + _stub_location_base(monkeypatch, state) + monkeypatch.setattr( + mcp, "list_mcp_services", lambda ws, token, parent: (sorted(services), None) + ) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, *a, **kw: configured.update({name: url}) or [], + ) + monkeypatch.setattr(mcp, "remove_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr(mcp, "save_state", lambda s: saved.append(s.copy())) + assert mcp.add_mcp_command(location=location, services=set(services)) == 0 + return configured, (saved[-1]["mcp_servers"] if saved else state.get("mcp_servers")) + + def test_free_leaf_gets_the_short_name(self, monkeypatch): + configured, _ = self._run(monkeypatch, {"main.default.exa", "main.default.github"}) + assert sorted(configured) == ["exa", "github"] + + def test_leaf_taken_in_ug_state_gets_the_dashed_name(self, monkeypatch): + held = {"name": "exa", "url": self._url("ml.tools.exa"), "auth": "proxy"} + held["clients"] = ["claude"] + configured, servers = self._run(monkeypatch, {"main.default.exa"}, {"mcp_servers": [held]}) + assert list(configured) == ["main-default-exa"] + assert held in servers + + def test_leaf_taken_by_a_managed_entry_gets_the_dashed_name(self, monkeypatch): + managed = {"name": "exa", "url": f"{WS}/api/2.0/mcp/external/exa", "clients": ["claude"]} + configured, _ = self._run( + monkeypatch, {"main.default.exa"}, {"managed_mcp_servers": [managed]} + ) + assert list(configured) == ["main-default-exa"] + + def test_same_service_already_managed_is_skipped_not_duplicated(self, monkeypatch): + managed = {"name": "exa", "url": self._url("main.default.exa"), "clients": ["claude"]} + configured, _ = self._run( + monkeypatch, {"main.default.exa"}, {"managed_mcp_servers": [managed]} + ) + assert configured == {} + + def test_leaf_in_native_claude_config_gets_the_dashed_name(self, monkeypatch): + # Any URL counts, even the same service registered by hand. + path = claude.claude_mcp_config_path() + path.write_text(json.dumps({"mcpServers": {"exa": {"url": self._url("main.default.exa")}}})) + configured, _ = self._run(monkeypatch, {"main.default.exa"}) + assert list(configured) == ["main-default-exa"] + + def test_leaf_in_native_codex_config_blocks_a_claude_only_run(self, monkeypatch): + codex.user_mcp_config_path().parent.mkdir(parents=True, exist_ok=True) + codex.user_mcp_config_path().write_text('[mcp_servers.exa]\nurl = "https://x.test/mcp"\n') + configured, _ = self._run(monkeypatch, {"main.default.exa", "main.default.github"}) + assert sorted(configured) == ["github", "main-default-exa"] + + def test_unreadable_native_config_makes_every_name_dashed(self, monkeypatch): + claude.claude_mcp_config_path().write_text("{not json") + configured, _ = self._run(monkeypatch, {"main.default.exa"}) + assert list(configured) == ["main-default-exa"] + + def test_services_sharing_a_leaf_in_one_run_both_get_dashed_names(self, monkeypatch): + names = mcp.mcp_service_entry_names( + ["main.default.exa", "ml.tools.exa", "main.default.x"], WS + ) + assert names == { + "main.default.exa": "main-default-exa", + "ml.tools.exa": "ml-tools-exa", + "main.default.x": "x", + } + + def test_rerun_never_renames_existing_entries(self, monkeypatch): + dashed = { + "name": "main-default-exa", + "url": self._url("main.default.exa"), + "auth": "proxy", + "clients": ["claude"], + } + short = {**dashed, "name": "github", "url": self._url("main.default.github")} + configured, servers = self._run( + monkeypatch, + {"main.default.exa", "main.default.github"}, + {"mcp_servers": [dashed, short]}, + ) + assert sorted(s["name"] for s in servers) == ["github", "main-default-exa"] + assert set(configured) <= {"github", "main-default-exa"} + + def test_managed_location_services_use_the_same_rule(self, monkeypatch): + monkeypatch.setattr(mcp, "get_databricks_token", lambda workspace, profile=None: "token") + monkeypatch.setattr( + mcp, "list_mcp_services", lambda ws, token, parent: (["main.default.exa"], None) + ) + working = mcp._resolve_managed_mcp_servers( + {"unity_catalog_location": "main.default"}, WS, None, ["claude"] + ) + assert [s["name"] for s in working] == ["exa"] + held = {"name": "exa", "url": f"{WS}/api/2.0/mcp/external/exa"} + working = mcp._resolve_managed_mcp_servers( + {"unity_catalog_location": "main.default"}, WS, None, ["claude"], [held] + ) + assert [s["name"] for s in working] == ["main-default-exa"] + + def test_leaf_held_by_another_selection_in_the_run_gets_the_dashed_name(self, monkeypatch): + # An external connection named `github` and `system.ai.github` picked together: both would + # resolve to `github`, so the service must take the dashed name and neither is dropped. + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + configured: dict[str, str] = {} + saved: list[dict] = [] + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) + monkeypatch.setattr(mcp, "discover_mcp_service_names", lambda workspace, profile=None: []) + monkeypatch.setattr( + mcp, + "prompt_for_mcp_server_choices", + lambda *a, **kw: [ + f"{mcp.MCP_ADD_PREFIX}{mcp.EXTERNAL_MCP_SELECTION_PREFIX}github", + f"{mcp.MCP_ADD_PREFIX}{mcp.MCP_SERVICE_SELECTION_PREFIX}system.ai.github", + ], + ) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, *a, **kw: configured.update({name: url}) or [], + ) + monkeypatch.setattr(mcp, "remove_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr(mcp, "save_state", lambda s: saved.append(s.copy())) + assert mcp.configure_mcp_command() == 0 + assert sorted(configured) == ["github", "system-ai-github"] + assert sorted(s["name"] for s in saved[-1]["mcp_servers"]) == ["github", "system-ai-github"] + + def test_picker_finds_a_service_registered_under_its_short_name(self): + urls = {"exa": self._url("main.default.exa")} + choice = mcp._mcp_service_choice( + "main.default.exa", {"exa"}, additive=False, entry_urls=urls + ) + assert choice.value == "exa" and choice.checked + assert mcp._mcp_service_choice("main.default.exa", {"exa"}, True, entry_urls=urls) is None + + def test_picker_ignores_an_unrelated_entry_with_the_leaf_name(self): + urls = {"exa": self._url("other.schema.exa")} + choice = mcp._mcp_service_choice( + "main.default.exa", {"exa"}, additive=False, entry_urls=urls + ) + assert choice.value.startswith(mcp.MCP_ADD_PREFIX)