bitfab 0.36.10 → 0.38.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 47a3e6b763bf08e8e4a25173414f1ed306dc7183ea443f0b71264645fe5b4c02
-  data.tar.gz: 2c629b4d599e25e124334f1d9ee61b5683316c7b849535ffa0aabe50324105eb
+  metadata.gz: e4bb59de5ab0ec989d85e8be14ee08ad935ada5145e01826b34657fc8c9e6623
+  data.tar.gz: 78b0255c5a3062470f4149baf632c94a5b18b430567485a1ba58f5b9cf9b9534
 SHA512:
-  metadata.gz: 30dbfa2213bd23581fa3387d47a600ad79ed62f8120ae11e237557b87aa28f8d1413162b0b9ab9f9ddc9d5489a0b884793e632a29538369e40753327384d6bf2
-  data.tar.gz: 7cb61050f45cc07f273f8bcc22f8ca8d0530d9b67386b22472906225eedc8d5937e2f8158a7810b0b211a1c0a558994c0c675f8e17a8efa633a882cc0d8fd596
+  metadata.gz: e127f198a39372bebaff5a0cf395b4b6b00f2ede4bf9768b8e3a4556ee00271ddf44fe4a40f16fcffac72d3def4d3c83c3f756801ba156edbe58ecbfd189f2e5
+  data.tar.gz: ca803b3df3783f6bb01d5eb5a518300334d0ccd4f31d43f3870c981779cd69710d076a1ded3f058c3b9e8aa65f643b3872dd60f99ac8e4394cd5f97855ec1970

data/README.md

@@ -301,7 +301,7 @@ client.replay(pipeline, :process, trace_function_key: "my-fn", mock: "marked")
 # "none": everything runs real code
 client.replay(pipeline, :process, trace_function_key: "my-fn", mock: "none")
 
-# "all": every child span returns its historical output
+# "all": every matched recorded child returns history; missing occurrences fail closed
 client.replay(pipeline, :process, trace_function_key: "my-fn", mock: "all")
 ```
 
@@ -331,7 +331,7 @@ class Pipeline
 end
 ```
 
-Use the default `mock: "marked"` behavior when you want to iterate on `process`'s logic without paying for the LLM call each run. Use `mock: "all"` for the cheapest possible replay (every child span returns its recorded output).
+Use the default `mock: "marked"` behavior when you want to iterate on `process`'s logic without paying for the LLM call each run. Use `mock: "all"` for the cheapest possible replay (every matched recorded child span returns its recorded output; a missing occurrence fails the item closed).
 
 ### Error Handling
 

data/lib/bitfab/client.rb

