Class: Bitfab::Client

Inherits:
Object
  • Object
show all
Defined in:
lib/bitfab/client.rb

Constant Summary collapse

SPAN_TYPES = 
%w[llm agent function guardrail handoff custom].freeze
UUID_PATTERN = 
/\A[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\z/i
MOCK_REPLAY_MISS =

Sentinel returned by check_mock_replay when this span should run real code (no mock active, wrong strategy, or no matching historical entry). Using a sentinel rather than nil/false avoids confusing legitimate mocked outputs (which may themselves be nil or false).

Object.new.freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(api_key: nil, service_url: nil, enabled: true, strict: false) ⇒ Client

Returns a new instance of Client.



32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
 
# File 'lib/bitfab/client.rb', line 32

def initialize(api_key: nil, service_url: nil, enabled: true, strict: false)
  @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.
  @explicitly_enabled = enabled
  @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)
  @pending_span_threads = {}
  @pending_span_mutex = Mutex.new
  # 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 = []
end

Instance Attribute Details

#service_url ⇒ Object (readonly)

Returns the value of attribute service_url.



30
31
32
 
# File 'lib/bitfab/client.rb', line 30

def service_url
  @service_url
end

Instance Method Details

#api_key ⇒ Object

The configured API key (a proc is resolved on read). Reflects what was passed to the client; the ENV fallback applied during actual tracing is not surfaced here, and reading this never warns.



56
57
58
 
# File 'lib/bitfab/client.rb', line 56

def api_key
  @api_key_config.respond_to?(:call) ? @api_key_config.call : @api_key_config
end

#clear_mock_overrides ⇒ void

This method returns an undefined value.

Remove all overrides registered via register_mock_override.



185
186
187
188
 
# File 'lib/bitfab/client.rb', line 185

def clear_mock_overrides
  @mock_overrides.clear
  nil
end

#enabled ⇒ Object

Effective tracing state, evaluated lazily: enabled only when not explicitly disabled AND a key resolves. Reading this resolves the key (and may emit the one-time empty-key warning), exactly as the first traced call would.



64
65
66
 
# File 'lib/bitfab/client.rb', line 64

def enabled
  tracing_enabled?
end

#execute_span(trace_function_key:, span_name:, span_type:, function_name:, args:, kwargs:, mock_on_replay: false) ⇒ Object

Execute a block inside a span context, sending trace data on completion. Called by Traceable, not intended for direct use.



230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
 
# File 'lib/bitfab/client.rb', line 230

