bitfab 0.10.5 → 0.12.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b789df310c91a7e9130f3bdb6a35e6237a8ddadede4df49a5f7dc4afc0525e91
-  data.tar.gz: 0cd79b9c9bc8bf15ef926625c0aa1aba8f3cf1798062c0d077e371d98b9f6aba
+  metadata.gz: 0e4ca1b81e502de48fe20835b4d072b05736ac5f4ff1a444d683627818648675
+  data.tar.gz: d7120ea5b06e8da151cd9ca26c9549b017458cc2d3b1fa9078ffa512c59d8606
 SHA512:
-  metadata.gz: fbf62eeabc5741c36ef1b867730f548039f20331d6591ccae7a66cc48882cd4638877edc1b80d57c3a326827edfbedff1d3ec226ed3b4e3358f9041e3405ea5f
-  data.tar.gz: 579cb384585a42160d5f1c9941cbf0762f53d35ffda7f109cb40c4c52a291a02d446e113ca664a2c84a08916481fe86912eb386a58956fb5651ff1a6dc108771
+  metadata.gz: b2cbfda78dfd7d726174ace8910f507ea3982d68bb19c402fcaa50b84a20eeddbaf7373932fbec81c93116d8056b41eef876f132aa27913cade2937506299e37
+  data.tar.gz: 41fa08a3eb1a36f21af3358570012625564473eeafd847ab0073c33e012fb341c2b42463a1247e73c929679a1075220c096a4e35a2ac5c7e81151268f3cbbd3c

data/README.md

@@ -276,6 +276,62 @@ client = ExternalHttpClient.new
 client.get("https://api.example.com")
 ```
 
+### Fluent API: `client.get_function`
+
+Bind a `trace_function_key` once and wrap multiple methods or classes against it. Mirrors `client.get_function` in the Python SDK and `client.getFunction` in TypeScript.
+
+```ruby
+fn = Bitfab.client.get_function("openai")
+
+fn.wrap(OpenAI::Client, :chat, name: "Chat", type: "llm")
+fn.wrap(OpenAI::Client, :embeddings, name: "Embed", type: "llm")
+```
+
+`#wrap` accepts the same options as `Bitfab::Traceable.wrap` (`name`, `type`, `mock_on_replay`), but the `trace_function_key` is fixed to the one bound on the `BitfabFunction`.
+
+### Replay with Mock Strategies
+
+Replay reruns historical traces through your code so you can compare outputs after an iteration. By default every child span runs real code — fine for offline traces, but expensive when children make paid LLM/API calls. Three strategies control whether child spans return their historical output instead of executing:
+
+```ruby
+# "none" (default): everything runs real code
+client.replay(pipeline, :process, trace_function_key: "my-fn", mock: "none")
+
+# "all": every child span returns its historical output
+client.replay(pipeline, :process, trace_function_key: "my-fn", mock: "all")
+
+# "marked": only spans tagged with `mock_on_replay: true` return historical output
+client.replay(pipeline, :process, trace_function_key: "my-fn", mock: "marked")
+```
+
+Tag the spans you want mocked at definition time:
+
+```ruby
+class Pipeline
+  include Bitfab::Traceable
+  bitfab_function "my-fn"
+
+  # mock_on_replay: true → returns historical output under mock: "marked"
+  bitfab_span :call_llm, type: "llm", mock_on_replay: true
+  def call_llm(prompt)
+    # paid OpenAI call — skip during replay
+  end
+
+  bitfab_span :transform, type: "function"
+  def transform(text)
+    # cheap, deterministic — keep running real
+  end
+
+  bitfab_span :process, type: "agent"
+  def process(text)
+    call_llm(text)
+    transform(text)
+  end
+end
+```
+
+Use `mock: "marked"` 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).
+
 ### Error Handling
 
 Errors are automatically captured and re-raised:

data/lib/bitfab/client.rb

@@ -13,6 +13,12 @@ module Bitfab
   class Client
     SPAN_TYPES = %w[llm agent function guardrail handoff custom].freeze
 
+    # Sentinel returned by check_mock_replay when this span should run real
+    # code (no mock active, wrong strategy, or no matching historical entry).
+    # 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
+
     attr_reader :api_key, :service_url, :enabled
 
     def initialize(api_key:, service_url: nil, enabled: true)
@@ -40,19 +46,41 @@ module Bitfab
     #   code change being tested in this replay (stored on the experiment)
     # @param code_change_files [Array<Hash>, nil] optional list of edited files,
     #   each as { path:, before:, after: } (use "" for new/deleted files)
+    # @param mock [String] mock strategy for child spans: "none" (default),
+    #   "all", or "marked". "all" mocks every child span; "marked" only mocks
+    #   spans declared with mock_on_replay: true.
     # @return [Hash] with :items, :test_run_id, :test_run_url
     def replay(receiver, method_name, trace_function_key:, limit: 5, trace_ids: nil, max_concurrency: 10,
-      code_change_description: nil, code_change_files: nil)
+      code_change_description: nil, code_change_files: nil, mock: "none")
       Replay.run(
         self, receiver, method_name,
         trace_function_key:, limit:, trace_ids:, max_concurrency:,
-        code_change_description:, code_change_files:
+        code_change_description:, code_change_files:, mock:
       )
     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
+    # then wrapping multiple methods or classes with that key. Mirrors
+    # `client.get_function(key)` in the Python SDK and `client.getFunction(key)`
+    # in the TypeScript SDK.
+    #
+    # @example
+    #   fn = Bitfab.client.get_function("order-processing")
+    #   fn.wrap(OrderService, :process_order, type: "function")
+    #   fn.wrap(OrderService, :validate_order, type: "guardrail")
+    #
+    # @param trace_function_key [String]
+    # @return [BitfabFunction]
+    def get_function(trace_function_key)
+      BitfabFunction.new(self, trace_function_key)
+    end
+
     # Execute a block inside a span context, sending trace data on completion.
     # Called by Traceable — not intended for direct use.
-    def execute_span(trace_function_key:, span_name:, span_type:, function_name:, args:, kwargs:)
+    def execute_span(trace_function_key:, span_name:, span_type:, function_name:, args:, kwargs:,
+      mock_on_replay: false)
       return yield unless @enabled
 
       parent = SpanContext.current
@@ -65,16 +93,52 @@ module Bitfab
       replay_ctx = ReplayContext.current
       resolved_test_run_id = replay_ctx&.dig(:test_run_id)
       resolved_input_source_span_id = replay_ctx&.dig(:input_source_span_id)
+      resolved_input_source_trace_id = replay_ctx&.dig(:input_source_trace_id)
 
       # Register trace state for root spans
       if is_root_span && !TraceState.get(trace_id)
-        TraceState.create(trace_id, test_run_id: resolved_test_run_id)
+        TraceState.create(
+          trace_id,
+          test_run_id: resolved_test_run_id,
+          input_source_trace_id: resolved_input_source_trace_id
+        )
       end
 
       if is_root_span
         @pending_span_mutex.synchronize { @pending_span_threads[trace_id] = [] }
       end
 
+      # Advance the per-(key, name) call counter for any non-root span under
+      # an active mock tree, even when this span won't itself be mocked.
+      # Unmarked spans must consume an index so subsequent marked siblings
+      # line up with `build_mock_tree`'s sequential numbering for the same
+      # (key, name) pair. Different (key, name) pairs have independent
+      # counters — they cannot shift each other.
+      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:
+        )
+        if mocked_output != MOCK_REPLAY_MISS
+          send_mocked_span(
+            trace_function_key:,
+            trace_id:,
+            span_id:,
+            parent_span_id:,
+            span_name:,
+            span_type:,
+            function_name:,
+            args:,
+            kwargs:,
+            mocked_output:,
+            started_at:,
+            test_run_id: resolved_test_run_id,
+            input_source_span_id: resolved_input_source_span_id
+          )
+          return mocked_output
+        end
+      end
+
       result = nil
       error = nil
       span_contexts = nil
@@ -228,6 +292,9 @@ module Bitfab
       if trace_state&.dig(:contexts)
         raw_trace["contexts"] = trace_state[:contexts]
       end
+      if trace_state&.dig(:input_source_trace_id)
+        raw_trace["input_source_trace_id"] = trace_state[:input_source_trace_id]
+      end
 
       payload = {
         "type" => "sdk-function",
@@ -296,5 +363,137 @@ module Bitfab
 
       @http_client.send_external_span(payload) # Returns the background thread
     end
+
+    # Advance the per-(key, name) call counter when this invocation is a
+    # non-root span under an active mock tree. Returns the call index this
+    # invocation owns, or nil when there's nothing to advance (root span, or
+    # no replay mock context). The counter MUST advance for every child span
+    # sharing the same (key, name) pair — including spans that won't be
+    # mocked — so unmarked spans don't silently shift subsequent marked
+    # spans' indices. Different (key, name) pairs have independent counters.
+    def advance_mock_counter(replay_ctx, trace_function_key, span_name, is_root_span:)
+      return nil if is_root_span
+      return nil unless replay_ctx&.dig(:mock_tree)
+
+      counters = replay_ctx[:call_counters]
+      counter_key = "#{trace_function_key}:#{span_name}"
+      call_index = counters[counter_key] || 0
+      counters[counter_key] = call_index + 1
+      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:)
+      strategy = replay_ctx[:mock_strategy]
+      case strategy
+      when "marked"
+        return MOCK_REPLAY_MISS unless mock_on_replay
+      when "all"
+        # All non-root spans are eligible
+      else
+        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]
+
+      # 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
+      end
+
+      output
+    end
+
+    # Record a span entry for a mocked invocation so the test run reflects the
+    # mocked execution. Mirrors send_span's payload shape but with the mocked
+    # output as the result and no error. The returned background thread is
+    # registered with @pending_span_threads so the root span's finalize joins
+    # it before sending trace completion; without this the trace completion
+    # can race ahead of the mocked span's HTTP send and the trace lands
+    # temporarily incomplete on the server.
+    def send_mocked_span(trace_function_key:, trace_id:, span_id:, parent_span_id:,
+      span_name:, span_type:, function_name:, args:, kwargs:, mocked_output:,
+      started_at:, test_run_id:, input_source_span_id:)
+      ended_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")
+      span_thread = send_span(
+        trace_function_key:,
+        trace_id:,
+        span_id:,
+        parent_span_id:,
+        span_name:,
+        span_type:,
+        function_name:,
+        contexts: nil,
+        prompt: nil,
+        args:,
+        kwargs:,
+        result: mocked_output,
+        error: nil,
+        started_at:,
+        ended_at:,
+        test_run_id:,
+        input_source_span_id:
+      )
+      # Mocked spans are always non-root (advance_mock_counter returns nil for
+      # root spans, so check_mock_replay never short-circuits them), so the
+      # thread always belongs in the parent's pending list, never standalone.
+      @pending_span_mutex.synchronize do
+        @pending_span_threads[trace_id] << span_thread if span_thread && @pending_span_threads.key?(trace_id)
+      end
+    rescue Exception # rubocop:disable Lint/RescueException
+      # Never crash the host app — mocked span recording is best-effort
+    end
+  end
+
+  # Fluent wrapper bound to a single trace_function_key. Mirrors
+  # `BitfabFunction` in the Python SDK and `BitfabFunction` in the TypeScript
+  # SDK — lets callers wrap multiple methods without repeating the key.
+  class BitfabFunction
+    attr_reader :trace_function_key
+
+    def initialize(client, trace_function_key)
+      @client = client
+      @trace_function_key = trace_function_key
+    end
+
+    # Wrap an existing method on a class with span tracing, binding this
+    # function's trace_function_key.
+    #
+    # Routes spans through the client this function was created from (matches
+    # Python's `BitfabFunction.span()` using `self._client.span(...)` and
+    # TypeScript's `BitfabFunction.withSpan()` using `this.client.withSpan(...)`),
+    # so non-global `Bitfab::Client` instances don't silently fall back to
+    # `Bitfab.client`.
+    #
+    # @example
+    #   fn = Bitfab.client.get_function("openai")
+    #   fn.wrap(OpenAI::Client, :chat, name: "Chat", type: "llm")
+    #
+    # @param klass [Class, Module] the class to wrap
+    # @param method_name [Symbol] the method to wrap
+    # @param name [String, nil] explicit span name (defaults to method name)
+    # @param type [String] span type
+    # @param mock_on_replay [Boolean] mark this span for the "marked" mock strategy
+    def wrap(klass, method_name, name: nil, type: "custom", mock_on_replay: false)
+      Bitfab::Traceable.wrap(
+        klass, method_name,
+        trace_function_key: @trace_function_key,
+        name:, type:, mock_on_replay:,
+        client: @client
+      )
+    end
   end
 end

data/lib/bitfab/http_client.rb

@@ -120,6 +120,17 @@ module Bitfab
       get("/api/sdk/externalSpans/#{span_id}", timeout: 30)
     end
 
+    # Fetch the span tree rooted at an external span. Blocking GET request.
+    # Used by replay when a mock strategy is active so child spans can be
+    # 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)
+    end
+
     # Mark a replay test run as completed. Blocking call.
     def complete_replay(test_run_id)
       request("/api/sdk/replay/complete", {"testRunId" => test_run_id}, timeout: 30)

data/lib/bitfab/replay.rb

@@ -4,6 +4,14 @@ require_relative "constants"
 require_relative "serialize"
 
 module Bitfab
+  # Replay mock strategies. Mirrors the Python and TypeScript SDKs.
+  #
+  # - "none"   — every child span runs real code (default)
+  # - "all"    — every child span returns its historical output
+  # - "marked" — only spans declared with mock_on_replay: true return historical
+  #              output; everything else runs real code
+  MOCK_STRATEGIES = %w[none all marked].freeze
+
   # Thread-local replay context management.
   module ReplayContext
     module_function
@@ -14,12 +22,26 @@ module Bitfab
 
     # Execute a block with replay context set on the current thread.
     # The context is automatically cleared when the block completes.
-    def with_context(test_run_id:, input_source_span_id: nil)
+    #
+    # @param test_run_id [String]
+    # @param input_source_span_id [String, nil]
+    # @param input_source_trace_id [String, nil]
+    # @param mock_tree [Hash{String => Hash}, nil] keyed by "#{key}:#{index}"
+    # @param mock_strategy [String, nil] one of MOCK_STRATEGIES
+    def with_context(test_run_id:, input_source_span_id: nil, input_source_trace_id: nil,
+      mock_tree: nil, mock_strategy: nil)
       previous = Thread.current[REPLAY_CONTEXT_KEY]
-      Thread.current[REPLAY_CONTEXT_KEY] = {
+      ctx = {
         test_run_id:,
-        input_source_span_id:
+        input_source_span_id:,
+        input_source_trace_id:
       }
+      if mock_tree
+        ctx[:mock_tree] = mock_tree
+        ctx[:mock_strategy] = mock_strategy || "none"
+        ctx[:call_counters] = {}
+      end
+      Thread.current[REPLAY_CONTEXT_KEY] = ctx
       yield
     ensure
       Thread.current[REPLAY_CONTEXT_KEY] = previous
@@ -46,9 +68,16 @@ module Bitfab
     #   code change being tested in this replay (stored on the experiment)
     # @param code_change_files [Array<Hash>, nil] optional list of edited files,
     #   each as { path:, before:, after: } (empty string for new/deleted files)
+    # @param mock [String] mock strategy for child spans: "none" (default),
+    #   "all", or "marked". "all" mocks every child span; "marked" only mocks
+    #   spans declared with mock_on_replay: true.
     # @return [Hash] with :items, :test_run_id, :test_run_url
     def run(client, receiver, method_name, trace_function_key:, limit: 5, trace_ids: nil, max_concurrency: 10,
-      code_change_description: nil, code_change_files: nil)
+      code_change_description: nil, code_change_files: nil, mock: "none")
+      unless MOCK_STRATEGIES.include?(mock.to_s)
+        raise ArgumentError, "Invalid mock strategy '#{mock}'. Must be one of: #{MOCK_STRATEGIES.join(", ")}"
+      end
+
       http_client = client.instance_variable_get(:@http_client)
 
       replay_data = http_client.start_replay(
@@ -63,7 +92,7 @@ module Bitfab
       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)
+        process_items(http_client, server_items, receiver, method_name, test_run_id, max_concurrency, mock.to_s)
       else
         []
       end
@@ -84,11 +113,13 @@ module Bitfab
     end
 
     # Process all replay items, optionally in parallel using threads.
-    def process_items(http_client, server_items, receiver, method_name, test_run_id, max_concurrency)
+    def process_items(http_client, server_items, receiver, method_name, test_run_id, max_concurrency, mock_strategy)
       concurrency = max_concurrency || server_items.length
 
       if concurrency <= 1
-        server_items.map { |item| process_single_item(http_client, item, receiver, method_name, test_run_id) }
+        server_items.map do |item|
+          process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy)
+        end
       else
         results_mutex = Mutex.new
         results = []
@@ -101,7 +132,7 @@ module Bitfab
               item, idx = work_mutex.synchronize { work_queue.shift }
               break unless item
 
-              result = process_single_item(http_client, item, receiver, method_name, test_run_id)
+              result = process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy)
               results_mutex.synchronize { results[idx] = result }
             end
           end
@@ -113,11 +144,67 @@ module Bitfab
     end
 
     # Fetch span data and execute a single replay item.
-    def process_single_item(http_client, server_item, receiver, method_name, test_run_id)
+    def process_single_item(http_client, server_item, receiver, method_name, test_run_id, mock_strategy)
       span = http_client.get_external_span(server_item["externalSpanId"])
       item_data = extract_span_data(span)
       metrics = extract_server_item_metrics(server_item)
-      execute_item(item_data, receiver, method_name, test_run_id, span["id"], metrics)
+
+      mock_tree = nil
+      if mock_strategy == "all" || mock_strategy == "marked"
+        tree = http_client.get_span_tree(server_item["externalSpanId"])
+        mock_tree = build_mock_tree(tree["root"] || {})
+      end
+
+      execute_item(
+        item_data,
+        receiver,
+        method_name,
+        test_run_id,
+        span["id"],
+        metrics,
+        input_source_trace_id: span["externalTraceId"],
+        mock_strategy:,
+        mock_tree:
+      )
+    end
+
+    # Walk the children of a root span tree node depth-first and build a
+    # lookup keyed by "#{trace_function_key}:#{span_name}:#{call_index}".
+    #
+    # The root node itself is excluded — at replay time the runtime root span
+    # never queries the mock tree.
+    #
+    # The compound (key, name) match disambiguates same-key spans that come
+    # from the fluent `client.get_function(key).wrap(...)` pattern: every
+    # wrapped method shares trace_function_key but differs in span_name. The
+    # counter is per-(key, name) pair so repeated same-name calls (including
+    # recursion) still order by occurrence. Mirrors the Python and TypeScript
+    # SDKs after HVT-2078 — keying by trace_function_key alone caused the
+    # wrong historical output for fluent-API span sets.
+    def build_mock_tree(root)
+      spans = {}
+      counters = {}
+
+      walk = lambda do |node|
+        key = node["traceFunctionKey"]
+        if key && !key.empty?
+          name = node["spanName"]
+          name = key if name.nil? || name.empty?
+          counter_key = "#{key}:#{name}"
+          index = counters[counter_key] || 0
+          counters[counter_key] = index + 1
+          spans["#{counter_key}:#{index}"] = {
+            source_span_id: node["sourceSpanId"],
+            output: node["output"],
+            output_meta: node["outputMeta"]
+          }
+        end
+        (node["children"] || []).each { |child| walk.call(child) }
+      end
+
+      (root["children"] || []).each { |child| walk.call(child) }
+
+      spans
     end
 
     # Extract input/output data from an external span's rawData.
@@ -155,13 +242,20 @@ module Bitfab
     end
 
     # 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 = {})
+    def execute_item(item, receiver, method_name, test_run_id, input_source_span_id = nil, metrics = {},
+      input_source_trace_id: nil, mock_strategy: "none", mock_tree: nil)
       args, kwargs = Serialize.deserialize_inputs(item)
 
       fn_result = nil
       fn_error = nil
 
-      ReplayContext.with_context(test_run_id:, input_source_span_id:) do
+      ReplayContext.with_context(
+        test_run_id:,
+        input_source_span_id:,
+        input_source_trace_id:,
+        mock_tree:,
+        mock_strategy:
+      ) do
         fn_result = if kwargs.empty?
           receiver.send(method_name, *args)
         else

data/lib/bitfab/span_context.rb

@@ -131,12 +131,13 @@ module Bitfab
       @states_mutex.synchronize { @states[trace_id] }
     end
 
-    def create(trace_id, test_run_id: nil)
+    def create(trace_id, test_run_id: nil, input_source_trace_id: nil)
       @states_mutex.synchronize do
         @states[trace_id] ||= {
           trace_id:,
           started_at: Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ"),
-          test_run_id:
+          test_run_id:,
+          input_source_trace_id:
         }.compact
       end
     end

data/lib/bitfab/traceable.rb

@@ -38,19 +38,32 @@ module Bitfab
     # @param trace_function_key [String] the trace function key
     # @param name [String, nil] explicit span name (defaults to method name)
     # @param type [String] span type: llm, agent, function, guardrail, handoff, custom
-    def self.wrap(klass, method_name, trace_function_key:, name: nil, type: "custom")
+    # @param mock_on_replay [Boolean] mark this span for the "marked" mock strategy.
+    #   When true, `client.replay(... mock: "marked")` returns this span's
+    #   historical output instead of executing the wrapped method.
+    # @param client [Bitfab::Client, nil] route spans through this specific client
+    #   instead of the global `Bitfab.client`. When nil (default), the wrapper
+    #   resolves `Bitfab.client` at each call (so `Bitfab.configure` / `reset!`
+    #   between calls keeps working). Used by `Bitfab::Client#get_function` to
+    #   preserve the bound client through the fluent wrapper, matching Python's
+    #   `BitfabFunction.span()` and TypeScript's `BitfabFunction.withSpan()`.
+    def self.wrap(klass, method_name, trace_function_key:, name: nil, type: "custom",
+      mock_on_replay: false, client: nil)
       span_name = name || method_name.to_s
       method_name_str = method_name.to_s
+      bound_client = client
 
       wrapper = Module.new do
         define_method(method_name) do |*args, **kwargs, &block|
-          Bitfab.client.send(:execute_span,
+          target_client = bound_client || Bitfab.client
+          target_client.send(:execute_span,
             trace_function_key:,
             span_name:,
             span_type: type,
             function_name: method_name_str,
             args:,
-            kwargs:) do
+            kwargs:,
+            mock_on_replay:) do
             super(*args, **kwargs, &block)
           end
         end
@@ -85,7 +98,10 @@ module Bitfab
       # @param trace_function_key [String, nil] trace function key (overrides class-level bitfab_function)
       # @param name [String, nil] explicit span name (defaults to method name)
       # @param type [String] span type: llm, agent, function, guardrail, handoff, custom
-      def bitfab_span(method_name, trace_function_key: nil, name: nil, type: "custom")
+      # @param mock_on_replay [Boolean] mark this span for the "marked" mock strategy.
+      #   When true, `client.replay(... mock: "marked")` returns this span's
+      #   historical output instead of executing the wrapped method.
+      def bitfab_span(method_name, trace_function_key: nil, name: nil, type: "custom", mock_on_replay: false)
         trace_function_key ||= @bitfab_function_key
         unless trace_function_key
           raise "No trace function key provided. Pass `trace_function_key:` to `bitfab_span` " \
@@ -94,14 +110,15 @@ module Bitfab
 
         # If the method already exists (inline or after-method style), wrap it immediately
         if method_defined?(method_name) || private_method_defined?(method_name)
-          _bitfab_wrap_method(method_name, trace_function_key:, name:, type:)
+          _bitfab_wrap_method(method_name, trace_function_key:, name:, type:, mock_on_replay:)
         else
           # Method doesn't exist yet (before-method style) — register for method_added hook
           @_bitfab_pending_spans ||= {}
           @_bitfab_pending_spans[method_name] = {
             trace_function_key:,
             name:,
-            type:
+            type:,
+            mock_on_replay:
           }
         end
       end
@@ -116,7 +133,7 @@ module Bitfab
         _bitfab_wrap_method(method_name, **config)
       end
 
-      def _bitfab_wrap_method(method_name, trace_function_key:, name: nil, type: "custom")
+      def _bitfab_wrap_method(method_name, trace_function_key:, name: nil, type: "custom", mock_on_replay: false)
         span_name = name || method_name.to_s
         method_name_str = method_name.to_s
 
@@ -128,7 +145,8 @@ module Bitfab
               span_type: type,
               function_name: method_name_str,
               args:,
-              kwargs:) do
+              kwargs:,
+              mock_on_replay:) do
               super(*args, **kwargs, &block)
             end
           end

data/lib/bitfab/version.rb

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

metadata

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