From 4ad930481433ed4226ba49e1a7c972d4345ef880 Mon Sep 17 00:00:00 2001 From: Juraj Majerik Date: Wed, 14 Oct 2026 07:41:29 +0200 Subject: [PATCH 1/5] fix(ai): share one request between concurrent fetches of the same prompt --- .sampo/changesets/prompts-single-flight.md | 5 ++ posthog/ai/prompts.py | 53 +++++++++++++++- posthog/test/ai/test_prompts.py | 73 ++++++++++++++++++++++ 3 files changed, 128 insertions(+), 3 deletions(-) create mode 100644 .sampo/changesets/prompts-single-flight.md diff --git a/.sampo/changesets/prompts-single-flight.md b/.sampo/changesets/prompts-single-flight.md new file mode 100644 index 00000000..6051ac9b --- /dev/null +++ b/.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. diff --git a/posthog/ai/prompts.py b/posthog/ai/prompts.py index 2cbae9db..e116c1c2 100644 --- a/posthog/ai/prompts.py +++ b/posthog/ai/prompts.py @@ -8,6 +8,7 @@ import logging import math import re +import threading import time import urllib.parse import warnings @@ -37,6 +38,15 @@ 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[Dict[str, Any]] = None + self.error: Optional[BaseException] = None + + @dataclass(frozen=True) class PromptResult: """Result of a prompt fetch with metadata about its source. @@ -270,6 +280,11 @@ 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() self._has_warned_deprecation = False self._client = posthog self._capture_errors = capture_errors @@ -568,7 +583,7 @@ def _get_internal( # Try to fetch from API try: - data = self._fetch_prompt_from_api(name, version, label) + data = self._fetch_prompt_once(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. @@ -604,8 +619,6 @@ def _get_internal( ) 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 +647,40 @@ def _get_internal( raise + def _fetch_prompt_once( + self, name: str, version: Optional[int], label: Optional[str] + ) -> Dict[str, Any]: + """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: + 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: + pending.result = self._fetch_prompt_from_api(name, version, label) + return pending.result + except BaseException as error: + pending.error = error + if isinstance(error, Exception): + self._maybe_capture_error( + error, name=name, version=version, label=label + ) + raise + finally: + with self._in_flight_lock: + self._in_flight.pop(cache_key, None) + pending.done.set() + def compile(self, prompt: str, variables: PromptVariables) -> str: """ Replace {{variableName}} placeholders with values. diff --git a/posthog/test/ai/test_prompts.py b/posthog/test/ai/test_prompts.py index 4e124f16..44246ce7 100644 --- a/posthog/test/ai/test_prompts.py +++ b/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,78 @@ 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") @patch("posthog.ai.prompts.log") def test_use_fallback_when_no_cache_and_fetch_fails_with_warning( From 327fa7936036528a5649a3eb8f5dcaa24de4b33a Mon Sep 17 00:00:00 2001 From: Juraj Majerik Date: Wed, 14 Oct 2026 10:21:53 +0200 Subject: [PATCH 2/5] fix(ai): cache before releasing the shared fetch, report errors after it, reset it after fork --- posthog/ai/prompts.py | 91 +++++++++++++++++++++------------ posthog/test/ai/test_prompts.py | 43 ++++++++++++++++ 2 files changed, 102 insertions(+), 32 deletions(-) diff --git a/posthog/ai/prompts.py b/posthog/ai/prompts.py index e116c1c2..46f0e1a7 100644 --- a/posthog/ai/prompts.py +++ b/posthog/ai/prompts.py @@ -7,11 +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 @@ -31,6 +33,10 @@ # The server's tightest prompt limit is per-minute, so a minute lets the bucket refill. DEFAULT_REFETCH_COOLDOWN_SECONDS = 60 MAX_REFETCH_COOLDOWN_SECONDS = 3600 +# Longest a caller waits on another thread's fetch: the 10s request timeout times +# the transport's three attempts. Past it the leader is assumed lost, for example +# to a fork, and the caller fetches for itself. +_IN_FLIGHT_WAIT_SECONDS = 30 PromptVariables = Dict[str, Union[str, int, float, bool]] PromptCacheKey = tuple[str, Optional[int], Optional[str]] @@ -285,6 +291,11 @@ def __init__( # not turn every cache miss or expiry into one request per caller. self._in_flight: Dict[PromptCacheKey, _InFlightFetch] = {} self._in_flight_lock = threading.Lock() + if hasattr(os, "register_at_fork"): + weak_self = weakref.ref(self) + os.register_at_fork( + after_in_child=lambda: Prompts._reset_in_flight_after_fork(weak_self) + ) self._has_warned_deprecation = False self._client = posthog self._capture_errors = capture_errors @@ -585,37 +596,13 @@ def _get_internal( try: data = self._fetch_prompt_once(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, - ) - return PromptResult( source="api", prompt=data["prompt"], name=data["name"], version=data["version"], label=data.get("label"), - config=copy.deepcopy(config), + config=copy.deepcopy(_extract_config(data)), ) except Exception as error: @@ -647,6 +634,16 @@ def _get_internal( raise + @staticmethod + def _reset_in_flight_after_fork(weak_self: "weakref.ref[Prompts]") -> 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. + self = weak_self() + if self is None: + return + self._in_flight = {} + self._in_flight_lock = threading.Lock() + def _fetch_prompt_once( self, name: str, version: Optional[int], label: Optional[str] ) -> Dict[str, Any]: @@ -660,26 +657,56 @@ def _fetch_prompt_once( self._in_flight[cache_key] = pending if not is_leader: - pending.done.wait() + if not pending.done.wait(_IN_FLIGHT_WAIT_SECONDS): + with self._in_flight_lock: + if self._in_flight.get(cache_key) is pending: + del self._in_flight[cache_key] + return self._fetch_prompt_once(name, version, label) if pending.error is not None: raise pending.error assert pending.result is not None return pending.result try: - pending.result = self._fetch_prompt_from_api(name, version, label) - return pending.result + 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"), + ) + + # 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] = CachedPrompt( + prompt=data["prompt"], + fetched_at=time.time(), + name=data["name"], + version=data["version"], + label=data.get("label"), + config=_extract_config(data), + ) + pending.result = data + return data except BaseException as error: pending.error = error - if isinstance(error, Exception): - self._maybe_capture_error( - error, name=name, version=version, label=label - ) 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: """ diff --git a/posthog/test/ai/test_prompts.py b/posthog/test/ai/test_prompts.py index 44246ce7..e64ac7e1 100644 --- a/posthog/test/ai/test_prompts.py +++ b/posthog/test/ai/test_prompts.py @@ -455,6 +455,49 @@ def failing_get(url, **kwargs): self.assertEqual({r.source for r in results}, {"code_fallback"}) posthog.capture_exception.assert_called_once() + @patch("posthog.ai.prompts._IN_FLIGHT_WAIT_SECONDS", 0.05) + @patch("posthog.ai.prompts._get_session") + def test_fetches_for_itself_when_the_in_flight_leader_is_gone( + self, mock_get_session + ): + from posthog.ai.prompts import _InFlightFetch, _cache_key + + 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() + + result = prompts.get("test-prompt", with_metadata=True) + + self.assertEqual(result.source, "api") + self.assertEqual(mock_get.call_count, 1) + self.assertEqual(prompts._in_flight, {}) + + @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 = [] + + def capture_exception(*args, **kwargs): + if nested_results: + return + 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( From efa529d3574258dbdddb6aaf73e9ad8cba8e1340 Mon Sep 17 00:00:00 2001 From: Juraj Majerik Date: Wed, 14 Oct 2026 10:22:20 +0200 Subject: [PATCH 3/5] test(ai): stop the re-entrant capture test from recursing --- posthog/test/ai/test_prompts.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/posthog/test/ai/test_prompts.py b/posthog/test/ai/test_prompts.py index e64ac7e1..82411c6f 100644 --- a/posthog/test/ai/test_prompts.py +++ b/posthog/test/ai/test_prompts.py @@ -483,10 +483,13 @@ def test_error_callback_that_fetches_the_same_prompt_does_not_deadlock( posthog = self.create_mock_posthog() prompts = Prompts(posthog, capture_errors=True) nested_results = [] + entered = threading.Event() def capture_exception(*args, **kwargs): - if nested_results: + # 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") ) From 246c2c08bbb136deaa6d06f971d95f3954996a6c Mon Sep 17 00:00:00 2001 From: Juraj Majerik Date: Wed, 14 Oct 2026 10:26:39 +0200 Subject: [PATCH 4/5] fix(ai): one fork hook for all prompt clients, keep a replacement fetch's entry --- posthog/ai/prompts.py | 39 +++++++++++++++++++++++---------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/posthog/ai/prompts.py b/posthog/ai/prompts.py index 46f0e1a7..97c5a736 100644 --- a/posthog/ai/prompts.py +++ b/posthog/ai/prompts.py @@ -53,6 +53,24 @@ def __init__(self) -> 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. @@ -291,11 +309,7 @@ def __init__( # not turn every cache miss or expiry into one request per caller. self._in_flight: Dict[PromptCacheKey, _InFlightFetch] = {} self._in_flight_lock = threading.Lock() - if hasattr(os, "register_at_fork"): - weak_self = weakref.ref(self) - os.register_at_fork( - after_in_child=lambda: Prompts._reset_in_flight_after_fork(weak_self) - ) + _instances.add(self) self._has_warned_deprecation = False self._client = posthog self._capture_errors = capture_errors @@ -634,16 +648,6 @@ def _get_internal( raise - @staticmethod - def _reset_in_flight_after_fork(weak_self: "weakref.ref[Prompts]") -> 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. - self = weak_self() - if self is None: - return - self._in_flight = {} - self._in_flight_lock = threading.Lock() - def _fetch_prompt_once( self, name: str, version: Optional[int], label: Optional[str] ) -> Dict[str, Any]: @@ -698,8 +702,11 @@ def _fetch_prompt_once( pending.error = error raise finally: + # A waiter that gave up on this fetch may have started its own; only + # remove the entry if it is still this one. with self._in_flight_lock: - self._in_flight.pop(cache_key, None) + if self._in_flight.get(cache_key) is pending: + del self._in_flight[cache_key] 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. From 7863971baf8046ece2b0785141d9670316af5946 Mon Sep 17 00:00:00 2001 From: Juraj Majerik Date: Wed, 14 Oct 2026 10:31:01 +0200 Subject: [PATCH 5/5] fix(ai): recheck the cache under the election lock and drop the timed wait --- posthog/ai/prompts.py | 56 +++++++++++++++++---------------- posthog/test/ai/test_prompts.py | 13 ++++---- 2 files changed, 36 insertions(+), 33 deletions(-) diff --git a/posthog/ai/prompts.py b/posthog/ai/prompts.py index 97c5a736..3a785f0c 100644 --- a/posthog/ai/prompts.py +++ b/posthog/ai/prompts.py @@ -33,10 +33,6 @@ # The server's tightest prompt limit is per-minute, so a minute lets the bucket refill. DEFAULT_REFETCH_COOLDOWN_SECONDS = 60 MAX_REFETCH_COOLDOWN_SECONDS = 3600 -# Longest a caller waits on another thread's fetch: the 10s request timeout times -# the transport's three attempts. Past it the leader is assumed lost, for example -# to a fork, and the caller fetches for itself. -_IN_FLIGHT_WAIT_SECONDS = 30 PromptVariables = Dict[str, Union[str, int, float, bool]] PromptCacheKey = tuple[str, Optional[int], Optional[str]] @@ -49,7 +45,7 @@ class _InFlightFetch: def __init__(self) -> None: self.done = threading.Event() - self.result: Optional[Dict[str, Any]] = None + self.result: Optional[CachedPrompt] = None self.error: Optional[BaseException] = None @@ -608,15 +604,15 @@ def _get_internal( # Try to fetch from API try: - data = self._fetch_prompt_once(name, version, label) + 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(_extract_config(data)), + prompt=fetched.prompt, + name=fetched.name, + version=fetched.version, + label=fetched.label, + config=copy.deepcopy(fetched.config), ) except Exception as error: @@ -649,11 +645,23 @@ def _get_internal( raise def _fetch_prompt_once( - self, name: str, version: Optional[int], label: Optional[str] - ) -> Dict[str, Any]: + 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: @@ -661,11 +669,7 @@ def _fetch_prompt_once( self._in_flight[cache_key] = pending if not is_leader: - if not pending.done.wait(_IN_FLIGHT_WAIT_SECONDS): - with self._in_flight_lock: - if self._in_flight.get(cache_key) is pending: - del self._in_flight[cache_key] - return self._fetch_prompt_once(name, version, label) + pending.done.wait() if pending.error is not None: raise pending.error assert pending.result is not None @@ -686,9 +690,7 @@ def _fetch_prompt_once( data.get("label"), ) - # 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] = CachedPrompt( + fetched = CachedPrompt( prompt=data["prompt"], fetched_at=time.time(), name=data["name"], @@ -696,17 +698,17 @@ def _fetch_prompt_once( label=data.get("label"), config=_extract_config(data), ) - pending.result = data - return 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: - # A waiter that gave up on this fetch may have started its own; only - # remove the entry if it is still this one. with self._in_flight_lock: - if self._in_flight.get(cache_key) is pending: - del self._in_flight[cache_key] + 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. diff --git a/posthog/test/ai/test_prompts.py b/posthog/test/ai/test_prompts.py index 82411c6f..7072bc90 100644 --- a/posthog/test/ai/test_prompts.py +++ b/posthog/test/ai/test_prompts.py @@ -455,12 +455,13 @@ def failing_get(url, **kwargs): self.assertEqual({r.source for r in results}, {"code_fallback"}) posthog.capture_exception.assert_called_once() - @patch("posthog.ai.prompts._IN_FLIGHT_WAIT_SECONDS", 0.05) @patch("posthog.ai.prompts._get_session") - def test_fetches_for_itself_when_the_in_flight_leader_is_gone( - self, mock_get_session - ): - from posthog.ai.prompts import _InFlightFetch, _cache_key + 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) @@ -468,11 +469,11 @@ def test_fetches_for_itself_when_the_in_flight_leader_is_gone( # 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) - self.assertEqual(prompts._in_flight, {}) @patch("posthog.ai.prompts._get_session") def test_error_callback_that_fetches_the_same_prompt_does_not_deadlock(