Skip to content
Open
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
12 changes: 12 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,17 @@ 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.

#### Remove configured servers

To unregister servers you've already configured, use `ucode mcp remove`:

```bash
ucode mcp remove
```

It shows the servers you currently have configured — each with the coding tools it's registered
on — and removes the ones you select from those tools. It needs no Databricks login.

### Skills (optional)

Configure Unity Catalog Skills for your coding tools with `ucode configure skills`:
Expand Down Expand Up @@ -267,6 +278,7 @@ pick the new config up on 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 remove` | Interactively unregister configured MCP servers from your coding tools |
| `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
18 changes: 18 additions & 0 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@
configure_mcp_command,
configure_skills_mcp_command,
purge_cross_workspace_mcp_residue,
remove_mcp_command,
revert_mcp_configs,
)
from ucode.skills_download import (
Expand Down Expand Up @@ -1097,6 +1098,23 @@ def mcp_add(
raise typer.Exit(130) from None


@mcp_app.command("remove")
def mcp_remove() -> 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.
"""
try:
remove_mcp_command()
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."""
Expand Down
90 changes: 85 additions & 5 deletions src/ucode/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -1712,11 +1712,21 @@ def prompt_for_mcp_search_sources(exclude_sources: set[str] | None = None) -> se
return {str(value) for value in selection}


def setup_mcp_clients(state: dict, section: str) -> tuple[str, str | None, list[str]]:
def setup_mcp_clients(
state: dict,
section: str,
*,
require_auth: bool = True,
action_note: str = "Configuring for",
) -> tuple[str, str | None, list[str]]:
"""Validate the workspace, resolve configured MCP clients, and prepare auth.

Returns ``(workspace, profile, clients)`` and prints the section header, the
"Configuring for" note, and a warning per configured-but-uninstalled client.
``action_note`` line, and a warning per configured-but-uninstalled client.

``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.
"""
workspace = state.get("workspace")
if not workspace:
Expand All @@ -1742,12 +1752,13 @@ def setup_mcp_clients(state: dict, section: str) -> tuple[str, str | None, list[
]

profile = state.get("profile")
apply_pat_environment(state)
ensure_databricks_auth(workspace, profile)
if require_auth:
apply_pat_environment(state)
ensure_databricks_auth(workspace, profile)

print_section(section)
client_names = ", ".join(str(MCP_CLIENTS[client]["display"]) for client in clients)
print_note(f"Configuring for: {client_names}")
print_note(f"{action_note}: {client_names}")
for client in missing_clients:
print_warning(
f"{MCP_CLIENTS[client]['display']} is configured in ucode but not installed; "
Expand Down Expand Up @@ -1966,6 +1977,75 @@ def _mcp_change_summary(added: list[str], removed: list[str], clients: list[str]
return f"{summary} {noun} across {client_names}" if client_names else f"{summary} {noun}"


def _prompt_for_mcp_removal(servers: list[dict]) -> list[str] | None:
"""Checklist of already-configured MCP servers to remove. Each item shows the
registered name and the tools it's currently on. Returns the selected server
names, ``None`` if cancelled (Ctrl-C), or ``[]`` if nothing was checked."""
choices: list[questionary.Choice | questionary.Separator] = []
for server in servers:
name = _server_name(server)
if not name:
continue
on_clients = [str(MCP_CLIENTS[c]["display"]) for c in _mcp_server_clients(server)]
title = f"{name} ({', '.join(on_clients)})" if on_clients else name
choices.append(questionary.Choice(title=title, value=name, checked=False))
if not choices:
return []
selection = _scrolling_checkbox(
"Remove MCP:",
choices=choices,
style=_picker_style(),
instruction="(space to toggle, ctrl-a all, enter to remove, type to filter)",
).ask()
if selection is None:
return None
return [str(value) for value in selection]


def remove_mcp_command() -> 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."""
state = load_state()
workspace, profile, clients = setup_mcp_clients(
state, "Remove MCP Servers", require_auth=False, action_note="Removing from"
)

original_mcp_servers = list(state.get("mcp_servers") or [])
removable = [s for s in original_mcp_servers if s.get("kind") != SKILLS_MCP_KIND]
if not removable:
print_note("No MCP servers are configured to remove.")
return 0

selection = _prompt_for_mcp_removal(removable)
if selection is None:
return 0
if not selection:
print_note("No MCP servers selected.")
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
]
changed = apply_mcp_server_changes(
original_mcp_servers,
working_mcp_servers,
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
save_state(state)
print_success(_mcp_change_summary([], sorted(remove_names), clients))
return 0


def _merge_clients(prior: list[str] | None, new: list[str]) -> list[str]:
"""Order-preserving union of a prior client list with newly-configured ones."""
prior = list(prior or [])
Expand Down
108 changes: 108 additions & 0 deletions tests/test_mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -386,6 +386,33 @@ def test_additive_picker_shows_configured_servers_as_disabled(self):
# A not-yet-configured server stays an addable, toggleable choice.
assert choices_by_title["Databricks SQL"].disabled is None

def test_removal_picker_lists_configured_servers_with_their_clients(self, monkeypatch):
checkbox_calls: list[dict] = []

class FakePrompt:
def ask(self):
return ["system-ai-github"]

def fake_checkbox(*args, **kwargs):
checkbox_calls.append(kwargs)
return FakePrompt()

monkeypatch.setattr(mcp, "_scrolling_checkbox", fake_checkbox)

result = mcp._prompt_for_mcp_removal(
[
{"name": "system-ai-github", "url": f"{WS}/x", "clients": ["claude", "codex"]},
{"name": "", "url": f"{WS}/y", "clients": ["claude"]}, # unnamed: skipped
]
)

assert result == ["system-ai-github"]
choices = checkbox_calls[0]["choices"]
# The unnamed server is skipped; the named one shows the tools it's on.
assert [c.title for c in choices] == ["system-ai-github (Claude Code, Codex)"]
assert [c.value for c in choices] == ["system-ai-github"]
assert all(c.checked is False for c in choices)

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"]
Expand Down Expand Up @@ -1918,6 +1945,87 @@ def test_empty_services_is_a_noop(self, monkeypatch):
assert called == []


class TestRemoveMcpCommand:
"""`ucode mcp remove` interactively unregisters configured MCP servers."""

GITHUB = {
"name": "system-ai-github",
"url": f"{WS}/ai-gateway/mcp-services/system.ai.github",
"auth": "proxy",
"clients": ["claude"],
}
SQL = {
"name": "databricks-sql",
"url": f"{WS}/api/2.0/mcp/sql",
"auth": "proxy",
"clients": ["claude"],
}

def test_removes_only_selected_servers(self, monkeypatch):
saved_states: list[dict] = []
removed: list[tuple[str, str]] = []
_stub_location_base(
monkeypatch,
{**CLAUDE_STATE, "mcp_servers": [self.GITHUB, self.SQL]},
)
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() == 0

assert removed == [("claude", "system-ai-github")]
assert [s["name"] for s in saved_states[-1]["mcp_servers"]] == ["databricks-sql"]

def test_cancel_makes_no_changes(self, monkeypatch):
saved_states: list[dict] = []
_stub_location_base(monkeypatch, {**CLAUDE_STATE, "mcp_servers": [self.SQL]})
monkeypatch.setattr(mcp, "_prompt_for_mcp_removal", lambda servers: None)
monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state))

assert mcp.remove_mcp_command() == 0
assert saved_states == []

def test_skills_connection_is_not_offered(self, monkeypatch):
offered: dict[str, list[str]] = {}
skills = {
"name": mcp.SKILLS_MCP_SERVER_NAME,
"kind": mcp.SKILLS_MCP_KIND,
"url": f"{WS}/ai-gateway/skills/?schema=main.default",
"auth": "env:OAUTH_TOKEN",
"clients": ["claude"],
}
_stub_location_base(
monkeypatch,
{**CLAUDE_STATE, "mcp_servers": [skills, self.GITHUB]},
)

def fake_prompt(servers):
offered["names"] = [s["name"] for s in servers]
return []

monkeypatch.setattr(mcp, "_prompt_for_mcp_removal", fake_prompt)

assert mcp.remove_mcp_command() == 0
# The skills connection is owned by `configure skills`, so it's never a
# removal candidate; only the real MCP server is offered.
assert offered["names"] == ["system-ai-github"]

def test_no_configured_servers_skips_the_picker(self, monkeypatch):
prompted: list[bool] = []
_stub_location_base(monkeypatch, {**CLAUDE_STATE, "mcp_servers": []})
monkeypatch.setattr(
mcp, "_prompt_for_mcp_removal", lambda servers: prompted.append(True) or []
)

assert mcp.remove_mcp_command() == 0
assert prompted == []


class TestConfigureMcpServicesSubset:
"""`--location <schema> --services a,b,...` configures exactly the named subset."""

Expand Down
Loading