diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index f855d086..4286c4e7 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_LOCATION_OVERRIDES_KEY = "skill_location_overrides" # 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",) @@ -1210,22 +1211,38 @@ def apply_managed_skills( ] if not desired and not prev_managed: return [] - # Preserve the developer's own locations, drop previously-managed ones no longer in the config, - # and add the current managed set. dict.fromkeys dedupes while keeping first-seen order. - current = _skill_mcp_locations(state) - developer_own = [loc for loc in current if loc not in prev_managed] - new_locations = list(dict.fromkeys([*developer_own, *desired])) - - original = list(state.get("mcp_servers") or []) - working = _resolve_skills_mcp_servers(workspace, [tool], new_locations, original) - changed = apply_mcp_server_changes( - original, working, [tool], workspace, profile, use_pat=use_pat - ) - if not (changed or original != working or prev_managed != desired): + entry = _skills_entry(list(state.get("mcp_servers") or [])) + default_locations = _skill_mcp_locations(state) + overrides = _skill_location_overrides(entry) + current = skill_locations_for_client(entry, tool) + entry_clients = set((entry or {}).get("clients") or []) + if not overrides and entry_clients <= {tool}: + # A single-client connection needs no per-agent additions until another client joins it. + developer_own = [loc for loc in default_locations if loc not in prev_managed] + default_locations = _union_locations(developer_own, desired) + else: + # Managed locations belong to this agent without absorbing the shared default into its list. + developer_own = [loc for loc in overrides.get(tool, []) if loc not in prev_managed] + _set_skill_location_override(overrides, tool, _union_locations(developer_own, desired)) + new_locations = _union_locations(default_locations, overrides.get(tool, [])) + needs_registration = tool not in entry_clients + if current == new_locations and prev_managed == desired and not needs_registration: return [] - state["mcp_servers"] = working + state["managed_skill_locations"] = desired - save_state(state) + if current != new_locations or needs_registration: + _update_skills_mcp( + state, + workspace, + profile, + [tool], + default_locations, + location_overrides=overrides, + print_summary=False, + use_pat=use_pat, + ) + else: + save_state(state) return desired @@ -2104,17 +2121,67 @@ 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_location_overrides(entry: dict | None) -> dict[str, list[str]]: + """Read normalized per-client skill additions from a persisted entry.""" + raw = (entry or {}).get(SKILL_LOCATION_OVERRIDES_KEY) + if not isinstance(raw, dict): + return {} return { + client: _dedupe_locations(locations) + for client, locations in raw.items() + if client in MCP_CLIENTS and isinstance(locations, list) + } + + +def skill_locations_for_client(entry: dict | None, client: str) -> list[str]: + """Return one client's effective skills scope from a persisted skills entry.""" + raw_default = (entry or {}).get("skill_locations") + default = _dedupe_locations(raw_default if isinstance(raw_default, list) else []) + additions = _skill_location_overrides(entry).get(client, []) + return _union_locations(default, additions) + + +def _set_skill_location_override( + overrides: dict[str, list[str]], client: str, locations: list[str] +) -> None: + """Persist one client's skill additions independently from the shared default.""" + # Explicit client additions remain durable even when the shared default currently contains them. + overrides[client] = _dedupe_locations(locations) + + +def _build_skills_entry( + workspace: str, + locations: list[str], + clients: list[str], + location_overrides: dict[str, list[str]] | None = None, +) -> dict: + """Build the skills-registry entry with a shared scope and client additions.""" + default = _dedupe_locations(locations) + normalized_overrides: dict[str, list[str]] = {} + for client, client_locations in (location_overrides or {}).items(): + if client in MCP_CLIENTS: + _set_skill_location_override(normalized_overrides, client, client_locations) + entry: dict = { "name": SKILLS_MCP_SERVER_NAME, "kind": SKILLS_MCP_KIND, - "skill_locations": list(locations), - "url": build_skills_mcp_url(workspace, locations), + "skill_locations": default, + "url": build_skills_mcp_url(workspace, default), "auth": "proxy", "clients": clients, } + if normalized_overrides: + entry[SKILL_LOCATION_OVERRIDES_KEY] = normalized_overrides + return entry + + +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( @@ -2122,6 +2189,7 @@ def _resolve_skills_mcp_servers( clients: list[str], locations: list[str], original_servers: list[dict], + location_overrides: dict[str, list[str]] | None = None, ) -> list[dict]: """Rebuild the MCP server list around exactly one skills entry. @@ -2131,14 +2199,17 @@ 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) + overrides = ( + _skill_location_overrides(prior) if location_overrides is None else location_overrides + ) 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, merged, overrides)] def _join_with_and(items: list[str]) -> str: @@ -2153,6 +2224,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 = [ @@ -2163,9 +2240,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()), list(entry.get("skill_locations") or [])) + 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." @@ -2173,22 +2265,65 @@ 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: list[str], + *, + location_overrides: dict[str, list[str]] | None = None, + print_summary: bool = True, + use_pat: bool | None = None, +) -> bool: + """Persist one skills entry and update only clients whose effective URL 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, original, location_overrides + ) + 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: + original_view = [] + if original_entry is not None and client in (original_entry.get("clients") or []): + original_view = [ + _build_skills_entry( + workspace, + skill_locations_for_client(original_entry, client), + [client], + ) + ] + working_view = [ + _build_skills_entry( + workspace, + skill_locations_for_client(working_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.""" + """Replace the shared skill scope while preserving per-client additions.""" state = load_state() workspace, profile, clients = setup_mcp_clients(state, "Skills MCP") _update_skills_mcp(state, workspace, profile, clients, locations) @@ -2197,8 +2332,9 @@ def configure_skills_mcp_command(locations: list[str]) -> int: 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( @@ -2213,6 +2349,7 @@ def register_schemaless_skills_connection( 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: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index d58062b6..1188d42a 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2459,6 +2459,25 @@ 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_replacing_default_preserves_client_additions(self, monkeypatch): + saved_states: list[dict] = [] + prior = mcp._resolve_skills_mcp_servers( + WS, + ["claude"], + ["old.default"], + [], + {"claude": ["claude.addition"]}, + ) + _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 entry["skill_locations"] == ["new.default"] + assert entry[mcp.SKILL_LOCATION_OVERRIDES_KEY] == {"claude": ["claude.addition"]} + def test_preserves_mcp_service_entries_across_set(self, monkeypatch): saved_states: list[dict] = [] service_entry = { @@ -2488,6 +2507,60 @@ 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_client_override_extends_legacy_default(self): + entry = mcp._build_skills_entry( + WS, + ["common.schema"], + ["claude", "codex"], + {"claude": ["claude.only"]}, + ) + + assert mcp.skill_locations_for_client(entry, "claude") == [ + "common.schema", + "claude.only", + ] + assert mcp.skill_locations_for_client(entry, "codex") == ["common.schema"] + + def test_equal_override_remains_durable_when_default_changes(self): + entry = mcp._build_skills_entry( + WS, + ["common.schema"], + ["claude", "codex"], + {"codex": ["common.schema"]}, + ) + + assert entry[mcp.SKILL_LOCATION_OVERRIDES_KEY] == {"codex": ["common.schema"]} + + changed_default = mcp._build_skills_entry( + WS, + [], + ["claude", "codex"], + mcp._skill_location_overrides(entry), + ) + assert mcp.skill_locations_for_client(changed_default, "claude") == [] + assert mcp.skill_locations_for_client(changed_default, "codex") == ["common.schema"] + + def test_default_and_override_are_deduped_into_effective_scope(self): + entry = { + "skill_locations": ["default.schema", "default.schema"], + mcp.SKILL_LOCATION_OVERRIDES_KEY: { + "codex": ["default.schema", "agent.schema", "agent.schema", ""], + }, + } + + assert mcp.skill_locations_for_client(entry, "claude") == ["default.schema"] + assert mcp.skill_locations_for_client(entry, "codex") == [ + "default.schema", + "agent.schema", + ] + class TestUnionLocations: def test_appends_new_after_existing(self): @@ -2883,6 +2956,31 @@ def test_registers_managed_locations_for_the_launching_tool(self, monkeypatch): assert entry["clients"] == ["claude"] assert state["managed_skill_locations"] == ["cat.sch"] + def test_keeps_managed_locations_in_the_launching_clients_additions(self, monkeypatch): + self._patch_apply(monkeypatch) + entry = mcp._resolve_skills_mcp_servers( + WS, + ["claude", "codex"], + ["shared.default"], + [], + {"codex": ["codex.own"]}, + )[0] + state = {"workspace": WS, "mcp_servers": [entry]} + + applied = mcp.apply_managed_skills(state, self._managed("managed.schema"), "claude", WS) + + assert applied == ["managed.schema"] + entry = self._skills_entry(state["mcp_servers"]) + assert entry["skill_locations"] == ["shared.default"] + assert entry[mcp.SKILL_LOCATION_OVERRIDES_KEY] == { + "codex": ["codex.own"], + "claude": ["managed.schema"], + } + assert mcp.skill_locations_for_client(entry, "claude") == [ + "shared.default", + "managed.schema", + ] + def test_preserves_developer_locations_and_drops_removed_managed_ones(self, monkeypatch): self._patch_apply(monkeypatch) # The developer configured `mine.own`; a prior launch applied `old.managed`, now dropped from @@ -2948,6 +3046,21 @@ def test_unchanged_managed_set_returns_empty(self, monkeypatch): # Same config, same tool already registered: no change, so no note-worthy locations returned. assert mcp.apply_managed_skills(state, self._managed("cat.sch"), "claude", WS) == [] + def test_registers_a_new_client_when_its_scope_already_matches(self, monkeypatch): + self._patch_apply(monkeypatch) + entry = mcp._resolve_skills_mcp_servers(WS, ["claude"], ["cat.sch"], [])[0] + state = { + "workspace": WS, + "managed_skill_locations": ["cat.sch"], + "mcp_servers": [entry], + } + + applied = mcp.apply_managed_skills(state, self._managed("cat.sch"), "codex", WS) + + assert applied == ["cat.sch"] + entry = self._skills_entry(state["mcp_servers"]) + assert entry["clients"] == ["claude", "codex"] + def test_non_client_tool_returns_empty(self, monkeypatch): monkeypatch.setattr( mcp, "apply_mcp_server_changes", lambda *a, **k: pytest.fail("should not apply")