fix(prompts): protect cache from stale background refreshes

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,13 +4035,15 @@ 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,
                         label=label,
                         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:

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:

tests/unit/test_prompt.py

@@ -1,3 +1,4 @@
+from threading import Event
 from unittest.mock import Mock, patch
 
 import pytest
@@ -581,6 +582,58 @@ 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()
+    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
+
+    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)
+
+    assert response_ready.is_set(), "Background fetch did not reach its response"
+    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