From cb544f11b855052626df137a3d1cfff95993a4f0 Mon Sep 17 00:00:00 2001 From: David Siqi Liu Date: Mon, 17 Aug 2026 15:12:46 -0400 Subject: [PATCH 1/6] Accept AI Gateway V3 during preflight --- src/ucode/cli.py | 4 +-- src/ucode/databricks.py | 61 ++++++++++++++++++++++++++++++-------- tests/test_cli.py | 8 ++--- tests/test_databricks.py | 63 ++++++++++++++++++++++++++++++++++++++++ tests/test_e2e.py | 10 +++---- 5 files changed, 123 insertions(+), 23 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index c2e38f4..3b56d9c 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -41,7 +41,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, @@ -543,7 +543,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 88935dd..9cf2d1d 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 @@ -25,7 +25,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 @@ -2852,6 +2852,50 @@ 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]: + 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 + + +def _probe_ai_gateway_v3(workspace: str, token: str) -> tuple[bool, str | None]: + 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 + + +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 ensure_ai_gateway(workspace: str, token: str) -> None: + """Pass if either AI Gateway V2 or V3 is available.""" + v2_ok, v2_reason = _probe_ai_gateway_v2(workspace, token) + if v2_ok: + return + if v2_reason and _looks_like_auth_failure(v2_reason): + _raise_ai_gateway_auth_failure(workspace, v2_reason) + + v3_ok, v3_reason = _probe_ai_gateway_v3(workspace, token) + if v3_ok: + return + if v3_reason and _looks_like_auth_failure(v3_reason): + _raise_ai_gateway_auth_failure(workspace, v3_reason) + + raise RuntimeError( + "Databricks AI Gateway is not enabled on this workspace: neither V2 " + f"({v2_reason or 'unknown error'}) nor V3 ({v3_reason or 'unknown error'}) is available. " + f"See {AI_GATEWAY_V2_DOCS_URL}" + ) + + def ensure_ai_gateway_v2(workspace: str, token: str) -> None: """Probe AI Gateway v2 and raise if unavailable. @@ -2865,19 +2909,12 @@ def ensure_ai_gateway_v2(workspace: str, token: str) -> None: - 404: AI Gateway V2 is not enabled on this workspace — point at the docs. - other (5xx, network errors): surface the reason verbatim. """ - hostname = workspace_hostname(workspace) - url = f"https://{hostname}/api/ai-gateway/v2/endpoints?page_size=1" - payload, reason = _http_get_json(url, token) - if payload is not None: + ok, reason = _probe_ai_gateway_v2(workspace, token) + if 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}" - ) + _raise_ai_gateway_auth_failure(workspace, reason_str) if "HTTP 404" in reason_str: raise RuntimeError( "Databricks Unity AI Gateway is not enabled on this workspace " diff --git a/tests/test_cli.py b/tests/test_cli.py index 55ab879..304d38d 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)) @@ -2207,7 +2207,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)) @@ -2267,7 +2267,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) @@ -2325,7 +2325,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 64b4dcc..50a18bc 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -43,6 +43,7 @@ ) WS = "https://example.databricks.com" +WS_HOST = "example.databricks.com" class _FakeResponse: @@ -1810,6 +1811,68 @@ def fake_run(args, **kwargs): list_databricks_apps(WS) +class TestEnsureAiGateway: + def test_v3_only_workspace_succeeds(self, monkeypatch): + calls: list[str] = [] + + def fake_get(url, token): + calls.append(url) + if "/api/ai-gateway/v2/endpoints" in url: + return None, "HTTP 404: Not Found" + return {"model_services": []}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + db_mod.ensure_ai_gateway(WS, "fake-token") + + assert calls == [ + f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1", + ] + + def test_v2_only_workspace_succeeds_without_v3_probe(self, monkeypatch): + calls: list[str] = [] + + def fake_get(url, token): + calls.append(url) + return {"endpoints": []}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + db_mod.ensure_ai_gateway(WS, "fake-token") + + assert calls == [f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1"] + + def test_neither_gateway_available_raises(self, monkeypatch): + reasons = iter(["HTTP 404: V2 missing", "HTTP 404: V3 missing"]) + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: (None, next(reasons)), + ) + + with pytest.raises(RuntimeError, match="neither V2") as excinfo: + db_mod.ensure_ai_gateway(WS, "fake-token") + + message = str(excinfo.value) + assert "HTTP 404: V2 missing" in message + assert "HTTP 404: V3 missing" in message + + def test_v2_auth_failure_does_not_probe_v3(self, monkeypatch): + calls: list[str] = [] + + def fake_get(url, token): + calls.append(url) + return None, "HTTP 401: Unauthorized" + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + with pytest.raises(RuntimeError, match="rejected"): + db_mod.ensure_ai_gateway(WS, "fake-token") + + assert calls == [f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1"] + + class TestEnsureAiGatewayV2: """Test ensure_ai_gateway_v2 without real network calls. 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) From 0d84269c3cca7154233a75de63ef37e86f6d7851 Mon Sep 17 00:00:00 2001 From: David Siqi Liu Date: Mon, 17 Aug 2026 15:33:40 -0400 Subject: [PATCH 2/6] Probe Gateway V3 after V2 forbidden --- src/ucode/databricks.py | 15 ++++++++++++++- tests/test_databricks.py | 18 ++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9cf2d1d..6a13aa2 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2880,7 +2880,7 @@ def ensure_ai_gateway(workspace: str, token: str) -> None: v2_ok, v2_reason = _probe_ai_gateway_v2(workspace, token) if v2_ok: return - if v2_reason and _looks_like_auth_failure(v2_reason): + if v2_reason and _looks_like_definitive_auth_failure(v2_reason): _raise_ai_gateway_auth_failure(workspace, v2_reason) v3_ok, v3_reason = _probe_ai_gateway_v3(workspace, token) @@ -2888,6 +2888,8 @@ def ensure_ai_gateway(workspace: str, token: str) -> None: return if v3_reason and _looks_like_auth_failure(v3_reason): _raise_ai_gateway_auth_failure(workspace, v3_reason) + if v2_reason and _looks_like_auth_failure(v2_reason): + _raise_ai_gateway_auth_failure(workspace, v2_reason) raise RuntimeError( "Databricks AI Gateway is not enabled on this workspace: neither V2 " @@ -2939,6 +2941,17 @@ def _looks_like_auth_failure(reason: str) -> bool: return False +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. + """ + if "HTTP 401" in reason: + return True + return "HTTP 400" in reason and "invalid token" in reason.lower() + + CODING_AGENT_RECOMMEND_MODEL_PATH = "/api/ai-gateway/v2/coding-agent-configs:recommendModel" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 50a18bc..970bd47 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1843,6 +1843,24 @@ def fake_get(url, token): assert calls == [f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1"] + def test_v2_forbidden_still_succeeds_when_v3_is_available(self, monkeypatch): + calls: list[str] = [] + + def fake_get(url, token): + calls.append(url) + if "/api/ai-gateway/v2/endpoints" in url: + return None, "HTTP 403: Forbidden" + return {"model_services": []}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + db_mod.ensure_ai_gateway(WS, "fake-token") + + assert calls == [ + f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1", + ] + def test_neither_gateway_available_raises(self, monkeypatch): reasons = iter(["HTTP 404: V2 missing", "HTTP 404: V3 missing"]) monkeypatch.setattr( From 795b28f3ec1b57844099c098088f32c59fafff99 Mon Sep 17 00:00:00 2001 From: David Siqi Liu Date: Mon, 17 Aug 2026 16:00:49 -0400 Subject: [PATCH 3/6] Prefer AI Gateway V3 during preflight --- src/ucode/databricks.py | 38 ++++++++++++++++++++++-------- tests/test_databricks.py | 50 ++++++++++++++++++++++++++++------------ 2 files changed, 63 insertions(+), 25 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 6a13aa2..f1a6589 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2875,25 +2875,39 @@ def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: ) +def _raise_ai_gateway_permission_failure( + workspace: str, v3_reason: str | None, v2_reason: str | None +) -> NoReturn: + raise RuntimeError( + f"Databricks AI Gateway access could not be verified on {workspace} because a probe " + f"was forbidden: V3 ({v3_reason or 'unknown error'}); " + f"V2 ({v2_reason or 'unknown error'}). The V3 probe requires permission to list " + "Unity Catalog model services. Verify USE CATALOG on `system`, USE SCHEMA on " + "`system.ai`, and the caller's workspace permissions." + ) + + def ensure_ai_gateway(workspace: str, token: str) -> None: """Pass if either AI Gateway V2 or V3 is available.""" - 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) - v3_ok, v3_reason = _probe_ai_gateway_v3(workspace, token) if v3_ok: return - if v3_reason and _looks_like_auth_failure(v3_reason): + if v3_reason and _looks_like_definitive_auth_failure(v3_reason): _raise_ai_gateway_auth_failure(workspace, v3_reason) - if v2_reason and _looks_like_auth_failure(v2_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 any( + reason and _looks_like_permission_failure(reason) for reason in (v3_reason, v2_reason) + ): + _raise_ai_gateway_permission_failure(workspace, v3_reason, v2_reason) raise RuntimeError( - "Databricks AI Gateway is not enabled on this workspace: neither V2 " - f"({v2_reason or 'unknown error'}) nor V3 ({v3_reason or 'unknown error'}) is available. " + "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}" ) @@ -2952,6 +2966,10 @@ def _looks_like_definitive_auth_failure(reason: str) -> bool: 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_databricks.py b/tests/test_databricks.py index 970bd47..e6825fd 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1812,13 +1812,11 @@ def fake_run(args, **kwargs): class TestEnsureAiGateway: - def test_v3_only_workspace_succeeds(self, monkeypatch): + def test_v3_only_workspace_succeeds_without_v2_probe(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): calls.append(url) - if "/api/ai-gateway/v2/endpoints" in url: - return None, "HTTP 404: Not Found" return {"model_services": []}, None monkeypatch.setattr(db_mod, "_http_get_json", fake_get) @@ -1826,57 +1824,61 @@ def fake_get(url, token): db_mod.ensure_ai_gateway(WS, "fake-token") assert calls == [ - f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", - f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1", + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1" ] - def test_v2_only_workspace_succeeds_without_v3_probe(self, monkeypatch): + def test_v2_only_workspace_succeeds_after_v3_probe(self, monkeypatch): calls: list[str] = [] 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 monkeypatch.setattr(db_mod, "_http_get_json", fake_get) db_mod.ensure_ai_gateway(WS, "fake-token") - assert calls == [f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1"] + 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", + ] - def test_v2_forbidden_still_succeeds_when_v3_is_available(self, monkeypatch): + def test_v3_forbidden_still_succeeds_when_v2_is_available(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): calls.append(url) - if "/api/ai-gateway/v2/endpoints" in url: + if "/api/2.1/unity-catalog/model-services" in url: return None, "HTTP 403: Forbidden" - return {"model_services": []}, None + return {"endpoints": []}, None monkeypatch.setattr(db_mod, "_http_get_json", fake_get) db_mod.ensure_ai_gateway(WS, "fake-token") assert calls == [ - f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", 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", ] def test_neither_gateway_available_raises(self, monkeypatch): - reasons = iter(["HTTP 404: V2 missing", "HTTP 404: V3 missing"]) + reasons = iter(["HTTP 404: V3 missing", "HTTP 404: V2 missing"]) monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token: (None, next(reasons)), ) - with pytest.raises(RuntimeError, match="neither V2") as excinfo: + with pytest.raises(RuntimeError, match="neither V3") as excinfo: db_mod.ensure_ai_gateway(WS, "fake-token") message = str(excinfo.value) assert "HTTP 404: V2 missing" in message assert "HTTP 404: V3 missing" in message - def test_v2_auth_failure_does_not_probe_v3(self, monkeypatch): + def test_v3_auth_failure_does_not_probe_v2(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): @@ -1888,7 +1890,25 @@ def fake_get(url, token): with pytest.raises(RuntimeError, match="rejected"): db_mod.ensure_ai_gateway(WS, "fake-token") - assert calls == [f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1"] + 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 class TestEnsureAiGatewayV2: From ec707ad8185e05d15b974c3beb8a188d2125a448 Mon Sep 17 00:00:00 2001 From: David Siqi Liu Date: Mon, 17 Aug 2026 16:05:37 -0400 Subject: [PATCH 4/6] Clarify gateway permission failures --- src/ucode/databricks.py | 31 ++++++++++++++++++++----------- tests/test_databricks.py | 15 +++++++++++++++ 2 files changed, 35 insertions(+), 11 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f1a6589..a16671e 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2875,15 +2875,24 @@ def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: ) -def _raise_ai_gateway_permission_failure( - workspace: str, v3_reason: str | None, v2_reason: str | None +def _raise_ai_gateway_v3_permission_failure( + workspace: str, v3_reason: str, v2_reason: str | None ) -> NoReturn: raise RuntimeError( - f"Databricks AI Gateway access could not be verified on {workspace} because a probe " - f"was forbidden: V3 ({v3_reason or 'unknown error'}); " - f"V2 ({v2_reason or 'unknown error'}). The V3 probe requires permission to list " - "Unity Catalog model services. Verify USE CATALOG on `system`, USE SCHEMA on " - "`system.ai`, and the caller's workspace permissions." + 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." ) @@ -2900,10 +2909,10 @@ def ensure_ai_gateway(workspace: str, token: str) -> None: return if v2_reason and _looks_like_definitive_auth_failure(v2_reason): _raise_ai_gateway_auth_failure(workspace, v2_reason) - if any( - reason and _looks_like_permission_failure(reason) for reason in (v3_reason, v2_reason) - ): - _raise_ai_gateway_permission_failure(workspace, v3_reason, 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 AI Gateway is not enabled on this workspace: neither V3 " diff --git a/tests/test_databricks.py b/tests/test_databricks.py index e6825fd..4b1a857 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1910,6 +1910,21 @@ def test_v3_forbidden_and_v2_unavailable_reports_permission_error(self, monkeypa 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)), + ) + + with pytest.raises(RuntimeError, match="workspace permissions") as excinfo: + db_mod.ensure_ai_gateway(WS, "fake-token") + + message = str(excinfo.value) + assert "V2 access could not be verified" in message + assert "USE SCHEMA" not in message + class TestEnsureAiGatewayV2: """Test ensure_ai_gateway_v2 without real network calls. From 1c534fa3390ff7237b666af06860120e8deae0f6 Mon Sep 17 00:00:00 2001 From: David Siqi Liu Date: Mon, 17 Aug 2026 16:14:26 -0400 Subject: [PATCH 5/6] Remove unused V2-only gateway check --- src/ucode/databricks.py | 43 -------------- tests/test_databricks.py | 122 +-------------------------------------- 2 files changed, 1 insertion(+), 164 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a16671e..c4cc515 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2921,49 +2921,6 @@ def ensure_ai_gateway(workspace: str, token: str) -> None: ) -def ensure_ai_gateway_v2(workspace: str, token: str) -> None: - """Probe AI Gateway v2 and raise if unavailable. - - 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. - """ - ok, reason = _probe_ai_gateway_v2(workspace, token) - if ok: - return - reason_str = reason or "unknown error" - if _looks_like_auth_failure(reason_str): - _raise_ai_gateway_auth_failure(workspace, reason_str) - 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}" - ) - raise RuntimeError( - "Databricks Unity AI Gateway probe failed on this workspace " - f"({reason_str}). 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. - - 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(): - return True - return False - - def _looks_like_definitive_auth_failure(reason: str) -> bool: """True when retrying another workspace API cannot rescue this token. diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 4b1a857..ce61037 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, @@ -1926,128 +1925,9 @@ def test_v2_forbidden_and_v3_unavailable_reports_permission_error(self, monkeypa assert "USE SCHEMA" not in message -class TestEnsureAiGatewayV2: - """Test ensure_ai_gateway_v2 without real network calls. - - 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. - """ - - @staticmethod - def _mock_json_response(body: str): - from unittest.mock import MagicMock - - 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 - - @staticmethod - def _http_error(code: int, msg: str, body: str = ""): - import io - from unittest.mock import MagicMock - from urllib.error import HTTPError - - fp = io.BytesIO(body.encode("utf-8")) if body else None - return HTTPError(url="", code=code, msg=msg, hdrs=MagicMock(), fp=fp) - - def test_raises_on_404(self): - from unittest.mock import patch - - 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 - - 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) - - def test_raises_on_401_with_auth_hint(self): - from unittest.mock import patch - - exc = self._http_error(401, "Unauthorized") - with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): - from ucode.databricks import ensure_ai_gateway_v2 - - 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 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 - - 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 - - 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 - - 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_raises_on_url_error(self): - from unittest.mock import patch - from urllib.error import URLError - - with patch( - "ucode.databricks.urllib_request.urlopen", - side_effect=URLError("connection refused"), - ): - from ucode.databricks import ensure_ai_gateway_v2 - - with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL): - ensure_ai_gateway_v2(WS, "fake-token") - - def test_succeeds_with_endpoints_list(self): - from unittest.mock import patch - - with patch( - "ucode.databricks.urllib_request.urlopen", - return_value=self._mock_json_response('{"endpoints": [{"name": "foo"}]}'), - ): - from ucode.databricks import ensure_ai_gateway_v2 - - ensure_ai_gateway_v2(WS, "fake-token") # should not raise - - def test_succeeds_with_empty_endpoints_list(self): - from unittest.mock import patch - - # 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 - - ensure_ai_gateway_v2(WS, "fake-token") # should not raise - - 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 From 07f571b8d4f25fc00ce26e830d7dd40a826baac3 Mon Sep 17 00:00:00 2001 From: David Siqi Liu Date: Tue, 18 Aug 2026 13:08:28 -0400 Subject: [PATCH 6/6] Fix CI formatting and MCP test ordering --- tests/test_databricks.py | 8 ++------ tests/test_mcp.py | 16 +++++++++------- 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 583e755..339735d 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1843,9 +1843,7 @@ def fake_get(url, token): db_mod.ensure_ai_gateway(WS, "fake-token") - assert calls == [ - f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1" - ] + assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1"] def test_v2_only_workspace_succeeds_after_v3_probe(self, monkeypatch): calls: list[str] = [] @@ -1910,9 +1908,7 @@ def fake_get(url, token): with pytest.raises(RuntimeError, match="rejected"): db_mod.ensure_ai_gateway(WS, "fake-token") - assert calls == [ - f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1" - ] + 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"]) 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",