bitfab 0.23.8 → 0.29.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 342b5de3b40edd60f060465123f21d1c5338063340b34ac7bb7067c7baedfe55
-  data.tar.gz: 2f73bec442f3570efe78253d28733942e00b46c03a26ce2731a72e778cd6d79f
+  metadata.gz: 1f26df032729a0c1fbad20c807e1581ec1c68bb9019e0b8f32a85237c785a6ed
+  data.tar.gz: 6acf7a7e153565cf4fb272b651f2740b4d3cb5722b67b6508552aa1b2e81d042
 SHA512:
-  metadata.gz: c76023257bcdd959cb49d6881faf05e85013b2d30c02826fd57a2a2aa8a72d72d7f8f10716f9f4b2d850a84dfb4a6ef7151077b88f43e2e1e7bd5b4c81ed23ff
-  data.tar.gz: 37fb2e1215527c780253dd35e10e4ff01d735336d068d5d8ac16a75d4fcc17e1749f7718cf8ffc550e965a27a9f0ee8fb49b2b847724c10a1ae031f0f1470574
+  metadata.gz: 5fdcbb34b7174b5308d46405c95735d0a986db51f1af6ddbeb26613ebcdd9c35fade8a40553b1f88b78dc699261146d8cd9d4d7abffda560fbf06f40527717ce
+  data.tar.gz: ea2cd76d972b4b67f1554e1baf28b2aeeaf8560d1329eeff1f6df721273ce6f6e27193d20e0c0995ea77077544a5ab6f64cbb06908276eb8a27c6a6d590171b7

data/lib/bitfab/client.rb

@@ -21,6 +21,12 @@ module Bitfab
     # outputs (which may themselves be nil or false).
     MOCK_REPLAY_MISS = Object.new.freeze
 
+    # Sentinel for an OMITTED override value in register_mock_override, so a
+    # forgotten `value:` raises instead of silently injecting nil. An explicit
+    # `value: nil` is a legitimate injected value and is distinct from this.
+    VALUE_UNSET = Object.new.freeze
+    private_constant :VALUE_UNSET
+
     attr_reader :service_url
 
     def initialize(api_key: nil, service_url: nil, enabled: true, strict: false)
@@ -38,6 +44,10 @@ module Bitfab
       @http_client = HttpClient.new(api_key: -> { resolve_api_key }, service_url: @service_url)
       @pending_span_threads = {}
       @pending_span_mutex = Mutex.new
+      # Mock overrides registered via register_mock_override, applied to every
+      # replay on this client (after any per-call mock_override). Instance
+      # state, no global; clear_mock_overrides resets it.
+      @mock_overrides = []
     end
 
     # The configured API key (a proc is resolved on read). Reflects what was
@@ -79,7 +89,18 @@ module Bitfab
     # @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
-    #   { trace_id:, source_span_id: }, and returns [new_args, new_kwargs].
+    #   { original_trace_id:, original_span_id: } (with deprecated
+    #   source_trace_id/source_span_id aliases), and returns [new_args, new_kwargs].
+    # @param mock_override [Hash, Array<Hash>, nil] optional selective mock
+    #   override(s), each a { match:, value: } hash. match is a callable:
+    #   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
+    #   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 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
@@ -89,15 +110,84 @@ module Bitfab
     # @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: nil, code_change_files: nil, experiment_group_id: nil, dataset_id: nil, mock: "marked",
-      adapt_inputs: nil, environment: nil, on_progress: nil)
+      adapt_inputs: nil, mock_override: nil, environment: 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:, mock:, adapt_inputs:, environment:,
+        code_change_description:, code_change_files:, experiment_group_id:, dataset_id:, mock:, adapt_inputs:,
+        mock_override:, environment:,
         on_progress:
       )
     end
 
