fix(prompts): share one request between concurrent fetches of the same prompt (#1059)

Author: jurajmajerik

.sampo/changesets/prompts-single-flight.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+`Prompts.get()` now makes one request when several threads ask for the same prompt at once, instead of one request per thread. The other threads wait for that result. With `capture_errors=True`, a shared failure is reported once instead of once per thread.

posthog/ai/prompts.py

@@ -7,10 +7,13 @@
 import copy
 import logging
 import math
+import os
 import re
+import threading
 import time
 import urllib.parse
 import warnings
+import weakref
 from dataclasses import dataclass
 from typing import Any, Dict, List, Literal, Optional, Union, overload
 
@@ -37,6 +40,33 @@
 PromptSource = Literal["api", "cache", "stale_cache", "code_fallback"]
 
 
+class _InFlightFetch:
+    """One network request that several callers wait on."""
+
+    def __init__(self) -> None:
+        self.done = threading.Event()
+        self.result: Optional[CachedPrompt] = None
+        self.error: Optional[BaseException] = None
+
+
+# Every live Prompts instance, so one fork hook can reset them all. Registered
+# once here, because os.register_at_fork keeps each callback for the process
+# lifetime and an app that builds many instances would accumulate them.
+_instances: "weakref.WeakSet[Prompts]" = weakref.WeakSet()
+
+
+def _reset_in_flight_after_fork() -> None:
+    # The child has no copy of the thread that was fetching, so an inherited
+    # in-flight entry would hold every later caller for that key.
+    for prompts in list(_instances):
+        prompts._in_flight = {}
+        prompts._in_flight_lock = threading.Lock()
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_reset_in_flight_after_fork)
+
+
 @dataclass(frozen=True)
 class PromptResult:
     """Result of a prompt fetch with metadata about its source.
@@ -270,6 +300,12 @@ def __init__(
             else DEFAULT_CACHE_TTL_SECONDS
         )
         self._cache: Dict[PromptCacheKey, CachedPrompt] = {}
+        # One request per cache key at a time. Threads that arrive while it is in
+        # flight wait for that result, so a worker with many concurrent callers does
+        # not turn every cache miss or expiry into one request per caller.
+        self._in_flight: Dict[PromptCacheKey, _InFlightFetch] = {}
+        self._in_flight_lock = threading.Lock()
+        _instances.add(self)
         self._has_warned_deprecation = False
         self._client = posthog
         self._capture_errors = capture_errors
@@ -568,44 +604,18 @@ def _get_internal(
 
         # Try to fetch from API
         try:
-            data = self._fetch_prompt_from_api(name, version, label)
-
-            # An older PostHog server ignores the label param and returns the latest
-            # version with no label field — surface that instead of failing silently.
-            if label is not None and data.get("label") != label:
-                log.warning(
-                    "[PostHog Prompts] Requested label %r for prompt %r but the server "
-                    "resolved %r. It may not support prompt labels yet and returned the "
-                    "latest version instead.",
-                    label,
-                    name,
-                    data.get("label"),
-                )
-
-            config = _extract_config(data)
-
-            # Update cache
-            self._cache[cache_key] = CachedPrompt(
-                prompt=data["prompt"],
-                fetched_at=time.time(),
-                name=data["name"],
-                version=data["version"],
-                label=data.get("label"),
-                config=config,
-            )
+            fetched = self._fetch_prompt_once(name, version, label, ttl=ttl)
 
             return PromptResult(
                 source="api",
-                prompt=data["prompt"],
-                name=data["name"],
-                version=data["version"],
-                label=data.get("label"),
-                config=copy.deepcopy(config),
+                prompt=fetched.prompt,
+                name=fetched.name,
+                version=fetched.version,
+                label=fetched.label,
+                config=copy.deepcopy(fetched.config),
             )
 
         except Exception as error:
-            self._maybe_capture_error(error, name=name, version=version, label=label)
-
             prompt_reference = _prompt_reference(name, version, label)
             # Return stale cache (with warning)
             if cached is not None:
@@ -634,6 +644,79 @@ def _get_internal(
 
             raise
 
+    def _fetch_prompt_once(
+        self,
+        name: str,
+        version: Optional[int],
+        label: Optional[str],
+        *,
+        ttl: float,
+    ) -> CachedPrompt:
+        """Fetch a prompt, sharing the request with other threads that want the same key."""
+        cache_key = _cache_key(name, version, label)
+        with self._in_flight_lock:
+            # Another thread may have finished a fetch between the caller's cache
+            # check and this lock. Its entry is fresh, so serve it instead of
+            # starting another request. A zero TTL is never fresh and still fetches.
+            cached = self._cache.get(cache_key)
+            if cached is not None and (time.time() - cached.fetched_at) < ttl:
+                return cached
+
+            pending = self._in_flight.get(cache_key)
+            is_leader = pending is None
+            if pending is None:
+                pending = _InFlightFetch()
+                self._in_flight[cache_key] = pending
+
+        if not is_leader:
+            pending.done.wait()
+            if pending.error is not None:
+                raise pending.error
+            assert pending.result is not None
+            return pending.result
+
+        try:
+            data = self._fetch_prompt_from_api(name, version, label)
+
+            # An older PostHog server ignores the label param and returns the latest
+            # version with no label field — surface that instead of failing silently.
+            if label is not None and data.get("label") != label:
+                log.warning(
+                    "[PostHog Prompts] Requested label %r for prompt %r but the server "
+                    "resolved %r. It may not support prompt labels yet and returned the "
+                    "latest version instead.",
+                    label,
+                    name,
+                    data.get("label"),
+                )
+
+            fetched = CachedPrompt(
+                prompt=data["prompt"],
+                fetched_at=time.time(),
+                name=data["name"],
+                version=data["version"],
+                label=data.get("label"),
+                config=_extract_config(data),
+            )
+            # Cached before the in-flight entry is released, so a caller that
+            # arrives in between finds the entry instead of starting a request.
+            self._cache[cache_key] = fetched
+            pending.result = fetched
+            return fetched
+        except BaseException as error:
+            pending.error = error
+            raise
+        finally:
+            with self._in_flight_lock:
+                self._in_flight.pop(cache_key, None)
+            pending.done.set()
+            # Reported after the waiters are released, so an error-tracking callback
+            # that fetches the same prompt cannot wait on this very request.
+            if isinstance(pending.error, Exception):
+                self._maybe_capture_error(
+                    pending.error, name=name, version=version, label=label
+                )
+
     def compile(self, prompt: str, variables: PromptVariables) -> str:
         """
         Replace {{variableName}} placeholders with values.

