Improve AI Gateway diagnostics during configure

src/ucode/cli.py

@@ -45,7 +45,6 @@
     discover_codex_models,
     discover_gemini_models,
     discover_model_services,
-    ensure_ai_gateway,
     ensure_databricks_auth,
     ensure_pat_bearer,
     find_profile_name_for_host,
@@ -57,6 +56,7 @@
     list_profile_entries,
     list_tool_provider_services,
     normalize_workspace_url,
+    probe_unity_gateway_capabilities,
     resolve_pat_token,
     resolve_provider_launch_model,
     run_databricks_login,
@@ -516,7 +516,7 @@ def configure_shared_state(
     fable_enabled: bool | None = None,
     databricks_ai_tools_enabled: bool | None = None,
 ) -> dict:
-    """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state.
+    """Log into Databricks, verify AI Gateway, fetch model lists, persist state.
 
     If tools is provided, only fetch models for those tools. Otherwise fetch all.
     If force_login is True, always run databricks auth login (used by explicit configure).
@@ -635,8 +635,9 @@ def configure_shared_state(
             state["profile"] = profile
     with spinner("Verifying Unity AI Gateway..."):
         token = get_databricks_token(workspace, profile)
-        ensure_ai_gateway(workspace, token)
+        model_service_probe = probe_unity_gateway_capabilities(workspace, token)
     print_success("Unity AI Gateway detected")
+    print_kv("Model service", model_service_probe.detail)
 
     want_claude = (
         fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools

src/ucode/databricks.py

@@ -51,7 +51,7 @@
 WINDOWS_DATABRICKS_INSTALL_URL = (
     "https://raw.githubusercontent.com/databricks/setup-cli/main/install.ps1"
 )
-AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta"
+AI_GATEWAY_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta"
 ANTHROPIC_MODELS_PATH = "/ai-gateway/anthropic/v1/models"
 # v1.0.0 is the release that ships `databricks aitools`.
 MIN_DATABRICKS_CLI_VERSION = (1, 0, 0)
@@ -3046,18 +3046,58 @@ def fetch_codex_models(workspace: str, token: str) -> list[str]:
     return models
 
 
-def _probe_ai_gateway_v2(workspace: str, token: str) -> tuple[bool, str | None]:
+class GatewayProbe(NamedTuple):
+    reachable: bool
+    detail: str
+    resource_available: bool = False
+
+
+def _version_neutral_gateway_detail(detail: str) -> str:
+    detail = re.sub(r"\bv3\b", "model service", detail, flags=re.IGNORECASE)
+    return re.sub(r"\bv2\b", "legacy endpoint", detail, flags=re.IGNORECASE)
+
+
+def _gateway_probe_result(
+    payload: dict | list | None,
+    reason: str | None,
+    collection_key: str,
+    resource_name: str,
+    empty_hint: str | None = None,
+) -> GatewayProbe:
+    if payload is None:
+        return GatewayProbe(False, _version_neutral_gateway_detail(reason or "unknown error"))
+    resources = payload.get(collection_key) if isinstance(payload, dict) else None
+    if resources:
+        return GatewayProbe(True, f"reachable, accessible {resource_name} returned", True)
+    detail = f"reachable, no accessible {resource_name}s returned"
+    if empty_hint:
+        detail = f"{detail}; {empty_hint}"
+    return GatewayProbe(True, detail)
+
+
+def _probe_ai_gateway_v2(workspace: str, token: str) -> GatewayProbe:
     hostname = workspace_hostname(workspace)
     url = f"https://{hostname}/api/ai-gateway/v2/endpoints?page_size=1"
     payload, reason = _http_get_json(url, token)
-    return payload is not None, reason
+    return _gateway_probe_result(
+        payload=payload,
+        reason=reason,
+        collection_key="endpoints",
+        resource_name="endpoint",
+    )
 
 
-def _probe_ai_gateway_v3(workspace: str, token: str) -> tuple[bool, str | None]:
+def _probe_ai_gateway_v3(workspace: str, token: str) -> GatewayProbe:
     hostname = workspace_hostname(workspace)
     url = f"https://{hostname}/api/2.1/unity-catalog/model-services?page_size=1"
     payload, reason = _http_get_json(url, token)
-    return payload is not None, reason
+    return _gateway_probe_result(
+        payload=payload,
+        reason=reason,
+        collection_key="model_services",
+        resource_name="model service",
+        empty_hint="check USE CATALOG on system, and USE SCHEMA and EXECUTE on system.ai",
+    )
 
 
 def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn:
@@ -3069,57 +3109,65 @@ def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn:
     )
 
 
-def _raise_ai_gateway_v3_permission_failure(
-    workspace: str, v3_reason: str, v2_reason: str | None
+def _raise_model_service_permission_failure(
+    workspace: str, model_service_reason: str, legacy_endpoint_reason: str
 ) -> NoReturn:
     raise RuntimeError(
-        f"Databricks AI Gateway V3 access could not be verified on {workspace} ({v3_reason}). "
-        f"The V2 fallback also failed ({v2_reason or 'unknown error'}). The V3 probe requires "
-        "permission to list Unity Catalog model services. Verify USE CATALOG on `system` and "
-        "USE SCHEMA on `system.ai`."
+        "Databricks Unity AI Gateway model service access could not be verified on "
+        f"{workspace} ({model_service_reason}). The legacy endpoint fallback also failed "
+        f"({legacy_endpoint_reason}). The model service probe requires permission to list "
+        "Unity Catalog model services. Verify USE CATALOG on `system`, and USE SCHEMA and "
+        "EXECUTE on `system.ai`."
     )
 
 
-def _raise_ai_gateway_v2_permission_failure(
-    workspace: str, v2_reason: str, v3_reason: str | None
+def _raise_legacy_endpoint_permission_failure(
+    workspace: str, legacy_endpoint_reason: str, model_service_reason: str
 ) -> NoReturn:
     raise RuntimeError(
-        f"Databricks AI Gateway V2 access could not be verified on {workspace} ({v2_reason}). "
-        f"The V3 probe also failed ({v3_reason or 'unknown error'}). Verify the caller's "
-        "workspace permissions for the AI Gateway V2 endpoints listing."
+        "Databricks Unity AI Gateway legacy endpoint access could not be verified on "
+        f"{workspace} ({legacy_endpoint_reason}). The model service probe also failed "
+        f"({model_service_reason}). Verify the caller's workspace permissions for the legacy "
+        "endpoints listing."
     )
 
 
-def ensure_ai_gateway(workspace: str, token: str) -> None:
-    """Pass if either AI Gateway V2 or V3 is available."""
-    v3_ok, v3_reason = _probe_ai_gateway_v3(workspace, token)
-    if v3_ok:
-        return
-    if v3_reason and _looks_like_definitive_auth_failure(v3_reason):
-        _raise_ai_gateway_auth_failure(workspace, v3_reason)
-
-    v2_ok, v2_reason = _probe_ai_gateway_v2(workspace, token)
-    if v2_ok:
-        return
-    if v2_reason and _looks_like_definitive_auth_failure(v2_reason):
-        _raise_ai_gateway_auth_failure(workspace, v2_reason)
-    if v3_reason and _looks_like_permission_failure(v3_reason):
-        _raise_ai_gateway_v3_permission_failure(workspace, v3_reason, v2_reason)
-    if v2_reason and _looks_like_permission_failure(v2_reason):
-        _raise_ai_gateway_v2_permission_failure(workspace, v2_reason, v3_reason)
+def probe_unity_gateway_capabilities(workspace: str, token: str) -> GatewayProbe:
+    """Return the model service probe after verifying an available gateway path."""
+    model_service_probe = _probe_ai_gateway_v3(workspace, token)
+    if not model_service_probe.reachable and _looks_like_definitive_auth_failure(
+        model_service_probe.detail
+    ):
+        _raise_ai_gateway_auth_failure(workspace, model_service_probe.detail)
+    if model_service_probe.resource_available:
+        return model_service_probe
+
+    legacy_endpoint_probe = _probe_ai_gateway_v2(workspace, token)
+    if legacy_endpoint_probe.reachable:
+        return model_service_probe
+    if _looks_like_definitive_auth_failure(legacy_endpoint_probe.detail):
+        _raise_ai_gateway_auth_failure(workspace, legacy_endpoint_probe.detail)
+    if _looks_like_permission_failure(model_service_probe.detail):
+        _raise_model_service_permission_failure(
+            workspace, model_service_probe.detail, legacy_endpoint_probe.detail
+        )
+    if _looks_like_permission_failure(legacy_endpoint_probe.detail):
+        _raise_legacy_endpoint_permission_failure(
+            workspace, legacy_endpoint_probe.detail, model_service_probe.detail
+        )
 
     raise RuntimeError(
-        "Databricks AI Gateway is not enabled on this workspace: neither V3 "
-        f"({v3_reason or 'unknown error'}) nor V2 ({v2_reason or 'unknown error'}) is available. "
-        f"See {AI_GATEWAY_V2_DOCS_URL}"
+        "Databricks Unity AI Gateway is not enabled on this workspace: neither model services "
+        f"({model_service_probe.detail}) nor legacy endpoints ({legacy_endpoint_probe.detail}) "
+        f"are available. See {AI_GATEWAY_DOCS_URL}"
     )
 
 
 def _looks_like_definitive_auth_failure(reason: str) -> bool:
     """True when retrying another workspace API cannot rescue this token.
 
-    A 403 can be endpoint-specific authorization, so the version-agnostic
-    preflight must still try V3 before surfacing it as an auth failure.
+    A 403 can be endpoint-specific authorization, so the preflight must still
+    try the fallback before surfacing it as an auth failure.
     """
     if "HTTP 401" in reason:
         return True

tests/test_cli.py

@@ -13,7 +13,9 @@
 import pytest
 from typer.testing import CliRunner
 
+import ucode.databricks as db_mod
 from ucode.cli import app
+from ucode.databricks import GatewayProbe
 
 _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")
 
@@ -29,6 +31,8 @@ def _strip_ansi(text: str) -> str:
 
 TOOLS = ["codex", "claude", "gemini", "opencode"]
 
+MODEL_SERVICE_PROBE = GatewayProbe(True, "reachable, accessible model service returned", True)
+
 
 def _jwt(expires_at: float) -> str:
     payload = base64.urlsafe_b64encode(json.dumps({"exp": expires_at}).encode()).decode()
@@ -2466,7 +2470,9 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None):
         monkeypatch.setattr(cli_mod, "resolve_pat_token", lambda p: pat_token)
         monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None)
         monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token")
-        monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: None)
+        monkeypatch.setattr(
+            cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE
+        )
         monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], None))
         monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None))
         monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None))
