langfuse-rb 0.10.0 → 0.11.0

data/lib/langfuse/prompt_cache_coordinator.rb

@@ -0,0 +1,288 @@
+# frozen_string_literal: true
+
+require_relative "prompt_fetch_result"
+require_relative "prompt_cache_events"
+
+module Langfuse
+  # Coordinates prompt fetch/cache behavior between the API transport and the
+  # configured cache backend. Both supported backends ({PromptCache} and
+  # {RailsCacheAdapter}) provide the full cache + SWR surface; only
+  # {RailsCacheAdapter} adds distributed-lock fetch, which is the one branch
+  # the dispatch needs to make.
+  #
+  # @api private
+  class PromptCacheCoordinator # rubocop:disable Metrics/ClassLength
+    # @param cache [PromptCache, RailsCacheAdapter, nil] Configured cache backend
+    # @param event_emitter [#emit_prompt_cache_event] Emitter for cache events
+    # @param fetch_prompt [#call] Callable that fetches prompt data from the API
+    # @return [PromptCacheCoordinator]
+    def initialize(cache:, event_emitter:, fetch_prompt:)
+      @cache = cache
+      @event_emitter = event_emitter
+      @fetch_prompt = fetch_prompt
+      @backend_name = compute_backend_name
+    end
+
+    # @return [String] Backend identifier reported in events and stats
+    attr_reader :backend_name
+
+    # Fetch a prompt and include cache metadata.
+    #
+    # @param name [String] Prompt name
+    # @param version [Integer, nil] Optional prompt version
+    # @param label [String, nil] Optional prompt label
+    # @param cache_ttl [Integer, nil] Optional TTL override (0 forces a bypass)
+    # @return [PromptFetchResult] Prompt data plus cache metadata
+    def get_prompt_result(name, version: nil, label: nil, cache_ttl: nil)
+      validate_fetch_options!(version, label, cache_ttl)
+      key = prompt_cache_key(name, version: version, label: label)
+
+      return fetch_uncached(key, CacheStatus::DISABLED) if @cache.nil?
+      return fetch_uncached(key, CacheStatus::BYPASS) if cache_ttl&.zero?
+
+      fetch_cached(key, cache_ttl)
+    end
+
+    # Refresh a prompt from the API, optionally writing through to cache.
+    #
+    # @param name [String] Prompt name
+    # @param version [Integer, nil] Optional prompt version
+    # @param label [String, nil] Optional prompt label
+    # @param cache_ttl [Integer, nil] Optional TTL override
+    # @return [PromptFetchResult] Prompt data plus cache metadata
+    def refresh_prompt(name, version: nil, label: nil, cache_ttl: nil)
+      validate_fetch_options!(version, label, cache_ttl)
+      key = prompt_cache_key(name, version: version, label: label)
+
+      emit(:refresh_start) { event_payload(key, CacheStatus::REFRESH, CacheSource::API) }
+      prompt_data = @fetch_prompt.call(name, version: version, label: label)
+      write_through(key, prompt_data, cache_ttl, status: CacheStatus::REFRESH) if @cache && !cache_ttl&.zero?
+      status = refresh_status(cache_ttl)
+      emit(:refresh_success) { event_payload(key, status, CacheSource::API) }
+      build_result(key, prompt_data, status, CacheSource::API)
+    rescue StandardError => e
+      emit(:refresh_failure) do
+        event_payload(key, CacheStatus::REFRESH, CacheSource::API,
+                      error_class: e.class.name, error_message: e.message)
+      end
+      raise
+    end
+
+    # Inspect the logical and generated cache keys for a prompt.
+    #
+    # @param name [String] Prompt name
+    # @param version [Integer, nil] Optional prompt version
+    # @param label [String, nil] Optional prompt label
+    # @return [PromptCacheKey] Logical and generated cache keys
+    def prompt_cache_key(name, version: nil, label: nil)
+      raise ArgumentError, "Cannot specify both version and label" if version && label
+
+      logical_key = PromptCache.build_key(name, version: version, label: label)
+      storage_key = @cache ? @cache.storage_key(logical_key, name: name) : logical_key
+      PromptCacheKey.new(name: name, version: version, label: label, logical_key: logical_key, storage_key: storage_key)
+    end
+
+    # Invalidate one exact logical prompt cache key.
+    #
+    # @param name [String] Prompt name
+    # @param version [Integer, nil] Optional prompt version
+    # @param label [String, nil] Optional prompt label
+    # @return [PromptCacheKey] Invalidated key
+    def invalidate_prompt_cache(name, version: nil, label: nil)
+      key = prompt_cache_key(name, version: version, label: label)
+      deleted = @cache ? @cache.delete(key.storage_key) : false
+      emit(:delete) { event_payload(key, CacheStatus::MISS, CacheSource::CACHE, deleted: deleted) }
+      emit(:invalidate) { event_payload(key, CacheStatus::MISS, CacheSource::CACHE, scope: :exact) }
+      key
+    end
+
+    # Invalidate all cached variants for one prompt name.
+    #
+    # @param name [String] Prompt name
+    # @return [Integer, nil] New generation, or nil when caching is disabled
+    def invalidate_prompt_cache_by_name(name)
+      emit_name_invalidation(name, mutation: false)
+    end
+
+    # Invalidate after prompt mutation (create/update). Distinct from manual
+    # invalidation so observers can tell the two apart.
+    #
+    # @param name [String] Prompt name
+    # @return [Integer, nil] New generation
+    def invalidate_after_mutation(name)
+      emit_name_invalidation(name, mutation: true)
+    end
+
+    # Logically clear the entire prompt cache namespace.
+    #
+    # @return [Integer, nil] New global generation, or nil when caching is disabled
+    def clear_prompt_cache
+      generation = @cache&.clear_logically
+      emit(:clear, backend: @backend_name, generation: generation)
+      generation
+    end
+
+    # @return [Hash] Prompt cache statistics
+    def prompt_cache_stats
+      @cache ? @cache.stats : disabled_stats
+    end
+
+    private
+
+    def validate_fetch_options!(version, label, cache_ttl)
+      raise ArgumentError, "Cannot specify both version and label" if version && label
+      return if cache_ttl.nil?
+      raise ArgumentError, "cache_ttl must be a non-negative Integer" unless cache_ttl.is_a?(Integer)
+      raise ArgumentError, "cache_ttl must be non-negative" if cache_ttl.negative?
+    end
+
+    def fetch_uncached(key, status)
+      prompt_data = @fetch_prompt.call(key.name, version: key.version, label: key.label)
+      build_result(key, prompt_data, status, CacheSource::API)
+    end
+
+    # Single dispatch: SWR > distributed lock > simple get/set.
+    def fetch_cached(key, cache_ttl)
+      return fetch_with_swr(key, cache_ttl) if @cache.swr_enabled?
+      return fetch_with_lock(key, cache_ttl) if @cache.is_a?(RailsCacheAdapter)
+
+      cached = @cache.get(key.storage_key)
+      return cache_hit(key, cached) if cached
+
+      fetch_and_cache(key, cache_ttl, swr: false)
+    end
+
+    def fetch_with_swr(key, cache_ttl)
+      entry = @cache.entry(key.storage_key)
+      return cache_hit(key, entry.data) if entry.respond_to?(:fresh?) && entry.fresh?
+
+      if entry.respond_to?(:stale?) && entry.stale?
+        emit(:stale_serve) { event_payload(key, CacheStatus::STALE, CacheSource::CACHE) }
+        schedule_refresh(key, cache_ttl)
+        return build_result(key, entry.data, CacheStatus::STALE, CacheSource::CACHE)
+      end
+
+      fetch_and_cache(key, cache_ttl, swr: true)
+    end
+
+    def fetch_with_lock(key, cache_ttl)
+      cached = @cache.get(key.storage_key)
+      return cache_hit(key, cached) if cached
+
+      emit(:miss) { event_payload(key, CacheStatus::MISS, CacheSource::API) }
+      fetched = false
+      prompt_data = @cache.fetch_with_lock(key.storage_key, ttl: cache_ttl) do
+        fetched = true
+        @fetch_prompt.call(key.name, version: key.version, label: key.label)
+      end
+      emit(:write) { event_payload(key, CacheStatus::MISS, CacheSource::API) } if fetched
+      status = fetched ? CacheStatus::MISS : CacheStatus::HIT
+      source = fetched ? CacheSource::API : CacheSource::CACHE
+      build_result(key, prompt_data, status, source)
+    end
+
+    def fetch_and_cache(key, cache_ttl, swr:)
+      emit(:miss) { event_payload(key, CacheStatus::MISS, CacheSource::API) }
+      prompt_data = @fetch_prompt.call(key.name, version: key.version, label: key.label)
+      write_through(key, prompt_data, cache_ttl, swr: swr)
+      build_result(key, prompt_data, CacheStatus::MISS, CacheSource::API)
+    end
+
+    def write_through(key, prompt_data, cache_ttl, swr: false, status: CacheStatus::MISS)
+      if swr
+        @cache.write_with_stale_while_revalidate(key.storage_key, prompt_data, ttl: cache_ttl)
+      else
+        @cache.set(key.storage_key, prompt_data, ttl: cache_ttl)
+      end
+      emit(:write) { event_payload(key, status, CacheSource::API) }
+    end
+
+    def cache_hit(key, prompt_data)
+      emit(:hit) { event_payload(key, CacheStatus::HIT, CacheSource::CACHE) }
+      build_result(key, prompt_data, CacheStatus::HIT, CacheSource::CACHE)
+    end
+
+    def schedule_refresh(key, cache_ttl)
+      scheduled = @cache.refresh_async(
+        key.storage_key,
+        ttl: cache_ttl,
+        on_success: ->(_value) { emit_refresh_success(key) },
+        on_failure: ->(error) { emit_refresh_failure(key, error) }
+      ) { @fetch_prompt.call(key.name, version: key.version, label: key.label) }
+      emit(:refresh_start) { event_payload(key, CacheStatus::STALE, CacheSource::CACHE) } if scheduled
+    end
+
+    def emit_refresh_success(key)
+      emit(:refresh_success) { event_payload(key, CacheStatus::REFRESH, CacheSource::API) }
+      emit(:write) { event_payload(key, CacheStatus::REFRESH, CacheSource::API) }
+    end
+
+    def emit_refresh_failure(key, error)
+      emit(:refresh_failure) do
+        event_payload(key, CacheStatus::STALE, CacheSource::CACHE,
+                      error_class: error.class.name, error_message: error.message)
+      end
+    end
+
+    def emit_name_invalidation(name, mutation:)
+      generation = @cache&.invalidate_name(name)
+      payload = { name: name, backend: @backend_name, generation: generation, scope: :name }
+      payload[:mutation] = true if mutation
+      emit(:invalidate, payload)
+      generation
+    end
+
+    def refresh_status(cache_ttl)
+      return CacheStatus::DISABLED unless @cache
+      return CacheStatus::BYPASS if cache_ttl&.zero?
+
+      CacheStatus::REFRESH
+    end
+
+    def build_result(key, prompt_data, status, source)
+      PromptFetchResult.new(
+        prompt: prompt_data,
+        logical_key: key.logical_key,
+        storage_key: key.storage_key,
+        cache_status: status,
+        source: source,
+        name: prompt_data["name"] || key.name,
+        version: prompt_data["version"] || key.version,
+        label: key.resolved_label
+      )
+    end
+
+    def emit(event, payload = nil, &)
+      @event_emitter.emit_prompt_cache_event(event, payload, &)
+    end
+
+    def event_payload(key, cache_status, source, **extra)
+      PromptCacheEvents.build_payload(
+        key,
+        cache_status: cache_status,
+        source: source,
+        backend: @backend_name,
+        extra: extra
+      )
+    end
+
+    def compute_backend_name
+      return CacheBackend::DISABLED unless @cache
+      return CacheBackend::RAILS if @cache.is_a?(RailsCacheAdapter)
+      return CacheBackend::MEMORY if @cache.is_a?(PromptCache)
+
+      @cache.class.name
+    end
+
+    def disabled_stats
+      {
+        backend: CacheBackend::DISABLED,
+        enabled: false,
+        current_generation_entries: nil,
+        orphaned_entries: nil,
+        total_entries: nil,
+        unsupported_counts: CacheBackend::UNSUPPORTED_COUNT_KEYS
+      }
+    end
+  end
+end

