fix(langfuse): defer eviction teardown during callbacks

litellm/integrations/langfuse/langfuse.py

@@ -459,20 +459,23 @@ class LangFuseLogger:
                 status_message=status_message,
             )
             verbose_logger.debug("OUTPUT IN LANGFUSE: %s; original: %s", output, response_obj)
-            trace_id, generation_id = self._log_langfuse_v2(
-                user_id=user_id,
-                metadata=metadata,
-                litellm_params=litellm_params,
-                output=output,
-                start_time=start_time,
-                end_time=end_time,
-                kwargs=kwargs,
-                optional_params=optional_params,
-                input=input,
-                response_obj=response_obj,
-                level=level,
-                litellm_call_id=litellm_call_id,
-            )
+            from litellm.integrations.langfuse.langfuse_sdk import lease_langfuse_client
+
+            with lease_langfuse_client(self.Langfuse):
+                trace_id, generation_id = self._log_langfuse_v2(
+                    user_id=user_id,
+                    metadata=metadata,
+                    litellm_params=litellm_params,
+                    output=output,
+                    start_time=start_time,
+                    end_time=end_time,
+                    kwargs=kwargs,
+                    optional_params=optional_params,
+                    input=input,
+                    response_obj=response_obj,
+                    level=level,
+                    litellm_call_id=litellm_call_id,
+                )
             verbose_logger.debug("Langfuse Layer Logging - final response object: %s", response_obj)
             verbose_logger.info("Langfuse Layer Logging - logging success")
 

litellm/integrations/langfuse/langfuse_sdk.py

@@ -4,7 +4,8 @@ import os
 import re
 import threading
 from base64 import b64encode
-from collections.abc import Mapping
+from collections.abc import Generator, Mapping
+from contextlib import contextmanager
 from datetime import datetime
 from hashlib import sha256
 from types import MappingProxyType
