bitfab 0.36.2 → 0.36.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 44af81a548095fe4bf8a347b6e58bb1c3bdcf2726becd78f029ec5c7c9507ca9
-  data.tar.gz: 4a00a95452cf45e5992e7eee4754850923eed8767713368a063fa1fed84d4b7c
+  metadata.gz: 5578e8fdcdd81610b4e65fa94a5e9be88339e816742f6ba7fea81b53867eedf5
+  data.tar.gz: 308d35969ac643d1ddcfb76c169a6108906c958275b392f7d159ea647b0bf63c
 SHA512:
-  metadata.gz: 9043c05ffb79c8830299127c39de0dc77b77dc3fa1fd2a573813696bcd9a7ba3fa95d7a06a7371ff596ed04781f439311f773bdff24afe055b252b36e0d36bb5
-  data.tar.gz: d73327657254114a22c091bb58503c17f1f8d0fcc19f0179c283b452df4c7240d446ff8d36962b16fcc47ea13e0315c80f554280b79f47b684f24ba50b11ce23
+  metadata.gz: 2c0511909223595a7792e780d0c77b8c0f5ef2d239e4b2a5dda7946a63fbe4801893ae02ba3392b74520ad9d82844e519d3319a900f22c30bef4857e616c0420
+  data.tar.gz: 0ae83e2f731bc2f6fa9f564b6fe8754b7fff82c5ae03a914067cca6fb75f4ddb10e935bba8976de6f0dabe7025812e49d9392e2ace61b5202c02d80b21fe00aa

data/lib/bitfab/client.rb

@@ -132,21 +132,27 @@ module Bitfab
     #   +:min_cu+/+:max_cu+ size the branch compute in Neon Compute Units and
     #   +:warmup_sql+ warms its cache. Read the resolved branch inside the
     #   replayed method with +Bitfab.current_replay_branch+.
-    # @param on_progress [#call, nil] optional callback invoked once per item as
-    #   it finishes, with a running-totals hash { completed:, total:, succeeded:,
-    #   errored: }. Use it to render replay progress (e.g. a terminal progress
-    #   bar). Replay does not know pass/fail yet, so the totals only distinguish
-    #   items whose method ran (:succeeded) from items that raised (:errored).
-    #   A raising callback never crashes the run.
+    # @param on_item_finish [#call, nil] optional callback invoked exactly once
+    #   per item as it finishes, always with a running-totals hash containing
+    #   { completed:, total:, succeeded:, errored:, item: }. It never receives a
+    #   whole-run completion event. Replay does not know pass/fail yet, so the
+    #   totals only distinguish items whose method ran (:succeeded) from items
+    #   that raised (:errored). A raising callback never crashes the run.
+    # @param on_progress [#call, nil] deprecated compatibility callback. It
+    #   receives the same per-item events plus the legacy item-less +"complete"+
+    #   event. Ignored when +on_item_finish+ is also provided.
+    # @param on_item_start [#call, nil] optional callback invoked when a worker
+    #   begins processing each item. Pair it with on_item_finish to distinguish
+    #   queued work from in-flight work. A raising callback never crashes the run.
     # @return [Hash] with :items, :test_run_id, :test_run_url
     def replay(receiver, method_name, trace_function_key:, limit: nil, trace_ids: nil, max_concurrency: 10,
       name: nil, code_change_description: Replay::CODE_CHANGE_UNSET, code_change_files: Replay::CODE_CHANGE_UNSET, experiment_group_id: nil, dataset_id: nil, grader_ids: nil, mock: "marked",
-      adapt_inputs: nil, mock_override: nil, db_branch: nil, on_progress: nil)
+      adapt_inputs: nil, mock_override: nil, db_branch: nil, on_item_start: nil, on_item_finish: nil, on_progress: nil)
       Replay.run(
         self, receiver, method_name,
         trace_function_key:, limit:, trace_ids:, name:, max_concurrency:,
         code_change_description:, code_change_files:, experiment_group_id:, dataset_id:, grader_ids:, mock:, adapt_inputs:,
-        mock_override:, db_branch:,
+        mock_override:, db_branch:, on_item_start:, on_item_finish:,
         on_progress:
       )
     end

