Disambiguate AI Gateway probe failures by status code

src/ucode/databricks.py

@@ -221,7 +221,13 @@ def _http_get_json(
         _debug(f"GET {url}", f"HTTP {exc.code} {exc.reason}")
         if _debug_enabled() and body:
             _debug("body", body[:4000])
-        return None, f"HTTP {exc.code} {exc.reason}"
+        reason = f"HTTP {exc.code} {exc.reason}"
+        # Surface the response body too — gateway auth failures return 400
+        # with body `Invalid Token`, which is invisible without this.
+        body_excerpt = body.strip()[:200]
+        if body_excerpt:
+            reason = f"{reason}: {body_excerpt}"
+        return None, reason
     except urllib_error.URLError as exc:
         _debug(f"GET {url}", f"URLError: {exc.reason}")
         return None, f"network error: {exc.reason}"
@@ -903,20 +909,50 @@ def ensure_ai_gateway_v2(workspace: str, token: str) -> None:
     Uses the dedicated v2 listing endpoint `GET /api/ai-gateway/v2/endpoints`:
     a 200 response (even with an empty list) means v2 is wired up on this
     workspace — a "no endpoints provisioned" case will surface naturally in
-    downstream discovery. 404 / 401 / 403 / network failures all raise a
-    clear error with the docs link instead of silently progressing.
+    downstream discovery. Failure branches:
+
+    - 401 / 403 / 400 with `Invalid Token`: the token is bad for *this*
+      workspace.
+    - 404: AI Gateway V2 is not enabled on this workspace — point at the docs.
+    - other (5xx, network errors): surface the reason verbatim.
     """
     hostname = workspace_hostname(workspace)
     url = f"https://{hostname}/api/ai-gateway/v2/endpoints?page_size=1"
     payload, reason = _http_get_json(url, token)
     if payload is not None:
         return
+    reason_str = reason or "unknown error"
+    if _looks_like_auth_failure(reason_str):
+        raise RuntimeError(
+            f"Databricks rejected the access token for {workspace} ({reason_str}). "
+            f"Try:\n"
+            f"  databricks auth logout --host {workspace}\n"
+            f"  databricks auth login --host {workspace}"
+        )
+    if "HTTP 404" in reason_str:
+        raise RuntimeError(
+            "Databricks Unity AI Gateway is not enabled on this workspace "
+            f"({reason_str}). See {AI_GATEWAY_V2_DOCS_URL}"
+        )
     raise RuntimeError(
-        "Databricks AI Gateway V2 is required but not available on this workspace "
-        f"({reason}). See {AI_GATEWAY_V2_DOCS_URL}"
+        "Databricks Unity AI Gateway probe failed on this workspace "
+        f"({reason_str}). See {AI_GATEWAY_V2_DOCS_URL}"
     )
 
 
+def _looks_like_auth_failure(reason: str) -> bool:
+    """True when the gateway response signals the token is not accepted.
+
+    Covers 401/403 directly and the gateway's 400 + `Invalid Token` body
+    (which happens when the bearer is valid but issued for a different
+    workspace)."""
+    if "HTTP 401" in reason or "HTTP 403" in reason:
+        return True
+    if "HTTP 400" in reason and "invalid token" in reason.lower():
+        return True
+    return False
+
+
 def discover_sql_warehouse_http_path(
     workspace: str,
     token: str,

tests/test_databricks.py

@@ -501,27 +501,68 @@ def _mock_json_response(body: str):
         mock_resp.read.return_value = body.encode("utf-8")
         return mock_resp
 
-    def test_raises_on_404(self):
-        from unittest.mock import MagicMock, patch
+    @staticmethod
+    def _http_error(code: int, msg: str, body: str = ""):
+        import io
+        from unittest.mock import MagicMock
         from urllib.error import HTTPError
 
-        exc = HTTPError(url="", code=404, msg="Not Found", hdrs=MagicMock(), fp=None)
+        fp = io.BytesIO(body.encode("utf-8")) if body else None
+        return HTTPError(url="", code=code, msg=msg, hdrs=MagicMock(), fp=fp)
+
+    def test_raises_on_404(self):
+        from unittest.mock import patch
+
+        exc = self._http_error(404, "Not Found")
         with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc):
             from ucode.databricks import ensure_ai_gateway_v2
 
-            with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL):
+            with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL) as excinfo:
                 ensure_ai_gateway_v2(WS, "fake-token")
+            assert "not enabled" in str(excinfo.value)
 
-    def test_raises_on_401(self):
-        from unittest.mock import MagicMock, patch
-        from urllib.error import HTTPError
+    def test_raises_on_401_with_auth_hint(self):
+        from unittest.mock import patch
+
+        exc = self._http_error(401, "Unauthorized")
+        with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc):
+            from ucode.databricks import ensure_ai_gateway_v2
+
+            with pytest.raises(RuntimeError, match="401") as excinfo:
+                ensure_ai_gateway_v2(WS, "fake-token")
+            message = str(excinfo.value)
+            assert "rejected" in message.lower()
+            assert "databricks auth login" in message
+
+    def test_raises_on_400_invalid_token_with_auth_hint(self):
+        """400 + body `Invalid Token` is the misleading-error case from issue #84."""
+        from unittest.mock import patch
+
+        exc = self._http_error(400, "Bad Request", body="Invalid Token")
+        with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc):
+            from ucode.databricks import ensure_ai_gateway_v2
 
