bitfab 0.51.9 → 0.57.0

data/lib/bitfab/seed.rb

@@ -0,0 +1,128 @@
+# frozen_string_literal: true
+
+require "uri"
+require_relative "seed_context"
+
+module Bitfab
+  # Supplies the root span for an explicit-key callable used by seed or replay.
+  class ManagedCallable
+    def initialize(client, key, callable)
+      raise ArgumentError, "Expected a callable function." unless callable.respond_to?(:call)
+
+      @client = client
+      @key = key
+      @callable = callable
+    end
+
+    def call(*args, **kwargs)
+      @client.execute_span(trace_function_key: @key, span_name: @key, span_type: "function",
+        function_name: @key, args:, kwargs:, managed_root: true) do
+        value = @callable.call(*args, **kwargs)
+        value.is_a?(Enumerator) ? value.to_a : value
+      end
+    end
+  end
+
+  class Client
+    SEED_INPUT_UNSET = Object.new.freeze
+
+    # Record a case without executing code, or execute a callable once as a seeded trace.
+    # A case supplies input: (a positional argument Array), expected:, and optional fn:
+    # for signature validation. A run supplies a callable and args:/kwargs:.
+    def seed_trace(trace_function_key, callable = nil, input: SEED_INPUT_UNSET, expected: nil, fn: nil,
+      args: [], kwargs: {}, metadata: nil, session_id: nil, name: nil, span_name: nil, span_type: "function")
+      raise ArgumentError, "Seed trace function key cannot be empty." if trace_function_key.to_s.strip.empty?
+      raise ArgumentError, "Seeding requires an API key." unless resolve_api_key
+
+      case_seed = !input.equal?(SEED_INPUT_UNSET)
+      if case_seed
+        raise ArgumentError, "A seed case input must be an Array of positional arguments." unless input.is_a?(Array)
+        raise ArgumentError, "A seed case cannot also execute a callable." if callable
+        if fn
+          validate_seed_signature(fn, input)
+          key = fn.is_a?(Method) && Traceable.trace_function_key_for(fn.receiver, fn.name)
+          raise ArgumentError, "Seed callable is traced under '#{key}', not '#{trace_function_key}'." if key && key != trace_function_key
+        end
+        args = input
+        kwargs = {}
+        callable = ->(*) { expected }
+      end
+      raise ArgumentError, "Seed requires a callable or input:." unless callable.respond_to?(:call)
+      declared_key = callable.is_a?(Method) && Traceable.trace_function_key_for(callable.receiver, callable.name)
+      if declared_key && declared_key != trace_function_key
+        raise ArgumentError, "Seed callable is traced under '#{declared_key}', not '#{trace_function_key}'."
+      end
+
+      trace_id = SecureRandom.uuid
+      state = TraceState.create(trace_id)
+      state.delete(:db_snapshot_ref)
+      state[:ingestion_type] = "seeded"
+      state[:trace_function_key] = trace_function_key
+      state[:metadata] = metadata.dup if metadata
+      state[:session_id] = session_id if session_id
+      state[:name] = name if name
+      previous_stack = Thread.current[SpanContext::STACK_KEY]
+      previous_bridge = Thread.current.thread_variable_get(SpanContext::FIBER_BRIDGE_STACK_KEY)
+      previous_replay = Thread.current.thread_variable_get(REPLAY_CONTEXT_KEY)
+      Thread.current[SpanContext::STACK_KEY] = []
+      Thread.current.thread_variable_set(SpanContext::FIBER_BRIDGE_STACK_KEY, [])
+      Thread.current.thread_variable_set(REPLAY_CONTEXT_KEY, nil)
+      SeedContext.with_context(trace_id:) do
+        if declared_key
+          value = callable.call(*args, **kwargs)
+          value.to_a if value.is_a?(Enumerator)
+        else
+          execute_span(trace_function_key:, span_name: span_name || trace_function_key,
+            span_type:, function_name: trace_function_key, args:, kwargs:, managed_root: true) do
+            value = callable.call(*args, **kwargs)
+            value.is_a?(Enumerator) ? value.to_a : value
+          end
+        end
+      end
+      raise "Seed execution did not record a root trace." if TraceState.get(trace_id)
+
+      trace_id
+    ensure
+      if trace_id
+        flush
+        TraceState.delete(trace_id)
+        Thread.current[SpanContext::STACK_KEY] = previous_stack
+        Thread.current.thread_variable_set(SpanContext::FIBER_BRIDGE_STACK_KEY, previous_bridge)
+        Thread.current.thread_variable_set(REPLAY_CONTEXT_KEY, previous_replay)
+      end
+    end
+
+    # Execute a stored case again, then adopt the successful run under its existing trace ID.
+    def reseed_trace(trace_function_key, callable, trace_id:)
+      endpoint = "/api/sdk/traces/#{URI.encode_www_form_component(trace_id)}/reseed"
+      source = @http_client.get("#{endpoint}Source")
+      unless source.fetch("traceFunctionKey") == trace_function_key
+        raise ArgumentError, "Trace #{trace_id} belongs to '#{source["traceFunctionKey"]}', not '#{trace_function_key}'."
+      end
+      args, kwargs = Serialize.deserialize_inputs({"input" => source["input"], "input_serialized" => source["inputSerialized"]})
+      run_trace_id = seed_trace(trace_function_key, callable, args:, kwargs:,
+        metadata: source["metadata"], session_id: source["sessionId"], name: source["name"])
+      adopted = @http_client.request(endpoint, {runTraceId: run_trace_id})
+      {trace_id: adopted.fetch("traceId"), previous_run_trace_id: adopted["previousRunTraceId"]}
+    end
+
+    private
+
+    def validate_seed_signature(callable, args)
+      raise ArgumentError, "Seed fn must be callable." unless callable.respond_to?(:call)
+
+      if callable.is_a?(Method)
+        while callable.owner.instance_variable_defined?(:@bitfab_span_method) && callable.super_method
+          callable = callable.super_method
+        end
+      end
+      parameters = callable.respond_to?(:parameters) ? callable.parameters : callable.method(:call).parameters
+      required = parameters.count { |kind, _| kind == :req }
+      maximum = parameters.count { |kind, _| [:req, :opt].include?(kind) }
+      variadic = parameters.any? { |kind, _| kind == :rest }
+      if args.length < required || (!variadic && args.length > maximum) || parameters.any? { |kind, _| kind == :keyreq }
+        raise ArgumentError, "Seed input does not match the callable's positional arguments."
+      end
+    end
+  end
+end

