fix: cancel media uploads when mask_otel_spans drops the batch

langfuse/_client/span_exporter.py

@@ -106,12 +106,14 @@ def __init__(
         self._mask_otel_spans = mask_otel_spans
 
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        enqueued_media_ids: set[str] = set()
         span_attributes = [
             (
                 span,
                 self._process_media_attributes(
                     span=span,
                     attributes=dict(span.attributes or {}),
+                    enqueued_media_ids=enqueued_media_ids,
                 ),
             )
             for span in spans
@@ -123,6 +125,14 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
             )
 
             if masked_span_attributes is None:
+                # The mask hook dropped the whole batch. The media bytes were
+                # already queued for upload during attribute processing, so
+                # cancel those uploads too -- "drops the whole export batch"
+                # must mean all of it, including the media behind the spans.
+                if enqueued_media_ids and self._media_manager is not None:
+                    self._media_manager.cancel_pending_uploads(
+                        media_ids=enqueued_media_ids
+                    )
                 return SpanExportResult.SUCCESS
 
             span_attributes = masked_span_attributes
@@ -144,7 +154,11 @@ def force_flush(self, timeout_millis: int = 30000) -> bool:
         return self._exporter.force_flush(timeout_millis=timeout_millis)
 
     def _process_media_attributes(
-        self, *, span: ReadableSpan, attributes: Dict[str, AttributeValue]
+        self,
+        *,
+        span: ReadableSpan,
+        attributes: Dict[str, AttributeValue],
+        enqueued_media_ids: Optional[set[str]] = None,
     ) -> Dict[str, AttributeValue]:
         if self._media_manager is None:
             return attributes
@@ -157,6 +171,7 @@ def _process_media_attributes(
                     span=span,
                     attribute_key=key,
                     value=value,
+                    enqueued_media_ids=enqueued_media_ids,
                 )
             except Exception as error:
                 langfuse_logger.warning(
@@ -175,12 +190,14 @@ def _process_media_attribute_value(
         span: ReadableSpan,
         attribute_key: str,
         value: AttributeValue,
+        enqueued_media_ids: Optional[set[str]] = None,
     ) -> AttributeValue:
         if isinstance(value, str):
             return self._process_media_string(
                 span=span,
                 attribute_key=attribute_key,
                 value=value,
+                enqueued_media_ids=enqueued_media_ids,
             )
 
         if _is_attribute_sequence(value):
@@ -193,6 +210,7 @@ def _process_media_attribute_value(
                         span=span,
                         attribute_key=attribute_key,
                         value=item,
+                        enqueued_media_ids=enqueued_media_ids,
                     )
                     if isinstance(item, str)
                     else item
@@ -208,6 +226,7 @@ def _process_media_string(
         span: ReadableSpan,
         attribute_key: str,
         value: str,
+        enqueued_media_ids: Optional[set[str]] = None,
     ) -> str:
         media_manager = cast(MediaManager, self._media_manager)
         field = _media_field_for_attribute(attribute_key)
@@ -221,6 +240,11 @@ def _process_media_string(
                 field=field,
             )
 
+            _collect_enqueued_media_ids(
+                value=processed_direct_value,
+                media_ids=enqueued_media_ids,
+            )
+
             direct_reference = _media_reference_string(processed_direct_value)
 
             if direct_reference is not None:
@@ -253,6 +277,11 @@ def _process_media_string(
             field=field,
         )
 
+        _collect_enqueued_media_ids(
+            value=processed_json_value,
+            media_ids=enqueued_media_ids,
+        )
+
         if processed_json_value == parsed_value:
             return value
 
@@ -563,6 +592,32 @@ def _media_reference_string(value: Any) -> Optional[str]:
     return value._reference_string
 
 
+def _collect_enqueued_media_ids(value: Any, media_ids: Optional[set[str]]) -> None:
+    """Collect media IDs enqueued for upload from a processed attribute value.
+
+    _find_and_process_media enqueues an upload job for every LangfuseMedia it
+    builds, but the processed value it returns is what the attribute becomes.
+    Walk it so the exporter knows which media uploads a batch triggered, and can
+    cancel them if the mask hook later drops the batch.
+    """
+    if media_ids is None:
+        return
+
+    if isinstance(value, LangfuseMedia):
+        if value._media_id is not None:
+            media_ids.add(value._media_id)
+        return
+
+    if isinstance(value, dict):
+        for item in value.values():
+            _collect_enqueued_media_ids(item, media_ids)
+        return
+
+    if isinstance(value, (list, tuple)):
+        for item in value:
+            _collect_enqueued_media_ids(item, media_ids)
+
+
 def _media_upload_failed_marker(content_type: Any) -> str:
     return f"<Upload handling failed for LangfuseMedia of type {content_type}>"
 

langfuse/_task_manager/media_manager.py

@@ -45,6 +45,7 @@ def __init__(
         self._enabled = os.environ.get(
             LANGFUSE_MEDIA_UPLOAD_ENABLED, "True"
         ).lower() not in ("false", "0")
+        self._cancelled_media_ids: set[str] = set()
 
     def reinitialize(
         self,
@@ -65,6 +66,14 @@ def process_next_media_upload(self) -> None:
                 self._queue.task_done()
                 return
 
+            if upload_job["media_id"] in self._cancelled_media_ids:
+                self._cancelled_media_ids.discard(upload_job["media_id"])
+                logger.debug(
+                    f"Media: Skipping cancelled upload for media_id={upload_job['media_id']} in trace_id={upload_job['trace_id']}"
+                )
+                self._queue.task_done()
+                return
+
             logger.debug(
                 f"Media: Processing upload for media_id={upload_job['media_id']} in trace_id={upload_job['trace_id']}"
             )
@@ -79,6 +88,16 @@ def process_next_media_upload(self) -> None:
             )
             self._queue.task_done()
 
+    def cancel_pending_uploads(self, *, media_ids: set[str]) -> None:
+        """Cancel media uploads enqueued for a batch that was dropped.
+
+        The mask hook runs after media attributes are processed, so upload jobs
+        for a batch are already queued by the time the hook decides to drop it.
+        Mark those media IDs so the consumer skips them instead of uploading the
+        bytes the hook just redacted.
+        """
+        self._cancelled_media_ids.update(media_ids)
+
     def signal_shutdown(self, *, count: int = 1) -> None:
         for _ in range(count):
             try:

tests/unit/test_mask_otel_spans.py

@@ -196,6 +196,50 @@ def mask_otel_spans(*, params: MaskOtelSpansParams):
     assert not media_queue.empty()
 
 
+def test_mask_otel_spans_drop_batch_cancels_enqueued_media_uploads():
+    # Media attributes are processed before the mask hook runs, so upload jobs
+    # are already queued when the hook drops the batch. The dropped batch must
+    # not upload the media behind it -- the consumer should skip those jobs.
+    exporter = InMemorySpanExporter()
+    media_manager, media_queue = _media_manager()
+    image_base64 = base64.b64encode(b"image-bytes").decode("utf-8")
+    uploads_processed: list[str] = []
+
+    def mask_otel_spans(*, params):
+        # The fail-closed pattern: a hook error drops the whole export batch.
+        raise RuntimeError("masking function blew up")
+
+    original_process_upload = media_manager._process_upload_media_job
+
+    def recording_process_upload(*, data):
+        uploads_processed.append(data["media_id"])
+
+    media_manager._process_upload_media_job = recording_process_upload
+
+    provider = _tracer_provider(
+        exporter=exporter,
+        media_manager=media_manager,
+        mask_otel_spans=mask_otel_spans,
+    )
+    tracer = provider.get_tracer("openinference.instrumentation.openai")
+
+    with tracer.start_as_current_span("third-party-media-span") as span:
+        span.set_attribute("gen_ai.prompt", f"data:image/jpeg;base64,{image_base64}")
+
+    provider.force_flush()
+
+    assert exporter.get_finished_spans() == []
+    assert not media_queue.empty()
+
+    # The consumer must skip the cancelled upload: drain the queue and verify
+    # nothing was actually uploaded.
+    media_manager.process_next_media_upload()
+    assert uploads_processed == []
+    assert media_queue.empty()
+
+    media_manager._process_upload_media_job = original_process_upload
+
+
 def test_export_stage_media_prefilter_skips_json_without_media_hints(monkeypatch):
     exporter = InMemorySpanExporter()
     media_manager, media_queue = _media_manager()