@@ -28,6 +29,7 @@ __all__ = (
     "acquire_langfuse_client",
     "build_isolated_tracer_provider",
     "evict_stale_langfuse_resources",
+    "lease_langfuse_client",
     "open_trace_context",
     "propagate_attributes",
     "register_langfuse_client",
@@ -221,6 +223,137 @@ _LIVE_CLIENTS_LOCK: Final = threading.Lock()
 _live_clients: Final[WeakKeyDictionary[LangfuseResourceManager, WeakSet]] = WeakKeyDictionary()
 
 
+class _LangfuseLifecycleState:
+    """How many callbacks are leasing one SDK resource bundle, and what eviction has queued behind them.
+
+    ``lock`` is never held across a teardown, which takes the SDK's own registry lock.
+    """
+
+    def __init__(self) -> None:
+        self.lock = threading.Lock()
+        self.active_leases = 0
+        self.teardown_in_progress = False
+        self.teardown_owner: int | None = None
+        self.pending_clients: set[Langfuse] = set()  # mutable-ok: eviction and callback threads queue into it
+
+    def open_lease(self) -> None:
+        with self.lock:
+            self.active_leases += 1
+
+    def claim_for_teardown(self, client: Langfuse) -> bool:
+        """Whether this thread owns ``client``'s teardown; a lease or another teardown in flight queues it instead."""
+        with self.lock:
+            if self.active_leases > 0 or self.teardown_in_progress:
+                self.pending_clients.add(client)
+                return False
+            self.teardown_in_progress = True
+            self.teardown_owner = threading.get_ident()
+            return True
+
+    def release_lease(self) -> tuple[Langfuse, ...]:
+        """Drop this lease and take ownership of the teardowns it was holding up, if it was the last one."""
+        with self.lock:
+            self.active_leases -= 1
+            if self.active_leases > 0 or self.teardown_in_progress or not self.pending_clients:
+                return ()
+            claimed: Final = tuple(self.pending_clients)
+            self.pending_clients.clear()
+            self.teardown_in_progress = True
+            self.teardown_owner = threading.get_ident()
+            return claimed
+
+    def next_teardown_batch(self) -> tuple[Langfuse, ...]:
+        """Whatever eviction queued while the last batch was draining, handing the ownership flag back when empty."""
+        with self.lock:
+            if self.active_leases == 0 and self.pending_clients:
+                claimed: Final = tuple(self.pending_clients)
+                self.pending_clients.clear()
+                return claimed
+            self.teardown_in_progress = False
+            self.teardown_owner = None
+            return ()
+
+    def requeue(self, clients: tuple[Langfuse, ...]) -> None:
+        with self.lock:
+            self.pending_clients.update(clients)
+
+    def end_teardown(self) -> None:
+        with self.lock:
+            if self.teardown_owner == threading.get_ident():
+                self.teardown_in_progress = False
+                self.teardown_owner = None
+
+
+_LIFECYCLE_STATES_LOCK: Final = threading.Lock()
+_LIFECYCLE_STATES: Final[WeakKeyDictionary[object, _LangfuseLifecycleState]] = WeakKeyDictionary()
+
+
+def _lifecycle_state(client: Langfuse) -> _LangfuseLifecycleState:
+    """One state per resource bundle, since teardown closes the provider every client on that bundle exports through."""
+    resources: Final = getattr(client, "_resources", None)
+    key: Final = client if resources is None else resources
+    with _LIFECYCLE_STATES_LOCK:
+        existing: Final = _LIFECYCLE_STATES.get(key)
+        if existing is not None:
+            return existing
+        created: Final = _LangfuseLifecycleState()
+        _LIFECYCLE_STATES[key] = created
+        return created
+
+
+@contextmanager
+def lease_langfuse_client(client: Langfuse) -> Generator[None]:
+    """Hold off cache eviction's teardown of ``client`` while the export inside is in flight.
+
+    Eviction reaches a client the cache handed a callback moments earlier, so closing the SDK client
+    and its tracer provider there drops the spans that callback is still writing. The lease protects
+    exactly the window it wraps: an eviction arriving inside it is deferred to the last lease exit.
+    Taking a lease never blocks; a teardown already running keeps running, because the spans of a
+    lease taken that late were lost before the lease began, and stalling every other callback in the
+    process would not bring them back. A client the registry hands out during the deferral registers
+    as a holder, and the reference count keeps its bundle alive from there.
+    """
+    state: Final = _lifecycle_state(client)
+    state.open_lease()
+    try:
+        yield
+    finally:
+        _run_teardowns(state, state.release_lease(), propagate_base_exception=False)
+
+
+def _run_teardowns(
+    state: _LangfuseLifecycleState,
+    clients: tuple[Langfuse, ...],
+    *,
+    propagate_base_exception: bool = True,
+) -> None:
+    """Tear down ``clients``, then whatever eviction queued meanwhile, and hand the flag back.
+
+    A failing ordinary teardown is logged and skipped rather than raised: the thread here is usually a
+    request callback that merely held the last lease, and its request must not fail on eviction's behalf.
+    Interrupts requeue the unfinished batch and normally propagate, while a callback exception already
+    in flight takes precedence over an eviction interrupt.
+    """
+    batch = clients  # rebind-ok: drains each batch queued while the previous one was being torn down
+    try:
+        from litellm._logging import verbose_logger
+
+        while batch:
+            for index, client in enumerate(batch):
+                try:
+                    _teardown_langfuse_client(client)
+                except Exception:
+                    verbose_logger.exception("Langfuse client teardown failed during cache eviction")
+                except BaseException:
+                    state.requeue(batch[index:])
+                    if propagate_base_exception:
+                        raise
+                    return
+            batch = state.next_teardown_batch()
+    finally:
+        state.end_teardown()
+
+
 def _evict_if_stale_locked(
     *, public_key: object, secret_key: object, base_url: object
 ) -> LangfuseResourceManager | None:
@@ -398,6 +531,19 @@ def shutdown_langfuse_client(client: Langfuse) -> None:
     tracer provider's export thread running and leaves the client in the
     registry, so a later request for the same key gets a dead client back.
 
+    A callback holding a lease on the client's bundle postpones all of this to
+    the moment that lease ends, so eviction cannot close the provider out from
+    under an export the lease is wrapping. See ``lease_langfuse_client``.
+    """
+    state: Final = _lifecycle_state(client)
+    if not state.claim_for_teardown(client):
+        return
+    _run_teardowns(state, (client,))
+
+
+def _teardown_langfuse_client(client: Langfuse) -> None:
+    """The blocking teardown behind ``shutdown_langfuse_client``.
+
     A client that shares its resources with another live client only flushes:
     shutting the shared provider down here would silence the other client for
     the rest of its life, as it did before the reference count existed.

tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py

@@ -6,6 +6,7 @@ call would otherwise record its own duration instead of the call's.
 """
 
 import json
+import threading
 from datetime import datetime, timedelta, timezone
 
 import opentelemetry.trace as otel_trace
@@ -25,9 +26,12 @@ from litellm.integrations.langfuse.langfuse_sdk import (
     AS_ROOT_ATTRIBUTE,
     PUBLIC_ATTRIBUTE,
     RELEASE_ATTRIBUTE,
+    _lifecycle_state,
     _litellm_built_providers,
+    _teardown_langfuse_client,
     build_isolated_tracer_provider,
     evict_stale_langfuse_resources,
+    lease_langfuse_client,
     open_trace_context,
     register_langfuse_client,
     resolve_observation_id,
@@ -102,7 +106,12 @@ def test_guardrail_span_with_float_timestamps_does_not_break_the_generation(clie
     context, claim_root = open_trace_context(client=lf, trace_id="9" * 32, parent_observation_id=None)
     guardrail_start = 1709294400.0
     start_child_span(
-        client=lf, context=context, name="guardrail", start_time=guardrail_start, claim_trace_root=claim_root, attributes={}
+        client=lf,
+        context=context,
+        name="guardrail",
+        start_time=guardrail_start,
+        claim_trace_root=claim_root,
+        attributes={},
     ).end(end_time=to_unix_nanos(guardrail_start + 2))
     start_generation(
         client=lf, context=context, name="gen", start_time=CALL_START, claim_trace_root=claim_root, attributes={}
@@ -143,7 +152,12 @@ def test_child_span_keeps_its_own_window_and_stays_a_sibling(client):
     context, claim_root = open_trace_context(client=lf, trace_id="d" * 32, parent_observation_id=None)
     guardrail_start = CALL_START + timedelta(seconds=1)
     start_child_span(
-        client=lf, context=context, name="guardrail", start_time=guardrail_start, claim_trace_root=claim_root, attributes={}
+        client=lf,
+        context=context,
+        name="guardrail",
+        start_time=guardrail_start,
+        claim_trace_root=claim_root,
+        attributes={},
     ).end(end_time=to_unix_nanos(guardrail_start + timedelta(seconds=2)))
     start_generation(
         client=lf, context=context, name="gen", start_time=CALL_START, claim_trace_root=claim_root, attributes={}
@@ -550,6 +564,153 @@ def test_evicting_a_client_that_shares_resources_keeps_the_other_exporting():
     assert LangfuseResourceManager._instances.get(PUBLIC_KEY) is first._resources
 
 
+def test_eviction_defers_teardown_until_active_callback_finishes():
+    """A cached client must keep exporting while its callback lease is active."""
+    exporter = InMemorySpanExporter()
+    provider = build_isolated_tracer_provider(environment=None, release=None)
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    client = Langfuse(
+        public_key=PUBLIC_KEY,
+        secret_key="sk-original",
+        host="http://127.0.0.1:1",
+        tracer_provider=provider,
+        span_exporter=exporter,
+    )
+    register_langfuse_client(client)
+
+    def evict() -> None:
+        shutdown_langfuse_client(client)
+
+    with lease_langfuse_client(client):
+        evictor = threading.Thread(target=evict)
+        evictor.start()
+        evictor.join(timeout=5)
+        assert not evictor.is_alive()
+        assert not exporter._stopped
+
+        context, claim_root = open_trace_context(client=client, trace_id="a" * 32, parent_observation_id=None)
+        start_generation(
+            client=client,
+            context=context,
+            name="active-callback",
+            start_time=None,
+            claim_trace_root=claim_root,
+            attributes={},
+        ).end()
+        client.flush()
+        assert any(span.name == "active-callback" for span in exporter.get_finished_spans())
+
+    evictor.join(timeout=5)
+    assert not evictor.is_alive()
+    assert exporter._stopped
+    assert LangfuseResourceManager._instances.get(PUBLIC_KEY) is not client._resources
+
+
+def test_teardown_failure_does_not_strand_queued_clients(monkeypatch):
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    first = Langfuse(
+        public_key=PUBLIC_KEY,
+        secret_key="sk-original",
+        host="http://127.0.0.1:1",
+        tracer_provider=provider,
+        span_exporter=exporter,
+    )
+    clients = (
+        first,
+        Langfuse(public_key=PUBLIC_KEY, secret_key="sk-original", host="http://127.0.0.1:1"),
+        Langfuse(public_key=PUBLIC_KEY, secret_key="sk-original", host="http://127.0.0.1:1"),
+    )
+    assert len({client._resources for client in clients}) == 1
+    for client in clients:
+        register_langfuse_client(client)
+    state = _lifecycle_state(clients[0])
+    original_teardown = _teardown_langfuse_client
+    calls = []
+
+    def teardown(client):
+        calls.append(client)
+        original_teardown(client)
+        if len(calls) == 1:
+            raise RuntimeError("teardown failed")
+
+    monkeypatch.setattr("litellm.integrations.langfuse.langfuse_sdk._teardown_langfuse_client", teardown)
+    with lease_langfuse_client(clients[0]):
+        for client in clients:
+            shutdown_langfuse_client(client)
+
+    assert len(calls) == 3
+    assert not state.pending_clients
+    assert not state.teardown_in_progress
+
+
+def test_queued_eviction_waits_for_the_last_of_two_overlapping_leases():
+    exporter = InMemorySpanExporter()
+    provider = build_isolated_tracer_provider(environment=None, release=None)
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    client = Langfuse(
+        public_key=PUBLIC_KEY,
+        secret_key="sk-original",
+        host="http://127.0.0.1:1",
+        tracer_provider=provider,
+        span_exporter=exporter,
+    )
+    register_langfuse_client(client)
+
+    with lease_langfuse_client(client):
+        with lease_langfuse_client(client):
+            shutdown_langfuse_client(client)
+        assert not exporter._stopped
+
+    assert exporter._stopped
+
+
+def test_a_client_adopted_during_deferred_teardown_keeps_exporting():
+    """The registry hands the same bundle back out while its teardown is queued behind a lease;
+    the holder count must degrade that teardown to a flush."""
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    _litellm_built_providers.add(provider)
+    evicted = Langfuse(
+        public_key=PUBLIC_KEY,
+        secret_key="sk-original",
+        host="http://127.0.0.1:1",
+        tracer_provider=provider,
+        span_exporter=exporter,
+    )
+    register_langfuse_client(evicted)
+
+    with lease_langfuse_client(evicted):
+        shutdown_langfuse_client(evicted)
+        adopter = Langfuse(public_key=PUBLIC_KEY, secret_key="sk-original", host="http://127.0.0.1:1")
+        assert adopter._resources is evicted._resources
+        register_langfuse_client(adopter)
+
+    assert _exports(adopter, exporter, "after-deferred-teardown")
+    assert LangfuseResourceManager._instances.get(PUBLIC_KEY) is adopter._resources
+
+
+def test_leases_on_one_client_do_not_serialise_callbacks():
+    """Every langfuse callback in the process shares one client, so leases must overlap."""
+    client = _lifecycle_client()
+    both_inside = threading.Barrier(2, timeout=5)
+
+    def hold_lease() -> None:
+        with lease_langfuse_client(client):
+            both_inside.wait()
+
+    holders = tuple(threading.Thread(target=hold_lease) for _ in range(2))
+    for holder in holders:
+        holder.start()
+    for holder in holders:
+        holder.join(timeout=5)
+
+    assert not any(holder.is_alive() for holder in holders)
+    assert not both_inside.broken
+
+
 def test_last_client_on_shared_resources_tears_them_down():
     first, second, exporter = _shared_resources_pair()
     shutdown_langfuse_client(second)

tests/test_litellm/integrations/test_langfuse.py

@@ -10,7 +10,7 @@ import pytest
 import litellm
 from litellm.integrations.langfuse import langfuse as langfuse_module
 from litellm.integrations.langfuse.langfuse import LangFuseLogger
-from litellm.integrations.langfuse.langfuse_sdk import resolve_trace_id
+from litellm.integrations.langfuse.langfuse_sdk import _lifecycle_state, resolve_trace_id
 
 
 # Import LangfuseUsageDetails directly from the module where it's defined
@@ -1386,6 +1386,34 @@ def _steering_logger():
     return logger, exporter
 
 
+def test_log_event_holds_a_client_lease_during_export():
+    logger, _ = _steering_logger()
+    state = _lifecycle_state(logger.Langfuse)
+
+    def assert_lease_is_active(**_: object) -> tuple[str, str]:
+        assert state.active_leases == 1
+        return "trace-id", "generation-id"
+
+    now = datetime.datetime.now()
+    with patch.object(logger, "_log_langfuse_v2", side_effect=assert_lease_is_active):
+        returned = logger.log_event_on_langfuse(
+            kwargs={
+                "call_type": "completion",
+                "litellm_params": {"metadata": {}},
+                "messages": [{"role": "user", "content": "the-input"}],
+                "optional_params": {},
+            },
+            response_obj=litellm.ModelResponse(
+                choices=[{"message": {"role": "assistant", "content": "the-output"}}]
+            ),
+            start_time=now,
+            end_time=now,
+        )
+
+    assert returned == {"trace_id": "trace-id", "generation_id": "generation-id"}
+    assert state.active_leases == 0
+
+
 def _emit(rig, *, metadata=None, headers=None):
     """``log_event_on_langfuse`` is the entry point that folds ``langfuse_*`` headers into metadata.