brute 0.4.1 → 1.0.0

data/lib/brute/middleware/retry.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
     # propagate immediately.
     #
     # Unlike forgecode's separate retry.rs, this middleware wraps the LLM call
-    # directly — it sees the error and retries without the orchestrator knowing.
+    # directly — it sees the error and retries without the agent loop knowing.
     #
     class Retry < Base
       DEFAULT_MAX_ATTEMPTS = 3
@@ -49,90 +47,111 @@ module Brute
   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
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
 
-      result = middleware.call(env)
-
-      expect(result).to eq(response)
-      expect(env[:metadata][:retry_attempt]).to eq(2)
-    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 "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
+  def mock_inner_app(response:)
+    calls = []
+    app = ->(env) { calls << env; response }
+    [app, calls]
+  end
 
-      result = middleware.call(env)
+  def flaky_inner_app(error_class, fail_count:, response:)
+    attempt = 0
+    ->(env) { attempt += 1; raise error_class, "transient" if attempt <= fail_count; response }
+  end
 
-      expect(result).to eq(response)
-      expect(env[:metadata][:retry_attempt]).to eq(1)
-    end
+  def no_sleep_retry(*args, **kwargs)
+    mw = Brute::Middleware::Retry.new(*args, **kwargs)
+    mw.define_singleton_method(:sleep) { |_| }
+    mw
+  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
+  it "returns the response on first successful call" do
+    response = MockResponse.new(content: "success")
+    app, calls = mock_inner_app(response: response)
+    middleware = Brute::Middleware::Retry.new(app)
+    result = middleware.call(build_env)
+    result.should == response
+  end
 
-      expect { middleware.call(env) }.to raise_error(LLM::RateLimitError, "rate limited")
-      expect(env[:metadata][:last_error]).to eq("rate limited")
-    end
+  it "calls inner app exactly once on success" do
+    response = MockResponse.new(content: "success")
+    app, calls = mock_inner_app(response: response)
+    Brute::Middleware::Retry.new(app).call(build_env)
+    calls.size.should == 1
+  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
+  it "retries on LLM::RateLimitError and succeeds" do
+    response = MockResponse.new(content: "success")
+    app = flaky_inner_app(LLM::RateLimitError, fail_count: 2, response: response)
+    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 2)
+    env = build_env
+    result = middleware.call(env)
+    result.should == response
+  end
 
-      expect { middleware.call(env) }.to raise_error(ArgumentError)
-      expect(call_count).to eq(1)
-    end
+  it "records retry_attempt in metadata after retries" do
+    response = MockResponse.new(content: "success")
+    app = flaky_inner_app(LLM::RateLimitError, fail_count: 2, response: response)
+    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 2)
+    env = build_env
+    middleware.call(env)
+    env[:metadata][:retry_attempt].should == 2
+  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
+  it "retries on LLM::ServerError and succeeds" do
+    response = MockResponse.new(content: "success")
+    app = flaky_inner_app(LLM::ServerError, fail_count: 1, response: response)
+    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 2)
+    result = middleware.call(build_env)
+    result.should == response
+  end
 
-      middleware.call(env)
+  it "re-raises after exhausting all attempts" do
+    app = ->(_env) { raise LLM::RateLimitError, "rate limited" }
+    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 2)
+    lambda { middleware.call(build_env) }.should.raise(LLM::RateLimitError)
+  end
 
-      # 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 "does not retry non-retryable errors" do
+    call_count = 0
+    app = ->(_env) { call_count += 1; raise ArgumentError, "bad input" }
+    middleware = Brute::Middleware::Retry.new(app)
+    lambda { middleware.call(build_env) }.should.raise(ArgumentError)
+  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
+  it "only calls inner app once for non-retryable errors" do
+    call_count = 0
+    app = ->(_env) { call_count += 1; raise ArgumentError, "bad input" }
+    middleware = Brute::Middleware::Retry.new(app)
+    begin; middleware.call(build_env); rescue ArgumentError; end
+    call_count.should == 1
+  end
 