@@ -16,7 +16,7 @@ module Bitfab
     UUID_PATTERN = /\A[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\z/i
 
     # Sentinel returned by check_mock_replay when this span should run real
-    # code (no mock active, wrong strategy, or no matching historical entry).
+    # code (no mock active or the base strategy did not select it).
     # Using a sentinel rather than nil/false avoids confusing legitimate mocked
     # outputs (which may themselves be nil or false).
     MOCK_REPLAY_MISS = Object.new.freeze
@@ -87,9 +87,10 @@ module Bitfab
     # @param receiver [Object, Class] an instance for instance methods, or a Class for class methods
     # @param method_name [Symbol] the method to replay
     # @param trace_function_key [String] the trace function key for this method
-    # @param limit [Integer, nil] maximum number of traces to replay (default: 5).
-    #   Mutually exclusive with trace_ids: an explicit ID list already
-    #   determines how many traces replay, so passing both raises.
+    # @param limit [Integer, nil] maximum number of traces to replay (default: 5;
+    #   maximum: 5,000). Ignored when trace_ids or dataset_id is passed because
+    #   either source already determines how many traces replay. Supplying
+    #   trace_ids also emits a warning.
     # @param trace_ids [Array<String>, nil] optional list of trace IDs to replay (max 100)
     # @param max_concurrency [Integer, nil] max threads for parallel replay (default: 10)
     # @param code_change_description [String, nil] optional rationale for the
@@ -111,7 +112,9 @@ module Bitfab
     #   same org and trace function or the server rejects the replay
     # @param mock [String] mock strategy for child spans: "marked" (default),
     #   "none", or "all". "marked" only mocks spans declared with
-    #   mock_on_replay: true; "all" mocks every child span.
+    #   mock_on_replay: true; "all" mocks every matched recorded child span. A selected
+    #   occurrence that is unavailable errors the item without executing the
+    #   real child.
     # @param adapt_inputs [#call, nil] optional hook to reshape recorded inputs
     #   onto the method's current signature when its shape changed after the
     #   traces were captured. Receives (args, kwargs, ctx) where ctx is
@@ -122,16 +125,17 @@ module Bitfab
     #   match.call(node) selects spans to substitute (node is
     #   { trace_function_key:, span_name:, type:, original_span_id: }). value is
     #   EITHER a flat value injected directly OR a callable invoked with
-    #   { node:, inputs:, get_original_output: }; its result becomes the span's
-    #   output (full replacement), so downstream real code runs against it. The
+    #   { node:, inputs:, kwargs:, get_original_output: }; its result becomes
+    #   the span's output (full replacement), so downstream real code runs against it. The
     #   first matching override wins. Per-call overrides take precedence over
     #   those registered via register_mock_override, and both take precedence
     #   over the base mock strategy. The root span is never overridden.
-    # @param db_branch [Hash, nil] passing it at all turns database branching
-    #   on, so +{}+ enables it with the mirror project's own sizing. Keys
-    #   +: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 db_branch [Boolean, Hash, nil] +true+ or a Hash enables database
+    #   branching; +false+/+nil+ leave it off. An empty Hash uses the mirror
+    #   project's own sizing. Keys +: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_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
@@ -281,7 +285,7 @@ module Bitfab
         )
         resolved_capture_when = "always"
       end
-      return yield if resolved_capture_when == "nested" && SpanContext.current.nil?
+      return yield if resolved_capture_when == "nested" && SpanContext.current.nil? && ReplayContext.current_span.nil?
 
       # Span setup runs before the user's block. Tracing is a side-channel, so
       # if anything here raises (id generation, trace-state bookkeeping, a
@@ -297,8 +301,8 @@ module Bitfab
       resolved_test_run_id = nil
       resolved_input_source_span_id = nil
       begin
-        parent = SpanContext.current
         replay_ctx = ReplayContext.current
+        parent = SpanContext.current || ReplayContext.current_span
         trace_id = parent ? parent[:trace_id] : (replay_ctx&.dig(:trace_id) || SecureRandom.uuid)
         span_id = SecureRandom.uuid
         parent_span_id = parent&.dig(:span_id)
@@ -420,7 +424,12 @@ module Bitfab
                 # Deprecated wire alias, kept so this SDK still reports usage
                 # against servers that predate the rename.
                 source_trace_id: replay_ctx[:source_bitfab_trace_id],
-                accessed: replay_ctx[:db_snapshot_accessed] == true
+                accessed: replay_ctx[:db_snapshot_accessed] == true,
+                # Echoed verbatim (camelCase inside) rather than re-cased into
+                # this record's snake_case: it is the server's own object
+                # coming back, and a translation layer here is one more thing
+                # to drift.
+                timings: replay_ctx[:db_branch_timings]
               }
             end
 
@@ -440,12 +449,14 @@ module Bitfab
       end
 
       begin
-        SpanContext.with_span(trace_id:, span_id:) do
-          result = yield
-        ensure
-          # Capture contexts before the span context is popped
-          span_contexts = SpanContext.current&.dig(:contexts)
-          span_prompt = SpanContext.current&.dig(:prompt)
+        ReplayContext.with_span(trace_id:, span_id:) do
+          SpanContext.with_span(trace_id:, span_id:) do
+            result = yield
+          ensure
+            # Capture contexts before the span context is popped
+            span_contexts = SpanContext.current&.dig(:contexts)
+            span_prompt = SpanContext.current&.dig(:prompt)
+          end
         end
       rescue => e
         error = e.message
