fix(ai): harden Prompts.get_all against old servers and malformed rows

.sampo/changesets/prompts-get-all-hardening.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+`Prompts.get_all` now fails loudly in two cases it previously papered over: a server that ignores the label filter but happens to have some labels on latest versions no longer produces a silently incomplete result, and a malformed row in the list response now raises the invalid-response error instead of being skipped. A rejected batch also no longer leaves partially cached prompts.

posthog/ai/prompts.py

@@ -112,22 +112,32 @@ def _is_prompt_api_response(data: Any) -> bool:
     )
 
 
-def _row_resolves_label(row: Dict[str, Any], label: str) -> bool:
-    """Check that the server resolved this list row through the requested label.
-
-    An older server ignores the label param on the list endpoint and returns the
-    latest version of every prompt. A row that was resolved through a label
-    carries a matching name and version entry in its all_labels field.
+def _row_label_state(
+    row: Dict[str, Any], label: str
+) -> Literal["resolved", "moved", "absent"]:
+    """Classify how a list row relates to the requested label, via all_labels.
+
+    'resolved': the row is the version the label points to.
+    'moved': the prompt carries the label, but on another version. Happens when
+    the label moves between the query and the response.
+    'absent': the prompt does not carry the label at any version. A server that
+    filters by label never returns such a row, so this means the server ignored
+    the label param (an older PostHog release) and served latest versions.
     """
     all_labels = row.get("all_labels")
     if not isinstance(all_labels, list):
-        return False
-    return any(
-        isinstance(entry, dict)
-        and entry.get("name") == label
-        and entry.get("version") == row.get("version")
-        for entry in all_labels
+        return "absent"
+    entry = next(
+        (
+            candidate
+            for candidate in all_labels
+            if isinstance(candidate, dict) and candidate.get("name") == label
+        ),
+        None,
     )
+    if entry is None:
+        return "absent"
+    return "resolved" if entry.get("version") == row.get("version") else "moved"
 
 
 def _is_same_origin(url: str, host: str) -> bool:
@@ -381,14 +391,55 @@ def get_all(self, *, label: str) -> Dict[str, PromptResult]:
             self._maybe_capture_error(error, name="*", version=None, label=label)
             raise
 
-        now = time.time()
-        results: Dict[str, PromptResult] = {}
+        # Validate every row before caching any, so a rejected batch leaves
+        # the cache untouched.
+        resolved_rows: List[Dict[str, Any]] = []
         skipped: List[str] = []
         for row in rows:
-            if not _is_prompt_api_response(row) or not _row_resolves_label(row, label):
-                skipped.append(str(row.get("name")) if isinstance(row, dict) else "?")
+            if not _is_prompt_api_response(row):
+                invalid_error = Exception(
+                    f'[PostHog Prompts] Invalid response format for prompts with label "{label}"'
+                )
+                self._maybe_capture_error(
+                    invalid_error, name="*", version=None, label=label
+                )
+                raise invalid_error
+
+            label_state = _row_label_state(row, label)
+            if label_state == "absent":
+                # Even one unlabeled row proves the server did not filter, and
+                # then rows that look resolved are only labels that happen to
+                # point at the latest version. A partial result here would hide
+                # the rest, so fail loudly instead.
+                compat_error = Exception(
+                    f'[PostHog Prompts] The server returned a prompt that does not carry label "{label}". '
+                    "It may not support fetching prompts by label on the list endpoint yet. "
+                    "Upgrade PostHog, or fetch prompts one by one with get()."
+                )
+                self._maybe_capture_error(
+                    compat_error, name="*", version=None, label=label
+                )
+                raise compat_error
+            if label_state == "moved":
+                skipped.append(row["name"])
                 continue
+            resolved_rows.append(row)
 
+        if rows and not resolved_rows:
+            # Every returned row was skipped as moved. One moved label is a
+            # mid-request race, but all of them means the server most likely
+            # ignored the label param and served latest versions.
+            compat_error = Exception(
+                f'[PostHog Prompts] The server returned prompts, but none resolve label "{label}". '
+                "It may not support fetching prompts by label on the list endpoint yet. "
+                "Upgrade PostHog, or fetch prompts one by one with get()."
+            )
+            self._maybe_capture_error(compat_error, name="*", version=None, label=label)
+            raise compat_error
+
+        now = time.time()
+        results: Dict[str, PromptResult] = {}
+        for row in resolved_rows:
             config = _extract_config(row)
             self._cache[_cache_key(row["name"], None, label)] = CachedPrompt(
                 prompt=row["prompt"],
@@ -407,19 +458,6 @@ def get_all(self, *, label: str) -> Dict[str, PromptResult]:
                 config=copy.deepcopy(config),
             )
 
-        if rows and not results:
-            # Nothing resolved the label, so the server most likely ignored the
-            # label param and served latest versions. Caching those under the
-            # label would be the silent wrong-version failure labels exist to
-            # prevent, so fail loudly instead.
-            compat_error = Exception(
-                f'[PostHog Prompts] The server returned prompts, but none resolve label "{label}". '
-                "It may not support fetching prompts by label on the list endpoint yet. "
-                "Upgrade PostHog, or fetch prompts one by one with get()."
-            )
-            self._maybe_capture_error(compat_error, name="*", version=None, label=label)
-            raise compat_error
-
         if skipped:
             log.warning(
                 "[PostHog Prompts] Skipped %d prompt(s) that did not resolve label %r: %s",

posthog/test/ai/test_prompts.py

@@ -1428,18 +1428,19 @@ def test_fetches_all_pages_and_seeds_the_cache(self, mock_get_session):
 
     @patch("posthog.ai.prompts._get_session")
     def test_raises_when_the_server_ignores_the_label(self, mock_get_session):
-        # An old server ignores ?label= and returns latest versions; none of the
-        # rows resolve the label, and caching them would serve wrong versions.
+        # An old server ignores ?label= and returns latest versions of every
+        # prompt, including prompts without the label. Even when some labels
+        # happen to point at latest, a partial result would hide the rest.
         mock_get = mock_get_session.return_value.get
-        row = self.labeled_row("prompt-a")
-        row["all_labels"] = []
-        mock_get.return_value = self.list_response([row])
+        looks_resolved = self.labeled_row("prompt-a")
+        unlabeled = {**self.labeled_row("prompt-b"), "all_labels": []}
+        mock_get.return_value = self.list_response([looks_resolved, unlabeled])
 
         prompts = Prompts(self.create_mock_posthog())
 
         with self.assertRaises(Exception) as ctx:
             prompts.get_all(label="production")
-        self.assertIn("none resolve label", str(ctx.exception))
+        self.assertIn("does not carry label", str(ctx.exception))
         self.assertEqual(prompts._cache, {})
 
     @patch("posthog.ai.prompts._get_session")
@@ -1458,6 +1459,46 @@ def test_skips_a_row_whose_label_moved_and_keeps_the_rest(self, mock_get_session
 
         self.assertEqual(list(results), ["prompt-b"])
 
+    @patch("posthog.ai.prompts._get_session")
+    def test_raises_when_every_row_was_skipped_as_moved(self, mock_get_session):
+        # An old server can serve latest versions while every prompt's label
+        # points at an earlier version. Each row then looks like a moved label;
+        # returning {} would report no labeled prompts despite them existing.
+        mock_get = mock_get_session.return_value.get
+        row_a = {
+            **self.labeled_row("prompt-a", version=2),
+            "all_labels": [{"name": "production", "version": 1}],
+        }
+        row_b = {
+            **self.labeled_row("prompt-b", version=3),
+            "all_labels": [{"name": "production", "version": 2}],
+        }
+        mock_get.return_value = self.list_response([row_a, row_b])
+
+        prompts = Prompts(self.create_mock_posthog())
+
+        with self.assertRaises(Exception) as ctx:
+            prompts.get_all(label="production")
+        self.assertIn("none resolve label", str(ctx.exception))
+        self.assertEqual(prompts._cache, {})
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_raises_on_a_malformed_row_and_caches_nothing(self, mock_get_session):
+        # A row failing response validation is a server error, not a moved
+        # label; returning the valid subset would hide it.
+        mock_get = mock_get_session.return_value.get
+        malformed = {**self.labeled_row("prompt-a"), "prompt": 42}
+        mock_get.return_value = self.list_response(
+            [self.labeled_row("prompt-b"), malformed]
+        )
+
+        prompts = Prompts(self.create_mock_posthog())
+
+        with self.assertRaises(Exception) as ctx:
+            prompts.get_all(label="production")
+        self.assertIn("Invalid response format", str(ctx.exception))
+        self.assertEqual(prompts._cache, {})
+
     @patch("posthog.ai.prompts._get_session")
     def test_refuses_a_pagination_link_off_the_configured_host(self, mock_get_session):
         mock_get = mock_get_session.return_value.get