brute 1.0.0 → 2.0.0

data/lib/brute/middleware/token_tracking.rb

@@ -1,124 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    # Tracks cumulative token usage across all LLM calls in a session.
-    #
-    # Runs POST-call: reads usage from the response and accumulates totals
-    # in env[:metadata]. Also records per-call usage for the most recent call.
-    #
-    class TokenTracking < Base
-      def initialize(app)
-        super(app)
-        @total_input = 0
-        @total_output = 0
-        @total_reasoning = 0
-        @call_count = 0
-      end
-
-      def call(env)
-        response = @app.call(env)
-
-        if response.respond_to?(:usage) && (usage = response.usage)
-          @total_input += usage.input_tokens.to_i
-          @total_output += usage.output_tokens.to_i
-          @total_reasoning += usage.reasoning_tokens.to_i
-          @call_count += 1
-
-          env[:metadata][:tokens] = {
-            total_input: @total_input,
-            total_output: @total_output,
-            total_reasoning: @total_reasoning,
-            total: @total_input + @total_output,
-            call_count: @call_count,
-            last_call: {
-              input: usage.input_tokens.to_i,
-              output: usage.output_tokens.to_i,
-              total: usage.total_tokens.to_i,
-            },
-          }
-        end
-
-        response
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../spec/support/mock_provider"
-  require_relative "../../../spec/support/mock_response"
-
-  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
-
-  def make_response
-    MockResponse.new(content: "hello",
-      usage: LLM::Usage.new(input_tokens: 100, output_tokens: 50, reasoning_tokens: 10, total_tokens: 160))
-  end
-
-  it "passes the response through unchanged" do
-    response = make_response
-    middleware = Brute::Middleware::TokenTracking.new(->(_env) { response })
-    result = middleware.call(build_env)
-    result.should == response
-  end
-
-  it "populates total_input tokens" do
-    middleware = Brute::Middleware::TokenTracking.new(->(_env) { make_response })
-    env = build_env
-    middleware.call(env)
-    env[:metadata][:tokens][:total_input].should == 100
-  end
-
-  it "populates total_output tokens" do
-    middleware = Brute::Middleware::TokenTracking.new(->(_env) { make_response })
-    env = build_env
-    middleware.call(env)
-    env[:metadata][:tokens][:total_output].should == 50
-  end
-
-  it "populates total_reasoning tokens" do
-    middleware = Brute::Middleware::TokenTracking.new(->(_env) { make_response })
-    env = build_env
-    middleware.call(env)
-    env[:metadata][:tokens][:total_reasoning].should == 10
-  end
-
-  it "populates call_count" do
-    middleware = Brute::Middleware::TokenTracking.new(->(_env) { make_response })
-    env = build_env
-    middleware.call(env)
-    env[:metadata][:tokens][:call_count].should == 1
-  end
-
-  it "accumulates token counts across multiple calls" do
-    middleware = Brute::Middleware::TokenTracking.new(->(_env) { make_response })
-    env = build_env
-    middleware.call(env)
-    middleware.call(env)
-    env[:metadata][:tokens][:total_input].should == 200
-  end
-
-  it "handles a response without usage gracefully" do
-    no_usage = Object.new
-    middleware = Brute::Middleware::TokenTracking.new(->(_env) { no_usage })
-    env = build_env
-    middleware.call(env)
-    env[:metadata][:tokens].should.be.nil
-  end
-
-  it "handles a response where usage returns nil" do
-    nil_usage = Struct.new(:usage).new(nil)
-    middleware = Brute::Middleware::TokenTracking.new(->(_env) { nil_usage })
-    env = build_env
-    middleware.call(env)
-    env[:metadata][:tokens].should.be.nil
-  end
-end

data/lib/brute/middleware/tool_error_tracking.rb