posthog/test/ai/test_prompts.py

@@ -1,3 +1,4 @@
+import threading
 import unittest
 import warnings
 from unittest.mock import MagicMock, patch
@@ -382,6 +383,125 @@ def test_hold_a_cooldown_after_each_failed_refetch(
         self.assertEqual(result.source, "api")
         self.assertEqual(mock_get.call_count, 4)
 
+    def _run_concurrent_gets(self, prompts, calls):
+        results = [None] * len(calls)
+        start_together = threading.Barrier(len(calls))
+
+        def run(index, kwargs):
+            start_together.wait(5)
+            results[index] = prompts.get("test-prompt", with_metadata=True, **kwargs)
+
+        threads = [
+            threading.Thread(target=run, args=(index, kwargs))
+            for index, kwargs in enumerate(calls)
+        ]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join(timeout=5)
+        return results
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_share_one_request_between_concurrent_callers(self, mock_get_session):
+        mock_get = mock_get_session.return_value.get
+        release = threading.Event()
+
+        def slow_get(url, **kwargs):
+            release.wait(5)
+            version = 2 if "version=2" in url else 1
+            return MockResponse(
+                json_data={**self.mock_prompt_response, "version": version}
+            )
+
+        mock_get.side_effect = slow_get
+        prompts = Prompts(self.create_mock_posthog())
+
+        calls = [{}] * 10 + [{"version": 2}]
+
+        # Release the network once every caller is waiting on it.
+        timer = threading.Timer(0.2, release.set)
+        timer.start()
+        results = self._run_concurrent_gets(prompts, calls)
+        timer.join()
+
+        self.assertEqual(mock_get.call_count, 2)
+        self.assertEqual({r.source for r in results}, {"api"})
+        self.assertEqual([r.version for r in results[:10]], [1] * 10)
+        self.assertEqual(results[10].version, 2)
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_concurrent_callers_share_one_failure_and_capture_it_once(
+        self, mock_get_session
+    ):
+        mock_get = mock_get_session.return_value.get
+        release = threading.Event()
+
+        def failing_get(url, **kwargs):
+            release.wait(5)
+            raise Exception("Network error")
+
+        mock_get.side_effect = failing_get
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog, capture_errors=True)
+
+        timer = threading.Timer(0.2, release.set)
+        timer.start()
+        results = self._run_concurrent_gets(
+            prompts, [{"fallback": "fallback prompt"}] * 10
+        )
+        timer.join()
+
+        self.assertEqual(mock_get.call_count, 1)
+        self.assertEqual({r.source for r in results}, {"code_fallback"})
+        posthog.capture_exception.assert_called_once()
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_fork_hook_clears_inherited_in_flight_entries(self, mock_get_session):
+        from posthog.ai.prompts import (
+            _InFlightFetch,
+            _cache_key,
+            _reset_in_flight_after_fork,
+        )
+
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(json_data=self.mock_prompt_response)
+        prompts = Prompts(self.create_mock_posthog())
+        # What a forked child inherits: an entry whose fetching thread no longer exists.
+        prompts._in_flight[_cache_key("test-prompt", None, None)] = _InFlightFetch()
+
+        _reset_in_flight_after_fork()
+        result = prompts.get("test-prompt", with_metadata=True)
+
+        self.assertEqual(result.source, "api")
+        self.assertEqual(mock_get.call_count, 1)
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_error_callback_that_fetches_the_same_prompt_does_not_deadlock(
+        self, mock_get_session
+    ):
+        mock_get = mock_get_session.return_value.get
+        mock_get.side_effect = Exception("Network error")
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog, capture_errors=True)
+        nested_results = []
+        entered = threading.Event()
+
+        def capture_exception(*args, **kwargs):
+            # The nested fetch fails too and reports again; only the first report fetches.
+            if entered.is_set():
+                return
+            entered.set()
+            nested_results.append(
+                prompts.get("test-prompt", with_metadata=True, fallback="nested")
+            )
+
+        posthog.capture_exception.side_effect = capture_exception
+
+        result = prompts.get("test-prompt", with_metadata=True, fallback="outer")
+
+        self.assertEqual(result.prompt, "outer")
+        self.assertEqual([r.prompt for r in nested_results], ["nested"])
+
     @patch("posthog.ai.prompts._get_session")
     @patch("posthog.ai.prompts.log")
     def test_use_fallback_when_no_cache_and_fetch_fails_with_warning(