def execute_span(trace_function_key:, span_name:, span_type:, function_name:, args:, kwargs:,
  mock_on_replay: 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?

  # Span setup runs before the user's block. Tracing is a side-channel, so
  # if anything here raises (id generation, trace-state bookkeeping, a
  # malformed replay mock tree) the user's method must still run. On failure
  # we clean up any partially registered trace state, warn once, and run the
  # block untraced. `trace_id` is declared out here so the rescue can clean
  # it up; the other locals stay visible to the real path below.
  trace_id = nil
  span_id = nil
  parent_span_id = nil
  is_root_span = nil
  started_at = nil
  resolved_test_run_id = nil
  resolved_input_source_span_id = nil
  begin
    parent = SpanContext.current
    replay_ctx = ReplayContext.current
    trace_id = parent ? parent[:trace_id] : (replay_ctx&.dig(:trace_id) || SecureRandom.uuid)
    span_id = SecureRandom.uuid
    parent_span_id = parent&.dig(:span_id)
    is_root_span = parent_span_id.nil?
    started_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")
    resolved_test_run_id = replay_ctx&.dig(:test_run_id)
    resolved_input_source_span_id = replay_ctx&.dig(:input_source_span_id)
    resolved_input_source_trace_id = replay_ctx&.dig(:input_source_trace_id)

    # Register trace state for root spans
    if is_root_span && !TraceState.get(trace_id)
      TraceState.create(
        trace_id,
        test_run_id: resolved_test_run_id,
        input_source_trace_id: resolved_input_source_trace_id
      )
    end

    if is_root_span
      @pending_span_mutex.synchronize { @pending_span_threads[trace_id] = [] }
    end

    # 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
    # line up with `build_mock_tree`'s sequential numbering for the same
    # (key, name) pair. Different (key, name) pairs have independent
    # counters: they cannot shift each other.
    call_index = advance_mock_counter(replay_ctx, trace_function_key, span_name, is_root_span:)
    if call_index
      mocked_output = check_mock_replay(
        replay_ctx, trace_function_key, span_name, call_index,
        span_type:, args:, kwargs:, mock_on_replay:
      )
      if mocked_output != MOCK_REPLAY_MISS
        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: resolved_test_run_id,
          input_source_span_id: resolved_input_source_span_id
        )
        return mocked_output
      end
    end
  rescue
    # Clean up any trace state this partial setup registered so it does not
    # leak.
    if trace_id
      TraceState.delete(trace_id)
      @pending_span_mutex.synchronize { @pending_span_threads.delete(trace_id) }
    end
    # During replay (a controlled eval) a setup failure must surface, not
    # silently run the block untraced: swallowing it would execute real code
    # with real side effects and skew the mock call counter, defeating the
    # replay. The never-crash fallback is for production hosts only.
    raise if ReplayContext.current

    Bitfab.warn_once(
      "span-setup:#{trace_function_key}",
      "span setup failed for '#{trace_function_key}'; this call runs untraced. " \
      "Your method still executes and returns normally."
    )
    return yield
  end

  result = nil
  error = nil
  span_contexts = nil
  span_prompt = nil
  finalized = false

  finalize = lambda do |final_result, final_error|
    # Never crash the host app due to span building/sending. Idempotent:
    # only the first call sends the span. Subsequent calls (e.g. from the
    # enumerator wrapper after iteration completes) are no-ops.
    next if finalized
    finalized = true

    begin
      ended_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")

      span_thread = send_span(
        trace_function_key:,
        trace_id:,
        span_id:,
        parent_span_id:,
        span_name:,
        span_type:,
        function_name:,
        contexts: span_contexts,
        prompt: span_prompt,
        args:,
        kwargs:,
        result: final_result,
        error: final_error,
        started_at:,
        ended_at:,
        test_run_id: resolved_test_run_id,
        input_source_span_id: resolved_input_source_span_id
      )

      if is_root_span
        pending = @pending_span_mutex.synchronize { @pending_span_threads.delete(trace_id) || [] }
        pending << span_thread if span_thread
        pending.each { |t| t.join(5) }

        # Built AFTER the wrapped method finished (finalize runs at root
        # span end), so :accessed reflects whether customer code obtained
        # the branch URL during this item. nil (key omitted) when no
        # lease was attached, so the server can distinguish "no branch"
        # from "branch ignored".
        lease = replay_ctx&.dig(:db_branch_lease)
        db_snapshot_usage = if lease
          {
            neon_branch_id: lease["neonBranchId"],
            snapshot_timestamp: lease["snapshotTimestamp"],
            source_trace_id: replay_ctx[:source_bitfab_trace_id],
            accessed: replay_ctx[:db_snapshot_accessed] == true
          }
        end

        completion_thread = send_trace_completion(
          trace_function_key:,
          trace_id:,
          started_at:,
          ended_at:,
          db_snapshot_usage:
        )

        # In replay, persistence is correctness: the replay runner joins
        # these threads before calling complete_replay, or the server's
        # trace-ID mapping races the uploads and every item's trace_id
        # comes back nil. The 5s join above is best-effort only; this
        # hands the full set (span uploads + trace completion) to the
        # runner. No-op outside replay, where sends stay fire-and-forget.
        persistence = ReplayContext.current&.dig(:pending_persistence)
        if persistence
          persistence.concat(pending)
          persistence << completion_thread if completion_thread
        end
      else
        @pending_span_mutex.synchronize do
          @pending_span_threads[trace_id] << span_thread if span_thread && @pending_span_threads.key?(trace_id)
        end
      end
    rescue Exception # rubocop:disable Lint/RescueException
      # Silently ignore: user's result/exception takes priority
      # Catches Exception (not just StandardError) to handle SystemStackError
      # from deeply nested serialization
    end
  end

  begin
    SpanContext.with_span(trace_id:, span_id:) do
      result = yield
    ensure
      # Capture contexts before the span context is popped
      span_contexts = SpanContext.current&.dig(:contexts)
      span_prompt = SpanContext.current&.dig(:prompt)
    end
  rescue => e
    error = e.message
    finalize.call(result, error)
    raise
  end

  # If the wrapped block returned an Enumerator (lazy iteration via
  # `enum_for`, `to_enum`, `Enumerator.new`, `[...].lazy.map(...)`, etc.),
  # the work hasn't actually run yet: the values are produced as the
  # caller iterates. Without special handling we'd close the span here
  # with `result == <the Enumerator object>`, and any nested `bitfab_span`
  # calls inside the enumerator body would see an empty span stack and
  # post their own root traces, fragmenting one logical workflow.
  #
  # Instead, hand the caller a wrapping Enumerator whose body restores
  # the parent span stack on the iterating fiber, drives the source,
  # collects yielded values as the span output, and finalizes the span
  # once iteration completes (or errors).
  #
  # Limitation: when the source enumerator itself runs its body in a
  # separate fiber (e.g. `Enumerator.new { |y| ... }` or `enum_for(...)`
  # without a block), nested `bitfab_span` calls inside that body fiber
  # still see an empty stack because `Thread.current[STACK_KEY]` is
  # fiber-local. Lazy chains over collections (`.lazy.map`) and ordinary
  # `each` callbacks DO run in the iterating fiber and nest correctly.
  if result.is_a?(Enumerator)
    return wrap_enumerator(result, trace_id:, span_id:, finalize:)
  end

  finalize.call(result, error)
  result
