brute 0.4.1 → 1.0.0

data/lib/brute/middleware/tool_error_tracking.rb

@@ -1,9 +1,7 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
@@ -15,7 +13,7 @@ module Brute
     # and counts failures and totals.
     #
     # When any tool exceeds max_failures, it sets env[:metadata][:tool_error_limit_reached]
-    # so the orchestrator can decide to stop.
+    # so the agent loop can decide to stop.
     #
     # Also stores env[:metadata][:tool_calls] with the cumulative number of
     # tool invocations in the current session.
@@ -44,6 +42,15 @@ module Brute
         env[:metadata][:tool_errors] = @errors.dup
         env[:metadata][:tool_error_limit_reached] = @errors.any? { |_, c| c >= @max_failures }
 
+        if env[:metadata][:tool_error_limit_reached]
+          failed_tool, fail_count = @errors.max_by { |_, c| c }
+          env[:should_exit] ||= {
+            reason:  "tool_error_limit_reached",
+            message: "Tool '#{failed_tool}' has failed #{fail_count} times (limit: #{@max_failures}). Stopping.",
+            source:  "ToolErrorTracking",
+          }
+        end
+
         @app.call(env)
       end
 
@@ -56,98 +63,117 @@ module Brute
   end
 end
 
-if __FILE__ == $0
-  require_relative "../../../spec/spec_helper"
-
-  RSpec.describe Brute::Middleware::ToolErrorTracking do
-    let(:response) { MockResponse.new(content: "tracked") }
-    let(:inner_app) { ->(_env) { response } }
-    let(:middleware) { described_class.new(inner_app, max_failures: 3) }
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
 
-    it "passes the response through" do
-      env = build_env
-      result = middleware.call(env)
-      expect(result).to eq(response)
-    end
-
-    it "reports zero tool calls when tool_results is nil" do
-      env = build_env(tool_results: nil)
-      middleware.call(env)
-
-      expect(env[:metadata][:tool_calls]).to eq(0)
-      expect(env[:metadata][:tool_errors]).to eq({})
-      expect(env[:metadata][:tool_error_limit_reached]).to be false
-    end
+  def build_env(**overrides)
+    { provider: MockProvider.new, model: nil, input: "test prompt", tools: [],
+      messages: [], stream: nil, params: {}, metadata: {}, callbacks: {},
+      tool_results: nil, streaming: false, should_exit: nil, pending_functions: [] }.merge(overrides)
+  end
 
-    it "counts total tool calls from tool_results" do
-      results = [
-        ["fs_read", { content: "data" }],
-        ["shell", { output: "ok" }],
-        ["fs_write", { success: true }],
-      ]
-      env = build_env(tool_results: results)
-      middleware.call(env)
+  def make_middleware(app = nil)
+    app ||= ->(_env) { MockResponse.new(content: "tracked") }
+    Brute::Middleware::ToolErrorTracking.new(app, max_failures: 3)
+  end
 
-      expect(env[:metadata][:tool_calls]).to eq(3)
-    end
+  it "passes the response through" do
+    response = MockResponse.new(content: "tracked")
+    app = ->(_env) { response }
+    result = make_middleware(app).call(build_env)
+    result.should == response
+  end
 
-    it "counts per-tool errors from results with error key" do
-      results = [
-        ["fs_read", { error: "not found" }],
-        ["fs_read", { error: "permission denied" }],
-        ["shell", { output: "ok" }],
-      ]
-      env = build_env(tool_results: results)
-      middleware.call(env)
+  it "reports zero tool calls when tool_results is nil" do
+    env = build_env(tool_results: nil)
+    make_middleware.call(env)
+    env[:metadata][:tool_calls].should == 0
+  end
 
-      expect(env[:metadata][:tool_errors]).to eq({ "fs_read" => 2 })
-    end
+  it "reports empty tool errors when tool_results is nil" do
+    env = build_env(tool_results: nil)
+    make_middleware.call(env)
+    env[:metadata][:tool_errors].should == {}
+  end
 
