From 5d83a9ef22a87b01a5945a8675769170ada79676 Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 19:42:43 +0000 Subject: [PATCH 1/6] Page the model-services probe past an ACL-emptied first page `ucode configure`'s model-service probe listed with page_size=1 and judged accessibility from that single page. Unity Catalog applies page_size before ACL filtering, so an early page can return empty with only a next_page_token while accessible model services remain behind the cursor. On the ai-devtools gateway workspace this made configure report "no accessible model services returned; check USE CATALOG ..." even though 42 model services were accessible. Follow the cursor (page_size=50, bounded) before reporting empty, and route a missing-OAuth-scope 403 to the re-login guidance instead of the UC-grant hint, which does not apply to a token-scope failure. Related: ES-2185388 Co-authored-by: Isaac --- src/ucode/databricks.py | 55 +++++++++++++++++------- tests/test_databricks.py | 92 ++++++++++++++++++++++++++++++++++++++-- 2 files changed, 127 insertions(+), 20 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 33e36c2e..f67734a5 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3062,17 +3062,13 @@ 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, _version_neutral_gateway_detail(reason or "unknown error")) resources = payload.get(collection_key) if isinstance(payload, dict) else None if resources: return GatewayProbe(True, f"reachable, accessible {resource_name} returned", True) - detail = f"reachable, no accessible {resource_name}s returned" - if empty_hint: - detail = f"{detail}; {empty_hint}" - return GatewayProbe(True, detail) + return GatewayProbe(True, f"reachable, no accessible {resource_name}s returned") def _probe_ai_gateway_v2(workspace: str, token: str) -> GatewayProbe: @@ -3087,17 +3083,39 @@ def _probe_ai_gateway_v2(workspace: str, token: str) -> GatewayProbe: ) +# The model-services listing applies `page_size` before ACL filtering, so an +# early page can come back empty (only a `next_page_token`) while accessible +# rows remain further in. Follow the cursor before concluding nothing is +# accessible, or a caller with many services still reads as empty. +_MODEL_SERVICE_PROBE_PAGE_SIZE = 50 +_MODEL_SERVICE_PROBE_MAX_PAGES = 20 +_MODEL_SERVICE_EMPTY_DETAIL = ( + "reachable, no accessible model services returned; " + "check USE CATALOG on system, and USE SCHEMA and EXECUTE on system.ai" +) + + 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=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", - ) + base = f"https://{hostname}/api/2.1/unity-catalog/model-services" + page_token: str | None = None + for page in range(_MODEL_SERVICE_PROBE_MAX_PAGES): + params: dict[str, object] = {"page_size": _MODEL_SERVICE_PROBE_PAGE_SIZE} + if page_token: + params["page_token"] = page_token + payload, reason = _http_get_json(f"{base}?{urlencode(params)}", token) + if payload is None: + if page == 0: + return GatewayProbe( + False, _version_neutral_gateway_detail(reason or "unknown error") + ) + break # Reachability already confirmed; stop paging on a later-page error. + if isinstance(payload, dict) and payload.get("model_services"): + return GatewayProbe(True, "reachable, accessible model service returned", True) + page_token = payload.get("next_page_token") if isinstance(payload, dict) else None + if not page_token: + break + return GatewayProbe(True, _MODEL_SERVICE_EMPTY_DETAIL) def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: @@ -3167,11 +3185,16 @@ 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 preflight must still - try the fallback before surfacing it as an auth failure. + try the fallback before surfacing it as an auth failure. A missing-OAuth- + scope 403 is the exception: it 403s every workspace API, so re-login (not a + UC grant) is the fix. """ if "HTTP 401" in reason: return True - return "HTTP 400" in reason and "invalid token" in reason.lower() + lowered = reason.lower() + if "HTTP 403" in reason and "required scopes" in lowered: + return True + return "HTTP 400" in reason and "invalid token" in lowered def _looks_like_permission_failure(reason: str) -> bool: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 0df389cc..3480cfb4 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1931,7 +1931,67 @@ def fake_get(url, token): True, "reachable, accessible model service returned", True ) assert calls == [ - 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=50", + ] + + def test_empty_first_page_follows_cursor_to_accessible_model_service(self, monkeypatch): + # UC applies page_size before ACL filtering, so an early page can come + # back empty (only a next_page_token) while accessible rows remain behind + # the cursor (ES-2185388). The probe must page before reporting empty. + calls: list[str] = [] + responses = iter( + [ + ({"next_page_token": "cursor-1"}, None), + ({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None), + ] + ) + + def fake_get(url, token): + calls.append(url) + return next(responses) + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + + assert model_service_probe == db_mod.GatewayProbe( + True, "reachable, accessible model service returned", True + ) + assert calls == [ + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50", + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50" + "&page_token=cursor-1", + ] + + def test_empty_pages_exhaust_cursor_before_reporting_no_model_service(self, monkeypatch): + calls: list[str] = [] + responses = iter( + [ + ({"next_page_token": "cursor-1"}, None), + ({"model_services": []}, None), + ] + ) + + def fake_get(url, token): + calls.append(url) + if "/api/ai-gateway/v2/endpoints" in url: + return {"endpoints": []}, None + return next(responses) + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + + 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", + ) + assert calls == [ + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50", + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50" + "&page_token=cursor-1", + f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", ] def test_empty_model_service_response_includes_permission_hint(self, monkeypatch): @@ -1964,7 +2024,7 @@ def fake_get(url, token): 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/2.1/unity-catalog/model-services?page_size=50", f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", ] @@ -1983,7 +2043,7 @@ def fake_get(url, token): 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/2.1/unity-catalog/model-services?page_size=50", f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", ] @@ -2042,7 +2102,31 @@ def fake_get(url, token): with pytest.raises(RuntimeError, match="rejected"): 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"] + assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50"] + + def test_missing_scope_403_routes_to_reauth_not_grants(self, monkeypatch): + # A token missing the OAuth scope 403s every workspace API, so the fix is + # re-login, not a UC grant. It must not reach the grant-hint message, nor + # probe the legacy endpoint. + calls: list[str] = [] + + def fake_get(url, token): + calls.append(url) + return None, ( + "HTTP 403 Forbidden: Provided OAuth token does not have required " + "scopes: unity-catalog" + ) + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + with pytest.raises(RuntimeError, match="rejected the access token") as excinfo: + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + + message = str(excinfo.value) + assert "databricks auth login" in message + assert "USE CATALOG" not in message + assert "USE SCHEMA" not in message + assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50"] def test_model_service_forbidden_and_legacy_unavailable_reports_permission_error( self, monkeypatch From 4ce03cdaf2286269d070374c1a5998c68d3f96ac Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 20:01:58 +0000 Subject: [PATCH 2/6] Address review: don't recreate the false-empty conclusion on early stops Only report "no accessible model services" after walking the listing to the end (no pending cursor). A later-page error or hitting the page cap with a cursor still pending is now reachable-but-inconclusive, not empty. Stop treating a model-service scope 403 as globally definitive: the model- service and legacy-endpoint APIs require different OAuth scopes, so attempt the legacy fallback first. Only when the fallback also fails is a missing-scope 403 surfaced, and then as re-login guidance rather than a UC-grant hint. Co-authored-by: Isaac --- src/ucode/databricks.py | 45 +++++++++++++++++++------ tests/test_databricks.py | 71 ++++++++++++++++++++++++++++++++++++---- 2 files changed, 100 insertions(+), 16 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f67734a5..7d979f9c 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3109,13 +3109,18 @@ def _probe_ai_gateway_v3(workspace: str, token: str) -> GatewayProbe: return GatewayProbe( False, _version_neutral_gateway_detail(reason or "unknown error") ) - break # Reachability already confirmed; stop paging on a later-page error. + # A later page failed: reachable, but the listing was not walked to + # the end, so we can't claim nothing is accessible. + return GatewayProbe(True, "reachable") if isinstance(payload, dict) and payload.get("model_services"): return GatewayProbe(True, "reachable, accessible model service returned", True) page_token = payload.get("next_page_token") if isinstance(payload, dict) else None if not page_token: - break - return GatewayProbe(True, _MODEL_SERVICE_EMPTY_DETAIL) + # Walked the whole listing without an accessible model service. + return GatewayProbe(True, _MODEL_SERVICE_EMPTY_DETAIL) + # Hit the page cap with a cursor still pending: reachable but inconclusive, + # so don't assert the listing is empty. + return GatewayProbe(True, "reachable") def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: @@ -3127,6 +3132,15 @@ def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: ) +def _raise_ai_gateway_scope_failure(workspace: str, reason: str) -> NoReturn: + raise RuntimeError( + f"The access token for {workspace} is missing an OAuth scope required by the " + f"AI Gateway APIs ({reason}). Re-authenticate to mint a token with the needed " + f"scopes:\n" + f" databricks auth login --host {workspace}" + ) + + def _raise_model_service_permission_failure( workspace: str, model_service_reason: str, legacy_endpoint_reason: str ) -> NoReturn: @@ -3165,6 +3179,12 @@ def probe_unity_gateway_capabilities(workspace: str, token: str) -> GatewayProbe return model_service_probe if _looks_like_definitive_auth_failure(legacy_endpoint_probe.detail): _raise_ai_gateway_auth_failure(workspace, legacy_endpoint_probe.detail) + # A missing-scope 403 that survived the fallback is a token problem, not a + # UC grant, so surface re-login guidance before the permission hints. + if _looks_like_scope_failure(model_service_probe.detail): + _raise_ai_gateway_scope_failure(workspace, model_service_probe.detail) + if _looks_like_scope_failure(legacy_endpoint_probe.detail): + _raise_ai_gateway_scope_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 @@ -3185,16 +3205,21 @@ 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 preflight must still - try the fallback before surfacing it as an auth failure. A missing-OAuth- - scope 403 is the exception: it 403s every workspace API, so re-login (not a - UC grant) is the fix. + try the fallback before surfacing it as an auth failure. """ if "HTTP 401" in reason: return True - lowered = reason.lower() - if "HTTP 403" in reason and "required scopes" in lowered: - return True - return "HTTP 400" in reason and "invalid token" in lowered + return "HTTP 400" in reason and "invalid token" in reason.lower() + + +def _looks_like_scope_failure(reason: str) -> bool: + """True for a 403 that reports the OAuth token is missing a required scope. + + The scopes the model-service and legacy-endpoint APIs require differ, so + this is only conclusive once both probes have failed on it -- re-login (not + a UC grant) is the fix. + """ + return "HTTP 403" in reason and "required scopes" in reason.lower() def _looks_like_permission_failure(reason: str) -> bool: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 3480cfb4..9c0c3e41 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2104,10 +2104,10 @@ def fake_get(url, token): assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50"] - def test_missing_scope_403_routes_to_reauth_not_grants(self, monkeypatch): - # A token missing the OAuth scope 403s every workspace API, so the fix is - # re-login, not a UC grant. It must not reach the grant-hint message, nor - # probe the legacy endpoint. + def test_missing_scope_403_on_both_paths_routes_to_reauth_not_grants(self, monkeypatch): + # A missing-scope 403 is only conclusive once the legacy fallback also + # fails on it: the fix is re-login, not a UC grant. It must probe the + # fallback first and must not reach the grant-hint message. calls: list[str] = [] def fake_get(url, token): @@ -2119,14 +2119,73 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - with pytest.raises(RuntimeError, match="rejected the access token") as excinfo: + with pytest.raises(RuntimeError, match="missing an OAuth scope") as excinfo: db_mod.probe_unity_gateway_capabilities(WS, "fake-token") message = str(excinfo.value) assert "databricks auth login" in message assert "USE CATALOG" not in message assert "USE SCHEMA" not in message - assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50"] + assert calls == [ + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50", + f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", + ] + + def test_model_service_scope_403_succeeds_when_legacy_reachable(self, monkeypatch): + # A token scoped for the legacy endpoint but not for model services must + # still succeed via the fallback rather than surfacing the scope 403. + calls: list[str] = [] + + def fake_get(url, token): + calls.append(url) + if "/api/ai-gateway/v2/endpoints" in url: + return {"endpoints": [{"name": "databricks-gpt-5"}]}, None + return None, ( + "HTTP 403 Forbidden: Provided OAuth token does not have required " + "scopes: unity-catalog" + ) + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + + assert not model_service_probe.reachable + assert "required scopes" in model_service_probe.detail + assert calls == [ + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50", + f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", + ] + + def test_probe_v3_later_page_error_is_reachable_not_empty(self, monkeypatch): + # A later-page error leaves the listing un-walked, so the probe must not + # conclude nothing is accessible. + responses = iter( + [ + ({"next_page_token": "cursor-1"}, None), + (None, "HTTP 500: Internal Server Error"), + ] + ) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: next(responses)) + + assert db_mod._probe_ai_gateway_v3(WS, "fake-token") == db_mod.GatewayProbe( + True, "reachable" + ) + + def test_probe_v3_page_cap_with_pending_cursor_is_reachable_not_empty(self, monkeypatch): + # Exhausting the page cap while a cursor is still pending is inconclusive, + # not empty. + calls: list[str] = [] + + def fake_get(url, token): + calls.append(url) + return {"next_page_token": "more"}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + assert db_mod._probe_ai_gateway_v3(WS, "fake-token") == db_mod.GatewayProbe( + True, "reachable" + ) + assert len(calls) == db_mod._MODEL_SERVICE_PROBE_MAX_PAGES def test_model_service_forbidden_and_legacy_unavailable_reports_permission_error( self, monkeypatch From c87ce63479df551456101f74eaf7ce00892f0dad Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 20:13:17 +0000 Subject: [PATCH 3/6] Address Codex review: inconclusive v3 is reachable; scope match is OAuth-only A reachable-but-inconclusive model-service probe (later-page error or page cap) no longer hard-fails as "gateway not enabled" when the legacy fallback is unavailable: the API answered, so it is enabled. GatewayProbe carries an explicit `conclusive` flag to separate "confirmed no accessible resource" from "reachable, unknown", and the coordinator returns the probe in the inconclusive case rather than raising. Scope the missing-scope classifier to the OAuth-token wording so a PAT's permission 403 is not misrouted to OAuth re-login guidance (which cannot fix a PAT) and instead falls through to the grant hint. Co-authored-by: Isaac --- src/ucode/databricks.py | 26 +++++++++++++++++++------- tests/test_databricks.py | 38 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 55 insertions(+), 9 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 7d979f9c..75dcc85c 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3050,6 +3050,10 @@ class GatewayProbe(NamedTuple): reachable: bool detail: str resource_available: bool = False + # False when the probe reached the API but couldn't reach a verdict on + # accessibility (e.g. a later-page error or the page cap). Distinguishes + # "reachable, confirmed no accessible resource" from "reachable, unknown". + conclusive: bool = True def _version_neutral_gateway_detail(detail: str) -> str: @@ -3111,7 +3115,7 @@ def _probe_ai_gateway_v3(workspace: str, token: str) -> GatewayProbe: ) # A later page failed: reachable, but the listing was not walked to # the end, so we can't claim nothing is accessible. - return GatewayProbe(True, "reachable") + return GatewayProbe(True, "reachable", conclusive=False) if isinstance(payload, dict) and payload.get("model_services"): return GatewayProbe(True, "reachable, accessible model service returned", True) page_token = payload.get("next_page_token") if isinstance(payload, dict) else None @@ -3120,7 +3124,7 @@ def _probe_ai_gateway_v3(workspace: str, token: str) -> GatewayProbe: return GatewayProbe(True, _MODEL_SERVICE_EMPTY_DETAIL) # Hit the page cap with a cursor still pending: reachable but inconclusive, # so don't assert the listing is empty. - return GatewayProbe(True, "reachable") + return GatewayProbe(True, "reachable", conclusive=False) def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: @@ -3177,6 +3181,11 @@ def probe_unity_gateway_capabilities(workspace: str, token: str) -> GatewayProbe legacy_endpoint_probe = _probe_ai_gateway_v2(workspace, token) if legacy_endpoint_probe.reachable: return model_service_probe + # The model-services API answered but the listing couldn't be walked to a + # verdict (later-page error or page cap): it is reachable, hence enabled, so + # don't let the unavailable fallback report the gateway as missing. + if model_service_probe.reachable and not model_service_probe.conclusive: + return model_service_probe if _looks_like_definitive_auth_failure(legacy_endpoint_probe.detail): _raise_ai_gateway_auth_failure(workspace, legacy_endpoint_probe.detail) # A missing-scope 403 that survived the fallback is a token problem, not a @@ -3213,13 +3222,16 @@ def _looks_like_definitive_auth_failure(reason: str) -> bool: def _looks_like_scope_failure(reason: str) -> bool: - """True for a 403 that reports the OAuth token is missing a required scope. + """True for a 403 that reports the OAuth *token* is missing a required scope. - The scopes the model-service and legacy-endpoint APIs require differ, so - this is only conclusive once both probes have failed on it -- re-login (not - a UC grant) is the fix. + Matched to the OAuth-token wording so a PAT's permission 403 -- which + re-login cannot fix -- is not misrouted to the re-login hint and instead + falls through to the grant guidance. The scopes the model-service and + legacy-endpoint APIs require differ, so this is only conclusive once both + probes have failed on it. """ - return "HTTP 403" in reason and "required scopes" in reason.lower() + lowered = reason.lower() + return "http 403" in lowered and "oauth token" in lowered and "required scopes" in lowered def _looks_like_permission_failure(reason: str) -> bool: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 9c0c3e41..b987fb68 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2168,7 +2168,7 @@ def test_probe_v3_later_page_error_is_reachable_not_empty(self, monkeypatch): monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: next(responses)) assert db_mod._probe_ai_gateway_v3(WS, "fake-token") == db_mod.GatewayProbe( - True, "reachable" + True, "reachable", conclusive=False ) def test_probe_v3_page_cap_with_pending_cursor_is_reachable_not_empty(self, monkeypatch): @@ -2183,10 +2183,44 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) assert db_mod._probe_ai_gateway_v3(WS, "fake-token") == db_mod.GatewayProbe( - True, "reachable" + True, "reachable", conclusive=False ) assert len(calls) == db_mod._MODEL_SERVICE_PROBE_MAX_PAGES + def test_inconclusive_model_service_probe_does_not_hard_fail_when_legacy_unavailable( + self, monkeypatch + ): + # A reachable-but-inconclusive v3 result means the API is enabled, so an + # unavailable legacy fallback must not report the gateway as missing. + v3_responses = iter( + [ + ({"next_page_token": "cursor-1"}, None), + (None, "HTTP 500: Internal Server Error"), + ] + ) + + def fake_get(url, token): + if "/api/ai-gateway/v2/endpoints" in url: + return None, "HTTP 404: AI Gateway V2 is not available for CSP-enabled workspaces" + return next(v3_responses) + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + assert db_mod.probe_unity_gateway_capabilities(WS, "fake-token") == db_mod.GatewayProbe( + True, "reachable", conclusive=False + ) + + def test_scope_failure_matches_oauth_token_but_not_pat(self): + # OAuth-token scope 403 -> re-login; a PAT permission 403 (no "OAuth + # token") must not match, so it falls through to the grant guidance. + assert db_mod._looks_like_scope_failure( + "HTTP 403 Forbidden: Provided OAuth token does not have required scopes: unity-catalog" + ) + assert not db_mod._looks_like_scope_failure( + "HTTP 403 Forbidden: Provided access token does not have required scopes" + ) + assert not db_mod._looks_like_scope_failure("HTTP 403: Missing Unity Catalog grants") + def test_model_service_forbidden_and_legacy_unavailable_reports_permission_error( self, monkeypatch ): From 2bac3135ca3093876a13974ff7b32359b7c2d318 Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 20:24:40 +0000 Subject: [PATCH 4/6] Drop inline comments from the probe change Co-authored-by: Isaac --- src/ucode/databricks.py | 17 ----------------- tests/test_databricks.py | 16 ---------------- 2 files changed, 33 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 75dcc85c..7fe3911c 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3050,9 +3050,6 @@ class GatewayProbe(NamedTuple): reachable: bool detail: str resource_available: bool = False - # False when the probe reached the API but couldn't reach a verdict on - # accessibility (e.g. a later-page error or the page cap). Distinguishes - # "reachable, confirmed no accessible resource" from "reachable, unknown". conclusive: bool = True @@ -3087,10 +3084,6 @@ def _probe_ai_gateway_v2(workspace: str, token: str) -> GatewayProbe: ) -# The model-services listing applies `page_size` before ACL filtering, so an -# early page can come back empty (only a `next_page_token`) while accessible -# rows remain further in. Follow the cursor before concluding nothing is -# accessible, or a caller with many services still reads as empty. _MODEL_SERVICE_PROBE_PAGE_SIZE = 50 _MODEL_SERVICE_PROBE_MAX_PAGES = 20 _MODEL_SERVICE_EMPTY_DETAIL = ( @@ -3113,17 +3106,12 @@ def _probe_ai_gateway_v3(workspace: str, token: str) -> GatewayProbe: return GatewayProbe( False, _version_neutral_gateway_detail(reason or "unknown error") ) - # A later page failed: reachable, but the listing was not walked to - # the end, so we can't claim nothing is accessible. return GatewayProbe(True, "reachable", conclusive=False) if isinstance(payload, dict) and payload.get("model_services"): return GatewayProbe(True, "reachable, accessible model service returned", True) page_token = payload.get("next_page_token") if isinstance(payload, dict) else None if not page_token: - # Walked the whole listing without an accessible model service. return GatewayProbe(True, _MODEL_SERVICE_EMPTY_DETAIL) - # Hit the page cap with a cursor still pending: reachable but inconclusive, - # so don't assert the listing is empty. return GatewayProbe(True, "reachable", conclusive=False) @@ -3181,15 +3169,10 @@ def probe_unity_gateway_capabilities(workspace: str, token: str) -> GatewayProbe legacy_endpoint_probe = _probe_ai_gateway_v2(workspace, token) if legacy_endpoint_probe.reachable: return model_service_probe - # The model-services API answered but the listing couldn't be walked to a - # verdict (later-page error or page cap): it is reachable, hence enabled, so - # don't let the unavailable fallback report the gateway as missing. if model_service_probe.reachable and not model_service_probe.conclusive: return model_service_probe if _looks_like_definitive_auth_failure(legacy_endpoint_probe.detail): _raise_ai_gateway_auth_failure(workspace, legacy_endpoint_probe.detail) - # A missing-scope 403 that survived the fallback is a token problem, not a - # UC grant, so surface re-login guidance before the permission hints. if _looks_like_scope_failure(model_service_probe.detail): _raise_ai_gateway_scope_failure(workspace, model_service_probe.detail) if _looks_like_scope_failure(legacy_endpoint_probe.detail): diff --git a/tests/test_databricks.py b/tests/test_databricks.py index b987fb68..836e9f8f 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1935,9 +1935,6 @@ def fake_get(url, token): ] def test_empty_first_page_follows_cursor_to_accessible_model_service(self, monkeypatch): - # UC applies page_size before ACL filtering, so an early page can come - # back empty (only a next_page_token) while accessible rows remain behind - # the cursor (ES-2185388). The probe must page before reporting empty. calls: list[str] = [] responses = iter( [ @@ -2105,9 +2102,6 @@ def fake_get(url, token): assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=50"] def test_missing_scope_403_on_both_paths_routes_to_reauth_not_grants(self, monkeypatch): - # A missing-scope 403 is only conclusive once the legacy fallback also - # fails on it: the fix is re-login, not a UC grant. It must probe the - # fallback first and must not reach the grant-hint message. calls: list[str] = [] def fake_get(url, token): @@ -2132,8 +2126,6 @@ def fake_get(url, token): ] def test_model_service_scope_403_succeeds_when_legacy_reachable(self, monkeypatch): - # A token scoped for the legacy endpoint but not for model services must - # still succeed via the fallback rather than surfacing the scope 403. calls: list[str] = [] def fake_get(url, token): @@ -2157,8 +2149,6 @@ def fake_get(url, token): ] def test_probe_v3_later_page_error_is_reachable_not_empty(self, monkeypatch): - # A later-page error leaves the listing un-walked, so the probe must not - # conclude nothing is accessible. responses = iter( [ ({"next_page_token": "cursor-1"}, None), @@ -2172,8 +2162,6 @@ def test_probe_v3_later_page_error_is_reachable_not_empty(self, monkeypatch): ) def test_probe_v3_page_cap_with_pending_cursor_is_reachable_not_empty(self, monkeypatch): - # Exhausting the page cap while a cursor is still pending is inconclusive, - # not empty. calls: list[str] = [] def fake_get(url, token): @@ -2190,8 +2178,6 @@ def fake_get(url, token): def test_inconclusive_model_service_probe_does_not_hard_fail_when_legacy_unavailable( self, monkeypatch ): - # A reachable-but-inconclusive v3 result means the API is enabled, so an - # unavailable legacy fallback must not report the gateway as missing. v3_responses = iter( [ ({"next_page_token": "cursor-1"}, None), @@ -2211,8 +2197,6 @@ def fake_get(url, token): ) def test_scope_failure_matches_oauth_token_but_not_pat(self): - # OAuth-token scope 403 -> re-login; a PAT permission 403 (no "OAuth - # token") must not match, so it falls through to the grant guidance. assert db_mod._looks_like_scope_failure( "HTTP 403 Forbidden: Provided OAuth token does not have required scopes: unity-catalog" ) From 001ece316316dfaf889234f75ee716e4a1ff2145 Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 21:04:35 +0000 Subject: [PATCH 5/6] Stay quiet on the gateway happy path When the model-service probe finds an accessible model service, configure now prints nothing about the gateway and moves on. It emits a single warning line only when no model service was detected but the gateway is still usable (or the result is inconclusive); a hard failure still raises. Drops the success header that would otherwise appear only when something is off. Co-authored-by: Isaac --- src/ucode/cli.py | 4 ++-- tests/test_cli.py | 15 +++++---------- 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6134c9f0..9e34c00e 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -637,8 +637,8 @@ def configure_shared_state( with spinner("Verifying Unity AI Gateway..."): token = get_databricks_token(workspace, profile) model_service_probe = probe_unity_gateway_capabilities(workspace, token) - print_success("Unity AI Gateway detected") - print_kv("Model service", model_service_probe.detail) + if not model_service_probe.resource_available: + print_warning(f"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/tests/test_cli.py b/tests/test_cli.py index 3f201b43..3788c583 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2495,22 +2495,18 @@ def test_use_pat_exports_bearer_and_skips_login(self, monkeypatch): assert state["use_pat"] is True assert saved and saved[-1]["use_pat"] is True - def test_prints_model_service_and_omits_unneeded_legacy_probe(self, monkeypatch, capsys): + def test_happy_path_prints_no_gateway_output(self, monkeypatch, capsys): cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") cli_mod.configure_shared_state(self.WS, profile="DEFAULT") output = _strip_ansi(capsys.readouterr().out) - assert "Model service: reachable, accessible model service returned" in output - assert "(Legacy) endpoints:" not in output + assert "Model service:" not in output + assert "Unity AI Gateway detected" not in output @pytest.mark.parametrize( ("responses", "expected_model_service"), [ - ( - [({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None)], - "reachable, accessible model service returned", - ), ( [({}, None), ({"endpoints": []}, None)], "reachable, no accessible model services returned; check USE CATALOG on system, " @@ -2525,12 +2521,11 @@ def test_prints_model_service_and_omits_unneeded_legacy_probe(self, monkeypatch, ), ], ids=[ - "model-service-resource", "model-service-empty-legacy-empty", "model-service-forbidden-legacy-resource", ], ) - def test_prints_local_gateway_probe_scenarios( + def test_prints_warning_when_model_service_not_detected( self, monkeypatch, capsys, @@ -2547,8 +2542,8 @@ def test_prints_local_gateway_probe_scenarios( cli_mod.configure_shared_state(self.WS, profile="DEFAULT") output = " ".join(_strip_ansi(capsys.readouterr().out).split()) - assert "Unity AI Gateway detected" in output assert f"Model service: {expected_model_service}" in output + assert "Unity AI Gateway detected" not in output assert "(Legacy) endpoints:" not in output assert "V2" not in output assert "V3" not in output From 5ac5efef4c902d189e3bac57aef0f1235a3addef Mon Sep 17 00:00:00 2001 From: David Liu Date: Wed, 2 Sep 2026 21:18:27 +0000 Subject: [PATCH 6/6] Close the gateway spinner with a success line on the happy path Per review: when a model service is accessible, print "Unity AI Gateway connected" as the spinner's success rather than nothing, while still dropping the redundant model-service detail line. The warning path is unchanged. Co-authored-by: Isaac --- src/ucode/cli.py | 4 +++- tests/test_cli.py | 8 ++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 9e34c00e..a8711ead 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -637,7 +637,9 @@ def configure_shared_state( with spinner("Verifying Unity AI Gateway..."): token = get_databricks_token(workspace, profile) model_service_probe = probe_unity_gateway_capabilities(workspace, token) - if not model_service_probe.resource_available: + if model_service_probe.resource_available: + print_success("Unity AI Gateway connected") + else: print_warning(f"Model service: {model_service_probe.detail}") want_claude = ( diff --git a/tests/test_cli.py b/tests/test_cli.py index 3788c583..d42c6a53 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2495,14 +2495,14 @@ 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_happy_path_prints_no_gateway_output(self, monkeypatch, capsys): + def test_happy_path_prints_success_without_model_service_detail(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 "Unity AI Gateway connected" in output assert "Model service:" not in output - assert "Unity AI Gateway detected" not in output @pytest.mark.parametrize( ("responses", "expected_model_service"), @@ -2543,7 +2543,7 @@ def test_prints_warning_when_model_service_not_detected( output = " ".join(_strip_ansi(capsys.readouterr().out).split()) assert f"Model service: {expected_model_service}" in output - assert "Unity AI Gateway detected" not in output + assert "Unity AI Gateway connected" not in output assert "(Legacy) endpoints:" not in output assert "V2" not in output assert "V3" not in output @@ -2591,7 +2591,7 @@ def test_local_gateway_probe_failures_do_not_print_success( cli_mod.configure_shared_state(self.WS, profile="DEFAULT") output = _strip_ansi(capsys.readouterr().out) - assert "Unity AI Gateway detected" not in output + assert "Unity AI Gateway connected" not in output message = str(excinfo.value) assert "v2" not in message.lower() assert "v3" not in message.lower()