end

#get_function(trace_function_key) ⇒ BitfabFunction

Get a function wrapper bound to a specific trace function key.

This provides a fluent API for binding a trace_function_key once and then wrapping multiple methods or classes with that key. Mirrors client.get_function(key) in the Python SDK and client.getFunction(key) in the TypeScript SDK.

Examples:

fn = Bitfab.client.get_function("order-processing")
fn.wrap(OrderService, :process_order, type: "function")
fn.wrap(OrderService, :validate_order, type: "guardrail")

Parameters:

  • trace_function_key (String)

Returns:



204
205
206
 
# File 'lib/bitfab/client.rb', line 204

def get_function(trace_function_key)
  BitfabFunction.new(self, trace_function_key)
end

#get_trace_span(trace_id, id: nil, name: nil, occurrence: "last") ⇒ Object

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 Integer index.

Raises:

  • (ArgumentError) 


212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
 
# File 'lib/bitfab/client.rb', line 212

def get_trace_span(trace_id, id: nil, name: nil, occurrence: "last")
  validate_trace_id(trace_id)
  raise ArgumentError, "Provide exactly one of id or name" if id.nil? == name.nil?
  validate_span_id(id) unless id.nil?
  if !name.nil? && (!name.is_a?(String) || name.empty?)
    raise ArgumentError, "name must be a non-empty string"
  end

  valid_occurrence = %w[first last].include?(occurrence) || (occurrence.is_a?(Integer) && occurrence >= 0)
  unless valid_occurrence
    raise ArgumentError, 'occurrence must be "first", "last", or a non-negative integer'
  end

  @http_client.get_trace_span(trace_id, id:, name:, occurrence:)
end

#register_mock_override(*positional, match: nil, value: VALUE_UNSET) ⇒ void

This method returns an undefined value.

Register a mock override applied to every subsequent replay on this client, so downstream real code runs against a value you supply for the matched span. Instance-scoped (no global state); call clear_mock_overrides to reset. Per-call replay(mock_override:) overrides take precedence, and both take precedence over the base mock strategy.

Accepts either the (match, value) positional form or the keyword form (a { match:, value: } hash is also accepted). match must be callable; value is either a flat value injected directly or a callable invoked with the ctx hash.

Examples:

keyword form (primary), flat value

client.register_mock_override(
  match: ->(node) { node[:trace_function_key] == "classify-intent" },
  value: {label: "refund"}
)

positional form (equivalent), callable value

client.register_mock_override(
  ->(node) { node[:trace_function_key] == "classify-intent" },
  ->(ctx) { {label: "refund", inputs: ctx[:inputs]} }
)

