brute 0.4.0 → 1.0.0

data/lib/brute/middleware/doom_loop_detection.rb

@@ -1,29 +1,39 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     # Detects when the agent is stuck repeating tool call patterns and injects
-    # a corrective warning into the context before the next LLM call.
+    # a corrective warning into the message history before the next LLM call.
     #
     # Runs PRE-call: inspects the conversation history for repeating tool call
-    # patterns. If detected, talks a warning message into the context so the
-    # LLM sees it as input alongside the normal tool results.
+    # patterns. If detected, appends a warning message so the LLM sees it as
+    # input alongside the normal tool results.
     #
     class DoomLoopDetection < Base
       def initialize(app, threshold: 3)
         super(app)
-        @detector = Brute::DoomLoopDetector.new(threshold: threshold)
+        @detector = Brute::Loop::DoomLoopDetector.new(threshold: threshold)
       end
 
       def call(env)
-        ctx = env[:context]
-        messages = ctx.messages.to_a
+        messages = env[:messages]
 
         if (reps = @detector.detect(messages))
           warning = @detector.warning_message(reps)
           # Inject the warning as a user message so the LLM sees it
-          ctx.talk(warning)
+          env[:messages] << LLM::Message.new(:user, warning)
           env[:metadata][:doom_loop_detected] = reps
+
+          # Signal the agent loop to exit after this LLM call completes.
+          # First-writer-wins: don't overwrite if another middleware already set it.
+          env[:should_exit] ||= {
+            reason:  "doom_loop_detected",
+            message: "Agent is stuck repeating the same tool calls (#{reps} repetitions).",
+            source:  "DoomLoopDetection",
+          }
         end
 
         @app.call(env)
@@ -31,3 +41,96 @@ 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
+
+  FakeFunc = Struct.new(:name, :arguments, keyword_init: true)
+
+  def assistant_msg_with_functions(function_list)
+    msg = LLM::Message.new(:assistant, "tool msg", {})
+    msg.define_singleton_method(:functions) { function_list }
+    msg
+  end
+
+  it "passes through when no doom loop is detected" do
+    inner_app = ->(_env) { MockResponse.new(content: "loop check") }
+    middleware = Brute::Middleware::DoomLoopDetection.new(inner_app, threshold: 3)
+    env = build_env
+    middleware.call(env)
+    env[:metadata][:doom_loop_detected].should.be.nil
+  end
+
+  it "detects consecutive identical tool calls" do
+    inner_app = ->(_env) { MockResponse.new(content: "loop check") }
+    fn = FakeFunc.new(name: "fs_read", arguments: '{"path":"x.rb"}')
+    messages = 4.times.map { assistant_msg_with_functions([fn]) }
+    middleware = Brute::Middleware::DoomLoopDetection.new(inner_app, threshold: 3)
+    env = build_env(messages: messages)
+    middleware.call(env)
+    env[:metadata][:doom_loop_detected].should.not.be.nil
+  end
+
+  it "does not trigger below the threshold" do
+    inner_app = ->(_env) { MockResponse.new(content: "loop check") }
+    fn = FakeFunc.new(name: "fs_read", arguments: '{"path":"x.rb"}')
+    messages = 2.times.map { assistant_msg_with_functions([fn]) }
+    middleware = Brute::Middleware::DoomLoopDetection.new(inner_app, threshold: 3)
+    env = build_env(messages: messages)
+    middleware.call(env)
+    env[:metadata][:doom_loop_detected].should.be.nil
+  end
+
+  it "sets should_exit reason when doom loop detected" do
+    inner_app = ->(_env) { MockResponse.new(content: "loop check") }
+    fn = FakeFunc.new(name: "fs_read", arguments: '{"path":"x.rb"}')
+    messages = 4.times.map { assistant_msg_with_functions([fn]) }
+    middleware = Brute::Middleware::DoomLoopDetection.new(inner_app, threshold: 3)
+    env = build_env(messages: messages)
+    middleware.call(env)
+    env[:should_exit][:reason].should == "doom_loop_detected"
+  end
+
+  it "does not set should_exit when no loop detected" do
+    inner_app = ->(_env) { MockResponse.new(content: "loop check") }
+    middleware = Brute::Middleware::DoomLoopDetection.new(inner_app, threshold: 3)
+    env = build_env
+    middleware.call(env)
+    env[:should_exit].should.be.nil
+  end
+
+  it "does not overwrite should_exit if already set" do
+    inner_app = ->(_env) { MockResponse.new(content: "loop check") }
+    fn = FakeFunc.new(name: "fs_read", arguments: '{"path":"x.rb"}')
+    messages = 4.times.map { assistant_msg_with_functions([fn]) }
+    middleware = Brute::Middleware::DoomLoopDetection.new(inner_app, threshold: 3)
+    existing = { reason: "other", message: "earlier", source: "Other" }
+    env = build_env(messages: messages, should_exit: existing)
+    middleware.call(env)
+    env[:should_exit][:reason].should == "other"
+  end
+
+  it "appends a warning message when loop detected" do
+    inner_app = ->(_env) { MockResponse.new(content: "loop check") }
+    fn = FakeFunc.new(name: "fs_read", arguments: '{"path":"x.rb"}')
+    messages = 4.times.map { assistant_msg_with_functions([fn]) }
+    middleware = Brute::Middleware::DoomLoopDetection.new(inner_app, threshold: 3)
+    env = build_env(messages: messages)
+    original_count = env[:messages].size
+    middleware.call(env)
+    env[:messages].size.should == original_count + 1
+  end
+
+  it "generates warning message with repetition count" do
+    detector = Brute::Loop::DoomLoopDetector.new(threshold: 3)
+    msg = detector.warning_message(5)
+    msg.should =~ /5 times/
+  end
+end

