diff --git a/README.md b/README.md index 7dbea331..ec48d64c 100644 --- a/README.md +++ b/README.md @@ -113,14 +113,25 @@ ug configure mcp ``` Add Databricks MCP servers to installed MCP-capable tools: Codex, Claude Code, Gemini CLI, OpenCode, GitHub Copilot CLI, and Cursor Agent. -Options are shown in this order: -- Discovered external MCP connections -- Databricks SQL -- Managed Databricks MCPs (Vector Search, UC Functions, etc.) -- Custom MCP server URL +The interactive picker discovers **MCP services** (the `system.ai.*` and workspace-wide +`.` Unity Catalog MCP services), plus Databricks SQL and a custom MCP server URL. -Discovered external MCP connections are listed directly. +V2 AI Gateway servers — Vector Search, UC Functions, external connections, Genie spaces, and +Databricks apps — are **not** offered in the picker, because consumer-only identities can't +reach the V2 AI Gateway. Workspace users add them non-interactively by naming them in +`--services` with a typed selector: + +```bash +ug mcp add --services vector-search:main.docs +ug mcp add --services uc-functions:main.tools +ug mcp add --services external:my-connection +ug mcp add --services genie-space: +ug mcp add --services app:my-app +``` + +These require workspace access; a consumer-only identity is gated at the AI Gateway (which +`ug` already hits when it sets up models), not by this command. Every Databricks MCP server is registered as a local **stdio** server that runs `ug mcp-proxy` — a small bridge (shipped with `ug`) between the coding tool and the Databricks diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 60d4f1cb..70e73f76 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1242,7 +1242,10 @@ def mcp_add( help="Register this comma-separated subset of MCP services (additively). Full names " "like `system.ai.github` work on their own; bare short names like `github` need " "--location to locate them. Omit --services to register the whole --location schema; " - 'an empty `--services ""` adds nothing (no-op).', + 'an empty `--services ""` adds nothing (no-op). V2 AI Gateway servers (not in the ' + "interactive picker) are added by naming them here: `vector-search:.`, " + "`uc-functions:.`, `external:`, `genie-space:`, or " + "`app:` (workspace access required).", ), ] = None, agents: Annotated[ @@ -3059,7 +3062,10 @@ def configure_mcp( "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.", + "(with --location) to remove all. V2 AI Gateway servers (not in the interactive " + "picker) are named directly: `vector-search:.`, " + "`uc-functions:.`, `external:`, `genie-space:`, or " + "`app:` (workspace access required).", ), ] = None, ) -> None: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index ee54d153..3ca8eec0 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1391,6 +1391,34 @@ def _extract_apps_payload(payload: object) -> list[dict]: raise RuntimeError("Databricks apps listing returned invalid JSON.") +class PermissionDeniedError(RuntimeError): + """A workspace API returned an authorization failure (HTTP 403 / permission denied). + + Callers use this only to skip V2 MCP discovery gracefully instead of aborting setup (see the + discovery wrappers in :mod:`ucode.mcp`), while other errors still surface. + + It deliberately does NOT try to distinguish a consumer-only identity (no `workspace-access` + entitlement) from a workspace user missing a grant on a specific resource: no service exposes + a signal that reliably tells them apart. The `workspace-access` entitlement is enforced with a + named 403 only on guarded AI Gateway / Model Serving *inference* and model-listing paths (which + ucode already exercises at model setup, so a consumer is gated there) — NOT on the Apps / UC / + Vector Search listing calls the MCP flow uses, which return an empty list or a generic ACL + denial for a consumer.""" + + +def _looks_like_cli_permission_error(stderr: str | None) -> bool: + """Whether a Databricks CLI stderr indicates an authorization failure. + + The CLI exit code is generic, so we match on the stable markers the CLI/API emit + for a denied workspace call rather than the status alone.""" + if not stderr: + return False + lowered = stderr.lower() + if "permission" in lowered and ("denied" in lowered or "insufficient" in lowered): + return True + return "403" in lowered or "not authorized" in lowered or "unauthorized" in lowered + + def list_databricks_apps(workspace: str, profile: str | None = None) -> list[dict]: env = build_databricks_cli_env(workspace) try: @@ -1412,6 +1440,11 @@ def list_databricks_apps(workspace: str, profile: str | None = None) -> list[dic ) return _extract_apps_payload(json.loads(result.stdout or "[]")) except subprocess.CalledProcessError as exc: + # A 403 here means the caller isn't authorized to list apps (a consumer-only identity, or + # a workspace user without apps permission); raise PermissionDeniedError so callers can skip + # discovery gracefully (AIGTWY-4471). Other CLI failures stay hard errors. + if _looks_like_cli_permission_error(exc.stderr): + raise PermissionDeniedError("Not authorized to list Databricks apps.") from exc raise RuntimeError("Failed to list Databricks apps via `databricks apps list`.") from exc except subprocess.TimeoutExpired as exc: raise RuntimeError("Timed out while listing Databricks apps.") from exc diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index f855d086..ad9c7818 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -29,6 +29,7 @@ from ucode.agents import copilot, cursor, gemini, opencode from ucode.config_io import restore_file from ucode.databricks import ( + PermissionDeniedError, apply_pat_environment, build_mcp_proxy_argv, build_mcp_service_url, @@ -1333,8 +1334,13 @@ def _discover_mcp_source(label: str, discover: Callable[[], list[Any]]) -> list[ try: with spinner(f"Discovering {label}..."): return discover() + except PermissionDeniedError: + # Consumer-only identities lack workspace access, so this source 403s for them. + # Skip it quietly (not as a scary warning) so setup completes (AIGTWY-4471). + print_note(f"Skipped {label} (no workspace access).") + return [] except (RuntimeError, OSError) as exc: - # Discovery is best-effort: a failure here (auth error, network timeout) + # Discovery is best-effort: a failure here (network timeout, transient error) # skips just this source so the rest of the picker still works. print_warning(f"Skipped {label} ({exc}).") return [] @@ -1366,6 +1372,10 @@ def message() -> str: try: with spinner(message): return discover(on_progress) + except PermissionDeniedError: + # See `_discover_mcp_source`: a consumer-only identity's 403 is a quiet skip. + print_note(f"Skipped {label} (no workspace access).") + return [] except (RuntimeError, OSError) as exc: print_warning(f"Skipped {label} ({exc}).") return [] @@ -1675,19 +1685,35 @@ def _resolve_location_mcp_servers( # The first wizard step lets the user choose which sources to search. Each is a -# (key, label, default_checked) triple. Vector Search and UC functions default -# off because they walk the workspace (endpoints/catalogs/schemas) and are slow; -# everything else is a cheap listing and defaults on. -MCP_SEARCH_SOURCES = ( - ("external", "External connections", True), - ("apps", "Databricks apps", True), - ("mcp-services", "MCP services", True), - ("genie", "Genie spaces", True), - ("vector-search", "Vector Search indexes (slower)", False), - ("uc-functions", "UC functions (slower)", False), +# (key, label, default_checked) triple. +# +# Only MCP services (the `/ai-gateway/mcp-services/` path) are offered interactively: +# it's the one source a consumer-only identity can reach. The V2 AI Gateway sources — +# external connections, Databricks apps, Genie spaces, Vector Search, and UC functions, +# all served under `/api/2.0/mcp/*` — were removed from the picker because consumer +# entitlements don't grant access to V2 AI Gateway features. Workspace users who still +# want one add it non-interactively with a typed `--services` selector (see +# `V2_MCP_SELECTOR_PREFIXES` and `_configure_v2_mcp_selectors`). +MCP_SEARCH_SOURCES = (("mcp-services", "MCP services", True),) + +# Typed `--services` selectors that name a V2 AI Gateway MCP server directly, e.g. +# `vector-search:main.docs` or `uc-functions:main.tools`. These bypass the interactive +# picker (which no longer offers V2 sources) so workspace users can still add them on +# request; a consumer-only identity is blocked with a clear error before registering. +V2_MCP_SELECTOR_PREFIXES = ( + VECTOR_SEARCH_SELECTION_PREFIX, + UC_FUNCTIONS_SELECTION_PREFIX, + EXTERNAL_MCP_SELECTION_PREFIX, + GENIE_SPACE_SELECTION_PREFIX, + APP_MCP_SELECTION_PREFIX, ) +def _is_v2_mcp_selector(service: str) -> bool: + """Whether a `--services` entry is a typed V2 MCP selector (see `V2_MCP_SELECTOR_PREFIXES`).""" + return service.startswith(V2_MCP_SELECTOR_PREFIXES) + + def prompt_for_mcp_search_sources(exclude_sources: set[str] | None = None) -> set[str] | None: """First wizard step: choose which sources to search. Returns the set of selected source keys, or `None` if the user cancelled (Ctrl-C). @@ -1815,6 +1841,74 @@ def add_mcp_command( return configure_mcp_command(location=location, services=services, append=True, agents=agents) +def _configure_v2_mcp_selectors( + selectors: list[str], + *, + append: bool, + agents: set[str] | None, +) -> int: + """Non-interactive add for V2 AI Gateway MCP servers named by typed `--services` + selectors (`vector-search:`/`uc-functions:`/`external:`/`genie-space:`/`app:`). + + The interactive picker no longer offers these sources; this is how a workspace user adds one + on request. These require workspace access, which consumer-only identities lack — but that's + enforced upstream at the AI Gateway (which ucode already hits during model setup), not here: + the listing calls this uses don't reliably signal consumer access (see `PermissionDeniedError`). + Registration mirrors the interactive add path: additive under ``append`` (`ucode mcp add`), an + exact replacement otherwise (`ucode configure mcp`), always preserving the skills connection.""" + state = load_state() + workspace, profile, clients = setup_mcp_clients( + state, "Add MCP Servers" if append else "MCP Servers", agents=agents + ) + + # `app:` selectors need the app's off-workspace URL, which only discovery knows. A 403 here + # means the caller can't list apps (no workspace access, or no apps permission). + available_app_servers: list[dict] = [] + if any(s.startswith(APP_MCP_SELECTION_PREFIX) for s in selectors): + try: + available_app_servers = discover_app_mcp_servers(workspace, profile) + except PermissionDeniedError as exc: + raise RuntimeError( + f"{exc} This needs workspace access to the Databricks apps listing; ask a " + "workspace admin if you're missing it." + ) from exc + + original_mcp_servers: list[dict] = list(state.get("mcp_servers") or []) + skills_servers = _skills_entries(original_mcp_servers) + picker_servers = [s for s in original_mcp_servers if s.get("kind") != SKILLS_MCP_KIND] + original_by_name = _servers_by_name(picker_servers) + + working_mcp_servers: list[dict] = list(skills_servers) + working_names: set[str] = set() + for selection in selectors: + entry_name, url = _resolve_mcp_selection(selection, workspace, available_app_servers) + if entry_name in working_names: + continue + working_mcp_servers.append( + {"name": entry_name, "url": url, "auth": "proxy", "clients": clients} + ) + working_names.add(entry_name) + + if append: + working_mcp_servers = _union_missing(original_mcp_servers, working_mcp_servers) + + changed = apply_mcp_server_changes( + original_mcp_servers, + working_mcp_servers, + clients, + workspace, + profile, + use_pat=bool(state.get("use_pat")), + ) + if changed or original_mcp_servers != working_mcp_servers: + state["mcp_servers"] = working_mcp_servers + save_state(state) + added = sorted(working_names - set(original_by_name)) + removed = [] if append else sorted(set(original_by_name) - working_names) + print_success(_mcp_change_summary(added, removed, clients)) + return 0 + + def configure_mcp_command( location: str | None = None, services: set[str] | None = None, @@ -1831,6 +1925,21 @@ def configure_mcp_command( final server list is unioned with the already-configured servers, so nothing outside the current selection is removed. ``agents`` scopes the operation to that subset of configured MCP clients.""" + if services is not None: + # A typed V2 MCP selector (`vector-search:main.docs`, `uc-functions:main.tools`, + # `external:conn`, `genie-space:`, `app:`) names a server the picker no + # longer offers. Route it through the dedicated non-interactive path so workspace + # users can still add it on request; consumer-only identities are blocked there. + v2_selectors = sorted(s for s in services if _is_v2_mcp_selector(s)) + if v2_selectors: + other = sorted(s for s in services if not _is_v2_mcp_selector(s)) + if other or location is not None: + raise RuntimeError( + "V2 MCP selectors (vector-search:/uc-functions:/external:/genie-space:/app:) " + "can't be combined with --location or plain MCP-service names in one call; add " + "them in a separate command." + ) + return _configure_v2_mcp_selectors(v2_selectors, append=append, agents=agents) 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 @@ -1890,11 +1999,18 @@ def configure_mcp_command( # Two-step wizard: (1) choose which sources to search, (2) pick servers from # the results. Pressing Left (←) in the picker returns to step 1, so the user - # can revise their source selection without restarting the command. + # can revise their source selection without restarting the command. When only + # one search source is available (MCP services — the V2 sources were removed), + # step 1 has nothing to choose, so skip it and go straight to the picker. + available_source_keys = [k for k, _, _ in MCP_SEARCH_SOURCES if k not in excluded_sources] + prompt_sources = len(available_source_keys) > 1 while True: - sources = prompt_for_mcp_search_sources(exclude_sources=excluded_sources) - if sources is None: - return 0 + if prompt_sources: + sources = prompt_for_mcp_search_sources(exclude_sources=excluded_sources) + if sources is None: + return 0 + else: + sources = set(available_source_keys) discovered = _discover_selected_mcp_sources(workspace, profile, sources) selections = prompt_for_mcp_server_choices( @@ -1905,7 +2021,7 @@ def configure_mcp_command( discovered["services"], discovered["vector_search"], discovered["uc_functions"], - allow_back=True, + allow_back=prompt_sources, additive=append, ) if selections is None: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 7debda88..a17657d5 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1951,6 +1951,27 @@ def fake_run(args, **kwargs): with pytest.raises(RuntimeError, match="invalid JSON"): list_databricks_apps(WS) + def test_permission_failure_raises_permission_denied_error(self, monkeypatch): + def fake_run(args, **kwargs): + raise subprocess.CalledProcessError( + 1, "databricks", stderr="Error: permission denied on apps" + ) + + monkeypatch.setattr(db_mod, "run", fake_run) + + with pytest.raises(db_mod.PermissionDeniedError): + list_databricks_apps(WS) + + def test_non_permission_cli_failure_stays_generic_runtime_error(self, monkeypatch): + def fake_run(args, **kwargs): + raise subprocess.CalledProcessError(1, "databricks", stderr="Error: connection reset") + + monkeypatch.setattr(db_mod, "run", fake_run) + + with pytest.raises(RuntimeError) as exc: + list_databricks_apps(WS) + assert not isinstance(exc.value, db_mod.PermissionDeniedError) + class TestProbeUnityGatewayCapabilities: def test_model_service_resource_skips_legacy_probe(self, monkeypatch): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index d58062b6..9100b778 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -616,6 +616,13 @@ def _patch_mcp_choices(monkeypatch, *values: str, categories: set[str] | None = # `categories`, which are unioned in. default_sources = {"external", "apps", "mcp-services", "genie"} selected_sources = default_sources | (categories or set()) + # Production ships a single search source (MCP services), so the source-selection + # step is skipped. These tests exercise the server-selection machinery (which still + # supports every server type via non-interactive add / managed configs), so restore a + # multi-source picker to keep that step live and use the stubbed source selection. + monkeypatch.setattr( + mcp, "MCP_SEARCH_SOURCES", tuple((key, key, True) for key in sorted(selected_sources)) + ) monkeypatch.setattr( mcp, "prompt_for_mcp_search_sources", lambda exclude_sources=None: selected_sources ) @@ -707,6 +714,12 @@ def test_back_reshows_source_screen(self, monkeypatch): 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"]) + # Multiple sources keep the source-selection step live (production ships one). + monkeypatch.setattr( + mcp, + "MCP_SEARCH_SOURCES", + (("external", "external", True), ("mcp-services", "MCP services", True)), + ) monkeypatch.setattr( mcp, "discover_external_mcp_connection_names", lambda workspace, profile=None: [] ) @@ -723,7 +736,7 @@ def test_back_reshows_source_screen(self, monkeypatch): def fake_sources(exclude_sources=None): source_calls.append(1) - return {"external", "apps", "mcp-services", "genie"} + return {"external", "mcp-services"} monkeypatch.setattr(mcp, "prompt_for_mcp_search_sources", fake_sources) @@ -745,6 +758,12 @@ def test_cancel_on_source_screen_exits(self, monkeypatch): 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"]) + # Multiple sources keep the source-selection step live (production ships one). + monkeypatch.setattr( + mcp, + "MCP_SEARCH_SOURCES", + (("external", "external", True), ("mcp-services", "MCP services", True)), + ) # Cancelling the first screen (None) returns without discovering anything. monkeypatch.setattr(mcp, "prompt_for_mcp_search_sources", lambda exclude_sources=None: None) monkeypatch.setattr( @@ -2968,6 +2987,17 @@ def fake_checkbox(*args, **kwargs): captured["choices"] = kwargs["choices"] return FakePrompt() + # Simulate a multi-source picker so exclusion is observable (production ships one). + monkeypatch.setattr( + mcp, + "MCP_SEARCH_SOURCES", + ( + ("external", "external", True), + ("apps", "apps", True), + ("mcp-services", "MCP services", True), + ("genie", "genie", True), + ), + ) monkeypatch.setattr(mcp, "_scrolling_checkbox", fake_checkbox) mcp.prompt_for_mcp_search_sources(exclude_sources=exclude) return [c.value for c in captured["choices"]] @@ -2998,3 +3028,126 @@ def test_workspace_relative_shapes_are_not_apps(self): def test_non_string_url_is_not_an_app(self): assert mcp._is_app_mcp_server({}) is False + + +class TestV2McpSelectors: + def test_is_v2_mcp_selector_recognizes_prefixes(self): + assert mcp._is_v2_mcp_selector("vector-search:main.docs") + assert mcp._is_v2_mcp_selector("uc-functions:main.tools") + assert mcp._is_v2_mcp_selector("external:my-conn") + assert mcp._is_v2_mcp_selector("genie-space:123") + assert mcp._is_v2_mcp_selector("app:my-app") + + def test_is_v2_mcp_selector_rejects_mcp_service_names(self): + assert not mcp._is_v2_mcp_selector("system.ai.github") + assert not mcp._is_v2_mcp_selector("github") + + def test_mcp_search_sources_only_offers_mcp_services(self): + # V2 AI Gateway sources are removed from the interactive picker. + assert [key for key, _, _ in mcp.MCP_SEARCH_SOURCES] == ["mcp-services"] + + def _base_mocks(self, monkeypatch): + monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_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, "get_databricks_token", lambda workspace, profile=None: "tok") + + def test_non_interactive_vector_search_add(self, monkeypatch): + saved_states: list[dict] = [] + configured: list[tuple[str, str, str]] = [] + self._base_mocks(monkeypatch) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, *a, **kw: configured.append((client, name, url)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert mcp.configure_mcp_command(services={"vector-search:main.docs"}) == 0 + + assert configured == [ + ( + "claude", + "databricks-vector-search-main-docs", + f"{WS}/api/2.0/mcp/vector-search/main/docs", + ) + ] + assert saved_states[-1]["mcp_servers"] == [ + { + "name": "databricks-vector-search-main-docs", + "url": f"{WS}/api/2.0/mcp/vector-search/main/docs", + "auth": "proxy", + "clients": ["claude"], + } + ] + + def test_app_add_permission_failure_is_actionable(self, monkeypatch): + self._base_mocks(monkeypatch) + + def deny(workspace, profile=None): + raise mcp.PermissionDeniedError("Not authorized to list Databricks apps.") + + monkeypatch.setattr(mcp, "discover_app_mcp_servers", deny) + monkeypatch.setattr(mcp, "save_state", lambda state: pytest.fail("must not save")) + + with pytest.raises(RuntimeError, match="workspace access"): + mcp.configure_mcp_command(services={"app:my-app"}) + + def test_v2_selector_cannot_combine_with_location(self, monkeypatch): + monkeypatch.setattr(mcp, "load_state", lambda: pytest.fail("must not reach load_state")) + with pytest.raises(RuntimeError, match="can't be combined"): + mcp.configure_mcp_command(location="system.ai", services={"vector-search:main.docs"}) + + def test_v2_selector_cannot_combine_with_plain_service(self, monkeypatch): + monkeypatch.setattr(mcp, "load_state", lambda: pytest.fail("must not reach load_state")) + with pytest.raises(RuntimeError, match="can't be combined"): + mcp.configure_mcp_command(services={"vector-search:main.docs", "system.ai.github"}) + + +class TestSingleSourceSkipsPrompt: + def test_source_prompt_skipped_and_no_back(self, monkeypatch): + monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_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, "discover_mcp_service_names", lambda workspace, profile=None: []) + monkeypatch.setattr( + mcp, + "discover_all_mcp_service_names", + lambda workspace, profile=None, on_progress=None: [], + ) + monkeypatch.setattr( + mcp, + "prompt_for_mcp_search_sources", + lambda exclude_sources=None: pytest.fail("source prompt must be skipped"), + ) + captured: dict = {} + + def fake_choices(*args, **kwargs): + captured["allow_back"] = kwargs.get("allow_back") + return [] + + monkeypatch.setattr(mcp, "prompt_for_mcp_server_choices", fake_choices) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + + assert mcp.configure_mcp_command() == 0 + assert captured["allow_back"] is False + + +class TestDiscoverySkipsPermissionErrors: + def test_discover_mcp_source_skips_permission_denied_quietly(self, monkeypatch, capsys): + def boom(): + raise mcp.PermissionDeniedError("no workspace access") + + assert mcp._discover_mcp_source("Databricks apps", boom) == [] + out = capsys.readouterr().out + assert "Skipped Databricks apps" in out + + def test_discover_mcp_source_warns_on_other_errors(self, monkeypatch, capsys): + def boom(): + raise RuntimeError("network down") + + assert mcp._discover_mcp_source("Genie spaces", boom) == [] + out = capsys.readouterr().out + assert "network down" in out