bitfab 0.51.9 → 0.51.10

data/lib/bitfab/http_client.rb

@@ -12,6 +12,9 @@ require_relative "transport"
 require_relative "trace_completion"
 require_relative "version"
 require_relative "warn_once"
+require_relative "simulation_plan"
+require_relative "span_origin"
+require_relative "thread_propagation"
 
 module Bitfab
   # What a caller learns about one tracked trace once it takes it back.
@@ -45,49 +48,91 @@ module Bitfab
       @carrier_seq = 0
       @transport = nil
       @closed = false
+      @close_started = false
+      @close_result = nil
+      @close_condition = ConditionVariable.new
+    end
+
+    def configure_simulation_plan(enabled: true)
+      @simulation_plan = SimulationPlan.new(self, enabled:)
+    end
+
+    def refresh_simulation_plan
+      @simulation_plan&.refresh
+    end
+
+    def get_simulation_plan
+      get("/api/sdk/sim-plan", timeout: SimulationPlan::READ_TIMEOUT, connection_close: true)
     end
 
     # Flush and permanently close this client's tracing transport.
     # Returns true when everything it queued was delivered within the deadline.
     def close(timeout: 30)
-      transport = @transport_mutex.synchronize do
-        current = @transport
-        @transport = nil
-        @closed = true
-        current
+      deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + [timeout, 0].max
+      @transport_mutex.synchronize do
+        while @close_started && @close_result.nil?
+          remaining = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          return false if remaining <= 0
+          @close_condition.wait(@transport_mutex, remaining)
+        end
+        return @close_result unless @close_result.nil?
+        @close_started = true
       end
-      return true if transport.nil?
 
-      transport.shutdown(timeout)
+      result = false
+      begin
+        remaining = [deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC), 0].max
+        released = @simulation_plan.nil? || @simulation_plan.release([remaining / 2.0, SimulationPlan::READ_TIMEOUT].min)
+        @simulation_plan&.stop
+        remaining = [deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC), 0].max
+        transport = @transport_mutex.synchronize do
+          current = @transport
+          @transport = nil
+          @closed = true
+          current
+        end
+        shutdown = transport.nil? || transport.shutdown(remaining)
+        result = released && shutdown
+      ensure
+        @transport_mutex.synchronize do
+          @close_result = result
+          @close_condition.broadcast
+        end
+      end
+      result
     end
 
     # Wait for the spans and traces this client queued to be delivered.
     def flush(timeout: 30)
+      deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + [timeout, 0].max
+      released = @simulation_plan.nil? || @simulation_plan.release([timeout / 2.0, SimulationPlan::READ_TIMEOUT].min)
+      timeout = [deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC), 0].max
       transport = @transport_mutex.synchronize { @transport }
-      return true if transport.nil?
+      return released if transport.nil?
 
-      transport.flush(timeout)
+      transport.flush(timeout) && released
     end
 
     # Make a POST request to the Bitfab API.
     # Returns parsed JSON response hash.
-    def request(endpoint, payload, timeout: nil, max_retries: 1, retry_delay: 0.1)
-      send_encoded(endpoint, Serialize.safe_generate(payload), timeout:, max_retries:, retry_delay:)
+    def request(endpoint, payload, timeout: nil, max_retries: 1, retry_delay: 0.1, method: "POST")
+      send_encoded(endpoint, Serialize.safe_generate(payload), timeout:, max_retries:, retry_delay:, method:)
     end
 
     # POST an already-encoded body. The span transport encodes its own batches,
     # so routing them back through #request would encode the same data twice.
-    def send_encoded(endpoint, body, timeout: nil, max_retries: 1, retry_delay: 0.1)
+    def send_encoded(endpoint, body, timeout: nil, max_retries: 1, retry_delay: 0.1, method: "POST")
       send_prepared(
         endpoint,
         Compress.prepare_request_body(body),
         timeout:,
         max_retries:,
-        retry_delay:
+        retry_delay:,
+        method:
       )
     end
 
