feat: propagate trace attributes onto all child spans on update

[hassiebp]

Important

Introduces correlation_context for trace attribute propagation to child spans, deprecating update_current_trace and update_trace.

• Behavior:
  • Adds correlation_context() in client.py to propagate trace attributes to child spans.
  • Deprecates update_current_trace() in client.py and update_trace() in span.py.
• Span Processing:
  • Updates LangfuseSpanProcessor in span_processor.py to propagate correlation context to spans.
• Utilities:
  • Adds get_attribute_key_from_correlation_context() in utils.py for attribute key mapping.
• Tests:
  • Adds tests in test_core_sdk.py for correlation_context() propagation of user_id, session_id, and metadata.

^{This description was created by for cf3ec39. You can customize this summary. It will automatically update as commits are pushed.}

Disclaimer: Experimental PR review

Greptile Overview

Updated On: 2025-10-07 17:24:05 UTC

Summary

This PR introduces a new `correlation_context` feature that enables automatic propagation of trace-level attributes (user_id, session_id, metadata) to all child spans within a given scope, replacing the manual approach of calling `update_trace()` on each span individually.

Key Changes Made:

New Context Manager API: A new langfuse.correlation_context() context manager allows developers to set trace attributes once and have them automatically applied to all nested spans and generations within that scope.

OpenTelemetry Integration: The implementation leverages OpenTelemetry's context propagation mechanisms, supporting both local context and baggage propagation modes. The baggage mode enables cross-service attribute propagation in distributed tracing scenarios.

Automatic Span Processing: The LangfuseSpanProcessor now includes an on_start() method that automatically reads correlation context from both OpenTelemetry baggage and local context, then applies relevant attributes to new spans as they're created.

Deprecation Path: The existing update_current_trace() and span.update_trace() methods are deprecated with clear warnings directing users to the new correlation context approach.

Support Infrastructure: New utility functions and constants support the attribute mapping between correlation context keys and OpenTelemetry span attributes, with special handling for well-known attributes like user_id and session_id.

This change addresses a common pain point in complex AI application tracing where maintaining consistent context across nested operations (like agent workflows, tool calls, etc.) was previously manual and error-prone. The new API follows OpenTelemetry best practices and provides a more ergonomic developer experience.

Important Files Changed

Changed Files

Filename	Score	Overview
langfuse/_client/client.py	4/5	Implements the new correlation_context context manager with OpenTelemetry integration and deprecates update_current_trace
langfuse/_client/span_processor.py	4/5	Adds automatic attribute propagation in on_start method to inherit correlation context on child spans
tests/test_core_sdk.py	4/5	Adds comprehensive test coverage for correlation context functionality including nested contexts and baggage propagation
langfuse/_client/utils.py	4/5	Introduces mapping utilities for correlation context to span attribute conversion
langfuse/_client/span.py	5/5	Adds deprecation warning for update_trace method with minor import reorganization
langfuse/_client/constants.py	5/5	Adds correlation context key constant for OpenTelemetry context storage
langfuse/_client/attributes.py	5/5	Minor formatting improvement removing unnecessary blank line in imports

Confidence score: 4/5

• This PR introduces a significant new feature with proper OpenTelemetry integration and comprehensive test coverage
• Score reflects the complexity of context propagation logic and potential edge cases in distributed tracing scenarios
• Pay close attention to the span processor implementation and baggage propagation behavior in production environments

Sequence Diagram

sequenceDiagram
    participant User
    participant Langfuse as Langfuse Client
    participant SpanProcessor as LangfuseSpanProcessor
    participant CorrelationContext as Correlation Context
    participant OTELSpan as OpenTelemetry Span
    participant API as Langfuse API

    User->>Langfuse: start_as_current_span(name="parent")
    Langfuse->>OTELSpan: create parent span
    
    User->>Langfuse: correlation_context({"user_id": "123", "session_id": "456"})
    Langfuse->>CorrelationContext: set context values
    CorrelationContext->>CorrelationContext: store in OTEL context
    
    User->>Langfuse: start_span(name="child")
    Langfuse->>OTELSpan: create child span
    SpanProcessor->>SpanProcessor: on_start()
    SpanProcessor->>CorrelationContext: get_value(LANGFUSE_CORRELATION_CONTEXT_KEY)
    CorrelationContext-->>SpanProcessor: return correlation context
    SpanProcessor->>SpanProcessor: process correlation context
    loop For each context key-value pair
        SpanProcessor->>SpanProcessor: get_attribute_key_from_correlation_context(key)
        SpanProcessor->>OTELSpan: set_attribute(attribute_key, value)
    end
    SpanProcessor->>SpanProcessor: propagate baggage attributes
    
    User->>Langfuse: flush()
    Langfuse->>SpanProcessor: export spans
    SpanProcessor->>API: send span data with propagated attributes
    
    Note over SpanProcessor: All child spans automatically inherit<br/>trace attributes from correlation context

Loading

[greptile-apps]

_{7 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}