data/lib/bitfab/replay.rb

@@ -12,6 +12,30 @@ require_relative "traceable"
 require_relative "transport"
 
 module Bitfab
+  # A requested replay database branch could not be resolved. The resolver code
+  # and original message remain structured for callers to inspect.
+  class DbBranchReplayError < StandardError
+    attr_reader :code, :original_trace_id
+
+    def initialize(code, message, original_trace_id)
+      super(message)
+      @code = code
+      @original_trace_id = original_trace_id
+    end
+  end
+
+  # Whole-run replay failure with every item collected before the run failed.
+  class ReplayError < StandardError
+    attr_reader :items, :test_run_id, :test_run_url
+
+    def initialize(message, items:, test_run_id:, test_run_url:)
+      super(message)
+      @items = items
+      @test_run_id = test_run_id
+      @test_run_url = test_run_url
+    end
+  end
+
   # Replay mock strategies. Mirrors the Python and TypeScript SDKs.
   #
   # - "marked" : only spans declared with mock_on_replay: true return historical
@@ -85,6 +109,50 @@ module Bitfab
 
     module_function
 
+    def json_safe(value)
+      case value
+      when Exception
+        serialized = {type: value.class.name, message: value.message, backtrace: value.backtrace}
+        if value.is_a?(DbBranchReplayError)
+          serialized[:code] = value.code
+          serialized[:original_trace_id] = value.original_trace_id
+          serialized[:cause] = json_safe(value.cause) if value.cause
+        end
+        serialized
+      when Hash
+        value.to_h { |key, child| [key, json_safe(child)] }
+      when Array
+        value.map { |child| json_safe(child) }
+      else
+        value
+      end
+    end
+
+    def replay_item_error_message(error)
+      return error.message unless error.is_a?(DbBranchReplayError)
+
+      "Replay requested a database branch for trace #{error.original_trace_id} but it " \
+        "could not be resolved (#{error.code}): #{error.message}. The method was not run, " \
+        "because replaying it against the live database would produce a result that looks " \
+        "valid but did not use the historical data you asked for."
+    end
+
+    def public_replay_items(items)
+      items.map { |item| item.except(:_sdk_trace_id) }
+    end
+
+    def preserve_replay_failure(items, test_run_id, test_run_url)
+      yield
+    rescue => cause
+      error = ReplayError.new(
+        cause.message,
+        items: public_replay_items(items),
+        test_run_id:,
+        test_run_url:
+      )
+      raise error, cause:
+    end
+
     # Replay historical traces through a method and create a test run.
     #
     # Fetches the last N traces for the given trace function key, re-runs each
@@ -133,11 +201,11 @@ module Bitfab
     #   and ceiling in Neon Compute Units, and +:warmup_sql+ warms the branch's
     #   cache before the replayed method sees it. Read the resolved branch
     #   inside the method with +Bitfab.current_replay_branch+.
-    # @param on_progress [#call, nil] optional callback invoked once per item as
+    # @param on_item_finish [#call, nil] optional callback invoked once per item as
     #   it finishes, with a running-totals hash { completed:, total:, succeeded:,
     #   errored:, item: } where item is { trace_id:, original_trace_id:,
     #   original_span_id:, error:, duration_ms: } for the single item that just
-    #   settled (source_trace_id/source_span_id remain as deprecated aliases).
+    #   finished (source_trace_id/source_span_id remain as deprecated aliases).
     #   trace_id is the new server replay trace id, written in after the run
     #   completes (nil during progress callbacks); original_trace_id is the
     #   ORIGINAL (historical) trace that was replayed, error is that item's replay
