add --services subset flag to configure mcp (#175)

`configure mcp` could only configure a whole `--location <catalog>.<schema>`.
Add `--services` to configure exactly a comma-separated subset, adding and
removing registered servers to match (so deselecting a service removes it).
Full names like `system.ai.github` work standalone (schema derived); bare
short names like `github` resolve against `--location`. An empty string
removes all; names not found at the location are warned about and skipped.

Co-authored-by: Isaac

Co-authored-by: Kecheng Cao <kecheng.cao@databricks.com>

Author: ckcuslife-source

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

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 <cat>.<schema>``.
 
     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 `<catalog>.<schema>`, 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
+        # `<catalog>.<schema>` 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.<name>`): {', '.join(bare)}"
+            )
+        if len(schemas) != 1:
+            raise RuntimeError(
+                "--services without --location must all share one `<catalog>.<schema>` "
+                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

tests/test_mcp.py

@@ -1429,6 +1429,219 @@ def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypa
         ]
 
 
+class TestConfigureMcpServicesSubset:
+    """`--location <schema> --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 `<catalog>.<schema>` 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]] = []