data/lib/brute/middleware/llm_call.rb

@@ -1,17 +1,35 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     # The terminal "app" in the pipeline — performs the actual LLM call.
     #
+    # Builds a fresh LLM::Context per call from env[:messages], makes the
+    # call, extracts new messages back into env[:messages], and stashes
+    # pending functions in env[:pending_functions].
+    #
     # When streaming, on_content fires incrementally via AgentStream.
     # When not streaming, fires on_content post-hoc with the full text.
     #
     class LLMCall
       def call(env)
-        ctx = env[:context]
+        ctx = build_context(env)
+
+        # Load existing conversation history into the ephemeral context
+        ctx.messages.concat(env[:messages])
+
         response = ctx.talk(env[:input])
 
+        # Extract new messages appended by talk() and store them
+        new_messages = ctx.messages.to_a.drop(env[:messages].size)
+        env[:messages].concat(new_messages)
+
+        # Stash pending functions for the agent loop
+        env[:pending_functions] = ctx.functions.to_a
+
         # Only fire on_content post-hoc when NOT streaming
         # (streaming delivers chunks incrementally via AgentStream)
         unless env[:streaming]
@@ -26,6 +44,14 @@ module Brute
 
       private
 
+      def build_context(env)
+        params = {}
+        params[:tools]  = env[:tools]   if env[:tools]&.any?
+        params[:stream] = env[:stream]  if env[:stream]
+        params[:model]  = env[:model]   if env[:model]
+        LLM::Context.new(env[:provider], **params)
+      end
+
       # Safely extract text content from an LLM response.
       # Returns nil when the response contains only tool calls (no assistant text),
       # which causes LLM::Contract::Completion#content to raise NoMethodError
