From b1313bcb6a46298b252a5146cd3075755f2b18b3 Mon Sep 17 00:00:00 2001 From: Anjali Sujithan Date: Wed, 24 Jun 2026 01:31:55 +0000 Subject: [PATCH 1/2] non-interactive ucode configure mcp --location --- src/ucode/cli.py | 15 +++- src/ucode/databricks.py | 28 +++--- src/ucode/mcp.py | 66 +++++++++++++- tests/test_databricks.py | 46 +++++++++- tests/test_mcp.py | 187 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 324 insertions(+), 18 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0d60916..a110f1f 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -888,10 +888,21 @@ def configure( @configure_app.command("mcp") -def configure_mcp() -> None: +def configure_mcp( + location: Annotated[ + str | None, + typer.Option( + "--location", + help="Non-interactive: replace registered MCPs with exactly the services " + "in the given Unity Catalog `.` (e.g. `system.ai`) and " + "exit without showing the picker. Any previously-registered MCPs outside " + "this location are removed.", + ), + ] = None, +) -> None: """Add Databricks MCP servers to installed coding tools.""" try: - configure_mcp_command() + configure_mcp_command(location=location) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 4b01120..580de49 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1212,16 +1212,16 @@ def discover_model_services( _MCP_SERVICE_NAME_PREFIX = "mcp-services/" -_MCP_SERVICE_REQUIRED_PREFIX = "system.ai." -def _mcp_service_full_name(service: dict) -> str | None: - """Extract `system.ai.` from one mcp-service entry, or None.""" +def _mcp_service_full_name(service: dict, required_prefix: str) -> str | None: + """Extract the full UC name from one mcp-service entry, or None if it + doesn't live under ``required_prefix`` or isn't ACTIVE.""" name = service.get("name") if not isinstance(name, str): return None name = name.strip().removeprefix(_MCP_SERVICE_NAME_PREFIX) - if not name.startswith(_MCP_SERVICE_REQUIRED_PREFIX): + if not name.startswith(required_prefix): return None status = ((service.get("config") or {}).get("connection") or {}).get("status") if status is not None and status != "ACTIVE": @@ -1229,30 +1229,32 @@ def _mcp_service_full_name(service: dict) -> str | None: return name -def list_mcp_services(workspace: str, token: str) -> tuple[list[str], str | None]: - """List Databricks-curated `system.ai.*` MCP services. +def list_mcp_services( + workspace: str, token: str, parent: str = "system.ai" +) -> tuple[list[str], str | None]: + """List UC MCP services under ``parent`` (a ``.`` ref). - The listing endpoint requires `?parent=schemas/system.ai`; without it the - request returns 499 with a truncated body (verified against e2-dogfood - 2026-06-11). Returns (full_names, reason). + A non-None string indicates the listing call itself failed. Callers can inspect + ``error`` for ``HTTP 404`` to distinguish "invalid location" from other failures. """ hostname = workspace_hostname(workspace) url = ( f"https://{hostname}/api/2.1/unity-catalog/mcp-services" - f"?{urlencode({'parent': 'schemas/system.ai'})}" + f"?{urlencode({'parent': f'schemas/{parent}'})}" ) payload, reason = _http_get_json(url, token, timeout=30) if payload is None: return [], reason + expected_prefix = parent + "." data = cast(dict, payload) if isinstance(payload, dict) else {} - names = [] + names: list[str] = [] for service in data.get("mcp_services") or []: if not isinstance(service, dict): continue - full_name = _mcp_service_full_name(service) + full_name = _mcp_service_full_name(service, expected_prefix) if full_name: names.append(full_name) - return sorted(set(names)), None if names else (reason or "no `system.ai.*` mcp services found") + return sorted(set(names)), None def build_mcp_service_url(workspace: str, full_name: str) -> str: diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 61a2315..f26f9f5 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -1103,7 +1103,57 @@ def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None: ) -def configure_mcp_command() -> int: +def _resolve_location_mcp_servers( + workspace: str, + profile: str | None, + clients: list[str], + location: str, + original_servers: list[dict], +) -> list[dict]: + """Build the desired MCP server list for ``--location .``. + + 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.""" + if location.count(".") != 1 or not all(part.strip() for part in location.split(".")): + raise RuntimeError(f"--location must be `.`, got `{location}`.") + + token = get_databricks_token(workspace, profile) + with spinner(f"Discovering MCP services in {location}..."): + names, reason = list_mcp_services(workspace, token, parent=location) + + if reason and reason.startswith("HTTP 404"): + raise RuntimeError( + f"Invalid location: `{location}` is not a valid Unity Catalog schema " + "in this workspace (or you lack USE permission on it)." + ) + if reason: + raise RuntimeError(f"Failed to list MCP services at `{location}`: {reason}") + if not names: + print_note(f"No MCP services exist at `{location}`.") + + original_by_name = _servers_by_name(original_servers) + working_servers: list[dict] = [] + for full_name in names: + entry_name = full_name.replace(".", "-") + original = original_by_name.get(entry_name) + original_clients = list((original or {}).get("clients") or []) + merged_clients = original_clients + [c for c in clients if c not in original_clients] + candidate = { + "name": entry_name, + "url": build_mcp_service_url(workspace, full_name), + "auth": f"env:{MCP_AUTH_TOKEN_ENV_VAR}", + "clients": merged_clients, + } + if original is not None and original == candidate: + working_servers.append(original.copy()) + else: + working_servers.append(candidate) + return working_servers + + +def configure_mcp_command(location: str | None = None) -> int: state = load_state() workspace = state.get("workspace") if not workspace: @@ -1141,6 +1191,20 @@ def configure_mcp_command() -> int: "skipping MCP config." ) + 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 + ) + changed = apply_mcp_server_changes( + original_mcp_servers_for_location, working_mcp_servers, clients + ) + if changed or original_mcp_servers_for_location != working_mcp_servers: + state["mcp_servers"] = working_mcp_servers + save_state(state) + print_success("Saved") + return 0 + available_external_mcp_names = _discover_mcp_source( "external connections", lambda: discover_external_mcp_connection_names(workspace, profile), diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 65f05d8..bac25d4 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -367,7 +367,7 @@ def test_http_failure_propagates_reason(self, monkeypatch): assert names == [] assert reason == "HTTP 500 Server Error" - def test_empty_payload_reports_no_results(self, monkeypatch): + def test_empty_payload_is_successful_with_no_reason(self, monkeypatch): monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token, timeout=30: ({"mcp_services": []}, None) ) @@ -375,7 +375,49 @@ def test_empty_payload_reports_no_results(self, monkeypatch): names, reason = db_mod.list_mcp_services(WS, "token") assert names == [] - assert reason and "no `system.ai.*`" in reason + assert reason is None + + def test_custom_parent_passes_through_to_url(self, monkeypatch): + captured: dict[str, str] = {} + + def fake_get(url, token, timeout=30): + captured["url"] = url + return {"mcp_services": []}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + db_mod.list_mcp_services(WS, "token", parent="main.svenwb") + + assert "parent=schemas%2Fmain.svenwb" in captured["url"] + + def test_custom_parent_filters_to_namespace(self, monkeypatch): + payload = { + "mcp_services": [ + {"name": "mcp-services/main.svenwb.github"}, + {"name": "mcp-services/main.svenwb.slack"}, + {"name": "mcp-services/system.ai.github"}, + ] + } + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: (payload, None) + ) + + names, reason = db_mod.list_mcp_services(WS, "token", parent="main.svenwb") + + assert reason is None + assert names == ["main.svenwb.github", "main.svenwb.slack"] + + def test_http_404_reason_surfaces_for_invalid_parent(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, timeout=30: (None, "HTTP 404 Not Found: NOT_FOUND"), + ) + + names, reason = db_mod.list_mcp_services(WS, "token", parent="nope.nope") + + assert names == [] + assert reason and reason.startswith("HTTP 404") def _foundation_models_payload(names): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 08b9705..6c7d5ea 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1242,6 +1242,193 @@ def test_removes_saved_server(self, monkeypatch): assert saved_states[-1]["mcp_servers"] == [] +def _stub_location_base(monkeypatch, state): + """Shared scaffolding for --location tests: no installed agents missing, + auth no-ops, no cross-workspace residue.""" + monkeypatch.setattr(mcp, "load_state", lambda: state) + monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + monkeypatch.setattr(mcp, "purge_cross_workspace_mcp_residue", lambda state, workspace: None) + monkeypatch.setattr(mcp, "get_databricks_token", lambda workspace, profile=None: "token") + + +class TestConfigureMcpFromLocation: + def test_rejects_malformed_location(self, monkeypatch): + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + monkeypatch.setattr(mcp, "list_mcp_services", lambda *a, **kw: ([], None)) + + for bad in ("system", "system.ai.extra", ".ai", "system.", ""): + try: + mcp.configure_mcp_command(location=bad) + except RuntimeError as exc: + assert "--location" in str(exc) + else: + raise AssertionError(f"expected RuntimeError for `{bad}`") + + def test_invalid_location_raises_with_clear_message(self, monkeypatch): + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + monkeypatch.setattr( + mcp, + "list_mcp_services", + lambda workspace, token, parent: ([], "HTTP 404 Not Found: NOT_FOUND"), + ) + try: + mcp.configure_mcp_command(location="nope.nope") + except RuntimeError as exc: + assert "Invalid location" in str(exc) + assert "nope.nope" in str(exc) + else: + raise AssertionError("expected RuntimeError") + + def test_other_listing_failure_surfaces_reason(self, monkeypatch): + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + monkeypatch.setattr( + mcp, + "list_mcp_services", + lambda workspace, token, parent: ([], "HTTP 500 Server Error"), + ) + try: + mcp.configure_mcp_command(location="system.ai") + except RuntimeError as exc: + assert "HTTP 500" in str(exc) + else: + raise AssertionError("expected RuntimeError") + + def test_registers_every_discovered_service(self, monkeypatch): + saved_states: list[dict] = [] + configured: list[tuple[str, str, str, dict]] = [] + seen: dict[str, str] = {} + picker_called: list[bool] = [] + _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, + "prompt_for_mcp_server_choices", + lambda *a, **kw: picker_called.append(True) or [], + ) + 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") == 0 + + assert seen == {"parent": "system.ai"} + assert picker_called == [] + assert [c[1] for c in configured] == ["system-ai-github", "system-ai-slack"] + assert configured[0][2] == f"{WS}/ai-gateway/mcp-services/system.ai.github" + assert saved_states[-1]["mcp_servers"] == [ + { + "name": "system-ai-github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + }, + { + "name": "system-ai-slack", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.slack", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + }, + ] + + def test_replaces_servers_outside_location(self, monkeypatch): + saved_states: list[dict] = [] + configured: list[tuple[str, str, str, dict]] = [] + removed: list[tuple[str, str]] = [] + outside_entry = { + "name": "databricks-sql", + "url": f"{WS}/api/2.0/mcp/sql", + "auth": "env:OAUTH_TOKEN", + "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, 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") == 0 + + assert removed == [("claude", "databricks-sql")] + 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": "env:OAUTH_TOKEN", + "clients": ["claude"], + }, + ] + + def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypatch): + """An entry registered before a second agent was configured should + get registered for that agent on the next --location run.""" + saved_states: list[dict] = [] + configured: list[tuple[str, str, str, dict]] = [] + existing = { + "name": "system-ai-github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + } + _stub_location_base( + monkeypatch, + { + "workspace": WS, + "available_tools": ["claude", "codex"], + "mcp_servers": [existing], + }, + ) + 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, 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") == 0 + + assert [c[0] for c in configured] == ["claude", "codex"] + assert saved_states[-1]["mcp_servers"] == [ + { + "name": "system-ai-github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude", "codex"], + } + ] + + class TestRevertMcpConfigs: def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkeypatch): removed: list[tuple[str, str]] = [] From 334797b91f1eedf8918fcbf3efb956671f38977b Mon Sep 17 00:00:00 2001 From: Anjali Sujithan Date: Wed, 24 Jun 2026 01:38:50 +0000 Subject: [PATCH 2/2] e2e tests for configure mcp --location flag --- tests/test_e2e_uc.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_e2e_uc.py b/tests/test_e2e_uc.py index 72fc2f1..32dbf63 100644 --- a/tests/test_e2e_uc.py +++ b/tests/test_e2e_uc.py @@ -65,6 +65,22 @@ def test_returns_only_system_ai_mcp_services(self, e2e_workspace, e2e_token): non_system = sorted({n for n in names if not n.startswith("system.ai.")}) assert not non_system, f"Non-system.ai entries leaked through: {non_system[:5]}" + def test_custom_parent_filters_server_side(self, e2e_workspace, e2e_token): + names, _ = list_mcp_services(e2e_workspace, e2e_token, parent="main.default") + if not names: + pytest.skip("No mcp-services in main.default on this workspace.") + outside = sorted({n for n in names if not n.startswith("main.default.")}) + assert not outside, f"Server returned entries outside main.default: {outside[:5]}" + + def test_invalid_parent_returns_http_404(self, e2e_workspace, e2e_token): + names, reason = list_mcp_services( + e2e_workspace, e2e_token, parent="nope_catalog.nope_schema" + ) + assert names == [] + assert reason and reason.startswith("HTTP 404"), ( + f"Expected HTTP 404 for bogus location, got: {reason}" + ) + # --------------------------------------------------------------------------- # `configure_shared_state` end-to-end: UC discovery is the default, with a