@@ -145,14 +213,22 @@ module Bitfab
     #   duration_ms is how long that one trace took to replay. Use it to
     #   render replay progress (e.g. a per-trace log). A raising callback never
     #   crashes the run.
+    # @param on_progress [#call, nil] deprecated compatibility callback. It
+    #   receives the same per-item finish events plus the legacy whole-run
+    #   +"complete"+ event. Ignored when +on_item_finish+ is also provided.
+    # @param on_item_start [#call, nil] optional callback invoked when a worker
+    #   begins processing each item, before replay setup and customer code run.
+    #   Pair it with on_item_finish to distinguish queued work from in-flight work.
+    #   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, grader_ids: nil, mock: "marked",
-      adapt_inputs: nil, mock_override: nil, db_branch: nil, on_progress: nil)
+      adapt_inputs: nil, mock_override: nil, db_branch: nil, on_item_start: nil, on_item_finish: nil, on_progress: nil)
       unless MOCK_STRATEGIES.include?(mock.to_s)
         raise ArgumentError, "Invalid mock strategy '#{mock}'. Must be one of: #{MOCK_STRATEGIES.join(", ")}"
       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?
         if trace_ids.length > 100
@@ -223,11 +299,13 @@ module Bitfab
       )
       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"] || []
 
       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_progress:,
+          adapt_inputs, include_db_branch_lease, resolved_db_branch_settings, on_item_start:,
+          on_item_finish: item_finish_callback,
           mock_overrides: resolved_overrides)
       else
         []
@@ -237,12 +315,16 @@ module Bitfab
       # to finalize once the server confirms every replay trace it queued: the
       # trace-ID mapping complete_replay builds would otherwise race the
       # in-flight batches and hand back nil for every item.
-      wait_for_replay_persistence(http_client, test_run_id, result_items.map { |item| item[:_sdk_trace_id] })
+      preserve_replay_failure(result_items, test_run_id, full_test_run_url) do
+        wait_for_replay_persistence(http_client, test_run_id, result_items.map { |item| item[:_sdk_trace_id] })
+      end
 
       # complete_replay finalizes the run and returns the token/diagnostic
       # mapping; its failures propagate loudly because a run that never
       # completed can't be finalized.
-      complete_response = http_client.complete_replay(test_run_id)
+      complete_response = preserve_replay_failure(result_items, test_run_id, full_test_run_url) do
+        http_client.complete_replay(test_run_id)
+      end
       trace_id_map = complete_response&.dig("traceIds")
       # Per-replay-trace token usage keyed by server trace id: the REPLAYED
       # run's tokens (span-aggregated server-side), used below to fill each
@@ -278,11 +360,18 @@ module Bitfab
         if completed_count.positive? && missing.length == completed_count
           trace_count = complete_response["traceCount"]
           server_count = trace_count.nil? ? "" : " The server persisted #{trace_count} trace(s) for this run."
-          raise "Replay completed but the server has no persisted trace for " \
+          cause = RuntimeError.new("Replay completed but the server has no persisted trace for " \
             "any of the #{completed_count} completed item(s) " \
             "(test_run_id #{test_run_id}).#{server_count} Trace uploads were " \
             "flushed, so either the uploads failed or the replayed method is " \
-            "not traced (no root span was emitted)."
+            "not traced (no root span was emitted).")
+          error = ReplayError.new(
+            cause.message,
+            items: public_replay_items(result_items),
+            test_run_id:,
+            test_run_url: full_test_run_url
+          )
+          raise error, cause:
         end
         # SOME completed items missing: per-item upload failure. Warn, but
         # return the run; the items that landed can still be labeled.
@@ -301,16 +390,14 @@ module Bitfab
       result = {
         items: result_items,
         test_run_id:,
-        test_run_url: "#{client.service_url}#{test_run_url}"
+        test_run_url: full_test_run_url
       }