-    def send_prepared(endpoint, prepared, timeout: nil, max_retries: 1, retry_delay: 0.1)
+    def send_prepared(endpoint, prepared, timeout: nil, max_retries: 1, retry_delay: 0.1, method: "POST")
       uri = URI("#{@service_url}#{endpoint}")
       request_timeout = timeout || @timeout
 
@@ -99,7 +144,7 @@ module Bitfab
         http.open_timeout = request_timeout
         http.read_timeout = request_timeout
 
-        req = Net::HTTP::Post.new(uri.path, headers)
+        req = Net::HTTPGenericRequest.new(method, true, true, uri.request_uri, headers)
         req["Content-Encoding"] = prepared.content_encoding if prepared.content_encoding
         req.body = prepared.body
 
@@ -128,18 +173,32 @@ module Bitfab
 
     # Queue an external span on this client's trace transport (fire-and-forget).
     def send_external_span(payload)
-      ref = carrier_ref(payload)
-      @trace_completion.record(ref.trace_id, ref.span_id) if ref&.span_id
-      trace_transport&.submit(
-        "external_span",
-        payload.merge("sdkVersion" => VERSION),
-        recorded_meta("external_span", payload, carrier_ref(payload))
-      )
+      if (data = payload.dig("rawSpan", "span_data")).is_a?(Hash) && !data.key?("runtime")
+        payload = payload.merge("rawSpan" => payload["rawSpan"].merge("span_data" => data.merge("runtime" => ThreadPropagation.runtime)))
+      end
+      if payload["rawSpan"].is_a?(Hash) && !payload["rawSpan"].key?("span_origin")
+        payload = payload.merge("rawSpan" => payload["rawSpan"].merge("span_origin" => Bitfab.make_span_origin("span")))
+      end
+      submit = ->(record) {
+        ref = carrier_ref(record)
+        @trace_completion.record(ref.trace_id, ref.span_id) if ref&.span_id
+        trace_transport&.submit(
+          "external_span",
+          record.merge("sdkVersion" => VERSION),
+          recorded_meta("external_span", record, ref)
+        )
+      }
+      @simulation_plan ? @simulation_plan.send_span(payload, &submit) : submit.call(payload)
+    end
+
+    def send_internal_trace(function_id, payload)
+      record = payload.merge("functionId" => function_id, "sdkVersion" => VERSION)
+      trace_transport&.submit("internal_trace", record, recorded_meta("internal_trace", record, nil))
     end
 
     # Make a GET request to the Bitfab API.
     # Returns parsed JSON response hash.
-    def get(endpoint, timeout: nil)
+    def get(endpoint, timeout: nil, connection_close: false)
       uri = URI("#{@service_url}#{endpoint}")
       request_timeout = timeout || @timeout
 
@@ -150,6 +209,7 @@ module Bitfab
 
       # request_uri (not path) so any query string on the endpoint survives.
       req = Net::HTTP::Get.new(uri.request_uri, headers)
+      req["Connection"] = "close" if connection_close
       response = http.request(req)
 
       unless response.is_a?(Net::HTTPSuccess)
@@ -189,7 +249,7 @@ module Bitfab
     #   same org and trace function or the server rejects the replay
     def start_replay(trace_function_key, limit, trace_ids: nil, code_change_description: nil,
       code_change_files: nil, experiment_group_id: nil, name: nil, include_db_branch_lease: false, dataset_ids: nil,
-      grader_ids: nil, db_branch_settings: nil)
+      grader_ids: nil, db_branch_settings: nil, attempts: 1, only_with_assertions: false, include_original_metadata: false)
       payload = {
         "traceFunctionKey" => trace_function_key
       }
@@ -211,6 +271,9 @@ module Bitfab
         end
       end
       payload["graderIds"] = grader_ids unless grader_ids.nil?