-    it "sets tool_error_limit_reached when a tool hits max_failures" do
-      results = [
-        ["fs_read", { error: "fail 1" }],
-        ["fs_read", { error: "fail 2" }],
-        ["fs_read", { error: "fail 3" }],
-      ]
-      env = build_env(tool_results: results)
-      middleware.call(env)
+  it "does not flag limit reached when tool_results is nil" do
+    env = build_env(tool_results: nil)
+    make_middleware.call(env)
+    env[:metadata][:tool_error_limit_reached].should.be.false
+  end
 
-      expect(env[:metadata][:tool_error_limit_reached]).to be true
-    end
+  it "counts total tool calls from tool_results" do
+    results = [["fs_read", { content: "data" }], ["shell", { output: "ok" }], ["fs_write", { success: true }]]
+    env = build_env(tool_results: results)
+    make_middleware.call(env)
+    env[:metadata][:tool_calls].should == 3
+  end
 
-    it "does not flag below the threshold" do
-      results = [
-        ["fs_read", { error: "fail 1" }],
-        ["fs_read", { error: "fail 2" }],
-      ]
-      env = build_env(tool_results: results)
-      middleware.call(env)
+  it "counts per-tool errors from results with error key" do
+    results = [["fs_read", { error: "not found" }], ["fs_read", { error: "denied" }], ["shell", { output: "ok" }]]
+    env = build_env(tool_results: results)
+    make_middleware.call(env)
+    env[:metadata][:tool_errors].should == { "fs_read" => 2 }
+  end
 
-      expect(env[:metadata][:tool_error_limit_reached]).to be false
-    end
+  it "sets tool_error_limit_reached when a tool hits max_failures" do
+    results = [["fs_read", { error: "1" }], ["fs_read", { error: "2" }], ["fs_read", { error: "3" }]]
+    env = build_env(tool_results: results)
+    make_middleware.call(env)
+    env[:metadata][:tool_error_limit_reached].should.be.true
+  end
 
-    it "accumulates counts across multiple calls" do
-      env1 = build_env(tool_results: [["fs_read", { error: "fail" }]])
-      middleware.call(env1)
+  it "does not flag below the threshold" do
+    results = [["fs_read", { error: "1" }], ["fs_read", { error: "2" }]]
+    env = build_env(tool_results: results)
+    make_middleware.call(env)
+    env[:metadata][:tool_error_limit_reached].should.be.false
+  end
 
-      env2 = build_env(tool_results: [["fs_read", { error: "fail again" }], ["shell", { output: "ok" }]])
-      middleware.call(env2)
+  it "accumulates counts across multiple calls" do
+    mw = make_middleware
+    mw.call(build_env(tool_results: [["fs_read", { error: "fail" }]]))
+    env2 = build_env(tool_results: [["fs_read", { error: "again" }], ["shell", { output: "ok" }]])
+    mw.call(env2)
+    env2[:metadata][:tool_calls].should == 3
+  end
 
-      expect(env2[:metadata][:tool_calls]).to eq(3) # 1 + 2
-      expect(env2[:metadata][:tool_errors]).to eq({ "fs_read" => 2 })
-    end
+  it "clears counters on reset!" do
+    mw = make_middleware
+    mw.call(build_env(tool_results: [["fs_read", { error: "fail" }]]))
+    mw.reset!
+    env2 = build_env(tool_results: nil)
+    mw.call(env2)
+    env2[:metadata][:tool_calls].should == 0
+  end
 
-    it "clears counters on reset!" do
-      env = build_env(tool_results: [["fs_read", { error: "fail" }]])
-      middleware.call(env)
+  it "sets should_exit reason when error limit reached" do
+    results = [["fs_read", { error: "1" }], ["fs_read", { error: "2" }], ["fs_read", { error: "3" }]]
+    env = build_env(tool_results: results)
+    make_middleware.call(env)
+    env[:should_exit][:reason].should == "tool_error_limit_reached"
+  end
 
