bitfab 0.57.6 → 0.57.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 34dfb7ca50f4c2dceaa4aa4ed46ee6e9426a5d8051a3602084fee0640a77c9e9
-  data.tar.gz: ad1e112ccdfc9a2cee2860b602e3edf3629686ef6228e81ac4b80f5132ae1fc1
+  metadata.gz: b92a57be2a0c05885f0f84f757ede302ab14aed53ad422c85ee2781591a39786
+  data.tar.gz: ddb9b050e2f66e9591c17be161c2ebd25d395b625704baee80e396801f5a9e25
 SHA512:
-  metadata.gz: 764ad6880ee3656d90d28ecb5571298cc54ab2bd1bc6a6b2b07399fd2bee8fd151ae3adcfdbf9f2bdc8f7c0d0fccbb1da29e33c93dc68353e11058be22dc9cc8
-  data.tar.gz: afadb371e2d697bab8aaac2a4f02a57491ab37745b6b756ca3302931915f6fec7eabcc80d4789e79f06a5f5b9dae9ab22627a1fc4f35a7f2f610d175da2c12ef
+  metadata.gz: 0542a3b4dcb64f183ce35f7c9f16ead26e858fa1fefb7889115842e552b4855842ba175bea19893912ec37547708403d115259a201b2e02a1e6cb5ffa84f79f6
+  data.tar.gz: a530cd08bc1cdbc42467ea640c8f64e3035569c6723032626284bd6397744a21f5149ef0bc2d32af50aab1b1bbb856e8ba7510d1eb3640824de4f6f4e4472c43

data/lib/bitfab/baml.rb

@@ -170,7 +170,6 @@ module Bitfab
   end
 
   class Client
-    # Execute a stored BAML prompt locally and return its parsed output.
     def call(method_name, **inputs)
       version = @http_client.request("/api/sdk/functions/lookup", {"name" => method_name})
       raise ArgumentError, "Function #{method_name.inspect} not found. Create it at #{@service_url}/functions" unless version["id"]
@@ -182,7 +181,9 @@ module Bitfab
       rescue
         result.to_s
       end
-      @http_client.send_internal_trace(version["id"], {"result" => output, "inputs" => inputs, "rawCollector" => collector, "source" => "ruby-sdk"})
+      if should_record? && !active_span_content_withheld?
+        @http_client.send_internal_trace(version["id"], {"result" => output, "inputs" => inputs, "rawCollector" => collector, "source" => "ruby-sdk"})
+      end
       result
     end
 

data/lib/bitfab/client.rb

@@ -49,14 +49,10 @@ module Bitfab
       @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
-      # env loaded) can't poison a later one.
       @resolved_api_key = nil
       @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, timeout:)
-      @http_client.configure_simulation_plan(enabled: simulation_plan)
+      @http_client.configure_simulation_plan(enabled: simulation_plan, start: !api_key_resolved_lazily?)
       @datasets = Datasets.new(@http_client)
       @assertion_categories = AssertionCategories.new(@http_client)
       @traces = Traces.new(@http_client)
@@ -381,11 +377,10 @@ module Bitfab
       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,
       result_recorded_elsewhere: 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 should_record?
 
+      @http_client.refresh_simulation_plan
+
       resolved_capture_when = capture_when.to_s
       unless ["always", "nested"].include?(resolved_capture_when)
         Bitfab.warn_once(
@@ -513,6 +508,12 @@ module Bitfab
       completed = false
       @http_client.trace_completion.start(trace_id, span_id)
       recording_override = managed_root ? {} : nil
+      content_withheld = lambda do
+        instrumentation = recording_override&.dig(:instrumentation) || ((surface == :trace) ? "trace" : "span")
+        content_off_by_simulation_plan?(
+          trace_id:, trace_function_key:, span_name: recording_override&.dig(:name) || span_name, instrumentation:
+        ) || (!parent_span_id.nil? && simulation_plan_unavailable?(instrumentation:))
+      end
 
       complete = lambda do |final_result, final_error|
         # Never crash the host app due to span building/sending. Idempotent:
@@ -528,9 +529,7 @@ module Bitfab
           recorded_error = final_error
           recorded_name = recording_override&.dig(:name) || span_name
           recorded_instrumentation = recording_override&.dig(:instrumentation) || ((surface == :trace) ? "trace" : "span")
-          content_off = content_off_by_simulation_plan?(
-            trace_id:, trace_function_key:, span_name: recorded_name, instrumentation: recorded_instrumentation
-          ) || (!parent_span_id.nil? && simulation_plan_unavailable?(instrumentation: recorded_instrumentation))
+          content_off = content_withheld.call
           if finalize && final_error.nil? && !(content_off && !result_recorded_elsewhere)
             begin
               recorded_result = finalize.call(final_result)
@@ -613,7 +612,7 @@ module Bitfab
             span_id:,
             explicit_span_receiver:,
             explicit_span_method_name: explicit_span_method_name&.to_sym,
-            surface:, recording_override:
+            surface:, recording_override:, content_withheld:
           ) do
             invoke = -> { around_body ? around_body.call { yield } : yield }
             result = if explicit_span_receiver