+    # Register a mock override applied to every subsequent replay on this
+    # client, so downstream real code runs against a value you supply for the
+    # matched span. Instance-scoped (no global state); call
+    # clear_mock_overrides to reset. Per-call replay(mock_override:) overrides
+    # take precedence, and both take precedence over the base mock strategy.
+    #
+    # Accepts either the (match, value) positional form or the keyword form
+    # (a { match:, value: } hash is also accepted). `match` must be callable;
+    # `value` is either a flat value injected directly or a callable invoked
+    # with the ctx hash.
+    #
+    # @example keyword form (primary), flat value
+    #   client.register_mock_override(
+    #     match: ->(node) { node[:trace_function_key] == "classify-intent" },
+    #     value: {label: "refund"}
+    #   )
+    #
+    # @example positional form (equivalent), callable value
+    #   client.register_mock_override(
+    #     ->(node) { node[:trace_function_key] == "classify-intent" },
+    #     ->(ctx) { {label: "refund", inputs: ctx[:inputs]} }
+    #   )
+    #
+    # @param positional [Array] either (match, value) or a single
+    #   { match:, value: } hash
+    # @param match [#call, nil] keyword form matcher
+    # @param value [Object, #call, nil] keyword form injected value or producer
+    # @return [void]
+    def register_mock_override(*positional, match: nil, value: VALUE_UNSET)
+      if match.nil? && value.equal?(VALUE_UNSET)
+        if positional.length == 1 && positional[0].is_a?(Hash)
+          override = positional[0]
+          match = override[:match]
+          # Distinguish an omitted :value from an explicit `value: nil`: the
+          # key's presence, not a nil read, is what marks it provided.
+          value = override.key?(:value) ? override[:value] : VALUE_UNSET
+        elsif positional.length == 2
+          match, value = positional
+        end
+      end
+
+      unless match.respond_to?(:call)
+        raise ArgumentError,
+          "register_mock_override requires a callable match. Pass (match, value) " \
+          "positionally, as keywords (match:, value:), or as a { match:, value: } hash. " \
+          "value may be a flat value or a callable."
+      end
+      # A forgotten value must not silently inject nil. An explicit nil is a
+      # legitimate injected value and passes this guard.
+      if value.equal?(VALUE_UNSET)
+        raise ArgumentError,
+          "register_mock_override requires a value (the second argument, or " \
+          "value:). It may be a flat value or a callable; pass value: nil " \
+          "explicitly to inject nil."
+      end
+
+      @mock_overrides << {match:, value:}
+      nil
+    end
+
+    # Remove all overrides registered via register_mock_override.
+    #
+    # @return [void]
+    def clear_mock_overrides
+      @mock_overrides.clear
+      nil
+    end
+
     # Get a function wrapper bound to a specific trace function key.
     #
     # This provides a fluent API for binding a trace_function_key once and