data/lib/langfuse/prompt_cache_events.rb

@@ -10,6 +10,29 @@ module Langfuse
     # ActiveSupport::Notifications event name used for prompt cache events.
     PROMPT_CACHE_NOTIFICATION = "prompt_cache.langfuse"
 
+    # Build a prompt cache event payload from a key, status, source, and backend.
+    # Shared by the ApiClient mixin and PromptCacheCoordinator so a payload-shape
+    # change can't drift between the two emitters.
+    #
+    # @param key [PromptCacheKey] Logical and storage cache key
+    # @param cache_status [Symbol] Cache status
+    # @param source [Symbol] Event source
+    # @param backend [String] Backend identifier
+    # @param extra [Hash] Additional payload fields
+    # @return [Hash] Event payload
+    def self.build_payload(key, cache_status:, source:, backend:, extra: {})
+      {
+        name: key.name,
+        version: key.version,
+        label: key.resolved_label,
+        logical_key: key.logical_key,
+        storage_key: key.storage_key,
+        backend: backend,
+        cache_status: cache_status,
+        source: source
+      }.merge(extra)
+    end
+
     # Configure prompt cache event dispatch. Wraps the observer once into a
     # 1-arg callable so the per-event hot path never re-checks arity.
     #
@@ -56,16 +79,13 @@ module Langfuse
 
     # @api private
     def event_payload(key, cache_status, source, extra = {})
