fix(prompts): address review on cooldown edge cases

Author: jurajmajerik

posthog/ai/prompts.py

@@ -14,6 +14,7 @@
 from dataclasses import dataclass
 from typing import Any, Dict, List, Literal, Optional, Union, overload
 
+from posthog.capture_v1 import _parse_retry_after
 from posthog.request import USER_AGENT, _get_session
 from posthog.utils import remove_trailing_slash
 
@@ -160,14 +161,9 @@ def _is_same_origin(url: str, host: str) -> bool:
 
 
 def _parse_retry_after_seconds(value: Optional[str]) -> Optional[float]:
-    """Read a Retry-After header of the delta-seconds form the API sends."""
-    if not value:
-        return None
-    try:
-        seconds = float(value.strip())
-    except ValueError:
-        return None
-    if not math.isfinite(seconds) or seconds <= 0:
+    """Read a Retry-After header (delta-seconds or HTTP-date) as a bounded cooldown."""
+    seconds = _parse_retry_after(value)
+    if seconds is None or not math.isfinite(seconds) or seconds <= 0:
         return None
     return min(seconds, MAX_REFETCH_COOLDOWN_SECONDS)
 
@@ -554,8 +550,13 @@ def _get_internal(
 
             # A failed refetch left this entry in cooldown. Serving it keeps one
             # throttled client from turning every later get() into another
-            # request, which is what holds it against the limit.
-            if cached.retry_not_before is not None and now < cached.retry_not_before:
+            # request, which is what holds it against the limit. A zero TTL is
+            # an explicit request to refetch on every read, so it opts out.
+            if (
+                ttl > 0
+                and cached.retry_not_before is not None
+                and now < cached.retry_not_before
+            ):
                 return PromptResult(
                     source="stale_cache",
                     prompt=cached.prompt,

posthog/test/ai/test_prompts.py

@@ -325,75 +325,62 @@ def test_use_stale_cache_on_fetch_failure_with_warning(
         warning_call = mock_log.warning.call_args
         self.assertIn("using stale cache", warning_call[0][0])
 
+    @parameterized.expand(
+        [
+            # The Retry-After cooldown (300s) outlives the 60s default, so the
+            # holds after second 1460 prove the server's value governs.
+            ("server_error", 500, None, 60.0),
+            ("rate_limited_retry_after", 429, {"Retry-After": "300"}, 300.0),
+        ]
+    )
     @patch("posthog.ai.prompts._get_session")
     @patch("posthog.ai.prompts.time.time")
-    def test_hold_a_cooldown_after_a_failed_refetch_then_retry(
-        self, mock_time, mock_get_session
+    def test_hold_a_cooldown_after_each_failed_refetch(
+        self, _scenario, status, headers, cooldown, mock_time, mock_get_session
     ):
         # Without the cooldown, one throttled client turns every later get()
         # into another network request until one succeeds.
         mock_get = mock_get_session.return_value.get
         mock_get.side_effect = [
             MockResponse(json_data=self.mock_prompt_response),
-            MockResponse(status_code=500, ok=False),
+            MockResponse(status_code=status, ok=False, headers=headers),
+            MockResponse(status_code=status, ok=False, headers=headers),
             MockResponse(json_data=self.mock_prompt_response),
         ]
         mock_time.return_value = 1000.0
 
         prompts = Prompts(self.create_mock_posthog())
         prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=False)
 
-        # Past TTL: the refetch fails, stale cache is served, cooldown starts.
+        # Past TTL: the refetch fails, stale cache is served, a cooldown starts.
         mock_time.return_value = 1400.0
         result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True)
         self.assertEqual(result.source, "stale_cache")
         self.assertEqual(mock_get.call_count, 2)
 
-        # Within the 60s cooldown: stale cache again, no network attempt.
-        mock_time.return_value = 1430.0
+        # Within the cooldown: stale cache again, no network attempt.
+        mock_time.return_value = 1400.0 + cooldown - 1
         result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True)
         self.assertEqual(result.source, "stale_cache")
         self.assertEqual(mock_get.call_count, 2)
 
-        # Past the cooldown: the network is retried and the cache refreshed.
-        mock_time.return_value = 1470.0
-        result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True)
-        self.assertEqual(result.source, "api")
-        self.assertEqual(mock_get.call_count, 3)
-
-    @patch("posthog.ai.prompts._get_session")
-    @patch("posthog.ai.prompts.time.time")
-    def test_hold_the_cooldown_for_the_retry_after_a_429_sends(
-        self, mock_time, mock_get_session
-    ):
-        # A rate-limited client must wait as long as the server told it to,
-        # not the 60s default.
-        mock_get = mock_get_session.return_value.get
-        mock_get.side_effect = [
-            MockResponse(json_data=self.mock_prompt_response),
-            MockResponse(status_code=429, ok=False, headers={"Retry-After": "300"}),
-            MockResponse(json_data=self.mock_prompt_response),
-        ]
-        mock_time.return_value = 1000.0
-
-        prompts = Prompts(self.create_mock_posthog())
-        prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=False)
-
-        mock_time.return_value = 1400.0
+        # Past the cooldown: the retry fails too and a new cooldown starts.
+        mock_time.return_value = 1400.0 + cooldown + 1
         result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True)
         self.assertEqual(result.source, "stale_cache")
+        self.assertEqual(mock_get.call_count, 3)
 
-        # Past the 60s default but within Retry-After: no network attempt.
-        mock_time.return_value = 1600.0
+        # Within the second cooldown: no network attempt.
+        mock_time.return_value = 1400.0 + 2 * cooldown
         result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True)
         self.assertEqual(result.source, "stale_cache")
-        self.assertEqual(mock_get.call_count, 2)
+        self.assertEqual(mock_get.call_count, 3)
 
-        # Past Retry-After: the network is retried.
-        mock_time.return_value = 1701.0
+        # Past the second cooldown: the network is retried and the cache refreshed.
+        mock_time.return_value = 1400.0 + 2 * cooldown + 2
         result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True)
         self.assertEqual(result.source, "api")
-        self.assertEqual(mock_get.call_count, 3)
+        self.assertEqual(mock_get.call_count, 4)
 
     @patch("posthog.ai.prompts._get_session")
     @patch("posthog.ai.prompts.log")
@@ -619,6 +606,18 @@ def test_default_cache_ttl_seconds_zero_disables_caching(
         prompts.get("test-prompt", with_metadata=False)
         self.assertEqual(mock_get.call_count, 2)
 
+        # A failed refetch must not start a cooldown here: a zero TTL is an
+        # explicit request to refetch on every read.
+        mock_get.side_effect = [
+            MockResponse(status_code=500, ok=False),
+            MockResponse(json_data=self.mock_prompt_response),
+        ]
+        result = prompts.get("test-prompt", with_metadata=True)
+        self.assertEqual(result.source, "stale_cache")
+        result = prompts.get("test-prompt", with_metadata=True)
+        self.assertEqual(result.source, "api")
+        self.assertEqual(mock_get.call_count, 4)
+
     @patch("posthog.ai.prompts._get_session")
     def test_url_encode_prompt_names_with_special_characters(self, mock_get_session):
         """Should URL-encode prompt names with special characters."""