@@ -39,3 +65,64 @@ 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 "calls the provider and returns a response" do
+    provider = MockProvider.new
+    middleware = Brute::Middleware::LLMCall.new
+    env = build_env(provider: provider, input: "hello", streaming: false)
+    response = middleware.call(env)
+    response.should.not.be.nil
+  end
+
+  it "records a call on the provider" do
+    provider = MockProvider.new
+    middleware = Brute::Middleware::LLMCall.new
+    env = build_env(provider: provider, input: "hello", streaming: false)
+    middleware.call(env)
+    provider.calls.size.should == 1
+  end
+
+  it "appends new messages to env[:messages]" do
+    provider = MockProvider.new
+    middleware = Brute::Middleware::LLMCall.new
+    env = build_env(provider: provider, input: "hello", streaming: false)
+    middleware.call(env)
+    env[:messages].should.not.be.empty
+  end
+
+  it "populates env[:pending_functions] as an Array" do
+    provider = MockProvider.new
+    middleware = Brute::Middleware::LLMCall.new
+    env = build_env(provider: provider, input: "hello", streaming: false)
+    middleware.call(env)
+    env[:pending_functions].should.be.kind_of(Array)
+  end
+
+  it "does not fire on_content callback when streaming" do
+    provider = MockProvider.new
+    middleware = Brute::Middleware::LLMCall.new
+    called = false
+    env = build_env(provider: provider, input: "hi", streaming: true, callbacks: { on_content: ->(_) { called = true } })
+    middleware.call(env)
+    called.should.be.false
+  end
+
+  it "preserves existing messages across calls" do
+    provider = MockProvider.new
+    middleware = Brute::Middleware::LLMCall.new
+    existing = LLM::Message.new(:user, "previous")
+    env = build_env(provider: provider, input: "hello", streaming: false, messages: [existing])
+    middleware.call(env)
+    env[:messages].first.should == existing
+  end
+end

data/lib/brute/middleware/message_tracking.rb

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     # Records every LLM exchange into a MessageStore in the OpenCode
@@ -15,7 +18,7 @@ module Brute
     #      corresponding tool parts to "completed" (or "error").
     #
     # The middleware also stores itself in env[:message_tracking] so the
-    # orchestrator can access the current assistant message ID for callbacks.
+    # agent loop can access the current assistant message ID for callbacks.
     #
     class MessageTracking < Base
       attr_reader :store
@@ -111,8 +114,7 @@ module Brute
       end
 
       def record_tool_calls(env)
-        ctx = env[:context]
-        functions = ctx.functions
+        functions = env[:pending_functions]
         return if functions.nil? || functions.empty?
 
         functions.each do |fn|
@@ -165,14 +167,11 @@ module Brute
       # ── Helpers ────────────────────────────────────────────────────
 
       # Resolve the actual model used for the request.
-      # Prefers the model set on the LLM::Context (which respects user overrides)
-      # and falls back to the provider's default_model.
+      # Prefers env[:model] (set by AgentTurn) and falls back to the
+      # provider's default_model.
       def resolve_model_name(env)
-        ctx = env[:context]
-        if ctx && ctx.instance_variable_defined?(:@params)
-          ctx_model = ctx.instance_variable_get(:@params)&.dig(:model)
-          return ctx_model.to_s if ctx_model
-        end
+        model = env[:model]
+        return model.to_s if model
 
         # Fall back to provider default
         env[:provider]&.respond_to?(:default_model) ? env[:provider].default_model.to_s : nil
@@ -207,3 +206,134 @@ module Brute
     end
   end
 end