-      {
-        name: key.name,
-        version: key.version,
-        label: key.resolved_label,
-        logical_key: key.logical_key,
-        storage_key: key.storage_key,
-        backend: cache_backend_name,
+      PromptCacheEvents.build_payload(
+        key,
         cache_status: cache_status,
-        source: source
-      }.merge(extra)
+        source: source,
+        backend: cache_backend_name,
+        extra: extra
+      )
     end
 
     # @api private
@@ -100,7 +120,8 @@ module Langfuse
     def wrap_cache_observer(observer)
       return nil if observer.nil?
 
-      if observer.method(:call).arity == 1
+      arity = observer.respond_to?(:arity) ? observer.arity : observer.method(:call).arity
+      if arity == 1
         ->(payload) { observer.call(payload) }
       else
         ->(payload) { observer.call(payload[:event], payload) }

data/lib/langfuse/prompt_variables.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+require "mustache"
+
+module Langfuse
+  # Extracts referenced variables from parsed Mustache templates.
+  #
+  # @api private
+  class PromptVariables
+    TAG_TYPES = %i[etag utag].freeze
+    SECTION_TYPES = %i[section inverted_section].freeze
+
+    class << self
+      # @api private
+      def extract(template)
+        tokens = Mustache::Template.new(template).tokens
+        collect(tokens, []).reject(&:empty?).uniq
+      end
+
+      private
+
+      def collect(tokens, scope)
+        tokens.each_with_object([]) do |token, variables|
+          next unless token.is_a?(Array)
+
+          variables.concat(token.first == :mustache ? from_tag(token, scope) : collect(token, scope))
+        end
+      end
+
+      def from_tag(token, scope)
+        return variable_path(token, scope) if TAG_TYPES.include?(token[1])
+        return section_paths(token, scope) if SECTION_TYPES.include?(token[1])
+
+        []
+      end
+
+      def variable_path(token, scope)
+        path = scoped_path(token, scope)
+        path.empty? ? [] : [path.join(".")]
+      end
+
+      def section_paths(token, scope)
+        section_path = scoped_path(token, scope)
+        body_scope = token[1] == :section ? section_path : scope
+        [section_path.join("."), *collect(token[4], body_scope)]
+      end
+
+      def scoped_path(token, scope)
+        segments = token.dig(2, 2)
+        segments == ["."] ? scope : scope + segments
+      end
+    end
+  end
+end