-      # Persist the enriched result two ways so the Bitfab plugin never has to
-      # parse the replay's stdout (which a dependency's logging can corrupt):
-      # write it to BITFAB_REPLAY_RESULT_PATH when the plugin set that env var,
-      # and stream a terminal "complete" progress event. The plugin prefers the
-      # streamed event and falls back to the file. The event routes through
-      # on_progress so only progress-reporting runs emit it.
+      # Persist the enriched result so the Bitfab plugin never has to parse the
+      # replay's stdout, which a dependency's logging can corrupt.
       write_replay_result_file(result)
-      if on_progress
+      # Preserve the legacy terminal event only for on_progress. on_item_finish
+      # is strictly item-scoped, so every invocation always includes an item.
+      if on_item_finish.nil? && on_progress
         errored = result_items.count { |item| !item[:error].nil? }
         total = result_items.length
         begin
@@ -543,7 +630,7 @@ module Bitfab
       begin
         dir = File.dirname(result_path)
         FileUtils.mkdir_p(dir) unless dir.nil? || dir.empty? || dir == "."
-        File.write(result_path, "#{JSON.pretty_generate(result)}\n")
+        File.write(result_path, "#{Bitfab.serialize_replay_result(result)}\n")
       rescue => e
         warn "Bitfab: failed to write replay result to BITFAB_REPLAY_RESULT_PATH (#{result_path}): #{e.message}"
       end
@@ -551,36 +638,56 @@ 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_progress: nil,
+      adapt_inputs = nil, include_db_branch_lease = false, db_branch_settings = nil, on_item_start: nil, on_item_finish: nil,
       mock_overrides: [])
       concurrency = max_concurrency || server_items.length
 
-      # Reports running totals once per item as it settles. In the parallel
-      # path it runs from worker threads, so the mutex both makes the counter
-      # updates safe and serializes the user's callback (never called
-      # concurrently). A raising callback is swallowed: progress UI must never
-      # crash the run.
+      # Lifecycle callbacks run from worker threads in the parallel path, so the
+      # mutex makes counter updates safe and serializes each callback. A raising
+      # callback is swallowed because progress UI must never crash the run.
       total = server_items.length
       progress_mutex = Mutex.new
+      started = 0
       completed = 0
       succeeded = 0
       errored = 0
-      # Each event carries the single item that just settled so a progress UI
+      # Each event carries the single item that just finished so a progress UI
       # can render per-trace pass/fail as the run streams. The item's :trace_id
       # is nil at this stage: the server replay trace id isn't known until run()
       # writes it in after complete_replay, and the client correlation id is
       # never surfaced. original_trace_id (the historical trace being replayed,
       # taken from the server item) is what a UI keys on to identify what just
-      # settled. source_trace_id/source_span_id are kept as deprecated aliases.
-      report = lambda do |result, original_trace_id, original_span_id, test_run_id|
-        return unless on_progress
+      # finished. source_trace_id/source_span_id are kept as deprecated aliases.
+      report_start = lambda do |original_trace_id, original_span_id|
+        progress_mutex.synchronize do
+          started += 1
+          next unless on_item_start
 
+          begin
+            on_item_start.call({
+              type: "started", test_run_id:, started:, completed:, total:, succeeded:, errored:,
+              item: {
+                original_trace_id:,
+                original_span_id:,
+                source_trace_id: original_trace_id,
+                source_span_id: original_span_id
+              }
+            })
+          rescue => e
+            warn "Bitfab: replay on_item_start callback raised: #{e.message}"
+          end
+        end
+      end
+
+      report = lambda do |result, original_trace_id, original_span_id, test_run_id|
         progress_mutex.synchronize do
           completed += 1
           error = result[:error]
           error.nil? ? (succeeded += 1) : (errored += 1)
+          next unless on_item_finish
+
           begin
