brute 0.4.0 → 0.4.1

data/lib/brute/middleware/otel/tool_results.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     module OTel
@@ -35,3 +40,90 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../../../spec/spec_helper"
+
+  RSpec.describe Brute::Middleware::OTel::ToolResults do
+    let(:response) { MockResponse.new(content: "processed") }
+    let(:inner_app) { ->(_env) { response } }
+    let(:middleware) { described_class.new(inner_app) }
+
+    it "passes the response through unchanged" do
+      env = build_env
+      result = middleware.call(env)
+      expect(result).to eq(response)
+    end
+
+    context "when env[:span] is nil" do
+      it "passes through without error" do
+        results = [["fs_read", { content: "data" }]]
+        env = build_env(tool_results: results)
+
+        result = middleware.call(env)
+        expect(result).to eq(response)
+      end
+    end
+
+    context "when env[:span] is present" do
+      let(:span) { mock_span }
+
+      it "does nothing when tool_results is nil" do
+        env = build_env(span: span, tool_results: nil)
+        middleware.call(env)
+
+        expect(span).not_to have_received(:add_event)
+        expect(span).not_to have_received(:set_attribute)
+      end
+
+      it "records a tool_result event per result" do
+        results = [
+          ["fs_read", { content: "file data" }],
+          ["shell", { output: "ok" }],
+        ]
+        env = build_env(span: span, tool_results: results)
+        middleware.call(env)
+
+        expect(span).to have_received(:set_attribute).with("brute.tool_results.count", 2)
+        expect(span).to have_received(:add_event).with(
+          "tool_result",
+          attributes: hash_including("tool.name" => "fs_read", "tool.status" => "ok")
+        )
+        expect(span).to have_received(:add_event).with(
+          "tool_result",
+          attributes: hash_including("tool.name" => "shell", "tool.status" => "ok")
+        )
+      end
+
+      it "records error status and message for failed tool results" do
+        results = [
+          ["fs_read", { error: "not found" }],
+        ]
+        env = build_env(span: span, tool_results: results)
+        middleware.call(env)
+
+        expect(span).to have_received(:add_event).with(
+          "tool_result",
+          attributes: hash_including(
+            "tool.name" => "fs_read",
+            "tool.status" => "error",
+            "tool.error" => "not found"
+          )
+        )
+      end
+
+      it "handles a mix of successful and failed results" do
+        results = [
+          ["fs_read", { content: "ok" }],
+          ["shell", { error: "exit code 1" }],
+          ["fs_write", { success: true }],
+        ]
+        env = build_env(span: span, tool_results: results)
+        middleware.call(env)
+
+        expect(span).to have_received(:set_attribute).with("brute.tool_results.count", 3)
+        expect(span).to have_received(:add_event).exactly(3).times
+      end
+    end
+  end
+end

data/lib/brute/middleware/otel.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     # OpenTelemetry instrumentation for the LLM pipeline.