@@ -704,6 +703,18 @@ module Bitfab
       nil
     end
 
+    def active_span_content_withheld?
+      return true if Subtree.current_session&.current_span_content_withheld?
+
+      SpanContext.current&.dig(:content_withheld)&.call == true
+    rescue
+      true
+    end
+
+    def api_key_resolved_lazily?
+      @api_key_config.respond_to?(:call) || @api_key_config.to_s.strip.empty?
+    end
+
     # Resolve the API key lazily, the first time a span actually needs it.
     #
     # The key is intentionally NOT read at construction: a client built when a
@@ -955,7 +966,7 @@ module Bitfab
         span_data["error_source"] = "code"
       end
       span_data["contexts"] = contexts if contexts&.any?
-      span_data["prompt"] = prompt if prompt
+      span_data["prompt"] = prompt if prompt && !content_off
 
       raw_span = {
         "id" => span_id,

data/lib/bitfab/connection_pool.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+require "net/http"
+
+module Bitfab
+  class ConnectionPool
+    KEEP_ALIVE_TIMEOUT_SECONDS = 30
+    MAX_IDLE_CONNECTIONS = 64
+
+    IdleConnection = Struct.new(:http, :released_at)
+    private_constant :IdleConnection
+
+    @registry = ObjectSpace::WeakMap.new
+    @registry_mutex = Mutex.new
+
+    class << self
+      def register(pool)
+        @registry_mutex.synchronize { @registry[pool] = true }
+      end
+
+      def close_all
+        pools = @registry_mutex.synchronize { @registry.keys }
+        pools.each(&:close)
+      end
+    end
+
+    def initialize
+      @pid = Process.pid
+      @mutex = Mutex.new
+      @idle = {}
+      self.class.register(self)
+    end
+
+    def with_connection(uri, timeout)
+      key = [uri.scheme, uri.host, uri.port]
+      owner_pid = Process.pid
+      reusable = false
+      http = checkout(key, uri, timeout)
+      response = yield http
+      reusable = true
+      response
+    ensure
+      release(key, http, owner_pid, reusable) if http
+    end
+
+    def close
+      drain.each { |http| finish(http) }
+    end
+
+    def idle_count
+      reset_after_fork
+      @mutex.synchronize { @idle.sum { |_key, connections| connections.size } }
+    end
+
+    private
+
+    def checkout(key, uri, timeout)
+      http = take_idle(key) || connect(uri, timeout)
+      http.read_timeout = timeout
+      http.write_timeout = timeout
+      http
+    end
+
+    def connect(uri, timeout)
+      http = Net::HTTP.new(uri.host, uri.port)
+      http.use_ssl = uri.scheme == "https"
+      http.open_timeout = timeout
+      http.read_timeout = timeout
+      http.write_timeout = timeout
+      http.keep_alive_timeout = KEEP_ALIVE_TIMEOUT_SECONDS
+      http.start
+    end
+
+    def take_idle(key)
+      reset_after_fork
+      now = monotonic_now
+      expired = []
+      http = @mutex.synchronize do
+        connections = @idle[key] || []
+        stale, fresh = connections.partition { |idle| now - idle.released_at >= KEEP_ALIVE_TIMEOUT_SECONDS }
+        expired.concat(stale.map(&:http))
+        @idle[key] = fresh
+        fresh.pop&.http
+      end
+      expired.each { |stale_http| finish(stale_http) }
+      http
+    end
+
+    def release(key, http, owner_pid, reusable)
+      return unless owner_pid == Process.pid
+
+      unless reusable && http.started?
+        finish(http)
+        return
+      end
+
+      reset_after_fork
+      pooled = @mutex.synchronize do
+        next false if @idle.sum { |_key, connections| connections.size } >= MAX_IDLE_CONNECTIONS
+
+        (@idle[key] ||= []) << IdleConnection.new(http, monotonic_now)
+        true
+      end
+      finish(http) unless pooled
+    end
+
+    def drain
+      return [] if reset_after_fork
+
+      @mutex.synchronize do
+        connections = @idle.values.flatten.map(&:http)
+        @idle = {}
+        connections
+      end
+    end
+
+    def reset_after_fork
+      return false if @pid == Process.pid
+
+      @pid = Process.pid
+      @mutex = Mutex.new
+      @idle = {}
+      true
+    end
+
+    def finish(http)
+      http.finish if http.started?
+    rescue
+      nil
+    end
+
+    def monotonic_now
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+  end
+end

data/lib/bitfab/http_client.rb

@@ -6,6 +6,7 @@ require "uri"
 
 require_relative "git_state"
 require_relative "compress"
+require_relative "connection_pool"
 require_relative "constants"
 require_relative "serialize"
 require_relative "transport"
@@ -37,6 +38,7 @@ module Bitfab
       @api_key = api_key
       @service_url = (service_url || DEFAULT_SERVICE_URL).chomp("/")
       @timeout = timeout
+      @connections = ConnectionPool.new
       @transport_mutex = Mutex.new
       # Eager for the same reason the transport's mutex is: `||=` racing between
       # the export threads recording acks would hand each a different mutex.
@@ -53,9 +55,9 @@ module Bitfab
       @close_condition = ConditionVariable.new
     end
 
-    def configure_simulation_plan(enabled: true)
+    def configure_simulation_plan(enabled: true, start: true)
       @simulation_plan = SimulationPlan.new(self, enabled:)
-      @simulation_plan.refresh
+      @simulation_plan.refresh if start
       @simulation_plan
     end
 
@@ -80,7 +82,7 @@ module Bitfab
     end
 
     def get_simulation_plan
-      get("/api/sdk/sim-plan", timeout: SimulationPlan::READ_TIMEOUT, connection_close: true)
+      get("/api/sdk/sim-plan", timeout: SimulationPlan::READ_TIMEOUT)
     end
 
     # Flush and permanently close this client's tracing transport.
@@ -112,6 +114,7 @@ module Bitfab
         shutdown = transport.nil? || transport.shutdown(remaining)
         result = released && shutdown
       ensure
+        @connections.close
         @transport_mutex.synchronize do
           @close_result = result
           @close_condition.broadcast
@@ -157,16 +160,11 @@ module Bitfab
       last_error = nil
 
       max_retries.times do |attempt|
-        http = Net::HTTP.new(uri.host, uri.port)
-        http.use_ssl = uri.scheme == "https"
-        http.open_timeout = request_timeout
-        http.read_timeout = request_timeout
-
         req = Net::HTTPGenericRequest.new(method, true, true, uri.request_uri, headers)
         req["Content-Encoding"] = prepared.content_encoding if prepared.content_encoding
         req.body = prepared.body
 
-        response = http.request(req)
+        response = @connections.with_connection(uri, request_timeout) { |http| http.request(req) }
 
         unless response.is_a?(Net::HTTPSuccess)
           raise Net::HTTPError.new("HTTP #{response.code}: #{response.body}", response)
@@ -216,19 +214,12 @@ module Bitfab
 
     # Make a GET request to the Bitfab API.
     # Returns parsed JSON response hash.
-    def get(endpoint, timeout: nil, connection_close: false)
+    def get(endpoint, timeout: nil)
       uri = URI("#{@service_url}#{endpoint}")
       request_timeout = timeout || @timeout
 
-      http = Net::HTTP.new(uri.host, uri.port)
-      http.use_ssl = uri.scheme == "https"
-      http.open_timeout = request_timeout
-      http.read_timeout = request_timeout
-
-      # request_uri (not path) so any query string on the endpoint survives.
       req = Net::HTTP::Get.new(uri.request_uri, headers)
-      req["Connection"] = "close" if connection_close
-      response = http.request(req)
+      response = @connections.with_connection(uri, request_timeout) { |http| http.request(req) }
 
       unless response.is_a?(Net::HTTPSuccess)
         raise Net::HTTPError.new("HTTP #{response.code}: #{response.body}", response)
@@ -742,5 +733,6 @@ module Bitfab
   at_exit do
     SimulationPlan.release_all(2, stop: true)
     Transport.shutdown_trace_transports(2)
+    ConnectionPool.close_all
   end
 end

data/lib/bitfab/simulation_plan.rb

@@ -2,7 +2,7 @@
 
 module Bitfab
   class SimulationPlan
-    CONTENT_KEYS = %w[input input_meta output output_meta input_serialized output_serialized].freeze
+    CONTENT_KEYS = %w[input input_meta output output_meta input_serialized output_serialized prompt].freeze
     FRAMEWORKS = %w[openai-agents langgraph claude-agent-sdk vercel-ai].freeze
     READ_TIMEOUT = 5.0
     REFRESH_SECONDS = 60.0
@@ -54,6 +54,7 @@ module Bitfab
       @draining = Set.new
       @stopped = false
       @first_read_finished = false
+      @first_read_deadline = nil
       self.class.register(self)
     end
 
@@ -144,14 +145,13 @@ module Bitfab
       @mutex.synchronize { @content_off.nil? && !reading? }
     end
 
-    def wait_for_first_read(timeout = READ_TIMEOUT)
-      reset_after_fork
+    def wait_for_first_read
+      refresh
       return if disabled?
 
-      deadline = now + [timeout, 0].max
       @mutex.synchronize do
-        while reading? && @thread
-          remaining = deadline - now
+        while reading? && @thread && @first_read_deadline
+          remaining = @first_read_deadline - now
           break unless remaining.positive?
           @condition.wait(@mutex, remaining)
         end
@@ -229,6 +229,7 @@ module Bitfab
     end
 
     def start_reader(delay)
+      @first_read_deadline ||= now + delay + READ_TIMEOUT
       @thread = Thread.new { read_until_loaded(delay) }
       @thread.name = "bitfab-sim-plan"
     rescue ThreadError
@@ -302,6 +303,7 @@ module Bitfab
       @thread = nil
       @held = []
       @draining = Set.new
+      @first_read_deadline = nil unless @first_read_finished
     end
   end
 end

data/lib/bitfab/span_context.rb

@@ -150,10 +150,11 @@ 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, surface: nil, recording_override: nil)
+    def with_span(trace_id:, span_id:, explicit_span_receiver: nil, explicit_span_method_name: nil, surface: nil, recording_override: nil, content_withheld: nil)
       entry = {trace_id:, span_id:}
       entry[:surface] = surface if surface
       entry[:recording_override] = recording_override if recording_override