+
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
+  require "tmpdir"
+  require "fileutils"
+
+  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 with_tracking
+    tmpdir = Dir.mktmpdir("brute_test_")
+    store = Brute::Store::MessageStore.new(session_id: "test-session", dir: tmpdir)
+    response = MockResponse.new(content: "Hello from the LLM")
+    inner_app = ->(_env) { response }
+    middleware = Brute::Middleware::MessageTracking.new(inner_app, store: store)
+    yield middleware, store, response
+  ensure
+    FileUtils.rm_rf(tmpdir)
+  end
+
+  it "records a user message on first call of a turn" do
+    with_tracking do |mw, store, _|
+      mw.call(build_env(input: "What is Ruby?", tool_results: nil))
+      user_msg = store.messages.find { |m| m[:info][:role] == "user" }
+      user_msg[:parts][0][:text].should == "What is Ruby?"
+    end
+  end
+
+  it "records only one user message per turn" do
+    with_tracking do |mw, store, _|
+      env = build_env(input: "Hello", tool_results: nil)
+      mw.call(env)
+      env[:tool_results] = [["read", "contents"]]
+      mw.call(env)
+      store.messages.select { |m| m[:info][:role] == "user" }.size.should == 1
+    end
+  end
+
+  it "records an assistant message after LLM call" do
+    with_tracking do |mw, store, _|
+      mw.call(build_env(input: "Hello", tool_results: nil))
+      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
+      asst.should.not.be.nil
+    end
+  end
+
+  it "captures text content as a text part" do
+    with_tracking do |mw, store, _|
+      mw.call(build_env(input: "Hello", tool_results: nil))
+      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
+      text_parts = asst[:parts].select { |p| p[:type] == "text" }
+      text_parts[0][:text].should == "Hello from the LLM"
+    end
+  end
+
+  it "captures token usage from response" do
+    with_tracking do |mw, store, _|
+      mw.call(build_env(input: "Hello", tool_results: nil))
+      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
+      asst[:info][:tokens][:input].should == 100
+    end
+  end
+
+  it "records tool calls as tool parts in running state" do
+    with_tracking do |mw, store, _|
+      fn = Struct.new(:id, :name, :arguments, keyword_init: true).new(id: "call_001", name: "read", arguments: { file_path: "/test" })
+      mw.call(build_env(input: "Read the file", tool_results: nil, pending_functions: [fn]))
+      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
+      tool_parts = asst[:parts].select { |p| p[:type] == "tool" }
+      tool_parts[0][:state][:status].should == "running"
+    end
+  end
+
+  it "updates tool parts when results arrive" do
+    with_tracking do |mw, store, _|
+      fn = Struct.new(:id, :name, :arguments, keyword_init: true).new(id: "call_001", name: "read", arguments: { file_path: "/test" })
+      env = build_env(input: "Read the file", tool_results: nil, pending_functions: [fn])
+      mw.call(env)
+      env[:pending_functions] = []
+      env[:tool_results] = [["read", "file contents here"]]
+      mw.call(env)
+      first_asst = store.messages.find { |m| m[:info][:role] == "assistant" }
+      tool_part = first_asst[:parts].find { |p| p[:type] == "tool" }
+      tool_part[:state][:status].should == "completed"
+    end
+  end
+
+  it "records provider default_model when no override" do
+    with_tracking do |mw, store, _|
+      mw.call(build_env(input: "Hello", tool_results: nil))
+      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
+      asst[:info][:modelID].should == "mock-model"
+    end
+  end
+
+  it "records overridden model when env[:model] is set" do
+    with_tracking do |mw, store, _|
+      mw.call(build_env(input: "Hello", tool_results: nil, model: "custom-haiku"))
+      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
+      asst[:info][:modelID].should == "custom-haiku"
+    end
+  end
+
+  it "stores itself in env[:message_tracking]" do
+    with_tracking do |mw, _, _|
+      env = build_env(input: "Hello", tool_results: nil)
+      mw.call(env)
+      env[:message_tracking].should == mw
+    end
+  end
+
+  it "returns the inner app response unchanged" do
+    with_tracking do |mw, _, response|
+      result = mw.call(build_env(input: "Hello", tool_results: nil))
+      result.should == response
+    end
+  end
+
+  it "adds a step-finish part to assistant messages" do
+    with_tracking do |mw, store, _|
+      mw.call(build_env(input: "Hello", tool_results: nil))
+      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
+      step_finish = asst[:parts].find { |p| p[:type] == "step-finish" }
+      step_finish[:reason].should == "stop"
+    end
+  end
+end

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

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     module OTel
@@ -23,13 +26,13 @@ module Brute
           return @app.call(env) unless defined?(::OpenTelemetry::SDK)
 
           provider_name = provider_type(env[:provider])