@@ -466,12 +477,6 @@ module Bitfab
       # collects yielded values as the span output, and finalizes the span
       # once iteration completes (or errors).
       #
-      # Limitation: when the source enumerator itself runs its body in a
-      # separate fiber (e.g. `Enumerator.new { |y| ... }` or `enum_for(...)`
-      # without a block), nested `bitfab_span` calls inside that body fiber
-      # still see an empty stack because `Thread.current[STACK_KEY]` is
-      # fiber-local. Lazy chains over collections (`.lazy.map`) and ordinary
-      # `each` callbacks DO run in the iterating fiber and nest correctly.
       if result.is_a?(Enumerator)
         return wrap_enumerator(result, trace_id:, span_id:, finalize:)
       end
@@ -544,19 +549,24 @@ module Bitfab
       Enumerator.new do |yielder|
         # Push our span onto this fiber's stack so anything that runs in this
         # fiber (including the user's lazy block and `each` callbacks) sees
-        # the right parent.
-        SpanContext.stack.push(span_entry)
-        begin
-          source.each do |value|
-            yielded << value
-            yielder << value
+        # the right parent. The fiber bridge carries the same parent into
+        # Enumerator.new / enum_for source fibers on this thread.
+        ReplayContext.with_span(trace_id:, span_id:) do
+          SpanContext.with_fiber_bridge(trace_id:, span_id:) do
+            SpanContext.stack.push(span_entry)
+            begin
+              source.each do |value|
+                yielded << value
+                yielder << value
+              end
+              finalize.call(yielded, nil)
+            rescue => e
+              finalize.call(yielded, e.message)
+              raise
+            ensure
+              SpanContext.stack.pop
+            end
           end
-          finalize.call(yielded, nil)
-        rescue => e
-          finalize.call(yielded, e.message)
-          raise
-        ensure
-          SpanContext.stack.pop
         end
       end
     end
@@ -612,6 +622,9 @@ module Bitfab
           usage["source_trace_id"] = db_snapshot_usage[:original_trace_id]
         end
         usage["accessed"] = db_snapshot_usage[:accessed]
+        if db_snapshot_usage[:timings]
+          usage["timings"] = db_snapshot_usage[:timings]
+        end
         raw_trace["db_snapshot_usage"] = usage
       end
 
@@ -799,7 +812,11 @@ module Bitfab
         return MOCK_REPLAY_MISS
       end
 
-      return MOCK_REPLAY_MISS unless mock_entry
+      unless mock_entry
+        raise StandardError,
+          "Replay selected span '#{trace_function_key}:#{span_name}' for mocking, " \
+          "but recorded occurrence #{call_index + 1} is unavailable. The real span was not executed."
+      end
 
       [resolve_recorded_output(mock_entry, replay_ctx), "recorded"]
     end
@@ -902,7 +919,8 @@ module Bitfab
     # @param capture_when [String, Symbol] "always" to allow a root trace, or
     #   "nested" to capture only when another Bitfab span is active. Unknown
     #   values warn once and default to "always"