+      entry[:content_withheld] = content_withheld if content_withheld
       if explicit_span_receiver
         entry[:explicit_span_receiver] = explicit_span_receiver
         entry[:explicit_span_method_name] = explicit_span_method_name

data/lib/bitfab/subtree.rb

@@ -12,6 +12,7 @@ module Bitfab
     DEFAULT_MAX_DEPTH = 30
     DEFAULT_MAX_SPANS = 2048
     DEFAULT_MAX_CAPTURED_SUBTREE_SPANS = 512
+    LIMIT_ORDER = %i[max_depth max_spans max_captured_subtree_spans].freeze
     SESSION_KEY = :bitfab_subtree_session
     DISPATCHER_KEY = :bitfab_subtree_dispatcher
     FIBER_CONTEXT_KEY = :bitfab_subtree_fiber_context
@@ -444,7 +445,8 @@ module Bitfab
         @method_metadata = Hash.new { |owners, owner| owners[owner] = {} }
         @spans_used = 1
         @captured_subtree_count = 0
-        @truncated = false
+        @limits_fired = Set.new
+        @dropped_spans = 0
         @exhausted = false
         @drained = false
         @deferred_frames = []
@@ -478,7 +480,7 @@ module Bitfab
         unless @deferred
           @fiber_parents.clear
           finish_open_spans
