Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 <dir>]` | Download a schema's skills to disk (under `<dir>`, 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) |
Expand Down
71 changes: 67 additions & 4 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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[
Expand All @@ -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
Expand All @@ -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
Expand Down
96 changes: 74 additions & 22 deletions src/ucode/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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:
Expand All @@ -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` "
Expand Down Expand Up @@ -1779,21 +1793,26 @@ 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.

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(
Expand All @@ -1802,14 +1821,16 @@ 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
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."""
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
# `<catalog>.<schema>` to configure is derived from them. Bare short names
Expand All @@ -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 [])
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down
Loading
Loading