diff --git a/README.md b/README.md index f26ea6b..3dfe8bb 100644 --- a/README.md +++ b/README.md @@ -164,12 +164,29 @@ the only difference is that it never removes servers outside the selection. In t picker, servers you already have configured are shown as `(already configured)` and can't be toggled off — you only pick new ones to add. +Pass `--agents` to target specific coding agents. Any named agent that isn't set up yet is +configured first (workspace + models), so this doubles as one-command setup: + +```bash +# Set up Claude Code (if needed) and register the server for it, in one command. +ucode mcp add --agents claude --services system.ai.slack + +# Target several agents at once. +ucode mcp add --agents claude,codex --location system.ai +``` + +Without `--agents`, the server is registered for every already-configured agent. + #### Remove configured servers To unregister servers you've already configured, use `ucode mcp remove`: ```bash ucode mcp remove + +# Remove only from specific agents. A server registered on several agents is +# unregistered from the named ones and kept on the rest. +ucode mcp remove --agents codex ``` It shows the servers you currently have configured — each with the coding tools it's registered @@ -292,7 +309,9 @@ their next ucode run. | `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 mcp add --agents claude --services system.ai.slack` | Set up the agent(s) if needed and register the server for them | | `ucode mcp remove` | Interactively unregister configured MCP servers from your coding tools | +| `ucode mcp remove --agents codex` | Unregister selected servers from specific agents only | | `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 dda657a..ca71dd6 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1104,6 +1104,36 @@ def _version_callback(value: bool) -> None: raise typer.Exit() +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 `ucode mcp add` has them as targets, and return their canonical + names. Mirrors `ucode configure --agents`: 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) + configure_workspace_command( + selected_tools=selected_tools, 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 + + @mcp_app.command("add") def mcp_add( location: Annotated[ @@ -1125,15 +1155,32 @@ def mcp_add( 'an empty `--services ""` adds nothing (no-op).', ), ] = None, + agents: Annotated[ + str | None, + typer.Option( + "--agents", + help="Comma-separated coding agents to register the server(s) for (e.g. " + "claude,codex,cursor). Any that aren't configured yet are set up first " + "(workspace + models), so this works as a one-command setup. Without --agents, " + "the server is registered for every already-configured agent.", + ), + ] = 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. + that are already configured, only registers new ones. Pass --agents to target + (and, if needed, set up) specific agents. """ selected = None if services is None else {s.strip() for s in services.split(",") if s.strip()} + requested_agents = ( + None + if agents is None + else ({a.strip().lower() for a in agents.split(",") if a.strip()} or None) + ) try: - add_mcp_command(location=location, services=selected) + scope = _configure_agents_for_mcp(sorted(requested_agents)) if requested_agents else None + add_mcp_command(location=location, services=selected, agents=scope) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None @@ -1143,14 +1190,30 @@ def mcp_add( @mcp_app.command("remove") -def mcp_remove() -> None: +def mcp_remove( + agents: Annotated[ + str | None, + typer.Option( + "--agents", + help="Comma-separated coding agents to remove the server(s) from (e.g. " + "claude,codex). A server registered on several agents is unregistered only " + "from the named ones and kept on the rest. Without --agents, a selected server " + "is removed from every agent it's on.", + ), + ] = None, +) -> None: """Remove configured Databricks MCP servers from your coding tools. Interactive: shows the servers you currently have configured and unregisters the ones you select. Needs no Databricks login. """ + requested_agents = ( + None + if agents is None + else ({a.strip().lower() for a in agents.split(",") if a.strip()} or None) + ) try: - remove_mcp_command() + remove_mcp_command(agents=requested_agents) 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 83ca6f4..75e3699 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -1718,6 +1718,7 @@ def setup_mcp_clients( *, require_auth: bool = True, action_note: str = "Configuring for", + agents: set[str] | None = None, ) -> tuple[str, str | None, list[str]]: """Validate the workspace, resolve configured MCP clients, and prepare auth. @@ -1727,6 +1728,11 @@ def setup_mcp_clients( ``require_auth`` forces a Databricks login (needed to register a server); the removal path passes ``False`` since unregistering a server is purely local and should work even when the workspace token has expired. + + ``agents`` (from ``--agents``) scopes the returned clients to that subset of + the configured MCP clients, so the operation touches only those agents instead + of every configured one. Requested agents that aren't configured/installed + raise a clear error. """ workspace = state.get("workspace") if not workspace: @@ -1741,6 +1747,14 @@ def setup_mcp_clients( "or GitHub Copilot CLI." ) clients = configured_mcp_clients(state, installed_clients) + if agents is not None: + missing = sorted(a for a in agents if a not in clients) + if missing: + raise RuntimeError( + f"Requested agent(s) not configured for MCP: {', '.join(missing)}. " + f"Configure them first with `ucode configure --agents {','.join(missing)}`." + ) + clients = [client for client in clients if client in agents] if not clients: raise RuntimeError( "No configured MCP-capable coding agents are installed. Run `ucode configure` " @@ -1779,6 +1793,7 @@ def _union_missing(base: list[dict], selected: list[dict]) -> list[dict]: def add_mcp_command( location: str | None = None, services: set[str] | None = None, + agents: set[str] | None = None, ) -> int: """`ucode mcp add`: register Databricks MCP servers WITHOUT removing any that are already configured. @@ -1786,14 +1801,18 @@ def add_mcp_command( 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.""" + selection. + + ``agents`` scopes the registration to that subset of configured MCP clients + (the agents must already be configured — the `--agents` CLI option sets up any + that aren't before calling this).""" 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) + return configure_mcp_command(location=location, services=services, append=True, agents=agents) def configure_mcp_command( @@ -1802,6 +1821,7 @@ def configure_mcp_command( *, exclude_sources: set[str] | None = None, append: bool = False, + agents: set[str] | None = None, ) -> 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 @@ -1809,7 +1829,8 @@ def configure_mcp_command( ``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.""" + outside the current selection is removed. ``agents`` scopes the operation to + that subset of configured MCP clients.""" 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 @@ -1829,7 +1850,7 @@ def configure_mcp_command( location = next(iter(schemas)) state = load_state() workspace, profile, clients = setup_mcp_clients( - state, "Add MCP Servers" if append else "MCP Servers" + state, "Add MCP Servers" if append else "MCP Servers", agents=agents ) original_mcp_servers_for_location: list[dict] = list(state.get("mcp_servers") or []) @@ -2002,22 +2023,33 @@ def _prompt_for_mcp_removal(servers: list[dict]) -> list[str] | None: return [str(value) for value in selection] -def remove_mcp_command() -> int: +def remove_mcp_command(agents: set[str] | None = None) -> int: """`ucode mcp remove`: interactively unregister configured MCP servers. Shows the servers currently configured (skills connections excluded — they're - owned by `configure skills`) and removes the ones you select from every coding - tool they're registered on. It never adds or reconfigures anything, and needs no - Databricks auth.""" + owned by `configure skills`) and removes the ones you select. It never adds or + reconfigures anything, and needs no Databricks auth. + + Without ``agents``, a selected server is removed from every coding tool it's + registered on. With ``agents`` (from ``--agents``), removal is scoped to those + agents: a server registered on several agents is unregistered only from the + named ones and kept on the rest; only servers registered on a named agent are + offered.""" state = load_state() workspace, profile, clients = setup_mcp_clients( - state, "Remove MCP Servers", require_auth=False, action_note="Removing from" + state, "Remove MCP Servers", require_auth=False, action_note="Removing from", agents=agents ) original_mcp_servers = list(state.get("mcp_servers") or []) - removable = [s for s in original_mcp_servers if s.get("kind") != SKILLS_MCP_KIND] + removable = [ + s + for s in original_mcp_servers + if s.get("kind") != SKILLS_MCP_KIND + and (agents is None or bool(set(_mcp_server_clients(s)) & agents)) + ] if not removable: - print_note("No MCP servers are configured to remove.") + scope = "" if agents is None else f" for {', '.join(sorted(agents))}" + print_note(f"No MCP servers are configured to remove{scope}.") return 0 selection = _prompt_for_mcp_removal(removable) @@ -2028,19 +2060,39 @@ def remove_mcp_command() -> int: return 0 remove_names = set(selection) - working_mcp_servers = [ - s for s in original_mcp_servers if (_server_name(s) or "") not in remove_names - ] + # Present each removed server to `apply_mcp_server_changes` with its client list + # narrowed to just the agents we're removing from (all recorded clients when + # `agents` is None), and drop it from the working list — so the machinery + # unregisters it from exactly those agents and no others. + removal_view: list[dict] = [] + for server in original_mcp_servers: + name = _server_name(server) + if name not in remove_names: + continue + recorded = _mcp_server_clients(server) + targets = recorded if agents is None else [c for c in recorded if c in agents] + if targets: + removal_view.append({**server, "clients": targets}) changed = apply_mcp_server_changes( - original_mcp_servers, - working_mcp_servers, - clients, - workspace, - profile, - use_pat=bool(state.get("use_pat")), + removal_view, [], clients, workspace, profile, use_pat=bool(state.get("use_pat")) ) - if changed or original_mcp_servers != working_mcp_servers: - state["mcp_servers"] = working_mcp_servers + + # Update saved state: drop a fully-removed server, or keep it with the named + # agents stripped from its client list when the removal was agent-scoped. + new_servers: list[dict] = [] + for server in original_mcp_servers: + name = _server_name(server) + if name not in remove_names: + new_servers.append(server) + continue + remaining = ( + [] if agents is None else [c for c in (server.get("clients") or []) if c not in agents] + ) + if remaining: + new_servers.append({**server, "clients": remaining}) + + if changed or new_servers != original_mcp_servers: + state["mcp_servers"] = new_servers save_state(state) print_success(_mcp_change_summary([], sorted(remove_names), clients)) return 0 diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 18f4e3e..69c60d0 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1946,6 +1946,46 @@ def test_empty_services_is_a_noop(self, monkeypatch): assert mcp.add_mcp_command(services=set()) == 0 assert called == [] + def test_agents_scopes_registration_to_named_agent(self, monkeypatch): + """With two agents configured, `agents={"claude"}` registers the server for + claude only and records only claude on the saved entry.""" + saved_states: list[dict] = [] + configured: list[tuple[str, str]] = [] + _stub_location_base( + monkeypatch, + {"workspace": WS, "available_tools": ["claude", "codex"], "mcp_servers": []}, + ) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude", "codex"]) + 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)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert mcp.add_mcp_command(location="system.ai", agents={"claude"}) == 0 + + assert configured == [("claude", "system-ai-github")] + assert saved_states[-1]["mcp_servers"][0]["clients"] == ["claude"] + + def test_agents_not_configured_raises(self, monkeypatch): + """`--agents` naming an agent that isn't configured for MCP is a clear error + (the CLI sets agents up first, so this guards the library entry point).""" + _stub_location_base(monkeypatch, {**CLAUDE_STATE, "mcp_servers": []}) + monkeypatch.setattr( + mcp, "list_mcp_services", lambda workspace, token, parent: (["system.ai.github"], None) + ) + try: + mcp.add_mcp_command(location="system.ai", agents={"gemini"}) + except RuntimeError as exc: + assert "gemini" in str(exc) + assert "not configured" in str(exc) + else: + raise AssertionError("expected RuntimeError for an unconfigured agent") + class TestRemoveMcpCommand: """`ucode mcp remove` interactively unregisters configured MCP servers.""" @@ -2027,6 +2067,57 @@ def test_no_configured_servers_skips_the_picker(self, monkeypatch): assert mcp.remove_mcp_command() == 0 assert prompted == [] + def test_agents_scopes_removal_to_named_agent(self, monkeypatch): + """`--agents claude` unregisters the server from claude only; a server that + was also on codex is kept there with claude stripped from its clients.""" + saved_states: list[dict] = [] + removed: list[tuple[str, str]] = [] + both = { + "name": "system-ai-github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "proxy", + "clients": ["claude", "codex"], + } + _stub_location_base( + monkeypatch, + {"workspace": WS, "available_tools": ["claude", "codex"], "mcp_servers": [both]}, + ) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude", "codex"]) + monkeypatch.setattr(mcp, "_prompt_for_mcp_removal", lambda servers: ["system-ai-github"]) + 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.remove_mcp_command(agents={"claude"}) == 0 + + # Only claude is unregistered; the state entry survives on codex. + assert removed == [("claude", "system-ai-github")] + assert saved_states[-1]["mcp_servers"] == [{**both, "clients": ["codex"]}] + + def test_agents_only_offers_servers_on_that_agent(self, monkeypatch): + prompted: list[bool] = [] + codex_only = { + "name": "system-ai-github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "proxy", + "clients": ["codex"], + } + _stub_location_base( + monkeypatch, + {"workspace": WS, "available_tools": ["claude", "codex"], "mcp_servers": [codex_only]}, + ) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude", "codex"]) + monkeypatch.setattr( + mcp, "_prompt_for_mcp_removal", lambda servers: prompted.append(True) or [] + ) + + # Nothing is registered on claude, so `--agents claude` has nothing to offer. + assert mcp.remove_mcp_command(agents={"claude"}) == 0 + assert prompted == [] + class TestConfigureMcpServicesSubset: """`--location --services a,b,...` configures exactly the named subset."""