@@ -2489,6 +2495,112 @@ def test_use_pat_exports_bearer_and_skips_login(self, monkeypatch):
         assert state["use_pat"] is True
         assert saved and saved[-1]["use_pat"] is True
 
+    def test_prints_model_service_and_omits_unneeded_legacy_probe(self, monkeypatch, capsys):
+        cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat")
+
+        cli_mod.configure_shared_state(self.WS, profile="DEFAULT")
+
+        output = _strip_ansi(capsys.readouterr().out)
+        assert "Model service: reachable, accessible model service returned" in output
+        assert "(Legacy) endpoints:" not in output
+
+    @pytest.mark.parametrize(
+        ("responses", "expected_model_service"),
+        [
+            (
+                [({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None)],
+                "reachable, accessible model service returned",
+            ),
+            (
+                [({}, None), ({"endpoints": []}, None)],
+                "reachable, no accessible model services returned; check USE CATALOG on system, "
+                "and USE SCHEMA and EXECUTE on system.ai",
+            ),
+            (
+                [
+                    (None, "HTTP 403 Forbidden"),
+                    ({"endpoints": [{"name": "databricks-gpt-5"}]}, None),
+                ],
+                "HTTP 403 Forbidden",
+            ),
+        ],
+        ids=[
+            "model-service-resource",
+            "model-service-empty-legacy-empty",
+            "model-service-forbidden-legacy-resource",
+        ],
+    )
+    def test_prints_local_gateway_probe_scenarios(
+        self,
+        monkeypatch,
+        capsys,
+        responses,
+        expected_model_service,
+    ):
+        cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat")
+        response_iter = iter(responses)
+        monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: next(response_iter))
+        monkeypatch.setattr(
+            cli_mod, "probe_unity_gateway_capabilities", db_mod.probe_unity_gateway_capabilities
+        )
+
+        cli_mod.configure_shared_state(self.WS, profile="DEFAULT")
+
+        output = " ".join(_strip_ansi(capsys.readouterr().out).split())
+        assert "Unity AI Gateway detected" in output
+        assert f"Model service: {expected_model_service}" in output
+        assert "(Legacy) endpoints:" not in output
+        assert "V2" not in output
+        assert "V3" not in output
+
+    @pytest.mark.parametrize(
+        ("responses", "error_match"),
+        [
+            (
+                [
+                    ({}, None),
+                    (
+                        None,
+                        "HTTP 404 Not Found: AI Gateway V2 is not available for CSP-enabled "
+                        "workspaces",
+                    ),
+                ],
+                "no accessible model services",
+            ),
+            (
+                [
+                    (None, "HTTP 404 Not Found: V3 unavailable"),
+                    (None, "HTTP 404 Not Found: V2 unavailable"),
+                ],
+                "neither model services",
+            ),
+            ([(None, "HTTP 401 Unauthorized")], "rejected the access token"),
+        ],
+        ids=[
+            "model-service-empty-legacy-unavailable",
+            "neither-path-reachable",
+            "invalid-token",
+        ],
+    )
+    def test_local_gateway_probe_failures_do_not_print_success(
+        self, monkeypatch, capsys, responses, error_match
+    ):
+        cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat")
+        response_iter = iter(responses)
+        monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: next(response_iter))
+        monkeypatch.setattr(
+            cli_mod, "probe_unity_gateway_capabilities", db_mod.probe_unity_gateway_capabilities
+        )
+
+        with pytest.raises(RuntimeError, match=error_match) as excinfo:
+            cli_mod.configure_shared_state(self.WS, profile="DEFAULT")
+
+        output = _strip_ansi(capsys.readouterr().out)
+        assert "Unity AI Gateway detected" not in output
+        message = str(excinfo.value)
+        assert "v2" not in message.lower()
+        assert "v3" not in message.lower()
+
     def test_use_pat_without_pat_profile_raises(self, monkeypatch):
         cli_mod, logins, _, _ = self._stub_deps(monkeypatch, pat_token=None)
 
