diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 59b4fb8a..cff0f108 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -101,6 +101,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",) @@ -2047,23 +2048,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. @@ -2074,14 +2126,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: @@ -2096,6 +2148,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 = [ @@ -2106,9 +2164,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." @@ -2116,32 +2189,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( @@ -2149,13 +2266,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: @@ -2166,9 +2285,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 5e1829bc..c28f0327 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2336,9 +2336,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 @@ -2364,7 +2370,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) @@ -2378,7 +2384,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 @@ -2390,7 +2398,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): @@ -2439,7 +2449,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())) @@ -2450,7 +2462,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())) @@ -2459,6 +2471,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 = { @@ -2467,7 +2494,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())) @@ -2481,13 +2510,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): @@ -2508,7 +2571,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) @@ -2518,7 +2583,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) @@ -2559,7 +2628,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"]) @@ -2584,7 +2655,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"])) @@ -2667,7 +2739,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 == [ @@ -2681,7 +2755,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]}