+      payload["attempts"] = attempts if attempts > 1
+      payload["onlyWithAssertions"] = true if only_with_assertions
+      payload["includeOriginalMetadata"] = true if include_original_metadata
       payload["dbBranchSettings"] = db_branch_settings unless db_branch_settings.nil?
       git_state = GitState.resolved
       payload["git"] = git_state if git_state
@@ -279,8 +342,9 @@ module Bitfab
       request("/api/sdk/replay/releaseDbBranchLease", {"neonBranchId" => neon_branch_id}, timeout: 30)
     end
 
-    def resolve_db_branch_lease(test_run_id, trace_id, db_branch_settings = nil)
+    def resolve_db_branch_lease(test_run_id, trace_id, db_branch_settings = nil, attempt = 0)
       payload = {"testRunId" => test_run_id, "traceId" => trace_id}
+      payload["attempt"] = attempt if attempt.positive?
       payload["dbBranchSettings"] = db_branch_settings unless db_branch_settings.nil?
       request(
         "/api/sdk/replay/resolveDbBranchLease",
@@ -291,15 +355,18 @@ module Bitfab
 
     # Queue an external trace on this client's trace transport (fire-and-forget).
     def send_external_trace(payload)
-      ref = carrier_ref(payload)
-      if payload["completed"] == true && ref
-        @trace_completion.close(ref.trace_id, dropped: payload["dropped"] == true) do |count|
-          submit_external_trace(payload.merge("expectedSpanCount" => count))
+      submit = ->(record) {
+        ref = carrier_ref(record)
+        if record["completed"] == true && ref
+          @trace_completion.close(ref.trace_id, dropped: record["dropped"] == true) do |count|
+            submit_external_trace(record.merge("expectedSpanCount" => count))
+          end
+        else
+          @trace_completion.open(ref.trace_id) if ref
+          submit_external_trace(record)
         end
-      else
-        @trace_completion.open(ref.trace_id) if ref
-        submit_external_trace(payload)
-      end
+      }
+      @simulation_plan ? @simulation_plan.send_trace(payload, &submit) : submit.call(payload)
     end
 
     def submit_external_trace(payload)
@@ -647,11 +714,15 @@ module Bitfab
     # Wait for queued spans and trace completions to reach the server, within
     # one total deadline. Returns true when everything landed in time.
     def flush_traces(timeout: 30)
-      Transport.flush_trace_transports([timeout, 0].max)
+      deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + [timeout, 0].max
+      released = SimulationPlan.release_all([timeout / 2.0, SimulationPlan::READ_TIMEOUT].min)
+      flushed = Transport.flush_trace_transports([deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC), 0].max)
+      released && flushed
     end
   end
 
   at_exit do
-    Transport.shutdown_trace_transports(4)
+    SimulationPlan.release_all(2, stop: true)
+    Transport.shutdown_trace_transports(2)
   end
 end