-      middleware.call(env)
+  it "records retry_delay in metadata" do
+    response = MockResponse.new(content: "success")
+    app = flaky_inner_app(LLM::RateLimitError, fail_count: 1, response: response)
+    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 3)
+    env = build_env
+    middleware.call(env)
+    env[:metadata][:retry_delay].should == 3
+  end
 
-      # base_delay ** attempts: 3**1 = 3
-      expect(env[:metadata][:retry_delay]).to eq(3)
-    end
+  it "tracks sleep delays for exponential backoff" do
+    response = MockResponse.new(content: "success")
+    app = flaky_inner_app(LLM::RateLimitError, fail_count: 2, response: response)
+    delays = []
+    mw = Brute::Middleware::Retry.new(app, max_attempts: 3, base_delay: 2)
+    mw.define_singleton_method(:sleep) { |d| delays << d }
+    mw.call(build_env)
+    delays.should == [2, 4]
   end
 end

data/lib/brute/middleware/session_persistence.rb

@@ -1,16 +1,15 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
     # Saves the conversation to disk after each LLM call.
     #
-    # Runs POST-call: delegates to Session#save. Failures are non-fatal —
-    # a broken session save should never crash the agent loop.
+    # Runs POST-call: serializes env[:messages] via Session#save_messages.
+    # Failures are non-fatal — a broken session save should never crash
+    # the agent loop.
     #
     class SessionPersistence < Base
       def initialize(app, session:)
@@ -22,7 +21,7 @@ module Brute
         response = @app.call(env)
 
         begin
-          @session.save(env[:context])
+          @session.save_messages(env[:messages])
         rescue => e
           warn "[brute] Session save failed: #{e.message}"
         end
@@ -33,39 +32,41 @@ module Brute
   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
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
 
-    it "does not propagate session save failures" do
-      allow(session).to receive(:save).and_raise(RuntimeError, "disk full")
-      env = build_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 { middleware.call(env) }.not_to raise_error
-    end
+  it "passes the response through unchanged" do
+    response = MockResponse.new(content: "saved response")
+    session = Struct.new(:saved) { def save_messages(m); self.saved = m; end }.new
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::SessionPersistence.new(inner_app, session: session)
+    result = middleware.call(build_env)
+    result.should == response
+  end
 
-    it "prints a warning to stderr on save failure" do
-      allow(session).to receive(:save).and_raise(RuntimeError, "disk full")
-      env = build_env
+  it "calls session.save_messages with env messages" do
+    response = MockResponse.new(content: "saved response")
+    session = Struct.new(:saved) { def save_messages(m); self.saved = m; end }.new
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::SessionPersistence.new(inner_app, session: session)
+    messages = [LLM::Message.new(:user, "hello")]
+    middleware.call(build_env(messages: messages))
+    session.saved.should == messages
+  end
 
-      expect { middleware.call(env) }.to output(/Session save failed: disk full/).to_stderr
-    end
+  it "does not propagate session save failures" do
+    response = MockResponse.new(content: "saved response")
+    session = Object.new
+    session.define_singleton_method(:save_messages) { |_| raise RuntimeError, "disk full" }
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::SessionPersistence.new(inner_app, session: session)
+    lambda { middleware.call(build_env) }.should.not.raise
   end
 end

data/lib/brute/middleware/token_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
@@ -50,74 +48,77 @@ module Brute
   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
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
 
-    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
+  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 "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
+  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 "accumulates token counts across multiple calls" do
-      env = build_env
-      middleware.call(env)
-      middleware.call(env)
+  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
 
-      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 "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 "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)
+  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
 
-      env = build_env
-      result = mw.call(env)
+  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
 
-      expect(result).to eq(no_usage_response)
-      expect(env[:metadata][:tokens]).to be_nil
-    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 "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)
+  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
 
-      env = build_env
-      mw.call(env)
+  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
 
-      expect(env[:metadata][:tokens]).to 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