diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a110f1f..dd7455a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -899,10 +899,24 @@ def configure_mcp( "this location are removed.", ), ] = None, + services: Annotated[ + str | None, + typer.Option( + "--services", + help="Configure exactly this comma-separated subset of MCP services (adding and " + "removing to match) instead of a whole schema. Full names like `system.ai.github` " + "work on their own; bare short names like `github` need --location to locate them. " + "Omit --services to configure the whole --location schema; pass an empty string " + "(with --location) to remove all.", + ), + ] = None, ) -> None: """Add Databricks MCP servers to installed coding tools.""" + # `--services` absent -> None (whole schema); present (even empty) -> the + # explicit subset, so `--services ""` deselects everything. + selected = None if services is None else {s.strip() for s in services.split(",") if s.strip()} try: - configure_mcp_command(location=location) + configure_mcp_command(location=location, services=selected) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index f26f9f5..6257458 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -1109,13 +1109,22 @@ def _resolve_location_mcp_servers( clients: list[str], location: str, original_servers: list[dict], + services: set[str] | None = None, ) -> list[dict]: """Build the desired MCP server list for ``--location .``. Strict replacement: the returned list is exactly the mcp-services discovered at ``location``. Any previously-registered MCP entries outside that location are removed by ``apply_mcp_server_changes``. Raises ``RuntimeError`` for an invalid - location (HTTP 404 from the listing API) or any other listing failure.""" + location (HTTP 404 from the listing API) or any other listing failure. + + When ``services`` is given, the discovered set is narrowed to exactly that + subset (matched by full name like ``system.ai.github`` or bare short name + like ``github``); names not found at ``location`` are skipped with a + warning rather than failing, so a saved selection that references a + since-removed service still configures the rest. An empty set selects + nothing (every previously-registered service in the location is removed). + ``None`` keeps the whole schema.""" if location.count(".") != 1 or not all(part.strip() for part in location.split(".")): raise RuntimeError(f"--location must be `.`, got `{location}`.") @@ -1133,6 +1142,21 @@ def _resolve_location_mcp_servers( if not names: print_note(f"No MCP services exist at `{location}`.") + if services is not None: + discovered_full = set(names) + discovered_short = {full_name.split(".")[-1] for full_name in names} + unknown = services - discovered_full - discovered_short + if unknown: + print_warning( + f"Ignoring requested MCP services not found in `{location}`: " + f"{', '.join(sorted(unknown))}." + ) + names = [ + full_name + for full_name in names + if full_name in services or full_name.split(".")[-1] in services + ] + original_by_name = _servers_by_name(original_servers) working_servers: list[dict] = [] for full_name in names: @@ -1153,7 +1177,24 @@ def _resolve_location_mcp_servers( return working_servers -def configure_mcp_command(location: str | None = None) -> int: +def configure_mcp_command(location: str | None = None, services: set[str] | None = None) -> int: + if services is not None and location is None: + # `--services` works standalone with full names (`system.ai.github`): the + # `.` to configure is derived from them. Bare short names + # (`github`) can't be located without `--location`. + schemas = {".".join(s.split(".")[:2]) for s in services if s.count(".") >= 2} + bare = sorted(s for s in services if s.count(".") < 2) + if bare: + raise RuntimeError( + "--services short names need --location (or pass full names like " + f"`system.ai.`): {', '.join(bare)}" + ) + if len(schemas) != 1: + raise RuntimeError( + "--services without --location must all share one `.` " + f"(got: {', '.join(sorted(schemas)) or 'none'}); pass --location instead." + ) + location = next(iter(schemas)) state = load_state() workspace = state.get("workspace") if not workspace: @@ -1194,7 +1235,7 @@ def configure_mcp_command(location: str | None = None) -> int: original_mcp_servers_for_location: list[dict] = list(state.get("mcp_servers") or []) if location is not None: working_mcp_servers = _resolve_location_mcp_servers( - workspace, profile, clients, location, original_mcp_servers_for_location + workspace, profile, clients, location, original_mcp_servers_for_location, services ) changed = apply_mcp_server_changes( original_mcp_servers_for_location, working_mcp_servers, clients diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 6c7d5ea..81333c8 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1429,6 +1429,219 @@ def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypa ] +class TestConfigureMcpServicesSubset: + """`--location --services a,b,...` configures exactly the named subset.""" + + def test_configures_only_the_requested_subset(self, monkeypatch): + configured: list[tuple[str, str, str, dict]] = [] + saved_states: list[dict] = [] + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + monkeypatch.setattr( + mcp, + "list_mcp_services", + lambda workspace, token, parent: ( + ["system.ai.github", "system.ai.slack", "system.ai.gmail"], + None, + ), + ) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert ( + mcp.configure_mcp_command( + location="system.ai", services={"system.ai.github", "system.ai.gmail"} + ) + == 0 + ) + + # slack is dropped; only the two requested services are configured. + assert sorted(c[1] for c in configured) == ["system-ai-github", "system-ai-gmail"] + assert sorted(s["name"] for s in saved_states[-1]["mcp_servers"]) == [ + "system-ai-github", + "system-ai-gmail", + ] + + def test_matches_bare_short_names(self, monkeypatch): + configured: list[tuple[str, str, str, dict]] = [] + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + monkeypatch.setattr( + mcp, + "list_mcp_services", + lambda workspace, token, parent: (["system.ai.github", "system.ai.slack"], None), + ) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + + assert mcp.configure_mcp_command(location="system.ai", services={"github"}) == 0 + + assert [c[1] for c in configured] == ["system-ai-github"] + + def test_unknown_requested_service_warns_and_skips(self, monkeypatch): + configured: list[tuple[str, str, str, dict]] = [] + warnings: list[str] = [] + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + monkeypatch.setattr( + mcp, + "list_mcp_services", + lambda workspace, token, parent: (["system.ai.github"], None), + ) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + monkeypatch.setattr(mcp, "print_warning", lambda msg: warnings.append(msg)) + + assert ( + mcp.configure_mcp_command( + location="system.ai", services={"system.ai.github", "system.ai.ghost"} + ) + == 0 + ) + + # The known service is still configured; the unknown one is reported, not fatal. + assert [c[1] for c in configured] == ["system-ai-github"] + assert any("system.ai.ghost" in w for w in warnings) + + def test_empty_services_removes_everything(self, monkeypatch): + existing = { + "name": "system-ai-github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + } + configured: list[tuple[str, str, str, dict]] = [] + removed: list[tuple[str, str]] = [] + saved_states: list[dict] = [] + _stub_location_base(monkeypatch, {**CLAUDE_STATE, "mcp_servers": [existing]}) + monkeypatch.setattr( + mcp, + "list_mcp_services", + lambda workspace, token, parent: (["system.ai.github"], None), + ) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + ) + monkeypatch.setattr( + mcp, + "remove_client_mcp_server", + lambda client, name: removed.append((client, name)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert mcp.configure_mcp_command(location="system.ai", services=set()) == 0 + + assert configured == [] + assert removed == [("claude", "system-ai-github")] + assert saved_states[-1]["mcp_servers"] == [] + + def test_adds_and_removes_to_match_new_selection(self, monkeypatch): + # The live case teammates want mid-session: started with github+slack, + # then the user deselects slack and selects gmail. + github = { + "name": "system-ai-github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + } + slack = { + "name": "system-ai-slack", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.slack", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + } + configured: list[tuple[str, str, str, dict]] = [] + removed: list[tuple[str, str]] = [] + saved_states: list[dict] = [] + _stub_location_base(monkeypatch, {**CLAUDE_STATE, "mcp_servers": [github, slack]}) + monkeypatch.setattr( + mcp, + "list_mcp_services", + lambda workspace, token, parent: ( + ["system.ai.github", "system.ai.slack", "system.ai.gmail"], + None, + ), + ) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + ) + monkeypatch.setattr( + mcp, + "remove_client_mcp_server", + lambda client, name: removed.append((client, name)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert ( + mcp.configure_mcp_command( + location="system.ai", services={"system.ai.github", "system.ai.gmail"} + ) + == 0 + ) + + # slack removed, gmail added, github untouched (entry unchanged). + assert removed == [("claude", "system-ai-slack")] + assert [c[1] for c in configured] == ["system-ai-gmail"] + assert sorted(s["name"] for s in saved_states[-1]["mcp_servers"]) == [ + "system-ai-github", + "system-ai-gmail", + ] + + def test_full_names_without_location_derive_the_schema(self, monkeypatch): + configured: list[tuple[str, str, str, dict]] = [] + seen: dict[str, str] = {} + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + + def fake_list(workspace, token, parent): + seen["parent"] = parent + return ["system.ai.github", "system.ai.slack"], None + + monkeypatch.setattr(mcp, "list_mcp_services", fake_list) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + + # No --location: the `.` is derived from the full names. + assert mcp.configure_mcp_command(services={"system.ai.github", "system.ai.slack"}) == 0 + + assert seen == {"parent": "system.ai"} + assert sorted(c[1] for c in configured) == ["system-ai-github", "system-ai-slack"] + + def test_short_name_without_location_raises(self): + try: + mcp.configure_mcp_command(services={"github"}) + except RuntimeError as exc: + assert "--location" in str(exc) + else: + raise AssertionError("expected RuntimeError for a bare short name without --location") + + def test_full_names_spanning_multiple_schemas_without_location_raises(self): + try: + mcp.configure_mcp_command(services={"system.ai.github", "other.cat.thing"}) + except RuntimeError as exc: + assert "--location" in str(exc) + else: + raise AssertionError( + "expected RuntimeError for multi-schema services without --location" + ) + + class TestRevertMcpConfigs: def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkeypatch): removed: list[tuple[str, str]] = []