feat(masking): batch and async span masking by iamkalio · Pull Request #1553 · langfuse/langfuse-python

iamkalio · 2026-03-05T23:26:09Z

Summary

First step toward batch/async span masking. Adds types, config, and the span wrapper; does not yet wire batch masking into the export path.

Changes

Types:BatchMaskFunction in langfuse/types.py
Config:batch_mask, mask_batch_size, use_async_masking on Langfuse and LangfuseResourceManager
Wrapper:MaskedAttributeSpanWrapper in langfuse/_client/span_processor.py to overlay masked attributes on spans before export
Tests:TestMaskedAttributeSpanWrapper in tests/test_otel.py

Follow-up (not in this PR)

Masking exporter that uses the wrapper and batch_mask
Skip sync mask when batch masking is enabled
Optional async masking (background thread)

Fixeslangfuse/langfuse#11931

Disclaimer: Experimental PR review

Greptile Summary

This PR introduces the infrastructure for batch and async span masking in the Langfuse Python SDK. It adds a BatchMaskFunction protocol type, three new configuration parameters (batch_mask, mask_batch_size, use_async_masking) threaded through Langfuse, LangfuseResourceManager, and LangfuseSpanProcessor, and the core MaskingSpanExporter + MaskedAttributeSpanWrapper classes in a new masking_exporter.py module. When batch_mask is configured, the OTLP exporter is wrapped so that span attributes are masked in batches (and optionally off the main thread) before being serialized and sent.

Key findings:

Critical:LangfuseOtelSpanAttributes.TRACE_METADATA is present in the codebase and set on trace spans, but is missing from _MASKABLE_ATTR_KEYS in masking_exporter.py. Trace-level metadata will not be passed to batch_mask, leaving PII potentially exposed.
Logic:force_flush waits up to timeout_sec for the worker's flush event and then calls self._span_exporter.force_flush(timeout_millis) with the full original timeout — callers can experience up to 2x the requested timeout.
Logic: The shared _flush_event is not safe for concurrent force_flush callers; one thread clearing the event can race with the worker signaling completion for another thread's flush.
Style: A dead assignment (chunk_backref) is computed but never referenced in the chunked batch masking loop.
Import rule: All five new test methods place from unittest.mock import MagicMock and the masking exporter imports inside the method body instead of at the module level.

Confidence Score: 2/5

Not safe to merge as-is due to a data-masking gap (TRACE_METADATA unmasked) and concurrency issues in force_flush.
The missing TRACE_METADATA key is a correctness/privacy bug — trace-level metadata will bypass batch masking entirely. The force_flush double-timeout and shared-event race condition are logic bugs that can cause data loss in async mode under concurrent flushes. These should be fixed before the masking exporter is wired into the export path.
Pay close attention to langfuse/_client/masking_exporter.py — specifically _MASKABLE_ATTR_KEYS, force_flush, and the concurrent flush handling in _worker_loop.

Important Files Changed

Filename	Overview
langfuse/_client/masking_exporter.py	New file implementing batch masking exporter; has a critical omission (TRACE_METADATA not in maskable keys), a double-timeout bug in force_flush, a concurrent flush race condition with the shared _flush_event, and a dead assignment (chunk_backref).
langfuse/_client/span_processor.py	Adds batch_mask, mask_batch_size, and use_async_masking parameters and wraps the OTLP exporter in MaskingSpanExporter when batch_mask is set; changes are minimal and correct.
langfuse/_client/span.py	Adds a guard in _mask_attribute to skip sync masking when batch_mask is set, deferring masking to export time; logic is correct.
langfuse/types.py	Adds BatchMaskFunction Protocol and updates all; well-documented and correctly typed.
langfuse/_client/client.py	Threads batch_mask, mask_batch_size, and use_async_masking through Langfuse constructor and into LangfuseResourceManager; straightforward plumbing with no issues.
langfuse/_client/resource_manager.py	Adds batch masking parameters to LangfuseResourceManager and forwards them to LangfuseSpanProcessor; also adds a tracer_provider.shutdown() call on teardown which is a useful improvement.
langfuse/_client/get_client.py	Passes batch_mask, mask_batch_size, and use_async_masking through _create_client_from_instance; correct and complete.
tests/test_otel.py	Good test coverage for MaskedAttributeSpanWrapper and MaskingSpanExporter (sync and async paths); imports are placed inside test methods instead of at the module level, violating the project import rule.