@@ -1,179 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    # Tracks per-tool error counts and total tool call count across LLM
-    # calls, and signals when the error ceiling is reached.
-    #
-    # This middleware doesn't execute tools itself — it inspects the tool
-    # results that were sent as input to the LLM call (env[:tool_results])
-    # and counts failures and totals.
-    #
-    # When any tool exceeds max_failures, it sets env[:metadata][:tool_error_limit_reached]
-    # 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.
-    #
-    class ToolErrorTracking < Base
-      DEFAULT_MAX_FAILURES = 3
-
-      def initialize(app, max_failures: DEFAULT_MAX_FAILURES)
-        super(app)
-        @max_failures = max_failures
-        @errors = Hash.new(0) # tool_name → count
-        @total_tool_calls = 0
-      end
-
-      def call(env)
-        # PRE: count errors and totals from tool results that are about to be sent
-        if (results = env[:tool_results])
-          @total_tool_calls += results.size
-
-          results.each do |name, result|
-            @errors[name] += 1 if result.is_a?(Hash) && result[:error]
-          end
-        end
-
-        env[:metadata][:tool_calls] = @total_tool_calls
-        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
-
-      # Reset counts (e.g., between user turns).
-      def reset!
-        @errors.clear
-        @total_tool_calls = 0
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../spec/support/mock_provider"
-  require_relative "../../../spec/support/mock_response"
-
-  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
-
-  def make_middleware(app = nil)
-    app ||= ->(_env) { MockResponse.new(content: "tracked") }
-    Brute::Middleware::ToolErrorTracking.new(app, max_failures: 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 "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
-
-  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 "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
-
-  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 "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
-
-  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 "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
-
-  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
-
-  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 "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
-
-  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
-
-  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
-
-  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,133 +0,0 @@
-# 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. 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 env[:messages]. It handles three
-    # cases:
-    #
-    #   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. pending_functions is empty (nil-choice bug) but the stream recorded
-    #      tool calls → inject synthetic message using stream metadata
-    #
-    class ToolUseGuard
-      def initialize(app)
-        @app = app
-      end
-
-      def call(env)
-        response = @app.call(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(env[:messages])
-
-        # Inject a synthetic assistant message for any uncovered tool calls.
-        uncovered = tool_data.reject { |td| covered_ids.include?(td[:id]) }
-        inject_synthetic!(env[:messages], uncovered) unless uncovered.empty?
-
-        response
-      end
-
-      private
-
-      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 = env[:stream]
-          if stream&.respond_to?(:pending_tool_calls)
-            data = stream.pending_tool_calls.dup
-            stream.clear_pending_tool_calls!
-            data
-          else
-            []
-          end
-        else
-          []
-        end
-      end
-
-      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!(messages, uncovered)
-        tool_calls = uncovered.map do |td|
-          LLM::Object.from(id: td[:id], name: td[:name], arguments: td[:arguments])
-        end
-        original_tool_calls = uncovered.map do |td|
-          { "type" => "tool_use", "id" => td[:id], "name" => td[:name], "input" => td[:arguments] || {} }
-        end
-
-        synthetic = LLM::Message.new(:assistant, "", {
-          tool_calls: tool_calls,
-          original_tool_calls: original_tool_calls,
-        })
-        messages << synthetic
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../spec/support/mock_provider"
-  require_relative "../../../spec/support/mock_response"
-
-  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 "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
-
-  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
-
-  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,124 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    # Logs timing and token usage for every LLM call, and tracks cumulative
-    # timing data in env[:metadata][:timing].
-    #
-    # As the outermost middleware, it sees the full pipeline elapsed time per
-    # 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 agent loop
-    # sets this on the first call of each run()).
-    #
-    # Stores in env[:metadata][:timing]:
-    #   total_elapsed:     wall-clock since the turn began (includes tool gaps)
-    #   total_llm_elapsed: cumulative time spent inside LLM calls only
-    #   llm_call_count:    number of LLM calls so far
-    #   last_call_elapsed: duration of the most recent LLM call
-    #
-    class Tracing < Base
-      def initialize(app, logger:)
-        super(app)
-        @logger = logger
-        @call_count = 0
-        @total_llm_elapsed = 0.0
-        @turn_start = nil
-      end
-
-      def call(env)
-        @call_count += 1
-
-        # Detect new turn: tool_results is nil on the first pipeline call
-        if env[:tool_results].nil?
-          @turn_start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-          @total_llm_elapsed = 0.0
-        end
-
-        messages = env[:messages]
-        @logger.debug("[brute] LLM call ##{@call_count} (#{messages.size} messages in context)")
-
-        start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-        response = @app.call(env)
-        now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-        elapsed = now - start
-
-        @total_llm_elapsed += elapsed
-
-        tokens = response.respond_to?(:usage) ? response.usage&.total_tokens : '?'
-        @logger.info("[brute] LLM response ##{@call_count}: #{tokens} tokens, #{elapsed.round(2)}s")
-
-        env[:metadata][:timing] = {
-          total_elapsed: now - (@turn_start || start),
-          total_llm_elapsed: @total_llm_elapsed,
-          llm_call_count: @call_count,
-          last_call_elapsed: elapsed
-        }
-
-        response
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../spec/support/mock_provider"
-  require_relative "../../../spec/support/mock_response"
-
-  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 "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 "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
-
-  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 "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
-
-  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

@@ -1,18 +0,0 @@
-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

data/lib/brute/orchestrator/turn.rb

@@ -1,105 +0,0 @@
-module Brute
-  class Orchestrator
-    class Turn
-      def initialize(env:, pending:)
-        @env = env
-        @pending = pending
-      end
-
-      def perform
-        @env.dig(:callbacks, :on_tool_call_start).then do |on_start|
-          on_start&.call(
-            @pending.map do |tool, _|
-              {
-                name: tool.name,
-                arguments: tool.arguments
-              }
-            end
-          )
-        end
-
-        execute_tool_calls.tap do |results|
-          errors.each do |_, err|
-            on_result = @env.dig(:callbacks, :on_tool_result)
-            on_result&.call(err.name, result_value(err))
-            results << err
-          end
-        end
-      end
-
-      def errors = @pending.select { |_, err| err }
-      def executable = @pending.reject { |_, err| err }.map(&:first)
-
-      def execute_tool_calls
-        if executable.empty?
-          []
-        else
-          # Questions block execution — they must complete before other tools
-          # run, since the LLM may need the answer to inform subsequent work.
-          # Execute any question tools first (sequentially), then dispatch
-          # the remaining tools concurrently.
-          questions, others = executable.partition { _1.name == "question" }
-
-          Array.new.tap do |results|
-            if questions.any?
-              results.concat(execute_sequential(questions))
-            end
-
-            if others.size <= 1
-              results.concat(execute_sequential(others))
-            else
-              results.concat(execute_parallel(others))
-            end
-          end
-        end
-      end
-
-      # Run a single tool call synchronously.
-      def execute_sequential(functions)
-        on_result = @env.dig(:callbacks, :on_tool_result)
-        on_question = @env.dig(:callbacks, :on_question)
-
-        functions.map do |fn|
-          Thread.current[:on_question] = on_question
-          result = fn.call
-          on_result&.call(fn.name, result_value(result))
-          result
-        end
-      end
-
-      # Run all pending tool calls concurrently via Async::Barrier.
-      #
-      # Each tool runs in its own fiber. File-mutating tools are safe because
-      # they go through FileMutationQueue, whose Mutex is fiber-scheduler-aware
-      # in Ruby 3.4 — a fiber blocked on a per-file mutex yields to other
-      # fibers instead of blocking the thread.
-      #
-      # The barrier is stored in @barrier so abort! can cancel in-flight tools.
-      #
-      def execute_parallel(functions)
-        on_result = @env.dig(:callbacks, :on_tool_result)
-        on_question = @env.dig(:callbacks, :on_question)
-
-        Array.new(functions.size).tap do |results|
-          Async do
-            @barrier = Async::Barrier.new
-
-            functions.each_with_index do |fn, i|
-              @barrier.async do
-                Thread.current[:on_question] = on_question
-                results[i] = fn.call
-                r = results[i]
-                on_result&.call(r.name, result_value(r))
-              end
-            end
-
-            @barrier.wait
-          ensure
-            @barrier&.stop
-            @barrier = nil
-          end
-        end
-      end
-    end
-  end
-end