feat(span-processor): mark app root spans (#1651)

Author: hassiebp

langfuse/_client/attributes.py

@@ -59,6 +59,7 @@ class LangfuseOtelSpanAttributes:
 
     # Internal
     AS_ROOT = "langfuse.internal.as_root"
+    IS_APP_ROOT = "langfuse.internal.is_app_root"
 
     # Experiments
     EXPERIMENT_ID = "langfuse.experiment.id"

langfuse/_client/client.py

@@ -27,6 +27,7 @@
 
 import backoff
 import httpx
+from opentelemetry import context as otel_context_api
 from opentelemetry import trace as otel_trace_api
 from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
 from opentelemetry.sdk.trace.export import SpanExporter
@@ -66,7 +67,9 @@
 )
 from langfuse._client.propagation import (
     PropagatedExperimentAttributes,
+    _detach_context_token_safely,
     _propagate_attributes,
+    _set_langfuse_trace_id_in_baggage,
 )
 from langfuse._client.resource_manager import LangfuseResourceManager
 from langfuse._client.span import (
@@ -1178,39 +1181,54 @@ def _start_as_current_otel_span_with_processed_media(
             name=name,
             end_on_exit=end_on_exit if end_on_exit is not None else True,
         ) as otel_span:
+            baggage_token = None
+
+            if otel_span.is_recording():
+                context_with_app_root_claim = _set_langfuse_trace_id_in_baggage(
+                    trace_id=self._get_otel_trace_id(otel_span),
+                    context=otel_context_api.get_current(),
+                )
+                baggage_token = otel_context_api.attach(context_with_app_root_claim)
+
             span_class = self._get_span_class(
                 as_type or "generation"
             )  # default was "generation"
-            common_args = {
-                "otel_span": otel_span,
-                "langfuse_client": self,
-                "environment": self._environment,
-                "release": self._release,
-                "input": input,
-                "output": output,
-                "metadata": metadata,
-                "version": version,
-                "level": level,
-                "status_message": status_message,
-            }
 
-            if span_class in [
-                LangfuseGeneration,
-                LangfuseEmbedding,
-            ]:
-                common_args.update(
-                    {
-                        "completion_start_time": completion_start_time,
-                        "model": model,
-                        "model_parameters": model_parameters,
-                        "usage_details": usage_details,
-                        "cost_details": cost_details,
-                        "prompt": prompt,
-                    }
-                )
-            # For span-like types (span, agent, tool, chain, retriever, evaluator, guardrail), no generation properties needed
+            try:
+                common_args = {
+                    "otel_span": otel_span,
+                    "langfuse_client": self,
+                    "environment": self._environment,
+                    "release": self._release,
+                    "input": input,
+                    "output": output,
+                    "metadata": metadata,
+                    "version": version,
+                    "level": level,
+                    "status_message": status_message,
+                }
+
+                if span_class in [
+                    LangfuseGeneration,
+                    LangfuseEmbedding,
+                ]:
+                    common_args.update(
+                        {
+                            "completion_start_time": completion_start_time,
+                            "model": model,
+                            "model_parameters": model_parameters,
+                            "usage_details": usage_details,
+                            "cost_details": cost_details,
+                            "prompt": prompt,
+                        }
+                    )
+                # For span-like types (span, agent, tool, chain, retriever, evaluator, guardrail), no generation properties needed
+
+                yield span_class(**common_args)  # type: ignore[arg-type]
 
-            yield span_class(**common_args)  # type: ignore[arg-type]
+            finally:
+                if baggage_token is not None:
+                    _detach_context_token_safely(baggage_token)
 
     def _get_current_otel_span(self) -> Optional[otel_trace_api.Span]:
         current_span = otel_trace_api.get_current_span()

langfuse/_client/propagation.py

@@ -316,6 +316,9 @@ def _get_propagated_attributes_from_context(
     # Handle baggage
     baggage_entries = baggage.get_all(context=context)
     for baggage_key, baggage_value in baggage_entries.items():
+        if baggage_key == LANGFUSE_TRACE_ID_BAGGAGE_KEY:
+            continue
+
         if baggage_key.startswith(LANGFUSE_BAGGAGE_PREFIX):
             span_key = _get_span_key_from_baggage_key(baggage_key)
 
@@ -471,12 +474,44 @@ def _get_propagated_context_key(key: str) -> str:
 
 
 LANGFUSE_BAGGAGE_PREFIX = "langfuse_"
+LANGFUSE_TRACE_ID_BAGGAGE_KEY = "langfuse_trace_id"
 
 
 def _get_propagated_baggage_key(key: str) -> str:
     return f"{LANGFUSE_BAGGAGE_PREFIX}{key}"
 
 
+def _get_langfuse_trace_id_from_baggage(
+    context: otel_context_api.Context,
+) -> Optional[str]:
+    value = otel_baggage_api.get_baggage(
+        name=LANGFUSE_TRACE_ID_BAGGAGE_KEY,
+        context=context,
+    )
+
+    if value is None:
+        return None
+
+    return str(value).lower()
+
+
+def _set_langfuse_trace_id_in_baggage(
+    *,
+    trace_id: str,
+    context: otel_context_api.Context,
+) -> otel_context_api.Context:
+    normalized_trace_id = trace_id.lower()
+
+    if _get_langfuse_trace_id_from_baggage(context) == normalized_trace_id:
+        return context
+
+    return otel_baggage_api.set_baggage(
+        name=LANGFUSE_TRACE_ID_BAGGAGE_KEY,
+        value=normalized_trace_id,
+        context=context,
+    )
+
+
 def _get_span_key_from_baggage_key(key: str) -> Optional[str]:
     if not key.startswith(LANGFUSE_BAGGAGE_PREFIX):
         return None