data/lib/langfuse/propagation.rb

@@ -6,7 +6,7 @@ module Langfuse
   # Attribute propagation utilities for Langfuse OpenTelemetry integration.
   #
   # This module provides the `propagate_attributes` method for setting trace-level
-  # attributes (user_id, session_id, metadata) that automatically propagate to all child spans
+  # attributes that automatically propagate to all child spans
   # within the context.
   #
   # @example Basic usage
@@ -21,29 +21,31 @@ module Langfuse
   #
   # rubocop:disable Metrics/ModuleLength
   module Propagation
+    # Baggage key prefix for cross-service propagation
+    BAGGAGE_PREFIX = "langfuse_"
+
+    # Baggage key that records which Langfuse trace already owns the application root
+    LANGFUSE_TRACE_ID_BAGGAGE_KEY = "#{BAGGAGE_PREFIX}trace_id".freeze
+
+    ENVIRONMENT_VALUE_PATTERN = /\A(?!langfuse)[a-z0-9_-]+\z/
+    private_constant :ENVIRONMENT_VALUE_PATTERN
+
     # Map of propagated attribute keys to span attribute keys
     SPAN_KEY_MAP = {
       "user_id" => OtelAttributes::TRACE_USER_ID,
       "session_id" => OtelAttributes::TRACE_SESSION_ID,
       "version" => OtelAttributes::VERSION,
       "tags" => OtelAttributes::TRACE_TAGS,
-      "metadata" => OtelAttributes::TRACE_METADATA
+      "metadata" => OtelAttributes::TRACE_METADATA,
+      "trace_name" => OtelAttributes::TRACE_NAME,
+      "release" => OtelAttributes::RELEASE,
+      "environment" => OtelAttributes::ENVIRONMENT
     }.freeze
 
     # OpenTelemetry context keys for propagated attributes