data/lib/brute/middleware/reasoning_normalizer.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     # Handles reasoning/thinking content across model switches.
@@ -96,3 +101,117 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../../spec/spec_helper"
+
+  RSpec.describe Brute::Middleware::ReasoningNormalizer do
+    let(:response) { MockResponse.new(content: "reasoned response") }
+    let(:inner_app) { ->(_env) { response } }
+
+    # Build a provider whose class name contains the given string.
+    def make_provider(type_name)
+      klass = Class.new do
+        define_method(:name) { :mock }
+        define_method(:default_model) { "mock-model" }
+        define_method(:user_role) { :user }
+        define_method(:system_role) { :system }
+        define_method(:assistant_role) { :assistant }
+        define_method(:tool_role) { :tool }
+        define_method(:tracer) { nil }
+        define_method(:tracer=) { |*| }
+        define_method(:complete) { |*_args, **_kw| MockResponse.new(content: "ok") }
+      end
+      # Override class name to trigger provider_type detection
+      klass.define_method(:class) do
+        c = super()
+        name_str = "LLM::#{type_name}"
+        c.define_singleton_method(:name) { name_str }
+        c
+      end
+      klass.new
+    end
+
+    context "with Anthropic provider and budget_tokens" do
+      it "injects thinking param into env[:params]" do
+        provider = make_provider("Anthropic")
+        middleware = described_class.new(inner_app, model_id: "claude-4", budget_tokens: 8000, enabled: true)
+        env = build_env(provider: provider, params: {})
+
+        middleware.call(env)
+
+        expect(env[:params][:thinking]).to eq({ type: "enabled", budget_tokens: 8000 })
+      end
+    end
+
+    context "with Anthropic provider without budget_tokens" do
+      it "does not inject thinking param" do
+        provider = make_provider("Anthropic")
+        middleware = described_class.new(inner_app, model_id: "claude-4", enabled: true)
+        env = build_env(provider: provider, params: {})
+
+        middleware.call(env)
+
+        expect(env[:params][:thinking]).to be_nil
+      end
+    end
+
+    context "with OpenAI provider" do
+      it "injects reasoning_effort based on effort level" do
+        provider = make_provider("OpenAI")
+        middleware = described_class.new(inner_app, model_id: "o3", effort: :high, enabled: true)
+        env = build_env(provider: provider, params: {})
+
+        middleware.call(env)
+
+        expect(env[:params][:reasoning_effort]).to eq("high")
+      end
+
+      it "maps effort levels correctly" do
+        provider = make_provider("OpenAI")
+
+        { low: "low", medium: "medium", high: "high", minimal: "low", max: "high" }.each do |effort, expected|
+          middleware = described_class.new(inner_app, model_id: "o3", effort: effort, enabled: true)
+          env = build_env(provider: provider, params: {})
+          middleware.call(env)
+          expect(env[:params][:reasoning_effort]).to eq(expected), "Expected effort #{effort} to map to #{expected}"
+        end
+      end
+    end
+
+    context "with unknown provider" do
+      it "does not inject any reasoning params" do
+        provider = make_provider("Mistral")
+        middleware = described_class.new(inner_app, model_id: "mistral-large", enabled: true)
+        env = build_env(provider: provider, params: {})
+
+        middleware.call(env)
+
+        expect(env[:params]).to eq({})
+      end
+    end
+
+    context "when disabled" do
+      it "does not inject reasoning params" do
+        provider = make_provider("Anthropic")
+        middleware = described_class.new(inner_app, model_id: "claude-4", budget_tokens: 8000, enabled: false)
+        env = build_env(provider: provider, params: {})
+
+        middleware.call(env)
+
+        expect(env[:params]).to eq({})
+      end
+    end
+
+    it "allows model_id to be updated mid-session" do
+      middleware = described_class.new(inner_app, model_id: "old-model", enabled: true)
+      middleware.model_id = "new-model"
+
+      provider = make_provider("OpenAI")
+      env = build_env(provider: provider, params: {})
+      middleware.call(env)
+
+      expect(env[:params][:reasoning_effort]).not_to be_nil
+    end
+  end
+end