-      middleware.reset!
+  it "sets should_exit source to ToolErrorTracking" do
+    results = [["fs_read", { error: "1" }], ["fs_read", { error: "2" }], ["fs_read", { error: "3" }]]
+    env = build_env(tool_results: results)
+    make_middleware.call(env)
+    env[:should_exit][:source].should == "ToolErrorTracking"
+  end
 
-      env2 = build_env(tool_results: nil)
-      middleware.call(env2)
+  it "does not set should_exit below the threshold" do
+    results = [["fs_read", { error: "1" }], ["fs_read", { error: "2" }]]
+    env = build_env(tool_results: results)
+    make_middleware.call(env)
+    env[:should_exit].should.be.nil
+  end
 
-      expect(env2[:metadata][:tool_calls]).to eq(0)
-      expect(env2[:metadata][:tool_errors]).to eq({})
-    end
+  it "does not overwrite should_exit if already set" do
+    results = [["fs_read", { error: "1" }], ["fs_read", { error: "2" }], ["fs_read", { error: "3" }]]
+    existing = { reason: "doom_loop_detected", message: "loop", source: "DoomLoopDetection" }
+    env = build_env(tool_results: results, should_exit: existing)
+    make_middleware.call(env)
+    env[:should_exit][:reason].should == "doom_loop_detected"
   end
 end

data/lib/brute/middleware/tool_use_guard.rb

@@ -1,24 +1,27 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     # Guards against tool-only LLM responses where the assistant message
     # is dropped from the context buffer.
     #
     # When the LLM responds with only tool_use blocks (no text), llm.rb's
-    # response adapter produces empty choices. Context#talk appends nil,
-    # BufferNilGuard strips it, and the assistant message carrying tool_use
-    # blocks is lost. This causes "unexpected tool_use_id" on the next call
-    # because tool_result references a tool_use that's missing from the buffer.
+    # response adapter produces empty choices. The assistant message carrying
+    # tool_use blocks may be lost. This causes "unexpected tool_use_id" on
+    # the next call because tool_result references a tool_use that's missing
+    # from the message history.
     #
     # This middleware runs post-call and ensures every pending tool_use ID
-    # is covered by an assistant message in the buffer. It handles three
+    # is covered by an assistant message in env[:messages]. It handles three
     # cases:
     #
-    #   1. ctx.functions is non-empty and the assistant message exists → no-op
-    #   2. ctx.functions is non-empty but the assistant message is missing
+    #   1. pending_functions is non-empty and the assistant message exists → no-op
+    #   2. pending_functions is non-empty but the assistant message is missing
     #      (or has different IDs) → inject synthetic message
-    #   3. ctx.functions is empty (nil-choice bug) but the stream recorded
+    #   3. pending_functions is empty (nil-choice bug) but the stream recorded
     #      tool calls → inject synthetic message using stream metadata
     #
     class ToolUseGuard
@@ -29,32 +32,30 @@ module Brute
       def call(env)
         response = @app.call(env)
 
-        ctx = env[:context]
-
-        # Collect pending tool data from ctx.functions (primary) or the
-        # stream's recorded metadata (fallback for nil-choice bug).
-        tool_data = collect_tool_data(ctx, env)
+        # Collect pending tool data from env[:pending_functions] (primary)
+        # or the stream's recorded metadata (fallback for nil-choice bug).
+        tool_data = collect_tool_data(env)
         return response if tool_data.empty?
 
         # Find all tool_use IDs already covered by assistant messages.
-        covered_ids = covered_tool_ids(ctx)
+        covered_ids = covered_tool_ids(env[:messages])
 
         # Inject a synthetic assistant message for any uncovered tool calls.
         uncovered = tool_data.reject { |td| covered_ids.include?(td[:id]) }
-        inject_synthetic!(ctx, uncovered) unless uncovered.empty?
+        inject_synthetic!(env[:messages], uncovered) unless uncovered.empty?
 
         response
       end
 
       private
 
