diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 5db605d..2d2f49c 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -42,7 +42,7 @@ discover_codex_models, discover_gemini_models, discover_model_services, - ensure_ai_gateway_v2, + ensure_ai_gateway, ensure_databricks_auth, ensure_pat_bearer, find_profile_name_for_host, @@ -545,7 +545,7 @@ def configure_shared_state( state["profile"] = profile with spinner("Verifying Unity AI Gateway..."): token = get_databricks_token(workspace, profile) - ensure_ai_gateway_v2(workspace, token) + ensure_ai_gateway(workspace, token) print_success("Unity AI Gateway detected") want_claude = ( diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 377c75c..dbca302 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1,5 +1,5 @@ """Databricks workspace integration: CLI auth, token retrieval, model -discovery, AI Gateway v2 enforcement, SQL warehouse discovery, URL builders.""" +discovery, AI Gateway checks, SQL warehouse discovery, URL builders.""" from __future__ import annotations @@ -26,7 +26,7 @@ ) from decimal import Decimal, InvalidOperation from pathlib import Path -from typing import Literal, NamedTuple, cast, overload +from typing import Literal, NamedTuple, NoReturn, cast, overload from urllib import error as urllib_error from urllib import request as urllib_request from urllib.parse import quote, urlencode, urlparse @@ -2881,54 +2881,88 @@ def fetch_codex_models(workspace: str, token: str) -> list[str]: return models -def ensure_ai_gateway_v2(workspace: str, token: str) -> None: - """Probe AI Gateway v2 and raise if unavailable. +def _probe_ai_gateway_v2(workspace: str, token: str) -> tuple[bool, str | None]: + 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 - Uses the dedicated v2 listing endpoint `GET /api/ai-gateway/v2/endpoints`: - a 200 response (even with an empty list) means v2 is wired up on this - workspace — a "no endpoints provisioned" case will surface naturally in - downstream discovery. Failure branches: - - 401 / 403 / 400 with `Invalid Token`: the token is bad for *this* - workspace. - - 404: AI Gateway V2 is not enabled on this workspace — point at the docs. - - other (5xx, network errors): surface the reason verbatim. - """ +def _probe_ai_gateway_v3(workspace: str, token: str) -> tuple[bool, str | None]: hostname = workspace_hostname(workspace) - url = f"https://{hostname}/api/ai-gateway/v2/endpoints?page_size=1" + url = f"https://{hostname}/api/2.1/unity-catalog/model-services?page_size=1" payload, reason = _http_get_json(url, token) - if payload is not None: + return payload is not None, reason + + +def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: + raise RuntimeError( + f"Databricks rejected the access token for {workspace} ({reason}). " + f"Try:\n" + f" databricks auth logout --host {workspace}\n" + f" databricks auth login --host {workspace}" + ) + + +def _raise_ai_gateway_v3_permission_failure( + workspace: str, v3_reason: str, v2_reason: str | None +) -> 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`." + ) + + +def _raise_ai_gateway_v2_permission_failure( + workspace: str, v2_reason: str, v3_reason: str | None +) -> 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." + ) + + +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 - reason_str = reason or "unknown error" - if _looks_like_auth_failure(reason_str): - raise RuntimeError( - f"Databricks rejected the access token for {workspace} ({reason_str}). " - f"Try:\n" - f" databricks auth logout --host {workspace}\n" - f" databricks auth login --host {workspace}" - ) - if "HTTP 404" in reason_str: - raise RuntimeError( - "Databricks Unity AI Gateway is not enabled on this workspace " - f"({reason_str}). See {AI_GATEWAY_V2_DOCS_URL}" - ) + 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) + raise RuntimeError( - "Databricks Unity AI Gateway probe failed on this workspace " - f"({reason_str}). See {AI_GATEWAY_V2_DOCS_URL}" + "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}" ) -def _looks_like_auth_failure(reason: str) -> bool: - """True when the gateway response signals the token is not accepted. +def _looks_like_definitive_auth_failure(reason: str) -> bool: + """True when retrying another workspace API cannot rescue this token. - Covers 401/403 directly and the gateway's 400 + `Invalid Token` body - (which happens when the bearer is valid but issued for a different - workspace).""" - if "HTTP 401" in reason or "HTTP 403" in reason: - return True - if "HTTP 400" in reason and "invalid token" in reason.lower(): + A 403 can be endpoint-specific authorization, so the version-agnostic + preflight must still try V3 before surfacing it as an auth failure. + """ + if "HTTP 401" in reason: return True - return False + return "HTTP 400" in reason and "invalid token" in reason.lower() + + +def _looks_like_permission_failure(reason: str) -> bool: + return "HTTP 403" in reason CODING_AGENT_RECOMMEND_MODEL_PATH = "/api/ai-gateway/v2/coding-agent-configs:recommendModel" diff --git a/tests/test_cli.py b/tests/test_cli.py index cc75427..2069cd2 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1918,7 +1918,7 @@ 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_v2", lambda w, t: None) + monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: None) 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)) @@ -2216,7 +2216,7 @@ 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_v2", lambda w, t: None) + monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: None) 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)) @@ -2276,7 +2276,7 @@ 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_v2", lambda w, t: None) + monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: None) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) monkeypatch.setattr(cli_mod, "save_state", lambda s: None) @@ -2334,7 +2334,7 @@ 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_v2", _boom("ensure_ai_gateway_v2")) + monkeypatch.setattr(cli_mod, "ensure_ai_gateway", _boom("ensure_ai_gateway")) 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") diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 7ef8b2d..339735d 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -12,7 +12,6 @@ import ucode.databricks as db_mod from ucode.databricks import ( - AI_GATEWAY_V2_DOCS_URL, CODING_AGENT_RECOMMEND_MODEL_PATH, _format_subprocess_result, _parse_databricks_cli_version, @@ -43,6 +42,7 @@ ) WS = "https://example.databricks.com" +WS_HOST = "example.databricks.com" class _FakeResponse: @@ -1831,128 +1831,120 @@ def fake_run(args, **kwargs): list_databricks_apps(WS) -class TestEnsureAiGatewayV2: - """Test ensure_ai_gateway_v2 without real network calls. +class TestEnsureAiGateway: + def test_v3_only_workspace_succeeds_without_v2_probe(self, monkeypatch): + calls: list[str] = [] - The probe is `GET /api/ai-gateway/v2/endpoints`: a successful JSON - response means v2 is wired up (even if `endpoints` is empty), while - 404/401/403/network errors all raise a RuntimeError with the docs URL. - """ + def fake_get(url, token): + calls.append(url) + return {"model_services": []}, None - @staticmethod - def _mock_json_response(body: str): - from unittest.mock import MagicMock + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - mock_resp = MagicMock() - mock_resp.__enter__ = lambda s: s - mock_resp.__exit__ = MagicMock(return_value=False) - mock_resp.read.return_value = body.encode("utf-8") - return mock_resp + db_mod.ensure_ai_gateway(WS, "fake-token") - @staticmethod - def _http_error(code: int, msg: str, body: str = ""): - import io - from unittest.mock import MagicMock - from urllib.error import HTTPError + assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1"] - fp = io.BytesIO(body.encode("utf-8")) if body else None - return HTTPError(url="", code=code, msg=msg, hdrs=MagicMock(), fp=fp) + def test_v2_only_workspace_succeeds_after_v3_probe(self, monkeypatch): + calls: list[str] = [] - def test_raises_on_404(self): - from unittest.mock import patch + def fake_get(url, token): + calls.append(url) + if "/api/2.1/unity-catalog/model-services" in url: + return None, "HTTP 404: Not Found" + return {"endpoints": []}, None - exc = self._http_error(404, "Not Found") - with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): - from ucode.databricks import ensure_ai_gateway_v2 + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL) as excinfo: - ensure_ai_gateway_v2(WS, "fake-token") - assert "not enabled" in str(excinfo.value) + db_mod.ensure_ai_gateway(WS, "fake-token") - def test_raises_on_401_with_auth_hint(self): - from unittest.mock import patch + assert calls == [ + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1", + f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", + ] - exc = self._http_error(401, "Unauthorized") - with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): - from ucode.databricks import ensure_ai_gateway_v2 + def test_v3_forbidden_still_succeeds_when_v2_is_available(self, monkeypatch): + calls: list[str] = [] - with pytest.raises(RuntimeError, match="401") as excinfo: - ensure_ai_gateway_v2(WS, "fake-token") - message = str(excinfo.value) - assert "rejected" in message.lower() - assert "databricks auth login" in message + def fake_get(url, token): + calls.append(url) + if "/api/2.1/unity-catalog/model-services" in url: + return None, "HTTP 403: Forbidden" + return {"endpoints": []}, None - def test_raises_on_400_invalid_token_with_auth_hint(self): - """400 + body `Invalid Token` is the misleading-error case from issue #84.""" - from unittest.mock import patch + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - exc = self._http_error(400, "Bad Request", body="Invalid Token") - with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): - from ucode.databricks import ensure_ai_gateway_v2 - - with pytest.raises(RuntimeError) as excinfo: - ensure_ai_gateway_v2(WS, "fake-token") - message = str(excinfo.value) - # The bug we are fixing: must NOT collapse to the generic - # "v2 not available" message — must call out the auth failure - # and point at re-login. - assert "Invalid Token" in message - assert "rejected" in message.lower() - assert "databricks auth login" in message - - def test_400_without_invalid_token_falls_through_to_generic(self): - """A 400 that is *not* an auth failure should still surface the body.""" - from unittest.mock import patch + db_mod.ensure_ai_gateway(WS, "fake-token") - exc = self._http_error(400, "Bad Request", body="some other detail") - with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): - from ucode.databricks import ensure_ai_gateway_v2 + assert calls == [ + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1", + f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", + ] - with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL) as excinfo: - ensure_ai_gateway_v2(WS, "fake-token") - assert "some other detail" in str(excinfo.value) + def test_neither_gateway_available_raises(self, monkeypatch): + reasons = iter(["HTTP 404: V3 missing", "HTTP 404: V2 missing"]) + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: (None, next(reasons)), + ) - def test_raises_on_url_error(self): - from unittest.mock import patch - from urllib.error import URLError + with pytest.raises(RuntimeError, match="neither V3") as excinfo: + db_mod.ensure_ai_gateway(WS, "fake-token") - with patch( - "ucode.databricks.urllib_request.urlopen", - side_effect=URLError("connection refused"), - ): - from ucode.databricks import ensure_ai_gateway_v2 + message = str(excinfo.value) + assert "HTTP 404: V2 missing" in message + assert "HTTP 404: V3 missing" in message - with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL): - ensure_ai_gateway_v2(WS, "fake-token") + def test_v3_auth_failure_does_not_probe_v2(self, monkeypatch): + calls: list[str] = [] - def test_succeeds_with_endpoints_list(self): - from unittest.mock import patch + def fake_get(url, token): + calls.append(url) + return None, "HTTP 401: Unauthorized" - with patch( - "ucode.databricks.urllib_request.urlopen", - return_value=self._mock_json_response('{"endpoints": [{"name": "foo"}]}'), - ): - from ucode.databricks import ensure_ai_gateway_v2 + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - ensure_ai_gateway_v2(WS, "fake-token") # should not raise + with pytest.raises(RuntimeError, match="rejected"): + db_mod.ensure_ai_gateway(WS, "fake-token") - def test_succeeds_with_empty_endpoints_list(self): - from unittest.mock import patch + assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1"] + + def test_v3_forbidden_and_v2_unavailable_reports_permission_error(self, monkeypatch): + reasons = iter(["HTTP 403: Missing Unity Catalog grants", "HTTP 404: V2 missing"]) + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: (None, next(reasons)), + ) + + with pytest.raises(RuntimeError, match="permission") as excinfo: + db_mod.ensure_ai_gateway(WS, "fake-token") + + message = str(excinfo.value) + assert "USE SCHEMA" in message + assert "rejected the access token" not in message + assert "not enabled" not in message + + def test_v2_forbidden_and_v3_unavailable_reports_permission_error(self, monkeypatch): + reasons = iter(["HTTP 404: V3 missing", "HTTP 403: V2 forbidden"]) + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: (None, next(reasons)), + ) - # A 200 with no endpoints still means v2 is wired up on this workspace — - # downstream discovery will surface "no models" with a clearer reason. - with patch( - "ucode.databricks.urllib_request.urlopen", - return_value=self._mock_json_response('{"endpoints": []}'), - ): - from ucode.databricks import ensure_ai_gateway_v2 + with pytest.raises(RuntimeError, match="workspace permissions") as excinfo: + db_mod.ensure_ai_gateway(WS, "fake-token") - ensure_ai_gateway_v2(WS, "fake-token") # should not raise + message = str(excinfo.value) + assert "V2 access could not be verified" in message + assert "USE SCHEMA" not in message class TestHttpGetJsonReason: """The `reason` string returned by `_http_get_json` must include the response body - so callers (e.g. ensure_ai_gateway_v2) can route on it. Before issue #84's fix + so callers (e.g. ensure_ai_gateway) can route on it. Before issue #84's fix the body was logged only when UCODE_DEBUG=1 and dropped from the bubbled error.""" @staticmethod diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 76b5087..e7183c3 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -26,7 +26,7 @@ build_shared_base_urls, build_tool_base_url, discover_sql_warehouses, - ensure_ai_gateway_v2, + ensure_ai_gateway, fetch_ai_gateway_claude_models, fetch_codex_models, fetch_gemini_models, @@ -131,13 +131,13 @@ def test_get_token_returns_non_empty_string(self, e2e_token): # --------------------------------------------------------------------------- -# AI Gateway v2 probe +# AI Gateway probe # --------------------------------------------------------------------------- -class TestAiGatewayV2: - def test_ensure_ai_gateway_v2_does_not_raise(self, e2e_workspace, e2e_token): - ensure_ai_gateway_v2(e2e_workspace, e2e_token) +class TestAiGateway: + def test_ensure_ai_gateway_does_not_raise(self, e2e_workspace, e2e_token): + ensure_ai_gateway(e2e_workspace, e2e_token) def test_workspace_hostname_resolves(self, e2e_workspace): hostname = workspace_hostname(e2e_workspace) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index d5286ce..f243a63 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -799,13 +799,15 @@ def fake_configure_client_mcp_server(client, name, url, *a, **kw): assert mcp.configure_mcp_command() == 0 - assert configured == [ - ("claude", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), - ("codex", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), - ("gemini", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), - ("opencode", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), - ("copilot", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), - ] + assert sorted(configured) == sorted( + [ + ("claude", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), + ("codex", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), + ("gemini", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), + ("opencode", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), + ("copilot", "github-mcp", f"{WS}/api/2.0/mcp/external/github-mcp"), + ] + ) assert saved_states[-1]["mcp_servers"] == [ { "name": "github-mcp",