data/lib/bitfab/labels.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require "uri"
+
+module Bitfab
+  class Labels
+    LABELS_PATH = "/api/sdk/traces/labels"
+    HUMAN_LABELS_PATH = "#{LABELS_PATH}/human"
+
+    def initialize(http_client)
+      @http_client = http_client
+    end
+
+    def save(label:, annotation:, trace_id: nil, original_trace_id: nil, attempt: nil,
+      confidence: nil, test_run_id: nil, assertion_id: nil)
+      update = label_target(trace_id:, original_trace_id:, attempt:, assertion_id:)
+      update["label"] = label
+      update["annotation"] = annotation
+      update["confidence"] = confidence unless confidence.nil?
+      save_all([update], test_run_id:).first
+    end
+
+    def save_all(labels, test_run_id: nil)
+      payload = {"labels" => labels}
+      payload["testRunId"] = test_run_id unless test_run_id.nil?
+      @http_client.request(LABELS_PATH, payload)["labels"]
+    end
+
+    def skip(trace_id: nil, original_trace_id: nil, attempt: nil, test_run_id: nil, assertion_id: nil)
+      update = label_target(trace_id:, original_trace_id:, attempt:, assertion_id:)
+      save_all([update.merge("skip" => true)], test_run_id:).first
+    end
+
+    def archive(trace_id: nil, original_trace_id: nil, attempt: nil, test_run_id: nil, assertion_id: nil)
+      update = label_target(trace_id:, original_trace_id:, attempt:, assertion_id:)
+      save_all([update.merge("archive" => true)], test_run_id:).first
+    end
+
+    def save_human(label:, annotation:, trace_id:, confidence: nil, assertion_id: nil)
+      update = {"traceId" => trace_id, "label" => label, "annotation" => annotation}
+      update["assertionId"] = assertion_id unless assertion_id.nil?
+      update["confidence"] = confidence unless confidence.nil?
+      save_human_all([update]).first
+    end
+
+    def save_human_all(labels)
+      @http_client.request(HUMAN_LABELS_PATH, {"labels" => labels})["labels"]
+    end
+
+    def get(trace_id)
+      get_all([trace_id]).first
+    end
+
+    def get_all(trace_ids)
+      return [] if trace_ids.empty?
+
+      query = URI.encode_www_form("traceIds" => trace_ids.join(","))
+      @http_client.get("#{LABELS_PATH}?#{query}")["labels"]
+    end
+
+    private
+
+    def label_target(trace_id:, original_trace_id:, attempt:, assertion_id:)
+      if trace_id.nil? == original_trace_id.nil?
+        raise ArgumentError, "Pass exactly one of trace_id or original_trace_id. " \
+          "Use original_trace_id for a replay verdict, with the test_run_id it ran under."
+      end
+
+      update = {}
+      update["assertionId"] = assertion_id unless assertion_id.nil?
+      if trace_id.nil?
+        update["originalTraceId"] = original_trace_id
+        update["attempt"] = attempt unless attempt.nil?
+      else
+        update["traceId"] = trace_id
+      end
+      update
+    end
+  end
+end

data/lib/bitfab/replay.rb

