From 486e324a397939e731dd620581a3d89f35139050 Mon Sep 17 00:00:00 2001 From: Arthur Jenoudet Date: Wed, 2 Sep 2026 17:34:40 +0000 Subject: [PATCH 1/3] skill: add additive skill configuration Add the ucode skill add command for additive MCP schema registration and skill downloads, including named and fully-qualified skill selection. --- README.md | 21 +++++++++ src/ucode/cli.py | 98 +++++++++++++++++++++++++++++++++++++++ src/ucode/mcp.py | 19 ++++++++ tests/test_cli.py | 115 ++++++++++++++++++++++++++++++++++++++++++++++ tests/test_mcp.py | 49 ++++++++++++++++++++ 5 files changed, 302 insertions(+) diff --git a/README.md b/README.md index 661ab410..26cdafb8 100644 --- a/README.md +++ b/README.md @@ -222,6 +222,24 @@ ucode configure skills --location main.default,ml.prod --mcp Each run prints the registered server, its URL, the configured agents, and its tools, and reminds you to run `ucode ` (existing agent sessions need a restart before the MCP tools load). +#### 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. + +```bash +# Add schemas to the skills MCP scope, keeping any already configured. +ucode skill add --location main.default,ml.prod --mcp + +# Download a schema's skills to disk, keeping existing downloads. +ucode skill add --location main.default + +# Download a named subset, by bare name (with --location) or fully-qualified name. +ucode skill add --location main.default --skills my-skill,other-skill +ucode skill add --skills main.default.my-skill,main.default.other-skill +``` + ### Managed config for a workspace (admins) Author the coding config your developers pick up automatically, instead of asking each of them to @@ -356,6 +374,9 @@ The output looks like: | `ucode configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | | `ucode configure skills --location main.default --skill my-skill` | Download only the named skill(s) from a schema (comma-separated for several) | | `ucode configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | +| `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` | 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) | | `ucode setup mcps` | Add or change the managed config's MCP servers | | `ucode setup skills [--location a.b,c.d]` | Add or change the managed config's skills | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6a1662ae..12142c1b 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -97,6 +97,7 @@ MCP_CLIENTS, SKILLS_MCP_KIND, add_mcp_command, + add_skills_command, apply_managed_mcp_servers, apply_managed_skills, configure_mcp_command, @@ -1159,6 +1160,8 @@ def revert() -> int: app.add_typer(configure_app, name="configure", help="Configure workspace and tool settings.") mcp_app = typer.Typer(add_completion=False, no_args_is_help=True) app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ucode.") +skill_app = typer.Typer(add_completion=False, no_args_is_help=True) +app.add_typer(skill_app, name="skill", help="Databricks Skills for your coding tools.") setup_app = typer.Typer(add_completion=False, no_args_is_help=False) app.add_typer( setup_app, @@ -1312,6 +1315,101 @@ def mcp_web_search_cmd() -> None: serve() +@skill_app.command("add") +def skills_add( + location: Annotated[ + str | None, + typer.Option( + "--location", help="Comma-separated `.` skill scopes to add." + ), + ] = None, + mcp: Annotated[ + bool, + typer.Option( + "--mcp", + help="Add the schemas to the skills MCP connection's scope instead of downloading.", + ), + ] = False, + path: Annotated[ + str | None, + typer.Option( + "--path", + help="(download) Existing absolute dir to download into; defaults to your home dir.", + ), + ] = None, + skills: Annotated[ + str | None, + typer.Option( + "--skills", + help="(download) Download only this comma-separated subset of skills instead of " + "every skill in the schema. Bare securable names (e.g. `my-skill`) need a single " + "--location; fully-qualified `..` names work on their own. " + "Not valid with --mcp.", + ), + ] = None, +) -> None: + """Add Databricks Skills to your coding tools, keeping any already configured. + + With ``--mcp``, adds the given schemas to the skills MCP connection's scope. + Otherwise downloads each schema's skills to disk (under ``--path``, or your home + dir), keeping already-downloaded skills. ``--skills`` narrows a download to a + subset of one schema's skills, by bare name (with ``--location``) or + fully-qualified ``..``. + """ + try: + locations = _parse_skill_locations(location) + requested_skills = ( + None if skills is None else {s.strip() for s in skills.split(",") if s.strip()} + ) + if mcp and path is not None: + raise RuntimeError("--path is not supported when using --mcp") + if mcp and requested_skills is not None: + raise RuntimeError("--skills is not supported when using --mcp") + if requested_skills is not None and not locations: + schemas = {".".join(s.split(".")[:2]) for s in requested_skills if s.count(".") >= 2} + bare = sorted(s for s in requested_skills if s.count(".") < 2) + if bare: + raise RuntimeError( + "--skills short names need --location (or pass full names like " + f"`..`): {', '.join(bare)}" + ) + if len(schemas) != 1: + raise RuntimeError( + "--skills without --location must all share one `.` " + f"(got: {', '.join(sorted(schemas)) or 'none'}); pass --location instead." + ) + locations = list(schemas) + if not locations: + raise RuntimeError("--location is required for `ucode skill add`.") + if requested_skills is not None and len(locations) != 1: + raise RuntimeError( + f"--skills requires a single --location (got: {', '.join(locations)})." + ) + mismatched_skills = sorted( + skill + for skill in requested_skills or set() + if skill.count(".") >= 2 and ".".join(skill.split(".")[:2]) != locations[0] + ) + if mismatched_skills: + raise RuntimeError( + f"--skills entries must match --location `{locations[0]}` " + f"(got: {', '.join(mismatched_skills)})." + ) + selected_skills = ( + None if requested_skills is None else {s.split(".")[-1] for s in requested_skills} + ) + if mcp: + add_skills_command(locations) + else: + configure_skills_download_command(locations, path=path, skills=selected_skills) + except (RuntimeError, ValueError) as exc: + print_err(str(exc)) + raise typer.Exit(1) from None + except KeyboardInterrupt: + print_err("Interrupted.") + raise typer.Exit(130) from None + + @app.command("mcp-proxy", hidden=True) def mcp_proxy_cmd( url: Annotated[ diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 75e3699d..f855d086 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -2210,3 +2210,22 @@ def register_schemaless_skills_connection( ``--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)) + + +def _union_locations(base: list[str], new: list[str]) -> list[str]: + have = set(base) + merged = list(base) + for location in new: + if location not in have: + merged.append(location) + have.add(location) + return merged + + +def add_skills_command(locations: list[str]) -> 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) + return 0 diff --git a/tests/test_cli.py b/tests/test_cli.py index 8bb1ea9b..8ff10329 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1047,6 +1047,121 @@ def test_path_without_location_exit_1(self): mock_download.assert_not_called() +class TestSkillsAddCommand: + """`ucode skill add` is the additive sibling of `configure skills`: `--mcp` + unions schemas into the connection scope, the default mode downloads.""" + + def test_mcp_flag_unions_locations(self): + with patch("ucode.cli.add_skills_command") as mock_add: + result = runner.invoke(app, ["skill", "add", "--location", "a.b", "--mcp"]) + assert result.exit_code == 0, result.output + mock_add.assert_called_once_with(["a.b"]) + + def test_comma_location_yields_multiple_schemas(self): + with patch("ucode.cli.add_skills_command") as mock_add: + result = runner.invoke(app, ["skill", "add", "--location", "a.b, c.d", "--mcp"]) + assert result.exit_code == 0, result.output + mock_add.assert_called_once_with(["a.b", "c.d"]) + + def test_default_mode_dispatches_download(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke(app, ["skill", "add", "--location", "a.b", "--path", "/tmp/s"]) + assert result.exit_code == 0, result.output + mock_download.assert_called_once_with(["a.b"], path="/tmp/s", skills=None) + + def test_skill_filter_dispatches_download_with_subset(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke(app, ["skill", "add", "--location", "a.b", "--skills", "s1, s2"]) + assert result.exit_code == 0, result.output + mock_download.assert_called_once_with(["a.b"], path=None, skills={"s1", "s2"}) + + def test_fully_qualified_skills_derive_location(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke(app, ["skill", "add", "--skills", "a.b.s1, a.b.s2"]) + assert result.exit_code == 0, result.output + mock_download.assert_called_once_with(["a.b"], path=None, skills={"s1", "s2"}) + + def test_fully_qualified_skills_match_explicit_location(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke( + app, ["skill", "add", "--location", "a.b", "--skills", "a.b.s1, a.b.s2"] + ) + assert result.exit_code == 0, result.output + mock_download.assert_called_once_with(["a.b"], path=None, skills={"s1", "s2"}) + + def test_fully_qualified_skills_must_match_explicit_location(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke( + app, ["skill", "add", "--location", "a.b", "--skills", "c.d.s1"] + ) + assert result.exit_code == 1 + assert "must match --location `a.b`" in _strip_ansi(result.output) + mock_download.assert_not_called() + + def test_bare_skills_without_location_exit_1(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke(app, ["skill", "add", "--skills", "s1"]) + assert result.exit_code == 1 + assert "--skills short names need --location" in _strip_ansi(result.output) + mock_download.assert_not_called() + + def test_fully_qualified_skills_across_schemas_exit_1(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke(app, ["skill", "add", "--skills", "a.b.s1, c.d.s2"]) + assert result.exit_code == 1 + assert "must all share one" in _strip_ansi(result.output) + mock_download.assert_not_called() + + def test_without_location_exit_1(self): + with ( + patch("ucode.cli.add_skills_command") as mock_add, + patch("ucode.cli.configure_skills_download_command") as mock_download, + ): + result = runner.invoke(app, ["skill", "add"]) + assert result.exit_code == 1 + assert "--location is required" in _strip_ansi(result.output) + mock_add.assert_not_called() + mock_download.assert_not_called() + + def test_skill_with_mcp_exit_1(self): + with ( + patch("ucode.cli.add_skills_command") as mock_add, + patch("ucode.cli.configure_skills_download_command") as mock_download, + ): + result = runner.invoke( + app, ["skill", "add", "--location", "a.b", "--mcp", "--skills", "s1"] + ) + assert result.exit_code == 1 + assert "--skills" in _strip_ansi(result.output) + mock_add.assert_not_called() + mock_download.assert_not_called() + + def test_path_with_mcp_exit_1(self): + with patch("ucode.cli.add_skills_command") as mock_add: + result = runner.invoke( + app, ["skill", "add", "--location", "a.b", "--mcp", "--path", "/tmp/s"] + ) + assert result.exit_code == 1 + assert "--path" in _strip_ansi(result.output) + mock_add.assert_not_called() + + def test_skill_with_multiple_locations_exit_1(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke( + app, ["skill", "add", "--location", "a.b, c.d", "--skills", "s1"] + ) + assert result.exit_code == 1 + assert "--skills requires a single --location" in _strip_ansi(result.output) + mock_download.assert_not_called() + + def test_malformed_location_exit_1(self): + with patch("ucode.cli.add_skills_command") as mock_add: + result = runner.invoke(app, ["skill", "add", "--location", "a.b.c", "--mcp"]) + assert result.exit_code == 1 + assert "--location" in _strip_ansi(result.output) + mock_add.assert_not_called() + + class TestApplyManagedSkills: """The launch path both registers the skills MCP connection and downloads bundles to disk.""" diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 69c60d0a..d58062b6 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2489,6 +2489,55 @@ def test_empty_when_no_skills_entry(self): assert mcp._skill_mcp_locations(_skills_state()) == [] +class TestUnionLocations: + def test_appends_new_after_existing(self): + assert mcp._union_locations(["a.b"], ["c.d"]) == ["a.b", "c.d"] + + def test_drops_locations_already_present(self): + assert mcp._union_locations(["a.b", "c.d"], ["c.d", "e.f"]) == ["a.b", "c.d", "e.f"] + + def test_empty_base_returns_new(self): + assert mcp._union_locations([], ["a.b", "c.d"]) == ["a.b", "c.d"] + + def test_drops_duplicate_new_locations(self): + assert mcp._union_locations(["a.b"], ["c.d", "c.d"]) == ["a.b", "c.d"] + + +class TestAddSkillsCommand: + """`ucode skill add --mcp` unions schemas into the connection scope rather + 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"], [])) + _stub_location_base(monkeypatch, state) + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr(mcp, "save_state", lambda s: None) + + assert mcp.add_skills_command(["B.b"]) == 0 + + 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"], [])) + _stub_location_base(monkeypatch, state) + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr(mcp, "save_state", lambda s: None) + + assert mcp.add_skills_command(["A.a"]) == 0 + + assert _find_skills(state["mcp_servers"])[0]["skill_locations"] == ["A.a", "B.b"] + + def test_registers_scope_from_empty_state(self, monkeypatch): + state = _skills_state() + _stub_location_base(monkeypatch, state) + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr(mcp, "save_state", lambda s: None) + + assert mcp.add_skills_command(["A.a"]) == 0 + + assert _find_skills(state["mcp_servers"])[0]["skill_locations"] == ["A.a"] + + class TestRegisterSchemalessSkillsConnection: def _stub(self, monkeypatch): saved_states: list[dict] = [] From d3c4a06c71111422720282caf6b18056bd40c78d Mon Sep 17 00:00:00 2001 From: Arthur Jenoudet Date: Wed, 2 Sep 2026 15:42:15 +0000 Subject: [PATCH 2/3] Add interactive metastore skill picker --- README.md | 7 +- src/ucode/cli.py | 13 ++-- src/ucode/skills_download.py | 95 +++++++++++++++++++++++++ tests/test_cli.py | 14 +++- tests/test_skills_download.py | 129 ++++++++++++++++++++++++++++++++++ 5 files changed, 251 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 26cdafb8..1af5cc34 100644 --- a/README.md +++ b/README.md @@ -226,9 +226,13 @@ you to run `ucode ` (existing agent sessions need a restart before the MC `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. +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. +ucode skill add + # Add schemas to the skills MCP scope, keeping any already configured. ucode skill add --location main.default,ml.prod --mcp @@ -374,6 +378,7 @@ The output looks like: | `ucode configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | | `ucode configure skills --location main.default --skill my-skill` | Download only the named skill(s) from a schema (comma-separated for several) | | `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` | 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) | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 12142c1b..a5d41c32 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -108,6 +108,7 @@ ) from ucode.skills_download import ( configure_skills_download_command, + configure_skills_download_interactive_command, download_managed_skills_on_launch, ) from ucode.smart_routing import v2 as smart_routing_v2 @@ -1352,9 +1353,10 @@ def skills_add( With ``--mcp``, adds the given schemas to the skills MCP connection's scope. Otherwise downloads each schema's skills to disk (under ``--path``, or your home - dir), keeping already-downloaded skills. ``--skills`` narrows a download to a - subset of one schema's skills, by bare name (with ``--location``) or - fully-qualified ``..``. + dir), keeping already-downloaded skills. With no selection flags, opens a searchable + metastore picker. ``--skills`` narrows a download to a subset of one schema's skills, + by bare name (with ``--location``) or fully-qualified + ``..``. """ try: locations = _parse_skill_locations(location) @@ -1380,7 +1382,10 @@ def skills_add( ) locations = list(schemas) if not locations: - raise RuntimeError("--location is required for `ucode skill add`.") + if mcp: + raise RuntimeError("--location is required when using --mcp.") + configure_skills_download_interactive_command(path=path) + return if requested_skills is not None and len(locations) != 1: raise RuntimeError( f"--skills requires a single --location (got: {', '.join(locations)})." diff --git a/src/ucode/skills_download.py b/src/ucode/skills_download.py index 2168b522..dab2daf4 100644 --- a/src/ucode/skills_download.py +++ b/src/ucode/skills_download.py @@ -21,6 +21,7 @@ print_success, print_warning, progress_bar, + prompt_for_multi_selection, prompt_yes_no, ) @@ -52,6 +53,18 @@ class SkillRef: bundle_name: str +@dataclass(frozen=True) +class MetastoreSkill: + """A finalized skill discovered without a schema filter.""" + + full_name: str + ref: SkillRef + + @property + def location(self) -> str: + return self.full_name.rsplit(".", 1)[0] + + def _non_empty_str(value: object) -> str | None: """``value`` when it is a non-empty string, else None.""" return value if isinstance(value, str) and value else None @@ -118,6 +131,58 @@ def list_schema_skills( return refs, None +def list_metastore_skills(workspace: str, token: str) -> tuple[list[MetastoreSkill], str | None]: + """List finalized, downloadable skills visible in the current metastore.""" + hostname = workspace_hostname(workspace) + base_url = f"https://{hostname}/api/2.1/unity-catalog/skills" + + skills: list[MetastoreSkill] = [] + page_token: str | None = None + while True: + url = base_url + if page_token: + url = f"{url}?{urlencode({'page_token': page_token})}" + payload, reason = _http_get_json(url, token, timeout=30) + if payload is None: + return [], reason + data = payload if isinstance(payload, dict) else {} + for skill in data.get("skills") or []: + if not isinstance(skill, dict): + continue + resource_name = _non_empty_str(skill.get("name")) + ref = _skill_ref(skill) + if ref is None: + continue + full_name = resource_name.removeprefix("skills/") if resource_name else None + if full_name is None or full_name.count(".") != 2: + print_warning( + f"Skipping `{resource_name or ''}`: expected a fully-qualified " + "`..` from the skills API." + ) + continue + skills.append(MetastoreSkill(full_name=full_name, ref=ref)) + page_token = data.get("next_page_token") + if not page_token: + return sorted(skills, key=lambda skill: skill.full_name.lower()), None + + +def prompt_for_skill_download( + skills: list[MetastoreSkill], +) -> list[MetastoreSkill] | None: + """Select metastore skills to download, or return None when cancelled.""" + by_name = {skill.full_name: skill for skill in skills} + options = [] + for skill in skills: + label = skill.full_name + if skill.ref.bundle_name != skill.ref.securable_name: + label = f"{label} (bundle: {skill.ref.bundle_name})" + options.append((skill.full_name, label)) + selected = prompt_for_multi_selection("Skills:", options, searchable=True) + if selected is None: + return None + return [by_name[name] for name in selected if name in by_name] + + def list_skill_files( workspace: str, token: str, catalog: str, schema: str, securable: str ) -> tuple[list[str], str | None]: @@ -450,3 +515,33 @@ def configure_skills_download_command( register_schemaless_skills_connection(state, workspace, profile, clients) return 0 + + +def configure_skills_download_interactive_command(*, path: str | None) -> int: + """Discover metastore skills, let the user select some, and download them.""" + state = load_state() + workspace, profile, clients = setup_mcp_clients(state, "Add Skills") + token = get_databricks_token(workspace, profile) + + available, reason = list_metastore_skills(workspace, token) + if reason: + raise RuntimeError(f"Could not list workspace skills: {reason}.") + if not available: + print_note("No finalized skills are available to download in this metastore.") + return 0 + + selected = prompt_for_skill_download(available) + if selected is None: + return 0 + if not selected: + print_note("No skills selected. Press space to toggle an item, then enter to download.") + return 0 + + selected_by_location: dict[str, set[str]] = {} + for skill in selected: + selected_by_location.setdefault(skill.location, set()).add(skill.ref.securable_name) + for location, securable_names in selected_by_location.items(): + download_skills(workspace, token, [location], path, securable_names) + + register_schemaless_skills_connection(state, workspace, profile, clients) + return 0 diff --git a/tests/test_cli.py b/tests/test_cli.py index 8ff10329..c8bde8c4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1112,16 +1112,26 @@ def test_fully_qualified_skills_across_schemas_exit_1(self): assert "must all share one" in _strip_ansi(result.output) mock_download.assert_not_called() - def test_without_location_exit_1(self): + def test_without_location_opens_download_picker(self): with ( patch("ucode.cli.add_skills_command") as mock_add, patch("ucode.cli.configure_skills_download_command") as mock_download, + patch("ucode.cli.configure_skills_download_interactive_command") as mock_interactive, ): result = runner.invoke(app, ["skill", "add"]) + + assert result.exit_code == 0, result.output + mock_add.assert_not_called() + mock_download.assert_not_called() + mock_interactive.assert_called_once_with(path=None) + + def test_mcp_without_location_exit_1(self): + with patch("ucode.cli.add_skills_command") as mock_add: + result = runner.invoke(app, ["skill", "add", "--mcp"]) + assert result.exit_code == 1 assert "--location is required" in _strip_ansi(result.output) mock_add.assert_not_called() - mock_download.assert_not_called() def test_skill_with_mcp_exit_1(self): with ( diff --git a/tests/test_skills_download.py b/tests/test_skills_download.py index c7927019..f333fce0 100644 --- a/tests/test_skills_download.py +++ b/tests/test_skills_download.py @@ -156,6 +156,59 @@ def test_http_failure_propagates_reason(self, monkeypatch): assert reason == "HTTP 500 Server Error" +class TestListMetastoreSkills: + def test_lists_finalized_skills_across_schemas_and_follows_pagination(self, monkeypatch): + pages = [ + { + "skills": [ + { + "name": "skills/ml.prod.triage", + "bundle_name": "triage", + "finalize_time": "t", + }, + {"name": "skills/ml.prod.draft", "bundle_name": "draft"}, + ], + "next_page_token": "next", + }, + { + "skills": [ + { + "name": "skills/main.default.pii", + "bundle_name": "pii-handling", + "finalize_time": "t", + } + ] + }, + ] + urls = [] + + def fake_get(url, token, timeout=30): + urls.append(url) + return pages.pop(0), None + + monkeypatch.setattr(sd, "_http_get_json", fake_get) + + skills, reason = sd.list_metastore_skills(WS, "token") + + assert reason is None + assert skills == [ + sd.MetastoreSkill("main.default.pii", ref("pii", "pii-handling")), + sd.MetastoreSkill("ml.prod.triage", ref("triage")), + ] + assert urls == [ + f"{WS}/api/2.1/unity-catalog/skills", + f"{WS}/api/2.1/unity-catalog/skills?page_token=next", + ] + + def test_http_failure_propagates_reason(self, monkeypatch): + monkeypatch.setattr(sd, "_http_get_json", lambda *a, **k: (None, "HTTP 500 Server Error")) + + skills, reason = sd.list_metastore_skills(WS, "token") + + assert skills == [] + assert reason == "HTTP 500 Server Error" + + class TestListSkillFiles: def test_lists_under_the_skills_place(self, monkeypatch): captured = {} @@ -744,3 +797,79 @@ def test_skills_filter_threads_through(self, monkeypatch): assert calls["download"] == (WS, "token", ["a.b"], None, {"triage"}) assert calls["register"] == (WS, "profile", ["claude"]) + + +class TestConfigureSkillsDownloadInteractiveCommand: + def _stub(self, monkeypatch): + calls: dict[str, object] = {"downloads": []} + state = {"state": True} + monkeypatch.setattr(sd, "load_state", lambda: state) + monkeypatch.setattr( + sd, "setup_mcp_clients", lambda actual, section: (WS, "profile", ["claude"]) + ) + monkeypatch.setattr(sd, "get_databricks_token", lambda ws, profile: "token") + monkeypatch.setattr( + sd, + "download_skills", + lambda ws, token, locations, path, skills: calls["downloads"].append( + (ws, token, locations, path, skills) + ), + ) + monkeypatch.setattr( + sd, + "register_schemaless_skills_connection", + lambda actual, ws, profile, clients: calls.update( + register=(actual, ws, profile, clients) + ), + ) + return state, calls + + def test_downloads_picker_selection_grouped_by_schema(self, monkeypatch): + state, calls = self._stub(monkeypatch) + available = [ + sd.MetastoreSkill("main.default.pii", ref("pii")), + sd.MetastoreSkill("main.default.triage", ref("triage")), + sd.MetastoreSkill("ml.prod.eval", ref("eval")), + ] + monkeypatch.setattr(sd, "list_metastore_skills", lambda *a: (available, None)) + monkeypatch.setattr( + sd, "prompt_for_skill_download", lambda actual: [available[0], available[2]] + ) + + assert sd.configure_skills_download_interactive_command(path="/tmp/project") == 0 + + assert calls["downloads"] == [ + (WS, "token", ["main.default"], "/tmp/project", {"pii"}), + (WS, "token", ["ml.prod"], "/tmp/project", {"eval"}), + ] + assert calls["register"] == (state, WS, "profile", ["claude"]) + + def test_cancel_is_a_noop(self, monkeypatch): + _, calls = self._stub(monkeypatch) + available = [sd.MetastoreSkill("main.default.pii", ref("pii"))] + monkeypatch.setattr(sd, "list_metastore_skills", lambda *a: (available, None)) + monkeypatch.setattr(sd, "prompt_for_skill_download", lambda actual: None) + + assert sd.configure_skills_download_interactive_command(path=None) == 0 + + assert calls["downloads"] == [] + assert "register" not in calls + + def test_picker_is_searchable_and_shows_bundle_name_when_different(self, monkeypatch): + captured = {} + skill = sd.MetastoreSkill("main.default.task", ref("task", "task-triage")) + monkeypatch.setattr( + sd, + "prompt_for_multi_selection", + lambda prompt, options, searchable: ( + captured.update(prompt=prompt, options=options, searchable=searchable) + or ["main.default.task"] + ), + ) + + assert sd.prompt_for_skill_download([skill]) == [skill] + assert captured == { + "prompt": "Skills:", + "options": [("main.default.task", "main.default.task (bundle: task-triage)")], + "searchable": True, + } From 6819458e7d18b8e86bbfeb9ef4afac754fe3fcdf Mon Sep 17 00:00:00 2001 From: Arthur Jenoudet Date: Wed, 2 Sep 2026 15:44:01 +0000 Subject: [PATCH 3/3] Support per-agent skill MCP scopes --- README.md | 12 ++- src/ucode/cli.py | 58 +++++++++--- src/ucode/mcp.py | 223 ++++++++++++++++++++++++++++++++++++++-------- tests/test_cli.py | 45 ++++++++++ tests/test_mcp.py | 39 ++++++++ 5 files changed, 328 insertions(+), 49 deletions(-) diff --git a/README.md b/README.md index 1af5cc34..66480bf8 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 a5d41c32..bcd1eb19 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -105,6 +105,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, @@ -1064,17 +1065,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 {} @@ -1348,6 +1363,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. @@ -1363,10 +1386,17 @@ 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: raise RuntimeError("--skills is not supported when using --mcp") + 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(s.split(".")[:2]) for s in requested_skills if s.count(".") >= 2} bare = sorted(s for s in requested_skills if s.count(".") < 2) @@ -1404,7 +1434,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 c8bde8c4..b8a25877 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1171,6 +1171,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.""" @@ -1323,6 +1345,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):