diff --git a/langfuse/_client/span_exporter.py b/langfuse/_client/span_exporter.py index bbdc569cf..8ce6d94f0 100644 --- a/langfuse/_client/span_exporter.py +++ b/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"" diff --git a/langfuse/_task_manager/media_manager.py b/langfuse/_task_manager/media_manager.py index 14dceec19..c06b7474e 100644 --- a/langfuse/_task_manager/media_manager.py +++ b/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: diff --git a/tests/unit/test_mask_otel_spans.py b/tests/unit/test_mask_otel_spans.py index 5d1d52ee4..26da5d2c9 100644 --- a/tests/unit/test_mask_otel_spans.py +++ b/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()