@@ -75,9 +75,10 @@ module Bitfab
     # thread. Child threads and processes intentionally do not inherit it.
     def with_context(test_run_id:, input_source_span_id: nil, input_source_trace_id: nil, trace_id: nil,
       mock_tree: nil, mock_strategy: nil, mock_overrides: nil, fetch_span_output: nil,
-      db_branch_lease: nil, db_branch_timings: nil, source_bitfab_trace_id: nil)
+      db_branch_lease: nil, db_branch_timings: nil, source_bitfab_trace_id: nil, replay_attempt: nil)
       previous = Thread.current.thread_variable_get(REPLAY_CONTEXT_KEY)
       ctx = {
+        replay_attempt:,
         test_run_id:,
         input_source_span_id:,
         input_source_trace_id:,
@@ -256,12 +257,27 @@ module Bitfab
     #   A raising callback never crashes the run.
     # @return [Hash] with :items, :test_run_id, :test_run_url
     def run(client, receiver, method_name, trace_function_key:, limit: nil, trace_ids: nil, name: nil,
-      max_concurrency: 10, code_change_description: CODE_CHANGE_UNSET, code_change_files: CODE_CHANGE_UNSET, experiment_group_id: nil,
-      dataset_id: nil, dataset_ids: nil, grader_ids: nil, mock: "marked",
+      concurrency: nil, max_concurrency: ReplayConcurrency::UNSET, code_change_description: CODE_CHANGE_UNSET, code_change_files: CODE_CHANGE_UNSET, experiment_group_id: nil,
+      dataset_id: nil, dataset_ids: nil, grader_ids: nil, mock: "marked", attempts: ReplayConcurrency::UNSET, only_with_assertions: false, dry_run: false,
       adapt_inputs: nil, mock_override: nil, db_branch: nil, on_item_start: nil, on_item_finish: nil, on_progress: nil)
+      concurrency_options = ReplayConcurrency.resolve(concurrency, attempts, max_concurrency)
+      attempts = concurrency_options.attempts
+      max_concurrency = concurrency_options.max_concurrency
+      process_executor = (concurrency_options.primitive == "process") ? ReplayProcesses.current : nil
+      if concurrency_options.primitive == "process" && process_executor.nil?
+        raise ArgumentError, "Process replay requires bitfab-replay --registry so each child can load the application"
+      end
       unless MOCK_STRATEGIES.include?(mock.to_s)
         raise ArgumentError, "Invalid mock strategy '#{mock}'. Must be one of: #{MOCK_STRATEGIES.join(", ")}"
       end
+      unless max_concurrency.nil?
+        unless max_concurrency.is_a?(Integer) && max_concurrency.positive?
+          raise ArgumentError, "max_concurrency must be a positive integer or nil for unlimited concurrency."
+        end
+      end
+      unless attempts.is_a?(Integer) && attempts.between?(1, 100)
+        raise ArgumentError, "attempts must be an integer from 1 to 100."
+      end
       item_finish_callback = on_item_finish || on_progress
       if trace_ids
         raise ArgumentError, "trace_ids must contain at least one trace ID." if trace_ids.empty?
@@ -281,7 +297,7 @@ module Bitfab
       # function, recorded under another). Only fires when the method's key is
       # introspectable; an untraced method falls through to the persistence
       # check in complete_replay. Mirrors the TypeScript/Python SDKs.
-      declared_key = Traceable.trace_function_key_for(receiver, method_name)
+      declared_key = method_name && Traceable.trace_function_key_for(receiver, method_name)
       if declared_key && declared_key != trace_function_key
         raise ArgumentError,
           "Method #{method_name} is traced under trace function key '#{declared_key}' but replay was " \
@@ -289,6 +305,13 @@ module Bitfab
           "or point at the method traced under '#{trace_function_key}'."
       end
 
+      raise ArgumentError, "Replay requires an API key." unless client.send(:resolve_api_key)
+      unless declared_key
+        callable = method_name ? receiver.method(method_name) : receiver
+        receiver = ManagedCallable.new(client, trace_function_key, callable)
+        method_name = :call
+      end
+
       http_client = client.instance_variable_get(:@http_client)
 
       # Resolved override list: per-call overrides FIRST (they win), then the
@@ -302,7 +325,7 @@ module Bitfab
       # the count), so it's omitted from the request entirely.
       effective_limit = trace_ids ? nil : (limit || 5)
 
-      include_db_branch_lease = db_branch_enabled?(db_branch)
+      include_db_branch_lease = db_branch_enabled?(db_branch) && !dry_run
 
       # code_change_files controls capture: omitted auto-captures, an array
       # wins, and explicit nil opts out. Preserve a caller-supplied description
@@ -329,22 +352,45 @@ module Bitfab
         include_db_branch_lease:,
         dataset_ids: resolved_dataset_ids,
         grader_ids:,
+        attempts:,
+        only_with_assertions:,
+        include_original_metadata: !adapt_inputs.nil?,
         db_branch_settings: resolved_db_branch_settings
       )
       test_run_id = replay_data["testRunId"]
       test_run_url = replay_data["testRunUrl"]
       full_test_run_url = "#{client.service_url}#{test_run_url}"