@@ -2787,7 +2899,9 @@ def _stub_external_deps(monkeypatch):
         monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None: None)
         monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None)
         monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token")
-        monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: None)
+        monkeypatch.setattr(
+            cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE
+        )
         monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], None))
         monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None))
         monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None))
@@ -2847,7 +2961,9 @@ def _stub(monkeypatch):
         monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: None)
         monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None)
         monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token")
-        monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: None)
+        monkeypatch.setattr(
+            cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE
+        )
         monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {})
         monkeypatch.setattr(cli_mod, "save_state", lambda s: None)
 
@@ -2905,7 +3021,11 @@ def _f(*a, **k):
         monkeypatch.setattr(cli_mod, "run_databricks_login", _boom("run_databricks_login"))
         monkeypatch.setattr(cli_mod, "ensure_pat_bearer", _boom("ensure_pat_bearer"))
         monkeypatch.setattr(cli_mod, "get_databricks_token", _boom("get_databricks_token"))
-        monkeypatch.setattr(cli_mod, "ensure_ai_gateway", _boom("ensure_ai_gateway"))
+        monkeypatch.setattr(
+            cli_mod,
+            "probe_unity_gateway_capabilities",
+            _boom("probe_unity_gateway_capabilities"),
+        )
         monkeypatch.setattr(cli_mod, "discover_model_services", _boom("discover_model_services"))
         monkeypatch.setattr(cli_mod, "discover_codex_models", _boom("discover_codex_models"))
         monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: "resolved")