-    CONTEXT_KEYS = {
-      "user_id" => OpenTelemetry::Context.create_key("langfuse_user_id"),
-      "session_id" => OpenTelemetry::Context.create_key("langfuse_session_id"),
-      "metadata" => OpenTelemetry::Context.create_key("langfuse_metadata"),
-      "version" => OpenTelemetry::Context.create_key("langfuse_version"),
-      "tags" => OpenTelemetry::Context.create_key("langfuse_tags")
-    }.freeze
-
-    # List of propagated attribute keys (derived from CONTEXT_KEYS)
-    PROPAGATED_ATTRIBUTES = CONTEXT_KEYS.keys.freeze
-
-    # Baggage key prefix for cross-service propagation
-    BAGGAGE_PREFIX = "langfuse_"
+    CONTEXT_KEYS = SPAN_KEY_MAP.keys.to_h do |key|
+      [key, OpenTelemetry::Context.create_key("#{BAGGAGE_PREFIX}#{key}")]
+    end.freeze
 
     # Propagate trace-level attributes to all spans created within this context.
     #
@@ -57,6 +59,9 @@ module Langfuse
     # @param metadata [Hash<String, String>, nil] Additional metadata (all values ≤200 characters)
     # @param version [String, nil] Version identifier (≤200 characters)
     # @param tags [Array<String>, nil] List of tags (each ≤200 characters)
+    # @param trace_name [String, nil] Trace name (≤200 characters)
+    # @param release [String, nil] Release identifier (≤200 characters)
+    # @param environment [String, nil] Lowercase environment identifier (≤40 characters)
     # @param as_baggage [Boolean] If true, propagates via OpenTelemetry baggage for cross-service propagation
     # @yield Block within which attributes are propagated
     # @return [Object] The result of the block
@@ -76,32 +81,29 @@ module Langfuse
     #     # All spans inherit these attributes
     #   end
     #
+    # rubocop:disable Metrics/ParameterLists
     def self.propagate_attributes(user_id: nil, session_id: nil, metadata: nil, version: nil, tags: nil,
-                                  as_baggage: false, &block)
+                                  trace_name: nil, release: nil, environment: nil, as_baggage: false, &block)
       raise ArgumentError, "Block required" unless block
 
-      _propagate_attributes(
-        user_id: user_id,
-        session_id: session_id,
-        metadata: metadata,
-        version: version,
-        tags: tags,
-        as_baggage: as_baggage,
-        &block
-      )
+      attributes = {
+        "user_id" => user_id, "session_id" => session_id, "metadata" => metadata,
+        "version" => version, "tags" => tags, "trace_name" => trace_name,
+        "release" => release, "environment" => environment
+      }
+      _propagate_attributes(attributes, as_baggage: as_baggage, &block)
     end
 
     # Internal implementation of propagate_attributes
     #
+    # @param attributes [Hash<String, Object>] Propagated attribute values
+    # @param as_baggage [Boolean] Whether to add the values to OpenTelemetry baggage
     # @api private
-    def self._propagate_attributes(user_id: nil, session_id: nil, metadata: nil, version: nil, tags: nil,
-                                   as_baggage: false, &)
+    def self._propagate_attributes(attributes, as_baggage:, &)
       current_context = OpenTelemetry::Context.current
       current_span = OpenTelemetry::Trace.current_span
 
-      # Process each propagated attribute using PROPAGATED_ATTRIBUTES constant
-      PROPAGATED_ATTRIBUTES.each do |key|
-        value = binding.local_variable_get(key.to_sym)
+      attributes.each do |key, value|
         next if value.nil?
         next if key == "tags" && value.empty?
 
@@ -120,6 +122,7 @@ module Langfuse
       # Execute block in new context
       OpenTelemetry::Context.with_current(current_context, &)
     end
+    # rubocop:enable Metrics/ParameterLists
 
     # Validate an attribute value based on its type
     #
@@ -140,6 +143,8 @@ module Langfuse
           validated_metadata[k.to_s] = v.to_s if _validate_string_value(v, "metadata.#{k}")
         end
         validated_metadata.any? ? validated_metadata : nil
+      when "environment"
+        _validate_environment_value(value)
       else
         _validate_propagated_value(value, key)
       end
@@ -157,7 +162,7 @@ module Langfuse
       propagated_attributes = _extract_baggage_attributes(context)
 
       # Handle OTEL context values
