fix(chat): Heroku-safe Gemini retries, 429 JSON, and tests

northdpole · northdpole · commit 658ef190f1ef · 2026-04-21T18:01:26.000+01:00
- Use shorter default retries/sleeps on Heroku web dynos (DYNO=web.*) so
  requests finish before the ~30s router timeout; keep env overrides.
- Log Gemini ClientError details (http_code, api status, message, details)
  for debugging 429s without exposing secrets.
- Treat google.genai ClientError HTTP status via .code (not .status_code)
  and return 503 JSON for rate limits from /rest/v1/completion.
- Add unit tests for retry helpers and completion route behavior.

Made-with: Cursor
diff --git a/application/prompt_client/vertex_prompt_client.py b/application/prompt_client/vertex_prompt_client.py
@@ -42,6 +42,7 @@ def _is_genai_rate_limit_error(err: Exception) -> bool:
     if status_code == 429:
         return True
 
+    # google.genai.errors.ClientError / APIError use ``code`` (HTTP status).
     code = getattr(err, "code", None)
     if code == 429:
         return True
@@ -57,6 +58,50 @@ def _is_genai_rate_limit_error(err: Exception) -> bool:
     return False
 
 
+def _log_genai_client_error(context: str, exc: genai.errors.ClientError) -> None:
+    """Log full Gemini error payload (no secrets) for support and debugging."""
+    logger.warning(
+        "Gemini ClientError [%s]: http_code=%s api_status=%s message=%s details=%r",
+        context,
+        getattr(exc, "code", None),
+        getattr(exc, "status", None),
+        getattr(exc, "message", None),
+        getattr(exc, "details", None),
+    )
+
+
+def _is_heroku_web_dyno() -> bool:
+    """Heroku sets DYNO=web.* on HTTP dynos; the router enforces a ~30s request window."""
+    dyno = os.environ.get("DYNO", "")
+    return dyno.startswith("web.")
+
+
+def _effective_gemini_generate_retry_settings() -> tuple[int, int]:
+    """(max_retries, sleep_seconds) for generate_content; env vars always override."""
+    if _is_heroku_web_dyno():
+        return (
+            int(os.environ.get("GEMINI_GENERATE_MAX_RETRIES", "1")),
+            int(os.environ.get("GEMINI_GENERATE_RETRY_SLEEP_SECONDS", "6")),
+        )
+    return (
+        int(os.environ.get("GEMINI_GENERATE_MAX_RETRIES", "3")),
+        int(os.environ.get("GEMINI_GENERATE_RETRY_SLEEP_SECONDS", "60")),
+    )
+
+
+def _effective_vertex_embed_retry_settings() -> tuple[int, int]:
+    """(max_retries, sleep_seconds) for embed_content; env vars always override."""
+    if _is_heroku_web_dyno():
+        return (
+            int(os.environ.get("VERTEX_EMBED_MAX_RETRIES", "1")),
+            int(os.environ.get("VERTEX_EMBED_RETRY_SLEEP_SECONDS", "6")),
+        )
+    return (
+        int(os.environ.get("VERTEX_EMBED_MAX_RETRIES", "3")),
+        int(os.environ.get("VERTEX_EMBED_RETRY_SLEEP_SECONDS", "60")),
+    )
+
+
 class VertexPromptClient:
     context = (
         'You are "OpenCRE Chat" a chatbot for security information that exists in opencre.org. '
@@ -96,17 +141,17 @@ def _with_genai_rate_limit_retry(
         Bounded retries for `generate_content` (SDK may retry briefly; this adds
         longer backoff for sustained quota pressure).
 
-        Configure via ``GEMINI_GENERATE_MAX_RETRIES`` (default 3) and
-        ``GEMINI_GENERATE_RETRY_SLEEP_SECONDS`` (default 60).
+        Configure via ``GEMINI_GENERATE_MAX_RETRIES`` and
+        ``GEMINI_GENERATE_RETRY_SLEEP_SECONDS``. On Heroku ``web.*`` dynos the
+        defaults are reduced so the request can finish before the ~30s router
+        timeout (long sleeps used to abort the Gunicorn worker mid-request).
         """
-        max_retries = int(os.environ.get("GEMINI_GENERATE_MAX_RETRIES", "3"))
-        retry_sleep_seconds = int(
-            os.environ.get("GEMINI_GENERATE_RETRY_SLEEP_SECONDS", "60")
-        )
+        max_retries, retry_sleep_seconds = _effective_gemini_generate_retry_settings()
         for attempt in range(max_retries + 1):
             try:
                 return fn()
             except genai.errors.ClientError as e:
+                _log_genai_client_error(context, e)
                 if not _is_genai_rate_limit_error(e) or attempt >= max_retries:
                     raise
                 logger.info(
@@ -150,10 +195,7 @@ def _truncate_one(t: str) -> str:
         texts: List[str] = text if is_batch else [_truncate_one(text)]  # type: ignore[arg-type]
         texts = [_truncate_one(t) for t in texts]
 
-        max_retries = int(os.environ.get("VERTEX_EMBED_MAX_RETRIES", "3"))
-        retry_sleep_seconds = int(
-            os.environ.get("VERTEX_EMBED_RETRY_SLEEP_SECONDS", "60")
-        )
+        max_retries, retry_sleep_seconds = _effective_vertex_embed_retry_settings()
 
         for attempt in range(max_retries + 1):
             try:
@@ -173,6 +215,7 @@ def _truncate_one(t: str) -> str:
                     return [emb.values for emb in result.embeddings]
                 return result.embeddings[0].values
             except genai.errors.ClientError as e:
+                _log_genai_client_error("embed_content", e)
                 if not _is_genai_rate_limit_error(e) or attempt >= max_retries:
                     raise
                 logger.info(
diff --git a/application/tests/chat_completion_test.py b/application/tests/chat_completion_test.py
@@ -0,0 +1,77 @@
+"""Tests for /rest/v1/completion (chatbot) error handling."""
+
+import json
+import os
+import unittest
+from unittest.mock import patch
+
+from google.genai import errors as genai_errors
+
+from application import create_app, sqla  # type: ignore
+
+
+class TestChatCompletion(unittest.TestCase):
+    def tearDown(self) -> None:
+        sqla.session.remove()
+        sqla.drop_all()
+        self.app_context.pop()
+        os.environ.pop("NO_LOGIN", None)
+
+    def setUp(self) -> None:
+        self.app = create_app(mode="test")
+        self.app_context = self.app.app_context()
+        self.app_context.push()
+        os.environ["INSECURE_REQUESTS"] = "True"
+        sqla.create_all()
+
+    def test_completion_returns_503_json_on_gemini_429(self) -> None:
+        os.environ["NO_LOGIN"] = "1"
+        err = genai_errors.ClientError(
+            429,
+            {
+                "error": {
+                    "code": 429,
+                    "message": "Resource exhausted",
+                    "status": "RESOURCE_EXHAUSTED",
+                }
+            },
+            None,
+        )
+        with patch(
+            "application.prompt_client.prompt_client.PromptHandler"
+        ) as mock_ph:
+            mock_ph.return_value.generate_text.side_effect = err
+            with self.app.test_client() as client:
+                response = client.post(
+                    "/rest/v1/completion",
+                    json={"prompt": "test"},
+                    content_type="application/json",
+                )
+        self.assertEqual(503, response.status_code)
+        data = json.loads(response.data)
+        self.assertIn("error", data)
+        self.assertIn("rate-limited", data["error"])
+
+    def test_completion_propagates_non_429_genai_error(self) -> None:
+        os.environ["NO_LOGIN"] = "1"
+        err = genai_errors.ClientError(
+            400,
+            {"error": {"code": 400, "message": "Bad request", "status": "INVALID_ARGUMENT"}},
+            None,
+        )
+        with patch(
+            "application.prompt_client.prompt_client.PromptHandler"
+        ) as mock_ph:
+            mock_ph.return_value.generate_text.side_effect = err
+            with self.app.test_client() as client:
+                with self.assertRaises(genai_errors.ClientError) as ctx:
+                    client.post(
+                        "/rest/v1/completion",
+                        json={"prompt": "test"},
+                        content_type="application/json",
+                    )
+                self.assertIs(ctx.exception, err)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/application/tests/vertex_prompt_client_test.py b/application/tests/vertex_prompt_client_test.py
@@ -0,0 +1,82 @@
+"""Tests for Gemini retry helpers and rate-limit detection."""
+
+import os
+import unittest
+
+from google.genai import errors as genai_errors
+
+from application.prompt_client.vertex_prompt_client import (
+    _effective_gemini_generate_retry_settings,
+    _effective_vertex_embed_retry_settings,
+    _is_genai_rate_limit_error,
+    _is_heroku_web_dyno,
+)
+
+
+class TestVertexPromptClientHelpers(unittest.TestCase):
+    def tearDown(self) -> None:
+        for key in (
+            "DYNO",
+            "GEMINI_GENERATE_MAX_RETRIES",
+            "GEMINI_GENERATE_RETRY_SLEEP_SECONDS",
+            "VERTEX_EMBED_MAX_RETRIES",
+            "VERTEX_EMBED_RETRY_SLEEP_SECONDS",
+        ):
+            os.environ.pop(key, None)
+
+    def test_is_heroku_web_dyno_true_for_web_prefix(self) -> None:
+        os.environ["DYNO"] = "web.1"
+        self.assertTrue(_is_heroku_web_dyno())
+
+    def test_is_heroku_web_dyno_false_for_worker(self) -> None:
+        os.environ["DYNO"] = "worker.1"
+        self.assertFalse(_is_heroku_web_dyno())
+
+    def test_is_heroku_web_dyno_false_when_unset(self) -> None:
+        self.assertFalse(_is_heroku_web_dyno())
+
+    def test_effective_gemini_retry_defaults_local(self) -> None:
+        self.assertEqual(_effective_gemini_generate_retry_settings(), (3, 60))
+
+    def test_effective_gemini_retry_defaults_heroku_web(self) -> None:
+        os.environ["DYNO"] = "web.1"
+        self.assertEqual(_effective_gemini_generate_retry_settings(), (1, 6))
+
+    def test_effective_gemini_retry_env_overrides_heroku(self) -> None:
+        os.environ["DYNO"] = "web.1"
+        os.environ["GEMINI_GENERATE_MAX_RETRIES"] = "2"
+        os.environ["GEMINI_GENERATE_RETRY_SLEEP_SECONDS"] = "10"
+        self.assertEqual(_effective_gemini_generate_retry_settings(), (2, 10))
+
+    def test_effective_embed_retry_defaults_local(self) -> None:
+        self.assertEqual(_effective_vertex_embed_retry_settings(), (3, 60))
+
+    def test_effective_embed_retry_defaults_heroku_web(self) -> None:
+        os.environ["DYNO"] = "web.1"
+        self.assertEqual(_effective_vertex_embed_retry_settings(), (1, 6))
+
+    def test_is_genai_rate_limit_error_recognizes_clienterror_code_429(self) -> None:
+        err = genai_errors.ClientError(
+            429,
+            {
+                "error": {
+                    "code": 429,
+                    "message": "Resource exhausted",
+                    "status": "RESOURCE_EXHAUSTED",
+                }
+            },
+            None,
+        )
+        self.assertTrue(_is_genai_rate_limit_error(err))
+
+    def test_is_genai_rate_limit_error_false_for_other_clienterror(self) -> None:
+        err = genai_errors.ClientError(
+            400,
+            {"error": {"code": 400, "message": "Bad request", "status": "INVALID_ARGUMENT"}},
+            None,
+        )
+        self.assertFalse(_is_genai_rate_limit_error(err))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/application/web/web_main.py b/application/web/web_main.py
@@ -918,10 +918,24 @@ def chat_cre() -> Any:
 
     database = db.Node_collection()
     # Lazy import to avoid loading heavy prompt/ML dependencies at web boot.
+    from google.genai import errors as genai_errors
     from application.prompt_client import prompt_client
 
     prompt = prompt_client.PromptHandler(database)
-    response = prompt.generate_text(message.get("prompt"))
+    try:
+        response = prompt.generate_text(message.get("prompt"))
+    except genai_errors.ClientError as e:
+        # google.genai APIError uses ``code`` (HTTP status), not ``status_code``.
+        if getattr(e, "code", None) == 429:
+            return jsonify(
+                {
+                    "error": (
+                        "The AI service is temporarily rate-limited. "
+                        "Please try again in a minute."
+                    )
+                }
+            ), 503
+        raise
     return jsonify(response)
 
 
