diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 61d60592..ef2dac68 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -107,8 +107,10 @@ add_mcp_command, add_skills_command, apply_managed_mcp_servers, + available_mcp_clients, configure_mcp_command, configure_skills_mcp_command, + configured_mcp_clients, purge_cross_workspace_mcp_residue, remove_mcp_command, revert_mcp_configs, @@ -1112,30 +1114,26 @@ def _configure_agents_for_mcp( requested: list[str], *, prompt_optional_updates: bool = True ) -> set[str]: """Ensure the named coding agents are set up (workspace + models) so a - subsequent `ug mcp add` has them as targets, and return their canonical - names. Mirrors `ug configure --agents`: model agents go through + subsequent `ug mcp add` / `ug skill add --mcp` has them as targets, and + return the full canonical name set. Agents already configured are left as-is; + only the rest are bootstrapped. Model agents go through configure_workspace_command (which installs binaries and configures models); Cursor is MCP-only, so it just needs workspace state established and rides along via MCP_ONLY_CLIENTS. Interactive — prompts for the workspace URL on first run.""" - wants_cursor = "cursor" in requested - model_agent_names = ",".join(a for a in requested if a != "cursor") - configured: set[str] = set() - if model_agent_names: - selected_tools = _parse_agents_option(model_agent_names) + scope = {a if a == "cursor" else normalize_tool(a) for a in requested} + ready = set(configured_mcp_clients(load_state(), available_mcp_clients())) + to_bootstrap = scope - ready + model_agents = sorted(a for a in to_bootstrap if a != "cursor") + if model_agents: configure_workspace_command( - selected_tools=selected_tools, prompt_optional_updates=prompt_optional_updates + selected_tools=model_agents, prompt_optional_updates=prompt_optional_updates ) - configured.update(selected_tools) - if wants_cursor: - # Establish workspace state for a Cursor-only run; when model agents were - # configured above the workspace is already set, so Cursor just rides along. - if not model_agent_names: - _configure_shared_workspace_states( - [_prompt_for_configuration(None)], tools=[], force_login=True - ) - configured.add("cursor") - return configured + if "cursor" in to_bootstrap and not model_agents: + _configure_shared_workspace_states( + [_prompt_for_configuration(None)], tools=[], force_login=True + ) + return scope def _configure_optional_setup(state: dict, tools: list[str]) -> None: @@ -1281,6 +1279,15 @@ 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. Without --agents, " + "every configured agent is updated.", + ), + ] = None, ) -> None: """Add Databricks Skills to your coding tools, keeping any already configured. @@ -1296,6 +1303,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: @@ -1316,6 +1328,9 @@ def skills_add( "`..` values " f"(invalid: {', '.join(sorted(invalid_skills))})." ) + # Downloaded skills use shared directory families, so only MCP scopes can be agent-scoped. + 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) @@ -1350,7 +1365,10 @@ 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 + ) + 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 de19d4db..451d6936 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -2310,10 +2310,14 @@ def _union_locations(base: list[str], new: list[str]) -> list[str]: return merged -def add_skills_command(locations: list[str]) -> int: - """Add ``locations`` to every configured client's skill scope, keeping any already configured.""" +def add_skills_command(locations: list[str], agents: set[str] | None = None) -> int: + """Add ``locations`` to each targeted client's skill scope, keeping any already configured. + + ``agents`` (from ``--agents``) scopes the update to that subset of configured clients; omitting + it targets every configured client. This mirrors ``ucode mcp add`` exactly: the client set is + the only thing ``--agents`` changes.""" state = load_state() - workspace, profile, clients = setup_mcp_clients(state, "Add Skills MCP") + workspace, profile, clients = setup_mcp_clients(state, "Add Skills MCP", agents=agents) locations_by_client = _skill_locations_by_client_from_state(state) for client in clients: locations_by_client[client] = _union_locations( diff --git a/tests/test_cli.py b/tests/test_cli.py index d026961d..38fa3671 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1421,13 +1421,13 @@ 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"]) + mock_add.assert_called_once_with(["a.b"], agents=None) 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"]) + mock_add.assert_called_once_with(["a.b", "c.d"], agents=None) def test_default_mode_dispatches_download(self): with patch("ucode.cli.configure_skills_download_command") as mock_download: @@ -1533,6 +1533,70 @@ def test_malformed_location_exit_1(self): assert "--location" in _strip_ansi(result.output) mock_add.assert_not_called() + def test_agents_scope_delegates_to_helper_and_forwards_returned_scope(self): + with ( + patch( + "ucode.cli._configure_agents_for_mcp", return_value={"claude", "codex"} + ) as configure, + patch("ucode.cli.add_skills_command") as mock_add, + ): + result = runner.invoke( + app, + ["skill", "add", "--location", "a.b", "--mcp", "--agents", "codex,claude"], + ) + + assert result.exit_code == 0, result.output + configure.assert_called_once_with(["claude", "codex"]) + mock_add.assert_called_once_with(["a.b"], agents={"claude", "codex"}) + + def test_empty_agents_folds_to_global_scope(self): + with ( + patch("ucode.cli._configure_agents_for_mcp") as configure, + patch("ucode.cli.add_skills_command") as mock_add, + ): + result = runner.invoke( + app, + ["skill", "add", "--location", "a.b", "--mcp", "--agents", ","], + ) + + assert result.exit_code == 0, result.output + configure.assert_not_called() + mock_add.assert_called_once_with(["a.b"], agents=None) + + 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 TestConfigureAgentsForMcp: + def test_bootstraps_only_unconfigured_and_returns_full_scope(self): + with ( + patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), + patch("ucode.cli.available_mcp_clients", return_value=["claude", "codex"]), + patch("ucode.cli.configured_mcp_clients", return_value=["claude"]), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + scope = cli_mod._configure_agents_for_mcp(["claude", "codex"]) + + assert scope == {"claude", "codex"} + mock_cfg.assert_called_once_with(selected_tools=["codex"], prompt_optional_updates=True) + + def test_all_configured_skips_bootstrap(self): + with ( + patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), + patch("ucode.cli.available_mcp_clients", return_value=["claude", "codex"]), + patch("ucode.cli.configured_mcp_clients", return_value=["claude", "codex"]), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + scope = cli_mod._configure_agents_for_mcp(["claude", "codex"]) + + assert scope == {"claude", "codex"} + mock_cfg.assert_not_called() + class TestManagedSkillsOnLaunch: """Managed skills are delivered by download only: the launch path downloads them and never diff --git a/tests/test_mcp.py b/tests/test_mcp.py index b1f10ee4..af6d0f93 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2254,6 +2254,65 @@ 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"], _by_client(["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 mcp.skill_locations_for_client(entry, "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")] + + def test_global_addition_reaches_every_client_and_keeps_divergence(self, monkeypatch): + prior = mcp._resolve_skills_mcp_servers( + WS, ["claude", "codex"], {"claude": ["A.a", "B.b"], "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 *a, **kw: []) + monkeypatch.setattr(mcp, "save_state", lambda s: None) + + assert mcp.add_skills_command(["C.c"]) == 0 + + entry = _find_skills(state["mcp_servers"])[0] + assert mcp.skill_locations_for_client(entry, "claude") == ["A.a", "B.b", "C.c"] + assert mcp.skill_locations_for_client(entry, "codex") == ["A.a", "C.c"] + + def test_agents_add_matching_existing_scope_is_a_noop(self, monkeypatch): + configured: list[tuple[str, str]] = [] + prior = mcp._resolve_skills_mcp_servers( + WS, ["claude", "codex"], _by_client(["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(["A.a"], agents={"claude"}) == 0 + + entry = _find_skills(state["mcp_servers"])[0] + assert mcp.skill_locations_for_client(entry, "claude") == ["A.a"] + assert configured == [] + class TestRegisterSchemalessSkillsConnection: def _stub(self, monkeypatch):