-    # @param mock_on_replay [Boolean] mark this span for the default "marked" mock strategy
+    # @param mock_on_replay [Boolean] mark this span for the default "marked"
+    #   mock strategy. A missing selected occurrence fails closed.
     def wrap(klass, method_name, name: nil, type: "custom", capture_when: "always", mock_on_replay: false)
       Bitfab::Traceable.wrap(
         klass, method_name,

data/lib/bitfab/mock_override.rb

@@ -31,10 +31,10 @@ module Bitfab
   #   nil is a legitimate injected value. To inject a proc as the literal
   #   output, wrap it in a callable value: value: ->(ctx) { the_proc }.
   #
-  # Unlike the TypeScript SDK, Ruby's span tree is fetched with outputs inline
-  # (get_span_tree), so get_original_output simply returns the inline recorded
-  # output. There is no lazy per-span fetch and Ruby is synchronous, so there
-  # is no async limitation.
+  # Unlike the TypeScript SDK, get_original_output is synchronous. Under
+  # non-"all" replay the first access may fetch the recorded output lazily;
+  # that fetch is memoized for the replay item. Under mock: "all" the output is
+  # inline even when overrides are present. Ruby has no async-span limitation.
   module MockOverride
     module_function
 

data/lib/bitfab/replay.rb

@@ -41,28 +41,47 @@ module Bitfab
   # - "marked" : only spans declared with mock_on_replay: true return historical
   #              output; everything else runs real code (default)
   # - "none"   : every child span runs real code
-  # - "all"    : every child span returns its historical output
+  # - "all"    : every matched recorded child span returns its historical
+  #              output; a missing occurrence fails the item closed
   MOCK_STRATEGIES = %w[none all marked].freeze
 
-  # Thread-local replay context management.
+  # Replay context shared by fibers on the current thread. Ruby's
+  # Thread#[] storage is fiber-local, so use true thread variables here: lazy
+  # Enumerator bodies run in another fiber but must retain replay mocking.
   module ReplayContext
     module_function
 
     def current
-      Thread.current[REPLAY_CONTEXT_KEY]
+      Thread.current.thread_variable_get(REPLAY_CONTEXT_KEY)
     end
 
-    # Execute a block with replay context set on the current thread.
-    # The context is automatically cleared when the block completes.
+    def current_span
+      current&.dig(:span_stack)&.last
+    end
+
+    def with_span(trace_id:, span_id:)
+      ctx = current
+      return yield unless ctx
+
+      entry = {trace_id:, span_id:}
+      ctx[:span_stack].push(entry)
+      yield
+    ensure
+      ctx[:span_stack].pop if ctx
+    end
+
+    # Execute a block with replay context set for every fiber on the current
+    # 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, source_bitfab_trace_id: nil)
-      previous = Thread.current[REPLAY_CONTEXT_KEY]
+      db_branch_lease: nil, db_branch_timings: nil, source_bitfab_trace_id: nil)
+      previous = Thread.current.thread_variable_get(REPLAY_CONTEXT_KEY)
       ctx = {
         test_run_id:,
         input_source_span_id:,
         input_source_trace_id:,
-        trace_id:
+        trace_id:,
+        span_stack: []
       }
       if mock_tree
         ctx[:mock_tree] = mock_tree
@@ -73,8 +92,9 @@ module Bitfab
         # is present. process_single_item fetches the tree whenever overrides
         # exist, even under mock: "none".
         ctx[:mock_overrides] = mock_overrides if mock_overrides && !mock_overrides.empty?
-        # Lazy per-span output fetcher for the payload-free ("marked"/override)
-        # path. Absent on the eager "all" path, whose outputs are inline.
+        # Lazy per-span output fetcher for a payload-free non-"all" tree path.
+        # Absent on the eager "all" path, whose outputs are inline even when
+        # overrides are present.
         ctx[:fetch_span_output] = fetch_span_output if fetch_span_output
       end
       # The per-trace DB branch (resolved server-side) and the Bitfab trace ID
@@ -93,11 +113,15 @@ module Bitfab
       # accessed for free, and the flag would stop separating "branch was
       # used" from "branch was offered".
       ctx[:db_branch_lease] = db_branch_lease if db_branch_lease
+      # Kept off ReplayBranch: customer code reads that mid-replay to reach the
+      # branch, and provisioning latency is a property of the run, not of the
+      # connection. It rides the context only to be echoed on the completion.
+      ctx[:db_branch_timings] = db_branch_timings if db_branch_timings
       ctx[:source_bitfab_trace_id] = source_bitfab_trace_id if source_bitfab_trace_id
