Restrict interactive MCP picker to MCP services for consumer access

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
+`<catalog>.<schema>` 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:<space-id>
+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

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:<catalog>.<schema>`, "
+            "`uc-functions:<catalog>.<schema>`, `external:<connection>`, `genie-space:<id>`, or "
+            "`app:<name>` (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:<catalog>.<schema>`, "
+            "`uc-functions:<catalog>.<schema>`, `external:<connection>`, `genie-space:<id>`, or "
+            "`app:<name>` (workspace access required).",
         ),
     ] = None,
 ) -> None:

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

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:<id>`, `app:<name>`) 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
         # `<catalog>.<schema>` 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:

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):