From d221c0a721059081dcec9a3036bdc76935d0b29a Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 14 Aug 2026 17:57:02 +0000 Subject: [PATCH] mcp: --mcp / configure mcp --location no longer clobbers other MCP servers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ucode configure --mcp system.ai.slack` (and `configure mcp --location ` / `--services`) removed every other registered MCP server, forcing a full re-run of `configure mcp` to get the built-in ones back. Cause: `_resolve_location_mcp_servers` returned only the services discovered at the target location (+ skills), and `apply_mcp_server_changes` deletes any original server absent from that working set — so external connections, Genie spaces, apps, Vector Search / UC functions, and mcp-services in *other* schemas all got wiped. Fix: scope the replacement to mcp-services *in the target location*. The resolver now carries through every original server that isn't a service in that location (matched via `_is_mcp_service_in_location`, which guards against nested names and schema-prefix lookalikes like `system.aiX`). The location's own services are still strictly replaced (a stale in-location service is removed), and skills stay owned by `configure skills`. Verified via the real resolver: `--mcp system.ai.github` against a config with a custom external + Genie + other-schema service preserves all three and adds the new one. test_mcp.py: 132 pass (rewrote the test that codified the old clobber; added out-of-location preservation, stale-in-location replacement, and `_is_mcp_service_in_location` edge-case coverage). Full suite 1801 passed (2 pre-existing live-gateway e2e failures unrelated). ruff + ty clean. Co-authored-by: Isaac --- src/ucode/mcp.py | 54 +++++++++++++++++---- tests/test_mcp.py | 120 +++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 153 insertions(+), 21 deletions(-) diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index a3db7906..eb87ac48 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -1582,6 +1582,24 @@ def _skills_entries(servers: list[dict]) -> list[dict]: return [s for s in servers if s.get("kind") == SKILLS_MCP_KIND] +_MCP_SERVICE_URL_MARKER = "/ai-gateway/mcp-services/" + + +def _is_mcp_service_in_location(server: dict, location: str) -> bool: + """True if ``server`` is an mcp-service registered at ``location``. + + mcp-service entries carry a URL ``.../ai-gateway/mcp-services/..``. + A server belongs to ``location`` when that full name is ``.`` + (exactly one segment past the schema) — so ``system.ai`` matches + ``system.ai.slack`` but not ``system.ai.sub.x`` or ``system.aiX.y``.""" + url = server.get("url") + if not isinstance(url, str) or _MCP_SERVICE_URL_MARKER not in url: + return False + full_name = url.split(_MCP_SERVICE_URL_MARKER, 1)[1].strip("/") + prefix = f"{location}." + return full_name.startswith(prefix) and "." not in full_name[len(prefix) :] + + def _resolve_location_mcp_servers( workspace: str, profile: str | None, @@ -1592,19 +1610,23 @@ def _resolve_location_mcp_servers( ) -> list[dict]: """Build the desired MCP server list for ``--location .``. - Strict replacement for mcp-services: the returned list is exactly the ones - discovered at ``location`` (any previously-registered mcp-service outside it - is removed by ``apply_mcp_server_changes``), plus any existing skills - connection, preserved untouched. Raises ``RuntimeError`` for an invalid - location (HTTP 404 from the listing API) or any other listing failure. + Replacement is scoped to the mcp-services *in ``location``*: the returned + list is the ones discovered there, plus every other original server carried + through untouched — external connections, Genie/apps/Vector-Search/UC + functions, mcp-services in *other* schemas, and skills. So configuring one + location never disturbs servers registered elsewhere. (Previously this + returned only ``location``'s services + skills, so a one-shot + ``ucode configure --mcp system.ai.slack`` wiped every other MCP server — + #the custom-mcp-clobber bug.) Raises ``RuntimeError`` for an invalid location + (HTTP 404 from the listing API) or any other listing failure. When ``services`` is given, the discovered set is narrowed to exactly that subset (matched by full name like ``system.ai.github`` or bare short name like ``github``); names not found at ``location`` are skipped with a warning rather than failing, so a saved selection that references a since-removed service still configures the rest. An empty set selects - nothing (every previously-registered service in the location is removed). - ``None`` keeps the whole schema.""" + nothing (every previously-registered service *in the location* is removed; + servers outside it are still preserved). ``None`` keeps the whole schema.""" if location.count(".") != 1 or not all(part.strip() for part in location.split(".")): raise RuntimeError(f"--location must be `.`, got `{location}`.") @@ -1654,7 +1676,23 @@ def _resolve_location_mcp_servers( working_servers.append(original.copy()) else: working_servers.append(candidate) - return [*working_servers, *_skills_entries(original_servers)] + + # Carry through every original server that isn't an mcp-service in *this* + # location: other schemas' services, external/genie/app/vector-search/UC + # servers, etc. Only the location's own services are (re)placed above; skills + # are re-appended by name below. Without this, configuring one location would + # remove all unrelated MCP servers (`apply_mcp_server_changes` deletes any + # original not present in the returned set). + replaced_names = {s["name"] for s in working_servers} + skills_names = {s.get("name") for s in _skills_entries(original_servers)} + preserved = [ + server + for server in original_servers + if server.get("name") not in replaced_names + and server.get("name") not in skills_names + and not _is_mcp_service_in_location(server, location) + ] + return [*working_servers, *preserved, *_skills_entries(original_servers)] # The first wizard step lets the user choose which sources to search. Each is a diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 2852f478..919b64b1 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1595,6 +1595,39 @@ def _stub_location_base(monkeypatch, state): monkeypatch.setattr(mcp, "get_databricks_token", lambda workspace, profile=None: "token") +class TestIsMcpServiceInLocation: + def _svc(self, full_name: str) -> dict: + return { + "name": full_name.replace(".", "-"), + "url": f"{WS}/ai-gateway/mcp-services/{full_name}", + } + + def test_matches_a_service_directly_in_the_location(self): + assert mcp._is_mcp_service_in_location(self._svc("system.ai.slack"), "system.ai") + + def test_does_not_match_a_nested_name(self): + # Only one segment past the schema counts; a deeper name isn't "in" it. + assert not mcp._is_mcp_service_in_location(self._svc("system.ai.sub.x"), "system.ai") + + def test_does_not_match_a_schema_prefix_lookalike(self): + # `system.aiX` must not match location `system.ai`. + assert not mcp._is_mcp_service_in_location(self._svc("system.aiX.y"), "system.ai") + + def test_does_not_match_another_schema(self): + assert not mcp._is_mcp_service_in_location(self._svc("main.tools.helper"), "system.ai") + + def test_non_service_entries_are_never_in_a_location(self): + for url in ( + f"{WS}/api/2.0/mcp/external/jira", + f"{WS}/api/2.0/mcp/genie/123", + f"{WS}/api/2.0/mcp/sql", + ): + assert not mcp._is_mcp_service_in_location({"name": "x", "url": url}, "system.ai") + + def test_entry_without_url_is_not_in_a_location(self): + assert not mcp._is_mcp_service_in_location({"name": "x"}, "system.ai") + + class TestConfigureMcpFromLocation: def test_rejects_malformed_location(self, monkeypatch): _stub_location_base(monkeypatch, {**CLAUDE_STATE}) @@ -1682,19 +1715,36 @@ def fake_list(workspace, token, parent): }, ] - def test_replaces_servers_outside_location(self, monkeypatch): + def test_preserves_servers_outside_the_location(self, monkeypatch): + # Regression: configuring one location must NOT remove unrelated MCP + # servers. Previously `ucode configure --mcp system.ai.slack` wiped the + # user's other servers (custom-mcp-clobber bug); they must be carried + # through untouched, with only the location's own services (re)placed. saved_states: list[dict] = [] - configured: list[tuple[str, str, str, dict]] = [] + configured: list[tuple[str, str, str]] = [] removed: list[tuple[str, str]] = [] - outside_entry = { + # A grab-bag of servers that live outside `system.ai`. + sql_entry = { "name": "databricks-sql", "url": f"{WS}/api/2.0/mcp/sql", "auth": "proxy", "clients": ["claude"], } + external_entry = { + "name": "jira-mcp", + "url": f"{WS}/api/2.0/mcp/external/jira-mcp", + "auth": "proxy", + "clients": ["claude"], + } + other_schema_entry = { + "name": "main-tools-helper", + "url": f"{WS}/ai-gateway/mcp-services/main.tools.helper", + "auth": "proxy", + "clients": ["claude"], + } _stub_location_base( monkeypatch, - {**CLAUDE_STATE, "mcp_servers": [outside_entry]}, + {**CLAUDE_STATE, "mcp_servers": [sql_entry, external_entry, other_schema_entry]}, ) monkeypatch.setattr( mcp, @@ -1715,16 +1765,60 @@ def test_replaces_servers_outside_location(self, monkeypatch): assert mcp.configure_mcp_command(location="system.ai") == 0 - assert removed == [("claude", "databricks-sql")] + # Nothing outside system.ai was removed. + assert removed == [] + # Only the location's own service was (re)configured. assert [c[1] for c in configured] == ["system-ai-github"] - assert saved_states[-1]["mcp_servers"] == [ - { - "name": "system-ai-github", - "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", - "auth": "proxy", - "clients": ["claude"], - }, - ] + # Saved state keeps the three outside servers plus the new one. + saved_names = {s["name"] for s in saved_states[-1]["mcp_servers"]} + assert saved_names == { + "databricks-sql", + "jira-mcp", + "main-tools-helper", + "system-ai-github", + } + + def test_replaces_stale_service_within_the_location(self, monkeypatch): + # The location's OWN services are still strictly replaced: a service + # previously registered under system.ai but no longer discovered there + # is removed, while out-of-location servers stay. + saved_states: list[dict] = [] + removed: list[tuple[str, str]] = [] + stale_in_location = { + "name": "system-ai-oldservice", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.oldservice", + "auth": "proxy", + "clients": ["claude"], + } + outside = { + "name": "jira-mcp", + "url": f"{WS}/api/2.0/mcp/external/jira-mcp", + "auth": "proxy", + "clients": ["claude"], + } + _stub_location_base( + monkeypatch, + {**CLAUDE_STATE, "mcp_servers": [stale_in_location, outside]}, + ) + monkeypatch.setattr( + mcp, + "list_mcp_services", + lambda workspace, token, parent: (["system.ai.github"], None), + ) + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr( + mcp, + "remove_client_mcp_server", + lambda client, name: removed.append((client, name)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert mcp.configure_mcp_command(location="system.ai") == 0 + + # Stale in-location service removed; out-of-location server untouched. + assert removed == [("claude", "system-ai-oldservice")] + saved_names = {s["name"] for s in saved_states[-1]["mcp_servers"]} + assert saved_names == {"jira-mcp", "system-ai-github"} def test_preserves_skills_connection(self, monkeypatch): """A skills connection is owned by `configure skills`, so `configure mcp