Sequence Diagram

sequenceDiagram
participant App
participant LangfuseSpanProcessor
participant MaskingSpanExporter
participant WorkerThread
participant OTLPSpanExporter
App->>LangfuseSpanProcessor: on_end(span)
LangfuseSpanProcessor->>LangfuseSpanProcessor: filter / batch
LangfuseSpanProcessor->>MaskingSpanExporter: export(spans)
alt Sync masking (use_async_masking=False)
MaskingSpanExporter->>MaskingSpanExporter: _collect_maskable_items(spans)
MaskingSpanExporter->>MaskingSpanExporter: batch_mask(items) [in chunks if mask_batch_size set]
MaskingSpanExporter->>MaskingSpanExporter: wrap each span in MaskedAttributeSpanWrapper
MaskingSpanExporter->>OTLPSpanExporter: export(wrapped_spans)
OTLPSpanExporter-->>MaskingSpanExporter: SpanExportResult
MaskingSpanExporter-->>LangfuseSpanProcessor: SpanExportResult
else Async masking (use_async_masking=True)
MaskingSpanExporter->>MaskingSpanExporter: queue.put(spans)
MaskingSpanExporter-->>LangfuseSpanProcessor: SUCCESS (immediate)
WorkerThread->>WorkerThread: _collect_maskable_items + batch_mask + wrap
WorkerThread->>OTLPSpanExporter: export(wrapped_spans)
end
App->>MaskingSpanExporter: force_flush(timeout_millis)
alt Async mode
MaskingSpanExporter->>MaskingSpanExporter: queue.put(_FLUSH_SENTINEL)
WorkerThread->>WorkerThread: drain queue, export remaining spans
WorkerThread->>MaskingSpanExporter: flush_event.set()
MaskingSpanExporter->>OTLPSpanExporter: force_flush(timeout_millis)
else Sync mode
MaskingSpanExporter->>OTLPSpanExporter: force_flush(timeout_millis)
end

_{Last reviewed commit: 9b21a71}

Greptile also left 3 inline comments on this PR.

_{(2/5) Greptile learns from your feedback when you react with thumbs up/down!}

Context used:

Rule used - Move imports to the top of the module instead of p... (source)

Learnt From
langfuse/langfuse-python#1387

langfuse/langfuse#11931) Foundation for batch/async span masking. Adds BatchMaskFunction and config (batch_mask, mask_batch_size, use_async_masking), and MaskedAttributeSpanWrapper in span_processor. Follow-up commits will wire the masking exporter and optional async path. Ref: langfuse/langfuse#11931

CLAassistant · 2026-03-05T23:26:18Z

All committers have signed the CLA.

…ut down

iamkalio · 2026-03-07T00:42:13Z

Hello @hassiebp

I opened a PR with a proposed approach to this issue. Before continuing and adding documentation, I wanted to confirm this aligns with the direction you had in mind.

greptile-apps · 2026-03-07T09:15:59Z

+_MASKABLE_ATTR_KEYS = (
+ LangfuseOtelSpanAttributes.TRACE_INPUT,
+ LangfuseOtelSpanAttributes.TRACE_OUTPUT,
+ LangfuseOtelSpanAttributes.OBSERVATION_INPUT,
+ LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT,
+ LangfuseOtelSpanAttributes.OBSERVATION_METADATA,
+)


