add non-interactive ucode configure mcp --location (#172)

* non-interactive ucode configure mcp --location

* e2e tests for configure mcp --location flag

Author: asujithan

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 `<catalog>.<schema>` (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

src/ucode/databricks.py

@@ -1212,47 +1212,49 @@ 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.<name>` 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":
         return 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 ``<catalog>.<schema>`` 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:

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 <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."""
+    if location.count(".") != 1 or not all(part.strip() for part in location.split(".")):
+        raise RuntimeError(f"--location must be `<catalog>.<schema>`, 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),

tests/test_databricks.py

@@ -367,15 +367,57 @@ 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)
         )
 
         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):

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