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_guard ⇒ Object

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