bitfab 0.51.9 → 0.51.10

data/lib/bitfab/client.rb

@@ -5,13 +5,18 @@ require "time"
 
 require_relative "constants"
 require_relative "env"
+require_relative "seed_context"
+require_relative "detached_trace"
 require_relative "http_client"
+require_relative "replay_concurrency"
+require_relative "replay_processes"
 require_relative "replay"
 require_relative "span_context"
 require_relative "subtree"
 require_relative "serialize"
 require_relative "timestamp"
 require_relative "warn_once"
+require_relative "thread_propagation"
 
 module Bitfab
   class Client
@@ -32,17 +37,16 @@ module Bitfab
 
     attr_reader :service_url
 
-    # Dataset operations for the authenticated organization.
-    attr_reader :datasets
+    attr_reader :datasets, :assertion_categories, :traces, :labels, :graders
 
-    # Trace lookup operations for the authenticated organization.
-    attr_reader :traces
-
-    def initialize(api_key: nil, service_url: nil, enabled: nil, strict: false)
+    def initialize(api_key: nil, service_url: nil, capture_enabled: nil, enabled: nil, strict: false, simulation_plan: true, timeout: 120, trace_across_threads: nil, env_vars: nil, baml_client: nil, db_snapshot: nil)
+      @db_snapshot = DbSnapshot.normalize_config(db_snapshot)
+      @env_vars = env_vars || ENV.to_h
+      @baml_client = baml_client
       @api_key_config = api_key
       @service_url = service_url || DEFAULT_SERVICE_URL
-      # The user's on/off intent; effective enabled also requires a resolved key.
-      @enabled_in_code = enabled
+      Bitfab.warn_once("enabled-deprecated", "Bitfab(enabled:) is deprecated; use capture_enabled: instead.") unless enabled.nil?
+      @enabled_in_code = capture_enabled.nil? ? enabled : capture_enabled
       @capture_configured_from_env = nil
       @strict = strict
       # Cached only once a non-empty key is found, so an early resolve (before
@@ -51,13 +55,19 @@ module Bitfab
       @api_key_warned = false
       # The key is NOT read here. HttpClient gets a proc so the key is resolved
       # at send time, after any env loading has run.
-      @http_client = HttpClient.new(api_key: -> { resolve_api_key }, service_url: @service_url)
+      @http_client = HttpClient.new(api_key: -> { resolve_api_key }, service_url: @service_url, timeout:)
+      @http_client.configure_simulation_plan(enabled: simulation_plan)
       @datasets = Datasets.new(@http_client)
+      @assertion_categories = AssertionCategories.new(@http_client)
       @traces = Traces.new(@http_client)
+      @labels = Labels.new(@http_client)
+      @graders = Graders.new(@http_client)
       # 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 = []
+      resolved_thread_propagation = trace_across_threads.nil? ? %w[1 true yes].include?(ENV.fetch("BITFAB_TRACE_ACROSS_THREADS", "").strip.downcase) : trace_across_threads
+      ThreadPropagation.install if resolved_thread_propagation
       CommitRef.start_resolution
     end
 
@@ -95,6 +105,10 @@ module Bitfab
       tracing_enabled?
     end
 
+    def capture_enabled
+      tracing_enabled?
+    end
+
     # Replay historical traces through a method and create a test run.
     #
     # @param receiver [Object, Class] an instance for instance methods, or a Class for class methods
@@ -173,8 +187,8 @@ module Bitfab
     #   begins processing each item. Pair it with on_item_finish to distinguish
     #   queued work from in-flight work. A raising callback never crashes the run.
     # @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: Replay::CODE_CHANGE_UNSET, code_change_files: Replay::CODE_CHANGE_UNSET, experiment_group_id: nil, dataset_id: nil, dataset_ids: nil, grader_ids: nil, mock: "marked",
