From 99452b5149449151cd0a68437f2ccaa63f5b36c5 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Thu, 3 Sep 2026 00:42:33 +0000 Subject: [PATCH 1/3] Restrict interactive MCP picker to MCP services for consumer access The interactive picker previously offered V2 AI Gateway sources (Vector Search, UC Functions, external connections, Genie, Databricks apps, all served under /api/2.0/mcp/*). Consumer-only identities lack the workspace-access entitlement and can't reach those, so the picker now offers only MCP services (the consumer-safe /ai-gateway/mcp-services path) and skips the source-selection step when it's the sole source. Workspace users still add a V2 server on request non-interactively by naming it in --services with a typed selector (vector-search:cat.schema, uc-functions:cat.schema, external:conn, genie-space:id, app:name). Consumer detection uses the authoritative signal: the AI Gateway's WorkspaceAccessGuard returns a 403 whose message names the `workspace-access` entitlement, which is the only thing distinguishing a consumer-only identity from a workspace user missing a specific grant. PermissionDeniedError carries consumer_only accordingly, discovery skips it gracefully, and the two cases get different, accurate error messages. Co-authored-by: Isaac --- README.md | 23 ++++-- src/ucode/cli.py | 10 ++- src/ucode/databricks.py | 66 ++++++++++++++++ src/ucode/mcp.py | 167 +++++++++++++++++++++++++++++++++++---- tests/test_databricks.py | 73 +++++++++++++++++ tests/test_mcp.py | 154 +++++++++++++++++++++++++++++++++++- 6 files changed, 468 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 7dbea331..28d6994a 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 +ucode mcp add --services vector-search:main.docs +ucode mcp add --services uc-functions:main.tools +ucode mcp add --services external:my-connection +ucode mcp add --services genie-space: +ucode mcp add --services app:my-app +``` + +A consumer-only identity that requests one of these is stopped with a clear error before any +config is written. 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..0df8cd5e 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1391,6 +1391,47 @@ def _extract_apps_payload(payload: object) -> list[dict]: raise RuntimeError("Databricks apps listing returned invalid JSON.") +# The AI Gateway's WorkspaceAccessGuard (ai-gateway/src/WorkspaceAccessGuard.scala) rejects a +# consumer-only identity — one without the `workspace-access` entitlement — with an HTTP 403 +# whose message names that entitlement. That message substring is the ONLY thing that tells a +# consumer-only identity apart from a workspace user who merely lacks a grant on one resource: +# the server returns the same 403 / JsonRpc FORBIDDEN for both otherwise (no distinct error code). +_WORKSPACE_ACCESS_MARKER = "workspace-access" + + +def _looks_like_consumer_access_failure(text: str | None) -> bool: + """Whether a 403 message indicates a consumer-only identity (missing the `workspace-access` + entitlement) rather than a per-resource permission denial.""" + return bool(text) and _WORKSPACE_ACCESS_MARKER in text.lower() + + +class PermissionDeniedError(RuntimeError): + """A workspace API returned an authorization failure (HTTP 403 / permission denied). + + ``consumer_only`` is ``True`` when the failure names the `workspace-access` entitlement — + i.e. a consumer-only identity with no workspace access at all — and ``False`` when it's a + workspace user who merely lacks a grant on the specific resource. Callers use this both to + skip V2 MCP discovery gracefully instead of blocking setup (see the discovery wrappers in + :mod:`ucode.mcp`) and to word the error for the right audience.""" + + def __init__(self, message: str, *, consumer_only: bool = False) -> None: + super().__init__(message) + self.consumer_only = consumer_only + + +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 +1453,15 @@ 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 is either a consumer-only identity (no workspace access) or a workspace + # user lacking apps permission; classify by the workspace-access marker so callers can + # skip discovery gracefully (AIGTWY-4471) and word the error correctly. Other CLI + # failures stay hard errors. + if _looks_like_cli_permission_error(exc.stderr): + raise PermissionDeniedError( + "Not authorized to list Databricks apps.", + consumer_only=_looks_like_consumer_access_failure(exc.stderr), + ) 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 @@ -3247,6 +3297,22 @@ def _looks_like_permission_failure(reason: str) -> bool: return "HTTP 403" in reason +def consumer_access_reason(workspace: str, token: str) -> str | None: + """Return a reason string when the identity is provably consumer-only, else ``None``. + + Consumer entitlements don't grant the `workspace-access` entitlement, so V2 AI Gateway + features (Vector Search, UC Functions, external/Genie/app MCP servers) can't work for them. + Probes the V2 AI Gateway and reports a failure ONLY when it carries the workspace-access + marker — the authoritative consumer signal. Best-effort and deliberately conservative: + returns ``None`` on any other outcome (reachable, empty listing, or an unrelated error) so a + workspace user is never wrongly blocked. Note the server-side guard is gated by the + `blockInferenceWithoutWorkspaceAccess` SAFE flag, so this fires only where that is ramped.""" + probe = _probe_ai_gateway_v2(workspace, token) + if not probe.reachable and _looks_like_consumer_access_failure(probe.detail): + return probe.detail + return None + + CODING_AGENT_RECOMMEND_MODEL_PATH = "/api/ai-gateway/v2/coding-agent-configs:recommendModel" diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index f855d086..6bc06ddf 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -29,10 +29,12 @@ 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, build_skills_mcp_url, + consumer_access_reason, ensure_databricks_auth, get_databricks_token, list_all_mcp_services, @@ -1333,8 +1335,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 +1373,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 +1686,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 +1842,92 @@ def add_mcp_command( return configure_mcp_command(location=location, services=services, append=True, agents=agents) +def _consumer_access_error(reason: str) -> str: + """The error shown when a consumer-only identity (no `workspace-access` entitlement) tries + to add a V2 AI Gateway MCP server.""" + return ( + "This identity has consumer-only access (no workspace-access entitlement), which can't " + "use V2 AI Gateway MCP servers (Vector Search, UC Functions, external connections, Genie, " + f"apps): {reason}. Ask a workspace admin for the workspace-access entitlement, or register " + "an MCP service instead (`ucode mcp add --location .`)." + ) + + +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. A provably consumer-only identity is blocked with a clear error + before any config is written (best-effort — see `consumer_access_reason`; the server-side + guard is SAFE-flag gated). 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 + ) + + token = get_databricks_token(workspace, profile) + reason = consumer_access_reason(workspace, token) + if reason is not None: + raise RuntimeError(_consumer_access_error(reason)) + + # `app:` selectors need the app's off-workspace URL, which only discovery knows. A 403 here + # is worded by whether it's a consumer identity or a workspace user missing 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: + if exc.consumer_only: + raise RuntimeError(_consumer_access_error(str(exc))) from exc + raise RuntimeError( + f"{exc} You have workspace access but lack permission to list Databricks apps; " + "ask the app owner to grant you access." + ) 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 +1944,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 +2018,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 +2040,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..ce839f60 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1951,6 +1951,79 @@ def fake_run(args, **kwargs): with pytest.raises(RuntimeError, match="invalid JSON"): list_databricks_apps(WS) + def test_resource_permission_failure_is_not_consumer_only(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) as exc: + list_databricks_apps(WS) + assert exc.value.consumer_only is False + + def test_workspace_access_failure_is_consumer_only(self, monkeypatch): + def fake_run(args, **kwargs): + raise subprocess.CalledProcessError( + 1, + "databricks", + stderr="Error: 403 The workspace-access entitlement is required.", + ) + + monkeypatch.setattr(db_mod, "run", fake_run) + + with pytest.raises(db_mod.PermissionDeniedError) as exc: + list_databricks_apps(WS) + assert exc.value.consumer_only is True + + 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 TestConsumerAccessReason: + def test_returns_reason_on_workspace_access_marker(self, monkeypatch): + detail = "HTTP 403: The workspace-access entitlement is required to query AI Gateway." + monkeypatch.setattr( + db_mod, + "_probe_ai_gateway_v2", + lambda workspace, token: db_mod.GatewayProbe(False, detail), + ) + assert db_mod.consumer_access_reason(WS, "tok") == detail + + def test_returns_none_when_reachable(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_probe_ai_gateway_v2", + lambda workspace, token: db_mod.GatewayProbe(True, "reachable", True), + ) + assert db_mod.consumer_access_reason(WS, "tok") is None + + def test_returns_none_on_permission_failure_without_marker(self, monkeypatch): + # A 403 that doesn't name the workspace-access entitlement is a resource-permission + # denial, not a consumer identity — don't wrongly block the workspace user. + monkeypatch.setattr( + db_mod, + "_probe_ai_gateway_v2", + lambda workspace, token: db_mod.GatewayProbe(False, "HTTP 403: missing USE_SCHEMA"), + ) + assert db_mod.consumer_access_reason(WS, "tok") is None + + def test_returns_none_on_non_permission_failure(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_probe_ai_gateway_v2", + lambda workspace, token: db_mod.GatewayProbe(False, "timeout (HTTP 500)"), + ) + assert db_mod.consumer_access_reason(WS, "tok") is None + 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..e5176980 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,125 @@ 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, "consumer_access_reason", lambda workspace, token: None) + 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_non_interactive_v2_add_blocks_consumer(self, monkeypatch): + self._base_mocks(monkeypatch) + monkeypatch.setattr( + mcp, "consumer_access_reason", lambda workspace, token: "unreachable (HTTP 403)" + ) + monkeypatch.setattr(mcp, "save_state", lambda state: pytest.fail("must not save on block")) + + with pytest.raises(RuntimeError, match="consumer-only access"): + mcp.configure_mcp_command(services={"uc-functions:main.tools"}) + + 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 From 7c3a6178982f1331428bea91d0121002620db16a Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 4 Sep 2026 21:55:55 +0000 Subject: [PATCH 2/3] Drop unreliable consumer-access detection from V2 MCP add MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Verified against universe: the `workspace-access` entitlement 403 marker is emitted only by guarded AI Gateway / Model Serving inference + model-listing paths (which ucode already exercises at model setup) — NOT by the Apps / UC / Vector Search / ai-gateway-v2-endpoints listing calls the MCP flow uses, which return an empty list or a generic ACL denial for a consumer. So there is no reliable consumer signal at MCP-add time, and a consumer is gated upstream. Remove the marker-based detection that can't work here: drop consumer_access_reason (it probed an unguarded listing and never fired) and the preemptive block, and stop classifying PermissionDeniedError as consumer_only. Keep the graceful skip on any 403; an app: add that 403s now gives an actionable "needs workspace access" error without misattributing the cause. Co-authored-by: Isaac --- README.md | 4 +-- src/ucode/databricks.py | 59 +++++++++------------------------------- src/ucode/mcp.py | 37 ++++++------------------- tests/test_databricks.py | 56 ++------------------------------------ tests/test_mcp.py | 17 ++++++------ 5 files changed, 35 insertions(+), 138 deletions(-) diff --git a/README.md b/README.md index 28d6994a..7e805e84 100644 --- a/README.md +++ b/README.md @@ -130,8 +130,8 @@ ucode mcp add --services genie-space: ucode mcp add --services app:my-app ``` -A consumer-only identity that requests one of these is stopped with a clear error before any -config is written. +These require workspace access; a consumer-only identity is gated at the AI Gateway (which +`ucode` 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/databricks.py b/src/ucode/databricks.py index 0df8cd5e..3ca8eec0 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1391,32 +1391,19 @@ def _extract_apps_payload(payload: object) -> list[dict]: raise RuntimeError("Databricks apps listing returned invalid JSON.") -# The AI Gateway's WorkspaceAccessGuard (ai-gateway/src/WorkspaceAccessGuard.scala) rejects a -# consumer-only identity — one without the `workspace-access` entitlement — with an HTTP 403 -# whose message names that entitlement. That message substring is the ONLY thing that tells a -# consumer-only identity apart from a workspace user who merely lacks a grant on one resource: -# the server returns the same 403 / JsonRpc FORBIDDEN for both otherwise (no distinct error code). -_WORKSPACE_ACCESS_MARKER = "workspace-access" - - -def _looks_like_consumer_access_failure(text: str | None) -> bool: - """Whether a 403 message indicates a consumer-only identity (missing the `workspace-access` - entitlement) rather than a per-resource permission denial.""" - return bool(text) and _WORKSPACE_ACCESS_MARKER in text.lower() - - class PermissionDeniedError(RuntimeError): """A workspace API returned an authorization failure (HTTP 403 / permission denied). - ``consumer_only`` is ``True`` when the failure names the `workspace-access` entitlement — - i.e. a consumer-only identity with no workspace access at all — and ``False`` when it's a - workspace user who merely lacks a grant on the specific resource. Callers use this both to - skip V2 MCP discovery gracefully instead of blocking setup (see the discovery wrappers in - :mod:`ucode.mcp`) and to word the error for the right audience.""" + 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. - def __init__(self, message: str, *, consumer_only: bool = False) -> None: - super().__init__(message) - self.consumer_only = consumer_only + 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: @@ -1453,15 +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 is either a consumer-only identity (no workspace access) or a workspace - # user lacking apps permission; classify by the workspace-access marker so callers can - # skip discovery gracefully (AIGTWY-4471) and word the error correctly. Other CLI - # failures stay hard errors. + # 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.", - consumer_only=_looks_like_consumer_access_failure(exc.stderr), - ) from exc + 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 @@ -3297,22 +3280,6 @@ def _looks_like_permission_failure(reason: str) -> bool: return "HTTP 403" in reason -def consumer_access_reason(workspace: str, token: str) -> str | None: - """Return a reason string when the identity is provably consumer-only, else ``None``. - - Consumer entitlements don't grant the `workspace-access` entitlement, so V2 AI Gateway - features (Vector Search, UC Functions, external/Genie/app MCP servers) can't work for them. - Probes the V2 AI Gateway and reports a failure ONLY when it carries the workspace-access - marker — the authoritative consumer signal. Best-effort and deliberately conservative: - returns ``None`` on any other outcome (reachable, empty listing, or an unrelated error) so a - workspace user is never wrongly blocked. Note the server-side guard is gated by the - `blockInferenceWithoutWorkspaceAccess` SAFE flag, so this fires only where that is ramped.""" - probe = _probe_ai_gateway_v2(workspace, token) - if not probe.reachable and _looks_like_consumer_access_failure(probe.detail): - return probe.detail - return None - - CODING_AGENT_RECOMMEND_MODEL_PATH = "/api/ai-gateway/v2/coding-agent-configs:recommendModel" diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 6bc06ddf..ad9c7818 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -34,7 +34,6 @@ build_mcp_proxy_argv, build_mcp_service_url, build_skills_mcp_url, - consumer_access_reason, ensure_databricks_auth, get_databricks_token, list_all_mcp_services, @@ -1842,17 +1841,6 @@ def add_mcp_command( return configure_mcp_command(location=location, services=services, append=True, agents=agents) -def _consumer_access_error(reason: str) -> str: - """The error shown when a consumer-only identity (no `workspace-access` entitlement) tries - to add a V2 AI Gateway MCP server.""" - return ( - "This identity has consumer-only access (no workspace-access entitlement), which can't " - "use V2 AI Gateway MCP servers (Vector Search, UC Functions, external connections, Genie, " - f"apps): {reason}. Ask a workspace admin for the workspace-access entitlement, or register " - "an MCP service instead (`ucode mcp add --location .`)." - ) - - def _configure_v2_mcp_selectors( selectors: list[str], *, @@ -1862,34 +1850,27 @@ def _configure_v2_mcp_selectors( """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. A provably consumer-only identity is blocked with a clear error - before any config is written (best-effort — see `consumer_access_reason`; the server-side - guard is SAFE-flag gated). Registration mirrors the interactive add path: additive under - ``append`` (`ucode mcp add`), an exact replacement otherwise (`ucode configure mcp`), - always preserving the skills connection.""" + 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 ) - token = get_databricks_token(workspace, profile) - reason = consumer_access_reason(workspace, token) - if reason is not None: - raise RuntimeError(_consumer_access_error(reason)) - # `app:` selectors need the app's off-workspace URL, which only discovery knows. A 403 here - # is worded by whether it's a consumer identity or a workspace user missing apps permission. + # 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: - if exc.consumer_only: - raise RuntimeError(_consumer_access_error(str(exc))) from exc raise RuntimeError( - f"{exc} You have workspace access but lack permission to list Databricks apps; " - "ask the app owner to grant you access." + 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 []) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index ce839f60..a17657d5 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1951,7 +1951,7 @@ def fake_run(args, **kwargs): with pytest.raises(RuntimeError, match="invalid JSON"): list_databricks_apps(WS) - def test_resource_permission_failure_is_not_consumer_only(self, monkeypatch): + 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" @@ -1959,23 +1959,8 @@ def fake_run(args, **kwargs): monkeypatch.setattr(db_mod, "run", fake_run) - with pytest.raises(db_mod.PermissionDeniedError) as exc: + with pytest.raises(db_mod.PermissionDeniedError): list_databricks_apps(WS) - assert exc.value.consumer_only is False - - def test_workspace_access_failure_is_consumer_only(self, monkeypatch): - def fake_run(args, **kwargs): - raise subprocess.CalledProcessError( - 1, - "databricks", - stderr="Error: 403 The workspace-access entitlement is required.", - ) - - monkeypatch.setattr(db_mod, "run", fake_run) - - with pytest.raises(db_mod.PermissionDeniedError) as exc: - list_databricks_apps(WS) - assert exc.value.consumer_only is True def test_non_permission_cli_failure_stays_generic_runtime_error(self, monkeypatch): def fake_run(args, **kwargs): @@ -1988,43 +1973,6 @@ def fake_run(args, **kwargs): assert not isinstance(exc.value, db_mod.PermissionDeniedError) -class TestConsumerAccessReason: - def test_returns_reason_on_workspace_access_marker(self, monkeypatch): - detail = "HTTP 403: The workspace-access entitlement is required to query AI Gateway." - monkeypatch.setattr( - db_mod, - "_probe_ai_gateway_v2", - lambda workspace, token: db_mod.GatewayProbe(False, detail), - ) - assert db_mod.consumer_access_reason(WS, "tok") == detail - - def test_returns_none_when_reachable(self, monkeypatch): - monkeypatch.setattr( - db_mod, - "_probe_ai_gateway_v2", - lambda workspace, token: db_mod.GatewayProbe(True, "reachable", True), - ) - assert db_mod.consumer_access_reason(WS, "tok") is None - - def test_returns_none_on_permission_failure_without_marker(self, monkeypatch): - # A 403 that doesn't name the workspace-access entitlement is a resource-permission - # denial, not a consumer identity — don't wrongly block the workspace user. - monkeypatch.setattr( - db_mod, - "_probe_ai_gateway_v2", - lambda workspace, token: db_mod.GatewayProbe(False, "HTTP 403: missing USE_SCHEMA"), - ) - assert db_mod.consumer_access_reason(WS, "tok") is None - - def test_returns_none_on_non_permission_failure(self, monkeypatch): - monkeypatch.setattr( - db_mod, - "_probe_ai_gateway_v2", - lambda workspace, token: db_mod.GatewayProbe(False, "timeout (HTTP 500)"), - ) - assert db_mod.consumer_access_reason(WS, "tok") is None - - class TestProbeUnityGatewayCapabilities: def test_model_service_resource_skips_legacy_probe(self, monkeypatch): calls: list[str] = [] diff --git a/tests/test_mcp.py b/tests/test_mcp.py index e5176980..9100b778 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -3057,7 +3057,6 @@ 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, "consumer_access_reason", lambda workspace, token: None) monkeypatch.setattr( mcp, "configure_client_mcp_server", @@ -3083,15 +3082,17 @@ def test_non_interactive_vector_search_add(self, monkeypatch): } ] - def test_non_interactive_v2_add_blocks_consumer(self, monkeypatch): + def test_app_add_permission_failure_is_actionable(self, monkeypatch): self._base_mocks(monkeypatch) - monkeypatch.setattr( - mcp, "consumer_access_reason", lambda workspace, token: "unreachable (HTTP 403)" - ) - monkeypatch.setattr(mcp, "save_state", lambda state: pytest.fail("must not save on block")) - with pytest.raises(RuntimeError, match="consumer-only access"): - mcp.configure_mcp_command(services={"uc-functions:main.tools"}) + 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")) From 1ff3fcc0cbedf8ee931838d9adca8bde2f1823c5 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 4 Sep 2026 22:11:41 +0000 Subject: [PATCH 3/3] Use ug branding in the V2 --services README block after rebase Co-authored-by: Isaac --- README.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 7e805e84..ec48d64c 100644 --- a/README.md +++ b/README.md @@ -123,15 +123,15 @@ reach the V2 AI Gateway. Workspace users add them non-interactively by naming th `--services` with a typed selector: ```bash -ucode mcp add --services vector-search:main.docs -ucode mcp add --services uc-functions:main.tools -ucode mcp add --services external:my-connection -ucode mcp add --services genie-space: -ucode mcp add --services app:my-app +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 -`ucode` already hits when it sets up models), not by this command. +`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