brute 4.3.2 → 5.0.0

data/lib/brute/message_transport/ruby_llm.rb

@@ -8,6 +8,12 @@ module Brute
   class MessageTransport
     class RubyLLM < MessageTransport
 
+      # What the provider reported about this call — the transport knows its
+      # own library's shape, so it knows which detector to ask.
+      def self.usage_metrics(message)
+        Brute::UsageDetection::RubyLLM.detect(message)
+      end
+
       # Brute::Message -> RubyLLM::Message (tool calls as ruby_llm's id-keyed hash).
       def self.dump(message)
         tool_calls = message.tool_calls&.to_h do |tc|

data/lib/brute/message_transport.rb

@@ -21,6 +21,14 @@ module Brute
       messages.map { |message| dump(message) }
     end
 
+    # Inbound: what the provider reported about this call, as a
+    # Brute::UsageDetection::Usage. Each transport knows its own library's
+    # response shape, so it is the one that knows which detector to ask.
+    # Answers nil for a library that reports no usage at all.
+    def self.usage_metrics(_result)
+      nil
+    end
+
     def initialize(result)
       @result = result
     end

data/lib/brute/middleware/000_base.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "brute"
+require "brute/hooks"
+
+module Brute
+  module Middleware
+    # The parent of every middleware: it takes the next app, does its work
+    # around it, and calls it.
+    #
+    #   class Shout < Brute::Middleware::Base
+    #     def call(env)
+    #       emit(ENTER_EVENT, env, self)
+    #       @app.call(env)
+    #     end
+    #   end
+    #
+    # Brute::Hooks is included, so the event names are first class in every
+    # subclass — ENTER_EVENT, not Brute::Hooks::ENTER_EVENT.
+    class Base
+      include Brute::Hooks
+
+      def initialize(app, *, **)
+        @app = app
+      end
+
+      # A layer that adds nothing passes the turn straight through.
+      def call(env) = @app.call(env)
+
+      private
+
+        attr_reader :app
+    end
+  end
+end
+
+__END__
+
+describe "brute/middleware/000_base" do
+  it "carries the app, passes the turn through, and makes the event names first class" do
+    passed = []
+    terminal = ->(env) { passed << env; env }
+
+    Brute::Middleware::Base.new(terminal).call({ turn: 1 }).should == { turn: 1 }
+    passed.should == [{ turn: 1 }]
+
+    # Whatever a subclass declares, the base swallows.
+    Brute::Middleware::Base.new(terminal, :extra, keyword: true).call({}).should == {}
+
+    # The event names resolve through the class itself, so a subclass writes
+    # ENTER_EVENT rather than Brute::Hooks::ENTER_EVENT.
+    Brute::Middleware::Base.const_get(:ENTER_EVENT).should == :enter
+    Brute::Middleware::SystemPrompt.const_get(:AFTER_TOOL_EVENT).should == :after_tool
+
+    # Every middleware in the chain descends from it.
+    Brute::Middleware::SystemPrompt.ancestors.should.include Brute::Middleware::Base
+    Brute::Middleware::ToolPipeline.ancestors.should.include Brute::Middleware::Base
+    Brute::Middleware::Loop::ToolResult.ancestors.should.include Brute::Middleware::Base
+  end
+end

data/lib/brute/middleware/002_session_log.rb

@@ -23,7 +23,7 @@ module Brute
     #     .use(Brute::Middleware::SystemPrompt)
     #     ...
     #
-    class SessionLog
+    class SessionLog < Brute::Middleware::Base
       def initialize(app, path:)
         @app  = app
         @path = path

data/lib/brute/middleware/004_summarize.rb

@@ -22,7 +22,7 @@ module Brute
     #   use ToolPipeline
     #   run ->(env) { ... }   # inline LLM call proc (see Brute.agent)
     #
-    class Summarize
+    class Summarize < Brute::Middleware::Base
       DEFAULT_PROMPT = "Provide your complete findings based on everything you've explored."
 
       def initialize(app, prompt: DEFAULT_PROMPT)

data/lib/brute/middleware/005_tracing.rb

@@ -21,7 +21,7 @@ module Brute
     #   llm_call_count:    number of LLM calls so far
     #   last_call_elapsed: duration of the most recent LLM call
     #