-          warn_if_truncated
+          record_truncation
         end
       end
 
@@ -491,7 +493,7 @@ module Bitfab
       end
 
       def handle(event)
-        return if @drained
+        return count_drained_call(event) if @drained
 
         if event.event == :raise
           record_error(event)
@@ -523,16 +525,21 @@ module Bitfab
       def reserve_node
         stack = current_stack(Fiber.current.object_id)
         if captured_depth(stack) >= @max_depth
-          @truncated = true
+          drop(:max_depth)
           return false
         end
         return true if take_slot(discovered: false)
 
-        @truncated = true
+        drop(:max_spans)
         @exhausted = true
         false
       end
 
+      def current_span_content_withheld?
+        entry = @stacks[Fiber.current.object_id]&.reverse_each&.find { |candidate| candidate.is_a?(OpenSpan) }
+        entry&.span_data&.dig("content_off_by_simulation_plan") == true
+      end
+
       def configured_entry(name:, type:, args:, kwargs:, test_run_id:)
         return unless reserve_node
 
@@ -624,11 +631,13 @@ module Bitfab
         fiber_id = Fiber.current.object_id
         if @suppressed_depths[fiber_id].to_i.positive?
           @suppressed_depths[fiber_id] += 1
+          count_dropped_call(event, resolved_path)
           return
         end
 
         if @exhausted
           suppress(fiber_id)
