bitfab 0.16.1 → 0.17.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0dfedcfce509695e2bb260c885fcfa460aaab8d9c89421379cd7572cd280eb95
-  data.tar.gz: e8aaba21b08b4a53c9f89319e7c7a9d9f0a911f0be657058d80fc84add6e4f6e
+  metadata.gz: 60fa4998513b5fa9f67950bf4abbbc253b324955732e3726531ca47856544ddb
+  data.tar.gz: ad45f27ff4c50a4a75329f48e6ba06f252eb74e425b1d19686b55f6c780ae6a0
 SHA512:
-  metadata.gz: 89851a89aae7ad24ac24d03335e3641b2715c5a841dfd76d302effd8ca07723d41cf0de6103884795fa96337c79671ed9ca5776a5b71c1aa805c3a87621a59b5
-  data.tar.gz: 3fb0c600868f499787e25dc424ce7643ee383f8da1e72fb99ec06a394949fc5ca5f8180e9f40b45791da8f2d4d1b562422bc4921d98643dbefb8c9d03e754893
+  metadata.gz: 1dd29ae11022f8cc0ad80b3d59e37d809843b85ab9cb48e0623afd01ee1a2d6b4e99e8ba91c996b7cad1acb589757ccf4a0291673a28d7a54eccffa098d5239e
+  data.tar.gz: efaf095342fee4995aa8d45c69773c7d268f5eb23e358606292c7dd3ba11a3958da7416606025f18c8c11b43c295b62ef0bcd80baf753c0070055ae0d3dcc0f3

data/lib/bitfab/client.rb

@@ -60,11 +60,11 @@ module Bitfab
     # @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,
       code_change_description: nil, code_change_files: nil, experiment_group_id: nil, mock: "none",
-      adapt_inputs: nil)
+      adapt_inputs: nil, environment: nil)
       Replay.run(
         self, receiver, method_name,
         trace_function_key:, limit:, trace_ids:, max_concurrency:,
-        code_change_description:, code_change_files:, experiment_group_id:, mock:, adapt_inputs:
+        code_change_description:, code_change_files:, experiment_group_id:, mock:, adapt_inputs:, environment:
       )
     end
 
@@ -315,6 +315,9 @@ module Bitfab
       if trace_state&.dig(:input_source_trace_id)
         raw_trace["input_source_trace_id"] = trace_state[:input_source_trace_id]
       end
+      if trace_state&.dig(:db_snapshot_ref)
+        raw_trace["db_snapshot_ref"] = trace_state[:db_snapshot_ref]
+      end
 
       payload = {
         "type" => "sdk-function",

data/lib/bitfab/db_snapshot.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Bitfab
+  # Per-trace database snapshot ref capture.
+  #
+  # 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'
+  # +DbSnapshotRef+ wire shape (camelCase key, since the server speaks the same
+  # protocol for every SDK).
+  module DbSnapshot
+    module_function
+
+    # Build a snapshot ref for one trace. Synchronous, no IO.
+    #
+    # 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).
+    #
+    # @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}
+    end
+  end
+end

data/lib/bitfab/http_client.rb

@@ -106,7 +106,7 @@ module Bitfab
     # @param experiment_group_id [String, nil] optional UUID grouping multiple
     #   replay runs into a single experiment batch
     def start_replay(trace_function_key, limit, trace_ids: nil, code_change_description: nil,
-      code_change_files: nil, experiment_group_id: nil)
+      code_change_files: nil, experiment_group_id: nil, include_db_branch_lease: false)
       payload = {
         "traceFunctionKey" => trace_function_key
       }
@@ -117,8 +117,14 @@ module Bitfab
       payload["codeChangeDescription"] = code_change_description unless code_change_description.nil?
       payload["codeChangeFiles"] = normalize_code_change_files(code_change_files) unless code_change_files.nil?
       payload["experimentGroupId"] = experiment_group_id unless experiment_group_id.nil?
-
-      request("/api/sdk/replay/start", payload, timeout: 30)
+      payload["includeDbBranchLease"] = true if include_db_branch_lease
+
+      # When DB branching is on, the server resolves a Neon preview branch per
+      # item (snapshot + restore + poll), which can run several seconds each.
+      # Use a generous timeout so the SDK doesn't give up before a healthy
+      # server finishes.
+      timeout = include_db_branch_lease ? 180 : 30
+      request("/api/sdk/replay/start", payload, timeout:)
     end
 
     # Fetch an external span by ID. Blocking GET request.
