feat!: Add per-execution runId, at-most-once tracking, and cross-process tracker resumption (#29)

Author: jsonbailey

examples/chatbot/aws-bedrock/hello_bedrock.rb

@@ -29,10 +29,13 @@ def initialize(ai_config, bedrock_client)
     @bedrock_client = bedrock_client
   end
 
+  # Returns [response_content, resumption_token]. The resumption token can be
+  # used to reconstruct a tracker for deferred operations like feedback.
   def ask_agent(question)
     @messages << LaunchDarkly::Server::AI::Message.new('user', question)
+    tracker = @ai_config.create_tracker
     begin
-      response = ai_config.tracker.track_bedrock_converse_metrics do
+      response = tracker.track_bedrock_converse_metrics do
         @bedrock_client.converse(
           map_converse_arguments(
             ai_config.model.name,
@@ -41,15 +44,17 @@ def ask_agent(question)
         )
       end
       @messages << LaunchDarkly::Server::AI::Message.new('assistant', response.output.message.content[0].text)
-      response.output.message.content[0].text
+      [response.output.message.content[0].text, tracker.resumption_token]
     rescue StandardError => e
-      "An error occured: #{e.message}"
+      ["An error occured: #{e.message}", nil]
     end
   end
 
-  def agent_was_helpful(helpful)
+  def agent_was_helpful(helpful, tracker)
+    return if tracker.nil?
+
     kind = helpful ? :positive : :negative
-    ai_config.tracker.track_feedback(kind: kind)
+    tracker.track_feedback(kind: kind)
   end
 
   def map_converse_arguments(model_id, messages)
@@ -89,7 +94,7 @@ def map_converse_arguments(model_id, messages)
 
 # Pass a default for improved resiliency when the flag is unavailable or LaunchDarkly is unreachable; omit for a disabled default.
 # Example:
-#   default = LaunchDarkly::Server::AI::AIConfig.new(
+#   default = LaunchDarkly::Server::AI::AIConfigDefault.new(
 #     enabled: true,
 #     model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'my-model'),
 #     provider: LaunchDarkly::Server::AI::ProviderConfig.new(name: 'my-provider'),
@@ -105,18 +110,27 @@ def map_converse_arguments(model_id, messages)
 
 chatbot = BedrockChatbot.new(ai_config, bedrock_client)
 
+last_resumption_token = nil
+
 loop do
   print "Ask a question (or type 'exit'): "
   question = gets&.chomp
   break if question.nil? || question.strip.downcase == 'exit'
 
-  response = chatbot.ask_agent(question)
+  response, last_resumption_token = chatbot.ask_agent(question)
   puts "AI Response: #{response}"
 end
 
 print "Was the chat helpful? [yes/no]: "
 feedback = gets&.chomp
 
-chatbot.agent_was_helpful(feedback == 'yes') unless feedback.nil?
+unless feedback.nil? || last_resumption_token.nil?
+  tracker = LaunchDarkly::Server::AI::AIConfigTracker.from_resumption_token(
+    token: last_resumption_token,
+    ld_client:,
+    context:
+  )
+  chatbot.agent_was_helpful(feedback == 'yes', tracker)
+end
 
 ld_client.close

examples/chatbot/openai/hello_openai.rb

@@ -35,26 +35,31 @@ def initialize(ai_config, openai_client)
     @openai_client = openai_client
   end
 
+  # Returns [response_content, resumption_token]. The resumption token can be
+  # used to reconstruct a tracker for deferred operations like feedback.
   def ask_agent(question)
     @messages << LaunchDarkly::Server::AI::Message.new('user', question)
+    tracker = @ai_config.create_tracker
     begin
-      completion = ai_config.tracker.track_openai_metrics do
+      completion = tracker.track_openai_metrics do
         @openai_client.chat.completions.create(
           model: ai_config.model.name,
           messages: @messages.map(&:to_h)
         )
       end
       response_content = completion[:choices][0][:message][:content]
       @messages << LaunchDarkly::Server::AI::Message.new('assistant', response_content)
-      response_content
+      [response_content, tracker.resumption_token]
     rescue StandardError => e
-      "An error occurred: #{e.message}"
+      ["An error occurred: #{e.message}", nil]
     end
   end
 
-  def agent_was_helpful(helpful)
+  def agent_was_helpful(helpful, tracker)
+    return if tracker.nil?
+
     kind = helpful ? :positive : :negative
-    ai_config.tracker.track_feedback(kind: kind)
+    tracker.track_feedback(kind: kind)
   end
 end
 
@@ -78,7 +83,7 @@ def agent_was_helpful(helpful)
 
 # Pass a default for improved resiliency when the flag is unavailable or LaunchDarkly is unreachable; omit for a disabled default.
 # Example:
-#   default = LaunchDarkly::Server::AI::AIConfig.new(
+#   default = LaunchDarkly::Server::AI::AIConfigDefault.new(
 #     enabled: true,
 #     model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'my-model'),
 #     provider: LaunchDarkly::Server::AI::ProviderConfig.new(name: 'my-provider'),
@@ -94,18 +99,27 @@ def agent_was_helpful(helpful)
 
 chatbot = Chatbot.new(ai_config, OpenAI::Client.new(api_key: openai_api_key))
 
+last_resumption_token = nil
+
 loop do
   print "Ask a question (or type 'exit'): "
   question = gets&.chomp
   break if question.nil? || question.strip.downcase == 'exit'
 
-  response = chatbot.ask_agent(question)
+  response, last_resumption_token = chatbot.ask_agent(question)
   puts "AI Response: #{response}"
 end
 
 print "Was the chat helpful? [yes/no]: "
 feedback = gets&.chomp
 
-chatbot.agent_was_helpful(feedback == 'yes') unless feedback.nil?
+unless feedback.nil? || last_resumption_token.nil?
+  tracker = LaunchDarkly::Server::AI::AIConfigTracker.from_resumption_token(
+    token: last_resumption_token,
+    ld_client:,
+    context:
+  )
+  chatbot.agent_was_helpful(feedback == 'yes', tracker)
+end
 
 ld_client.close

launchdarkly-server-sdk-ai.gemspec

@@ -19,6 +19,7 @@ Gem::Specification.new do |spec|
   spec.require_paths = ['lib']
   spec.required_ruby_version = '>= 3.1.0'
 
+  spec.add_dependency 'base64'
   spec.add_dependency 'launchdarkly-server-sdk', '~> 8.5'
   spec.add_dependency 'logger'
   spec.add_dependency 'mustache', '~> 1.1'

lib/server/ai/ai_config_tracker.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'base64'
+require 'json'
 require 'ldclient-rb'
 
 module LaunchDarkly
@@ -39,7 +41,14 @@ def initialize
       end
 
       #
-      # The AIConfigTracker class is used to track AI configuration usage.
+      # The AIConfigTracker records metrics for a single AI run. Unless
+      # otherwise noted, the tracker's methods are not safe for concurrent use.
+      #
+      # All events a tracker emits share a runId (a UUIDv4) so LaunchDarkly can
+      # correlate them in metrics views. See individual track methods for their
+      # specific semantics. Call create_tracker on the AI Config to start a new
+      # run. A resumption token preserves the runId, so events emitted by a
+      # tracker reconstructed in another process share the original runId.
       #
       class AIConfigTracker
         attr_reader :ld_client, :config_key, :context, :variation_key, :version, :summary, :model_name, :provider_name
@@ -55,7 +64,7 @@ class AIConfigTracker
         # @param provider_name [String] The name of the AI provider
         # @param context [LDContext] The context used for the flag evaluation
         #
-        def initialize(ld_client:, variation_key:, config_key:, version:, model_name:, provider_name:, context:)
+        def initialize(ld_client:, run_id:, config_key:, variation_key:, version:, context:, model_name:, provider_name:)
           @ld_client = ld_client
           @variation_key = variation_key
           @config_key = config_key
@@ -64,14 +73,62 @@ def initialize(ld_client:, variation_key:, config_key:, version:, model_name:, p
           @provider_name = provider_name
           @context = context
           @summary = MetricSummary.new
+          @run_id = run_id
+          @logger = LaunchDarkly::Server::AI.default_logger
+        end
+
+        #
+        # Returns a URL-safe Base64-encoded JSON token that can be used to reconstruct
+        # a tracker in a different process (e.g. for deferred feedback).
+        #
+        # The token contains: runId, configKey, variationKey, version.
+        # modelName and providerName are NOT included.
+        #
+        # @return [String] the resumption token
+        #
+        def resumption_token
+          payload = { runId: @run_id, configKey: @config_key }
+          payload[:variationKey] = @variation_key if @variation_key && !@variation_key.empty?
+          payload[:version] = @version
+          Base64.urlsafe_encode64(JSON.generate(payload), padding: false)
+        end
+
+        #
+        # Reconstructs a tracker from a resumption token.
+        #
+        # @param token [String] A URL-safe Base64-encoded JSON resumption token
+        # @param ld_client [LDClient] The LaunchDarkly client instance
+        # @param context [LDContext] The context for track events
+        # @return [AIConfigTracker] A new tracker instance
+        #
+        def self.from_resumption_token(token:, ld_client:, context:)
+          json = Base64.urlsafe_decode64(token)
+          payload = JSON.parse(json)
+
+          new(
+            ld_client: ld_client,
+            run_id: payload['runId'],
+            config_key: payload['configKey'],
+            variation_key: payload.fetch('variationKey', ''),
+            version: payload['version'],
+            context: context,
+            model_name: '',
+            provider_name: ''
+          )
         end
 
         #
-        # Track the duration of an AI operation
+        # Track the duration of an AI run.
+        #
+        # Records at most once per Tracker; further calls are ignored.
         #
         # @param duration [Integer] The duration in milliseconds
         #
         def track_duration(duration)
+          unless @summary.duration.nil?
+            @logger&.warn("Skipping track_duration: duration already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}")
+            return
+          end
           @summary.duration = duration
           @ld_client.track(
             '$ld:ai:duration:total',
@@ -96,11 +153,18 @@ def track_duration_of(&block)
         end
 
         #
-        # Track time to first token
+        # Track time to first token.
+        #
+        # Records at most once per Tracker; further calls are ignored.
         #
         # @param duration [Integer] The duration in milliseconds
         #
         def track_time_to_first_token(time_to_first_token)
+          unless @summary.time_to_first_token.nil?
+            @logger&.warn("Skipping track_time_to_first_token: time-to-first-token already recorded on this tracker. " \
+                          "Call create_tracker on the AI Config for a new run. #{flag_data}")
+            return
+          end
           @summary.time_to_first_token = time_to_first_token
           @ld_client.track(
             '$ld:ai:tokens:ttf',
@@ -111,11 +175,17 @@ def track_time_to_first_token(time_to_first_token)
         end
 
         #
-        # Track user feedback
+        # Track user feedback.
+        #
+        # Records at most once per Tracker; further calls are ignored.
         #
         # @param kind [Symbol] The kind of feedback (:positive or :negative)
         #
         def track_feedback(kind:)
+          unless @summary.feedback.nil?
+            @logger&.warn("Skipping track_feedback: feedback already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}")
+            return
+          end
           @summary.feedback = kind
           event_name = kind == :positive ? '$ld:ai:feedback:user:positive' : '$ld:ai:feedback:user:negative'
           @ld_client.track(
@@ -127,9 +197,17 @@ def track_feedback(kind:)
         end
 
         #
-        # Track a successful AI generation
+        # Track a successful AI generation.
+        #
+        # Records at most once per Tracker. track_success and track_error share
+        # state; only one of the two can record per Tracker, and subsequent
+        # calls are ignored.
         #
         def track_success
+          unless @summary.success.nil?
+            @logger&.warn("Skipping track_success: success/error already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}")
+            return
+          end
           @summary.success = true
           @ld_client.track(
             '$ld:ai:generation:success',
@@ -140,9 +218,17 @@ def track_success
         end
 
         #
-        # Track an error in AI generation
+        # Track an error in AI generation.
+        #
+        # Records at most once per Tracker. track_success and track_error share
+        # state; only one of the two can record per Tracker, and subsequent
+        # calls are ignored.
         #
         def track_error
+          unless @summary.success.nil?
+            @logger&.warn("Skipping track_error: success/error already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}")
+            return
+          end
           @summary.success = false
           @ld_client.track(
             '$ld:ai:generation:error',
@@ -153,11 +239,17 @@ def track_error
         end
 
         #
-        # Track token usage
+        # Track token usage.
+        #
+        # Records at most once per Tracker; further calls are ignored.
         #
         # @param token_usage [TokenUsage] An object containing token usage details
         #
         def track_tokens(token_usage)
+          unless @summary.usage.nil?
+            @logger&.warn("Skipping track_tokens: token usage already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}")
+            return
+          end
           @summary.usage = token_usage
           if token_usage.total.positive?
             @ld_client.track(
@@ -191,6 +283,10 @@ def track_tokens(token_usage)
         # If the provided block raises, this method will also raise.
         # A failed operation will not have any token usage data.
         #
+        # Subsequent calls re-run the inner block but emit only metrics not
+        # already recorded on this Tracker. Call create_tracker on the AI
+        # Config to start a new run.
+        #
         # @yield The block to track.
         # @return The result of the tracked block.
         #
@@ -208,6 +304,10 @@ def track_openai_metrics(&block)
         # Track AWS Bedrock conversation operations.
         # This method tracks the duration, token usage, and success/error status.
         #
+        # Subsequent calls re-run the inner block but emit only metrics not
+        # already recorded on this Tracker. Call create_tracker on the AI
+        # Config to start a new run.
+        #
         # @yield The block to track.
         # @return [Hash] The original response hash.
         #
@@ -222,13 +322,15 @@ def track_bedrock_converse_metrics(&block)
         end
 
         private def flag_data
-          {
-            variationKey: @variation_key,
+          data = {
+            runId: @run_id,
             configKey: @config_key,
             version: @version,
             modelName: @model_name,
             providerName: @provider_name,
           }
+          data[:variationKey] = @variation_key if @variation_key && !@variation_key.empty?
+          data
         end
 
         private def openai_to_token_usage(usage)