data/lib/bitfab/seed_cli.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module Bitfab
+  module SeedCli
+    module_function
+
+    def parse_cases(raw, path, run)
+      text = raw.strip
+      raise ArgumentError, "Seed file '#{path}' is empty." if text.empty?
+
+      cases = text.start_with?("[") ? JSON.parse(text) : text.lines.reject { |line| line.strip.empty? }.map { |line| JSON.parse(line) }
+      cases.each_with_index do |value, index|
+        unless value.is_a?(Hash) && value["input"].is_a?(Array)
+          raise ArgumentError, "Seed case #{index} must be an object with an input array."
+        end
+        if run && value.key?("expected")
+          raise ArgumentError, "Seed case #{index} carries expected, which --run does not record. Remove expected or drop --run."
+        end
+      end
+      cases
+    end
+
+    def seed_from_registry(registry, pipeline, cases, run: false)
+      entry = registry.fetch(pipeline)
+      callable = entry.method_name ? entry.receiver.method(entry.method_name) : entry.receiver
+      trace_ids = cases.map do |value|
+        options = {metadata: value["metadata"], session_id: value["sessionId"], name: value["name"]}
+        if run
+          raise ArgumentError, "Run-based seed cases cannot carry expected." if value.key?("expected")
+
+          entry.client.seed_trace(entry.trace_function_key, callable, args: value.fetch("input"), **options)
+        else
+          entry.client.seed_trace(entry.trace_function_key, input: value.fetch("input"), expected: value["expected"], fn: callable, **options)
+        end
+      end
+      {pipeline:, traceFunctionKey: entry.trace_function_key, traceIds: trace_ids}
+    end
+
+    def reseed_from_registry(registry, pipeline, trace_ids)
+      entry = registry.fetch(pipeline)
+      callable = entry.method_name ? entry.receiver.method(entry.method_name) : entry.receiver
+      reseeded = trace_ids.map do |trace_id|
+        value = entry.client.reseed_trace(entry.trace_function_key, callable, trace_id:)
+        {traceId: value[:trace_id], previousRunTraceId: value[:previous_run_trace_id]}
+      end
+      {pipeline:, traceFunctionKey: entry.trace_function_key, reseeded:}
+    end
+
+    def run(registry, argv: ARGV, stdout: $stdout, stderr: $stderr)
+      args = ReplayCli.parse(registry, argv.dup)
+      if (!!args[:seed] == !!args[:from_trace]) || (args[:from_trace] && args[:run])
+        raise ArgumentError, "Supply --cases PATH [--run] or --from-trace IDS."
+      end
+      pipeline = args.fetch(:pipeline)
+      result = if args[:from_trace]
+        stderr.puts "[seed] Re-seeding #{args[:from_trace].length} trace(s) through #{pipeline}..."
+        reseed_from_registry(registry, pipeline, args[:from_trace])
+      else
+        cases = parse_cases(File.read(args[:seed]), args[:seed], args[:run])
+        stderr.puts "[seed] Recording #{cases.length} case(s) through #{pipeline}..."
+        seed_from_registry(registry, pipeline, cases, run: !!args[:run])
+      end
+      stdout.puts JSON.generate(result)
+      result
+    end
+  end
+end