-          model = begin; env[:context].model; rescue; nil; end
+          model = env[:model] || (env[:provider].default_model rescue nil)
           span_name = model ? "llm.call #{model}" : "llm.call"
 
           attributes = {
             "brute.provider" => provider_name,
             "brute.streaming" => !!env[:streaming],
-            "brute.context_messages" => env[:context].messages.to_a.size,
+            "brute.context_messages" => env[:messages].size,
           }
           attributes["brute.model"] = model.to_s if model
           attributes["brute.session_id"] = env[:metadata][:session_id].to_s if env.dig(:metadata, :session_id)
@@ -73,3 +76,30 @@ 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 through when OpenTelemetry::SDK is not defined" do
+    response = MockResponse.new(content: "hello from LLM")
+    middleware = Brute::Middleware::OTel::Span.new(->(_env) { response })
+    env = build_env
+    result = middleware.call(env)
+    result.should == response
+  end
+
+  it "env[:span] is nil when OTel is not defined" do
+    response = MockResponse.new(content: "hello from LLM")
+    middleware = Brute::Middleware::OTel::Span.new(->(_env) { response })
+    env = build_env
+    middleware.call(env)
+    env[:span].should.be.nil
+  end
+end

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

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     module OTel
@@ -28,3 +31,38 @@ 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::OTel::TokenUsage.new(->(_env) { response })
+    result = middleware.call(build_env)
+    result.should == response
+  end
+
+  it "passes through without error when span is nil" do
+    response = make_response
+    middleware = Brute::Middleware::OTel::TokenUsage.new(->(_env) { response })
+    lambda { middleware.call(build_env) }.should.not.raise
+  end
+
+  it "handles a response without usage gracefully" do
+    no_usage = Object.new
+    middleware = Brute::Middleware::OTel::TokenUsage.new(->(_env) { no_usage })
+    lambda { middleware.call(build_env) }.should.not.raise
+  end
+end

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

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     module OTel
@@ -15,7 +18,7 @@ module Brute
 
           span = env[:span]
           if span
-            functions = env[:context].functions
+            functions = env[:pending_functions]
             if functions && !functions.empty?
               span.set_attribute("brute.tool_calls.count", functions.size)
 
@@ -37,3 +40,29 @@ 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: "here's my plan")
+    middleware = Brute::Middleware::OTel::ToolCalls.new(->(_env) { response })
+    result = middleware.call(build_env)
+    result.should == response
+  end
+
+  it "passes through without error when span is nil with pending functions" do
+    response = MockResponse.new(content: "here's my plan")
+    fn = Struct.new(:name, :id, :arguments, keyword_init: true).new(name: "fs_read", id: "tc_001", arguments: { "path" => "/tmp" })
+    middleware = Brute::Middleware::OTel::ToolCalls.new(->(_env) { response })
+    result = middleware.call(build_env(pending_functions: [fn]))
+    result.should == response
+  end
+end

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

@@ -1,11 +1,14 @@
 # frozen_string_literal: true
 
+require "bundler/setup"
+require "brute"
+
 module Brute
   module Middleware
     module OTel
       # Records tool results being sent back to the LLM as span events.
       #
-      # Runs PRE-call: when env[:tool_results] is present, the orchestrator
+      # Runs PRE-call: when env[:tool_results] is present, the agent loop
       # is sending tool execution results back to the LLM. Each result gets
       # a span event with the tool name and success/error status.
       #
@@ -35,3 +38,28 @@ 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: "processed")
+    middleware = Brute::Middleware::OTel::ToolResults.new(->(_env) { response })
+    result = middleware.call(build_env)
+    result.should == response
+  end
+
+  it "passes through without error when span is nil" do
+    response = MockResponse.new(content: "processed")
+    middleware = Brute::Middleware::OTel::ToolResults.new(->(_env) { response })
+    result = middleware.call(build_env(tool_results: [["fs_read", { content: "data" }]]))
+    result.should == response
+  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.