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
 
# 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)
  # 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.



54
55
56
 
# File 'lib/bitfab/client.rb', line 54

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.



217
218
219
220
 
# File 'lib/bitfab/client.rb', line 217

def clear_mock_overrides
  @mock_overrides.clear
  nil
end

#close(timeout: 30) ⇒ Boolean

Flush and permanently close this client's tracing transport, releasing its batch worker. Long-running processes that build transient clients should call this; a shared client is closed at process exit.

Parameters:

  • timeout (Numeric) (defaults to: 30) —

    maximum total seconds to wait

Returns:

  • (Boolean) —

    true when everything queued was delivered in time



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

def close(timeout: 30)
  @http_client.close(timeout:)
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.



81
82
83
 
# File 'lib/bitfab/client.rb', line 81

def enabled
  tracing_enabled?
end

#execute_span(trace_function_key:, span_name:, span_type:, function_name:, args:, kwargs:, capture_when: "always", 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.



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
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
 
# File 'lib/bitfab/client.rb', line 262

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

  resolved_capture_when = capture_when.to_s
  unless %w[always nested].include?(resolved_capture_when)
    Bitfab.warn_once(
      "invalid-capture-when:#{trace_function_key}",
      "unknown capture_when value #{capture_when.inspect}; defaulting to \"always\". " \
      "Valid values: \"always\", \"nested\"."
    )
    resolved_capture_when = "always"
  end
  return yield if resolved_capture_when == "nested" && SpanContext.current.nil?

  # 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

    # 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.
    TraceState.delete(trace_id) if trace_id
    # 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")

      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
        # 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"],
            original_trace_id: replay_ctx[:source_bitfab_trace_id],
            # Deprecated wire alias, kept so this SDK still reports usage
            # against servers that predate the rename.
            source_trace_id: replay_ctx[:source_bitfab_trace_id],
            accessed: replay_ctx[:db_snapshot_accessed] == true
          }
        end

        send_trace_completion(
          trace_function_key:,
          trace_id:,
          started_at:,
          ended_at:,
          db_snapshot_usage:
        )
      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

#flush(timeout: 30) ⇒ Boolean

Wait for the spans and traces this client queued to be delivered, without closing it.

Parameters:

  • timeout (Numeric) (defaults to: 30) —

    maximum total seconds to wait

Returns:

  • (Boolean) —

    true when everything queued was delivered in time



73
74
75
 
# File 'lib/bitfab/client.rb', line 73

def flush(timeout: 30)
  @http_client.flush(timeout:)
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:



236
237
238
 
# File 'lib/bitfab/client.rb', line 236

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) 


244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
 
# File 'lib/bitfab/client.rb', line 244

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



182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
 
# File 'lib/bitfab/client.rb', line 182

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: Replay::CODE_CHANGE_UNSET, code_change_files: Replay::CODE_CHANGE_UNSET, experiment_group_id: nil, dataset_id: nil, grader_ids: nil, mock: "marked", adapt_inputs: nil, mock_override: nil, db_branch: 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: Replay::CODE_CHANGE_UNSET) —

    optional rationale for the code change being tested in this replay (stored on the experiment). When supplied without code_change_files, this text is preserved while the SDK automatically captures the files. Pass nil when the replay should have no description; use code_change_files: nil to suppress automatic file capture.

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

    optional list of edited files, each as { path:, before:, after: } (use "" for new/deleted files). Omit to capture automatically, or pass nil to suppress file capture.

  • 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

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

    optional UUIDs of graders attached directly to this experiment, graded as the union with the dataset's runnable graders at completion; each must be an active/live grader in the same org and trace function or the server rejects the replay

  • 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 { original_trace_id:, original_span_id: } (with deprecated source_trace_id/source_span_id aliases), 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.

  • db_branch (Hash, nil) (defaults to: nil) —

    passing it at all turns database branching on, so {} enables it with the mirror project's own sizing. Keys +:min_cu+/+:max_cu+ size the branch compute in Neon Compute Units and :warmup_sql warms its cache. Read the resolved branch inside the replayed method with Bitfab.current_replay_branch.

  • 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



142
143
144
145
146
147
148
149
150
151
152
 
# File 'lib/bitfab/client.rb', line 142

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, grader_ids: nil, mock: "marked",
  adapt_inputs: nil, mock_override: nil, db_branch: 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:, grader_ids:, mock:, adapt_inputs:,
    mock_override:, db_branch:,
    on_progress:
  )
end