bitfab 0.20.0 → 0.20.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9624757032370e1e318121cabcf2bfa984c7bcdcd07c9ecd5dcee449fef57a7c
-  data.tar.gz: c3925bbae804c16d21c3f373ec6cde4de340ee645fc6423505fc0b5d5f5e638e
+  metadata.gz: d2428cc39dcd97d154e36ba9792628478e86a4b4dde5101c79d37ebb1635378f
+  data.tar.gz: effa0dba40674696ad3d9155615d9d9e2e40a3fe7f9aed9c8ee97516985e13c6
 SHA512:
-  metadata.gz: d2c33cd4a3432d79703685f00bb8c8278e2bebbd694b3a695f1dee646680a5af11370f233635f04590eeb0d5cedee850c496b40bee2fc176c34cf22fb721cde8
-  data.tar.gz: c64be06a07ca8ea807bca23e0f96acb15e1647945b33c3ea731f10c0b52759870b88637d86b680aee07caac7836e3818946f2fb85f0979d83b3ea89ef8ea1704
+  metadata.gz: 5f172f0178829a08315ae1caf3de6ff4a22389b2f83dbf449721ca53ca3a466580af00c15bd5279dffc2c9d074aee2afd0e157733f9257987460ef8d101350c9
+  data.tar.gz: 9c6af19d4ebac35d13280aeaa89ac6104cbff3020e24813e292e4170c6524016a0df943025f3ced5510345c863e4b86f9f027fe281912e99a6438f863dde9d28

data/lib/bitfab/replay.rb

@@ -2,6 +2,7 @@
 
 require_relative "constants"
 require_relative "serialize"
+require_relative "traceable"
 
 module Bitfab
   # Replay mock strategies. Mirrors the Python and TypeScript SDKs.
@@ -117,6 +118,21 @@ module Bitfab
           "determines how many traces replay."
       end
 
+      # Reject a trace_function_key that contradicts the method's declared key.
+      # replay() fetches historical traces by trace_function_key but records the
+      # replayed spans under the method's own bitfab_span key (via send below),
+      # so a mismatch produces an incoherent test run (traces fetched for one
+      # function, recorded under another). Only fires when the method's key is
+      # introspectable; an untraced method falls through to the persistence
+      # check in complete_replay. Mirrors the TypeScript/Python SDKs.
+      declared_key = Traceable.trace_function_key_for(receiver, method_name)
+      if declared_key && declared_key != trace_function_key
+        raise ArgumentError,
+          "Method #{method_name} is traced under trace function key '#{declared_key}' but replay was " \
+          "called with '#{trace_function_key}'. Pass trace_function_key: '#{declared_key}' to replay it, " \
+          "or point at the method traced under '#{trace_function_key}'."
+      end
+
       http_client = client.instance_variable_get(:@http_client)
 
       # limit is meaningless with explicit trace_ids (the ID list determines
@@ -154,6 +170,10 @@ module Bitfab
       # loudly.
       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
+      # run's tokens (span-aggregated server-side), used below to fill each
+      # item's :tokens.
+      replay_tokens = complete_response&.dig("tokens") || {}
 
       if trace_id_map.nil?
         # Older servers don't return the mapping. Preserve the legacy
@@ -186,6 +206,9 @@ module Bitfab
             completed_count += 1
             missing << item[:trace_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?
@@ -370,27 +393,34 @@ 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.
+    # Pull durationMs / model from the start-replay server item. Nil-safe
+    # defaults so older servers without these fields still produce a consistent
+    # shape. Tokens are intentionally NOT read from the start item (it carries
+    # 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.
     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:,
+        tokens: nil,
         model: server_item["model"]
       }
     end
 
+    # Normalize a complete-replay tokens hash (string-keyed JSON) into the
+    # symbol-keyed shape the replay item exposes. Nil when the server reported
+    # no token data for this trace.
+    def normalize_tokens(raw_tokens)
+      return nil unless raw_tokens.is_a?(Hash)
+
+      {
+        input: raw_tokens["input"],
+        output: raw_tokens["output"],
+        cached: raw_tokens["cached"],
+        total: raw_tokens["total"]
+      }
+    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 = {},
       input_source_trace_id: nil, mock_strategy: "none", mock_tree: nil, adapt_inputs: nil, adapt_ctx: nil,

data/lib/bitfab/traceable.rb

@@ -69,9 +69,43 @@ module Bitfab
         end
       end
 
+      wrapper.instance_variable_set(:@bitfab_span_method, method_name_str)
+      wrapper.instance_variable_set(:@bitfab_span_key, trace_function_key)
       klass.prepend(wrapper)
     end
 
+    # Resolve the trace function key a traced method records spans under, or
+    # nil if the method carries no Bitfab span wrapper. Walks the receiver's
+    # ancestor chain for the wrapper module that bitfab_span / Traceable.wrap
+    # prepended. Instances resolve via the class. A Class/Module receiver
+    # resolves either way: singleton-class ancestors first (class/module
+    # methods wrapped on the singleton class), then the class's own ancestors
+    # (methods whose wrapper was prepended onto the class itself), so a
+    # class-method replay is guarded regardless of which surface carries the
+    # wrapper. Singleton-first keeps the method actually dispatched by
+    # `receiver.send` (a class method shadows a same-named instance wrapper).
+    #
+    # replay() uses this to reject a trace_function_key that contradicts the
+    # method's declared key: a mismatch fetches one function's historical
+    # traces but re-records the replay under the method's own key, producing
+    # an incoherent test run. Mirrors the TypeScript and Python SDKs, which
+    # throw on the same key mismatch.
+    def self.trace_function_key_for(receiver, method_name)
+      name = method_name.to_s
+      ancestors = if receiver.is_a?(Module)
+        receiver.singleton_class.ancestors + receiver.ancestors
+      else
+        receiver.class.ancestors
+      end
+      ancestors.each do |mod|
+        next unless mod.instance_variable_defined?(:@bitfab_span_method)
+        next unless mod.instance_variable_get(:@bitfab_span_method) == name
+
+        return mod.instance_variable_get(:@bitfab_span_key)
+      end
+      nil
+    end
+
     module ClassMethods
       # Set the trace function key for this class.
       # All spans declared in this class will be grouped under this key.
@@ -152,6 +186,8 @@ module Bitfab
           end
         end
 
+        wrapper.instance_variable_set(:@bitfab_span_method, method_name_str)
+        wrapper.instance_variable_set(:@bitfab_span_key, trace_function_key)
         prepend(wrapper)
       end
     end

data/lib/bitfab/version.rb

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

metadata

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