-        exc = HTTPError(url="", code=401, msg="Unauthorized", hdrs=MagicMock(), fp=None)
+            with pytest.raises(RuntimeError) as excinfo:
+                ensure_ai_gateway_v2(WS, "fake-token")
+            message = str(excinfo.value)
+            # The bug we are fixing: must NOT collapse to the generic
+            # "v2 not available" message — must call out the auth failure
+            # and point at re-login.
+            assert "Invalid Token" in message
+            assert "rejected" in message.lower()
+            assert "databricks auth login" in message
+
+    def test_400_without_invalid_token_falls_through_to_generic(self):
+        """A 400 that is *not* an auth failure should still surface the body."""
+        from unittest.mock import patch
+
+        exc = self._http_error(400, "Bad Request", body="some other detail")
         with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc):
             from ucode.databricks import ensure_ai_gateway_v2
 
-            with pytest.raises(RuntimeError, match="401"):
+            with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL) as excinfo:
                 ensure_ai_gateway_v2(WS, "fake-token")
+            assert "some other detail" in str(excinfo.value)
 
     def test_raises_on_url_error(self):
         from unittest.mock import patch
@@ -561,6 +602,44 @@ def test_succeeds_with_empty_endpoints_list(self):
             ensure_ai_gateway_v2(WS, "fake-token")  # should not raise
 
 
+class TestHttpGetJsonReason:
+    """The `reason` string returned by `_http_get_json` must include the response body
+    so callers (e.g. ensure_ai_gateway_v2) can route on it. Before issue #84's fix
+    the body was logged only when UCODE_DEBUG=1 and dropped from the bubbled error."""
+
+    @staticmethod
+    def _http_error(code: int, msg: str, body: str = ""):
+        import io
+        from unittest.mock import MagicMock
+        from urllib.error import HTTPError
+
+        fp = io.BytesIO(body.encode("utf-8")) if body else None
+        return HTTPError(url="", code=code, msg=msg, hdrs=MagicMock(), fp=fp)
+
+    def test_reason_includes_body_on_http_error(self):
+        from unittest.mock import patch
+
+        from ucode.databricks import _http_get_json
+
+        exc = self._http_error(400, "Bad Request", body="Invalid Token")
+        with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc):
+            payload, reason = _http_get_json("https://x/y", "tok")
+        assert payload is None
+        assert "HTTP 400" in reason
+        assert "Invalid Token" in reason
+
+    def test_reason_without_body_is_status_only(self):
+        from unittest.mock import patch
+
+        from ucode.databricks import _http_get_json
+
+        exc = self._http_error(404, "Not Found")
+        with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc):
+            payload, reason = _http_get_json("https://x/y", "tok")
+        assert payload is None
+        assert reason == "HTTP 404 Not Found"
+
+
 class TestParseDatabricksCliVersion:
     def test_parses_standard_format(self):
         assert _parse_databricks_cli_version("Databricks CLI v0.299.2") == (0, 299, 2)