-      server_items = replay_data["items"] || []
+      server_items = attempts.times.flat_map do |attempt|
+        (replay_data["items"] || []).map do |source|
+          item = source.merge("attempt" => attempt)
+          if attempt.positive?
+            item.delete("dbBranchLease")
+            item.delete("dbBranchLeaseError")
+          end
+          item
+        end
+      end
 
       result_items = if server_items.any?
         process_items(http_client, server_items, receiver, method_name, test_run_id, max_concurrency, mock.to_s,
           adapt_inputs, include_db_branch_lease, resolved_db_branch_settings, on_item_start:,
           on_item_finish: item_finish_callback,
-          mock_overrides: resolved_overrides)
+          mock_overrides: resolved_overrides, dry_run:, process_executor:)
       else
         []
       end
 
+      if dry_run
+        begin
+          http_client.complete_replay(test_run_id)
+        rescue => e
+          warn "Bitfab: could not finalize dry run: #{e.message}"
+        end
+        result = {items: public_replay_items(result_items), test_run_id:, test_run_url: full_test_run_url, attempts:}
+        write_replay_result_file(result)
+        return result
+      end
+
       # Spans and completions ride a batched transport, so the run is only safe
       # to finalize once the server confirms every replay trace it queued: the
       # trace-ID mapping complete_replay builds would otherwise race the
@@ -440,6 +486,7 @@ module Bitfab
       result_items.each { |item| item.delete(:_sdk_trace_id) }
 
       result = {
+        attempts:,
         items: result_items,
         test_run_id:,
         test_run_url: full_test_run_url
@@ -698,7 +745,7 @@ module Bitfab
     # Process all replay items, optionally in parallel using threads.
     def process_items(http_client, server_items, receiver, method_name, test_run_id, max_concurrency, mock_strategy,
       adapt_inputs = nil, include_db_branch_lease = false, db_branch_settings = nil, on_item_start: nil, on_item_finish: nil,
-      mock_overrides: [])
+      mock_overrides: [], dry_run: false, process_executor: nil)
       concurrency = max_concurrency || server_items.length
 
       # Lifecycle callbacks run from worker threads in the parallel path, so the
@@ -719,7 +766,7 @@ module Bitfab
       # surfaced. original_trace_id (the historical trace being replayed, taken
       # from the server item) is what a UI keys on to identify what just finished.
       # source_trace_id/source_span_id are kept as deprecated aliases.
-      report_start = lambda do |original_trace_id, original_span_id|
+      report_start = lambda do |original_trace_id, original_span_id, attempt|
         progress_mutex.synchronize do
           started += 1
           next unless on_item_start
