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
15 changes: 13 additions & 2 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
28 changes: 15 additions & 13 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
66 changes: 65 additions & 1 deletion src/ucode/mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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),
Expand Down
46 changes: 44 additions & 2 deletions tests/test_databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
16 changes: 16 additions & 0 deletions tests/test_e2e_uc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading