bitfab 0.9.0 → 0.10.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 78b34044dcc028414fb2f6fbe020f98bf0ce1a953507c8939bc12d8c15630b09
-  data.tar.gz: 174aa2eb61c2aca0a11078f2361b4392d7612f2a582e3dd4046608cbbb539a27
+  metadata.gz: e5d6df92ce06e91ff963c8dc0e78af3a0f591c51b53330c1c8630bb5d66e7663
+  data.tar.gz: 199d4883cfad921563b3b5786d2db5b017d4b618e000c7d259ed9f69be3eb7ba
 SHA512:
-  metadata.gz: d2eeec89057ecdee5f7f5eae5342ad9218aa190fc472aa01d3dbf0323262b33165faed2d61d6ae0be9bf7da3af9a7e88902e0fa426344b53751b78d0c1f5cf18
-  data.tar.gz: 9726da71f354723007ff6b9bf38466544e93a1ae21f0b02e2b258827697d4996caea781d1c062c7386dd5508eddad097fbc85f4382798e07cd11a0535a6e1a7a
+  metadata.gz: 02c87809df7f4039dd608d3e60ddc0ff85b791d834ea7dc5557a368987d5dba1a174705074a26a9df80d4a8bb7c4ff7493f4c8ac03beb457629d896768a70567
+  data.tar.gz: 579f94120498658ceeb5129b3a21a9d1477d816be9216963c168e0c8747aef0687298f632539f33ace2960db53a19a9ab7847f776d27d971e1b006ad1ebbc909

data/lib/bitfab/client.rb

@@ -53,9 +53,13 @@ module Bitfab
       is_root_span = parent_span_id.nil?
       started_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")
 
+      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)
+
       # Register trace state for root spans
       if is_root_span && !TraceState.get(trace_id)
-        TraceState.create(trace_id)
+        TraceState.create(trace_id, test_run_id: resolved_test_run_id)
       end
 
       if is_root_span
@@ -66,27 +70,18 @@ module Bitfab
       error = nil
       span_contexts = nil
       span_prompt = nil
+      finalized = false
+
+      finalize = lambda do |final_result, final_error|
+        # Never crash the host app due to span building/sending. Idempotent —
+        # only the first call sends the span. Subsequent calls (e.g. from the
+        # enumerator wrapper after iteration completes) are no-ops.
+        next if finalized
+        finalized = true
 
-      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)
-        end
-      rescue => e
-        error = e.message
-        raise
-      ensure
-        # Never crash the host app due to span building/sending
         begin
           ended_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")
 
-          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)
-
           span_thread = send_span(
             trace_function_key:,
             trace_id:,
@@ -99,8 +94,8 @@ module Bitfab
             prompt: span_prompt,
             args:,
             kwargs:,
-            result:,
-            error:,
+            result: final_result,
+            error: final_error,
             started_at:,
             ended_at:,
             test_run_id: resolved_test_run_id,
@@ -130,11 +125,78 @@ module Bitfab
         end
       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)
+        end
+      rescue => e
+        error = e.message
+        finalize.call(result, error)
+        raise
+      end
+
+      # If the wrapped block returned an Enumerator (lazy iteration via
+      # `enum_for`, `to_enum`, `Enumerator.new`, `[...].lazy.map(...)`, etc.),
+      # the work hasn't actually run yet — the values are produced as the
+      # caller iterates. Without special handling we'd close the span here
+      # with `result == <the Enumerator object>`, and any nested `bitfab_span`
+      # calls inside the enumerator body would see an empty span stack and
+      # post their own root traces, fragmenting one logical workflow.
+      #
+      # Instead, hand the caller a wrapping Enumerator whose body restores
+      # the parent span stack on the iterating fiber, drives the source,
+      # 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
+
+      finalize.call(result, error)
       result
     end
 
     private
 
+    # Build an Enumerator that drives `source`, restoring `[trace_id, span_id]`
+    # on the iterating fiber so nested `bitfab_span` calls inside lazy / `each`
+    # callbacks nest under the parent span. Yielded values are collected as the
+    # span output. The span is sent exactly once — when iteration finishes,
+    # raises, or the wrapper is `.close`d.
+    def wrap_enumerator(source, trace_id:, span_id:, finalize:)
+      span_entry = {trace_id:, span_id:}
+      yielded = []
+
+      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
+          end
+          finalize.call(yielded, nil)
+        rescue => e
+          finalize.call(yielded, e.message)
+          raise
+        ensure
+          SpanContext.stack.pop
+        end
+      end
+    end
+
     def validate_span_type!(type)
       return if SPAN_TYPES.include?(type.to_s)
 
@@ -169,6 +231,9 @@ module Bitfab
       if trace_state&.dig(:session_id)
         payload["sessionId"] = trace_state[:session_id]
       end
+      if trace_state&.dig(:test_run_id)
+        payload["testRunId"] = trace_state[:test_run_id]
+      end
 
       @http_client.send_external_trace(payload)
 

data/lib/bitfab/replay.rb

@@ -105,7 +105,8 @@ module Bitfab
     def process_single_item(http_client, server_item, receiver, method_name, test_run_id)
       span = http_client.get_external_span(server_item["externalSpanId"])
       item_data = extract_span_data(span)
-      execute_item(item_data, receiver, method_name, test_run_id, span["id"])
+      metrics = extract_server_item_metrics(server_item)
+      execute_item(item_data, receiver, method_name, test_run_id, span["id"], metrics)
     end
 
     # Extract input/output data from an external span's rawData.
@@ -121,8 +122,29 @@ module Bitfab
       }
     end
 
+    # Pull durationMs / tokens / model from the start-replay server item.
+    # Normalizes to symbol-keyed tokens hash and nil-safe defaults so older
+    # servers without these fields still produce a consistent shape.
+    def extract_server_item_metrics(server_item)
+      raw_tokens = server_item["tokens"]
+      tokens = if raw_tokens.is_a?(Hash)
+        {
+          input: raw_tokens["input"],
+          output: raw_tokens["output"],
+          cached: raw_tokens["cached"],
+          total: raw_tokens["total"]
+        }
+      end
+
+      {
+        duration_ms: server_item["durationMs"],
+        tokens:,
+        model: server_item["model"]
+      }
+    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)
+    def execute_item(item, receiver, method_name, test_run_id, input_source_span_id = nil, metrics = {})
       args, kwargs = Serialize.deserialize_inputs(item)
 
       fn_result = nil
@@ -142,7 +164,10 @@ module Bitfab
         input: args,
         result: fn_result,
         original_output: item["output"],
-        error: fn_error
+        error: fn_error,
+        duration_ms: metrics[:duration_ms],
+        tokens: metrics[:tokens],
+        model: metrics[:model]
       }
     end
   end

data/lib/bitfab/span_context.rb

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

data/lib/bitfab/version.rb

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

metadata

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