+    def replay(receiver, method_name = nil, trace_function_key:, limit: nil, trace_ids: nil, concurrency: nil, max_concurrency: ReplayConcurrency::UNSET,
+      name: nil, code_change_description: Replay::CODE_CHANGE_UNSET, code_change_files: Replay::CODE_CHANGE_UNSET, experiment_group_id: nil, dataset_id: nil, dataset_ids: nil, grader_ids: nil, mock: "marked", attempts: ReplayConcurrency::UNSET, only_with_assertions: false, dry_run: false,
       adapt_inputs: nil, mock_override: nil, db_branch: nil, on_item_start: nil, on_item_finish: nil, on_progress: nil)
       Replay.run(
         self,
@@ -184,7 +198,11 @@ module Bitfab
         limit:,
         trace_ids:,
         name:,
+        concurrency:,
         max_concurrency:,
+        attempts:,
+        only_with_assertions:,
+        dry_run:,
         code_change_description:,
         code_change_files:,
         experiment_group_id:,
@@ -307,13 +325,14 @@ module Bitfab
     # @param trace_function_key [String]
     # @return [BitfabFunction]
     def get_function(trace_function_key)
+      @http_client.refresh_simulation_plan
       BitfabFunction.new(self, trace_function_key)
     end
 
     # Configure an existing method only when it is discovered beneath a
     # bitfab_trace root owned by this client.
     def node(klass, method_name, name: nil, type: "custom", capture: true,
-      test_run_id: nil, mock_on_replay: false, finalize: nil)
+      test_run_id: nil, mock_on_replay: nil, finalize: nil)
       Bitfab::Traceable.node(
         klass,
         method_name,
@@ -327,6 +346,13 @@ module Bitfab
       )
     end
 
+    # Return a handle that updates a persisted trace without an active span.
+    # Each handle operation waits for the server and raises on rejection.
+    def get_trace(trace_id)
+      validate_trace_id(trace_id)
+      DetachedTrace.new(self, @http_client, trace_id)
+    end
+
     # Fetch one persisted span without loading the full trace.
     # Exactly one of id or name is required. Name lookups return the last
     # matching span by default; occurrence also accepts "first" or a zero-based
@@ -352,11 +378,12 @@ module Bitfab
     # Called by Traceable, not intended for direct use.
     def execute_span(trace_function_key:, span_name:, span_type:, function_name:, args:, kwargs:,
       capture_when: "always", mock_on_replay: false, independent_root: false,
-      explicit_span_receiver: nil, explicit_span_method_name: nil, test_run_id: nil, finalize: nil)
+      explicit_span_receiver: nil, explicit_span_method_name: nil, test_run_id: nil, finalize: nil,
+      surface: nil, on_start: nil, on_complete: nil, around_body: nil, span_data: {}, managed_root: false)
       # Decide at CALL time, not construction. The key may be set after the
       # client is built (env loaded later), so re-checking per call lets a
       # late-resolved key take effect.
-      return yield unless tracing_enabled?
+      return yield unless should_record?
 
       resolved_capture_when = capture_when.to_s
       unless ["always", "nested"].include?(resolved_capture_when)
@@ -395,7 +422,7 @@ module Bitfab
         trace_id = if independent_root && ambient_parent
           SecureRandom.uuid
         else
-          parent&.dig(:trace_id) || replay_ctx&.dig(:trace_id) || SecureRandom.uuid
+          parent&.dig(:trace_id) || SeedContext.current&.dig(:trace_id) || replay_ctx&.dig(:trace_id) || SecureRandom.uuid
         end
         span_id = SecureRandom.uuid
         parent_span_id = parent&.dig(:span_id)
@@ -409,11 +436,16 @@ module Bitfab
         if is_root_span && !TraceState.get(trace_id)
           TraceState.create(
             trace_id,
+            trace_function_key:,
             test_run_id: resolved_test_run_id,
-            input_source_trace_id: resolved_input_source_trace_id
+            input_source_trace_id: resolved_input_source_trace_id,
+            replay_attempt: replay_ctx&.dig(:replay_attempt),
+            db_snapshot: @db_snapshot
           )
         end
 
+        on_start&.call(trace_id:, span_id:)
+
         # 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
@@ -448,8 +480,10 @@ module Bitfab
               started_at:,
               test_run_id: resolved_test_run_id,
               input_source_span_id: resolved_input_source_span_id,
-              mock_source:
+              mock_source:,
+              instrumentation: (surface == :trace) ? "trace" : "span"
             )
+            on_complete&.call(result: mocked_output, error: nil, ended_at: Bitfab.now_iso_timestamp)
             return mocked_output
           end
         end
