fix(resource-manager): prevent stale clients from queueing after shutdown

langfuse/_client/resource_manager.py

@@ -134,10 +134,10 @@ def __new__(
         id_generator: Optional[IdGenerator] = None,
         span_exporter: Optional[SpanExporter] = None,
     ) -> "LangfuseResourceManager":
-        if public_key in cls._instances:
-            return cls._instances[public_key]
-
         with cls._lock:
+            if public_key in cls._instances:
+                return cls._instances[public_key]
+
             if public_key not in cls._instances:
                 instance = super(LangfuseResourceManager, cls).__new__(cls)
 
@@ -209,6 +209,8 @@ def _initialize_instance(
         self.mask_otel_spans = mask_otel_spans
         self.environment = environment
         self._shutdown = False
+        self._shutdown_event = threading.Event()
+        self._admission_lock = threading.Lock()
 
         # Store additional client settings for get_client() to use
         self.timeout = timeout
@@ -420,6 +422,18 @@ def _at_fork_reinit(self) -> None:
         # the lock is class-level state needed by the child (e.g. to create a new client)
         # even if this particular instance was already shut down.
         LangfuseResourceManager._lock = threading.RLock()
+        # Recreate shutdown event and admission lock for child process
+        if hasattr(self, "_shutdown_event"):
+            self._shutdown_event = threading.Event()
+            if self._shutdown:
+                self._shutdown_event.set()
+        if hasattr(self, "_admission_lock"):
+            self._admission_lock = threading.Lock()
+        # Recreate media manager state lock if held at fork
+        if hasattr(self, "_media_manager") and hasattr(
+            self._media_manager, "_state_lock"
+        ):
+            self._media_manager._state_lock = threading.Lock()
 
         if self._shutdown:
             return
@@ -508,10 +522,16 @@ def add_score_task(self, event: dict, *, force_sample: bool = False) -> None:
             )
 
             if should_sample:
-                langfuse_logger.debug(
-                    f"Score: Enqueuing event type={event['type']} for trace_id={event['body'].trace_id} name={event['body'].name} value={event['body'].value}"
-                )
-                self._score_ingestion_queue.put(event, block=False)
+                with self._admission_lock:
+                    if self._shutdown:
+                        langfuse_logger.warning(
+                            "Score: Dropping event because the Langfuse client has already been shut down."
+                        )
+                        return
+                    langfuse_logger.debug(
+                        f"Score: Enqueuing event type={event['type']} for trace_id={event['body'].trace_id} name={event['body'].name} value={event['body'].value}"
+                    )
+                    self._score_ingestion_queue.put(event, block=False)
 
         except Full:
             langfuse_logger.warning(
@@ -531,10 +551,16 @@ def add_trace_task(
         event: dict,
     ) -> None:
         try:
-            langfuse_logger.debug(
-                f"Trace: Enqueuing event type={event['type']} for trace_id={event['body'].id}"
-            )
-            self._score_ingestion_queue.put(event, block=False)
+            with self._admission_lock:
+                if self._shutdown:
+                    langfuse_logger.warning(
+                        "Trace: Dropping event because the Langfuse client has already been shut down."
+                    )
+                    return
+                langfuse_logger.debug(
+                    f"Trace: Enqueuing event type={event['type']} for trace_id={event['body'].id}"
+                )
+                self._score_ingestion_queue.put(event, block=False)
 
         except Full:
             langfuse_logger.warning(
@@ -612,13 +638,40 @@ def flush(self) -> None:
         langfuse_logger.debug("Successfully flushed media upload queue")
 
     def shutdown(self) -> None:
-        self._shutdown = True
+        is_first = False
+        with self._admission_lock:
+            if self._shutdown:
+                # Another thread already started shutdown; wait for it
+                pass
+            else:
+                self._shutdown = True
+                is_first = True
+
+        if not is_first:
+            # Wait indefinitely for the first shutdown to complete.
+            # A timeout would hide an incomplete shutdown and violate
+            # the invariant that shutdown() only returns after resources
+            # are drained and consumers stopped.
+            self._shutdown_event.wait()
+            return
 
         # Unregister the atexit handler first
-        atexit.unregister(self.shutdown)
+        try:
+            atexit.unregister(self.shutdown)
+        except Exception:
+            pass
 
-        self.flush()
-        self._stop_and_join_consumer_threads()
+        try:
+            self.flush()
+        finally:
+            try:
+                # Block further media enqueues that could be stranded
+                # after flush but before consumers are stopped. Media discovered
+                # during the preceding force_flush() was already enqueued.
+                self._media_manager.begin_shutdown()
+                self._stop_and_join_consumer_threads()
+            finally:
+                self._shutdown_event.set()
 
 
 def _init_tracer_provider(

langfuse/_task_manager/media_manager.py

@@ -1,4 +1,5 @@
 import os
+import threading
 import time
 from queue import Empty, Full, Queue
 from typing import Any, Callable, Optional, TypeVar, cast
@@ -45,6 +46,8 @@ def __init__(
         self._enabled = os.environ.get(
             LANGFUSE_MEDIA_UPLOAD_ENABLED, "True"
         ).lower() not in ("false", "0")
+        self._shutdown = False
+        self._state_lock = threading.Lock()
 
     def reinitialize(
         self,
@@ -53,9 +56,15 @@ def reinitialize(
         httpx_client: httpx.Client,
         media_upload_queue: Queue,
     ) -> None:
-        self._api_client = api_client
-        self._httpx_client = httpx_client
-        self._queue = media_upload_queue
+        with self._state_lock:
+            self._api_client = api_client
+            self._httpx_client = httpx_client
+            self._queue = media_upload_queue
+            self._shutdown = False
+
+    def begin_shutdown(self) -> None:
+        with self._state_lock:
+            self._shutdown = True
 
     def process_next_media_upload(self) -> None:
         try:
@@ -98,6 +107,12 @@ def _find_and_process_media(
     ) -> Any:
         if not self._enabled:
             return data
+        with self._state_lock:
+            if self._shutdown:
+                logger.warning(
+                    "Media: Skipping upload because the Langfuse client has already been shut down."
+                )
+                return data
 
         seen = set()
         max_levels = 10
@@ -279,10 +294,16 @@ def _process_media(
                 field=field,
             )
 
-            self._queue.put(
-                item=upload_media_job,
-                block=False,
-            )
+            with self._state_lock:
+                if self._shutdown:
+                    logger.warning(
+                        f"Media: Skipping upload for media_id={media._media_id} because the Langfuse client has already been shut down."
+                    )
+                    return
+                self._queue.put(
+                    item=upload_media_job,
+                    block=False,
+                )
             logger.debug(
                 f"Queue: Enqueued media ID {media._media_id} for upload processing | trace_id={trace_id} | field={field}"
             )

tests/unit/test_media.py

@@ -263,6 +263,7 @@ def test_resolve_media_references_uses_configured_httpx_client():
         "https://example.com/test.jpg", timeout=fetch_timeout_seconds
     )
 
+
 def test_init_with_urlsafe_base64_data_uri():
     original_bytes = b"\xfb\xff"
     urlsafe_base64 = base64.urlsafe_b64encode(original_bytes).decode()
@@ -275,3 +276,78 @@ def test_init_with_urlsafe_base64_data_uri():
     assert media._content_type == "application/octet-stream"
     assert media._content_bytes == original_bytes
 
+
+def test_1799_media_admission_toctou_and_begin_shutdown():
+    """Media check+put must be atomic – no job stranded after begin_shutdown."""
+    import threading
+    from queue import Queue
+    from unittest.mock import Mock
+
+    from langfuse._task_manager.media_manager import MediaManager
+
+    q: Queue = Queue()
+    mm = MediaManager(
+        api_client=Mock(),
+        httpx_client=Mock(),
+        media_upload_queue=q,
+    )
+
+    # Stub upload so queue drain is instant
+    mm._process_upload_media_job = Mock(return_value=None)  # type: ignore[attr-defined]
+
+    # Deterministic barrier: block inside queue.put while holding _state_lock
+    orig_put = q.put
+    block_in_put = threading.Event()
+    release_put = threading.Event()
+
+    def blocking_put(*args, **kwargs):
+        block_in_put.set()
+        assert release_put.wait(timeout=2)
+        return orig_put(*args, **kwargs)
+
+    q.put = blocking_put  # type: ignore[assignment]
+
+    # Create a valid media object
+    media = LangfuseMedia(content_bytes=b"hello", content_type="text/plain")
+    assert media._media_id is not None
+
+    producer = threading.Thread(
+        target=lambda: mm._process_media(
+            media=media, trace_id="0" * 32, observation_id="0" * 16, field="input"
+        )
+    )
+    producer.start()
+    assert block_in_put.wait(timeout=2)
+
+    # begin_shutdown must wait for the put's lock, not interleave and strand
+    shutdown_done = threading.Event()
+
+    def do_shutdown():
+        mm.begin_shutdown()
+        shutdown_done.set()
+
+    t_shutdown = threading.Thread(target=do_shutdown)
+    t_shutdown.start()
+    # shutdown should be blocked while producer holds _state_lock inside put
+    assert not shutdown_done.wait(timeout=0.3)
+
+    release_put.set()
+    producer.join(timeout=2)
+    t_shutdown.join(timeout=2)
+
+    # Producer's job was enqueued atomically; simulate flush drain via task_done
+    assert q.unfinished_tasks == 1
+    assert q.qsize() == 1
+    # Drain as flush would (join would wait for task_done)
+    q.get()
+    q.task_done()
+    q.join()
+    assert q.unfinished_tasks == 0
+
+    # After shutdown, further admission is dropped, not stranded
+    media2 = LangfuseMedia(content_bytes=b"world", content_type="text/plain")
+    mm._process_media(
+        media=media2, trace_id="0" * 32, observation_id="0" * 16, field="input"
+    )
+    assert q.unfinished_tasks == 0
+    assert q.qsize() == 0