-    class Tracing
+    class Tracing < Brute::Middleware::Base
       def initialize(app, logger:)
         @app = app
 

data/lib/brute/middleware/006_loop.rb

@@ -21,7 +21,7 @@ module Brute
     #     !env[:should_exit] && env[:messages].last&.role == :tool
     #   end
     #
-    class Loop
+    class Loop < Brute::Middleware::Base
       def initialize(app, condition = nil, &block)
         @app       = app
         @condition = condition || block

data/lib/brute/middleware/008_checkpoint.rb

@@ -33,7 +33,7 @@ module Brute
     # System messages are not persisted (SystemPrompt re-adds them each
     # turn); restored history is inserted after any leading system message
     # and before the current turn's input.
-    class Checkpoint
+    class Checkpoint < Brute::Middleware::Base
       def initialize(app, path:, resume: nil)
         @app    = app
         @path   = path

data/lib/brute/middleware/010_max_iterations.rb

@@ -11,7 +11,7 @@ module Brute
     # stating that maximum iterations have been reached. This causes
     # Loop::ToolResult to exit its loop naturally (last message is not :tool).
     #
-    class MaxIterations
+    class MaxIterations < Brute::Middleware::Base
 
       DEFAULT_MAX_ITERATIONS = 100
 

data/lib/brute/middleware/020_system_prompt.rb

@@ -25,7 +25,7 @@ module Brute
     # message (e.g. from session.system(...)), so manually-set system
     # prompts are respected.
     #
-    class SystemPrompt
+    class SystemPrompt < Brute::Middleware::Base
       def initialize(app, system_prompt: Brute::SystemPrompt.default)
         @app = app
         @system_prompt = system_prompt

data/lib/brute/middleware/025_skills.rb

@@ -23,7 +23,7 @@ module Brute
     #
     # Place it before Middleware::SystemPrompt in the stack. It never touches
     # env[:messages] itself.
-    class Skills
+    class Skills < Brute::Middleware::Base
       def initialize(app, skills: [])
         @app = app
         @skills = skills

data/lib/brute/middleware/040_compaction_check.rb

@@ -15,7 +15,7 @@ module Brute
     # Or an LLM that doesn't support it can just use the messages
     # that come after the compaction
     #
-    class CompactionCheck
+    class CompactionCheck < Brute::Middleware::Base
       def initialize(app, compactor: nil, system_prompt:, **compactor_opts)
         @app = app
         @compactor = compactor

data/lib/brute/middleware/060_questions.rb

@@ -5,7 +5,7 @@ require "brute"
 
 module Brute
   module Middleware
-    class Question
+    class Question < Brute::Middleware::Base
       def initialize(app)
         @app = app
       end

data/lib/brute/middleware/070_tool_pipeline.rb

@@ -8,7 +8,7 @@ require "async/barrier"
 
 module Brute
   module Middleware
-    class ToolPipeline
+    class ToolPipeline < Brute::Middleware::Base
       def initialize(app, tools: [])
         @app   = app
         @tools = tools
@@ -56,34 +56,37 @@ module Brute
                   name:      name.to_s,
                   arguments: args,
                   result:    nil,
+                  denied:    nil,
                   events:    env[:events],
                   metadata:  {},
                   turn_env:  env,
                 }
-                if (hooks = env[:hooks])
-                  responses = hooks.emit(:before_tool, call_env).compact
-                  call_env[:result] = responses.last if call_env[:result].nil? && !responses.empty?
-
-                  if call_env[:result].nil?
-                    denial = hooks.emit(:approve_tool, call_env).find { |r| r == false || r.is_a?(String) }
-                    unless denial.nil?
-                      call_env[:result] = denial.is_a?(String) ? denial : %(Tool call to "#{name}" was denied.)
-                    end
+                # A subscriber takes part by mutating the call env: set
+                # :result to answer without executing, set :denied to refuse.
+                emit(BEFORE_TOOL_EVENT, env, call_env)
+
+                if call_env[:result].nil?
+                  emit(APPROVE_TOOL_EVENT, env, call_env)
+
+                  if (denial = call_env[:denied])
+                    call_env[:result] = denial.is_a?(String) ? denial : %(Tool call to "#{name}" was denied.)
                   end
                 end
 
-                result = if call_env[:result].nil?
-                           available_tools[name].call(call_env[:arguments])
-                         else
-                           call_env[:result]
-                         end
+                # Only the tool's own execution is timed: a call that
+                # before_tool answered, or approve_tool denied, never ran.
+                result = call_env[:result]
 
-                if (hooks = env[:hooks])
-                  call_env[:result] = result
-                  hooks.emit(:after_tool, call_env)
-                  result = call_env[:result]
+                if result.nil?
+                  emit(TOOL_DURATION_EVENT, env, call_env) do
+                    result = available_tools[name].call(call_env[:arguments])
+                  end
                 end
 
+                call_env[:result] = result
+                emit(AFTER_TOOL_EVENT, env, call_env)
+                result = call_env[:result]
+
                 # Coerce to String so Hash results (e.g. Shell's
                 # {stdout:, stderr:, exit_code:}) serialize predictably.
                 if result.is_a?(String)
@@ -176,8 +179,18 @@ describe "brute/middleware/070_tool_pipeline" do
 
   # --- lifecycle hooks (Brute::Hooks) ---
 
-  def hook_env(hooks)
-    { messages: Brute.log, events: [], hooks: hooks }
+  # A layer only gets its emit from the builder that made it, so a hook spec
+  # builds a real pipeline rather than instantiating the middleware alone.
+  def hooked(inner, tools:, &subscribe)
+    pipeline = Brute::Turn::Pipeline.new
+    pipeline.use Brute::Middleware::ToolPipeline, tools: tools
+    pipeline.run(Object.new.tap { |o| o.define_singleton_method(:call, &inner) })
+    subscribe.call(pipeline)
+    pipeline
+  end
+
+  def hook_env
+    { messages: Brute.log, events: [] }
   end
 
   it "before_tool may rewrite arguments and short-circuit with a result" do
@@ -186,20 +199,19 @@ describe "brute/middleware/070_tool_pipeline" do
       env[:messages] << Brute::Message.new(role: :assistant, content: "",
         tool_calls: [{ id: "tc1", name: "echo", arguments: { "text" => "orig" } }])
     end
-    hooks = Brute::Hooks.new
-    hooks.on(:before_tool) { |call| call[:arguments] = { text: "rewritten" }; nil }
-    mw = Brute::Middleware::ToolPipeline.new(inner, tools: [tool])
-    env = hook_env(hooks)
+
+    pipeline = hooked(inner, tools: [tool]) do |p|
+      p.on(:before_tool) { |_env, call| call[:arguments] = { text: "rewritten" } }
+    end
+    env = hook_env
     env[:messages].user("hi")
-    mw.call(env)
+    pipeline.call(env)
     env[:messages].last.content.should == "ran:rewritten"
 
-    hooks2 = Brute::Hooks.new
-    hooks2.on(:before_tool) { |_call| "canned" }
-    mw2 = Brute::Middleware::ToolPipeline.new(inner, tools: [tool])
-    env2 = hook_env(hooks2)
+    canned = hooked(inner, tools: [tool]) { |p| p.on(:before_tool) { |_env, call| call[:result] = "canned" } }
+    env2 = hook_env
     env2[:messages].user("hi")
-    mw2.call(env2)
+    canned.call(env2)
     env2[:messages].last.content.should == "canned" # never executed
   end
 
@@ -209,18 +221,17 @@ describe "brute/middleware/070_tool_pipeline" do
       env[:messages] << Brute::Message.new(role: :assistant, content: "",
         tool_calls: [{ id: "tc1", name: "exec", arguments: {} }])
     end
-    hooks = Brute::Hooks.new
-    hooks.on(:approve_tool) { |_call| false }
-    env = hook_env(hooks)
+
+    denied = hooked(inner, tools: [tool]) { |p| p.on(:approve_tool) { |_env, call| call[:denied] = true } }
+    env = hook_env
     env[:messages].user("hi")
-    Brute::Middleware::ToolPipeline.new(inner, tools: [tool]).call(env)
+    denied.call(env)
     env[:messages].last.content.should == %(Tool call to "exec" was denied.)
 
-    hooks2 = Brute::Hooks.new
-    hooks2.on(:approve_tool) { |_call| "denied by policy" }
-    env2 = hook_env(hooks2)
+    by_policy = hooked(inner, tools: [tool]) { |p| p.on(:approve_tool) { |_env, call| call[:denied] = "denied by policy" } }
+    env2 = hook_env
     env2[:messages].user("hi")
-    Brute::Middleware::ToolPipeline.new(inner, tools: [tool]).call(env2)
+    by_policy.call(env2)
     env2[:messages].last.content.should == "denied by policy"
   end
 
@@ -230,11 +241,13 @@ describe "brute/middleware/070_tool_pipeline" do
       env[:messages] << Brute::Message.new(role: :assistant, content: "",
         tool_calls: [{ id: "tc1", name: "echo", arguments: {} }])
     end
-    hooks = Brute::Hooks.new
-    hooks.on(:after_tool) { |call| call[:result] = "rewrote(#{call[:result]})" }
-    env = hook_env(hooks)
+
+    pipeline = hooked(inner, tools: [tool]) do |p|
+      p.on(:after_tool) { |_env, call| call[:result] = "rewrote(#{call[:result]})" }
+    end
+    env = hook_env
     env[:messages].user("hi")
-    Brute::Middleware::ToolPipeline.new(inner, tools: [tool]).call(env)
+    pipeline.call(env)
     env[:messages].last.content.should == "rewrote(raw)"
   end
 
@@ -262,13 +275,13 @@ describe "brute/middleware/070_tool_pipeline" do
     inner = ->(env) {
       env[:messages] << Brute::Message.new(role: :assistant, content: "", tool_calls: tool_calls)
     }
-    mw = Brute::Middleware::ToolPipeline.new(inner, tools: [big_tool])
+    pipeline = hooked(inner, tools: [big_tool]) { |_p| nil }
     env = {
       messages: Brute.log,
       events: [],
     }
     env[:messages].user("hello")
-    mw.call(env)
+    pipeline.call(env)
 
     tool_msg = env[:messages].select { |m| m.role == :tool }.last
     tool_msg.content.lines.size.should.be < 2100
@@ -299,13 +312,13 @@ describe "brute/middleware/070_tool_pipeline" do
     inner = ->(env) {
       env[:messages] << Brute::Message.new(role: :assistant, content: "", tool_calls: tool_calls)
     }
-    mw = Brute::Middleware::ToolPipeline.new(inner, tools: [pre_truncated_tool])
+    pipeline = hooked(inner, tools: [pre_truncated_tool]) { |_p| nil }
     env = {
       messages: Brute.log,
       events: [],
     }
     env[:messages].user("hello")
-    mw.call(env)
+    pipeline.call(env)
 
     tool_msg = env[:messages].select { |m| m.role == :tool }.last
     # Should contain exactly one truncation marker, not two

data/lib/brute/middleware/event_handler.rb

@@ -5,7 +5,7 @@ require 'brute'
 
 module Brute
   module Middleware
-    class EventHandler
+    class EventHandler < Brute::Middleware::Base
       def initialize(app, handler_class:, **opts)
         @app = app
         @handler_class = handler_class

data/lib/brute/middleware/user_queue.rb

@@ -5,7 +5,7 @@ require "brute"
 
 module Brute
   module Middleware
-    class UserQueue
+    class UserQueue < Brute::Middleware::Base
 
       # Useful for testing...
       # App will keep looping till all inputs are drained.

data/lib/brute/turn/agent_pipeline.rb

@@ -45,13 +45,12 @@ module Brute
           events:            events,
           metadata:          {},
           current_iteration: 1,
-          hooks:             hooks,
         }
-        hooks.emit(:turn_start, env)
+        hooks.emit(TURN_START_EVENT, env)
         begin
-          build.call(env)
+          hooks.emit(TURN_DURATION_EVENT, env) { build.call(env) }
         ensure
-          hooks.emit(:turn_end, env)
+          hooks.emit(TURN_END_EVENT, env)
         end
         env
       end