-      def collect_tool_data(ctx, env)
-        functions = ctx.functions
+      def collect_tool_data(env)
+        functions = env[:pending_functions]
         if functions && !functions.empty?
           functions.map { |fn| { id: fn.id, name: fn.name, arguments: fn.arguments } }
         elsif env[:streaming]
-          stream = resolve_stream(ctx)
-          if stream
+          stream = env[:stream]
+          if stream&.respond_to?(:pending_tool_calls)
             data = stream.pending_tool_calls.dup
             stream.clear_pending_tool_calls!
             data
@@ -66,19 +67,14 @@ module Brute
         end
       end
 
-      def resolve_stream(ctx)
-        stream = ctx.instance_variable_get(:@params)&.dig(:stream)
-        stream if stream.respond_to?(:pending_tool_calls)
-      end
-
-      def covered_tool_ids(ctx)
-        ctx.messages.to_a
+      def covered_tool_ids(messages)
+        messages
           .select { |m| m.role.to_s == "assistant" && m.tool_call? }
           .flat_map { |m| (m.extra.original_tool_calls || []).map { |tc| tc["id"] } }
           .to_set
       end
 
-      def inject_synthetic!(ctx, uncovered)
+      def inject_synthetic!(messages, uncovered)
         tool_calls = uncovered.map do |td|
           LLM::Object.from(id: td[:id], name: td[:name], arguments: td[:arguments])
         end
@@ -90,77 +86,48 @@ module Brute
           tool_calls: tool_calls,
           original_tool_calls: original_tool_calls,
         })
-        ctx.messages.concat([synthetic])
+        messages << synthetic
       end
     end
   end
 end
 
-if __FILE__ == $0
-  require_relative "../../../spec/spec_helper"
-
-  RSpec.describe Brute::Middleware::ToolUseGuard do
-    let(:provider) { MockProvider.new }
-
-    # Helper: build a response that produces pending tool calls (functions) in the context.
-    def make_tool_response(tool_calls:)
-      MockResponse.new(content: "", tool_calls: tool_calls)
-    end
-
-    it "passes the response through when there are no pending functions" do
-      response = MockResponse.new(content: "no tools")
-      allow(provider).to receive(:complete).and_return(response)
-
-      ctx = LLM::Context.new(provider, tools: [])
-      prompt = ctx.prompt { |p| p.system("sys"); p.user("hi") }
-
-      inner_app = ->(_env) { ctx.talk(prompt); response }
-      middleware = described_class.new(inner_app)
-      env = build_env(context: ctx, provider: provider)
-
-      result = middleware.call(env)
-      expect(result).to eq(response)
-    end
-
-    it "does not inject a synthetic message when the assistant message already has tool_call?" do
-      tool_calls = [{ id: "toolu_1", name: "fs_read", arguments: { "path" => "test.rb" } }]
-      response = make_tool_response(tool_calls: tool_calls)
-      allow(provider).to receive(:complete).and_return(response)
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
 
-      ctx = LLM::Context.new(provider, tools: [])
-      prompt = ctx.prompt { |p| p.system("sys"); p.user("read it") }
-
-      inner_app = ->(_env) { ctx.talk(prompt); response }
-      middleware = described_class.new(inner_app)
-      env = build_env(context: ctx, provider: provider)
-
-      middleware.call(env)
-
-      messages = ctx.messages.to_a
-      assistant_msgs = messages.select { |m| m.role.to_s == "assistant" }
-      # Should only have the original assistant message, no synthetic
-      expect(assistant_msgs.size).to eq(1)
-    end
-
-    it "injects a synthetic assistant message when tool calls exist but assistant is missing" do
-      tool_calls = [{ id: "toolu_1", name: "fs_read", arguments: { "path" => "test.rb" } }]
-      response = MockResponse.new(content: "")
-      # Simulate the bug: choices[-1] is nil, so no assistant message stored
-      allow(response).to receive(:choices).and_return([nil])
-      allow(provider).to receive(:complete).and_return(response)
-
-      ctx = LLM::Context.new(provider, tools: [])
-      prompt = ctx.prompt { |p| p.system("sys"); p.user("read it") }
+  def build_env(**overrides)
+    { provider: MockProvider.new, model: nil, input: "test prompt", tools: [],
+      messages: [], stream: nil, params: {}, metadata: {}, callbacks: {},
+      tool_results: nil, streaming: false, should_exit: nil, pending_functions: [] }.merge(overrides)
+  end
 
