diff --git a/README.md b/README.md index d0a457c..17a4589 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,28 @@ ucode configure --agents claude --mcp system.ai.slack `--mcp` also works without `--agents` for MCP-only clients (it configures just the workspace, then registers the servers); pass a comma-separated list to register several at once. +#### Add servers without replacing existing ones + +`ucode configure mcp` **replaces** the registered MCP servers with your selection — anything +outside a `--location`/`--services` scope (or left unchecked in the picker) is removed. To +**add** servers while leaving everything already configured in place, use `ucode mcp add`: + +```bash +# Register a whole schema's services, keeping any servers already configured. +ucode mcp add --location system.ai + +# Register just a subset (same name rules as `configure mcp --services`). +ucode mcp add --services system.ai.slack,system.ai.github + +# No arguments launches the same interactive picker, but never removes servers. +ucode mcp add +``` + +`ucode mcp add` takes the same `--location` and `--services` options as `ucode configure mcp`; +the only difference is that it never removes servers outside the selection. In the interactive +picker, servers you already have configured are shown as `(already configured)` and can't be +toggled off — you only pick new ones to add. + ### Skills (optional) Configure Unity Catalog Skills for your coding tools with `ucode configure skills`: @@ -243,6 +265,8 @@ pick the new config up on their next ucode run. | `ucode configure --skip-validate` | Write configs without sending a test message through each agent | | `ucode configure --agents claude,codex,pi --skip-unavailable` | Configure the requested agents that are available; skip the rest with a warning | | `ucode configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command | +| `ucode mcp add --location system.ai` | Register a schema's MCP servers, keeping any already configured (additive; never removes) | +| `ucode mcp add --services system.ai.slack` | Register specific MCP server(s) without removing existing ones | | `ucode configure skills` | Register the skills MCP connection (utility tools only); no skills download | | `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) | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 35b6a24..5db605d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -85,6 +85,7 @@ from ucode.mcp import ( MCP_CLIENTS, SKILLS_MCP_KIND, + add_mcp_command, apply_managed_mcp_servers, apply_managed_skills, configure_mcp_command, @@ -1058,6 +1059,44 @@ def _version_callback(value: bool) -> None: raise typer.Exit() +@mcp_app.command("add") +def mcp_add( + location: Annotated[ + str | None, + typer.Option( + "--location", + help="Non-interactive: register the MCP services in the given Unity Catalog " + "`.` (e.g. `system.ai`) and exit without showing the picker. " + "Servers already configured outside this location are kept.", + ), + ] = None, + services: Annotated[ + str | None, + typer.Option( + "--services", + help="Register this comma-separated subset of MCP services (additively). Full names " + "like `system.ai.github` work on their own; bare short names like `github` need " + "--location to locate them. Omit --services to register the whole --location schema; " + 'an empty `--services ""` adds nothing (no-op).', + ), + ] = None, +) -> None: + """Add Databricks MCP servers to installed coding tools. + + Like `ucode configure mcp`, but purely additive: it never removes MCP servers + that are already configured, only registers new ones. + """ + selected = None if services is None else {s.strip() for s in services.split(",") if s.strip()} + try: + add_mcp_command(location=location, services=selected) + except RuntimeError as exc: + print_err(str(exc)) + raise typer.Exit(1) from None + except KeyboardInterrupt: + print_err("Interrupted.") + raise typer.Exit(130) from None + + @mcp_app.command("web-search") def mcp_web_search_cmd() -> None: """Run the web_search MCP server over stdio. Invoked as a subprocess by Claude Code.""" diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index a3db790..24ace6c 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -880,15 +880,27 @@ def build_mcp_picker_choices( available_mcp_service_names: list[str] | None = None, available_vector_search_servers: list[dict] | None = None, available_uc_functions_servers: list[dict] | None = None, + additive: bool = False, ) -> list[questionary.Choice | questionary.Separator]: original_by_name = _servers_by_name(original_servers) known_names = set(original_by_name) + def known_choice(name: str, title: str | None = None) -> questionary.Choice: + # `ucode mcp add` (additive) never removes an already-configured server, so + # show it as a non-toggleable note rather than a pre-checked box whose + # unchecking would be silently ignored. `configure mcp` (replace) keeps it a + # pre-checked toggle so unchecking removes it. + if additive: + return questionary.Choice( + title=title or name, value=name, disabled="already configured" + ) + return _server_choice(name, True, title) + choices: list[questionary.Choice | questionary.Separator] = [] displayed_names: set[str] = set() if "databricks-sql" in known_names: - choices.append(_server_choice("databricks-sql", True, "Databricks SQL")) + choices.append(known_choice("databricks-sql", "Databricks SQL")) else: choices.append(_add_choice(SQL_MCP_VALUE, "Databricks SQL")) displayed_names.add("databricks-sql") @@ -900,7 +912,7 @@ def build_mcp_picker_choices( registered_as = name.replace(".", "-") display_title = f"MCP: {name}" if registered_as in known_names: - choices.append(_server_choice(registered_as, True, display_title)) + choices.append(known_choice(registered_as, display_title)) else: choices.append(_add_choice(f"{MCP_SERVICE_SELECTION_PREFIX}{name}", display_title)) displayed_names.add(registered_as) @@ -908,7 +920,7 @@ def build_mcp_picker_choices( for name in available_external_names: display_title = f"Connection: {name}" if name in known_names: - choices.append(_server_choice(name, True, display_title)) + choices.append(known_choice(name, display_title)) else: choices.append(_add_choice(f"{EXTERNAL_MCP_SELECTION_PREFIX}{name}", display_title)) displayed_names.add(name) @@ -920,7 +932,7 @@ def build_mcp_picker_choices( continue display_title = f"Genie: {title}" if isinstance(title, str) and title else name if name in known_names: - choices.append(_server_choice(name, True, display_title)) + choices.append(known_choice(name, display_title)) else: choices.append( _add_choice( @@ -937,7 +949,7 @@ def build_mcp_picker_choices( continue display_title = f"App: {title}" if isinstance(title, str) and title else name if name in known_names: - choices.append(_server_choice(name, True, display_title)) + choices.append(known_choice(name, display_title)) else: choices.append( _add_choice( @@ -955,7 +967,7 @@ def build_mcp_picker_choices( continue display_title = f"Vector Search: {catalog}.{schema}" if name in known_names: - choices.append(_server_choice(name, True, display_title)) + choices.append(known_choice(name, display_title)) else: choices.append( _add_choice( @@ -973,7 +985,7 @@ def build_mcp_picker_choices( continue display_title = f"UC Functions: {catalog}.{schema}" if name in known_names: - choices.append(_server_choice(name, True, display_title)) + choices.append(known_choice(name, display_title)) else: choices.append( _add_choice( @@ -984,7 +996,7 @@ def build_mcp_picker_choices( displayed_names.add(name) for name in sorted(known_names - displayed_names): - choices.append(_server_choice(name, True)) + choices.append(known_choice(name)) return choices @@ -997,10 +1009,14 @@ def prompt_for_mcp_server_choices( available_vector_search_servers: list[dict] | None = None, available_uc_functions_servers: list[dict] | None = None, allow_back: bool = False, + additive: bool = False, ) -> list[str] | None | _Back: """Show the MCP server picker. Returns the list of selected values, `None` if cancelled (Ctrl-C), or `_BACK` if `allow_back` and the user pressed Left - to return to the previous wizard step.""" + to return to the previous wizard step. + + ``additive`` (``ucode mcp add``) shows already-configured servers as + non-toggleable notes instead of pre-checked, removable boxes.""" instruction = "(space to toggle, ctrl-a all, enter to save, type to filter)" if allow_back: instruction = "(space to toggle, ctrl-a all, ← back, enter to save, type to filter)" @@ -1014,6 +1030,7 @@ def prompt_for_mcp_server_choices( available_mcp_service_names, available_vector_search_servers, available_uc_functions_servers, + additive=additive, ), style=_picker_style(), instruction=instruction, @@ -1739,15 +1756,49 @@ def setup_mcp_clients(state: dict, section: str) -> tuple[str, str | None, list[ return workspace, profile, clients +def _union_missing(base: list[dict], selected: list[dict]) -> list[dict]: + """Return ``selected`` followed by every ``base`` server whose name isn't + already in it. Used by ``ucode mcp add`` so registering new servers never + removes ones that are already configured (append semantics).""" + have = _servers_by_name(selected) + extra = [s for s in base if (_server_name(s) or "") not in have] + return [*selected, *extra] + + +def add_mcp_command( + location: str | None = None, + services: set[str] | None = None, +) -> int: + """`ucode mcp add`: register Databricks MCP servers WITHOUT removing any that + are already configured. + + Uses the same discovery and options as `configure mcp` — the interactive + picker, or the non-interactive `--location`/`--services` paths — but is purely + additive: unlike `configure mcp`, it never removes servers outside the + selection.""" + if services is not None and not services: + # An empty `--services` selects nothing. For `configure mcp` that means + # "remove all"; for the additive `add` there is simply nothing to register, + # so it's a no-op (and doesn't need --location the way a real subset does). + print_note("No MCP services given to add (empty --services); nothing to do.") + return 0 + return configure_mcp_command(location=location, services=services, append=True) + + def configure_mcp_command( location: str | None = None, services: set[str] | None = None, *, exclude_sources: set[str] | None = None, + append: bool = False, ) -> int: """Interactive MCP picker. ``exclude_sources`` hides search sources the caller can't use — `ucode setup` passes ``{"apps"}`` because a managed config can't carry an app's off-workspace - host, so an app picked here would be silently dropped from the published config.""" + host, so an app picked here would be silently dropped from the published config. + + ``append`` (used by `ucode mcp add`) makes the command purely additive: the + final server list is unioned with the already-configured servers, so nothing + outside the current selection is removed.""" 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 @@ -1766,13 +1817,19 @@ def configure_mcp_command( ) location = next(iter(schemas)) state = load_state() - workspace, profile, clients = setup_mcp_clients(state, "MCP Servers") + workspace, profile, clients = setup_mcp_clients( + state, "Add MCP Servers" if append else "MCP Servers" + ) 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, services ) + if append: + working_mcp_servers = _union_missing( + original_mcp_servers_for_location, working_mcp_servers + ) changed = apply_mcp_server_changes( original_mcp_servers_for_location, working_mcp_servers, @@ -1817,6 +1874,7 @@ def configure_mcp_command( discovered["vector_search"], discovered["uc_functions"], allow_back=True, + additive=append, ) if selections is None: return 0 @@ -1866,6 +1924,9 @@ def configure_mcp_command( ) working_names.add(entry_name) + if append: + working_mcp_servers = _union_missing(original_mcp_servers, working_mcp_servers) + changed = apply_mcp_server_changes( original_mcp_servers, working_mcp_servers, @@ -1878,7 +1939,8 @@ def configure_mcp_command( state["mcp_servers"] = working_mcp_servers save_state(state) added = sorted(working_names - set(original_by_name)) - removed = sorted(set(original_by_name) - working_names) + # `add` never removes; the union above re-keeps unselected servers. + removed = [] if append else sorted(set(original_by_name) - working_names) print_success(_mcp_change_summary(added, removed, clients)) elif not selections and not original_mcp_servers: # User submitted the picker without toggling anything --> make it clear nothing was selected diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 2852f47..d5286ce 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -369,6 +369,23 @@ def test_picker_marks_configured_servers(self): assert choices_by_title["Connection: github-mcp"].checked is True assert choices_by_title["Databricks SQL"].checked is False + def test_additive_picker_shows_configured_servers_as_disabled(self): + """In `ucode mcp add` mode an already-configured server can't be removed, so + it's shown as a non-toggleable note rather than a pre-checked box.""" + choices = mcp.build_mcp_picker_choices( + ["github-mcp"], + [], + [], + [{"name": "github-mcp", "url": f"{WS}/api/2.0/mcp/external/github-mcp"}], + additive=True, + ) + choices_by_title = {choice.title: choice for choice in choices} + configured = choices_by_title["Connection: github-mcp"] + assert configured.disabled == "already configured" + assert configured.checked is False + # A not-yet-configured server stays an addable, toggleable choice. + assert choices_by_title["Databricks SQL"].disabled is None + def test_picker_keeps_databricks_sql_when_nothing_discovered(self): choices = mcp.build_mcp_picker_choices([], [], [], []) assert [choice.title for choice in choices] == ["Databricks SQL"] @@ -1806,6 +1823,101 @@ def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypa ] +class TestAddMcpCommand: + """`ucode mcp add` (append) registers new servers without removing existing ones.""" + + def test_keeps_servers_outside_location(self, monkeypatch): + """Unlike `configure mcp --location`, `mcp add --location` preserves any + server outside the location instead of removing it.""" + saved_states: list[dict] = [] + configured: list[tuple[str, str, str]] = [] + removed: list[tuple[str, str]] = [] + outside_entry = { + "name": "databricks-sql", + "url": f"{WS}/api/2.0/mcp/sql", + "auth": "proxy", + "clients": ["claude"], + } + _stub_location_base( + monkeypatch, + {**CLAUDE_STATE, "mcp_servers": [outside_entry]}, + ) + 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, *a, **kw: configured.append((client, name, url)) 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.add_mcp_command(location="system.ai") == 0 + + # Nothing is removed; the new service is added and the outside one kept. + assert removed == [] + assert [c[1] for c in configured] == ["system-ai-github"] + assert saved_states[-1]["mcp_servers"] == [ + { + "name": "system-ai-github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "proxy", + "clients": ["claude"], + }, + outside_entry, + ] + + def test_services_subset_keeps_others_in_location(self, monkeypatch): + """`mcp add --services` registers the named subset while leaving other + already-registered services in the same schema untouched.""" + saved_states: list[dict] = [] + removed: list[tuple[str, str]] = [] + existing = { + "name": "system-ai-slack", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.slack", + "auth": "proxy", + "clients": ["claude"], + } + _stub_location_base( + monkeypatch, + {**CLAUDE_STATE, "mcp_servers": [existing]}, + ) + 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 *a, **kw: []) + 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.add_mcp_command(location="system.ai", services={"github"}) == 0 + + assert removed == [] + names = [s["name"] for s in saved_states[-1]["mcp_servers"]] + assert names == ["system-ai-github", "system-ai-slack"] + + def test_empty_services_is_a_noop(self, monkeypatch): + """`mcp add --services ""` has nothing to add, so it's a no-op that never + reaches configuration (and doesn't need --location the way a subset does).""" + called: list[bool] = [] + monkeypatch.setattr(mcp, "load_state", lambda: called.append(True) or {}) + + assert mcp.add_mcp_command(services=set()) == 0 + assert called == [] + + class TestConfigureMcpServicesSubset: """`--location --services a,b,...` configures exactly the named subset."""