-            on_progress.call({
+            on_item_finish.call({
               test_run_id:, completed:, total:, succeeded:, errored:,
               item: {
                 trace_id: result[:trace_id],
@@ -593,6 +700,8 @@ module Bitfab
                 result: result[:result],
                 original_output: result[:original_output],
                 error:,
+                trace_error: result[:trace_error],
+                replay_error: result[:replay_error],
                 duration_ms: result[:duration_ms],
                 tokens: result[:tokens],
                 model: result[:model],
@@ -600,13 +709,14 @@ module Bitfab
               }
             })
           rescue => e
-            warn "Bitfab: replay on_progress callback raised: #{e.message}"
+            warn "Bitfab: replay on_item_finish callback raised: #{e.message}"
           end
         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.call(result, original_trace_id_of(item), original_span_id_of(item), test_run_id)
@@ -624,6 +734,7 @@ 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:)
               results_mutex.synchronize { results[idx] = result }
@@ -664,16 +775,26 @@ module Bitfab
       lease_error = include_db_branch_lease ? server_item["dbBranchLeaseError"] : nil
       db_snapshot_ref = server_item["dbSnapshotRef"]
       if include_db_branch_lease && lease.nil? && lease_error.nil?
-        resolved = http_client.resolve_db_branch_lease(test_run_id, original_trace_id, db_branch_settings)
+        begin
+          resolved = http_client.resolve_db_branch_lease(test_run_id, original_trace_id, db_branch_settings)
+        rescue => cause
+          error = DbBranchReplayError.new(
+            "lease_request_failed",
+            "Bitfab could not request the database branch: #{cause.message}",
+            original_trace_id
+          )
+          raise error, cause:
+        end
         lease = resolved["lease"]
         lease_error = resolved["leaseError"]
         db_snapshot_ref = resolved["dbSnapshotRef"] || db_snapshot_ref
       end
       if lease_error
-        raise "Replay requested a database branch for trace #{original_trace_id} but it " \
-          "could not be resolved (#{lease_error["code"]}): #{lease_error["message"]}. " \
-          "The method was not run, because replaying it against the live database would " \
-          "produce a result that looks valid but did not use the historical data you asked for."
+        raise DbBranchReplayError.new(
+          lease_error["code"].to_s,
+          lease_error["message"].to_s,
+          original_trace_id
+        )
       end
 
       span = http_client.get_external_span(original_span_id, replay_view: true)
@@ -752,7 +873,9 @@ module Bitfab
         input: [],
         result: nil,
         original_output: nil,
-        error: e.message,
+        error: replay_item_error_message(e),
+        trace_error: nil,
+        replay_error: e,
         duration_ms: metrics&.dig(:duration_ms),
         tokens: metrics&.dig(:tokens),
         model: metrics&.dig(:model),
@@ -941,51 +1064,64 @@ module Bitfab
 
       fn_result = nil
       fn_error = nil
+      replay_error = nil
       # Client-side correlation id that tags this item's replay spans so the
       # server can echo back the row id it minted (resolved in run()'s
       # complete-replay loop). Carried on the item under :_sdk_trace_id, never
       # surfaced as the public :trace_id.
       sdk_trace_id = SecureRandom.uuid
 
-      ReplayContext.with_context(
-        test_run_id:,
-        input_source_span_id:,
-        input_source_trace_id:,
-        trace_id: sdk_trace_id,
-        mock_tree:,
-        mock_strategy:,
-        mock_overrides:,
-        fetch_span_output:,
-        db_branch_lease:,
-        source_bitfab_trace_id:
-      ) do
-        # Reshape recorded inputs onto the current signature when an adapter is
-        # supplied. Inside the rescue so a raising adapter surfaces on this
-        # item's :error instead of crashing the run; args is reported on :input.
-        if adapt_inputs
-          ctx = adapt_ctx || {
-            original_trace_id: nil,
-            original_span_id: input_source_span_id,
-            # Deprecated aliases for original_trace_id/original_span_id.
-            source_trace_id: nil,
-            source_span_id: input_source_span_id
-          }
-          args, kwargs = adapt_inputs.call(args, kwargs, ctx)
-        end
-        fn_result = if kwargs.empty?
-          receiver.send(method_name, *args)
-        else
-          receiver.send(method_name, *args, **kwargs)
+      begin
+        ReplayContext.with_context(
+          test_run_id:,
+          input_source_span_id:,
+          input_source_trace_id:,
+          trace_id: sdk_trace_id,
+          mock_tree:,
+          mock_strategy:,
+          mock_overrides:,
+          fetch_span_output:,
+          db_branch_lease:,
+          source_bitfab_trace_id:
+        ) do
+          begin
+            if adapt_inputs
+              ctx = adapt_ctx || {
+                original_trace_id: nil,
+                original_span_id: input_source_span_id,
+                # Deprecated aliases for original_trace_id/original_span_id.
+                source_trace_id: nil,
+                source_span_id: input_source_span_id
+              }
+              args, kwargs = adapt_inputs.call(args, kwargs, ctx)
+            end
+          rescue => e
+            replay_error = e
+          end
+          if replay_error.nil?
+            begin
+              fn_result = if kwargs.empty?
+                receiver.send(method_name, *args)
+              else
+                receiver.send(method_name, *args, **kwargs)
+              end
+            rescue => e
+              fn_error = e
+            end
+          end
         end
       rescue => e
-        fn_error = e.message
+        replay_error = e
       end
 
+      item_error = fn_error || replay_error
       {
         input: args,
         result: fn_result,
         original_output: item["output"],
-        error: fn_error,
+        error: item_error&.message,
+        trace_error: fn_error,
+        replay_error:,
         duration_ms: metrics[:duration_ms],
         tokens: metrics[:tokens],
         model: metrics[:model],

data/lib/bitfab/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Bitfab
-  VERSION = "0.36.2"
+  VERSION = "0.36.4"
 end

data/lib/bitfab.rb

@@ -144,24 +144,27 @@ module Bitfab
       CurrentTrace.new(entry[:trace_id])
     end
 
-    # A ready-made on_progress callback for replay: writes one @@bitfab:progress
-    # line per trace to stderr, which the Bitfab plugin polls to report live
-    # progress while the replay runs in the background, so replay scripts never
-    # hand-format the protocol. stdout stays the ReplayResult JSON. Never raises
-    # (progress must not crash a run). The progress hash's item is forwarded
-    # verbatim through to_json, so the plugin sees the source trace id, local
-    # replay trace id, outputs, error, and duration of the item that just settled.
+    # A ready-made replay lifecycle reporter. Pass it as both on_item_start and
+    # on_item_finish so the Bitfab plugin sees in-flight and finished items. It writes
+    # @@bitfab:progress lines to stderr, leaving stdout for ReplayResult JSON, and
+    # never raises because progress reporting must not crash a run.
     #
     # @example
-    #   Bitfab.client.replay(..., on_progress: Bitfab.method(:report_replay_progress))
+    #   reporter = Bitfab.method(:report_replay_progress)
+    #   Bitfab.client.replay(..., on_item_start: reporter, on_item_finish: reporter)
     #
     # @param progress [Hash] running totals with keys completed, total,
     #   succeeded, errored, and item (the object replay already passes to
-    #   on_progress)
+    #   on_item_finish)
     def report_replay_progress(progress)
-      warn "#{BITFAB_PROGRESS_PREFIX}#{progress.to_json}"
+      warn "#{BITFAB_PROGRESS_PREFIX}#{Replay.json_safe(progress).to_json}"
     rescue
       nil
     end
+
+    # Serialize a replay result while retaining structured exception fields.
+    def serialize_replay_result(result)
+      JSON.pretty_generate(Replay.json_safe(result))
+    end
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bitfab
 version: !ruby/object:Gem::Version
-  version: 0.36.2
+  version: 0.36.4
 platform: ruby
 authors:
 - Harvest Team