-      inner_app = ->(_env) do
-        ctx.talk(prompt)
-        response
-      end
+  it "passes the response through when there are no pending functions" do
+    response = MockResponse.new(content: "no tools")
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::ToolUseGuard.new(inner_app)
+    result = middleware.call(build_env(pending_functions: []))
+    result.should == response
+  end
 
-      middleware = described_class.new(inner_app)
-      env = build_env(context: ctx, provider: provider)
+  it "injects a synthetic assistant message when tool calls exist but assistant is missing" do
+    fn = Struct.new(:id, :name, :arguments, keyword_init: true)
+           .new(id: "toolu_1", name: "fs_read", arguments: { "path" => "test.rb" })
+    response = MockResponse.new(content: "")
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::ToolUseGuard.new(inner_app)
+    env = build_env(messages: [], pending_functions: [fn])
+    lambda { middleware.call(env) }.should.not.raise
+  end
 
-      expect { middleware.call(env) }.not_to raise_error
-    end
+  it "creates one assistant message for uncovered tool calls" do
+    fn = Struct.new(:id, :name, :arguments, keyword_init: true)
+           .new(id: "toolu_1", name: "fs_read", arguments: { "path" => "test.rb" })
+    response = MockResponse.new(content: "")
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::ToolUseGuard.new(inner_app)
+    env = build_env(messages: [], pending_functions: [fn])
+    middleware.call(env)
+    env[:messages].select { |m| m.role.to_s == "assistant" }.size.should == 1
   end
 end

data/lib/brute/middleware/tracing.rb

@@ -1,9 +1,7 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
@@ -14,7 +12,7 @@ module Brute
     # call. It also tracks total wall-clock time across all calls in a turn
     # (including tool execution gaps between LLM calls).
     #
-    # A new turn is detected when env[:tool_results] is nil (the orchestrator
+    # A new turn is detected when env[:tool_results] is nil (the agent loop
     # sets this on the first call of each run()).
     #
     # Stores in env[:metadata][:timing]:
@@ -41,7 +39,7 @@ module Brute
           @total_llm_elapsed = 0.0
         end
 
-        messages = env[:context].messages.to_a
+        messages = env[:messages]
         @logger.debug("[brute] LLM call ##{@call_count} (#{messages.size} messages in context)")
 
         start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
@@ -67,68 +65,60 @@ module Brute
   end
 end
 
-if __FILE__ == $0
-  require_relative "../../../spec/spec_helper"
-
-  RSpec.describe Brute::Middleware::Tracing do
-    let(:response) { MockResponse.new(content: "traced response") }
-    let(:inner_app) { ->(_env) { response } }
-    let(:log_output) { StringIO.new }
-    let(:logger) { Logger.new(log_output) }
-    let(:middleware) { described_class.new(inner_app, logger: logger) }
-
-    it "passes the response through unchanged" do
-      env = build_env(tool_results: nil)
-      result = middleware.call(env)
-      expect(result).to eq(response)
-    end
-
-    it "populates env[:metadata][:timing] with all required keys" do
-      env = build_env(tool_results: nil)
-      middleware.call(env)
-
-      timing = env[:metadata][:timing]
-      expect(timing).to include(
-        :total_elapsed,
-        :total_llm_elapsed,
-        :llm_call_count,
-        :last_call_elapsed
-      )
-      expect(timing[:llm_call_count]).to eq(1)
-      expect(timing[:last_call_elapsed]).to be >= 0
-      expect(timing[:total_llm_elapsed]).to be >= 0
-    end
-
-    it "resets turn timing when tool_results is nil (new turn)" do
-      env = build_env(tool_results: nil)
-      middleware.call(env)
-      first_elapsed = env[:metadata][:timing][:total_llm_elapsed]
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
 