@@ -477,6 +511,7 @@ module Bitfab
       span_prompt = nil
       completed = false
       @http_client.trace_completion.start(trace_id, span_id)
+      recording_override = managed_root ? {} : nil
 
       complete = lambda do |final_result, final_error|
         # Never crash the host app due to span building/sending. Idempotent:
@@ -499,14 +534,18 @@ module Bitfab
             end
           end
 
+          on_complete&.call(result: recorded_result, error: recorded_error, ended_at:)
+
           send_span(
+            span_data:,
+            instrumentation: recording_override&.dig(:instrumentation) || ((surface == :trace) ? "trace" : "span"),
             trace_function_key:,
             trace_id:,
             span_id:,
             parent_span_id:,
-            span_name:,
-            span_type:,
-            function_name:,
+            span_name: recording_override&.dig(:name) || span_name,
+            span_type: recording_override&.dig(:type) || span_type,
+            function_name: recording_override&.dig(:function_name) || function_name,
             contexts: span_contexts,
             prompt: span_prompt,
             args:,
@@ -567,17 +606,19 @@ module Bitfab
             trace_id:,
             span_id:,
             explicit_span_receiver:,
-            explicit_span_method_name: explicit_span_method_name&.to_sym
+            explicit_span_method_name: explicit_span_method_name&.to_sym,
+            surface:, recording_override:
           ) do
+            invoke = -> { around_body ? around_body.call { yield } : yield }
             result = if explicit_span_receiver
               Subtree.with_borrowed_span(
                 trace_id:,
                 span_id:,
                 receiver: explicit_span_receiver,
                 method_name: explicit_span_method_name
-              ) { yield }
+              ) { invoke.call }
             else
-              yield
+              invoke.call
             end
           ensure
             # Capture contexts before the span context is popped
@@ -610,6 +651,8 @@ module Bitfab
           trace_id:,
           span_id:,
           complete:,
+          surface:,
+          around_body:,
           explicit_span_receiver:,
           explicit_span_method_name:
         )
@@ -673,6 +716,13 @@ module Bitfab
 
     # Whether tracing should run, decided lazily at call time. An explicit
     # `enabled: false` short-circuits without ever touching the key.
+    def should_record?
+      active_run = ReplayContext.current || SeedContext.current
+      return false unless capture_configured? || active_run
+
+      !resolve_api_key.nil?
+    end
+
     def tracing_enabled?
       return false unless capture_configured?
 
@@ -685,8 +735,8 @@ module Bitfab
     # 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:, complete:, explicit_span_receiver: nil,
-      explicit_span_method_name: nil)
-      span_entry = {trace_id:, span_id:}
+      explicit_span_method_name: nil, surface: nil, around_body: nil)
+      span_entry = {trace_id:, span_id:, surface:}
       if explicit_span_receiver
         span_entry[:explicit_span_receiver] = explicit_span_receiver
         span_entry[:explicit_span_method_name] = explicit_span_method_name&.to_sym
@@ -699,7 +749,7 @@ module Bitfab
         # the right parent. The fiber bridge carries the same parent into
         # Enumerator.new / enum_for source fibers on this thread.
         ReplayContext.with_span(trace_id:, span_id:) do
-          SpanContext.with_fiber_bridge(trace_id:, span_id:) do
+          SpanContext.with_fiber_bridge(trace_id:, span_id:, surface:) do
             SpanContext.stack.push(span_entry)
             begin
               iterate = lambda do
@@ -708,6 +758,8 @@ module Bitfab
                   yielder << value
                 end
               end
+              raw_iterate = iterate
+              iterate = -> { around_body.call(&raw_iterate) } if around_body
               if explicit_span_receiver
                 Subtree.with_borrowed_span(
                   trace_id:,
@@ -755,6 +807,7 @@ module Bitfab
         "ended_at" => ended_at
       }
 
+      raw_trace["ingestion_type"] = trace_state[:ingestion_type] if trace_state&.dig(:ingestion_type)
       if trace_state&.dig(:name)
         raw_trace["name"] = trace_state[:name]
       end
@@ -764,6 +817,9 @@ module Bitfab
       if trace_state&.dig(:contexts)
         raw_trace["contexts"] = trace_state[:contexts]
       end