TRACE_METADATA missing from maskable keys
LangfuseOtelSpanAttributes.TRACE_METADATA ("langfuse.trace.metadata") exists and is set on trace-level spans (confirmed in attributes.py line 173), but it is absent from _MASKABLE_ATTR_KEYS. This means trace metadata will never be passed to batch_mask, leaving potentially sensitive data unmasked even when batch masking is configured.
Suggested change
_MASKABLE_ATTR_KEYS= (
LangfuseOtelSpanAttributes.TRACE_INPUT,
LangfuseOtelSpanAttributes.TRACE_OUTPUT,
LangfuseOtelSpanAttributes.OBSERVATION_INPUT,
LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT,
LangfuseOtelSpanAttributes.OBSERVATION_METADATA,
)
_MASKABLE_ATTR_KEYS= (
LangfuseOtelSpanAttributes.TRACE_INPUT,
LangfuseOtelSpanAttributes.TRACE_OUTPUT,
LangfuseOtelSpanAttributes.TRACE_METADATA,
LangfuseOtelSpanAttributes.OBSERVATION_INPUT,
LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT,
LangfuseOtelSpanAttributes.OBSERVATION_METADATA,
)

greptile-apps · 2026-03-07T09:16:00Z

+ timeout_sec = (timeout_millis / 1000.0) if timeout_millis is not None else 10.0
+ if not self._flush_event.wait(timeout=timeout_sec):
+ return False
+ if hasattr(self._span_exporter, "force_flush"):
+ return self._span_exporter.force_flush(timeout_millis)


force_flush can wait up to 2x the requested timeout
After waiting up to timeout_sec for the worker's flush event, the code immediately calls self._span_exporter.force_flush(timeout_millis) with the full original timeout. In the worst case, a caller requesting a 5-second flush could wait up to 10 seconds before the call returns. The inner force_flush should receive the remaining time budget.
Suggested change
timeout_sec= (timeout_millis/1000.0) iftimeout_millisisnotNoneelse10.0
ifnotself._flush_event.wait(timeout=timeout_sec):
returnFalse
ifhasattr(self._span_exporter, "force_flush"):
returnself._span_exporter.force_flush(timeout_millis)
timeout_sec= (timeout_millis/1000.0) iftimeout_millisisnotNoneelse10.0
start=time.monotonic()
ifnotself._flush_event.wait(timeout=timeout_sec):
returnFalse
ifhasattr(self._span_exporter, "force_flush"):
elapsed_ms=int((time.monotonic() -start) *1000)
remaining_ms=max(0, timeout_millis-elapsed_ms) iftimeout_millisisnotNoneelseNone
returnself._span_exporter.force_flush(remaining_ms)
returnTrue
Note: this also requires adding import time at the top of the module.

greptile-apps · 2026-03-07T09:16:01Z

+ if self._queue is None or self._flush_event is None or self._closed:
+ return True
+
+ self._flush_event.clear()
+ try:
+ self._queue.put(_FLUSH_SENTINEL, timeout=self._QUEUE_PUT_TIMEOUT_SEC)
+ except Full:
+ return False
+ timeout_sec = (timeout_millis / 1000.0) if timeout_millis is not None else 10.0
+ if not self._flush_event.wait(timeout=timeout_sec):
+ return False
+ if hasattr(self._span_exporter, "force_flush"):
+ return self._span_exporter.force_flush(timeout_millis)


Shared _flush_event causes a race condition under concurrent flushes
_flush_event is a single threading.Event shared by all callers of force_flush. If two threads call force_flush concurrently, the sequence below can cause the first caller to time out even though the flush completed:
Thread A calls _flush_event.clear() and enqueues _FLUSH_SENTINEL.
Worker processes Thread A's sentinel and calls _flush_event.set().
Thread B calls _flush_event.clear() — this clears the event before Thread A calls _flush_event.wait().
Thread A's _flush_event.wait() now blocks until its own timeout expires.
Consider using per-flush threading.Event objects (e.g. by passing the event alongside the sentinel in the queue) so each caller waits only on their own flush completion signal.