@@ -142,6 +148,13 @@ module Bitfab
       request("/api/sdk/replay/complete", {"testRunId" => test_run_id}, timeout: 30)
     end
 
+    # Release a previously-resolved DB branch by deleting its Neon branch.
+    # Blocking call. Idempotent server-side (a missing branch is treated as
+    # already released).
+    def release_db_branch_lease(neon_branch_id)
+      request("/api/sdk/replay/releaseDbBranchLease", {"neonBranchId" => neon_branch_id}, timeout: 30)
+    end
+
     # Send an external trace (fire-and-forget in background thread).
     def send_external_trace(payload)
       merged = payload.merge("sdkVersion" => VERSION)

data/lib/bitfab/replay.rb

@@ -27,7 +27,8 @@ module Bitfab
     # threads (span uploads + trace completion) so the replay runner can join
     # them before complete_replay builds the trace-ID mapping.
     def with_context(test_run_id:, input_source_span_id: nil, input_source_trace_id: nil, trace_id: nil,
-      mock_tree: nil, mock_strategy: nil, pending_persistence: nil)
+      mock_tree: nil, mock_strategy: nil, pending_persistence: nil, db_branch_lease: nil,
+      source_bitfab_trace_id: nil)
       previous = Thread.current[REPLAY_CONTEXT_KEY]
       ctx = {
         test_run_id:,
@@ -41,6 +42,11 @@ module Bitfab
         ctx[:mock_strategy] = mock_strategy || "none"
         ctx[:call_counters] = {}
       end
+      # The per-trace DB branch (resolved server-side) and the Bitfab trace ID
+      # it belongs to ride on the context so ReplayEnvironment can read them
+      # inside the replayed method.
+      ctx[:db_branch_lease] = db_branch_lease if db_branch_lease
+      ctx[:source_bitfab_trace_id] = source_bitfab_trace_id if source_bitfab_trace_id
       Thread.current[REPLAY_CONTEXT_KEY] = ctx
       yield
     ensure
@@ -84,7 +90,7 @@ module Bitfab
     # @return [Hash] with :items, :test_run_id, :test_run_url
     def run(client, receiver, method_name, trace_function_key:, limit: nil, trace_ids: nil, max_concurrency: 10,
       code_change_description: nil, code_change_files: nil, experiment_group_id: nil, mock: "none",
-      adapt_inputs: nil)
+      adapt_inputs: nil, environment: nil)
       unless MOCK_STRATEGIES.include?(mock.to_s)
         raise ArgumentError, "Invalid mock strategy '#{mock}'. Must be one of: #{MOCK_STRATEGIES.join(", ")}"
       end
@@ -105,13 +111,16 @@ module Bitfab
       # the count), so it's omitted from the request entirely.
       effective_limit = trace_ids ? nil : (limit || 5)
 
+      include_db_branch_lease = !environment.nil?
+
       replay_data = http_client.start_replay(
         trace_function_key,
         effective_limit,
         trace_ids:,
         code_change_description:,
         code_change_files:,
-        experiment_group_id:
+        experiment_group_id:,
+        include_db_branch_lease:
       )
       test_run_id = replay_data["testRunId"]
       test_run_url = replay_data["testRunUrl"]
@@ -119,7 +128,7 @@ module Bitfab
 
       result_items = if server_items.any?
         process_items(http_client, server_items, receiver, method_name, test_run_id, max_concurrency, mock.to_s,
-          adapt_inputs)
+          adapt_inputs, include_db_branch_lease)
       else
         []
       end
@@ -192,12 +201,13 @@ module Bitfab
 
     # Process all replay items, optionally in parallel using threads.
     def process_items(http_client, server_items, receiver, method_name, test_run_id, max_concurrency, mock_strategy,
-      adapt_inputs = nil)
+      adapt_inputs = nil, include_db_branch_lease = false)
       concurrency = max_concurrency || server_items.length
 
       if concurrency <= 1
         server_items.map do |item|
-          process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy, adapt_inputs)
+          process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy, adapt_inputs,
+            include_db_branch_lease)
         end
       else
         results_mutex = Mutex.new
