brute 0.4.0 → 1.0.0

data/lib/brute/middleware/reasoning_normalizer.rb

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     # Handles reasoning/thinking content across model switches.
@@ -96,3 +99,94 @@ module Brute
     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_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
+    klass.define_method(:class) do
+      c = super()
+      name_str = "LLM::#{type_name}"
+      c.define_singleton_method(:name) { name_str }
+      c
+    end
+    klass.new
+  end
+
+  inner_app = ->(_env) { MockResponse.new(content: "reasoned response") }
+
+  it "injects thinking param for Anthropic with budget_tokens" do
+    provider = make_provider("Anthropic")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "claude-4", budget_tokens: 8000, enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:thinking].should == { type: "enabled", budget_tokens: 8000 }
+  end
+
+  it "does not inject thinking param for Anthropic without budget_tokens" do
+    provider = make_provider("Anthropic")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "claude-4", enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:thinking].should.be.nil
+  end
+
+  it "injects reasoning_effort for OpenAI" do
+    provider = make_provider("OpenAI")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "o3", effort: :high, enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:reasoning_effort].should == "high"
+  end
+
+  it "maps low effort correctly for OpenAI" do
+    provider = make_provider("OpenAI")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "o3", effort: :low, enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:reasoning_effort].should == "low"
+  end
+
+  it "does not inject params for unknown provider" do
+    provider = make_provider("Mistral")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "mistral-large", enabled: true)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params].should == {}
+  end
+
+  it "does not inject params when disabled" do
+    provider = make_provider("Anthropic")
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "claude-4", budget_tokens: 8000, enabled: false)
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params].should == {}
+  end
+
+  it "allows model_id to be updated mid-session" do
+    middleware = Brute::Middleware::ReasoningNormalizer.new(inner_app, model_id: "old", enabled: true)
+    middleware.model_id = "new"
+    provider = make_provider("OpenAI")
+    env = build_env(provider: provider, params: {})
+    middleware.call(env)
+    env[:params][:reasoning_effort].should.not.be.nil
+  end
+end

data/lib/brute/middleware/retry.rb

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     # Retries the inner call on transient LLM errors with exponential backoff.
@@ -9,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
@@ -43,3 +46,112 @@ module Brute
     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 mock_inner_app(response:)
+    calls = []
+    app = ->(env) { calls << env; response }
+    [app, calls]
+  end
+
+  def flaky_inner_app(error_class, fail_count:, response:)
+    attempt = 0
+    ->(env) { attempt += 1; raise error_class, "transient" if attempt <= fail_count; response }
+  end
+
+  def no_sleep_retry(*args, **kwargs)
+    mw = Brute::Middleware::Retry.new(*args, **kwargs)
+    mw.define_singleton_method(:sleep) { |_| }
+    mw
+  end
+
+  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
+
+  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 "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
+
+  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 "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
+
+  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
+
+  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 "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
+
+  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
+
+  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,11 +1,15 @@
 # frozen_string_literal: true
 
+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:)
@@ -17,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
@@ -27,3 +31,42 @@ module Brute
     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: "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 "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
+
+  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,5 +1,8 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     # Tracks cumulative token usage across all LLM calls in a session.
@@ -44,3 +47,78 @@ module Brute
     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,5 +1,8 @@
 # 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
@@ -10,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.
@@ -39,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
 
@@ -50,3 +62,118 @@ module Brute
     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,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,8 +86,48 @@ module Brute
           tool_calls: tool_calls,
           original_tool_calls: original_tool_calls,
         })
-        ctx.messages.concat([synthetic])
+        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