-      # Simulate continuation within the same turn (tool_results present)
-      env[:tool_results] = [["read", { content: "file data" }]]
-      middleware.call(env)
+  def build_env(**overrides)
+    { provider: MockProvider.new, model: nil, input: "test prompt", tools: [],
+      messages: [], stream: nil, params: {}, metadata: {}, callbacks: {},
+      tool_results: nil, streaming: false, should_exit: nil, pending_functions: [] }.merge(overrides)
+  end
 
-      expect(env[:metadata][:timing][:llm_call_count]).to eq(2)
-      expect(env[:metadata][:timing][:total_llm_elapsed]).to be >= first_elapsed
-    end
+  it "passes the response through unchanged" do
+    response = MockResponse.new(content: "traced response")
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::Tracing.new(inner_app, logger: Logger.new(StringIO.new))
+    result = middleware.call(build_env(tool_results: nil))
+    result.should == response
+  end
 
-    it "accumulates call count across multiple calls" do
-      env = build_env(tool_results: nil)
-      middleware.call(env)
-      env[:tool_results] = [["read", {}]]
-      middleware.call(env)
-      middleware.call(env)
+  it "populates timing with llm_call_count" do
+    response = MockResponse.new(content: "traced response")
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::Tracing.new(inner_app, logger: Logger.new(StringIO.new))
+    env = build_env(tool_results: nil)
+    middleware.call(env)
+    env[:metadata][:timing][:llm_call_count].should == 1
+  end
 
-      expect(env[:metadata][:timing][:llm_call_count]).to eq(3)
-    end
+  it "populates timing with non-negative last_call_elapsed" do
+    response = MockResponse.new(content: "traced response")
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::Tracing.new(inner_app, logger: Logger.new(StringIO.new))
+    env = build_env(tool_results: nil)
+    middleware.call(env)
+    (env[:metadata][:timing][:last_call_elapsed] >= 0).should.be.true
+  end
 
-    it "logs debug and info messages" do
-      env = build_env(tool_results: nil)
-      middleware.call(env)
+  it "accumulates call count across multiple calls" do
+    response = MockResponse.new(content: "traced response")
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::Tracing.new(inner_app, logger: Logger.new(StringIO.new))
+    env = build_env(tool_results: nil)
+    middleware.call(env)
+    env[:tool_results] = [["read", {}]]
+    middleware.call(env)
+    middleware.call(env)
+    env[:metadata][:timing][:llm_call_count].should == 3
+  end
 
-      log_text = log_output.string
-      expect(log_text).to include("LLM call #1")
-      expect(log_text).to include("LLM response #1")
-    end
+  it "logs LLM call and response messages" do
+    response = MockResponse.new(content: "traced response")
+    inner_app = ->(_env) { response }
+    log_output = StringIO.new
+    middleware = Brute::Middleware::Tracing.new(inner_app, logger: Logger.new(log_output))
+    middleware.call(build_env(tool_results: nil))
+    log_output.string.should =~ /LLM call #1/
   end
 end

data/lib/brute/middleware.rb

@@ -0,0 +1,18 @@
+require_relative 'middleware/base'
+require_relative 'middleware/llm_call'
+require_relative 'middleware/retry'
+require_relative 'middleware/doom_loop_detection'
+require_relative 'middleware/token_tracking'
+require_relative 'middleware/compaction_check'
+require_relative 'middleware/session_persistence'
+require_relative 'middleware/message_tracking'
+require_relative 'middleware/tracing'
+require_relative 'middleware/tool_error_tracking'
+require_relative 'middleware/reasoning_normalizer'
+require_relative "middleware/tool_use_guard"
+require_relative "middleware/otel"
+
+module Brute
+  module Middleware
+  end
+end