From 12dbab200f4fce3cecf10c4eee1b9d97333be43e Mon Sep 17 00:00:00 2001 From: Xiang Shen Date: Thu, 23 Jul 2026 02:20:33 +0000 Subject: [PATCH] Refactor: extract shared MCP setup helpers (no behavior change) Prep for the upcoming `ucode configure skills` rework. Pure, behavior-preserving refactor of src/ucode/mcp.py: - Extract `_setup_mcp_clients(state, section)` from `configure_mcp_command` (workspace validation, cross-workspace purge, installed/configured client resolution, auth, and section/notes/warnings), returning (workspace, profile, clients). - Extract `_split_catalog_schema(location, flag)` from `_resolve_location_mcp_servers`; the error message now names the flag. - Add an `always_load` keyword to `build_mcp_http_entry`; default off, no caller uses it yet. Existing tests stay green; adds unit tests for `_split_catalog_schema` and `build_mcp_http_entry(always_load=)`. Co-authored-by: Isaac --- src/ucode/mcp.py | 69 ++++++++++++++++++++++++++++++----------------- tests/test_mcp.py | 22 +++++++++++++++ 2 files changed, 67 insertions(+), 24 deletions(-) diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 6257458..17ec439 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -96,14 +96,17 @@ ) -def build_mcp_http_entry(url: str) -> dict: - return { +def build_mcp_http_entry(url: str, *, always_load: bool = False) -> dict: + entry: dict[str, Any] = { "type": "http", "url": url, "headers": { "Authorization": f"Bearer ${{{MCP_AUTH_TOKEN_ENV_VAR}}}", }, } + if always_load: + entry["alwaysLoad"] = True + return entry def add_claude_mcp_server(name: str, entry: dict, scope: str = MCP_USER_SCOPE) -> None: @@ -1103,6 +1106,15 @@ def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None: ) +def _split_catalog_schema(location: str, flag: str) -> tuple[str, str]: + """Split ``.`` into its two non-empty parts.""" + parts = location.split(".") + if len(parts) != 2 or not all(part.strip() for part in parts): + raise RuntimeError(f"{flag} must be `.`, got `{location}`.") + catalog, schema = parts + return catalog, schema + + def _resolve_location_mcp_servers( workspace: str, profile: str | None, @@ -1125,8 +1137,7 @@ def _resolve_location_mcp_servers( 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 `.`, got `{location}`.") + _split_catalog_schema(location, "--location") token = get_databricks_token(workspace, profile) with spinner(f"Discovering MCP services in {location}..."): @@ -1177,25 +1188,12 @@ def _resolve_location_mcp_servers( return working_servers -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 - # `.` 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.`): {', '.join(bare)}" - ) - if len(schemas) != 1: - raise RuntimeError( - "--services without --location must all share one `.` " - f"(got: {', '.join(sorted(schemas)) or 'none'}); pass --location instead." - ) - location = next(iter(schemas)) - state = load_state() +def _setup_mcp_clients(state: dict, section: str) -> 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. + """ workspace = state.get("workspace") if not workspace: raise RuntimeError("Workspace is not configured. Run `ucode configure` first.") @@ -1223,7 +1221,7 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None apply_pat_environment(state) ensure_databricks_auth(workspace, profile) - print_section("MCP Servers") + print_section(section) client_names = ", ".join(str(MCP_CLIENTS[client]["display"]) for client in clients) print_note(f"Configuring for: {client_names}") for client in missing_clients: @@ -1231,6 +1229,29 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None f"{MCP_CLIENTS[client]['display']} is configured in ucode but not installed; " "skipping MCP config." ) + return workspace, profile, clients + + +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 + # `.` 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.`): {', '.join(bare)}" + ) + if len(schemas) != 1: + raise RuntimeError( + "--services without --location must all share one `.` " + f"(got: {', '.join(sorted(schemas)) or 'none'}); pass --location instead." + ) + location = next(iter(schemas)) + state = load_state() + workspace, profile, clients = _setup_mcp_clients(state, "MCP Servers") original_mcp_servers_for_location: list[dict] = list(state.get("mcp_servers") or []) if location is not None: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 81333c8..8d4bd9b 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -25,6 +25,28 @@ def test_uses_oauth_token_header_reference(self): assert "oauth" not in entry assert "headersHelper" not in entry + def test_omits_always_load_by_default(self): + entry = mcp.build_mcp_http_entry(f"{WS}/ai-gateway/skills/") + assert "alwaysLoad" not in entry + + def test_sets_always_load_when_requested(self): + entry = mcp.build_mcp_http_entry(f"{WS}/ai-gateway/skills/", always_load=True) + assert entry["alwaysLoad"] is True + + +class TestSplitCatalogSchema: + def test_accepts_catalog_schema(self): + assert mcp._split_catalog_schema("a.b", "--location") == ("a", "b") + + def test_rejects_malformed_locations(self): + for bad in ("a", "a.b.c", ".b", "a.", ""): + try: + mcp._split_catalog_schema(bad, "--flag") + except RuntimeError as exc: + assert "--flag" in str(exc) + else: + raise AssertionError(f"expected RuntimeError for `{bad}`") + class TestAddClaudeMcpServer: def test_adds_user_scoped_json(self, monkeypatch):