AqueelAhmedV · 2026-03-07T10:39:46Z

@iamkalio looks good! we have something similar running in production (wrapping the internal exporter post Langfuse() init), good to see this supported natively.

OndrejSlama · 2026-04-21T14:57:15Z

Any update on this?

manul-ai · 2026-05-06T10:15:53Z

Hi, it would be really helpful to have it merged

iamkalio changed the title ~~feat(masking): add batch masking config and MaskedAttributeSpanWrapper (langfuse/langfuse#11931)~~feat(masking): batch and async span masking (langfuse/langfuse#11931)Mar 5, 2026

iamkalio changed the title ~~feat(masking): batch and async span masking (langfuse/langfuse#11931)~~feat(masking): batch and async span maskingMar 5, 2026

iamkalio added 4 commits March 7, 2026 00:20

added batch masking exporter and wire into span processor
15159f9

skip sync mask when batch_mask is set
5d4f215

run batch masking in background thread when use_async_masking is True
c5b61f2

ensure shutdown calls tracer provider so async masking exporter is sh…
9b21a71
…ut down

iamkalio marked this pull request as ready for review March 7, 2026 09:12

greptile-appsBot reviewed Mar 7, 2026
View reviewed changes

hassiebp mentioned this pull request May 8, 2026
fix(openai): handle Azure stream chunks without delta #1648
Merged

Navigation Menu

Uh oh!

feat(masking): batch and async span masking - #1553

feat(masking): batch and async span masking#1553
iamkalio wants to merge 5 commits into
langfuse:mainfrom
iamkalio:fix/pyhton-sdk-masking

iamkalio commented Mar 5, 2026•
edited by greptile-apps Bot
Loading

Uh oh!

CLAassistant commented Mar 5, 2026•
edited
Loading

Uh oh!

iamkalio commented Mar 7, 2026

Uh oh!

greptile-appsBotMar 7, 2026

Uh oh!

greptile-appsBotMar 7, 2026

Uh oh!

greptile-appsBotMar 7, 2026

Uh oh!

AqueelAhmedV commented Mar 7, 2026

Uh oh!

OndrejSlama commented Apr 21, 2026

Uh oh!

manul-ai commented May 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

-timeout_sec= (timeout_millis/1000.0) iftimeout_millisisnotNoneelse10.0
-ifnotself._flush_event.wait(timeout=timeout_sec):
-returnFalse
-ifhasattr(self._span_exporter, "force_flush"):
-returnself._span_exporter.force_flush(timeout_millis)
+timeout_sec= (timeout_millis/1000.0) iftimeout_millisisnotNoneelse10.0
+start=time.monotonic()
+ifnotself._flush_event.wait(timeout=timeout_sec):
+returnFalse
+ifhasattr(self._span_exporter, "force_flush"):
+elapsed_ms=int((time.monotonic() -start) *1000)
+remaining_ms=max(0, timeout_millis-elapsed_ms) iftimeout_millisisnotNoneelseNone
+returnself._span_exporter.force_flush(remaining_ms)
+returnTrue

Uh oh!

Conversation

iamkalio commented Mar 5, 2026• edited by greptile-apps BotLoading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Follow-up (not in this PR)

Disclaimer: Experimental PR review

Greptile Summary

Confidence Score: 2/5

Important Files Changed

Sequence Diagram

Uh oh!

CLAassistant commented Mar 5, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

iamkalio commented Mar 7, 2026

Uh oh!

greptile-appsBotMar 7, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-appsBotMar 7, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-appsBotMar 7, 2026

Choose a reason for hiding this comment

Uh oh!

AqueelAhmedV commented Mar 7, 2026

Uh oh!

OndrejSlama commented Apr 21, 2026

Uh oh!

manul-ai commented May 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

iamkalio commented Mar 5, 2026•
edited by greptile-apps Bot
Loading

CLAassistant commented Mar 5, 2026•
edited
Loading