-      Thread.current[REPLAY_CONTEXT_KEY] = ctx
+      Thread.current.thread_variable_set(REPLAY_CONTEXT_KEY, ctx)
       yield
     ensure
-      Thread.current[REPLAY_CONTEXT_KEY] = previous
+      Thread.current.thread_variable_set(REPLAY_CONTEXT_KEY, previous)
     end
   end
 
@@ -163,8 +187,9 @@ module Bitfab
     # @param method_name [Symbol] the method to replay
     # @param trace_function_key [String] the trace function key for this method
     # @param limit [Integer, nil] maximum number of traces to replay (default: 5).
-    #   Ignored when trace_ids is passed (with a warning): an explicit ID list
-    #   already determines how many traces replay.
+    #   Ignored when trace_ids or dataset_id is passed because either source
+    #   already determines how many traces replay. Supplying trace_ids also
+    #   emits a warning.
     # @param trace_ids [Array<String>, nil] optional list of trace IDs to replay (max 100)
     # @param name [String, nil] optional display name for the resulting experiment/test run
     # @param max_concurrency [Integer, nil] max threads for parallel replay (default: 10)
@@ -187,7 +212,7 @@ module Bitfab
     #   same org and trace function or the server rejects the replay
     # @param mock [String] mock strategy for child spans: "marked" (default),
     #   "none", or "all". "marked" only mocks spans declared with
-    #   mock_on_replay: true; "all" mocks every child span.
+    #   mock_on_replay: true; "all" mocks every matched recorded child span.
     # @param adapt_inputs [#call, nil] optional hook to reshape recorded inputs
     #   onto the method's current signature when its shape changed after the
     #   traces were captured. Receives (args, kwargs, ctx) where ctx is
@@ -206,11 +231,11 @@ module Bitfab
     #   errored:, item: } where item is { trace_id:, original_trace_id:,
     #   original_span_id:, error:, duration_ms: } for the single item that just
     #   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
-    #   error or nil, and
-    #   duration_ms is how long that one trace took to replay. Use it to
+    #   trace_id is the new server replay trace id, available after that item's
+    #   trace is flushed (nil only when delivery could not be confirmed);
+    #   original_trace_id is the ORIGINAL (historical) trace that was replayed,
+    #   error is that item's replay error or nil, and duration_ms is the source
+    #   trace's recorded duration. 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
@@ -744,9 +769,13 @@ module Bitfab
                 trace_error: result[:trace_error],
                 replay_error: result[:replay_error],
                 duration_ms: result[:duration_ms],
+                original_duration_ms: result[:original_duration_ms],
+                original_tokens: result[:original_tokens],
+                original_model: result[:original_model],
                 tokens: result[:tokens],
                 model: result[:model],
-                db_snapshot_ref: result[:db_snapshot_ref]
+                db_snapshot_ref: result[:db_snapshot_ref],
+                db_branch_timings: result[:db_branch_timings]
               }
             })
           rescue => e
@@ -807,14 +836,17 @@ module Bitfab
       # (only when include_db_branch_lease was sent). Release it in the +ensure+
       # below so any raise (span fetch, mock-tree build, or the replayed
       # method) frees the Neon resource. Items whose source trace had no
-      # snapshot ref arrive without a lease (env.active? is false), and
-      # replaying those against the live database is correct.
+      # snapshot ref arrive without a lease, so the app uses its normal DB
+      # path. Unsafe calls on that path still require replay mocking.
       lease = include_db_branch_lease ? server_item["dbBranchLease"] : nil
       # A resolve that was ATTEMPTED and failed is different: the caller asked
       # for a branch, so running their method against live data would produce a
       # result that looks valid and is not. Fail the item instead, loudly.
       lease_error = include_db_branch_lease ? server_item["dbBranchLeaseError"] : nil
       db_snapshot_ref = server_item["dbSnapshotRef"]