@@ -212,7 +222,7 @@ module Bitfab
               break unless item
 
               result = process_single_item(http_client, item, receiver, method_name, test_run_id, mock_strategy,
-                adapt_inputs)
+                adapt_inputs, include_db_branch_lease)
               results_mutex.synchronize { results[idx] = result }
             end
           end
@@ -230,8 +240,15 @@ module Bitfab
     # than propagated, so one bad trace never aborts the whole replay run
     # (mirrors the TypeScript and Python SDKs' per-item rescue).
     def process_single_item(http_client, server_item, receiver, method_name, test_run_id, mock_strategy,
-      adapt_inputs = nil)
+      adapt_inputs = nil, include_db_branch_lease = false)
       metrics = extract_server_item_metrics(server_item)
+      # The server resolves a Neon preview branch per item during /replay/start
+      # (only when include_db_branch_lease was sent). Release it in the +ensure+
+      # below so any raise — span fetch, mock-tree build, or the replayed
+      # method — frees the Neon resource. Items whose source trace had no
+      # snapshot ref, or whose resolve failed server-side, arrive without a
+      # lease (env.active? is false for those).
+      lease = include_db_branch_lease ? server_item["dbBranchLease"] : nil
 
       span = http_client.get_external_span(server_item["externalSpanId"])
       item_data = extract_span_data(span)
@@ -255,7 +272,10 @@ module Bitfab
         mock_strategy:,
         mock_tree:,
         adapt_inputs:,
-        adapt_ctx:
+        adapt_ctx:,
+        db_branch_lease: lease,
+        source_bitfab_trace_id: server_item["traceId"],
+        db_snapshot_ref: server_item["dbSnapshotRef"]
       )
     rescue => e
       warn "Bitfab: replay item for span #{server_item["externalSpanId"]} failed before execution: #{e.message}"
@@ -267,8 +287,22 @@ module Bitfab
         duration_ms: metrics&.dig(:duration_ms),
         tokens: metrics&.dig(:tokens),
         model: metrics&.dig(:model),
-        trace_id: nil
+        trace_id: nil,
+        db_snapshot_ref: server_item["dbSnapshotRef"]
       }
+    ensure
+      release_db_branch_lease(http_client, lease) if lease
+    end
+
+    # Delete the per-item Neon preview branch. Best-effort: a failure is warned
+    # but never raised — the server-side TTL janitor reaps orphans.
+    def release_db_branch_lease(http_client, lease)
+      neon_branch_id = lease["neonBranchId"]
+      return unless neon_branch_id
+
+      http_client.release_db_branch_lease(neon_branch_id)
+    rescue => e
+      warn "Bitfab: failed to release DB branch #{neon_branch_id} (TTL janitor will catch it): #{e.message}"
     end
 
     # Walk the children of a root span tree node depth-first and build a
@@ -346,7 +380,8 @@ module Bitfab
 
     # 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)
+      input_source_trace_id: nil, mock_strategy: "none", mock_tree: nil, adapt_inputs: nil, adapt_ctx: nil,
+      db_branch_lease: nil, source_bitfab_trace_id: nil, db_snapshot_ref: nil)
       args, kwargs = Serialize.deserialize_inputs(item)
 
       fn_result = nil
@@ -365,7 +400,9 @@ module Bitfab
         trace_id: sdk_trace_id,
         mock_tree:,
         mock_strategy:,
-        pending_persistence:
+        pending_persistence:,
+        db_branch_lease:,
+        source_bitfab_trace_id:
       ) do
         # Reshape recorded inputs onto the current signature when an adapter is
         # supplied. Inside the rescue so a raising adapter surfaces on this
@@ -397,7 +434,8 @@ module Bitfab
         duration_ms: metrics[:duration_ms],
         tokens: metrics[:tokens],
         model: metrics[:model],
-        trace_id: sdk_trace_id
+        trace_id: sdk_trace_id,
+        db_snapshot_ref:
       }
     end
   end