data/lib/brute/middleware/retry.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     # Retries the inner call on transient LLM errors with exponential backoff.
@@ -43,3 +48,91 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../../spec/spec_helper"
+
+  RSpec.describe Brute::Middleware::Retry do
+    let(:response) { MockResponse.new(content: "success") }
+
+    it "returns the response on first successful call" do
+      app, calls = mock_inner_app(response: response)
+      middleware = described_class.new(app)
+      env = build_env
+
+      result = middleware.call(env)
+
+      expect(result).to eq(response)
+      expect(calls.size).to eq(1)
+    end
+
+    it "retries on LLM::RateLimitError and succeeds" do
+      app = flaky_inner_app(LLM::RateLimitError, fail_count: 2, response: response)
+      middleware = described_class.new(app, max_attempts: 3, base_delay: 2)
+      allow(middleware).to receive(:sleep)
+      env = build_env
+
+      result = middleware.call(env)
+
+      expect(result).to eq(response)
+      expect(env[:metadata][:retry_attempt]).to eq(2)
+    end
+
+    it "retries on LLM::ServerError and succeeds" do
+      app = flaky_inner_app(LLM::ServerError, fail_count: 1, response: response)
+      middleware = described_class.new(app, max_attempts: 3, base_delay: 2)
+      allow(middleware).to receive(:sleep)
+      env = build_env
+
+      result = middleware.call(env)
+
+      expect(result).to eq(response)
+      expect(env[:metadata][:retry_attempt]).to eq(1)
+    end
+
+    it "re-raises after exhausting all attempts" do
+      app = failing_inner_app(LLM::RateLimitError, message: "rate limited")
+      middleware = described_class.new(app, max_attempts: 3, base_delay: 2)
+      allow(middleware).to receive(:sleep)
+      env = build_env
+
+      expect { middleware.call(env) }.to raise_error(LLM::RateLimitError, "rate limited")
+      expect(env[:metadata][:last_error]).to eq("rate limited")
+    end
+
+    it "does not retry non-retryable errors" do
+      call_count = 0
+      app = ->(_env) { call_count += 1; raise ArgumentError, "bad input" }
+      middleware = described_class.new(app)
+      env = build_env
+
+      expect { middleware.call(env) }.to raise_error(ArgumentError)
+      expect(call_count).to eq(1)
+    end
+
+    it "sleeps with exponential backoff delays" do
+      app = flaky_inner_app(LLM::RateLimitError, fail_count: 2, response: response)
+      middleware = described_class.new(app, max_attempts: 3, base_delay: 2)
+      allow(middleware).to receive(:sleep)
+      env = build_env
+
+      middleware.call(env)
+
+      # base_delay ** attempts: 2**1 = 2, 2**2 = 4
+      expect(middleware).to have_received(:sleep).with(2).ordered
+      expect(middleware).to have_received(:sleep).with(4).ordered
+    end
+
+    it "records retry_delay in metadata" do
+      app = flaky_inner_app(LLM::RateLimitError, fail_count: 1, response: response)
+      middleware = described_class.new(app, max_attempts: 3, base_delay: 3)
+      allow(middleware).to receive(:sleep)
+      env = build_env
+
+      middleware.call(env)
+
+      # base_delay ** attempts: 3**1 = 3
+      expect(env[:metadata][:retry_delay]).to eq(3)
+    end
+  end
+end

data/lib/brute/middleware/session_persistence.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     # Saves the conversation to disk after each LLM call.
@@ -27,3 +32,40 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../../spec/spec_helper"
+
+  RSpec.describe Brute::Middleware::SessionPersistence do
+    let(:response) { MockResponse.new(content: "saved response") }
+    let(:inner_app) { ->(_env) { response } }
+    let(:session) { double("session", save: nil) }
+    let(:middleware) { described_class.new(inner_app, session: session) }
+
+    it "passes the response through unchanged" do
+      env = build_env
+      result = middleware.call(env)
+      expect(result).to eq(response)
+    end
+
+    it "calls session.save with the context after a successful LLM call" do
+      env = build_env
+      middleware.call(env)
+      expect(session).to have_received(:save).with(env[:context])
+    end
+
+    it "does not propagate session save failures" do
+      allow(session).to receive(:save).and_raise(RuntimeError, "disk full")
+      env = build_env
+
+      expect { middleware.call(env) }.not_to raise_error
+    end
+
+    it "prints a warning to stderr on save failure" do
+      allow(session).to receive(:save).and_raise(RuntimeError, "disk full")
+      env = build_env
+
+      expect { middleware.call(env) }.to output(/Session save failed: disk full/).to_stderr
+    end
+  end
+end