data/lib/bitfab/seed_context.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Bitfab
+  module SeedContext
+    KEY = :bitfab_seed_context
+
+    module_function
+
+    def current
+      Thread.current.thread_variable_get(KEY)
+    end
+
+    def with_context(trace_id:)
+      previous = current
+      Thread.current.thread_variable_set(KEY, {trace_id:})
+      yield
+    ensure
+      Thread.current.thread_variable_set(KEY, previous)
+    end
+  end
+end

data/lib/bitfab/simulation_plan.rb

@@ -0,0 +1,307 @@
+# frozen_string_literal: true
+
+module Bitfab
+  class SimulationPlan
+    CONTENT_KEYS = %w[input input_meta output output_meta input_serialized output_serialized].freeze
+    FRAMEWORKS = %w[openai-agents langgraph claude-agent-sdk vercel-ai].freeze
+    READ_TIMEOUT = 5.0
+    REFRESH_SECONDS = 60.0
+    RETRY_SECONDS = 10.0
+    MAX_HELD = 1000
+    @plans = ObjectSpace::WeakMap.new
+    @plans_mutex = Mutex.new
+    @plans_pid = Process.pid
+
+    class << self
+      def register(plan)
+        reset_registry_after_fork
+        @plans_mutex.synchronize { @plans[plan] = true }
+      end
+
+      def release_all(timeout, stop: false)
+        reset_registry_after_fork
+        deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + [timeout, 0].max
+        plans = @plans_mutex.synchronize { @plans.keys }
+        released = true
+        plans.each do |plan|
+          result = plan.release([deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC), 0].max)
+          result = plan.stop if stop
+          released = result && released
+        end
+        released
+      end
+
+      private
+
+      def reset_registry_after_fork
+        return if @plans_pid == Process.pid
+
+        @plans_pid = Process.pid
+        @plans_mutex = Mutex.new
+      end
+    end
+
+    def initialize(reader, enabled: true)
+      @reader = reader
+      @enabled = enabled
+      @pid = Process.pid
+      @mutex = Mutex.new
+      @condition = ConditionVariable.new
+      @content_off = nil
+      @refresh_after = 0.0
+      @thread = nil
+      @held = []
+      @draining = Set.new
+      @stopped = false
+      @first_read_finished = false
+      self.class.register(self)
+    end
+
+    def disabled?
+      !@enabled || !ENV.fetch("BITFAB_DISABLE_SIM_PLAN", "").strip.empty?
+    end
+
+    def refresh
+      reset_after_fork
+      return if disabled? || @stopped
+
+      @mutex.synchronize { start_reader(0) if @thread.nil? && now >= @refresh_after }
+    end
+
+    def send_span(payload, &submit)
+      refresh
+      data = payload.dig("rawSpan", "span_data")
+      key = payload["rootTraceFunctionKey"] || payload["traceFunctionKey"]
+      if disabled? || key.nil? || !data.is_a?(Hash) || framework?(payload)
+        submit.call(payload)
+        return
+      end
+      entry = [payload, key, submit]
+      ready = @mutex.synchronize { reading? ? hold(entry) : entry }
+      deliver(*ready) if ready
+    end
+
+    def send_trace(payload, &submit)
+      refresh
+      entry = [payload, nil, submit]
+      ready = @mutex.synchronize do
+        id = trace_id(payload)
+        if !disabled? && !@stopped && id && (@draining.include?(id) || @held.any? { |held| held[1] && trace_id(held[0]) == id })
+          hold(entry)
+        else
+          entry
+        end
+      end
+      deliver(*ready) if ready
+    end
+
+    def release(timeout)
+      refresh
+      deadline = now + [timeout, 0].max
+      @mutex.synchronize do
+        while @thread && (!@held.empty? || !@draining.empty?)
+          remaining = deadline - now
+          break unless remaining.positive?
+          @condition.wait(@mutex, remaining)
+        end
+        if @content_off.nil? && !@held.empty?
+          Bitfab.warn_once("sim-plan-never-loaded", "#{@held.length} records remain held because the simulation plan has not loaded")
+        end
+        @held.empty? && @draining.empty?
+      end
+    end
+
+    def stop
+      reset_after_fork
+      queued = @mutex.synchronize do
+        @stopped = true
+        @condition.broadcast
+        entries, @held = @held, []
+        entries
+      end
+      queued.each { |entry| deliver_safely(*entry) }
+      @mutex.synchronize { @held.empty? && @draining.empty? }
+    end
+
+    def content_off?(key, name)
+      reset_after_fork
+      return false if disabled? || @stopped
+
+      @mutex.synchronize { @content_off&.fetch(key, nil)&.include?(name) == true }
+    end
+
+    def content_off_names?(key)
+      reset_after_fork
+      return false if disabled? || @stopped
+
+      @mutex.synchronize { @content_off&.fetch(key, nil)&.any? == true }
+    end
+
+    def unavailable?
+      reset_after_fork
+      return false if disabled?
+
+      @mutex.synchronize { @content_off.nil? && !reading? }
+    end
+
+    def wait_for_first_read(timeout = READ_TIMEOUT)
+      reset_after_fork
+      return if disabled?
+
+      deadline = now + [timeout, 0].max
+      @mutex.synchronize do
+        while reading? && @thread
+          remaining = deadline - now
+          break unless remaining.positive?
+          @condition.wait(@mutex, remaining)
+        end
+      end
+    end
+
+    def apply(payload, key)
+      return payload if framework?(payload)
+
+      names = @mutex.synchronize { @content_off&.fetch(key, nil) }
+      data = payload.dig("rawSpan", "span_data")
+      return payload unless data.is_a?(Hash) && names&.include?(data["name"])
+
+      without_content(payload, data)
+    end
+
+    private
+
+    def now
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+
+    def reading?
+      @content_off.nil? && !@first_read_finished && !@stopped
+    end
+
+    def trace_id(payload)
+      payload["traceId"] || payload["sourceTraceId"] || payload.dig("externalTrace", "id") || payload["id"]
+    end
+
+    def framework?(payload)
+      FRAMEWORKS.include?(payload.dig("rawSpan", "span_origin", "instrumentation", "name"))
+    end
+
+    def root?(payload)
+      payload.dig("rawSpan", "parent_id").nil?
+    end
+
+    def hold(entry)
+      @held << entry
+      start_reader([@refresh_after - now, 0].max) if @thread.nil? && !@stopped && !disabled?
+      @held.shift if @held.length > MAX_HELD
+    end
+
+    def deliver(payload, key, submit)
+      submit.call(prepare(payload, key))
+    end
+
+    def deliver_safely(payload, key, submit)
+      deliver(payload, key, submit)
+    rescue => error
+      Bitfab.warn_once("sim-plan-held-span-dropped", "Held record could not be sent: #{error}")
+    end
+
+    def prepare(payload, key)
+      return payload if key.nil? || disabled?
+      return apply(payload, key) if @mutex.synchronize { @content_off }
+
+      without_unread_content(payload)
+    end
+
+    def without_unread_content(payload)
+      data = payload.dig("rawSpan", "span_data")
+      return payload if !data.is_a?(Hash) || root?(payload) || framework?(payload)
+
+      Bitfab.warn_once("sim-plan-unread-content",
+        "Sending spans without inputs and outputs because the simulation plan could not be read")
+      without_content(payload, data)
+    end
+
+    def without_content(payload, data)
+      kept = data.except(*CONTENT_KEYS)
+      kept["content_off_by_simulation_plan"] = true
+      payload.merge("rawSpan" => payload["rawSpan"].merge("span_data" => kept))
+    end
+
+    def start_reader(delay)
+      @thread = Thread.new { read_until_loaded(delay) }
+      @thread.name = "bitfab-sim-plan"
+    rescue ThreadError
+      @first_read_finished = true
+      @refresh_after = now + RETRY_SECONDS
+    end
+
+    def pause(delay)
+      @mutex.synchronize do
+        deadline = now + delay
+        while !@stopped && (remaining = deadline - now).positive?
+          @condition.wait(@mutex, remaining)
+        end
+      end
+    end
+
+    def read_once
+      body = @reader.get_simulation_plan
+      unless body.is_a?(Hash) && body["nodes"].is_a?(Array)
+        Bitfab.warn_once("sim-plan-unreadable", "Simulation plan response was not understood; sending spans without inputs and outputs until it loads")
+        return
+      end
+      names = {}
+      body["nodes"].each do |node|
+        next unless node.is_a?(Hash) && node["traceFunctionKey"].is_a?(String) && node["name"].is_a?(String) && node["captureContent"] == false
+        (names[node["traceFunctionKey"]] ||= Set.new).add(node["name"])
+      end
+      names
+    rescue => error
+      return {} if error.respond_to?(:response) && error.response&.code.to_i == 404
+
+      Bitfab.warn_once("sim-plan-unavailable", "Could not read simulation plan: #{error}; sending spans without inputs and outputs until it loads")
+      nil
+    end
+
+    def read_until_loaded(delay)
+      pause(delay) if delay.positive?
+      parsed = (disabled? || @stopped) ? nil : read_once
+      entries = @mutex.synchronize do
+        @first_read_finished = true
+        @content_off = parsed if parsed
+        @refresh_after = now + (parsed ? REFRESH_SECONDS : RETRY_SECONDS)
+        queued, @held = @held, []
+        @draining = queued.filter_map { |payload, key, _| trace_id(payload) if key }.to_set
+        @condition.broadcast
+        queued
+      end
+      until entries.empty?
+        entries.each { |entry| deliver_safely(*entry) }
+        entries = @mutex.synchronize do
+          queued, @held = @held, []
+          @draining.clear if queued.empty?
+          queued
+        end
+      end
+    ensure
+      @mutex.synchronize do
+        @draining.clear
+        @thread = nil
+        @first_read_finished = true
+        @condition.broadcast
+      end
+    end
+
+    def reset_after_fork
+      return if @pid == Process.pid
+
+      @pid = Process.pid
+      @mutex = Mutex.new
+      @condition = ConditionVariable.new
+      @thread = nil
+      @held = []
+      @draining = Set.new
+    end
+  end
+end