+      # Reported whichever way the resolve went, so it is tracked separately
+      # from both the lease and the error rather than hanging off either.
+      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)
@@ -829,6 +861,7 @@ module Bitfab
         lease = resolved["lease"]
         lease_error = resolved["leaseError"]
         db_snapshot_ref = resolved["dbSnapshotRef"] || db_snapshot_ref
+        db_branch_timings = resolved["timings"] || db_branch_timings
       end
       if lease_error
         raise DbBranchReplayError.new(
@@ -847,10 +880,11 @@ module Bitfab
       # the tree must be fetched for them to fire even under mock: "none".
       #
       # Only mock: "all" needs every span's recorded output inline (it mocks
-      # every child), so it fetches an eager tree. "marked" and overrides mock
-      # only a few spans, so they fetch a payload-free tree (includeOutputs=
-      # false) and pull each mocked span's output lazily by externalSpanId,
-      # never dragging down outputs that no span consumes.
+      # every matched recorded child), so it fetches an eager tree even when
+      # overrides exist.
+      # Non-"all" runs that need a tree ("marked", or "none" with overrides)
+      # fetch it payload-free (includeOutputs=false) and pull each mocked span's
+      # output lazily by externalSpanId, never dragging down unused outputs.
       overrides_present = !mock_overrides.nil? && !mock_overrides.empty?
       include_outputs = mock_strategy == "all"
       mock_tree = nil
@@ -864,14 +898,10 @@ module Bitfab
           mock_tree = build_mock_tree(tree["root"] || {})
         rescue Exception => e # rubocop:disable Lint/RescueException
           raise if e.is_a?(SystemExit) || e.is_a?(SignalException)
-          # "all" and overrides both depend on the tree: "all" mocks every span
-          # from it, and overrides gate on its call-counter machinery. If the
-          # fetch fails, surface the error on the item rather than silently
-          # running everything real with the overrides dropped. Only bare
-          # "marked" can fall back to real execution (its marked spans just
-          # re-run). Mirrors the Python SDK's has_overrides re-raise.
-          raise if mock_strategy == "all" || overrides_present
-          mock_tree = nil
+          raise unless mock_strategy == "marked" && !overrides_present
+          # Keep an active empty tree so the root can still run, while any
+          # marked child fails closed at its call site.
+          mock_tree = {}
         end
       end
 
@@ -905,6 +935,7 @@ module Bitfab
         adapt_inputs:,
         adapt_ctx:,
         db_branch_lease: lease,
+        db_branch_timings:,
         source_bitfab_trace_id: original_trace_id,
         db_snapshot_ref:,
         http_client:
@@ -918,16 +949,20 @@ module Bitfab
         error: replay_item_error_message(e),
         trace_error: nil,
         replay_error: e,
-        duration_ms: metrics&.dig(:duration_ms),
+        duration_ms: nil,
+        original_duration_ms: metrics&.dig(:original_duration_ms),
+        original_tokens: metrics&.dig(:original_tokens),
+        original_model: metrics&.dig(:original_model),
         tokens: metrics&.dig(:tokens),
-        model: metrics&.dig(:model),
+        model: metrics&.dig(:original_model),
         trace_id: nil,
         original_trace_id:,
         original_span_id:,
         # Deprecated aliases for original_trace_id/original_span_id.
         source_trace_id: original_trace_id,
         source_span_id: original_span_id,
-        db_snapshot_ref:
+        db_snapshot_ref:,
+        db_branch_timings:
       }
     ensure
       release_db_branch_lease(http_client, lease) if lease
@@ -1041,11 +1076,15 @@ module Bitfab
     # the ORIGINAL trace's tokens); the replayed run's tokens are filled in by
     # run() from the complete-replay response once spans are aggregated
     # server-side, and stay nil here and on older servers.
+    # The original trace's measurements, falling back to the unprefixed keys a
+    # server that predates the rename sends. :tokens stays nil here: it is the
+    # REPLAYED run's, filled in from the complete-replay response.
     def extract_server_item_metrics(server_item)
       {
-        duration_ms: server_item["durationMs"],
-        tokens: nil,
-        model: server_item["model"]
+        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"],
+        tokens: nil
       }
     end
 