+      unless trace_state&.dig(:replay_attempt).nil?
+        raw_trace["replay_attempt"] = trace_state[:replay_attempt]
+      end
       if trace_state&.dig(:input_source_trace_id)
         raw_trace["input_source_trace_id"] = trace_state[:input_source_trace_id]
       end
@@ -822,7 +878,8 @@ module Bitfab
     def send_span(trace_function_key:, trace_id:, span_id:, parent_span_id:,
       span_name:, span_type:, function_name:, contexts:, prompt:, args:, kwargs:, result:, error:,
       started_at:, ended_at:, test_run_id: nil, input_source_span_id: nil, mocked: false,
-      mock_target: nil, mock_source: nil)
+      mock_target: nil, mock_source: nil, span_data: {}, instrumentation: "span")
+      extra_span_data = span_data
       # If drop() was called on this trace, suppress the span PAYLOAD upload for
       # every span that completes after the flag was set. The trace completion
       # still rides out with dropped: true, so the server scrubs any sibling
@@ -849,6 +906,7 @@ module Bitfab
         "output" => human_output,
         "function_name" => function_name
       }
+      span_data.merge!(extra_span_data)
       span_data["input_serialized"] = marshalled_input if marshalled_input
       span_data["output_serialized"] = marshalled_output if marshalled_output
       if error
@@ -863,7 +921,8 @@ module Bitfab
         "trace_id" => trace_id,
         "started_at" => started_at,
         "ended_at" => ended_at,
-        "span_data" => span_data
+        "span_data" => span_data,
+        "span_origin" => Bitfab.make_span_origin(instrumentation)
       }
       raw_span["parent_id"] = parent_span_id if parent_span_id
       raw_span["input_source_span_id"] = input_source_span_id if input_source_span_id
@@ -875,6 +934,7 @@ module Bitfab
         "source" => "ruby-sdk-function",
         "sourceTraceId" => trace_id,
         "traceFunctionKey" => trace_function_key,
+        "rootTraceFunctionKey" => TraceState.get(trace_id)&.dig(:trace_function_key) || trace_function_key,
         "rawSpan" => raw_span
       }
       payload["testRunId"] = test_run_id if test_run_id
@@ -1029,7 +1089,7 @@ module Bitfab
     # output as the result and no error.
     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:, mock_source:)
+      started_at:, test_run_id:, input_source_span_id:, mock_source:, instrumentation: "span")
       ended_at = Bitfab.now_iso_timestamp
       send_span(
         trace_function_key:,
@@ -1050,6 +1110,7 @@ module Bitfab
         test_run_id:,
         input_source_span_id:,
         mocked: true,
+        instrumentation:,
         # Both paths short-circuit the call, so the target is always the
         # output; only where the value came from differs.
         mock_target: "output",
@@ -1093,7 +1154,7 @@ module Bitfab
     #   values warn once and default to "always"
     # @param mock_on_replay [Boolean] mark this span for the default "marked"
     #   mock strategy. A missing selected occurrence fails closed.
-    def wrap(klass, method_name, name: nil, type: "custom", capture_when: "always", mock_on_replay: false)
+    def wrap(klass, method_name, name: nil, type: "custom", capture_when: "always", mock_on_replay: false, finalize: nil, test_run_id: nil)
       Bitfab::Traceable.wrap(
         klass,
         method_name,
@@ -1102,6 +1163,8 @@ module Bitfab
         type:,
         capture_when:,
         mock_on_replay:,
+        finalize:,
+        test_run_id:,
         client: @client
       )
     end
@@ -1109,7 +1172,7 @@ module Bitfab
     # Experimentally trace an existing method and its first-party call subtree.
     def trace(klass, method_name, name: nil, type: "custom",
       max_depth: Subtree::DEFAULT_MAX_DEPTH, max_spans: Subtree::DEFAULT_MAX_SPANS,
-      exclude: [], include_wrappers: false)
+      exclude: [], include_wrappers: false, mock_on_replay_default: false)
       Bitfab::Traceable.trace(
         klass, method_name,
         trace_function_key: @trace_function_key,
@@ -1119,6 +1182,7 @@ module Bitfab
         max_spans:,
         exclude:,
         include_wrappers:,
+        mock_on_replay_default:,
         client: @client
       )
     end

