fix: normalize frozen Langfuse spans and update SDK test usage

mobilerun/agent/utils/tracing_setup.py

@@ -200,9 +200,10 @@ def _fixed_encoder(obj):
             _fixed_encoder._mobilerun_pydantic_v2 = True
             _handler._encoder = _fixed_encoder
 
-        # STEP 4: Register preprocessing before Langfuse registers its exporter.
+        # STEP 4: Normalize private snapshots inside the SDK's processor chain.
         from mobilerun.telemetry.langfuse_processor import (
             LangfuseSpanProcessor,
+            _LangfuseTracerProvider,
             set_current_agent,
         )
 
@@ -214,7 +215,6 @@ def _fixed_encoder(obj):
             or _langfuse_tracer_provider is not tracer_provider
         ):
             _langfuse_preprocessor = LangfuseSpanProcessor()
-            tracer_provider.add_span_processor(_langfuse_preprocessor)
             _langfuse_tracer_provider = tracer_provider
 
         # STEP 5: The public client owns the single exporter, queue, media uploads,
@@ -228,7 +228,9 @@ def _fixed_encoder(obj):
             public_key=public_key,
             secret_key=secret_key,
             base_url=base_url,
-            tracer_provider=tracer_provider,
+            tracer_provider=_LangfuseTracerProvider(
+                tracer_provider, _langfuse_preprocessor
+            ),
             should_export_span=_export_all_spans,
         )
 
@@ -287,27 +289,32 @@ def _export_all_spans(_span) -> bool:
 
 
 def _provider_has_langfuse_processor(tracer_provider: object) -> bool:
-    """Detect an exporter installed before Mobilerun's required preprocessor."""
+    """Detect both directly registered and wrapped Langfuse processors."""
+    from mobilerun.telemetry.langfuse_processor import _LangfuseSpanProcessor
+
     return any(
-        type(processor).__module__.startswith("langfuse.")
+        isinstance(processor, _LangfuseSpanProcessor)
+        or type(processor).__module__.startswith("langfuse.")
         for processor in _provider_span_processors(tracer_provider)
     )
 
 
 def _provider_has_owned_langfuse_pipeline(tracer_provider: object) -> bool:
-    """Confirm exactly one SDK exporter follows Mobilerun's preprocessor."""
-    processors = _provider_span_processors(tracer_provider)
-    try:
-        preprocessor_index = processors.index(_langfuse_preprocessor)
-    except ValueError:
-        return False
+    """Confirm exactly one SDK processor uses our snapshot normalizer."""
+    from mobilerun.telemetry.langfuse_processor import _LangfuseSpanProcessor
 
-    exporter_indexes = [
-        index
-        for index, processor in enumerate(processors)
-        if type(processor).__module__.startswith("langfuse.")
+    processors = [
+        processor
+        for processor in _provider_span_processors(tracer_provider)
+        if isinstance(processor, _LangfuseSpanProcessor)
+        or type(processor).__module__.startswith("langfuse.")
     ]
-    return len(exporter_indexes) == 1 and preprocessor_index < exporter_indexes[0]
+    return (
+        len(processors) == 1
+        and isinstance(processors[0], _LangfuseSpanProcessor)
+        and processors[0].normalizer is _langfuse_preprocessor
+        and type(processors[0].processor).__module__.startswith("langfuse.")
+    )
 
 
 def _provider_span_processors(tracer_provider: object) -> tuple[object, ...]:

mobilerun/telemetry/langfuse_processor.py

@@ -1,17 +1,18 @@
 """OpenTelemetry span preprocessing for the Langfuse integration.
 
-Langfuse owns span export, batching, and media upload. This processor runs
-before the Langfuse processor and only normalizes Mobilerun/OpenInference spans
-into the attributes understood by Langfuse.
+Langfuse owns span export, batching, and media upload. Its processor receives
+a normalized export snapshot; other processors retain the original ended span.
 """
 
 import base64
 import json
 import logging
 from contextvars import ContextVar
+from copy import copy
 from typing import TYPE_CHECKING, Any, Optional
 
 from opentelemetry import trace
+from opentelemetry.attributes import BoundedAttributes
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 
@@ -62,12 +63,67 @@ def get_last_step_span_context() -> Optional[Context]:
     return _last_step_span_context.get()
 
 
+class _LangfuseSpanProcessor(SpanProcessor):
+    """Wrap the SDK processor without adding an exporter, queue, or uploader."""
+
+    def __init__(self, processor: SpanProcessor, normalizer: "LangfuseSpanProcessor"):
+        self.processor = processor
+        self.normalizer = normalizer
+
+    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+        self.normalizer.on_start(span, parent_context)
+        self.processor.on_start(span, parent_context)
+
+    def _on_ending(self, span: Span) -> None:
+        on_ending = getattr(self.processor, "_on_ending", None)
+        if on_ending is not None:
+            on_ending(span)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        # Ended spans can already be frozen before _on_ending runs. A shallow
+        # copy retains all span metadata, including bounded events/links and
+        # their dropped counts; only attributes need independent writable storage.
+        snapshot = copy(span)
+        snapshot._attributes = BoundedAttributes(
+            attributes=dict(span.attributes or {}),
+            immutable=False,
+            extended_attributes=getattr(
+                span._attributes, "_extended_attributes", False
+            ),
+        )
+        snapshot._attributes.dropped = span.dropped_attributes
+        self.normalizer.on_end(snapshot)
+        snapshot._attributes._immutable = True
+        self.processor.on_end(snapshot)
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return self.processor.force_flush(timeout_millis)
+
+    def shutdown(self) -> None:
+        self.processor.shutdown()
+
+
+class _LangfuseTracerProvider:
+    """Let the public Langfuse client register its processor on our provider."""
+
+    def __init__(self, provider, normalizer: "LangfuseSpanProcessor"):
+        self._provider = provider
+        self._normalizer = normalizer
+
+    def add_span_processor(self, processor: SpanProcessor) -> None:
+        self._provider.add_span_processor(
+            _LangfuseSpanProcessor(processor, self._normalizer)
+        )
+
+    def __getattr__(self, name):
+        return getattr(self._provider, name)
+
+
 class LangfuseSpanProcessor(SpanProcessor):
-    """Normalize spans before Langfuse's public OTel exporter sees them.
+    """Enrich live spans and normalize writable Langfuse export snapshots.
 
     This processor deliberately does not export, batch, upload, or own threads.
-    Register it before constructing the public :class:`langfuse.Langfuse`
-    client so the client's processor receives the normalized span.
+    Use it through _LangfuseTracerProvider so on_end only receives private copies.
     """
 
     def __init__(self, agent: Optional["MobileAgent"] = None) -> None:

tests/test_grok_api.py

@@ -74,9 +74,12 @@ def _xai_completed_response(*, usage: ResponseUsage) -> Response:
 
 
 def _xai_usage() -> ResponseUsage:
+    input_details = {"cached_tokens": 0}
+    if "cache_write_tokens" in InputTokensDetails.model_fields:
+        input_details["cache_write_tokens"] = 0
     return ResponseUsage(
         input_tokens=7,
-        input_tokens_details=InputTokensDetails(cached_tokens=0),
+        input_tokens_details=InputTokensDetails(**input_details),
         output_tokens=4,
         output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
         total_tokens=11,