-      PROPAGATED_ATTRIBUTES.each do |key|
+      SPAN_KEY_MAP.each_key do |key|
         context_key = _get_propagated_context_key(key)
         value = context.value(context_key)
 
@@ -165,7 +170,10 @@ module Langfuse
 
         span_key = _get_propagated_span_key(key)
 
-        if key == "metadata" && value.is_a?(Hash)
+        if key == "environment"
+          validated_environment = _validate_environment_value(value)
+          propagated_attributes[span_key] = validated_environment if validated_environment
+        elsif key == "metadata" && value.is_a?(Hash)
           value.each do |k, v|
             metadata_key = "#{OtelAttributes::TRACE_METADATA}.#{k}"
             propagated_attributes[metadata_key] = v.to_s
@@ -318,6 +326,29 @@ module Langfuse
     end
     # rubocop:enable Naming/PredicateMethod
 
+    # Validate a propagated environment value against the cross-SDK contract.
+    #
+    # @param value [Object] Environment value to validate
+    # @return [String, nil] Validated environment or nil
+    # @api private
+    def self._validate_environment_value(value)
+      return _drop_environment("value is not a string") unless value.is_a?(String)
+      return _drop_environment("value is over 40 characters (#{value.length} chars)") if value.length > 40
+
+      return value if ENVIRONMENT_VALUE_PATTERN.match?(value)
+
+      _drop_environment(
+        "must use lowercase letters, numbers, hyphens, or underscores and must not start with 'langfuse'"
+      )
+    end
+
+    def self._drop_environment(reason)
+      Langfuse.configuration.logger.warn(
+        "Langfuse: Propagated attribute 'environment' #{reason}. Dropping value."
+      )
+      nil
+    end
+
     # Get context key for a propagated attribute
     #
     # @param key [String] Attribute key (user_id, session_id, etc.)
@@ -378,6 +409,40 @@ module Langfuse
       defined?(OpenTelemetry::Baggage)
     end
 
+    # Get the Langfuse trace claim from OpenTelemetry baggage.
+    #
+    # @param context [OpenTelemetry::Context] Context to inspect
+    # @return [String, nil] Lowercase trace ID or nil
+    # @api private
+    def self._get_langfuse_trace_id_from_baggage(context)
+      return nil unless baggage_available?
+
+      OpenTelemetry::Baggage.values(context: context)[LANGFUSE_TRACE_ID_BAGGAGE_KEY]&.to_s&.downcase
+    rescue StandardError => e
+      Langfuse.configuration.logger.debug("Langfuse: Trace baggage read failed: #{e.message}")
+      nil
+    end
+
+    # Set the Langfuse trace claim in OpenTelemetry baggage.
+    #
+    # @param trace_id [String] Lowercase hexadecimal trace ID
+    # @param context [OpenTelemetry::Context] Context to extend
+    # @return [OpenTelemetry::Context] Context with the trace claim
+    # @api private
+    def self._set_langfuse_trace_id_in_baggage(trace_id, context:)
+      return context unless baggage_available?
+
+      normalized_trace_id = trace_id.downcase
+      return context if _get_langfuse_trace_id_from_baggage(context) == normalized_trace_id
+
+      OpenTelemetry::Baggage.set_value(
+        LANGFUSE_TRACE_ID_BAGGAGE_KEY, normalized_trace_id, context: context
+      )
+    rescue StandardError => e
+      Langfuse.configuration.logger.debug("Langfuse: Trace baggage write failed: #{e.message}")
+      context
+    end
+
     # Extract propagated attributes from baggage
     #
     # @param context [OpenTelemetry::Context] The context to read baggage from
@@ -399,7 +464,7 @@ module Langfuse
 
         attributes[span_key] = _parse_baggage_value(span_key, baggage_value)
       end
-      attributes
+      attributes.compact
     rescue StandardError => e
       Langfuse.configuration.logger.debug("Langfuse: Baggage extraction failed: #{e.message}")
       {}
@@ -413,7 +478,9 @@ module Langfuse
     #
     # @api private
     def self._parse_baggage_value(span_key, baggage_value)
-      if span_key == OtelAttributes::TRACE_TAGS && baggage_value.is_a?(String)
+      if span_key == OtelAttributes::ENVIRONMENT
+        _validate_environment_value(baggage_value)
+      elsif span_key == OtelAttributes::TRACE_TAGS && baggage_value.is_a?(String)
         baggage_value.split(",")
       else
         baggage_value.to_s