From 97e3649d44ae4dd84460356c2e9c14f5f966816f Mon Sep 17 00:00:00 2001 From: Xiang Shen Date: Tue, 8 Sep 2026 04:46:43 +0000 Subject: [PATCH 1/2] Support per-client skill MCP scopes Store the skills MCP connection's scope as a per-client map (`skill_locations_by_client`) instead of a single flat `skill_locations` list, so `skill add`/`skill remove` can target individual agents. Reads fall back to the flat list, mirrored to every client, for connections written by older builds; `skill_locations` is kept as the union mirror. Co-authored-by: Arthur Jenoudet Co-authored-by: Isaac --- src/ucode/mcp.py | 183 ++++++++++++++++++++++++++++++++++++++-------- tests/test_mcp.py | 104 ++++++++++++++++++++++---- 2 files changed, 243 insertions(+), 44 deletions(-) diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index b40a8ed6..47f5262a 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -100,6 +100,7 @@ class _Back: } SKILLS_MCP_KIND = "skills" SKILLS_MCP_SERVER_NAME = "databricks-skill-registry" +SKILL_LOCATIONS_BY_CLIENT_KEY = "skill_locations_by_client" # MCP-only clients ucode never launches for model routing, so they never land in # `available_tools`; they're eligible for MCP config purely on being installed. MCP_ONLY_CLIENTS = ("cursor",) @@ -2058,23 +2059,74 @@ def _merge_clients(prior: list[str] | None, new: list[str]) -> list[str]: return prior + [c for c in new if c not in prior] -def _build_skills_entry(workspace: str, locations: list[str], clients: list[str]) -> dict: - """Canonical single skills-registry entry. ``skill_locations`` is the source - of truth; the URL is always derived from it, never parsed back.""" +def _dedupe_locations(locations: list[str]) -> list[str]: + """Return valid locations once each, preserving their input order.""" + return list(dict.fromkeys(loc for loc in locations if isinstance(loc, str) and loc)) + + +def _skill_locations_by_client(entry: dict | None) -> dict[str, list[str]]: + """Per-client skill locations. Reads the stored per-client map when present; otherwise derives + it from a legacy flat ``skill_locations``, mirrored to every client, so reads work on both shapes.""" + stored = (entry or {}).get(SKILL_LOCATIONS_BY_CLIENT_KEY) + if isinstance(stored, dict): + return { + client: _dedupe_locations(locations) + for client, locations in stored.items() + if client in MCP_CLIENTS and isinstance(locations, list) + } + flat = (entry or {}).get("skill_locations") + flat = _dedupe_locations(flat if isinstance(flat, list) else []) + return { + client: list(flat) + for client in ((entry or {}).get("clients") or []) + if client in MCP_CLIENTS + } + + +def _skill_locations_by_client_from_state(state: dict) -> dict[str, list[str]]: + return _skill_locations_by_client(_skills_entry(list(state.get("mcp_servers") or []))) + + +def skill_locations_for_client(entry: dict | None, client: str) -> list[str]: + """One client's skills scope from a persisted skills entry.""" + return _skill_locations_by_client(entry).get(client, []) + + +def _build_skills_entry( + workspace: str, + locations_by_client: dict[str, list[str]], + clients: list[str], +) -> dict: + """Build the skills-registry entry from a per-client developer scope. ``skill_locations`` mirrors + the union across clients so legacy readers and a downgrade to a flat-scope build stay coherent.""" + by_client = { + client: _dedupe_locations(locations) + for client, locations in (locations_by_client or {}).items() + if client in MCP_CLIENTS and _dedupe_locations(locations) + } + mirror: list[str] = [] + for locations in by_client.values(): + mirror = _union_locations(mirror, locations) return { "name": SKILLS_MCP_SERVER_NAME, "kind": SKILLS_MCP_KIND, - "skill_locations": list(locations), - "url": build_skills_mcp_url(workspace, locations), + "skill_locations": mirror, + SKILL_LOCATIONS_BY_CLIENT_KEY: by_client, + "url": build_skills_mcp_url(workspace, mirror), "auth": "proxy", "clients": clients, } +def _skills_entry(servers: list[dict]) -> dict | None: + """Return the skills-registry entry, if one is present.""" + return next((server for server in servers if server.get("kind") == SKILLS_MCP_KIND), None) + + def _resolve_skills_mcp_servers( workspace: str, clients: list[str], - locations: list[str], + locations_by_client: dict[str, list[str]], original_servers: list[dict], ) -> list[dict]: """Rebuild the MCP server list around exactly one skills entry. @@ -2085,14 +2137,14 @@ def _resolve_skills_mcp_servers( else, and appends one rebuilt entry whose clients merge the prior skills entry's clients with ``clients``. """ - prior = next((s for s in original_servers if s.get("kind") == SKILLS_MCP_KIND), None) + prior = _skills_entry(original_servers) merged = _merge_clients((prior or {}).get("clients"), clients) kept = [ s for s in original_servers if s.get("kind") != SKILLS_MCP_KIND and _server_name(s) != SKILLS_MCP_SERVER_NAME ] - return [*kept, _build_skills_entry(workspace, locations, merged)] + return [*kept, _build_skills_entry(workspace, locations_by_client, merged)] def _join_with_and(items: list[str]) -> str: @@ -2107,6 +2159,12 @@ def _skills_tools_description(locations: list[str]) -> str: return f"UC skill utility tools + skills tools in schema {_join_with_and(locations)}" +def _skills_workspace(entry: dict) -> str: + """Extract the workspace base URL from a skills-registry entry.""" + url = str(entry.get("url") or "") + return url.split("/ai-gateway/skills/", 1)[0] + + def _print_skills_summary(entry: dict) -> None: """Report the registered skills connection and how to start using it.""" clients = [ @@ -2117,9 +2175,24 @@ def _print_skills_summary(entry: dict) -> None: console.print() print_success("Skills MCP registered") print_kv("Server", str(entry.get("name") or SKILLS_MCP_SERVER_NAME)) - print_kv("URL", str(entry.get("url") or "")) - print_kv("Configured", ", ".join(clients) if clients else "none") - print_kv("Tools", _skills_tools_description(entry.get("skill_locations") or [])) + scopes = { + client: skill_locations_for_client(entry, client) + for client in (entry.get("clients") or []) + if client in MCP_CLIENTS + } + distinct_scopes = {tuple(locations) for locations in scopes.values()} + if len(distinct_scopes) <= 1: + locations = next(iter(scopes.values()), []) + print_kv("URL", build_skills_mcp_url(_skills_workspace(entry), locations)) + print_kv("Configured", ", ".join(clients) if clients else "none") + print_kv("Tools", _skills_tools_description(locations)) + else: + print_kv("Configured", ", ".join(clients) if clients else "none") + workspace = _skills_workspace(entry) + for client, locations in scopes.items(): + display = str(MCP_CLIENTS[client]["display"]) + print_kv(f"{display} URL", build_skills_mcp_url(workspace, locations)) + print_kv(f"{display} tools", _skills_tools_description(locations)) print_note( "Run `ucode ` to use the skills MCP. For existing sessions, " "restart the agent for the skills to take effect." @@ -2127,32 +2200,76 @@ def _print_skills_summary(entry: dict) -> None: def _update_skills_mcp( - state: dict, workspace: str, profile: str | None, clients: list[str], locations: list[str] -) -> None: - """Rebuild the single skills connection for ``locations`` and persist it.""" + state: dict, + workspace: str, + profile: str | None, + clients: list[str], + locations_by_client: dict[str, list[str]], + *, + print_summary: bool = True, + use_pat: bool | None = None, +) -> bool: + """Persist one skills entry and update only clients whose scope changed.""" original = list(state.get("mcp_servers") or []) - working = _resolve_skills_mcp_servers(workspace, clients, locations, original) - changed = apply_mcp_server_changes(original, working, clients, workspace, profile) + working = _resolve_skills_mcp_servers(workspace, clients, locations_by_client, original) + original_entry = _skills_entry(original) + working_entry = _skills_entry(working) + if working_entry is None: + raise RuntimeError("Failed to build the Skills MCP connection.") + + changed = False + for client in clients: + working_view = [ + _build_skills_entry( + workspace, + {client: skill_locations_for_client(working_entry, client)}, + [client], + ) + ] + original_view = [] + if original_entry is not None and client in (original_entry.get("clients") or []): + original_view = [ + _build_skills_entry( + workspace, + {client: skill_locations_for_client(original_entry, client)}, + [client], + ) + ] + changed = ( + apply_mcp_server_changes( + original_view, + working_view, + [client], + workspace, + profile, + use_pat=bool(state.get("use_pat")) if use_pat is None else use_pat, + ) + or changed + ) if changed or original != working: state["mcp_servers"] = working save_state(state) - entry = next(s for s in working if s.get("kind") == SKILLS_MCP_KIND) - _print_skills_summary(entry) + if print_summary: + _print_skills_summary(working_entry) + return changed or original != working def configure_skills_mcp_command(locations: list[str]) -> int: - """Set the skills MCP connection's ``skill_locations`` to exactly ``locations``, - replacing any previous set.""" + """Set every configured client's skill scope to ``locations``.""" state = load_state() workspace, profile, clients = setup_mcp_clients(state, "Skills MCP") - _update_skills_mcp(state, workspace, profile, clients, locations) + locations_by_client = _skill_locations_by_client_from_state(state) + for client in clients: + locations_by_client[client] = list(locations) + _update_skills_mcp(state, workspace, profile, clients, locations_by_client) return 0 def _skill_mcp_locations(state: dict) -> list[str]: """The skills MCP connection's ``skill_locations``, or ``[]`` if none exists.""" - entry = next(iter(_skills_entries(list(state.get("mcp_servers") or []))), None) - return list((entry or {}).get("skill_locations") or []) + entry = _skills_entry(list(state.get("mcp_servers") or [])) + locations = (entry or {}).get("skill_locations") + return _dedupe_locations(locations if isinstance(locations, list) else []) def register_schemaless_skills_connection( @@ -2160,13 +2277,15 @@ def register_schemaless_skills_connection( ) -> None: """Register/keep the skills MCP connection without changing its schema set. - Download mode calls this after writing files: it preserves any prior - ``--mcp`` ``skill_locations`` and otherwise registers the bare schema-less - route (utility tools only).""" - _update_skills_mcp(state, workspace, profile, clients, _skill_mcp_locations(state)) + Download mode calls this after writing files: it preserves each client's prior + ``--mcp`` scope and otherwise registers the bare schema-less route (utility tools only).""" + _update_skills_mcp( + state, workspace, profile, clients, _skill_locations_by_client_from_state(state) + ) def _union_locations(base: list[str], new: list[str]) -> list[str]: + """Return an order-preserving union of two skill-location lists.""" have = set(base) merged = list(base) for location in new: @@ -2177,9 +2296,13 @@ def _union_locations(base: list[str], new: list[str]) -> list[str]: def add_skills_command(locations: list[str]) -> int: - """Add ``locations`` to the skills MCP connection's scope, keeping any already configured.""" + """Add ``locations`` to every configured client's skill scope, keeping any already configured.""" state = load_state() workspace, profile, clients = setup_mcp_clients(state, "Add Skills MCP") - merged = _union_locations(_skill_mcp_locations(state), locations) - _update_skills_mcp(state, workspace, profile, clients, merged) + locations_by_client = _skill_locations_by_client_from_state(state) + for client in clients: + locations_by_client[client] = _union_locations( + locations_by_client.get(client, []), locations + ) + _update_skills_mcp(state, workspace, profile, clients, locations_by_client) return 0 diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 73b28f39..bf289465 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1939,9 +1939,15 @@ def _find_skills(servers): return [s for s in servers if s.get("kind") == mcp.SKILLS_MCP_KIND] +def _by_client(clients, locations): + return {client: list(locations) for client in clients} + + class TestResolveSkillsMcpServers: def test_builds_single_canonical_entry(self): - servers = mcp._resolve_skills_mcp_servers(WS, ["claude"], ["main.default"], []) + servers = mcp._resolve_skills_mcp_servers( + WS, ["claude"], _by_client(["claude"], ["main.default"]), [] + ) assert _find_skills(servers) == servers entry = servers[0] assert entry["name"] == mcp.SKILLS_MCP_SERVER_NAME @@ -1967,7 +1973,7 @@ def test_keeps_other_entries_and_rebuilds_to_one_skills_entry(self): "clients": ["codex"], } servers = mcp._resolve_skills_mcp_servers( - WS, ["claude"], ["a.b"], [service_entry, stale_skills] + WS, ["claude"], _by_client(["claude"], ["a.b"]), [service_entry, stale_skills] ) assert service_entry in servers skills = _find_skills(servers) @@ -1981,7 +1987,9 @@ def test_drops_entry_matching_skills_server_name_even_without_kind(self): "url": f"{WS}/ai-gateway/skills/", "clients": ["claude"], } - servers = mcp._resolve_skills_mcp_servers(WS, ["claude"], ["a.b"], [old_named]) + servers = mcp._resolve_skills_mcp_servers( + WS, ["claude"], _by_client(["claude"], ["a.b"]), [old_named] + ) assert len(servers) == 1 assert servers[0]["kind"] == mcp.SKILLS_MCP_KIND @@ -1993,7 +2001,9 @@ def test_url_derives_from_locations_not_stale_url(self): "url": f"{WS}/ai-gateway/skills/?schema=stale.value", "clients": ["claude"], } - servers = mcp._resolve_skills_mcp_servers(WS, ["claude"], ["new.two"], [stale]) + servers = mcp._resolve_skills_mcp_servers( + WS, ["claude"], _by_client(["claude"], ["new.two"]), [stale] + ) assert servers[0]["url"] == f"{WS}/ai-gateway/skills/?schema=new.two" def test_empty_locations_yields_bare_route(self): @@ -2042,7 +2052,9 @@ def test_set_on_empty_registers_connection(self, monkeypatch): def test_location_replaces_prior_set(self, monkeypatch): saved_states: list[dict] = [] - prior = mcp._resolve_skills_mcp_servers(WS, ["claude"], ["A.a", "B.b"], []) + prior = mcp._resolve_skills_mcp_servers( + WS, ["claude"], _by_client(["claude"], ["A.a", "B.b"]), [] + ) _stub_location_base(monkeypatch, _skills_state(prior)) monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -2053,7 +2065,7 @@ def test_location_replaces_prior_set(self, monkeypatch): def test_multiple_locations_set_in_order(self, monkeypatch): saved_states: list[dict] = [] - prior = mcp._resolve_skills_mcp_servers(WS, ["claude"], ["A.a"], []) + prior = mcp._resolve_skills_mcp_servers(WS, ["claude"], _by_client(["claude"], ["A.a"]), []) _stub_location_base(monkeypatch, _skills_state(prior)) monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -2062,6 +2074,21 @@ def test_multiple_locations_set_in_order(self, monkeypatch): assert _find_skills(saved_states[-1]["mcp_servers"])[0]["skill_locations"] == ["X.x", "Y.y"] + def test_replaces_scope_for_configured_clients_only(self, monkeypatch): + saved_states: list[dict] = [] + prior = mcp._resolve_skills_mcp_servers( + WS, ["claude", "codex"], {"claude": ["claude.old"], "codex": ["codex.kept"]}, [] + ) + _stub_location_base(monkeypatch, _skills_state(prior)) + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert mcp.configure_skills_mcp_command(["new.default"]) == 0 + + entry = _find_skills(saved_states[-1]["mcp_servers"])[0] + assert mcp.skill_locations_for_client(entry, "claude") == ["new.default"] + assert mcp.skill_locations_for_client(entry, "codex") == ["codex.kept"] + def test_preserves_mcp_service_entries_across_set(self, monkeypatch): saved_states: list[dict] = [] service_entry = { @@ -2070,7 +2097,9 @@ def test_preserves_mcp_service_entries_across_set(self, monkeypatch): "auth": "env:OAUTH_TOKEN", "clients": ["claude"], } - prior = mcp._resolve_skills_mcp_servers(WS, ["claude"], ["A.a"], [service_entry]) + prior = mcp._resolve_skills_mcp_servers( + WS, ["claude"], _by_client(["claude"], ["A.a"]), [service_entry] + ) _stub_location_base(monkeypatch, _skills_state(prior)) monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -2084,13 +2113,47 @@ def test_preserves_mcp_service_entries_across_set(self, monkeypatch): class TestSkillMcpLocations: def test_reads_locations_off_skills_entry(self): - state = _skills_state(mcp._resolve_skills_mcp_servers(WS, ["claude"], ["A.a", "B.b"], [])) + state = _skills_state( + mcp._resolve_skills_mcp_servers( + WS, ["claude"], _by_client(["claude"], ["A.a", "B.b"]), [] + ) + ) assert mcp._skill_mcp_locations(state) == ["A.a", "B.b"] def test_empty_when_no_skills_entry(self): assert mcp._skill_mcp_locations(_skills_state([])) == [] assert mcp._skill_mcp_locations(_skills_state()) == [] + def test_ignores_malformed_default_locations(self): + entry = {"kind": mcp.SKILLS_MCP_KIND, "skill_locations": "not-a-list"} + state = _skills_state([entry]) + + assert mcp._skill_mcp_locations(state) == [] + assert mcp.skill_locations_for_client(entry, "claude") == [] + + def test_per_client_scopes_are_independent(self): + entry = mcp._build_skills_entry( + WS, + {"claude": ["common.schema", "claude.only"], "codex": ["common.schema"]}, + ["claude", "codex"], + ) + + assert mcp.skill_locations_for_client(entry, "claude") == [ + "common.schema", + "claude.only", + ] + assert mcp.skill_locations_for_client(entry, "codex") == ["common.schema"] + + def test_legacy_flat_scope_mirrors_to_every_client(self): + entry = { + "kind": mcp.SKILLS_MCP_KIND, + "skill_locations": ["a.b", "c.d"], + "clients": ["claude", "codex"], + } + + assert mcp.skill_locations_for_client(entry, "claude") == ["a.b", "c.d"] + assert mcp.skill_locations_for_client(entry, "codex") == ["a.b", "c.d"] + class TestUnionLocations: def test_appends_new_after_existing(self): @@ -2111,7 +2174,9 @@ class TestAddSkillsCommand: than replacing it (unlike `configure_skills_mcp_command`).""" def test_unions_into_existing_scope(self, monkeypatch): - state = _skills_state(mcp._resolve_skills_mcp_servers(WS, ["claude"], ["A.a"], [])) + state = _skills_state( + mcp._resolve_skills_mcp_servers(WS, ["claude"], _by_client(["claude"], ["A.a"]), []) + ) _stub_location_base(monkeypatch, state) monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) monkeypatch.setattr(mcp, "save_state", lambda s: None) @@ -2121,7 +2186,11 @@ def test_unions_into_existing_scope(self, monkeypatch): assert _find_skills(state["mcp_servers"])[0]["skill_locations"] == ["A.a", "B.b"] def test_existing_schema_leaves_scope_unchanged(self, monkeypatch): - state = _skills_state(mcp._resolve_skills_mcp_servers(WS, ["claude"], ["A.a", "B.b"], [])) + state = _skills_state( + mcp._resolve_skills_mcp_servers( + WS, ["claude"], _by_client(["claude"], ["A.a", "B.b"]), [] + ) + ) _stub_location_base(monkeypatch, state) monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) monkeypatch.setattr(mcp, "save_state", lambda s: None) @@ -2162,7 +2231,9 @@ def test_registers_bare_route_when_none_exists(self, monkeypatch): def test_preserves_prior_mcp_location_set(self, monkeypatch): self._stub(monkeypatch) - prior = mcp._resolve_skills_mcp_servers(WS, ["claude"], ["X.x", "Y.y"], []) + prior = mcp._resolve_skills_mcp_servers( + WS, ["claude"], _by_client(["claude"], ["X.x", "Y.y"]), [] + ) state = _skills_state(prior) mcp.register_schemaless_skills_connection(state, WS, None, ["claude"]) @@ -2187,7 +2258,8 @@ def test_multiple_schemas_joined_with_and(self): class TestPrintSkillsSummary: def _entry(self, locations): - return mcp._resolve_skills_mcp_servers(WS, ["claude", "codex"], locations, [])[0] + clients = ["claude", "codex"] + return mcp._resolve_skills_mcp_servers(WS, clients, _by_client(clients, locations), [])[0] def test_reports_scoped_connection(self, capsys): mcp._print_skills_summary(self._entry(["main.default"])) @@ -2270,7 +2342,9 @@ def test_removes_skills_registry_across_its_clients(self, monkeypatch): ) monkeypatch.setattr(mcp, "restore_file", lambda *a, **kw: False) - skills_entry = mcp._resolve_skills_mcp_servers(WS, ["claude", "codex"], ["a.b"], [])[0] + skills_entry = mcp._resolve_skills_mcp_servers( + WS, ["claude", "codex"], _by_client(["claude", "codex"], ["a.b"]), [] + )[0] mcp.revert_mcp_configs({"mcp_servers": [skills_entry]}) assert removed == [ @@ -2284,7 +2358,9 @@ def test_drops_foreign_workspace_skills_entry(self, monkeypatch): removed: list[tuple[str, str]] = [] saved_states: list[dict] = [] foreign = "https://other.databricks.com" - skills_entry = mcp._resolve_skills_mcp_servers(foreign, ["claude"], ["a.b"], [])[0] + skills_entry = mcp._resolve_skills_mcp_servers( + foreign, ["claude"], _by_client(["claude"], ["a.b"]), [] + )[0] # The skills URL carries a `?schema=` query; its host must still parse. assert mcp._mcp_entry_url_host(skills_entry) == "other.databricks.com" state = {"mcp_servers": [skills_entry]} From 4a160bf06b5162df33989cfd58244e1465941e74 Mon Sep 17 00:00:00 2001 From: Xiang Shen Date: Fri, 11 Sep 2026 00:15:46 +0000 Subject: [PATCH 2/2] Register skills MCP for all clients in one concurrent batch Extract the per-client work runner out of apply_mcp_server_changes into _run_client_work, and add apply_skills_mcp_changes so _update_skills_mcp registers every client in a single concurrent batch (one spinner), each with its own scoped URL, instead of the serial per-client loop that issued one apply_mcp_server_changes call (and spinner) per client. Co-authored-by: Isaac --- src/ucode/mcp.py | 111 ++++++++++++++++++++++++++-------------------- tests/test_mcp.py | 45 +++++++++++++++++++ 2 files changed, 108 insertions(+), 48 deletions(-) diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 47f5262a..de19d4db 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -1383,9 +1383,30 @@ def apply_mcp_server_changes( ) changed = True + _run_client_work(work) + return changed + + +class _Counter: + """Thread-safe monotonic counter for cross-thread progress reporting.""" + + def __init__(self) -> None: + self._value = 0 + self._lock = threading.Lock() + + def increment(self) -> None: + with self._lock: + self._value += 1 + + def value(self) -> int: + with self._lock: + return self._value + + +def _run_client_work(work: dict[str, list[Callable[[], object]]]) -> None: total_ops = sum(len(ops) for ops in work.values()) if total_ops == 0: - return changed + return completed = _Counter() @@ -1404,24 +1425,6 @@ def message() -> str: for future in as_completed(futures): future.result() - return changed - - -class _Counter: - """Thread-safe monotonic counter for cross-thread progress reporting.""" - - def __init__(self) -> None: - self._value = 0 - self._lock = threading.Lock() - - def increment(self) -> None: - with self._lock: - self._value += 1 - - def value(self) -> int: - with self._lock: - return self._value - def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None: installed = set(available_mcp_clients()) @@ -2199,6 +2202,39 @@ def _print_skills_summary(entry: dict) -> None: ) +def apply_skills_mcp_changes( + original_entry: dict | None, + working_entry: dict, + clients: list[str], + workspace: str, + profile: str | None = None, + *, + use_pat: bool = False, +) -> bool: + """Register the skills connection for every client in one concurrent batch, each with its own scoped URL.""" + configured_before = set(original_entry.get("clients") or []) if original_entry else set() + work: dict[str, list[Callable[[], object]]] = {} + changed = False + for client in clients: + locations = skill_locations_for_client(working_entry, client) + unchanged = ( + client in configured_before + and skill_locations_for_client(original_entry, client) == locations + ) + if unchanged: + continue + url = build_skills_mcp_url(workspace, locations) + work[client] = [ + lambda c=client, u=url: configure_client_mcp_server( + c, SKILLS_MCP_SERVER_NAME, u, workspace, profile, use_pat=use_pat, always_load=True + ) + ] + changed = True + + _run_client_work(work) + return changed + + def _update_skills_mcp( state: dict, workspace: str, @@ -2217,35 +2253,14 @@ def _update_skills_mcp( if working_entry is None: raise RuntimeError("Failed to build the Skills MCP connection.") - changed = False - for client in clients: - working_view = [ - _build_skills_entry( - workspace, - {client: skill_locations_for_client(working_entry, client)}, - [client], - ) - ] - original_view = [] - if original_entry is not None and client in (original_entry.get("clients") or []): - original_view = [ - _build_skills_entry( - workspace, - {client: skill_locations_for_client(original_entry, client)}, - [client], - ) - ] - changed = ( - apply_mcp_server_changes( - original_view, - working_view, - [client], - workspace, - profile, - use_pat=bool(state.get("use_pat")) if use_pat is None else use_pat, - ) - or changed - ) + changed = apply_skills_mcp_changes( + original_entry, + working_entry, + clients, + workspace, + profile, + use_pat=bool(state.get("use_pat")) if use_pat is None else use_pat, + ) if changed or original != working: state["mcp_servers"] = working save_state(state) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index bf289465..b1f10ee4 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -547,6 +547,51 @@ def test_no_ops_returns_false_without_spinner(self, monkeypatch): assert mcp.apply_mcp_server_changes(servers, servers, ["claude"], WS) is False +class TestApplySkillsMcpChanges: + def _entry(self, by_client): + return mcp._build_skills_entry(WS, by_client, list(by_client)) + + def test_divergent_scopes_configure_each_client_in_one_batch(self, monkeypatch): + configured: list[tuple[str, str, object]] = [] + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, *a, **kw: ( + configured.append((client, url, kw.get("always_load"))) or [] + ), + ) + batches: list[list[str]] = [] + run = mcp._run_client_work + monkeypatch.setattr( + mcp, "_run_client_work", lambda work: batches.append(sorted(work)) or run(work) + ) + + working = self._entry({"claude": ["a.b"], "codex": ["c.d"]}) + changed = mcp.apply_skills_mcp_changes(None, working, ["claude", "codex"], WS) + + assert changed is True + assert batches == [["claude", "codex"]] + urls = {client: url for client, url, _ in configured} + assert urls["claude"] == mcp.build_skills_mcp_url(WS, ["a.b"]) + assert urls["codex"] == mcp.build_skills_mcp_url(WS, ["c.d"]) + assert all(always_load is True for *_, always_load in configured) + + def test_skips_clients_whose_scope_is_unchanged(self, monkeypatch): + configured: list[str] = [] + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, *a, **kw: configured.append(client) or [], + ) + original = self._entry({"claude": ["a.b"], "codex": ["c.d"]}) + working = self._entry({"claude": ["a.b"], "codex": ["c.d", "e.f"]}) + + changed = mcp.apply_skills_mcp_changes(original, working, ["claude", "codex"], WS) + + assert changed is True + assert configured == ["codex"] + + class TestConfigureMcpCommand: def test_skips_existing_server_state_by_name(self, monkeypatch): saved_states: list[dict] = []