Module: Braintrust::Contrib::Setup

Defined in:
lib/braintrust/contrib/setup.rb

Overview

Automatic instrumentation setup for LLM libraries.

Intercepts ‘require` calls to detect when LLM libraries are loaded, then patches them automatically. The main challenge is doing this safely with zeitwerk (Rails’ autoloader) which also hooks into require.

## The Zeitwerk Problem

Zeitwerk uses ‘alias_method :zeitwerk_original_require, :require`. If we prepend to Kernel before zeitwerk loads, zeitwerk captures our method as its “original”, creating an infinite loop.

## Solution: Two-Phase Hook

┌─────────────────────────────────────────────────────────────────────────┐
│                         Setup.run! called                               │
└─────────────────────────────────────────────────────────────────────────┘
                                   │
         ┌─────────────────────────┼─────────────────────────┐
         ▼                         ▼                         ▼
┌─────────────────┐     ┌─────────────────────┐    ┌─────────────────────┐
│ Rails loaded?   │     │ Zeitwerk loaded?    │    │ Neither loaded yet  │
│ (Rails::Railtie)│     │                     │    │                     │
└────────┬────────┘     └──────────┬──────────┘    └──────────┬──────────┘
         │                         │                          │
         ▼                         ▼                          ▼
┌─────────────────┐     ┌─────────────────────┐    ┌─────────────────────┐
│ Install Railtie │     │ Install prepend     │    │ Install watcher     │
│ (after_init)    │     │ hook directly       │    │ hook (alias_method) │
└─────────────────┘     └─────────────────────┘    └──────────┬──────────┘
         │                         │                          │
         │                         │               ┌──────────┴──────────┐
         │                         │               ▼                     ▼
         │                         │      ┌──────────────┐    ┌──────────────────┐
         │                         │      │ Rails loads? │    │ Zeitwerk loads?  │
         │                         │      │ → Railtie    │    │ → Prepend hook   │
         │                         │      └──────────────┘    └──────────────────┘
         │                         │               │                     │
         ▼                         ▼               ▼                     ▼
┌─────────────────────────────────────────────────────────────────────────┐
│                    LLM library loads → patch!                           │
└─────────────────────────────────────────────────────────────────────────┘

The watcher hook uses alias_method which zeitwerk captures harmlessly (alias chains work correctly). Once zeitwerk/Rails loads, we upgrade to the better approach: prepend (takes precedence, ‘super` chains through zeitwerk) or Railtie (patches after all gems loaded via after_initialize).

Constant Summary collapse

REENTRANCY_KEY =
:braintrust_in_require_hook

Class Method Summary collapse

Class Method Details

.install_railtie!Object



112
113
114
115
116
# File 'lib/braintrust/contrib/setup.rb', line 112

def install_railtie!
  return if @railtie_installed
  @railtie_installed = true
  require_relative "rails/railtie"
end

.install_require_hook!Object



118
119
120
121
122
# File 'lib/braintrust/contrib/setup.rb', line 118

def install_require_hook!
  return if @require_hook_installed
  @require_hook_installed = true
  Kernel.prepend(RequireHook)
end

.on_require(path) ⇒ Object

Called after each require to check if we should patch anything.



83
84
85
86
87
88
89
90
91
92
93
94
95
96
# File 'lib/braintrust/contrib/setup.rb', line 83

def on_require(path)
  return unless @registry

  @registry.integrations_for_require_path(path).each do |integration|
    next unless integration.available? && integration.compatible?
    next if @only && !@only.include?(integration.integration_name)
    next if @except&.include?(integration.integration_name)

    Braintrust::Log.debug("Contrib::Setup: patching #{integration.integration_name}")
    integration.patch!
  end
rescue => e
  Braintrust::Log.error("Auto-instrument failed: #{e.message}")
end

.railtie_installed?Boolean

Returns:

  • (Boolean)


109
# File 'lib/braintrust/contrib/setup.rb', line 109

def railtie_installed? = @railtie_installed

.require_hook_installed?Boolean

Returns:

  • (Boolean)


110
# File 'lib/braintrust/contrib/setup.rb', line 110

def require_hook_installed? = @require_hook_installed

.run!Object

Main entry point. Call once per process.



60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
# File 'lib/braintrust/contrib/setup.rb', line 60

def run!
  unless Internal::Env.auto_instrument
    Braintrust::Log.debug("Contrib::Setup: auto-instrumentation disabled via environment")
    return
  end

  @registry = Contrib::Registry.instance
  @only = Internal::Env.instrument_only
  @except = Internal::Env.instrument_except

  if defined?(::Rails::Railtie)
    Braintrust::Log.debug("Contrib::Setup: using Rails railtie hook")
    install_railtie!
  elsif defined?(::Zeitwerk)
    Braintrust::Log.debug("Contrib::Setup: using require hook (zeitwerk detected)")
    install_require_hook!
  else
    Braintrust::Log.debug("Contrib::Setup: using watcher hook")
    install_watcher_hook!
  end
end

.with_reentrancy_guardObject

Execute block with reentrancy protection (prevents infinite loops).



99
100
101
102
103
104
105
106
107
# File 'lib/braintrust/contrib/setup.rb', line 99

def with_reentrancy_guard
  return if Thread.current[REENTRANCY_KEY]
  Thread.current[REENTRANCY_KEY] = true
  yield
rescue => e
  Braintrust::Log.error("Auto-instrument failed: #{e.message}")
ensure
  Thread.current[REENTRANCY_KEY] = false
end