data/lib/bitfab/db_snapshot.rb

@@ -6,8 +6,8 @@ module Bitfab
   # Every root trace carries a snapshot ref that pins the DB state at trace
   # open by wall-clock timestamp. Capturing the timestamp is free (no IO) and
   # harmless, so it happens on every trace regardless of configuration: that
-  # lets any trace be replayed against a historical branch later. The provider
-  # is resolved at replay time. Mirrors the TypeScript and Python SDKs'
+  # lets any trace be replayed against a historical branch later. The optional configured provider
+  # is captured too; otherwise it is resolved at replay time. Mirrors the TypeScript and Python SDKs'
   # +DbSnapshotRef+ wire shape (camelCase key, since the server speaks the same
   # protocol for every SDK).
   module DbSnapshot
@@ -17,13 +17,26 @@ module Bitfab
     #
     # Stores only the wall clock the SDK observed immediately before invoking
     # the wrapped function; the server-side resolver uses that as the snapshot
-    # timestamp. No provider is captured (it is resolved at replay time).
+    # timestamp. A configured provider is included; otherwise it is resolved at replay time.
     #
     # @param sdk_wall_clock_before_fn [String] ISO wall-clock timestamp the SDK
     #   observed immediately before invoking the wrapped function.
     # @return [Hash] snapshot ref with a single camelCase +sdkWallClockBeforeFn+ key.
-    def build_snapshot_ref(sdk_wall_clock_before_fn)
-      {"sdkWallClockBeforeFn" => sdk_wall_clock_before_fn}
+    def build_snapshot_ref(sdk_wall_clock_before_fn, config: nil)
+      ref = {"sdkWallClockBeforeFn" => sdk_wall_clock_before_fn}
+      ref["provider"] = config.fetch("provider") if config
+      ref
+    end
+
+    def normalize_config(config)
+      return nil if config.nil?
+      raise ArgumentError, "db_snapshot must be a Hash with provider: neon." unless config.is_a?(Hash)
+
+      provider = config[:provider] || config["provider"]
+      unless provider.to_s == "neon"
+        raise ArgumentError, "db_snapshot.provider #{provider.inspect} is not supported. Supported providers: neon."
+      end
+      {"provider" => "neon"}.freeze
     end
   end
 end

data/lib/bitfab/detached_trace.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module Bitfab
+  class DetachedTrace
+    attr_reader :trace_id
+
+    def initialize(client, http_client, trace_id)
+      @client = client
+      @http_client = http_client
+      @trace_id = trace_id
+    end
+
+    def add_context(context)
+      return unless context.is_a?(Hash)
+
+      patch("appendContexts" => [context])
+    end
+
+    def set_metadata(metadata)
+      return unless metadata.is_a?(Hash)
+
+      patch("mergeMetadata" => metadata)
+    end
+
+    def set_session_id(session_id)
+      return unless session_id.is_a?(String) && !session_id.empty?
+
+      patch("setSessionId" => session_id)
+    end
+
+    def set_name(name)
+      return unless name.is_a?(String) && !name.empty?
+
+      patch("setName" => name)
+    end
+
+    private
+
+    def patch(payload)
+      return unless @client.enabled
+
+      @http_client.request("/api/sdk/traces/#{@trace_id}", payload, method: "PATCH", timeout: 10)
+      nil
+    end
+  end
+end

data/lib/bitfab/graders.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require "uri"
+
+module Bitfab
+  class Graders
+    def initialize(http_client)
+      @http_client = http_client
+    end
+
+    def get_labels(trace_ids: nil, grader_id: nil, limit: nil)
+      if (trace_ids.nil? || trace_ids.empty?) && grader_id.nil?
+        raise ArgumentError, "Pass trace_ids, grader_id, or both: there is nothing to read otherwise."
+      end
+
+      params = {}
+      params["traceIds"] = trace_ids.join(",") unless trace_ids.nil? || trace_ids.empty?
+      params["graderId"] = grader_id unless grader_id.nil?
+      params["limit"] = limit unless limit.nil?
+      @http_client.get("/api/sdk/graderLabels?#{URI.encode_www_form(params)}")["labels"]
+    end
+  end
+end