From 48e98934653b0818ab1ad0c1ed22ecfc87e82aae Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 13:53:05 +0000 Subject: [PATCH 1/7] Show AI Gateway versions during configure --- src/ucode/cli.py | 6 ++-- src/ucode/databricks.py | 65 +++++++++++++++++++++++++++------------- tests/test_cli.py | 21 +++++++++++-- tests/test_databricks.py | 44 +++++++++++++++++++++++---- 4 files changed, 104 insertions(+), 32 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6a1662ae..94680677 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -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,10 @@ def configure_shared_state( state["profile"] = profile with spinner("Verifying Unity AI Gateway..."): token = get_databricks_token(workspace, profile) - ensure_ai_gateway(workspace, token) + gateway_capabilities = ensure_ai_gateway(workspace, token) print_success("Unity AI Gateway detected") + print_kv("V3 model services", gateway_capabilities.v3.detail) + print_kv("V2 endpoints", gateway_capabilities.v2.detail) want_claude = ( fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f875a5c7..5e5c006f 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3046,18 +3046,42 @@ 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 + + +class GatewayCapabilities(NamedTuple): + v3: GatewayProbe + v2: GatewayProbe + + +def _gateway_probe_result( + payload: dict | list | None, + reason: str | None, + collection_key: str, + resource_name: str, +) -> GatewayProbe: + if payload is None: + return GatewayProbe(False, 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") + return GatewayProbe(True, f"reachable, no accessible {resource_name}s returned") + + +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, reason, "endpoints", "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, reason, "model_services", "model service") def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: @@ -3090,27 +3114,26 @@ def _raise_ai_gateway_v2_permission_failure( ) -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) +def ensure_ai_gateway(workspace: str, token: str) -> GatewayCapabilities: + """Return both gateway probe results if either V2 or V3 is reachable.""" + v3 = _probe_ai_gateway_v3(workspace, token) + if not v3.reachable and _looks_like_definitive_auth_failure(v3.detail): + _raise_ai_gateway_auth_failure(workspace, v3.detail) - 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) + v2 = _probe_ai_gateway_v2(workspace, token) + capabilities = GatewayCapabilities(v3=v3, v2=v2) + if v3.reachable or v2.reachable: + return capabilities + if _looks_like_definitive_auth_failure(v2.detail): + _raise_ai_gateway_auth_failure(workspace, v2.detail) + if _looks_like_permission_failure(v3.detail): + _raise_ai_gateway_v3_permission_failure(workspace, v3.detail, v2.detail) + if _looks_like_permission_failure(v2.detail): + _raise_ai_gateway_v2_permission_failure(workspace, v2.detail, v3.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"({v3.detail}) nor V2 ({v2.detail}) is available. " f"See {AI_GATEWAY_V2_DOCS_URL}" ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 8bb1ea9b..56cf777d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -14,6 +14,7 @@ from typer.testing import CliRunner from ucode.cli import app +from ucode.databricks import GatewayCapabilities, GatewayProbe _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") @@ -29,6 +30,11 @@ def _strip_ansi(text: str) -> str: TOOLS = ["codex", "claude", "gemini", "opencode"] +GATEWAY_CAPABILITIES = GatewayCapabilities( + v3=GatewayProbe(True, "reachable, accessible model service returned"), + v2=GatewayProbe(False, "HTTP 404 Not Found"), +) + def _jwt(expires_at: float) -> str: payload = base64.urlsafe_b64encode(json.dumps({"exp": expires_at}).encode()).decode() @@ -2466,7 +2472,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", lambda w, t: None) + monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: GATEWAY_CAPABILITIES) 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,15 @@ 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_v3_and_v2_gateway_capabilities(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 "V3 model services: reachable, accessible model service returned" in output + assert "V2 endpoints: HTTP 404 Not Found" in output + def test_use_pat_without_pat_profile_raises(self, monkeypatch): cli_mod, logins, _, _ = self._stub_deps(monkeypatch, pat_token=None) @@ -2787,7 +2802,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", lambda w, t: None) + monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: GATEWAY_CAPABILITIES) 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 +2862,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", lambda w, t: None) + monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: GATEWAY_CAPABILITIES) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) monkeypatch.setattr(cli_mod, "save_state", lambda s: None) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 2d9f61a5..2307e32a 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1916,18 +1916,42 @@ def fake_run(args, **kwargs): class TestEnsureAiGateway: - def test_v3_only_workspace_succeeds_without_v2_probe(self, monkeypatch): + def test_reports_v3_resource_and_empty_v2_listing(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): calls.append(url) - return {"model_services": []}, None + if "/api/2.1/unity-catalog/model-services" in url: + return {"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None + return {"endpoints": []}, None monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - db_mod.ensure_ai_gateway(WS, "fake-token") + capabilities = 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 capabilities.v3 == db_mod.GatewayProbe( + True, "reachable, accessible model service returned" + ) + assert capabilities.v2 == db_mod.GatewayProbe( + True, "reachable, no accessible endpoints returned" + ) + 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_empty_v3_response_is_reachable_but_not_reported_as_accessible(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: ({}, None), + ) + + capabilities = db_mod.ensure_ai_gateway(WS, "fake-token") + + assert capabilities.v3 == db_mod.GatewayProbe( + True, "reachable, no accessible model services returned" + ) def test_v2_only_workspace_succeeds_after_v3_probe(self, monkeypatch): calls: list[str] = [] @@ -1940,8 +1964,12 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - db_mod.ensure_ai_gateway(WS, "fake-token") + capabilities = db_mod.ensure_ai_gateway(WS, "fake-token") + assert capabilities.v3 == db_mod.GatewayProbe(False, "HTTP 404: Not Found") + assert capabilities.v2 == db_mod.GatewayProbe( + True, "reachable, no accessible endpoints returned" + ) 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", @@ -1958,8 +1986,12 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - db_mod.ensure_ai_gateway(WS, "fake-token") + capabilities = db_mod.ensure_ai_gateway(WS, "fake-token") + assert capabilities.v3 == db_mod.GatewayProbe(False, "HTTP 403: Forbidden") + assert capabilities.v2 == db_mod.GatewayProbe( + True, "reachable, no accessible endpoints returned" + ) 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", From 959900d9655ddaabce1657886f41cdfa9458e891 Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 14:27:58 +0000 Subject: [PATCH 2/7] Add guidance for empty V3 model services --- src/ucode/databricks.py | 14 ++++++++++++-- tests/test_databricks.py | 6 ++++-- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 5e5c006f..2f31292a 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3061,13 +3061,17 @@ def _gateway_probe_result( reason: str | None, collection_key: str, resource_name: str, + empty_hint: str | None = None, ) -> GatewayProbe: if payload is None: return GatewayProbe(False, 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") - return GatewayProbe(True, f"reachable, no accessible {resource_name}s returned") + 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: @@ -3081,7 +3085,13 @@ 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 _gateway_probe_result(payload, reason, "model_services", "model service") + return _gateway_probe_result( + payload, + reason, + "model_services", + "model service", + "check USE CATALOG on system, and USE SCHEMA and EXECUTE on system.ai", + ) def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 2307e32a..e55d4d1b 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1940,7 +1940,7 @@ def fake_get(url, token): f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", ] - def test_empty_v3_response_is_reachable_but_not_reported_as_accessible(self, monkeypatch): + def test_empty_v3_response_includes_permission_hint(self, monkeypatch): monkeypatch.setattr( db_mod, "_http_get_json", @@ -1950,7 +1950,9 @@ def test_empty_v3_response_is_reachable_but_not_reported_as_accessible(self, mon capabilities = db_mod.ensure_ai_gateway(WS, "fake-token") assert capabilities.v3 == db_mod.GatewayProbe( - True, "reachable, no accessible model services returned" + True, + "reachable, no accessible model services returned; check USE CATALOG on system, and " + "USE SCHEMA and EXECUTE on system.ai", ) def test_v2_only_workspace_succeeds_after_v3_probe(self, monkeypatch): From 807bb6d4714fda840633a0fc8ef6d530eaa9d2e1 Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 15:24:57 +0000 Subject: [PATCH 3/7] Isolate publish CLI tests from bootstrap --- tests/test_managed_wizard.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 013cdc9d..0fcbebac 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -2857,8 +2857,11 @@ def test_publish_file_flag_is_forwarded(self): assert publish.call_args.kwargs["file_path"] == "/tmp/cfg.json" def test_publish_error_exits_nonzero_with_a_message(self): - with patch.object( - cli_mod, "publish_command", side_effect=RuntimeError("no config authored") + with ( + patch("ucode.cli.install_databricks_cli"), + patch.object( + cli_mod, "publish_command", side_effect=RuntimeError("no config authored") + ), ): result = runner.invoke(app, ["publish"]) assert result.exit_code == 1 @@ -2866,7 +2869,10 @@ def test_publish_error_exits_nonzero_with_a_message(self): def test_successful_publish_exits_zero(self): # Same trap as `setup`: `typer.Exit` subclasses RuntimeError, so raising it inside the # command's try block would report success as "ERROR 0". - with patch.object(cli_mod, "publish_command", return_value=0): + with ( + patch("ucode.cli.install_databricks_cli"), + patch.object(cli_mod, "publish_command", return_value=0), + ): result = runner.invoke(app, ["publish"]) assert result.exit_code == 0 assert "ERROR" not in result.output From cfc0deadcad323e12f5e7948f9257ebdd2fa6923 Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 15:31:07 +0000 Subject: [PATCH 4/7] Cover local gateway probe scenarios --- tests/test_cli.py | 94 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) diff --git a/tests/test_cli.py b/tests/test_cli.py index 56cf777d..16f5520d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -13,6 +13,7 @@ import pytest from typer.testing import CliRunner +import ucode.databricks as db_mod from ucode.cli import app from ucode.databricks import GatewayCapabilities, GatewayProbe @@ -2504,6 +2505,99 @@ def test_prints_v3_and_v2_gateway_capabilities(self, monkeypatch, capsys): assert "V3 model services: reachable, accessible model service returned" in output assert "V2 endpoints: HTTP 404 Not Found" in output + @pytest.mark.parametrize( + ("v3_response", "v2_response", "expected_v3", "expected_v2"), + [ + ( + ({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None), + ({"endpoints": [{"name": "databricks-gpt-5"}]}, None), + "reachable, accessible model service returned", + "reachable, accessible endpoint returned", + ), + ( + ({}, None), + ({"endpoints": []}, None), + "reachable, no accessible model services returned; check USE CATALOG on system, " + "and USE SCHEMA and EXECUTE on system.ai", + "reachable, no accessible endpoints returned", + ), + ( + ({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None), + (None, "HTTP 404 Not Found: FEATURE_DISABLED"), + "reachable, accessible model service returned", + "HTTP 404 Not Found: FEATURE_DISABLED", + ), + ( + ({}, None), + (None, "HTTP 404 Not Found: FEATURE_DISABLED"), + "reachable, no accessible model services returned; check USE CATALOG on system, " + "and USE SCHEMA and EXECUTE on system.ai", + "HTTP 404 Not Found: FEATURE_DISABLED", + ), + ( + (None, "HTTP 403 Forbidden"), + ({"endpoints": [{"name": "databricks-gpt-5"}]}, None), + "HTTP 403 Forbidden", + "reachable, accessible endpoint returned", + ), + ], + ids=[ + "both-return-resources", + "both-return-empty-listings", + "v3-resource-v2-unavailable", + "v3-empty-v2-unavailable", + "v3-forbidden-v2-resource", + ], + ) + def test_prints_local_gateway_probe_scenarios( + self, + monkeypatch, + capsys, + v3_response, + v2_response, + expected_v3, + expected_v2, + ): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + responses = iter([v3_response, v2_response]) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: next(responses)) + monkeypatch.setattr(cli_mod, "ensure_ai_gateway", db_mod.ensure_ai_gateway) + + 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"V3 model services: {expected_v3}" in output + assert f"V2 endpoints: {expected_v2}" in output + + @pytest.mark.parametrize( + ("responses", "error_match"), + [ + ( + [ + (None, "HTTP 404 Not Found: V3 missing"), + (None, "HTTP 404 Not Found: V2 missing"), + ], + "neither V3", + ), + ([(None, "HTTP 401 Unauthorized")], "rejected the access token"), + ], + ids=["neither-version-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, "ensure_ai_gateway", db_mod.ensure_ai_gateway) + + with pytest.raises(RuntimeError, match=error_match): + cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + output = _strip_ansi(capsys.readouterr().out) + assert "Unity AI Gateway detected" not in output + def test_use_pat_without_pat_profile_raises(self, monkeypatch): cli_mod, logins, _, _ = self._stub_deps(monkeypatch, pat_token=None) From 0bc2d73ea0bc360135557a1107dce1b165c56e43 Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 15:58:56 +0000 Subject: [PATCH 5/7] Probe legacy endpoints only as fallback --- src/ucode/cli.py | 5 ++-- src/ucode/databricks.py | 9 ++++-- tests/test_cli.py | 62 ++++++++++++++++++---------------------- tests/test_databricks.py | 13 +++------ 4 files changed, 41 insertions(+), 48 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 94680677..aa9473b2 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -637,8 +637,9 @@ def configure_shared_state( token = get_databricks_token(workspace, profile) gateway_capabilities = ensure_ai_gateway(workspace, token) print_success("Unity AI Gateway detected") - print_kv("V3 model services", gateway_capabilities.v3.detail) - print_kv("V2 endpoints", gateway_capabilities.v2.detail) + print_kv("Model service", gateway_capabilities.v3.detail) + if gateway_capabilities.v2 is not None: + print_kv("(Legacy) endpoints", gateway_capabilities.v2.detail) want_claude = ( fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 2f31292a..4abe4ff3 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3049,11 +3049,12 @@ def fetch_codex_models(workspace: str, token: str) -> list[str]: class GatewayProbe(NamedTuple): reachable: bool detail: str + resource_available: bool = False class GatewayCapabilities(NamedTuple): v3: GatewayProbe - v2: GatewayProbe + v2: GatewayProbe | None def _gateway_probe_result( @@ -3067,7 +3068,7 @@ def _gateway_probe_result( return GatewayProbe(False, 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") + 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}" @@ -3125,10 +3126,12 @@ def _raise_ai_gateway_v2_permission_failure( def ensure_ai_gateway(workspace: str, token: str) -> GatewayCapabilities: - """Return both gateway probe results if either V2 or V3 is reachable.""" + """Return gateway probe results, checking V2 only when V3 has no accessible resource.""" v3 = _probe_ai_gateway_v3(workspace, token) if not v3.reachable and _looks_like_definitive_auth_failure(v3.detail): _raise_ai_gateway_auth_failure(workspace, v3.detail) + if v3.resource_available: + return GatewayCapabilities(v3=v3, v2=None) v2 = _probe_ai_gateway_v2(workspace, token) capabilities = GatewayCapabilities(v3=v3, v2=v2) diff --git a/tests/test_cli.py b/tests/test_cli.py index 16f5520d..e24be330 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -32,8 +32,8 @@ def _strip_ansi(text: str) -> str: TOOLS = ["codex", "claude", "gemini", "opencode"] GATEWAY_CAPABILITIES = GatewayCapabilities( - v3=GatewayProbe(True, "reachable, accessible model service returned"), - v2=GatewayProbe(False, "HTTP 404 Not Found"), + v3=GatewayProbe(True, "reachable, accessible model service returned", True), + v2=None, ) @@ -2496,79 +2496,73 @@ 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_v3_and_v2_gateway_capabilities(self, monkeypatch, capsys): + 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 "V3 model services: reachable, accessible model service returned" in output - assert "V2 endpoints: HTTP 404 Not Found" in output + assert "Model service: reachable, accessible model service returned" in output + assert "(Legacy) endpoints:" not in output @pytest.mark.parametrize( - ("v3_response", "v2_response", "expected_v3", "expected_v2"), + ("responses", "expected_model_service", "expected_legacy_endpoints"), [ ( - ({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None), - ({"endpoints": [{"name": "databricks-gpt-5"}]}, None), + [({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None)], "reachable, accessible model service returned", - "reachable, accessible endpoint returned", + None, ), ( - ({}, None), - ({"endpoints": []}, None), + [({}, None), ({"endpoints": []}, None)], "reachable, no accessible model services returned; check USE CATALOG on system, " "and USE SCHEMA and EXECUTE on system.ai", "reachable, no accessible endpoints returned", ), ( - ({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None), - (None, "HTTP 404 Not Found: FEATURE_DISABLED"), - "reachable, accessible model service returned", - "HTTP 404 Not Found: FEATURE_DISABLED", - ), - ( - ({}, None), - (None, "HTTP 404 Not Found: FEATURE_DISABLED"), + [({}, None), (None, "HTTP 404 Not Found: FEATURE_DISABLED")], "reachable, no accessible model services returned; check USE CATALOG on system, " "and USE SCHEMA and EXECUTE on system.ai", "HTTP 404 Not Found: FEATURE_DISABLED", ), ( - (None, "HTTP 403 Forbidden"), - ({"endpoints": [{"name": "databricks-gpt-5"}]}, None), + [ + (None, "HTTP 403 Forbidden"), + ({"endpoints": [{"name": "databricks-gpt-5"}]}, None), + ], "HTTP 403 Forbidden", "reachable, accessible endpoint returned", ), ], ids=[ - "both-return-resources", - "both-return-empty-listings", - "v3-resource-v2-unavailable", - "v3-empty-v2-unavailable", - "v3-forbidden-v2-resource", + "model-service-resource", + "model-service-empty-legacy-empty", + "model-service-empty-legacy-unavailable", + "model-service-forbidden-legacy-resource", ], ) def test_prints_local_gateway_probe_scenarios( self, monkeypatch, capsys, - v3_response, - v2_response, - expected_v3, - expected_v2, + responses, + expected_model_service, + expected_legacy_endpoints, ): cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") - responses = iter([v3_response, v2_response]) - monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: next(responses)) + response_iter = iter(responses) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: next(response_iter)) monkeypatch.setattr(cli_mod, "ensure_ai_gateway", db_mod.ensure_ai_gateway) 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"V3 model services: {expected_v3}" in output - assert f"V2 endpoints: {expected_v2}" in output + assert f"Model service: {expected_model_service}" in output + if expected_legacy_endpoints is None: + assert "(Legacy) endpoints:" not in output + else: + assert f"(Legacy) endpoints: {expected_legacy_endpoints}" in output @pytest.mark.parametrize( ("responses", "error_match"), diff --git a/tests/test_databricks.py b/tests/test_databricks.py index e55d4d1b..06ff8211 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1916,28 +1916,23 @@ def fake_run(args, **kwargs): class TestEnsureAiGateway: - def test_reports_v3_resource_and_empty_v2_listing(self, monkeypatch): + def test_v3_resource_skips_v2_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 {"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None - return {"endpoints": []}, None + return {"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None monkeypatch.setattr(db_mod, "_http_get_json", fake_get) capabilities = db_mod.ensure_ai_gateway(WS, "fake-token") assert capabilities.v3 == db_mod.GatewayProbe( - True, "reachable, accessible model service returned" - ) - assert capabilities.v2 == db_mod.GatewayProbe( - True, "reachable, no accessible endpoints returned" + True, "reachable, accessible model service returned", True ) + assert capabilities.v2 is None 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_empty_v3_response_includes_permission_hint(self, monkeypatch): From 3f710d1d4dbef7f3aa7cabd432d2559bb4c2f446 Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 16:07:38 +0000 Subject: [PATCH 6/7] Simplify gateway fallback result --- src/ucode/databricks.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 4abe4ff3..4869dd49 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3126,7 +3126,7 @@ def _raise_ai_gateway_v2_permission_failure( def ensure_ai_gateway(workspace: str, token: str) -> GatewayCapabilities: - """Return gateway probe results, checking V2 only when V3 has no accessible resource.""" + """Probe model services first, falling back to legacy endpoints when needed.""" v3 = _probe_ai_gateway_v3(workspace, token) if not v3.reachable and _looks_like_definitive_auth_failure(v3.detail): _raise_ai_gateway_auth_failure(workspace, v3.detail) @@ -3134,9 +3134,8 @@ def ensure_ai_gateway(workspace: str, token: str) -> GatewayCapabilities: return GatewayCapabilities(v3=v3, v2=None) v2 = _probe_ai_gateway_v2(workspace, token) - capabilities = GatewayCapabilities(v3=v3, v2=v2) if v3.reachable or v2.reachable: - return capabilities + return GatewayCapabilities(v3=v3, v2=v2) if _looks_like_definitive_auth_failure(v2.detail): _raise_ai_gateway_auth_failure(workspace, v2.detail) if _looks_like_permission_failure(v3.detail): From 7a2446b927008446534ace61e1df14f43a09b263 Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 17:05:33 +0000 Subject: [PATCH 7/7] Address gateway capability review feedback --- src/ucode/cli.py | 8 ++- src/ucode/databricks.py | 101 +++++++++++++++++++---------------- tests/test_cli.py | 81 ++++++++++++++++------------ tests/test_databricks.py | 110 ++++++++++++++++++++++++++------------- tests/test_e2e.py | 6 +-- 5 files changed, 187 insertions(+), 119 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index aa9473b2..d0789093 100644 --- a/src/ucode/cli.py +++ b/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, @@ -635,11 +635,9 @@ def configure_shared_state( state["profile"] = profile with spinner("Verifying Unity AI Gateway..."): token = get_databricks_token(workspace, profile) - gateway_capabilities = ensure_ai_gateway(workspace, token) + model_service_probe = probe_unity_gateway_capabilities(workspace, token) print_success("Unity AI Gateway detected") - print_kv("Model service", gateway_capabilities.v3.detail) - if gateway_capabilities.v2 is not None: - print_kv("(Legacy) endpoints", gateway_capabilities.v2.detail) + 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 diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 4869dd49..33e36c2e 100644 --- a/src/ucode/databricks.py +++ b/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) @@ -3052,9 +3052,9 @@ class GatewayProbe(NamedTuple): resource_available: bool = False -class GatewayCapabilities(NamedTuple): - v3: GatewayProbe - v2: GatewayProbe | None +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( @@ -3065,7 +3065,7 @@ def _gateway_probe_result( empty_hint: str | None = None, ) -> GatewayProbe: if payload is None: - return GatewayProbe(False, reason or "unknown error") + 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) @@ -3079,7 +3079,12 @@ 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 _gateway_probe_result(payload, reason, "endpoints", "endpoint") + return _gateway_probe_result( + payload=payload, + reason=reason, + collection_key="endpoints", + resource_name="endpoint", + ) def _probe_ai_gateway_v3(workspace: str, token: str) -> GatewayProbe: @@ -3087,11 +3092,11 @@ def _probe_ai_gateway_v3(workspace: str, token: str) -> GatewayProbe: url = f"https://{hostname}/api/2.1/unity-catalog/model-services?page_size=1" payload, reason = _http_get_json(url, token) return _gateway_probe_result( - payload, - reason, - "model_services", - "model service", - "check USE CATALOG on system, and USE SCHEMA and EXECUTE on system.ai", + 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", ) @@ -3104,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) -> GatewayCapabilities: - """Probe model services first, falling back to legacy endpoints when needed.""" - v3 = _probe_ai_gateway_v3(workspace, token) - if not v3.reachable and _looks_like_definitive_auth_failure(v3.detail): - _raise_ai_gateway_auth_failure(workspace, v3.detail) - if v3.resource_available: - return GatewayCapabilities(v3=v3, v2=None) - - v2 = _probe_ai_gateway_v2(workspace, token) - if v3.reachable or v2.reachable: - return GatewayCapabilities(v3=v3, v2=v2) - if _looks_like_definitive_auth_failure(v2.detail): - _raise_ai_gateway_auth_failure(workspace, v2.detail) - if _looks_like_permission_failure(v3.detail): - _raise_ai_gateway_v3_permission_failure(workspace, v3.detail, v2.detail) - if _looks_like_permission_failure(v2.detail): - _raise_ai_gateway_v2_permission_failure(workspace, v2.detail, v3.detail) +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.detail}) nor V2 ({v2.detail}) 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 diff --git a/tests/test_cli.py b/tests/test_cli.py index e24be330..a59f3015 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -15,7 +15,7 @@ import ucode.databricks as db_mod from ucode.cli import app -from ucode.databricks import GatewayCapabilities, GatewayProbe +from ucode.databricks import GatewayProbe _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") @@ -31,10 +31,7 @@ def _strip_ansi(text: str) -> str: TOOLS = ["codex", "claude", "gemini", "opencode"] -GATEWAY_CAPABILITIES = GatewayCapabilities( - v3=GatewayProbe(True, "reachable, accessible model service returned", True), - v2=None, -) +MODEL_SERVICE_PROBE = GatewayProbe(True, "reachable, accessible model service returned", True) def _jwt(expires_at: float) -> str: @@ -2473,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: GATEWAY_CAPABILITIES) + 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)) @@ -2506,24 +2505,16 @@ def test_prints_model_service_and_omits_unneeded_legacy_probe(self, monkeypatch, assert "(Legacy) endpoints:" not in output @pytest.mark.parametrize( - ("responses", "expected_model_service", "expected_legacy_endpoints"), + ("responses", "expected_model_service"), [ ( [({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None)], "reachable, accessible model service returned", - None, ), ( [({}, None), ({"endpoints": []}, None)], "reachable, no accessible model services returned; check USE CATALOG on system, " "and USE SCHEMA and EXECUTE on system.ai", - "reachable, no accessible endpoints returned", - ), - ( - [({}, None), (None, "HTTP 404 Not Found: FEATURE_DISABLED")], - "reachable, no accessible model services returned; check USE CATALOG on system, " - "and USE SCHEMA and EXECUTE on system.ai", - "HTTP 404 Not Found: FEATURE_DISABLED", ), ( [ @@ -2531,13 +2522,11 @@ def test_prints_model_service_and_omits_unneeded_legacy_probe(self, monkeypatch, ({"endpoints": [{"name": "databricks-gpt-5"}]}, None), ], "HTTP 403 Forbidden", - "reachable, accessible endpoint returned", ), ], ids=[ "model-service-resource", "model-service-empty-legacy-empty", - "model-service-empty-legacy-unavailable", "model-service-forbidden-legacy-resource", ], ) @@ -2547,36 +2536,51 @@ def test_prints_local_gateway_probe_scenarios( capsys, responses, expected_model_service, - expected_legacy_endpoints, ): 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, "ensure_ai_gateway", db_mod.ensure_ai_gateway) + 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 - if expected_legacy_endpoints is None: - assert "(Legacy) endpoints:" not in output - else: - assert f"(Legacy) endpoints: {expected_legacy_endpoints}" 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, "HTTP 404 Not Found: V3 missing"), - (None, "HTTP 404 Not Found: V2 missing"), + ({}, None), + ( + None, + "HTTP 404 Not Found: AI Gateway V2 is not available for CSP-enabled " + "workspaces", + ), ], - "neither V3", + "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=["neither-version-reachable", "invalid-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 @@ -2584,13 +2588,18 @@ def test_local_gateway_probe_failures_do_not_print_success( 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, "ensure_ai_gateway", db_mod.ensure_ai_gateway) + monkeypatch.setattr( + cli_mod, "probe_unity_gateway_capabilities", db_mod.probe_unity_gateway_capabilities + ) - with pytest.raises(RuntimeError, match=error_match): + 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) @@ -2890,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: GATEWAY_CAPABILITIES) + 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)) @@ -2950,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: GATEWAY_CAPABILITIES) + 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) @@ -3008,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") diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 06ff8211..0df389cc 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1915,8 +1915,8 @@ def fake_run(args, **kwargs): list_databricks_apps(WS) -class TestEnsureAiGateway: - def test_v3_resource_skips_v2_probe(self, monkeypatch): +class TestProbeUnityGatewayCapabilities: + def test_model_service_resource_skips_legacy_probe(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): @@ -1925,32 +1925,31 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - capabilities = db_mod.ensure_ai_gateway(WS, "fake-token") + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") - assert capabilities.v3 == db_mod.GatewayProbe( + assert model_service_probe == db_mod.GatewayProbe( True, "reachable, accessible model service returned", True ) - assert capabilities.v2 is None assert calls == [ f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1", ] - def test_empty_v3_response_includes_permission_hint(self, monkeypatch): + def test_empty_model_service_response_includes_permission_hint(self, monkeypatch): monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token: ({}, None), ) - capabilities = db_mod.ensure_ai_gateway(WS, "fake-token") + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") - assert capabilities.v3 == db_mod.GatewayProbe( + assert model_service_probe == db_mod.GatewayProbe( True, "reachable, no accessible model services returned; check USE CATALOG on system, and " "USE SCHEMA and EXECUTE on system.ai", ) - def test_v2_only_workspace_succeeds_after_v3_probe(self, monkeypatch): + def test_legacy_only_workspace_returns_model_service_probe(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): @@ -1961,18 +1960,15 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - capabilities = db_mod.ensure_ai_gateway(WS, "fake-token") + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") - assert capabilities.v3 == db_mod.GatewayProbe(False, "HTTP 404: Not Found") - assert capabilities.v2 == db_mod.GatewayProbe( - True, "reachable, no accessible endpoints returned" - ) + assert model_service_probe == db_mod.GatewayProbe(False, "HTTP 404: Not Found") 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_v3_forbidden_still_succeeds_when_v2_is_available(self, monkeypatch): + def test_model_service_forbidden_still_succeeds_when_legacy_is_available(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): @@ -1983,33 +1979,58 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - capabilities = db_mod.ensure_ai_gateway(WS, "fake-token") + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") - assert capabilities.v3 == db_mod.GatewayProbe(False, "HTTP 403: Forbidden") - assert capabilities.v2 == db_mod.GatewayProbe( - True, "reachable, no accessible endpoints returned" - ) + assert model_service_probe == db_mod.GatewayProbe(False, "HTTP 403: Forbidden") 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_empty_model_service_requires_reachable_legacy_fallback(self, monkeypatch): + responses = iter( + [ + ({}, None), + (None, "HTTP 404: AI Gateway V2 is not available for CSP-enabled workspaces"), + ] + ) + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: next(responses), + ) + + with pytest.raises(RuntimeError, match="no accessible model services") as excinfo: + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + + message = str(excinfo.value) + assert "HTTP 404: AI Gateway legacy endpoint is not available" in message + assert "v2" not in message.lower() + assert "v3" not in message.lower() + def test_neither_gateway_available_raises(self, monkeypatch): - reasons = iter(["HTTP 404: V3 missing", "HTTP 404: V2 missing"]) + reasons = iter( + [ + "HTTP 404: V3 unavailable", + "HTTP 404: V2 unavailable", + ] + ) monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token: (None, next(reasons)), ) - with pytest.raises(RuntimeError, match="neither V3") as excinfo: - db_mod.ensure_ai_gateway(WS, "fake-token") + with pytest.raises(RuntimeError, match="neither model services") as excinfo: + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") message = str(excinfo.value) - assert "HTTP 404: V2 missing" in message - assert "HTTP 404: V3 missing" in message + assert "HTTP 404: model service unavailable" in message + assert "HTTP 404: legacy endpoint unavailable" in message + assert "v2" not in message.lower() + assert "v3" not in message.lower() - def test_v3_auth_failure_does_not_probe_v2(self, monkeypatch): + def test_model_service_auth_failure_does_not_probe_legacy(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): @@ -2019,12 +2040,19 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) with pytest.raises(RuntimeError, match="rejected"): - db_mod.ensure_ai_gateway(WS, "fake-token") + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") 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"]) + def test_model_service_forbidden_and_legacy_unavailable_reports_permission_error( + self, monkeypatch + ): + reasons = iter( + [ + "HTTP 403: Missing Unity Catalog grants", + "HTTP 404: legacy endpoints unavailable", + ] + ) monkeypatch.setattr( db_mod, "_http_get_json", @@ -2032,15 +2060,25 @@ def test_v3_forbidden_and_v2_unavailable_reports_permission_error(self, monkeypa ) with pytest.raises(RuntimeError, match="permission") as excinfo: - db_mod.ensure_ai_gateway(WS, "fake-token") + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") message = str(excinfo.value) assert "USE SCHEMA" in message + assert "EXECUTE" in message + assert "v2" not in message.lower() + assert "v3" not in message.lower() 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"]) + def test_legacy_forbidden_and_model_service_unavailable_reports_permission_error( + self, monkeypatch + ): + reasons = iter( + [ + "HTTP 404: V3 unavailable", + "HTTP 403: V2 forbidden", + ] + ) monkeypatch.setattr( db_mod, "_http_get_json", @@ -2048,16 +2086,18 @@ def test_v2_forbidden_and_v3_unavailable_reports_permission_error(self, monkeypa ) with pytest.raises(RuntimeError, match="workspace permissions") as excinfo: - db_mod.ensure_ai_gateway(WS, "fake-token") + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") message = str(excinfo.value) - assert "V2 access could not be verified" in message + assert "legacy endpoint access could not be verified" in message + assert "v2" not in message.lower() + assert "v3" not in message.lower() 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) can route on it. Before issue #84's fix + so callers (e.g. the Unity Gateway capability probe) 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 3b9319eb..a578e5b2 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -27,7 +27,6 @@ build_tool_base_url, discover_model_services, discover_sql_warehouses, - ensure_ai_gateway, fetch_ai_gateway_claude_models, fetch_codex_models, fetch_gemini_models, @@ -35,6 +34,7 @@ is_model_provider_feature_unavailable, list_model_provider_services, list_tool_provider_services, + probe_unity_gateway_capabilities, service_usable_for_tool, workspace_hostname, ) @@ -137,8 +137,8 @@ def test_get_token_returns_non_empty_string(self, 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_probe_unity_gateway_capabilities_does_not_raise(self, e2e_workspace, e2e_token): + probe_unity_gateway_capabilities(e2e_workspace, e2e_token) def test_workspace_hostname_resolves(self, e2e_workspace): hostname = workspace_hostname(e2e_workspace)