diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index ca7108412f0..70bd32f19b4 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -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. diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py index 1902c4a55e5..6f8394fd870 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py @@ -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