@@ -1126,13 +1165,15 @@ module Bitfab
     # Execute a single replay item: deserialize inputs, call method with replay context.
     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, source_bitfab_trace_id: nil,
-      db_snapshot_ref: nil, http_client: 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)
       args, kwargs = Serialize.deserialize_inputs(item)
 
       fn_result = nil
       fn_error = nil
       replay_error = nil
+      replay_duration_ms = nil
+      replay_started = 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
@@ -1154,6 +1195,7 @@ module Bitfab
           mock_overrides:,
           fetch_span_output:,
           db_branch_lease:,
+          db_branch_timings:,
           source_bitfab_trace_id:
         ) do
           begin
@@ -1172,17 +1214,26 @@ module Bitfab
           end
           if replay_error.nil?
             begin
+              replay_started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
               fn_result = if kwargs.empty?
                 receiver.send(method_name, *args)
               else
                 receiver.send(method_name, *args, **kwargs)
               end
+              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.
+              if replay_started
+                replay_duration_ms = ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - replay_started) * 1000).round
+              end
               fn_error = e
             end
           end
         end
       rescue => e
+        if replay_started
+          replay_duration_ms = ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - replay_started) * 1000).round
+        end
         replay_error = e
       end
 
@@ -1194,9 +1245,12 @@ module Bitfab
         error: item_error&.message,
         trace_error: fn_error,
         replay_error:,
-        duration_ms: metrics[:duration_ms],
+        duration_ms: replay_duration_ms,
+        original_duration_ms: metrics[:original_duration_ms],
+        original_tokens: metrics[:original_tokens],
+        original_model: metrics[:original_model],
         tokens: metrics[:tokens],
-        model: metrics[:model],
+        model: metrics[:original_model],
         # Written in by run() from the complete-replay response once the server
         # has minted this replay trace's row. Nil until then: the client-side
         # correlation id (below) is never surfaced as the public :trace_id.
@@ -1207,7 +1261,8 @@ module Bitfab
         # Deprecated aliases for original_trace_id/original_span_id.
         source_trace_id: source_bitfab_trace_id,
         source_span_id: input_source_span_id,
-        db_snapshot_ref:
+        db_snapshot_ref:,
+        db_branch_timings:
       }
     end
   end

data/lib/bitfab/span_context.rb

@@ -114,6 +114,7 @@ module Bitfab
   # Each entry is a Hash with :trace_id and :span_id keys.
   module SpanContext
     STACK_KEY = :__bitfab_span_stack
+    FIBER_BRIDGE_STACK_KEY = :__bitfab_span_fiber_bridge_stack
 
     module_function
 
@@ -122,7 +123,18 @@ module Bitfab
     end
 
     def current
-      stack.last
+      stack.last || fiber_bridge_stack.last
+    end
+
+    # Enumerator.new and enum_for run their source body in another Fiber on
+    # the same thread. Bridge only the active Enumerator parent across that
+    # boundary; ordinary fiber-local span stacks remain isolated.
+    def with_fiber_bridge(trace_id:, span_id:)
+      entry = {trace_id:, span_id:}
+      fiber_bridge_stack.push(entry)
+      yield
+    ensure
+      fiber_bridge_stack.pop
     end
 
     # Execute a block with a new span pushed onto the stack.
@@ -134,6 +146,15 @@ module Bitfab
     ensure
       stack.pop
     end
+
+    def fiber_bridge_stack
+      current_thread = Thread.current
+      current_thread.thread_variable_get(FIBER_BRIDGE_STACK_KEY) || begin
+        value = []
+        current_thread.thread_variable_set(FIBER_BRIDGE_STACK_KEY, value)
+        value
+      end
+    end
   end
 
   # Global storage for trace states (trace_id -> state hash)

data/lib/bitfab/version.rb

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

metadata

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