data/lib/bitfab/span_context.rb

@@ -139,8 +139,9 @@ module Bitfab
     # Enumerator.new and enum_for run their source body in another Fiber on
     # the same thread. Bridge only the active Enumerator parent across that
     # boundary; ordinary fiber-local span stacks remain isolated.
-    def with_fiber_bridge(trace_id:, span_id:)
+    def with_fiber_bridge(trace_id:, span_id:, surface: nil)
       entry = {trace_id:, span_id:}
+      entry[:surface] = surface if surface
       fiber_bridge_stack.push(entry)
       yield
     ensure
@@ -149,8 +150,10 @@ module Bitfab
 
     # Execute a block with a new span pushed onto the stack.
     # The span is automatically popped when the block completes.
-    def with_span(trace_id:, span_id:, explicit_span_receiver: nil, explicit_span_method_name: nil)
+    def with_span(trace_id:, span_id:, explicit_span_receiver: nil, explicit_span_method_name: nil, surface: nil, recording_override: nil)
       entry = {trace_id:, span_id:}
+      entry[:surface] = surface if surface
+      entry[:recording_override] = recording_override if recording_override
       if explicit_span_receiver
         entry[:explicit_span_receiver] = explicit_span_receiver
         entry[:explicit_span_method_name] = explicit_span_method_name