Parameters:

  • positional (Array) —

    either (match, value) or a single { match:, value: } hash

  • match (#call, nil) (defaults to: nil) —

    keyword form matcher

  • value (Object, #call, nil) (defaults to: VALUE_UNSET) —

    keyword form injected value or producer



150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
 
# File 'lib/bitfab/client.rb', line 150

def register_mock_override(*positional, match: nil, value: VALUE_UNSET)
  if match.nil? && value.equal?(VALUE_UNSET)
    if positional.length == 1 && positional[0].is_a?(Hash)
      override = positional[0]
      match = override[:match]
      # Distinguish an omitted :value from an explicit `value: nil`: the
      # key's presence, not a nil read, is what marks it provided.
      value = override.key?(:value) ? override[:value] : VALUE_UNSET
    elsif positional.length == 2
      match, value = positional
    end
  end

  unless match.respond_to?(:call)
    raise ArgumentError,
      "register_mock_override requires a callable match. Pass (match, value) " \
      "positionally, as keywords (match:, value:), or as a { match:, value: } hash. " \
      "value may be a flat value or a callable."
  end
  # A forgotten value must not silently inject nil. An explicit nil is a
  # legitimate injected value and passes this guard.
  if value.equal?(VALUE_UNSET)
    raise ArgumentError,
      "register_mock_override requires a value (the second argument, or " \
      "value:). It may be a flat value or a callable; pass value: nil " \
      "explicitly to inject nil."
  end

  @mock_overrides << {match:, value:}
  nil
end

#replay(receiver, method_name, trace_function_key:, limit: nil, trace_ids: nil, max_concurrency: 10, name: nil, code_change_description: nil, code_change_files: nil, experiment_group_id: nil, dataset_id: nil, mock: "marked", adapt_inputs: nil, mock_override: nil, environment: nil, on_progress: nil) ⇒ Hash

Replay historical traces through a method and create a test run.

Parameters:

  • receiver (Object, Class) —

    an instance for instance methods, or a Class for class methods

  • method_name (Symbol) —

    the method to replay

  • trace_function_key (String) —

    the trace function key for this method

  • limit (Integer, nil) (defaults to: nil) —

    maximum number of traces to replay (default: 5). Mutually exclusive with trace_ids: an explicit ID list already determines how many traces replay, so passing both raises.

  • trace_ids (Array<String>, nil) (defaults to: nil) —

    optional list of trace IDs to replay (max 100)

  • max_concurrency (Integer, nil) (defaults to: 10) —

    max threads for parallel replay (default: 10)

  • code_change_description (String, nil) (defaults to: nil) —

    optional rationale for the code change being tested in this replay (stored on the experiment)

  • code_change_files (Array<Hash>, nil) (defaults to: nil) —

    optional list of edited files, each as { path:, before:, after: } (use "" for new/deleted files)

  • experiment_group_id (String, nil) (defaults to: nil) —

    optional UUID grouping multiple replay runs into a single experiment batch

  • dataset_id (String, nil) (defaults to: nil) —

    optional UUID of the dataset this replay runs against, stored on the resulting experiment for durable attribution

  • mock (String) (defaults to: "marked") —

    mock strategy for child spans: "marked" (default), "none", or "all". "marked" only mocks spans declared with mock_on_replay: true; "all" mocks every child span.

  • adapt_inputs (#call, nil) (defaults to: nil) —

    optional hook to reshape recorded inputs onto the method's current signature when its shape changed after the traces were captured. Receives (args, kwargs, ctx) where ctx is { trace_id:, source_span_id: }, and returns [new_args, new_kwargs].

  • mock_override (Hash, Array<Hash>, nil) (defaults to: nil) —

    optional selective mock override(s), each a { match:, value: } hash. match is a callable: match.call(node) selects spans to substitute (node is { trace_function_key:, span_name:, type:, original_span_id: }). value is EITHER a flat value injected directly OR a callable invoked with { node:, inputs:, get_original_output: }; its result becomes the span's output (full replacement), so downstream real code runs against it. The first matching override wins. Per-call overrides take precedence over those registered via register_mock_override, and both take precedence over the base mock strategy. The root span is never overridden.

  • on_progress (#call, nil) (defaults to: nil) —

    optional callback invoked once per item as it finishes, with a running-totals hash { completed:, total:, succeeded:, errored: }. Use it to render replay progress (e.g. a terminal progress bar). Replay does not know pass/fail yet, so the totals only distinguish items whose method ran (:succeeded) from items that raised (:errored). A raising callback never crashes the run.

Returns:

  • (Hash) —

    with :items, :test_run_id, :test_run_url



110
111
112
113
114
115
116
117
118
119
120
 
# File 'lib/bitfab/client.rb', line 110

def replay(receiver, method_name, trace_function_key:, limit: nil, trace_ids: nil, max_concurrency: 10,
  name: nil, code_change_description: nil, code_change_files: nil, experiment_group_id: nil, dataset_id: nil, mock: "marked",
  adapt_inputs: nil, mock_override: nil, environment: nil, on_progress: nil)
  Replay.run(
    self, receiver, method_name,
    trace_function_key:, limit:, trace_ids:, name:, max_concurrency:,
    code_change_description:, code_change_files:, experiment_group_id:, dataset_id:, mock:, adapt_inputs:,
    mock_override:, environment:,
    on_progress:
  )
end