data/lib/brute/turn/pipeline.rb

@@ -2,11 +2,14 @@
 
 require "bundler/setup"
 require "brute"
+require "brute/hooks"
 
 module Brute
   module Turn
     class Pipeline < ::Rack::Builder
       module Chainable
+        include Brute::Hooks
+
         # Enables the following syntax:
         #
         #   Brute.agent
@@ -17,8 +20,49 @@ module Brute
         #       your_llm_library.complete("How to make money?")
         #     }
         #
-        def use(...) = tap { super }
-        def run(...) = tap { super }
+        # Every layer announces itself: :middleware_added when it goes on the
+        # stack, then :enter and :exit around its own work on each turn. Both
+        # of those carry the middleware instance as self, and :exit fires from
+        # an ensure so a layer that raises is still reported.
+        def use(middleware, *args, &block)
+          tap do
+            super
+            hooks.emit(MIDDLEWARE_ADDED_EVENT, {}, middleware, *args)
+
+            builder = @use.pop
+            @use << lambda do |app|
+              bind_emitter(builder.call(app)).tap do |layer|
+                layer.define_singleton_method(:call) do |env|
+                  emit(ENTER_EVENT, env, self)
+
+                  # The layer's work is :duration's block, so its subscribers
+                  # are handed how long this layer took — its own work plus
+                  # everything below it. :exit then marks the layer done,
+                  # from an ensure so it lands even when the work raised.
+                  result = nil
+                  begin
+                    emit(DURATION_EVENT, env, self) { result = super(env) }
+                  ensure
+                    emit(EXIT_EVENT, env, self)
+                  end
+
+                  result
+                end
+              end
+            end
+          end
+        end
+        # Keeps a trailing hash flagged as keywords through `super`, the way
+        # Rack::Builder#use does — without it `use Mw, label: "x"` arrives at
+        # the middleware's constructor as a positional Hash.
+        ruby2_keywords(:use) if respond_to?(:ruby2_keywords, true)
+
+        def run(app = nil, &block)
+          tap do
+            bind_emitter(app) unless app.nil?
+            super
+          end
+        end
 
         # Subscribe a lifecycle hook (see Brute::Hooks):
         #
@@ -29,6 +73,29 @@ module Brute
         #     .on(:approve_tool) { |call| call[:name] != "exec" }
         #
         def on(...) = tap { hooks.on(...) }
+
+        private
+
+          # Bind an emit() to this builder's store, so a layer fires events at
+          # the pipeline it belongs to rather than fishing them out of env.
+          def bind_emitter(object)
+            if object.is_a?(Proc)
+              warn("brute: #{self.class} was given a lambda, which cannot be given an emit — its events will not fire")
+              return object
+            end
+
+            if object.respond_to?(:emit)
+              raise ArgumentError, "#{object.class} already defines #emit, so the pipeline will not bind its own over it"
+            end
+
+            store = hooks
+            object.define_singleton_method(:emit) do |event, env, *extras, &work|
+              store.emit(event, env, *extras, &work)
+            end
+            object
+          end
+
+        public
       end
 
       include Chainable
@@ -66,6 +133,88 @@ end
 __END__
 
 describe "brute/turn/pipeline" do
