diff --git a/src/ucode/cli.py b/src/ucode/cli.py index abe73f57..826a6f1f 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -637,8 +637,10 @@ 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 model_service_probe.resource_available: + print_success("Unity AI Gateway connected") + else: + 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/src/ucode/databricks.py b/src/ucode/databricks.py index 33e36c2e..7fe3911c 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3050,6 +3050,7 @@ class GatewayProbe(NamedTuple): reachable: bool detail: str resource_available: bool = False + conclusive: bool = True def _version_neutral_gateway_detail(detail: str) -> str: @@ -3062,17 +3063,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 +3084,35 @@ def _probe_ai_gateway_v2(workspace: str, token: str) -> GatewayProbe: ) +_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") + ) + 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: + return GatewayProbe(True, _MODEL_SERVICE_EMPTY_DETAIL) + return GatewayProbe(True, "reachable", conclusive=False) def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: @@ -3109,6 +3124,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: @@ -3145,8 +3169,14 @@ 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 + 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) + 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 @@ -3174,6 +3204,19 @@ def _looks_like_definitive_auth_failure(reason: str) -> bool: 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. + + 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. + """ + 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: return "HTTP 403" in reason diff --git a/tests/test_cli.py b/tests/test_cli.py index dca05bcd..2a51231e 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2504,22 +2504,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_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 "Model service: reachable, accessible model service returned" in output - assert "(Legacy) endpoints:" not in output + assert "Unity AI Gateway connected" in output + assert "Model service:" 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, " @@ -2534,12 +2530,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, @@ -2556,8 +2551,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 connected" not in output assert "(Legacy) endpoints:" not in output assert "V2" not in output assert "V3" not in output @@ -2605,7 +2600,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() diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 0df389cc..836e9f8f 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1931,7 +1931,64 @@ 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): + 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 +2021,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 +2040,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 +2099,111 @@ 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_on_both_paths_routes_to_reauth_not_grants(self, monkeypatch): + 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="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", + f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", + ] + + def test_model_service_scope_403_succeeds_when_legacy_reachable(self, monkeypatch): + 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): + 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", conclusive=False + ) + + def test_probe_v3_page_cap_with_pending_cursor_is_reachable_not_empty(self, monkeypatch): + 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", 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 + ): + 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): + 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