data/lib/bitfab/replay_environment.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+module Bitfab
+  # Per-trace environment exposed to customer code during replay.
+  #
+  # The customer instantiates one +ReplayEnvironment+ and passes it to
+  # +client.replay(environment: ...)+. Inside the replayed method they read
+  # +env.database_url+ (and friends) to pick up the per-trace branch URL the
+  # Bitfab service resolved from the source trace's snapshot reference.
+  #
+  # Outside replay, reading +env.database_url+ raises. Customer code uses the
+  # env only on the replay path; live request code keeps reading
+  # +ENV["DATABASE_URL"]+ the normal way.
+  #
+  # Concurrency-safe: the readers resolve through the thread-local replay
+  # context, so each in-flight replay item sees its own per-trace values even
+  # when the SDK runs items across worker threads.
+  #
+  # Internally the resolved per-item state is a DB branch lease (the SDK <->
+  # server protocol term). We expose its useful fields directly here so
+  # customer code never sees the word.
+  class ReplayEnvironment
+    # The per-trace branch URL for the item currently being replayed.
+    # Raises if read outside a replay item.
+    def database_url
+      require_snapshot.fetch(:database_url)
+    end
+
+    # When the per-trace branch URL stops being valid. ISO-8601.
+    def expires_at
+      require_snapshot.fetch(:expires_at)
+    end
+
+    # Deep link to the branch in the provider console, if available.
+    def provider_console_url
+      require_snapshot[:provider_console_url]
+    end
+
+    # True if the branch is read-only. Customer code can use this to skip write
+    # operations during replay when the provider returned a read-only lease.
+    def read_only
+      require_snapshot[:read_only]
+    end
+
+    # The historical trace ID that produced the input for this replay item.
+    def trace_id
+      require_snapshot.fetch(:trace_id)
+    end
+
+    # True when read inside a replay item that has a resolved branch.
+    def active?
+      !read_snapshot.nil?
+    end
+
+    # Non-raising variant for callers that handle the inactive case. Returns a
+    # symbol-keyed hash or nil.
+    def snapshot
+      read_snapshot
+    end
+
+    private
+
+    def read_snapshot
+      ctx = ReplayContext.current
+      return nil unless ctx
+
+      lease = ctx[:db_branch_lease]
+      return nil unless lease
+
+      # Surface the Bitfab trace ID (what the customer sees in the dashboard),
+      # falling back to the external trace ID only if the Bitfab ID is somehow
+      # absent — keeps replays from external sources working until the
+      # source-system path is fully wired.
+      trace_id = ctx[:source_bitfab_trace_id] || ctx[:input_source_trace_id]
+      return nil unless trace_id
+
+      {
+        database_url: lease["databaseUrl"],
+        expires_at: lease["expiresAt"],
+        provider_console_url: lease["providerConsoleUrl"],
+        read_only: lease["readOnly"],
+        trace_id:
+      }
+    end
+
+    def require_snapshot
+      snap = read_snapshot
+      return snap if snap
+
+      raise "ReplayEnvironment accessed outside of a replay item. Pass it to " \
+        "client.replay(environment: ...) and only read it inside the replayed method."
+    end
+  end
+end

data/lib/bitfab/span_context.rb

@@ -133,12 +133,20 @@ module Bitfab
 
     def create(trace_id, test_run_id: nil, input_source_trace_id: nil)
       @states_mutex.synchronize do
-        @states[trace_id] ||= {
-          trace_id:,
-          started_at: Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ"),
-          test_run_id:,
-          input_source_trace_id:
-        }.compact
+        @states[trace_id] ||= begin
+          started_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")
+          {
+            trace_id:,
+            started_at:,
+            test_run_id:,
+            input_source_trace_id:,
+            # 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)
+          }.compact
+        end
       end
     end
 

data/lib/bitfab/version.rb

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

data/lib/bitfab.rb

@@ -3,9 +3,11 @@
 require_relative "bitfab/version"
 require_relative "bitfab/constants"
 require_relative "bitfab/serialize"
+require_relative "bitfab/db_snapshot"
 require_relative "bitfab/span_context"
 require_relative "bitfab/http_client"
 require_relative "bitfab/replay"
+require_relative "bitfab/replay_environment"
 require_relative "bitfab/client"
 require_relative "bitfab/traceable"
 

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bitfab
 version: !ruby/object:Gem::Version
-  version: 0.16.1
+  version: 0.17.0
 platform: ruby
 authors:
 - Harvest Team
@@ -119,8 +119,10 @@ files:
 - lib/bitfab.rb
 - lib/bitfab/client.rb
 - lib/bitfab/constants.rb
+- lib/bitfab/db_snapshot.rb
 - lib/bitfab/http_client.rb
 - lib/bitfab/replay.rb
+- lib/bitfab/replay_environment.rb
 - lib/bitfab/serialize.rb
 - lib/bitfab/span_context.rb
 - lib/bitfab/traceable.rb