diff --git a/README.md b/README.md index 2d905346..4324c772 100644 --- a/README.md +++ b/README.md @@ -225,9 +225,11 @@ you to run `ucode ` (existing agent sessions need a restart before the MC #### Add skill scopes without replacing existing ones `ucode skill add` registers skills additively, keeping anything already configured. With `--mcp` it -adds the schemas to the connection's scope, otherwise it downloads their skills to disk. `--skills` -narrows a download to a subset of one schema's skills. With no selection flags, a searchable picker -lists finalized skills visible in the metastore. +adds the schemas to every configured agent's scope, or only to `--agents` when supplied. Agents that +are not configured yet are set up first. Without `--mcp`, it downloads skills to disk; download mode +always writes both directory families and does not accept `--agents`. `--skills` narrows a download +to a subset of one schema's skills. With no selection flags, a searchable picker lists finalized +skills visible in the metastore. ```bash # Browse the metastore and choose skills to download. @@ -236,6 +238,9 @@ ucode skill add # Add schemas to the skills MCP scope, keeping any already configured. ucode skill add --location main.default,ml.prod --mcp +# Add a schema only to selected agents. +ucode skill add --location main.default --mcp --agents claude,codex + # Download a schema's skills to disk, keeping existing downloads. ucode skill add --location main.default @@ -380,6 +385,7 @@ The output looks like: | `ucode configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | | `ucode skill add` | Interactively choose finalized metastore skills to download | | `ucode skill add --location main.default --mcp` | Add schemas to the skills MCP scope, keeping any already configured (additive; never replaces) | +| `ucode skill add --location main.default --mcp --agents claude` | Set up selected agents if needed and add schemas only to their MCP scopes | | `ucode skill add --location main.default` | Download a schema's skills to disk without removing existing downloads | | `ucode skill add --skills main.default.my-skill` | Download a named subset of skills (bare names need `--location`; fully-qualified names stand alone) | | `ucode setup` | Author the managed config's agents and models (workspace admins only) | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 8b997ba5..ab8ea506 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -106,6 +106,7 @@ purge_cross_workspace_mcp_residue, remove_mcp_command, revert_mcp_configs, + skill_locations_for_client, ) from ucode.skills_download import ( configure_skills_download_command, @@ -1068,17 +1069,31 @@ def status() -> int: if not skill_mcp_entry: print_kv("Skills", "not configured") else: - locations = skill_mcp_entry.get("skill_locations") or [] - print_kv( - "Skill MCP Locations", - ", ".join(locations) if locations else "none — utility tools only", - ) - configured_agents = [ - str(MCP_CLIENTS[client]["display"]) - for client in (skill_mcp_entry.get("clients") or []) - if client in MCP_CLIENTS + configured_clients = [ + client for client in (skill_mcp_entry.get("clients") or []) if client in MCP_CLIENTS ] - print_kv("Configured", ", ".join(configured_agents) if configured_agents else "none") + scopes = { + client: skill_locations_for_client(skill_mcp_entry, client) + for client in configured_clients + } + if len({tuple(locations) for locations in scopes.values()}) <= 1: + locations = next( + iter(scopes.values()), list(skill_mcp_entry.get("skill_locations") or []) + ) + print_kv( + "Skill MCP Locations", + ", ".join(locations) if locations else "none — utility tools only", + ) + configured_agents = [ + str(MCP_CLIENTS[client]["display"]) for client in configured_clients + ] + print_kv("Configured", ", ".join(configured_agents) if configured_agents else "none") + else: + for client, locations in scopes.items(): + print_kv( + f"{MCP_CLIENTS[client]['display']} skill MCP locations", + ", ".join(locations) if locations else "none — utility tools only", + ) print_heading("Tracing") tracing = state.get("tracing") or {} @@ -1353,6 +1368,14 @@ def skills_add( "Not valid with --mcp.", ), ] = None, + agents: Annotated[ + str | None, + typer.Option( + "--agents", + help="(--mcp only) Comma-separated coding agents whose skills MCP scope should " + "be updated. Any that aren't configured yet are set up first.", + ), + ] = None, ) -> None: """Add Databricks Skills to your coding tools, keeping any already configured. @@ -1369,6 +1392,11 @@ def skills_add( requested_skills = ( None if skills is None else {s.strip() for s in skills.split(",") if s.strip()} ) + requested_agents = ( + None + if agents is None + else ({agent.strip().lower() for agent in agents.split(",") if agent.strip()} or None) + ) if mcp and path is not None: raise RuntimeError("--path is not supported when using --mcp") if mcp and requested_skills is not None: @@ -1389,6 +1417,8 @@ def skills_add( "`..` values " f"(invalid: {', '.join(sorted(invalid_skills))})." ) + if not mcp and agents is not None: + raise RuntimeError("--agents is only supported when using --mcp") if requested_skills is not None and not locations: schemas = {".".join(parts[:2]) for parts in qualified_skill_parts.values()} bare = sorted(skill for skill in requested_skills if skill not in qualified_skill_parts) @@ -1426,7 +1456,13 @@ def skills_add( None if requested_skills is None else {s.split(".")[-1] for s in requested_skills} ) if mcp: - add_skills_command(locations) + scope = ( + _configure_agents_for_mcp(sorted(requested_agents)) if requested_agents else None + ) + if scope is None: + add_skills_command(locations) + else: + add_skills_command(locations, agents=scope) else: configure_skills_download_command(locations, path=path, skills=selected_skills) except (RuntimeError, ValueError) as exc: diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index f855d086..1f1a49ce 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,36 @@ 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) + # Preserve this agent's own locations, drop previously-managed ones no longer in the config, + # and add the current managed set. Other agents may intentionally have different scopes. + entry = _skills_entry(list(state.get("mcp_servers") or [])) + current = skill_locations_for_client(entry, tool) 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): + if current == new_locations and prev_managed == desired: return [] - state["mcp_servers"] = working + + default_locations = _skill_mcp_locations(state) + overrides = _skill_location_overrides(entry) + entry_clients = set((entry or {}).get("clients") or []) + if not overrides and entry_clients <= {tool}: + default_locations = new_locations + else: + _set_skill_location_override(overrides, tool, new_locations, default_locations) state["managed_skill_locations"] = desired - save_state(state) + if current != new_locations: + _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 +2119,64 @@ 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 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]]: + 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.""" + default = _dedupe_locations(list((entry or {}).get("skill_locations") or [])) + return _skill_location_overrides(entry).get(client, default) + + +def _set_skill_location_override( + overrides: dict[str, list[str]], client: str, locations: list[str], default: list[str] +) -> None: + normalized = _dedupe_locations(locations) + if normalized == default: + overrides.pop(client, None) + else: + overrides[client] = normalized + + +def _build_skills_entry( + workspace: str, + locations: list[str], + clients: list[str], + location_overrides: dict[str, list[str]] | None = None, +) -> dict: + """Build the single skills-registry entry with a common scope and sparse overrides.""" + 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, default) + 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 next((server for server in servers if server.get("kind") == SKILLS_MCP_KIND), None) def _resolve_skills_mcp_servers( @@ -2122,6 +2184,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 +2194,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 +2219,11 @@ 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: + 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 +2234,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,17 +2259,60 @@ 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) + assert working_entry is not None + + 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: @@ -2191,13 +2320,13 @@ def configure_skills_mcp_command(locations: list[str]) -> int: replacing any previous set.""" state = load_state() workspace, profile, clients = setup_mcp_clients(state, "Skills MCP") - _update_skills_mcp(state, workspace, profile, clients, locations) + _update_skills_mcp(state, workspace, profile, clients, locations, location_overrides={}) 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) + entry = _skills_entry(list(state.get("mcp_servers") or [])) return list((entry or {}).get("skill_locations") or []) @@ -2222,10 +2351,34 @@ def _union_locations(base: list[str], new: list[str]) -> list[str]: return merged -def add_skills_command(locations: list[str]) -> int: +def add_skills_command(locations: list[str], agents: set[str] | None = None) -> int: """Add ``locations`` to the skills MCP connection's 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) + workspace, profile, clients = setup_mcp_clients(state, "Add Skills MCP", agents=agents) + entry = _skills_entry(list(state.get("mcp_servers") or [])) + default = _skill_mcp_locations(state) + overrides = _skill_location_overrides(entry) + if agents is None: + merged = _union_locations(default, locations) + overrides = { + client: _union_locations(client_locations, locations) + for client, client_locations in overrides.items() + } + else: + merged = default + for client in clients: + _set_skill_location_override( + overrides, + client, + _union_locations(skill_locations_for_client(entry, client), locations), + default, + ) + _update_skills_mcp( + state, + workspace, + profile, + clients, + merged, + location_overrides=overrides, + ) return 0 diff --git a/tests/test_cli.py b/tests/test_cli.py index 6efa8322..f44e8d2c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1190,6 +1190,28 @@ def test_malformed_location_exit_1(self): assert "--location" in _strip_ansi(result.output) mock_add.assert_not_called() + def test_agents_scope_is_configured_and_forwarded_for_mcp(self): + with ( + patch("ucode.cli._configure_agents_for_mcp", return_value={"claude"}) as configure, + patch("ucode.cli.add_skills_command") as mock_add, + ): + result = runner.invoke( + app, + ["skill", "add", "--location", "a.b", "--mcp", "--agents", "claude"], + ) + + assert result.exit_code == 0, result.output + configure.assert_called_once_with(["claude"]) + mock_add.assert_called_once_with(["a.b"], agents={"claude"}) + + def test_agents_is_rejected_for_download_mode(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke(app, ["skill", "add", "--location", "a.b", "--agents", "claude"]) + + assert result.exit_code == 1 + assert "--agents is only supported when using --mcp" in _strip_ansi(result.output) + mock_download.assert_not_called() + class TestApplyManagedSkills: """The launch path both registers the skills MCP connection and downloads bundles to disk.""" @@ -1342,6 +1364,29 @@ def test_skills_entry_absent_from_per_client_mcp_lines(self): assert "databricks-skill-registry" not in line assert "Skill MCP Locations: main.default" in out + def test_renders_per_agent_locations_when_scopes_diverge(self): + state = { + **MINIMAL_STATE, + "mcp_servers": [ + { + "name": "databricks-skill-registry", + "kind": "skills", + "skill_locations": ["main.default"], + "skill_location_overrides": {"claude": ["main.default", "claude.only"]}, + "url": "https://example.databricks.com/ai-gateway/skills/?schema=main.default", + "auth": "proxy", + "clients": ["claude", "codex"], + } + ], + } + + result = self._run(state) + + assert result.exit_code == 0, result.output + out = _strip_ansi(result.output) + assert "Claude Code skill MCP locations: main.default, claude.only" in out + assert "Codex skill MCP locations: main.default" in out + class TestRevert: def test_reverts_mcp_configs_before_clearing_state(self): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index d58062b6..e604c9c5 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2488,6 +2488,20 @@ def test_empty_when_no_skills_entry(self): assert mcp._skill_mcp_locations(_skills_state([])) == [] assert mcp._skill_mcp_locations(_skills_state()) == [] + def test_client_override_takes_precedence_over_legacy_default(self): + entry = mcp._build_skills_entry( + WS, + ["common.schema"], + ["claude", "codex"], + {"claude": ["common.schema", "claude.only"]}, + ) + + assert mcp.skill_locations_for_client(entry, "claude") == [ + "common.schema", + "claude.only", + ] + assert mcp.skill_locations_for_client(entry, "codex") == ["common.schema"] + class TestUnionLocations: def test_appends_new_after_existing(self): @@ -2537,6 +2551,31 @@ def test_registers_scope_from_empty_state(self, monkeypatch): assert _find_skills(state["mcp_servers"])[0]["skill_locations"] == ["A.a"] + def test_agents_updates_only_selected_client_scope(self, monkeypatch): + configured: list[tuple[str, str]] = [] + prior = mcp._resolve_skills_mcp_servers(WS, ["claude", "codex"], ["A.a"], []) + state = { + "workspace": WS, + "available_tools": ["claude", "codex"], + "mcp_servers": prior, + } + _stub_location_base(monkeypatch, state) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude", "codex"]) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, *a, **kw: configured.append((client, url)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda s: None) + + assert mcp.add_skills_command(["B.b"], agents={"claude"}) == 0 + + entry = _find_skills(state["mcp_servers"])[0] + assert entry["skill_locations"] == ["A.a"] + assert entry[mcp.SKILL_LOCATION_OVERRIDES_KEY] == {"claude": ["A.a", "B.b"]} + assert mcp.skill_locations_for_client(entry, "codex") == ["A.a"] + assert configured == [("claude", f"{WS}/ai-gateway/skills/?schema=A.a&schema=B.b")] + class TestRegisterSchemalessSkillsConnection: def _stub(self, monkeypatch):