@@ -728,6 +775,7 @@ module Bitfab
             on_item_start.call({
               type: "started", test_run_id:, started:, completed:, total:, succeeded:, errored:,
               item: {
+                attempt:,
                 original_trace_id:,
                 original_span_id:,
                 source_trace_id: original_trace_id,
@@ -759,7 +807,7 @@ module Bitfab
         # Deliver this item's trace now, then read its server id back off the
         # ingest response. A flush failure never crashes the run: the id stays
         # nil and the end-of-run barrier remains the authority on persistence.
-        flush_finished_item_trace.call
+        flush_finished_item_trace.call unless dry_run
         result[:trace_id] = http_client.peek_server_trace_id(result[:_sdk_trace_id]) || result[:trace_id]
         progress_mutex.synchronize do
           completed += 1
@@ -771,6 +819,8 @@ module Bitfab
             on_item_finish.call({
               test_run_id:, completed:, total:, succeeded:, errored:,
               item: {
+                attempt: result[:attempt],
+                ingestion_type: result[:ingestion_type],
                 trace_id: result[:trace_id],
                 original_trace_id:,
                 original_span_id:,
@@ -801,11 +851,19 @@ module Bitfab
         end
       end
 
+      execute = lambda do |item|
+        if process_executor
+          process_executor.call(item, test_run_id, mock_strategy, include_db_branch_lease, db_branch_settings, dry_run:)
+        else
+          process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy,
+            adapt_inputs, include_db_branch_lease, db_branch_settings, mock_overrides:, dry_run:)
+        end
+      end
+
       if concurrency <= 1
         server_items.map do |item|
-          report_start.call(original_trace_id_of(item), original_span_id_of(item))
-          result = process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy,
-            adapt_inputs, include_db_branch_lease, db_branch_settings, mock_overrides:)
+          report_start.call(original_trace_id_of(item), original_span_id_of(item), item["attempt"] || 0)
+          result = execute.call(item)
           report.call(result, original_trace_id_of(item), original_span_id_of(item), test_run_id)
           result
         end
@@ -821,9 +879,8 @@ module Bitfab
               item, idx = work_mutex.synchronize { work_queue.shift }
               break unless item
 
-              report_start.call(original_trace_id_of(item), original_span_id_of(item))
-              result = process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy,
-                adapt_inputs, include_db_branch_lease, db_branch_settings, mock_overrides:)
+              report_start.call(original_trace_id_of(item), original_span_id_of(item), item["attempt"] || 0)
+              result = execute.call(item)
               results_mutex.synchronize { results[idx] = result }
               report.call(result, original_trace_id_of(item), original_span_id_of(item), test_run_id)
             end
@@ -842,8 +899,9 @@ module Bitfab
     # than propagated, so one bad trace never aborts the whole replay run
     # (mirrors the TypeScript and Python SDKs' per-item rescue).
     def process_single_item(http_client, server_item, receiver, method_name, test_run_id, mock_strategy,
-      adapt_inputs = nil, include_db_branch_lease = false, db_branch_settings = nil, mock_overrides: [])
+      adapt_inputs = nil, include_db_branch_lease = false, db_branch_settings = nil, mock_overrides: [], dry_run: false)
       metrics = extract_server_item_metrics(server_item)
+      attempt = server_item["attempt"] || 0
       # The ORIGINAL (historical) trace/span this item replays. Canonical
       # keys are originalTraceId/originalSpanId; older servers send them under
       # the deprecated sourceTraceId/sourceSpanId aliases (see *_of helpers).
@@ -866,7 +924,7 @@ module Bitfab
       db_branch_timings = include_db_branch_lease ? server_item["dbBranchTimings"] : nil
       if include_db_branch_lease && lease.nil? && lease_error.nil?
         begin
-          resolved = http_client.resolve_db_branch_lease(test_run_id, original_trace_id, db_branch_settings)
+          resolved = http_client.resolve_db_branch_lease(test_run_id, original_trace_id, db_branch_settings, attempt)
         rescue => cause
           error = DbBranchReplayError.new(
             "lease_request_failed",
@@ -905,7 +963,7 @@ module Bitfab
       overrides_present = !mock_overrides.nil? && !mock_overrides.empty?
       include_outputs = mock_strategy == "all"
       mock_tree = nil
-      if mock_strategy == "all" || mock_strategy == "marked" || overrides_present
+      if !dry_run && (mock_strategy == "all" || mock_strategy == "marked" || overrides_present)
         begin
           tree = http_client.get_span_tree(
             original_span_id,
@@ -930,6 +988,7 @@ module Bitfab
       fetch_span_output = (mock_tree && !include_outputs) ? build_span_output_fetcher(http_client) : nil
 
       adapt_ctx = {
+        metadata: server_item["originalMetadata"].is_a?(Hash) ? server_item["originalMetadata"].dup : {},
         original_trace_id:,
         original_span_id:,
         # Deprecated aliases for original_trace_id/original_span_id.
@@ -951,6 +1010,8 @@ module Bitfab
         fetch_span_output:,
         adapt_inputs:,
         adapt_ctx:,
+        attempt:,
+        dry_run:,
         db_branch_lease: lease,
         db_branch_timings:,
         source_bitfab_trace_id: original_trace_id,
@@ -960,6 +1021,8 @@ module Bitfab
     rescue => e
       warn "Bitfab: replay item for span #{original_span_id} failed before execution: #{e.message}"
       {
+        attempt:,
+        ingestion_type: server_item["ingestionType"],
         input: [],
         result: nil,
         original_output: nil,
@@ -1103,6 +1166,7 @@ module Bitfab
         original_duration_ms: server_item["originalDurationMs"] || server_item["durationMs"],
         original_tokens: server_item["originalTokens"] || server_item["tokens"],
         original_model: server_item["originalModel"] || server_item["model"],
+        ingestion_type: server_item["ingestionType"],
         tokens: nil
       }
     end
@@ -1149,6 +1213,7 @@ module Bitfab
         status = http_client.get_replay_status(test_run_id, expected_span_counts)
         ready = status["traceIds"]
         ready = {} unless ready.is_a?(Hash)
+        read_back_trace_ids.merge!(ready)
         missing = expected_span_counts.keys - ready.keys
         return read_back_trace_ids if missing.empty?
         break if Otel.monotonic_now >= deadline
@@ -1185,7 +1250,7 @@ module Bitfab
     def execute_item(item, receiver, method_name, test_run_id, input_source_span_id = nil, metrics = {},
       input_source_trace_id: nil, mock_strategy: "marked", mock_tree: nil, mock_overrides: nil,
       fetch_span_output: nil, adapt_inputs: nil, adapt_ctx: nil, db_branch_lease: nil, db_branch_timings: nil,
-      source_bitfab_trace_id: nil, db_snapshot_ref: nil, http_client: nil)
+      source_bitfab_trace_id: nil, db_snapshot_ref: nil, http_client: nil, attempt: 0, dry_run: false)
       args, kwargs = Serialize.deserialize_inputs(item)
 
       fn_result = nil
@@ -1201,10 +1266,11 @@ module Bitfab
       # Declared before this item's first span is submitted: the transport
       # records delivery only for traces someone asked about, so anything
       # submitted before this would go untracked.
-      http_client&.track_trace_deliveries([sdk_trace_id])
+      http_client&.track_trace_deliveries([sdk_trace_id]) unless dry_run
 
       begin
         ReplayContext.with_context(
+          replay_attempt: attempt,
           test_run_id:,
           input_source_span_id:,
           input_source_trace_id:,
@@ -1231,7 +1297,7 @@ module Bitfab
           rescue => e
             replay_error = e
           end
-          if replay_error.nil?
+          if replay_error.nil? && !dry_run
             begin
               replay_started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
               fn_result = if kwargs.empty?
@@ -1239,6 +1305,7 @@ module Bitfab
               else
                 receiver.send(method_name, *args, **kwargs)
               end
+              fn_result = fn_result.to_a if fn_result.is_a?(Enumerator)
               replay_duration_ms = ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - replay_started) * 1000).round
             rescue => e
               # The method ran and raised, so it still has a duration.
@@ -1258,7 +1325,9 @@ module Bitfab
 
       item_error = fn_error || replay_error
       {
-        input: args,
+        attempt:,
+        ingestion_type: metrics[:ingestion_type],
+        input: kwargs.empty? ? args : [*args, kwargs],
         result: fn_result,
         original_output: item["output"],
         error: item_error&.message,

data/lib/bitfab/replay_cli.rb

@@ -14,7 +14,13 @@ module Bitfab
       registry_path, replay_argv = parse(argv)
       loader = registry_loader || method(:load_registry)
       registry = loader.call(registry_path)
-      ReplayCli.run(registry, argv: replay_argv, stdout:, stderr:)
+      if replay_argv.first == "--execute-item"
+        raise ArgumentError, "Expected one process assignment path" unless replay_argv.length == 2
+        return ReplayProcesses.execute_assignment(registry, replay_argv[1], stderr:)
+      end
+      ReplayProcesses.with_registry(registry_path, replay_argv, stderr:) do
+        ReplayCli.run(registry, argv: replay_argv, stdout:, stderr:)
+      end
     end
 
     def main(argv: ARGV, stdout: $stdout, stderr: $stderr)