fix(langfuse): end the generation when a child span fails, take the client slot last, keep prompt cache keys structured

Generation spans now end in a finally block so a bad guardrail or provider entry cannot strand the trace. The logger acquires its export channel and REST client before counting a client slot and releases the channel synchronously if the REST client fails to build, so retries after a bad config do not exhaust the budget. LANGFUSE_TIMEOUT accepts decimals for the REST client like it already did for OTLP export. The prompt cache keys on (name, version, label) so a missing label and the literal label None stay apart

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/integrations/langfuse/langfuse.py

@@ -284,7 +284,6 @@ class LangFuseLogger:
                 f"\033[91mLangfuse not installed, try running 'pip install langfuse' to fix this error: {e}\033[0m"
             ) from e
         raise_if_unsupported_langfuse_version(self.langfuse_sdk_version)
-        from litellm.integrations.langfuse.langfuse_sdk import acquire_langfuse_tracing
 
         self.public_key, self.secret_key, self.langfuse_host = resolve_langfuse_credentials(
             langfuse_public_key=langfuse_public_key,
@@ -310,16 +309,9 @@ class LangFuseLogger:
             self.langfuse_client = self._http_handler.client
             self.is_mock_mode = False
 
-        self.api_client: LangfuseApiClient = self.safe_init_langfuse_client()
-        self.tracing: LangfuseTracing = acquire_langfuse_tracing(
-            public_key=str(self.public_key),
-            secret_key=str(self.secret_key),
-            base_url=self.langfuse_host,
-            environment=self.langfuse_environment,
-            release=self.langfuse_release,
-            flush_interval=self.langfuse_flush_interval,
-            mock_mode=self.is_mock_mode,
-        )
+        self.api_client: LangfuseApiClient
+        self.tracing: LangfuseTracing
+        self.api_client, self.tracing = self.safe_init_langfuse_client()
 
         # set the current langfuse project id in the environ
         # this is used by Alerting to link to the correct project
@@ -336,8 +328,8 @@ class LangFuseLogger:
 
         warn_if_upstream_langfuse_configured()
 
-    def safe_init_langfuse_client(self) -> LangfuseApiClient:
-        """Build the REST client while the process is under its logger budget.
+    def safe_init_langfuse_client(self) -> "tuple[LangfuseApiClient, LangfuseTracing]":
+        """Build the REST client and export channel while the process is under its logger budget.
 
         The budget dates from the SDK client, which started a consumer thread per instance and once
         pinned a CPU at 100% when many were built; it still bounds the number of per-key loggers.
@@ -346,17 +338,34 @@ class LangFuseLogger:
             raise Exception(
                 f"Max langfuse clients reached: {litellm.initialized_langfuse_clients} is greater than {MAX_LANGFUSE_INITIALIZED_CLIENTS}"
             )
-        from litellm.integrations.langfuse.langfuse_sdk import build_langfuse_client
-
-        api_client: Final = build_langfuse_client(
-            public_key=self.public_key,
-            secret_key=self.secret_key,
-            base_url=self.langfuse_host,
-            httpx_client=self.langfuse_client,
+        from litellm.integrations.langfuse.langfuse_sdk import (
+            acquire_langfuse_tracing,
+            build_langfuse_client,
+            release_langfuse_tracing,
         )
+
+        tracing: Final = acquire_langfuse_tracing(
+            public_key=str(self.public_key),
+            secret_key=str(self.secret_key),
+            base_url=self.langfuse_host,
+            environment=self.langfuse_environment,
+            release=self.langfuse_release,
+            flush_interval=self.langfuse_flush_interval,
+            mock_mode=self.is_mock_mode,
+        )
+        try:
+            api_client: Final = build_langfuse_client(
+                public_key=self.public_key,
+                secret_key=self.secret_key,
+                base_url=self.langfuse_host,
+                httpx_client=self.langfuse_client,
+            )
+        except Exception:
+            release_langfuse_tracing(tracing, grace_seconds=0.0)
+            raise
         litellm.initialized_langfuse_clients += 1
         verbose_logger.debug("Created langfuse client number %s", litellm.initialized_langfuse_clients)
-        return api_client
+        return api_client, tracing
 
     def flush(self) -> None:
         """Push every queued observation to Langfuse before the process goes away."""
@@ -916,11 +925,15 @@ class LangFuseLogger:
                 public=trace_public,
                 attributes=MappingProxyType({**generation_attributes, **trace_level_attributes}),
             )
-            log_provider_specific_information_as_span(tracing=self.tracing, parent=generation, enrichments=enrichments)
-            self._log_guardrail_information_as_span(
-                tracing=self.tracing, parent=generation, standard_logging_object=standard_logging_object
-            )
-            generation.end(end_time)
+            try:
+                log_provider_specific_information_as_span(
+                    tracing=self.tracing, parent=generation, enrichments=enrichments
+                )
+                self._log_guardrail_information_as_span(
+                    tracing=self.tracing, parent=generation, standard_logging_object=standard_logging_object
+                )
+            finally:
+                generation.end(end_time)
 
             # log_event_on_langfuse tuple-unpacks this and re-wraps it in the dict callers cache.
             # The observation id is the requested generation_id after resolve_observation_id.

litellm/integrations/langfuse/langfuse_sdk.py

@@ -742,6 +742,9 @@ class _CachedPrompt:
     fetched_at: float
 
 
+_PromptKey = tuple[str, int | None, str | None]
+
+
 def _prompt_client(prompt: Prompt) -> PromptClient:
     return ChatPromptClient(prompt) if isinstance(prompt, Prompt_Chat) else TextPromptClient(prompt)
 
@@ -762,7 +765,8 @@ class LangfuseApiClient:
     def __init__(self, api: LangfuseAPI, *, prompt_cache_ttl_seconds: float) -> None:
         self.api: Final = api
         self.prompt_cache_ttl_seconds: Final = prompt_cache_ttl_seconds
-        self._prompts: Final[dict[str, _CachedPrompt]] = {}  # mutable-ok: per-client prompt cache, guarded by _lock
+        # mutable-ok: per-client prompt cache, guarded by _lock
+        self._prompts: Final[dict[_PromptKey, _CachedPrompt]] = {}
         self._lock: Final = threading.Lock()
 
     def auth_check(self) -> bool:
@@ -777,7 +781,7 @@ class LangfuseApiClient:
         return projects[0].id if projects else None
 
     def get_prompt(self, name: str, *, label: str | None = None, version: int | None = None) -> PromptClient:
-        key: Final = f"{name}:version:{version}" if version is not None else f"{name}:label:{label}"
+        key: Final[_PromptKey] = (name, version, label)
         with self._lock:
             cached: Final = self._prompts.get(key)
         if cached is not None and monotonic() - cached.fetched_at < self.prompt_cache_ttl_seconds:
@@ -815,7 +819,7 @@ def build_langfuse_client(
             x_langfuse_sdk_version=version("langfuse"),
             x_langfuse_public_key=public_key,
             httpx_client=httpx_client,
-            timeout=int(os.getenv("LANGFUSE_TIMEOUT", "5")),
+            timeout=float(os.getenv("LANGFUSE_TIMEOUT", "5")),
         ),
         prompt_cache_ttl_seconds=float(os.getenv("LANGFUSE_PROMPT_CACHE_DEFAULT_TTL_SECONDS", "60")),
     )

tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py

@@ -29,6 +29,7 @@ from litellm.integrations.langfuse.langfuse import (
 )
 from litellm.integrations.langfuse.langfuse_sdk import (
     DiscardingSpanExporter,
+    LangfuseApiClient,
     LangfuseSpanExporter,
     LangfuseTracing,
     _build_span_exporter,
@@ -1067,3 +1068,40 @@ def test_export_endpoint_never_doubles_the_slash_or_leaves_the_configured_host(
     exporter = _build_span_exporter(public_key="pk", secret_key="sk", base_url=base_url)
 
     assert exporter.endpoint == expected
+
+
+class _RecordingPromptsApi:
+    """Answers ``prompts.get`` with a text prompt that names the label it was asked for."""
+
+    def __init__(self) -> None:
+        self.prompts = self
+        self.requests: list[tuple[str, int | None, str | None]] = []  # mutable-ok: test-side call log
+
+    def get(self, name: str, *, version: int | None, label: str | None):
+        from langfuse.api import Prompt_Text
+
+        self.requests.append((name, version, label))
+        return Prompt_Text(
+            name=name,
+            version=version or 1,
+            config={},
+            labels=[label or "production"],
+            tags=[],
+            prompt=f"label={label!r}",
+        )
+
+
+def test_prompt_cache_keeps_a_missing_label_apart_from_the_label_named_none():
+    """A prompt labelled ``"None"`` and the unlabelled default are different prompts in Langfuse
+    and must not answer each other's requests from the cache."""
+    api = _RecordingPromptsApi()
+    client = LangfuseApiClient(api, prompt_cache_ttl_seconds=60)  # pyright: ignore[reportArgumentType]  # duck-typed prompts API
+
+    unlabelled = client.get_prompt("greeting")
+    named_none = client.get_prompt("greeting", label="None")
+    cached_unlabelled = client.get_prompt("greeting")
+
+    assert unlabelled.prompt == "label=None"
+    assert named_none.prompt == "label='None'"
+    assert cached_unlabelled is unlabelled
+    assert api.requests == [("greeting", None, None), ("greeting", None, "None")]

tests/test_litellm/integrations/test_langfuse.py

@@ -633,6 +633,24 @@ class TestLangfuseUsageDetails(unittest.TestCase):
             assert child.context.trace_id == generation.context.trace_id
             assert "langfuse.trace.name" not in child.attributes
 
+    def test_generation_is_exported_when_a_child_span_fails(self):
+        """v2 buffered the generation in one call, so a bad guardrail entry could not lose it;
+        the OTel generation is open until ``end()`` and must still be ended when a child raises."""
+        self._drive_with_canary(
+            guardrail_information=[
+                {
+                    "guardrail_name": "pii-post",
+                    "guardrail_mode": "post_call",
+                    "start_time": "not-a-timestamp",
+                    "end_time": 1704110403.0,
+                }
+            ],
+        )
+
+        [generation] = [span for span in self.span_exporter.get_finished_spans() if span.name.startswith("litellm-")]
+        assert generation.attributes["langfuse.trace.name"] == "canary-trace"
+        assert self.exported_spans_named("guardrail") == []
+
     def test_caller_cannot_spoof_an_allowlisted_identity_field(self):
         """
         Request metadata never reaches the blob, so a caller naming user_api_key_alias
@@ -2220,6 +2238,35 @@ def test_stopped_logger_hands_its_export_channel_back(monkeypatch):
     assert acquire_same_credentials() is not logger.tracing, "stop() did not give the logger's hold back"
 
 
+def test_logger_that_fails_to_build_takes_no_slot_and_no_channel(monkeypatch):
+    """Each failed retry for the same dynamic credentials would otherwise eat a client slot and a
+    holder on the channel, so fixing the configuration could not bring Langfuse logging back."""
+    from litellm.integrations.langfuse.langfuse_sdk import acquire_langfuse_tracing, release_langfuse_tracing
+
+    monkeypatch.setenv("LANGFUSE_TIMEOUT", "5.5")
+    probe = _build_langfuse_logger(monkeypatch, langfuse_public_key="pk-failed-build")
+
+    def acquire_same_credentials():
+        return acquire_langfuse_tracing(
+            public_key="pk-failed-build",
+            secret_key="sk-lit5228",
+            base_url=_UNREACHABLE_HOST,
+            environment=probe.langfuse_environment,
+            release=probe.langfuse_release,
+            flush_interval=probe.langfuse_flush_interval,
+            mock_mode=False,
+        )
+
+    monkeypatch.setenv("LANGFUSE_TIMEOUT", "not-a-number")
+    with pytest.raises(ValueError, match="not-a-number"):
+        _build_langfuse_logger(monkeypatch, langfuse_public_key="pk-failed-build")
+    assert litellm.initialized_langfuse_clients == 0
+
+    monkeypatch.setenv("LANGFUSE_TIMEOUT", "5.5")
+    release_langfuse_tracing(probe.tracing, grace_seconds=0.0)
+    assert acquire_same_credentials() is not probe.tracing, "the failed build left a holder on the channel"
+
+
 def test_int_steering_values_reach_langfuse_as_strings():
     """Langfuse models user, session and version as strings; v2's pydantic coerced ints for the caller."""
     rig = _steering_logger()