@@ -182,20 +185,22 @@ module Bitfab
       @states_mutex.synchronize { @states[trace_id] }
     end
 
-    def create(trace_id, test_run_id: nil, input_source_trace_id: nil)
+    def create(trace_id, test_run_id: nil, input_source_trace_id: nil, replay_attempt: nil, trace_function_key: nil, db_snapshot: nil)
       @states_mutex.synchronize do
         @states[trace_id] ||= begin
           started_at = Bitfab.now_iso_timestamp
           {
             trace_id:,
+            trace_function_key:,
             started_at:,
             test_run_id:,
             input_source_trace_id:,
+            replay_attempt:,
             # Capture the wall clock now, before the wrapped function runs.
             # Stored on every trace (no IO, harmless) so any trace can later be
             # replayed against a historical branch; the provider is resolved at
             # replay time.
-            db_snapshot_ref: DbSnapshot.build_snapshot_ref(started_at)
+            db_snapshot_ref: DbSnapshot.build_snapshot_ref(started_at, config: db_snapshot)
           }.compact
         end
       end

data/lib/bitfab/span_origin.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+require_relative "version"
+
+module Bitfab
+  def self.make_span_origin(instrumentation)
+    {
+      "name" => "bitfab.sdk.ruby",
+      "version" => VERSION,
+      "instrumentation" => {"name" => instrumentation}
+    }
+  end
+end