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
16 changes: 15 additions & 1 deletion src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
47 changes: 44 additions & 3 deletions src/ucode/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}`.")

Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand Down
213 changes: 213 additions & 0 deletions tests/test_mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]] = []
Expand Down
Loading