fix(langfuse): retire the provider a credential rotation abandons

Rotating the secret or host for one public key evicted the registry entry but left the tracer provider litellm built for it running, so repeated rotations leaked an export thread each time. The evicted bundle is now retired after the registry lock is released, and only when no live client is still on it.
This commit is contained in:
Yucheng Zhu 2026-09-01 14:39:06 -07:00
parent d4a1cd312f
commit 442d85a23c
2 changed files with 80 additions and 17 deletions

View file

@ -213,25 +213,51 @@ class DiscardingSpanExporter(SpanExporter):
return True
_LIVE_CLIENTS_LOCK: Final = threading.Lock()
# litellm clients still using each SDK resource bundle; the bundle is torn down with the last one.
# Both sides are weak so a throwaway client (a health probe, an alerting lookup) that is simply
# garbage-collected stops holding the bundle open rather than inflating a counter forever.
_live_clients: Final[WeakKeyDictionary[LangfuseResourceManager, WeakSet]] = WeakKeyDictionary()
def _evict_if_stale_locked(
*, public_key: object, secret_key: object, base_url: object
) -> LangfuseResourceManager | None:
"""Assumes ``LangfuseResourceManager._lock`` is held; returns the still-valid cached bundle, if any."""
) -> tuple[LangfuseResourceManager | None, LangfuseResourceManager | None]:
"""Assumes ``LangfuseResourceManager._lock`` is held; returns the still-valid bundle and the evicted one."""
if not public_key:
return None
return None, None
cached: Final = LangfuseResourceManager._instances.get(public_key) # pyright: ignore[reportPrivateUsage] # registry has no public accessor
if cached is None:
return None
return None, None
if getattr(cached, "secret_key", None) == secret_key and getattr(cached, "base_url", None) == base_url:
return cached
LangfuseResourceManager._instances.pop(public_key, None) # pyright: ignore[reportPrivateUsage] # registry has no public accessor
return None
return cached, None
return None, LangfuseResourceManager._instances.pop(public_key, None) # pyright: ignore[reportPrivateUsage] # registry has no public accessor
def _shutdown_abandoned_provider(resources: LangfuseResourceManager | None) -> None:
"""Retire the provider litellm built for an evicted bundle, once no live client is still on it.
Runs after the registry lock is released: provider shutdown flushes and joins the export
thread, which must never happen under the lock every other client init contends on. A
client still holding these resources tears them down itself in ``shutdown_langfuse_client``.
"""
if resources is None:
return
with _LIVE_CLIENTS_LOCK:
holders: Final = _live_clients.get(resources)
if holders is not None and len(holders) > 0:
return
_live_clients.pop(resources, None)
provider: Final = getattr(resources, "tracer_provider", None)
if provider is not None and provider in _litellm_built_providers:
provider.shutdown()
def evict_stale_langfuse_resources(*, public_key: str | None, secret_key: str | None, base_url: str | None) -> None:
"""Drop a cached client whose credentials no longer match the ones being requested."""
with LangfuseResourceManager._lock: # pyright: ignore[reportPrivateUsage] # registry has no public accessor
_evict_if_stale_locked(public_key=public_key, secret_key=secret_key, base_url=base_url)
_, abandoned = _evict_if_stale_locked(public_key=public_key, secret_key=secret_key, base_url=base_url)
_shutdown_abandoned_provider(abandoned)
def _build_verified_span_exporter(*, public_key: object, secret_key: object, base_url: object) -> SpanExporter | None:
@ -300,7 +326,7 @@ def acquire_langfuse_client(
)
)
with LangfuseResourceManager._lock: # pyright: ignore[reportPrivateUsage] # registry has no public accessor
cached: Final = _evict_if_stale_locked(
cached, abandoned = _evict_if_stale_locked(
public_key=public_key,
secret_key=parameters.get("secret_key"),
base_url=parameters.get("base_url"),
@ -313,16 +339,10 @@ def acquire_langfuse_client(
span_exporter=span_exporter,
)
register_langfuse_client(client)
_shutdown_abandoned_provider(abandoned)
return client
_LIVE_CLIENTS_LOCK: Final = threading.Lock()
# litellm clients still using each SDK resource bundle; the bundle is torn down with the last one.
# Both sides are weak so a throwaway client (a health probe, an alerting lookup) that is simply
# garbage-collected stops holding the bundle open rather than inflating a counter forever.
_live_clients: Final[WeakKeyDictionary[LangfuseResourceManager, WeakSet]] = WeakKeyDictionary()
def register_langfuse_client(client: Langfuse) -> None:
"""Track the client against the SDK resources it ended up with.

View file

@ -23,9 +23,9 @@ from litellm.integrations.langfuse.langfuse import (
)
from litellm.integrations.langfuse.langfuse_sdk import (
AS_ROOT_ATTRIBUTE,
_litellm_built_providers,
PUBLIC_ATTRIBUTE,
RELEASE_ATTRIBUTE,
_litellm_built_providers,
build_isolated_tracer_provider,
evict_stale_langfuse_resources,
open_trace_context,
@ -572,6 +572,49 @@ def test_shutdown_of_a_stale_client_does_not_deregister_the_live_one():
assert LangfuseResourceManager._instances.get(PUBLIC_KEY) is live._resources
def _rotation_provider():
"""A litellm-built provider on the lifecycle public key, exporting in memory."""
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)
assert _exports(client, exporter, "before-rotation")
return client, exporter
def test_rotation_retires_the_provider_no_client_is_left_on():
"""Prompt management builds throwaway clients from request credentials.
Alternating the secret for one public key evicts a bundle nobody holds any
more, and its export thread has to go with it or every rotation leaks one.
"""
import gc
client, exporter = _rotation_provider()
del client
gc.collect()
evict_stale_langfuse_resources(public_key=PUBLIC_KEY, secret_key="sk-rotated", base_url="http://127.0.0.1:2")
assert exporter._stopped
def test_rotation_keeps_a_still_live_client_exporting():
"""The evicted bundle is only retired when nothing is on it; a live logger must survive."""
client, exporter = _rotation_provider()
evict_stale_langfuse_resources(public_key=PUBLIC_KEY, secret_key="sk-rotated", base_url="http://127.0.0.1:2")
assert _exports(client, exporter, "after-rotation")
def test_ssl_exporter_is_only_built_with_custom_tls_material(monkeypatch, tmp_path):
"""v4 exports over its own OTLP channel, so litellm's CA bundle must be rebuilt onto it."""
import litellm