From dcc970e69fc17c6a9f8272e46028d25e79e1bb64 Mon Sep 17 00:00:00 2001 From: Aafimalek <91587101+Aafimalek@users.noreply.github.com> Date: Fri, 6 Nov 2026 13:28:20 +0530 Subject: [PATCH 1/2] fix(prompts): protect cache from stale background refreshes --- langfuse/_client/client.py | 17 ++++++++--- langfuse/_utils/prompt_cache.py | 19 +++++++++++-- tests/unit/test_prompt.py | 50 +++++++++++++++++++++++++++++++++ 3 files changed, 80 insertions(+), 6 deletions(-) diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py index 5d128567c..41755aa45 100644 --- a/langfuse/_client/client.py +++ b/langfuse/_client/client.py @@ -97,7 +97,7 @@ from langfuse._utils import _get_timestamp, json_path from langfuse._utils.environment import get_common_release_envs from langfuse._utils.parse_error import handle_fern_exception -from langfuse._utils.prompt_cache import PromptCache +from langfuse._utils.prompt_cache import PromptCache, PromptCacheItem from langfuse.api import ( AsyncLangfuseAPI, CreateChatPromptRequest, @@ -4035,6 +4035,7 @@ def get_prompt( ) def refresh_task() -> None: + """Refresh only the cache entry that scheduled this task.""" self._fetch_prompt_and_update_cache( name, version=version, @@ -4042,6 +4043,7 @@ def refresh_task() -> None: ttl_seconds=cache_ttl_seconds, max_retries=bounded_max_retries, fetch_timeout_seconds=fetch_timeout_seconds, + expected_cache_item=cached_prompt, ) self._resources.prompt_cache.add_refresh_prompt_task_if_current( @@ -4076,7 +4078,9 @@ def _fetch_prompt_and_update_cache( ttl_seconds: Optional[int] = None, max_retries: int, fetch_timeout_seconds: Optional[int], + expected_cache_item: Optional[PromptCacheItem] = None, ) -> PromptClient: + """Fetch a prompt and update the cache if the expected entry is current.""" cache_key = PromptCache.generate_cache_key(name, version=version, label=label) langfuse_logger.debug("Fetching prompt '%s' from server...", cache_key) @@ -4106,16 +4110,21 @@ def fetch_prompts() -> Any: prompt = TextPromptClient(prompt_response) if self._resources is not None: - self._resources.prompt_cache.set(cache_key, prompt, ttl_seconds) + self._resources.prompt_cache.set( + cache_key, prompt, ttl_seconds, expected_item=expected_cache_item + ) return prompt except NotFoundError as not_found_error: langfuse_logger.warning( - "Prompt '%s' not found during refresh, evicting from cache.", cache_key + "Prompt '%s' not found during fetch.", + cache_key, ) if self._resources is not None: - self._resources.prompt_cache.delete(cache_key) + self._resources.prompt_cache.delete( + cache_key, expected_item=expected_cache_item + ) raise not_found_error except Exception as e: diff --git a/langfuse/_utils/prompt_cache.py b/langfuse/_utils/prompt_cache.py index e286dca49..c895aa04d 100644 --- a/langfuse/_utils/prompt_cache.py +++ b/langfuse/_utils/prompt_cache.py @@ -169,15 +169,30 @@ def get(self, key: str) -> Optional[PromptCacheItem]: with self._lock: return self._cache.get(key, None) - def set(self, key: str, value: PromptClient, ttl_seconds: Optional[int]) -> None: + def set( + self, + key: str, + value: PromptClient, + ttl_seconds: Optional[int], + *, + expected_item: Optional[PromptCacheItem] = None, + ) -> None: + """Cache a prompt if the expected entry has not changed.""" if ttl_seconds is None: ttl_seconds = DEFAULT_PROMPT_CACHE_TTL_SECONDS with self._lock: + if expected_item is not None and self._cache.get(key) is not expected_item: + return self._cache[key] = PromptCacheItem(value, ttl_seconds) - def delete(self, key: str) -> None: + def delete( + self, key: str, *, expected_item: Optional[PromptCacheItem] = None + ) -> None: + """Remove a prompt if the expected entry has not changed.""" with self._lock: + if expected_item is not None and self._cache.get(key) is not expected_item: + return self._cache.pop(key, None) def invalidate(self, prompt_name: str) -> None: diff --git a/tests/unit/test_prompt.py b/tests/unit/test_prompt.py index eadfb8221..f346cb519 100644 --- a/tests/unit/test_prompt.py +++ b/tests/unit/test_prompt.py @@ -1,3 +1,4 @@ +from threading import Event from unittest.mock import Mock, patch import pytest @@ -581,6 +582,55 @@ def test_skip_redundant_refresh_when_cache_already_updated( add_task_mock.assert_not_called() +@pytest.mark.parametrize("not_found", [False, True]) +@pytest.mark.parametrize("replace", [False, True]) +@patch.object(PromptCacheItem, "get_epoch_seconds") +def test_old_refresh_does_not_overwrite_or_delete_new_cache_entry( + mock_time, langfuse: Langfuse, not_found: bool, replace: bool +) -> None: + """An in-flight refresh must not change a newer prompt cache entry.""" + name = "refresh-race" + cache = langfuse._resources.prompt_cache + key = PromptCache.generate_cache_key(name, version=None, label=None) + old = Prompt_Text( + name=name, version=1, prompt="old", labels=[], type="text", config={}, tags=[] + ) + new = Prompt_Text( + name=name, version=2, prompt="new", labels=[], type="text", config={}, tags=[] + ) + started = Event() + release = Event() + + def delayed_fetch(*args, **kwargs): + """Hold the old response until a newer cache entry has been stored.""" + started.set() + assert release.wait(timeout=10) + if not_found: + raise NotFoundError({"message": "not found"}) + return old + + mock_time.return_value = 0 + cache.set(key, TextPromptClient(old), None) + mock_time.return_value = DEFAULT_PROMPT_CACHE_TTL_SECONDS + 1 + langfuse.api.prompts.get.side_effect = delayed_fetch + try: + assert langfuse.get_prompt(name, max_retries=0).version == 1 + assert started.wait(timeout=10) + cache.invalidate(name) + if replace: + cache.set(key, TextPromptClient(new), None) + finally: + release.set() + wait_for_prompt_refresh(langfuse) + + current = cache.get(key) + if replace: + assert current is not None + assert current.value.version == 2 + else: + assert current is None + + @patch.object(PromptCacheItem, "get_epoch_seconds") def test_get_fresh_prompt_when_expired_cache_default_ttl(mock_time, langfuse: Langfuse): mock_time.return_value = 0 From a1bb799af44c0253a3235b2dd20801802fc8b58c Mon Sep 17 00:00:00 2001 From: Aafimalek <91587101+Aafimalek@users.noreply.github.com> Date: Fri, 6 Nov 2026 14:07:17 +0530 Subject: [PATCH 2/2] test(prompts): verify background refresh reaches response --- tests/unit/test_prompt.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/test_prompt.py b/tests/unit/test_prompt.py index f346cb519..300be58ed 100644 --- a/tests/unit/test_prompt.py +++ b/tests/unit/test_prompt.py @@ -600,11 +600,13 @@ def test_old_refresh_does_not_overwrite_or_delete_new_cache_entry( ) started = Event() release = Event() + response_ready = Event() def delayed_fetch(*args, **kwargs): """Hold the old response until a newer cache entry has been stored.""" started.set() assert release.wait(timeout=10) + response_ready.set() if not_found: raise NotFoundError({"message": "not found"}) return old @@ -623,6 +625,7 @@ def delayed_fetch(*args, **kwargs): release.set() wait_for_prompt_refresh(langfuse) + assert response_ready.is_set(), "Background fetch did not reach its response" current = cache.get(key) if replace: assert current is not None