+          count_dropped_call(event, resolved_path)
           return
         end
 
@@ -638,13 +647,13 @@ module Bitfab
 
         stack = current_stack(fiber_id)
         if captured_depth(stack) >= @max_depth
-          @truncated = true
+          drop(:max_depth)
           suppress(fiber_id)
           return
         end
         content_off = content_off?(metadata.name)
         if (limit = limit_reached(discovered: true, content_off:))
-          @truncated = true
+          drop(limit)
           stop_discovery(fiber_id) unless limit == :max_captured_subtree_spans && content_off_names?
           return
         end
@@ -927,9 +936,35 @@ module Bitfab
         false
       end
 
-      def warn_if_truncated
-        return unless @truncated
+      def drop(limit)
+        @limits_fired << limit
+        @dropped_spans += 1
+      end
+
+      def count_dropped_call(event, resolved_path)
+        metadata = metadata_for(event, resolved_path)
+        return if metadata.excluded || (!@include_wrappers && metadata.passthrough)
+
+        @dropped_spans += 1
+      end
+
+      def count_drained_call(event)
+        return unless event.event == :call
+        return if skip_boundary_call?(event)
 
+        resolved_path = owned_path(event.path)
+        count_dropped_call(event, resolved_path) if resolved_path
+      rescue Exception # rubocop:disable Lint/RescueException
+        nil
+      end
+
+      def record_truncation
+        return if @limits_fired.empty?
+
+        CurrentTrace.new(@trace_id).set_metadata(
+          "bitfab.truncated_by" => LIMIT_ORDER.select { |limit| @limits_fired.include?(limit) }.join(","),
+          "bitfab.dropped_spans" => @dropped_spans
+        )
         Bitfab.warn_once(
           "subtree-truncated:#{@trace_function_key}",
           "'#{@trace_function_key}' hit a subtree capture limit (max_depth=#{@max_depth}, " \

data/lib/bitfab/version.rb

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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bitfab
 version: !ruby/object:Gem::Version
-  version: 0.57.6
+  version: 0.57.7
 platform: ruby
 authors:
 - Harvest Team
@@ -161,6 +161,7 @@ files:
 - lib/bitfab/client.rb
 - lib/bitfab/commit_ref.rb
 - lib/bitfab/compress.rb
+- lib/bitfab/connection_pool.rb
 - lib/bitfab/constants.rb
 - lib/bitfab/datasets.rb
 - lib/bitfab/db_snapshot.rb