data/lib/brute/middleware/token_tracking.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     # Tracks cumulative token usage across all LLM calls in a session.
@@ -44,3 +49,75 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../../spec/spec_helper"
+
+  RSpec.describe Brute::Middleware::TokenTracking do
+    let(:response) do
+      MockResponse.new(
+        content: "hello",
+        usage: LLM::Usage.new(input_tokens: 100, output_tokens: 50, reasoning_tokens: 10, total_tokens: 160)
+      )
+    end
+
+    let(:inner_app) { ->(_env) { response } }
+    let(:middleware) { described_class.new(inner_app) }
+
+    it "passes the response through unchanged" do
+      env = build_env
+      result = middleware.call(env)
+      expect(result).to eq(response)
+    end
+
+    it "populates env[:metadata][:tokens] with correct values" do
+      env = build_env
+      middleware.call(env)
+
+      tokens = env[:metadata][:tokens]
+      expect(tokens[:total_input]).to eq(100)
+      expect(tokens[:total_output]).to eq(50)
+      expect(tokens[:total_reasoning]).to eq(10)
+      expect(tokens[:total]).to eq(150) # input + output
+      expect(tokens[:call_count]).to eq(1)
+      expect(tokens[:last_call]).to eq(input: 100, output: 50, total: 160)
+    end
+
+    it "accumulates token counts across multiple calls" do
+      env = build_env
+      middleware.call(env)
+      middleware.call(env)
+
+      tokens = env[:metadata][:tokens]
+      expect(tokens[:total_input]).to eq(200)
+      expect(tokens[:total_output]).to eq(100)
+      expect(tokens[:total_reasoning]).to eq(20)
+      expect(tokens[:call_count]).to eq(2)
+    end
+
+    it "handles a response without usage gracefully" do
+      no_usage_response = double("response")
+      allow(no_usage_response).to receive(:respond_to?).with(:usage).and_return(false)
+      app = ->(_env) { no_usage_response }
+      mw = described_class.new(app)
+
+      env = build_env
+      result = mw.call(env)
+
+      expect(result).to eq(no_usage_response)
+      expect(env[:metadata][:tokens]).to be_nil
+    end
+
+    it "handles a response where usage returns nil" do
+      nil_usage_response = double("response", usage: nil)
+      allow(nil_usage_response).to receive(:respond_to?).with(:usage).and_return(true)
+      app = ->(_env) { nil_usage_response }
+      mw = described_class.new(app)
+
+      env = build_env
+      mw.call(env)
+
+      expect(env[:metadata][:tokens]).to be_nil
+    end
+  end
+end

data/lib/brute/middleware/tool_error_tracking.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     # Tracks per-tool error counts and total tool call count across LLM
@@ -50,3 +55,99 @@ module Brute
     end
   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) }
+
+    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
+
+    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)
+
+      expect(env[:metadata][:tool_calls]).to eq(3)
+    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)
+
+      expect(env[:metadata][:tool_errors]).to eq({ "fs_read" => 2 })
+    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)
+
+      expect(env[:metadata][:tool_error_limit_reached]).to be true
+    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)
+
+      expect(env[:metadata][:tool_error_limit_reached]).to be false
+    end
+
+    it "accumulates counts across multiple calls" do
+      env1 = build_env(tool_results: [["fs_read", { error: "fail" }]])
+      middleware.call(env1)
+
+      env2 = build_env(tool_results: [["fs_read", { error: "fail again" }], ["shell", { output: "ok" }]])
+      middleware.call(env2)
+
+      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
+      env = build_env(tool_results: [["fs_read", { error: "fail" }]])
+      middleware.call(env)
+
+      middleware.reset!
+
+      env2 = build_env(tool_results: nil)
+      middleware.call(env2)
+
+      expect(env2[:metadata][:tool_calls]).to eq(0)
+      expect(env2[:metadata][:tool_errors]).to eq({})
+    end
+  end
+end

data/lib/brute/middleware/tool_use_guard.rb

@@ -95,3 +95,72 @@ module Brute
     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)
+
+      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") }
+
+      inner_app = ->(_env) do
+        ctx.talk(prompt)
+        response
+      end
+
+      middleware = described_class.new(inner_app)
+      env = build_env(context: ctx, provider: provider)
+
+      expect { middleware.call(env) }.not_to raise_error
+    end
+  end
+end