+  it "announces a layer when it is added, then on the way in and out of every turn" do
+    seen = []
+    hooks = Brute::Hooks.new
+    hooks.on(Brute::Hooks::MIDDLEWARE_ADDED_EVENT) { |env, mw, *args| seen << [:added, env, mw, args] }
+    hooks.on(Brute::Hooks::ENTER_EVENT) { |_env, layer| seen << [:enter, layer.class] }
+    hooks.on(Brute::Hooks::DURATION_EVENT) { |_env, started, finished, layer| seen << [:duration, layer.class, finished >= started] }
+    hooks.on(Brute::Hooks::EXIT_EVENT) { |_env, layer| seen << [:exit, layer.class] }
+
+    labeller = Class.new do
+      def initialize(app, label:); @app = app; @label = label; end
+      def call(env); @app.call(env); end
+    end
+
+    pipeline = Brute::Turn::Pipeline.new
+    pipeline.instance_variable_set(:@hooks, hooks)
+    pipeline.use labeller, label: "outer"
+    pipeline.run ->(env) { env }
+    pipeline.call({ hooks: hooks })
+
+    seen.first.should == [:added, {}, labeller, [{ label: "outer" }]]
+    seen[1].should == [:enter, labeller]
+    seen[2].should == [:duration, labeller, true]
+    seen[3].should == [:exit, labeller]
+  end
+
+  it "reports a layer that raises through :exit anyway" do
+    seen = []
+
+    boom = Class.new do
+      def initialize(app); @app = app; end
+      def call(_env); raise "boom"; end
+    end
+
+    pipeline = Brute::Turn::Pipeline.new
+    pipeline.on(Brute::Hooks::EXIT_EVENT) { |_env, layer| seen << layer.class }
+    pipeline.use boom
+    pipeline.run Object.new.tap { |o| o.define_singleton_method(:call) { |env| env } }
+    begin
+      pipeline.call({})
+    rescue RuntimeError
+      nil
+    end
+
+    seen.should == [boom]
+  end
+
+  it "binds emit to its own store, warns for a lambda, and refuses to shadow an existing emit" do
+    seen = []
+    quiet = Class.new do
+      def initialize(app); @app = app; end
+      def call(env); emit(:ping, env, :from_layer); @app.call(env); end
+    end
+
+    pipeline = Brute::Turn::Pipeline.new
+    pipeline.on(:ping) { |env, extra| seen << [env, extra] }
+    pipeline.use quiet
+    pipeline.run Object.new.tap { |o| o.define_singleton_method(:call) { |env| env } }
+    pipeline.call({})
+    seen.should == [[{}, :from_layer]]
+
+    # A lambda cannot reach the bound method, so it is warned about, not patched.
+    warned = []
+    lambda_pipeline = Brute::Turn::Pipeline.new
+    lambda_pipeline.define_singleton_method(:warn) { |message| warned << message }
+    lambda_pipeline.run ->(env) { env }
+    warned.size.should == 1
+    warned.first.should.match(/lambda/)
+
+    # An object that already answers to emit is left alone, loudly.
+    emitter = Class.new do
+      def initialize(app); @app = app; end
+      def emit(*); end
+      def call(env); env; end
+    end
+    # The layer is constructed when the stack is built, so that is when the
+    # collision is discovered.
+    shadowing = Brute::Turn::Pipeline.new
+    shadowing.use emitter
+    shadowing.run Object.new.tap { |o| o.define_singleton_method(:call) { |env| env } }
+    should.raise(ArgumentError) { shadowing.call({}) }
+  end
+
   it "builds and calls a chain" do
     seen = []
     inc = Class.new do

data/lib/brute/usage_detection/lang_chain.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "brute"
+require "brute/usage_detection/usage"
+
+module Brute
+  module UsageDetection
+    # langchainrb has no single usage object: each provider's response
+    # subclass answers prompt_tokens / completion_tokens / total_tokens by
+    # digging its own raw shape, and some answer none of them.
+    module LangChain
+      def self.detect(response)
+        detected = Usage.new(
+          input:  read(response, :prompt_tokens),
+          output: read(response, :completion_tokens),
+          total:  read(response, :total_tokens),
+          raw:    response,
+        )
+
+        detected.empty? ? nil : detected
+      end
+
+      def self.read(response, name) = response.respond_to?(name) ? response.public_send(name) : nil
+    end
+  end
+end
+
+__END__
+
+describe "brute/usage_detection/lang_chain" do
+  LangChainResponse = Struct.new(:prompt_tokens, :completion_tokens, :total_tokens) unless defined?(LangChainResponse)
+
+  it "reads the per-provider token methods and answers nil when the response has none" do
+    usage = Brute::UsageDetection::LangChain.detect(LangChainResponse.new(10, 5, 15))
+    usage.input.should == 10
+    usage.output.should == 5
+    usage.total.should == 15
+
+    # A provider subclass that reports nothing.
+    Brute::UsageDetection::LangChain.detect(LangChainResponse.new(nil, nil, nil)).should.be.nil
+    Brute::UsageDetection::LangChain.detect(Object.new).should.be.nil
+  end
+end