@@ -192,7 +282,8 @@ module Bitfab
         call_index = advance_mock_counter(replay_ctx, trace_function_key, span_name, is_root_span:)
         if call_index
           mocked_output = check_mock_replay(
-            replay_ctx, trace_function_key, span_name, call_index, mock_on_replay:
+            replay_ctx, trace_function_key, span_name, call_index,
+            span_type:, args:, kwargs:, mock_on_replay:
           )
           if mocked_output != MOCK_REPLAY_MISS
             send_mocked_span(
@@ -285,6 +376,9 @@ module Bitfab
               {
                 neon_branch_id: lease["neonBranchId"],
                 snapshot_timestamp: lease["snapshotTimestamp"],
+                original_trace_id: replay_ctx[:source_bitfab_trace_id],
+                # 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
               }
@@ -452,10 +546,11 @@ module Bitfab
     # db_snapshot_usage: replay DB branch usage record, present only when a
     # lease was attached to the replay item. Serialized as `db_snapshot_usage`
     # on the raw trace so the server can stamp the trace's metadata at ingest:
-    # { neon_branch_id:, snapshot_timestamp: (optional), source_trace_id:
-    # (optional), accessed: } with :accessed true if customer code obtained
-    # the branch URL and false if it ignored it. nil outside replay or when no
-    # lease was attached, in which case the key is omitted entirely.
+    # { neon_branch_id:, snapshot_timestamp: (optional), original_trace_id:
+    # (optional, with its deprecated source_trace_id alias), accessed: } with
+    # :accessed true if customer code obtained the branch URL and false if it
+    # ignored it. nil outside replay or when no lease was attached, in which
+    # case the key is omitted entirely.
     def send_trace_completion(trace_function_key:, trace_id:, started_at:, ended_at:, db_snapshot_usage: nil)
       trace_state = TraceState.get(trace_id)
       trace_started_at = trace_state&.dig(:started_at) || started_at
@@ -483,8 +578,11 @@ module Bitfab
         if db_snapshot_usage[:snapshot_timestamp]
           usage["snapshot_timestamp"] = db_snapshot_usage[:snapshot_timestamp]
         end
-        if db_snapshot_usage[:source_trace_id]
-          usage["source_trace_id"] = db_snapshot_usage[:source_trace_id]
+        if db_snapshot_usage[:original_trace_id]
+          usage["original_trace_id"] = db_snapshot_usage[:original_trace_id]
+          # Deprecated wire alias, kept so this SDK still reports usage
+          # against servers that predate the rename.
+          usage["source_trace_id"] = db_snapshot_usage[:original_trace_id]
         end
         usage["accessed"] = db_snapshot_usage[:accessed]
         raw_trace["db_snapshot_usage"] = usage
@@ -617,10 +715,53 @@ module Bitfab
       call_index
     end
 
-    # Decide whether this child span should be short-circuited to its recorded
-    # output. Returns MOCK_REPLAY_MISS when the span should run real code,
-    # otherwise returns the deserialized historical output.
-    def check_mock_replay(replay_ctx, trace_function_key, span_name, call_index, mock_on_replay:)
+    # Decide whether this child span should be short-circuited during replay.
+    # Returns MOCK_REPLAY_MISS when the span should run real code, otherwise
+    # returns the value to inject (a selective override's produced value, or the
+    # deserialized historical output under the base "all"/"marked" strategy).
+    #
+    # Precedence: a matching mock override wins over the base strategy (per-call
+    # overrides already precede registered ones in the resolved list). A span no
+    # override matches falls through to the marked/all logic unchanged. An
+    # override that matches ALWAYS injects, even when it returns nil, so a
+    # legitimately-nil injected value is never mistaken for a miss.
+    def check_mock_replay(replay_ctx, trace_function_key, span_name, call_index, span_type:, args:, kwargs:, mock_on_replay:)
+      mock_entry = replay_ctx[:mock_tree]["#{trace_function_key}:#{span_name}:#{call_index}"]
+
+      # 1) Selective overrides. The matcher runs on structural metadata only;
+      # the recorded output is fetched (lazily, by externalSpanId) only if the
+      # processor actually asks for it via get_original_output, so a pure
+      # synthetic override triggers zero output fetches.
+      overrides = replay_ctx[:mock_overrides]
+      if overrides&.any?
+        node = {
+          trace_function_key:,
+          span_name:,
+          type: span_type,
+          original_span_id: mock_entry && mock_entry[:source_span_id]
+        }
+        override = overrides.find { |o| o[:match].call(node) }
+        if override
+          value = override[:value]
+          # value is EITHER a flat value (injected directly) OR a callable
+          # invoked with the ctx hash. The result IS the span's output (full
+          # replacement). nil is a legitimate injected value: because a matching
+          # override always injects, execute_span short-circuits even when this
+          # is nil (nil != MOCK_REPLAY_MISS). To inject a proc as the literal
+          # output, wrap it in a callable value.
+          return value unless value.respond_to?(:call)
+
+          get_original_output = lambda do
+            unless mock_entry
+              raise "No recorded span to source output for '#{trace_function_key}'."
+            end
+            resolve_recorded_output(mock_entry, replay_ctx)
+          end
+          return value.call({node:, inputs: args, kwargs:, get_original_output:})
+        end
+      end
+
+      # 2) Base strategy: replay recorded output for eligible spans.
       strategy = replay_ctx[:mock_strategy]
       case strategy
       when "marked"
@@ -631,25 +772,40 @@ module Bitfab
         return MOCK_REPLAY_MISS
       end
 
-      mock_entry = replay_ctx[:mock_tree]["#{trace_function_key}:#{span_name}:#{call_index}"]
       return MOCK_REPLAY_MISS unless mock_entry
 
-      output = mock_entry[:output]
-      output_meta = mock_entry[:output_meta]
+      resolve_recorded_output(mock_entry, replay_ctx)
+    end
 
-      # Type-preserving deserialization when the server included Ruby-side
-      # Marshal+Base64 metadata. Falls back to the JSON output silently: the
-      # spanTree endpoint currently returns superjson/jsonpickle-shaped meta,
-      # which Ruby cannot reconstruct.
-      if output_meta.is_a?(String) && !output_meta.empty?
-        begin
-          output = Serialize.unmarshal_value(output_meta)
-        rescue
-          # Fall through to the JSON output
-        end
+    # Resolve a mock-tree entry's recorded output. Prefers an inline output when
+    # the tree carried one (eager "all", or an older server that ignores
+    # includeOutputs and returns outputs); otherwise lazily fetches it by
+    # externalSpanId through the per-item memoized fetcher on the replay
+    # context, so only the spans actually mocked pull their output.
+    #
+    # Deserialization is type-preserving when Ruby-side Marshal+Base64 metadata
+    # is present, falling back to the raw JSON output when the meta is a shape
+    # Ruby cannot reconstruct (e.g. superjson/jsonpickle from another SDK).
+    def resolve_recorded_output(mock_entry, replay_ctx)
+      # Presence of the :output (or :output_meta) key means the tree carried the
+      # output inline - including a legitimately nil inline output. A
+      # payload-free entry has neither key and is fetched lazily below.
+      if mock_entry.key?(:output) || mock_entry.key?(:output_meta)
+        return Serialize.deserialize_output(mock_entry[:output], mock_entry[:output_meta])
+      end
+
+      fetcher = replay_ctx[:fetch_span_output]
+      external_span_id = mock_entry[:external_span_id]
+      if fetcher && external_span_id
+        return fetcher.call(external_span_id)
       end
 
-      output
+      # No inline output and nothing to lazily fetch. This is a span that never
+      # recorded an output (e.g. an errored or incomplete child): inject nil
+      # rather than raising, so it does not fail the whole replay item. Mirrors
+      # the Python and TypeScript SDKs, which return the (absent) recorded output
+      # as null here.
+      nil
     end
 
     # Record a span entry for a mocked invocation so the test run reflects the

data/lib/bitfab/http_client.rb

@@ -80,6 +80,7 @@ module Bitfab
       http.open_timeout = request_timeout
       http.read_timeout = request_timeout
 
+      # request_uri (not path) so any query string on the endpoint survives.
       req = Net::HTTP::Get.new(uri.request_uri, headers)
       response = http.request(req)
 
@@ -158,10 +159,16 @@ module Bitfab
     # matched against their historical outputs.
     #
     # Returns a hash shaped { "root" => SpanTreeNode } where each node has
-    # sourceSpanId, traceFunctionKey, spanName, type, output, optional
-    # outputMeta, and children.
-    def get_span_tree(external_span_id)
-      get("/api/sdk/replay/spanTree/#{external_span_id}", timeout: 30)
+    # sourceSpanId, externalSpanId, traceFunctionKey, spanName, type, and
+    # children. When include_outputs is true each node also carries its recorded
+    # output (and optional outputMeta); when false the server omits the output
+    # payloads so only the spans actually mocked are fetched later (by
+    # externalSpanId), avoiding dragging down every span's output when only a
+    # few are mocked.
+    def get_span_tree(external_span_id, include_outputs: true)
+      endpoint = "/api/sdk/replay/spanTree/#{external_span_id}"
+      endpoint += "?includeOutputs=false" unless include_outputs
+      get(endpoint, timeout: 30)
     end
 
     # Mark a replay test run as completed. Blocking call.

data/lib/bitfab/mock_override.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module Bitfab
+  # Selective mock overrides for replay. Mirrors the TypeScript SDK's
+  # mockOverride feature.
+  #
+  # A mock override injects a custom value into a specific span (node) during
+  # replay: the matched span short-circuits its real execution and returns the
+  # value the processor produces, so downstream real code runs against the
+  # substituted output. This is a third mock mode alongside "run real code" and
+  # "replay recorded output" (see MOCK_STRATEGIES in replay.rb).
+  #
+  # An override is a (match, value) pair, represented as a symbol-keyed hash
+  # { match:, value: }:
+  #
+  # - match is a callable (proc/lambda): match.call(node) -> boolean. node is a
+  #   structural-metadata hash { trace_function_key:, span_name:, type:,
+  #   original_span_id: } (symbol keys; original_span_id may be nil when the
+  #   live span has no recorded counterpart). Matching must not depend on the
+  #   recorded output.
+  #
+  # - value is EITHER a flat value OR a callable. A flat value is injected
+  #   directly (e.g. value: { label: "refund" }). A callable is invoked with
+  #   the ctx hash: value.call(ctx) where ctx is
+  #   { node:, inputs:, kwargs:, get_original_output: } - inputs is the live
+  #   replay positional args array, kwargs is the live keyword args hash (empty
+  #   when the call used none), and get_original_output is a proc returning this
+  #   span's recorded output (deserialized; raises if there is no recorded
+  #   counterpart). Either
+  #   way the resulting value IS the span's output (full replacement, no merge);
+  #   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.
+  module MockOverride
+    module_function
+
+    # Normalize the per-call `mock_override` option (a single override hash, an
+    # array of them, or nil) into an array. First match wins downstream, so
+    # order is preserved.
+    #
+    # Each override is validated: it must be a hash with a callable `:match` and
+    # must include a `:value` key. A forgotten `:value` raises rather than
+    # silently short-circuiting matched spans with nil (an explicit `value: nil`
+    # is a legitimate injected value and passes). Mirrors the guard on
+    # `register_mock_override`, and the Python/TypeScript SDKs, where the per-call
+    # override object requires a value.
+    #
+    # @param mock_override [Hash, Array<Hash>, nil]
+    # @return [Array<Hash>]
+    def normalize(mock_override)
+      return [] if mock_override.nil?
+
+      overrides = mock_override.is_a?(Array) ? mock_override : [mock_override]
+      overrides.each do |override|
+        unless override.is_a?(Hash) && override[:match].respond_to?(:call)
+          raise ArgumentError,
+            "mock_override requires a callable :match. Pass a " \
+            "{ match:, value: } hash (or an array of them). value may be a " \
+            "flat value or a callable."
+        end
+        unless override.key?(:value)
+          raise ArgumentError,
+            "mock_override requires a :value (a flat value or a callable); " \
+            "pass value: nil explicitly to inject nil."
+        end
+      end
+      overrides
+    end
+  end
+end

data/lib/bitfab/replay.rb

@@ -4,6 +4,7 @@ require "fileutils"
 require "json"
 
 require_relative "constants"
+require_relative "mock_override"
 require_relative "serialize"
 require_relative "traceable"
 
@@ -31,8 +32,8 @@ module Bitfab
     # threads (span uploads + trace completion) so the replay runner can join
     # them before complete_replay builds the trace-ID mapping.
     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, pending_persistence: nil, db_branch_lease: nil,
-      source_bitfab_trace_id: nil)
+      mock_tree: nil, mock_strategy: nil, mock_overrides: nil, fetch_span_output: nil, pending_persistence: nil,
+      db_branch_lease: nil, source_bitfab_trace_id: nil)
       previous = Thread.current[REPLAY_CONTEXT_KEY]
       ctx = {
         test_run_id:,
@@ -45,6 +46,14 @@ module Bitfab
         ctx[:mock_tree] = mock_tree
         ctx[:mock_strategy] = mock_strategy || "marked"
         ctx[:call_counters] = {}
+        # Selective mock overrides ride alongside the tree: they gate on the
+        # same non-root call-counter machinery, so they only fire when a tree
+        # 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.
+        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
       # it belongs to ride on the context so ReplayEnvironment can read them
@@ -101,14 +110,19 @@ module Bitfab
     # @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
-    #   { trace_id:, source_span_id: }, and returns [new_args, new_kwargs]. Runs
-    #   per item inside the same rescue as the method, so a raising adapter sets
-    #   that item's :error rather than crashing the run.
+    #   { original_trace_id:, original_span_id: } (with deprecated
+    #   source_trace_id/source_span_id aliases), and returns [new_args, new_kwargs].
+    #   Runs per item inside the same rescue as the method, so a raising adapter
+    #   sets that item's :error rather than crashing the run.
     # @param on_progress [#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:, error:, duration_ms: } for
-    #   the single item that just settled. trace_id is the SOURCE (historical)
-    #   trace that was replayed, error is that item's replay error or nil, and
+    #   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).
+    #   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
     #   render replay progress (e.g. a per-trace log). A raising callback never
     #   crashes the run.
@@ -116,7 +130,7 @@ module Bitfab
     def run(client, receiver, method_name, trace_function_key:, limit: nil, trace_ids: nil, name: nil,
       max_concurrency: 10, code_change_description: nil, code_change_files: nil, experiment_group_id: nil,
       dataset_id: nil, mock: "marked",
-      adapt_inputs: nil, environment: nil, on_progress: nil)
+      adapt_inputs: nil, mock_override: nil, environment: 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
@@ -148,6 +162,13 @@ module Bitfab
 
       http_client = client.instance_variable_get(:@http_client)
 
+      # Resolved override list: per-call overrides FIRST (they win), then the
+      # client's registered overrides. First matching override wins downstream,
+      # so this ordering makes per-call take precedence over registered, and
+      # both take precedence over the base mock strategy.
+      registered_overrides = client.instance_variable_get(:@mock_overrides) || []
+      resolved_overrides = MockOverride.normalize(mock_override) + registered_overrides
+
       # limit is meaningless with explicit trace_ids (the ID list determines
       # the count), so it's omitted from the request entirely.
       effective_limit = trace_ids ? nil : (limit || 5)
@@ -171,17 +192,16 @@ module Bitfab
 
       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, on_progress:)
+          adapt_inputs, include_db_branch_lease, on_progress:, mock_overrides: resolved_overrides)
       else
         []
       end
 
       # Every item joined its own trace-persistence threads (span uploads +
       # completion) in execute_item, so all replay traces are on the server
-      # by now: no flush needed, and complete_replay's trace-ID mapping is
-      # deterministic. complete_replay failures propagate: a missing mapping
-      # means verdicts can't be persisted, which callers must hear about
-      # loudly.
+      # by now: no flush needed. 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)
       trace_id_map = complete_response&.dig("traceIds")
       # Per-replay-trace token usage keyed by server trace id: the REPLAYED
@@ -189,66 +209,85 @@ module Bitfab
       # item's :tokens.
       replay_tokens = complete_response&.dig("tokens") || {}
 
-      if trace_id_map.nil?
-        # Older servers don't return the mapping. Preserve the legacy
-        # nil-trace_id behavior but say why.
-        warn "Bitfab: server did not return replay trace IDs; item trace_id " \
-          "will be nil (server upgrade required for verdict persistence)"
-        result_items.each { |item| item[:trace_id] = nil }
-      else
-        # Map each item's locally-generated trace ID to the server's trace
-        # row ID. A completed item with no mapping means its trace was sent
-        # but the server has no record: a nil trace_id blocks verdict
-        # persistence and the Studio experiments view downstream, so this
-        # must never be silent.
-        #
-        # Severity splits on scope:
-        # - ALL completed items missing: systemic (the replayed method is
-        #   not traced, or uploads are wholesale broken). Raise; the run's
-        #   results are unusable for persistence.
-        # - SOME completed items missing: per-item upload failure (transient
-        #   network blip, one oversized payload). Nil those items and warn
-        #   loudly, but return the run so callers can persist verdicts for
-        #   the items that landed.
+      # trace_id_map maps each item's client-side correlation id (:_sdk_trace_id,
+      # which tagged that item's spans during the run) to the server's trace row
+      # id. We use it to write the real server replay trace id into item[:trace_id]
+      # now that the row exists, attach each item's server-aggregated token usage,
+      # and detect a systemic upload failure early. Verdict persistence does NOT
+      # use this map: it is keyed by the original-trace lineage (:original_trace_id +
+      # test_run_id), which needs no client-held server id. Older servers that omit
+      # the map yield no replay tokens and leave item[:trace_id] nil.
+      unless trace_id_map.nil?
         missing = []
         completed_count = 0
         result_items.each do |item|
-          next unless item[:trace_id]
-
-          mapped = trace_id_map[item[:trace_id]]
+          local_id = item[:_sdk_trace_id]
+          mapped = local_id && trace_id_map[local_id]
+          # Write the real server replay trace id in as it comes back; the item
+          # held nil until now (the client correlation id is never surfaced).
+          item[:trace_id] = mapped
           if item[:error].nil?
             completed_count += 1
-            missing << item[:trace_id] if mapped.nil?
+            missing << local_id if mapped.nil?
           end
-          # Pull this item's replayed-run tokens by its server trace id, before
-          # :trace_id is overwritten with that id below.
           item[:tokens] = normalize_tokens(replay_tokens[mapped]) if mapped
-          item[:trace_id] = mapped
         end
-        if missing.any?
+        # ALL completed items missing: systemic (the replayed method is not
+        # traced, or uploads are wholesale broken). Raise; the run's traces
+        # never persisted, so nothing can be labeled.
+        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."
-          if missing.length == completed_count
-            raise "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 " \
-              "joined, so either the uploads failed or the replayed method is " \
-              "not traced (no root span was emitted)."
-          end
+          raise "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 " \
+            "joined, so either the uploads failed or the replayed method is " \
+            "not traced (no root span was emitted)."
+        end
+        # SOME completed items missing: per-item upload failure. Warn, but
+        # return the run; the items that landed can still be labeled.
+        if missing.any?
           warn "Bitfab: server has no persisted trace for #{missing.length} of " \
             "#{completed_count} completed replay item(s) " \
-            "(test_run_id #{test_run_id}).#{server_count} Their trace_id is nil " \
-            "and verdicts cannot be persisted for them. Missing: #{missing.join(", ")}"
+            "(test_run_id #{test_run_id}). Their replay token usage is " \
+            "unavailable and they cannot be labeled."
         end
       end
 
-      replay_result = {
+      # Strip the internal correlation handle so returned items expose only the
+      # public shape (runs against older servers that omit the map too).
+      result_items.each { |item| item.delete(:_sdk_trace_id) }
+
+      result = {
         items: result_items,
         test_run_id:,
         test_run_url: "#{client.service_url}#{test_run_url}"
       }
-      write_replay_result_file(replay_result)
-      replay_result
+      # 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.
+      write_replay_result_file(result)
+      if on_progress
+        errored = result_items.count { |item| !item[:error].nil? }
+        total = result_items.length
+        begin
+          on_progress.call({
+            type: "complete",
+            test_run_id:,
+            completed: total,
+            total:,
+            succeeded: total - errored,
+            errored:,
+            result:
+          })
+        rescue => e
+          warn "Bitfab: replay on_progress callback raised: #{e.message}"
+        end
+      end
+      result
     end
 
     def write_replay_result_file(result)
@@ -266,7 +305,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, on_progress: nil)
+      adapt_inputs = nil, include_db_branch_lease = false, on_progress: nil, mock_overrides: [])
       concurrency = max_concurrency || server_items.length
 
       # Reports running totals once per item as it settles. In the parallel
@@ -280,11 +319,13 @@ module Bitfab
       succeeded = 0
       errored = 0
       # Each event carries the single item that just settled so a progress UI
-      # can render per-trace pass/fail as the run streams. trace_id is the
-      # SOURCE (historical) trace that was replayed, taken from the server
-      # item: the result's own :trace_id at this stage is the new replay
-      # trace id (assigned in run() after complete_replay), not the source.
-      report = lambda do |result, source_trace_id, test_run_id|
+      # 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
 
         progress_mutex.synchronize do
@@ -295,8 +336,12 @@ module Bitfab
             on_progress.call({
               test_run_id:, completed:, total:, succeeded:, errored:,
               item: {
-                trace_id: source_trace_id,
-                replay_trace_id: result[:trace_id],
+                trace_id: result[:trace_id],
+                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,
                 input: result[:input],
                 result: result[:result],
                 original_output: result[:original_output],
@@ -316,8 +361,8 @@ module Bitfab
       if concurrency <= 1
         server_items.map do |item|
           result = process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy,
-            adapt_inputs, include_db_branch_lease)
-          report.call(result, item["traceId"], test_run_id)
+            adapt_inputs, include_db_branch_lease, mock_overrides:)
+          report.call(result, original_trace_id_of(item), original_span_id_of(item), test_run_id)
           result
         end
       else
@@ -333,9 +378,9 @@ module Bitfab
               break unless item
 
               result = process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy,
-                adapt_inputs, include_db_branch_lease)
+                adapt_inputs, include_db_branch_lease, mock_overrides:)
               results_mutex.synchronize { results[idx] = result }
-              report.call(result, item["traceId"], test_run_id)
+              report.call(result, original_trace_id_of(item), original_span_id_of(item), test_run_id)
             end
           end
         end
@@ -352,8 +397,13 @@ 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)
+      adapt_inputs = nil, include_db_branch_lease = false, mock_overrides: [])
       metrics = extract_server_item_metrics(server_item)
+      # 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).
+      original_trace_id = original_trace_id_of(server_item)
+      original_span_id = original_span_id_of(server_item)
       # The server resolves a Neon preview branch per item during /replay/start
       # (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
@@ -362,22 +412,53 @@ module Bitfab
       # lease (env.active? is false for those).
       lease = include_db_branch_lease ? server_item["dbBranchLease"] : nil
 
-      span = http_client.get_external_span(server_item["externalSpanId"])
+      span = http_client.get_external_span(original_span_id)
       item_data = extract_span_data(span)
 
+      # Fetch the span tree when the base strategy needs recorded outputs
+      # ("all"/"marked") OR when selective mock overrides are present. Overrides
+      # gate on the same non-root call-counter machinery the tree drives, so
+      # 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.
+      overrides_present = !mock_overrides.nil? && !mock_overrides.empty?
+      include_outputs = mock_strategy == "all"
       mock_tree = nil
-      if mock_strategy == "all" || mock_strategy == "marked"
+      if mock_strategy == "all" || mock_strategy == "marked" || overrides_present
         begin
-          tree = http_client.get_span_tree(server_item["externalSpanId"])
+          tree = http_client.get_span_tree(original_span_id, include_outputs:)
           mock_tree = build_mock_tree(tree["root"] || {})
         rescue Exception => e # rubocop:disable Lint/RescueException
           raise if e.is_a?(SystemExit) || e.is_a?(SignalException)
-          raise if mock_strategy == "all"
+          # "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
         end
       end
 
-      adapt_ctx = {trace_id: server_item["traceId"], source_span_id: server_item["externalSpanId"]}
+      # Memoized per-item lazy fetcher, present only on the payload-free path
+      # (not eager "all", whose outputs are already inline). Passed to the
+      # replay context so resolve_recorded_output can pull a mocked span's
+      # output on demand, once per span even if read twice (base mock + an
+      # override's get_original_output).
+      fetch_span_output = (mock_tree && !include_outputs) ? build_span_output_fetcher(http_client) : nil
+
+      adapt_ctx = {
+        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
+      }
 
       execute_item(
         item_data,
@@ -389,14 +470,16 @@ module Bitfab
         input_source_trace_id: span["externalTraceId"],
         mock_strategy:,
         mock_tree:,
+        mock_overrides:,
+        fetch_span_output:,
         adapt_inputs:,
         adapt_ctx:,
         db_branch_lease: lease,
-        source_bitfab_trace_id: server_item["traceId"],
+        source_bitfab_trace_id: original_trace_id,
         db_snapshot_ref: server_item["dbSnapshotRef"]
       )
     rescue => e
-      warn "Bitfab: replay item for span #{server_item["externalSpanId"]} failed before execution: #{e.message}"
+      warn "Bitfab: replay item for span #{original_span_id} failed before execution: #{e.message}"
       {
         input: [],
         result: nil,
@@ -406,12 +489,28 @@ module Bitfab
         tokens: metrics&.dig(:tokens),
         model: metrics&.dig(: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: server_item["dbSnapshotRef"]
       }
     ensure
       release_db_branch_lease(http_client, lease) if lease
     end
 
+    # The ORIGINAL (historical) trace/span an item replays. Canonical server
+    # keys are originalTraceId/originalSpanId; older servers send them under
+    # the deprecated sourceTraceId/sourceSpanId aliases.
+    def original_trace_id_of(server_item)
+      server_item["originalTraceId"] || server_item["sourceTraceId"]
+    end
+
+    def original_span_id_of(server_item)
+      server_item["originalSpanId"] || server_item["sourceSpanId"]
+    end
+
     # Delete the per-item Neon preview branch. Best-effort: a failure is warned
     # but never raised: the server-side TTL janitor reaps orphans.
     def release_db_branch_lease(http_client, lease)
@@ -448,11 +547,18 @@ module Bitfab
           counter_key = "#{key}:#{name}"
           index = counters[counter_key] || 0
           counters[counter_key] = index + 1
-          spans["#{counter_key}:#{index}"] = {
+          # externalSpanId is what the lazy path fetches the recorded output by
+          # when the tree came back payload-free (includeOutputs=false). output
+          # / outputMeta are copied ONLY when the node actually carries them, so
+          # the presence of the :output key distinguishes an inline nil output
+          # from a payload-free entry (mirrors TS's undefined-vs-null check).
+          entry = {
             source_span_id: node["sourceSpanId"],
-            output: node["output"],
-            output_meta: node["outputMeta"]
+            external_span_id: node["externalSpanId"]
           }
+          entry[:output] = node["output"] if node.key?("output")
+          entry[:output_meta] = node["outputMeta"] if node.key?("outputMeta")
+          spans["#{counter_key}:#{index}"] = entry
         end
         (node["children"] || []).each { |child| walk.call(child) }
       end
@@ -462,6 +568,27 @@ module Bitfab
       spans
     end
 
+    # Build a memoized lazy fetcher for a single replay item. The returned
+    # lambda takes an externalSpanId and returns that span's deserialized
+    # recorded output, fetching it via get_external_span at most once per id
+    # (a span read twice - as a base "marked" mock and via an override's
+    # get_original_output - fetches once). The cache is per item: each item
+    # replays synchronously on its own thread, so no lock is needed.
+    def build_span_output_fetcher(http_client)
+      cache = {}
+      lambda do |external_span_id|
+        return cache[external_span_id] if cache.key?(external_span_id)
+
+        span = http_client.get_external_span(external_span_id)
+        span_data = (span["rawData"] || {})["span_data"] || {}
+        # Prefer the Ruby Marshal payload (output_serialized) written by this
+        # SDK; fall back to another SDK's output_meta, then the raw JSON output.
+        meta = span_data["output_serialized"]
+        meta = span_data["output_meta"] if meta.nil? || meta == ""
+        cache[external_span_id] = Serialize.deserialize_output(span_data["output"], meta)
+      end
+    end
+
     # Extract input/output data from an external span's rawData.
     def extract_span_data(span)
       raw_data = span["rawData"] || {}
@@ -505,12 +632,17 @@ 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, adapt_inputs: nil, adapt_ctx: nil,
-      db_branch_lease: nil, source_bitfab_trace_id: nil, db_snapshot_ref: nil)
+      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)
       args, kwargs = Serialize.deserialize_inputs(item)
 
       fn_result = nil
       fn_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
       # Collects the root span's persistence threads (span uploads + trace
       # completion). Joined below so this item's trace is on the server
@@ -525,6 +657,8 @@ module Bitfab
         trace_id: sdk_trace_id,
         mock_tree:,
         mock_strategy:,
+        mock_overrides:,
+        fetch_span_output:,
         pending_persistence:,
         db_branch_lease:,
         source_bitfab_trace_id:
@@ -533,7 +667,13 @@ module Bitfab
         # 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 || {trace_id: nil, source_span_id: input_source_span_id}
+          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?
@@ -559,7 +699,16 @@ module Bitfab
         duration_ms: metrics[:duration_ms],
         tokens: metrics[:tokens],
         model: metrics[:model],
-        trace_id: sdk_trace_id,
+        # 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.
+        trace_id: nil,
+        _sdk_trace_id: sdk_trace_id,
+        original_trace_id: source_bitfab_trace_id,
+        original_span_id: input_source_span_id,
+        # 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:
       }
     end

data/lib/bitfab/serialize.rb

@@ -182,6 +182,27 @@ module Bitfab
       Marshal.load(Base64.strict_decode64(encoded)) # rubocop:disable Security/MarshalLoad
     end
 
+    # Deserialize a recorded span output. Prefers the Ruby-side Marshal+Base64
+    # `meta` when present (type-preserving); falls back to the raw JSON output
+    # when meta is absent or a shape Ruby cannot reconstruct (e.g. the
+    # superjson/jsonpickle meta another SDK's span may carry). Shared by the
+    # inline mock-tree path and the lazy per-span fetch so both deserialize
+    # identically.
+    #
+    # @param raw_output [Object] the human-readable JSON output
+    # @param meta [String, nil] the Marshal+Base64 metadata, if any
+    # @return [Object] the reconstructed output, or raw_output on fallback
+    def deserialize_output(raw_output, meta)
+      if meta.is_a?(String) && !meta.empty?
+        begin
+          return unmarshal_value(meta)
+        rescue
+          # Fall through to the raw JSON output
+        end
+      end
+      raw_output
+    end
+
     # Deserialize replay inputs from a span's data into [args, kwargs].
     #
     # Prefers Marshal-serialized `inputSerialized` for type preservation,

data/lib/bitfab/version.rb

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

data/lib/bitfab.rb

@@ -9,6 +9,7 @@ require_relative "bitfab/serialize"
 require_relative "bitfab/db_snapshot"
 require_relative "bitfab/span_context"
 require_relative "bitfab/http_client"
+require_relative "bitfab/mock_override"
 require_relative "bitfab/replay"
 require_relative "bitfab/replay_environment"
 require_relative "bitfab/client"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bitfab
 version: !ruby/object:Gem::Version
-  version: 0.23.8
+  version: 0.29.1
 platform: ruby
 authors:
 - Harvest Team
@@ -121,6 +121,7 @@ files:
 - lib/bitfab/constants.rb
 - lib/bitfab/db_snapshot.rb
 - lib/bitfab/http_client.rb
+- lib/bitfab/mock_override.rb
 - lib/bitfab/replay.rb
 - lib/bitfab/replay_environment.rb
 - lib/bitfab/serialize.rb