brute 0.4.0 → 0.4.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e6fa4c53a825578634b110724522c021f089595e75e80faea05b5c53697010dd
-  data.tar.gz: 1cff09cf5e255928aada4f09a11c2f77ccf873839ee4f6d0ba24bc12beaefeba
+  metadata.gz: 87cc119e0dc26d2499af1ceab1fe26c0ea2a6b9685c2acb9718a29d155959971
+  data.tar.gz: 7887f2d6d2a3680cf660c93bbd64999f45e28afcf1f86fbd4fe52605fc0fae28
 SHA512:
-  metadata.gz: 795a6b851f47daba23755f8791f98c4c54f1c738704748767e70ff0bf25b797dca15fc25892642b7b46c7f6c8acab83d5dd110b0741e4252e8e8b1ce8798ffa1
-  data.tar.gz: 827d9628e7d5142fe1eaabc5e3de47cf04468afa5e1985a9af6b7ccc16e471ce35236953d3b746e988ef34a779df3cd4b1e6821ca9cd45815fc302785d8d1a00
+  metadata.gz: 4b04baad572cd024f4e7cf0b6b42fb3b8794e773918e3d4ee21d751ec0a296e4047903fdf290c68dc54aa9faf0556363b216c3e8f1eb23ce11c6e646bd7f14ca
+  data.tar.gz: c719e091120b55f5f0f93149c0502dc39cfdb984ae141a71260d1f1a8116824a6718e63639a092726bc83fe199c65b2cf1f98412c82991b13e06d297b488f092

data/lib/brute/agent_stream.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   # Bridges llm.rb's streaming callbacks to the host application.
   #
@@ -61,3 +66,116 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../spec/spec_helper"
+
+  RSpec.describe Brute::AgentStream do
+    # Build a mock tool that quacks like LLM::Function.
+    def mock_tool(id:, name:, arguments: {})
+      instance_double(LLM::Function,
+        id: id,
+        name: name,
+        arguments: arguments,
+      )
+    end
+
+    describe "#on_tool_call" do
+      it "records tool/error pair in pending_tools without spawning threads" do
+        stream = described_class.new
+        tool = mock_tool(id: "toolu_1", name: "read")
+
+        stream.on_tool_call(tool, nil)
+
+        expect(stream.pending_tools.size).to eq(1)
+        recorded_tool, recorded_error = stream.pending_tools.first
+        expect(recorded_tool).to eq(tool)
+        expect(recorded_error).to be_nil
+      end
+
+      it "records error tools in pending_tools" do
+        stream = described_class.new
+        tool = mock_tool(id: "toolu_err", name: "bad_tool")
+        error = LLM::Function::Return.new("toolu_err", "bad_tool", { error: true })
+
+        stream.on_tool_call(tool, error)
+
+        expect(stream.pending_tools.size).to eq(1)
+        _, recorded_error = stream.pending_tools.first
+        expect(recorded_error).to eq(error)
+      end
+
+      it "records pending tool call metadata for ToolUseGuard" do
+        stream = described_class.new
+        tool = mock_tool(
+          id: "toolu_abc",
+          name: "read",
+          arguments: { "file_path" => "test.rb" },
+        )
+
+        stream.on_tool_call(tool, nil)
+
+        calls = stream.pending_tool_calls
+        expect(calls).not_to be_empty
+        expect(calls.first).to include(
+          id: "toolu_abc",
+          name: "read",
+          arguments: { "file_path" => "test.rb" },
+        )
+      end
+
+      it "records metadata for multiple tool calls" do
+        stream = described_class.new
+        tool1 = mock_tool(id: "toolu_1", name: "read", arguments: { "file_path" => "a.rb" })
+        tool2 = mock_tool(id: "toolu_2", name: "write", arguments: { "file_path" => "b.rb" })
+
+        stream.on_tool_call(tool1, nil)
+        stream.on_tool_call(tool2, nil)
+
+        expect(stream.pending_tool_calls.size).to eq(2)
+        expect(stream.pending_tool_calls.map { |c| c[:id] }).to eq(["toolu_1", "toolu_2"])
+
+        expect(stream.pending_tools.size).to eq(2)
+        expect(stream.pending_tools.map { |t, _| t }).to eq([tool1, tool2])
+      end
+    end
+
+    describe "#clear_pending_tool_calls! and #clear_pending_tools!" do
+      it "empties both pending_tool_calls and pending_tools" do
+        stream = described_class.new
+        tool = mock_tool(id: "toolu_1", name: "read")
+
+        stream.on_tool_call(tool, nil)
+        expect(stream.pending_tool_calls).not_to be_empty
+        expect(stream.pending_tools).not_to be_empty
+
+        stream.clear_pending_tool_calls!
+        stream.clear_pending_tools!
+        expect(stream.pending_tool_calls).to be_empty
+        expect(stream.pending_tools).to be_empty
+      end
+    end
+
+    describe "#on_content" do
+      it "fires the content callback" do
+        received = nil
+        stream = described_class.new(on_content: ->(text) { received = text })
+
+        stream.on_content("hello")
+
+        expect(received).to eq("hello")
+      end
+    end
+
+    describe "#on_reasoning_content" do
+      it "fires the reasoning callback" do
+        received = nil
+        stream = described_class.new(on_reasoning: ->(text) { received = text })
+
+        stream.on_reasoning_content("thinking...")
+
+        expect(received).to eq("thinking...")
+      end
+    end
+  end
+end

data/lib/brute/diff.rb

@@ -24,3 +24,37 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../spec/spec_helper"
+
+  RSpec.describe Brute::Diff do
+    describe ".unified" do
+      it "generates a unified diff for changed content" do
+        old = "line1\nold\nline3\n"
+        new_text = "line1\nnew\nline3\n"
+        diff = described_class.unified(old, new_text)
+        expect(diff).to include("-old")
+        expect(diff).to include("+new")
+        expect(diff).to include("@@")
+      end
+
+      it "returns empty string for identical content" do
+        text = "same\ncontent\n"
+        expect(described_class.unified(text, text)).to eq("")
+      end
+
+      it "handles empty old content (new file)" do
+        diff = described_class.unified("", "new\ncontent\n")
+        expect(diff).to include("+new")
+        expect(diff).to include("+content")
+      end
+
+      it "handles empty new content (deleted file)" do
+        diff = described_class.unified("old\ncontent\n", "")
+        expect(diff).to include("-old")
+        expect(diff).to include("-content")
+      end
+    end
+  end
+end

data/lib/brute/message_store.rb

@@ -267,3 +267,197 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../spec/spec_helper"
+
+  require "tmpdir"
+
+  RSpec.describe Brute::MessageStore do
+    let(:tmpdir) { Dir.mktmpdir("brute_test_") }
+    let(:session_id) { "test-session-123" }
+    let(:store) { described_class.new(session_id: session_id, dir: tmpdir) }
+
+    after { FileUtils.rm_rf(tmpdir) }
+
+    describe "#append_user" do
+      it "creates a user message with text part" do
+        id = store.append_user(text: "Hello")
+
+        msg = store.message(id)
+        expect(msg[:info][:role]).to eq("user")
+        expect(msg[:info][:sessionID]).to eq(session_id)
+        expect(msg[:parts].size).to eq(1)
+        expect(msg[:parts][0][:type]).to eq("text")
+        expect(msg[:parts][0][:text]).to eq("Hello")
+      end
+
+      it "generates sequential message IDs" do
+        id1 = store.append_user(text: "First")
+        id2 = store.append_user(text: "Second")
+
+        expect(id1).to eq("msg_0001")
+        expect(id2).to eq("msg_0002")
+      end
+
+      it "persists to disk as JSON" do
+        id = store.append_user(text: "Persisted")
+
+        path = File.join(tmpdir, "#{id}.json")
+        expect(File.exist?(path)).to be true
+
+        data = JSON.parse(File.read(path), symbolize_names: true)
+        expect(data[:info][:role]).to eq("user")
+        expect(data[:parts][0][:text]).to eq("Persisted")
+      end
+    end
+
+    describe "#append_assistant" do
+      it "creates an assistant message" do
+        user_id = store.append_user(text: "Hi")
+        asst_id = store.append_assistant(parent_id: user_id, model_id: "claude", provider_id: "anthropic")
+
+        msg = store.message(asst_id)
+        expect(msg[:info][:role]).to eq("assistant")
+        expect(msg[:info][:parentID]).to eq(user_id)
+        expect(msg[:info][:modelID]).to eq("claude")
+        expect(msg[:info][:providerID]).to eq("anthropic")
+        expect(msg[:info][:tokens]).to include(input: 0, output: 0)
+        expect(msg[:parts]).to be_empty
+      end
+    end
+
+    describe "#add_text_part" do
+      it "appends a text part to an existing message" do
+        asst_id = store.append_assistant
+
+        store.add_text_part(message_id: asst_id, text: "Here is my response")
+
+        msg = store.message(asst_id)
+        expect(msg[:parts].size).to eq(1)
+        expect(msg[:parts][0][:type]).to eq("text")
+        expect(msg[:parts][0][:text]).to eq("Here is my response")
+      end
+    end
+
+    describe "#add_tool_part / #complete_tool_part / #error_tool_part" do
+      it "tracks tool lifecycle: running → completed" do
+        asst_id = store.append_assistant
+
+        store.add_tool_part(
+          message_id: asst_id,
+          tool: "read",
+          call_id: "call_001",
+          input: { file_path: "/tmp/test.rb" },
+        )
+
+        msg = store.message(asst_id)
+        tool_part = msg[:parts].find { |p| p[:type] == "tool" }
+        expect(tool_part[:tool]).to eq("read")
+        expect(tool_part[:state][:status]).to eq("running")
+
+        store.complete_tool_part(
+          message_id: asst_id,
+          call_id: "call_001",
+          output: "file contents here",
+        )
+
+        msg = store.message(asst_id)
+        tool_part = msg[:parts].find { |p| p[:type] == "tool" }
+        expect(tool_part[:state][:status]).to eq("completed")
+        expect(tool_part[:state][:output]).to eq("file contents here")
+        expect(tool_part[:state][:time][:end]).to be_a(Integer)
+      end
+
+      it "tracks tool lifecycle: running → error" do
+        asst_id = store.append_assistant
+
+        store.add_tool_part(
+          message_id: asst_id,
+          tool: "shell",
+          call_id: "call_002",
+          input: { command: "rm -rf /" },
+        )
+
+        store.error_tool_part(
+          message_id: asst_id,
+          call_id: "call_002",
+          error: "permission denied",
+        )
+
+        msg = store.message(asst_id)
+        tool_part = msg[:parts].find { |p| p[:type] == "tool" }
+        expect(tool_part[:state][:status]).to eq("error")
+        expect(tool_part[:state][:error]).to eq("permission denied")
+      end
+    end
+
+    describe "#complete_assistant" do
+      it "sets completion time and token counts" do
+        asst_id = store.append_assistant
+
+        store.complete_assistant(
+          message_id: asst_id,
+          tokens: { input: 100, output: 50, reasoning: 10, cache: { read: 20, write: 5 } },
+        )
+
+        msg = store.message(asst_id)
+        expect(msg[:info][:time][:completed]).to be_a(Integer)
+        expect(msg[:info][:tokens][:input]).to eq(100)
+        expect(msg[:info][:tokens][:output]).to eq(50)
+        expect(msg[:info][:tokens][:reasoning]).to eq(10)
+      end
+    end
+
+    describe "#messages" do
+      it "returns all messages in order" do
+        store.append_user(text: "Q1")
+        store.append_assistant
+        store.append_user(text: "Q2")
+
+        msgs = store.messages
+        expect(msgs.size).to eq(3)
+        expect(msgs[0][:info][:role]).to eq("user")
+        expect(msgs[1][:info][:role]).to eq("assistant")
+        expect(msgs[2][:info][:role]).to eq("user")
+      end
+    end
+
+    describe "#count" do
+      it "returns the number of stored messages" do
+        expect(store.count).to eq(0)
+
+        store.append_user(text: "Q1")
+        expect(store.count).to eq(1)
+
+        store.append_assistant
+        expect(store.count).to eq(2)
+      end
+    end
+
+    describe "loading from disk" do
+      it "restores messages from existing files" do
+        store.append_user(text: "Persisted Q")
+        asst_id = store.append_assistant(model_id: "claude")
+        store.add_text_part(message_id: asst_id, text: "Persisted A")
+
+        # Create a new store from the same directory
+        store2 = described_class.new(session_id: session_id, dir: tmpdir)
+
+        expect(store2.count).to eq(2)
+        expect(store2.messages[0][:parts][0][:text]).to eq("Persisted Q")
+        expect(store2.messages[1][:parts][0][:text]).to eq("Persisted A")
+      end
+
+      it "continues sequence numbering from loaded messages" do
+        store.append_user(text: "Q1")
+        store.append_user(text: "Q2")
+
+        store2 = described_class.new(session_id: session_id, dir: tmpdir)
+        id = store2.append_user(text: "Q3")
+
+        expect(id).to eq("msg_0003")
+      end
+    end
+  end
+end

data/lib/brute/middleware/compaction_check.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     # Checks context size after each LLM call and triggers compaction
@@ -57,3 +62,131 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../../spec/spec_helper"
+
+  RSpec.describe Brute::Middleware::CompactionCheck do
+    let(:response) { MockResponse.new(content: "compaction response") }
+    let(:inner_app) { ->(_env) { response } }
+    let(:compactor) { double("compactor") }
+    let(:system_prompt) { "You are a helpful assistant." }
+    let(:tools) { [] }
+    let(:middleware) do
+      described_class.new(inner_app, compactor: compactor, system_prompt: system_prompt, tools: tools)
+    end
+
+    it "passes the response through when compaction is not needed" do
+      allow(compactor).to receive(:should_compact?).and_return(false)
+      env = build_env
+
+      result = middleware.call(env)
+
+      expect(result).to eq(response)
+      expect(env[:metadata][:compaction]).to be_nil
+    end
+
+    it "does not replace context when compaction is not triggered" do
+      allow(compactor).to receive(:should_compact?).and_return(false)
+      env = build_env
+      original_ctx = env[:context]
+
+      middleware.call(env)
+
+      expect(env[:context]).to equal(original_ctx)
+    end
+
+    it "triggers compaction and rebuilds context when threshold is exceeded" do
+      allow(compactor).to receive(:should_compact?).and_return(true)
+      allow(compactor).to receive(:compact).and_return(["Summary of conversation", []])
+
+      provider = MockProvider.new
+      ctx = LLM::Context.new(provider, tools: [])
+      prompt = ctx.prompt { |p| p.system("sys"); p.user("hello") }
+      ctx.talk(prompt)
+
+      env = build_env(context: ctx, provider: provider)
+      middleware.call(env)
+
+      expect(env[:metadata][:compaction]).to include(:messages_before, :timestamp)
+      expect(env[:context]).not_to equal(ctx)
+    end
+
+    it "handles compactor returning nil gracefully" do
+      allow(compactor).to receive(:should_compact?).and_return(true)
+      allow(compactor).to receive(:compact).and_return(nil)
+
+      env = build_env
+      original_ctx = env[:context]
+
+      middleware.call(env)
+
+      expect(env[:context]).to equal(original_ctx)
+      expect(env[:metadata][:compaction]).to be_nil
+    end
+
+    context "when streaming is enabled" do
+      let(:stream) { double("AgentStream") }
+
+      let(:middleware_with_stream) do
+        described_class.new(inner_app,
+          compactor: compactor,
+          system_prompt: system_prompt,
+          tools: tools,
+          stream: stream,
+        )
+      end
+
+      it "preserves the stream parameter on the rebuilt context" do
+        allow(compactor).to receive(:should_compact?).and_return(true)
+        allow(compactor).to receive(:compact).and_return(["Summary of conversation", []])
+
+        provider = MockProvider.new
+        original_ctx = LLM::Context.new(provider, tools: [], stream: stream)
+        prompt = original_ctx.prompt { |p| p.system("sys"); p.user("hello") }
+        original_ctx.talk(prompt)
+
+        env = build_env(context: original_ctx, provider: provider, streaming: true)
+        middleware_with_stream.call(env)
+
+        new_ctx = env[:context]
+        expect(new_ctx).not_to equal(original_ctx)
+
+        ctx_params = new_ctx.instance_variable_get(:@params)
+        expect(ctx_params[:stream]).to eq(stream),
+          "Expected rebuilt context to have stream: #{stream.inspect} " \
+          "in @params, but got: #{ctx_params[:stream].inspect}. " \
+          "This causes on_content callbacks to silently stop firing after compaction."
+      end
+
+      it "fires on_content callback on the rebuilt context when streaming" do
+        received_content = nil
+        callback = ->(text) { received_content = text }
+
+        allow(compactor).to receive(:should_compact?).and_return(true)
+        allow(compactor).to receive(:compact).and_return(["Summary", []])
+
+        provider = MockProvider.new
+        original_ctx = LLM::Context.new(provider, tools: [], stream: stream)
+        prompt = original_ctx.prompt { |p| p.system("sys"); p.user("hello") }
+        original_ctx.talk(prompt)
+
+        env = build_env(
+          context: original_ctx,
+          provider: provider,
+          streaming: true,
+          callbacks: { on_content: callback },
+        )
+        middleware_with_stream.call(env)
+
+        new_ctx = env[:context]
+
+        ctx_params = new_ctx.instance_variable_get(:@params)
+        expect(ctx_params).to have_key(:stream),
+          "Rebuilt context is missing :stream in @params. " \
+          "LLMCall will skip the on_content fallback because env[:streaming] is true, " \
+          "so content from the next LLM call will be silently dropped."
+      end
+    end
+  end
+end

data/lib/brute/middleware/doom_loop_detection.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     # Detects when the agent is stuck repeating tool call patterns and injects
@@ -31,3 +36,98 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../../spec/spec_helper"
+
+  RSpec.describe Brute::Middleware::DoomLoopDetection do
+    let(:response) { MockResponse.new(content: "loop check") }
+    let(:inner_app) { ->(_env) { response } }
+
+    # Build a fake assistant message whose .functions returns the given list.
+    def assistant_msg_with_functions(function_list)
+      msg = LLM::Message.new(:assistant, "tool msg", {})
+      allow(msg).to receive(:functions).and_return(function_list)
+      msg
+    end
+
+    def fake_function(name:, arguments:)
+      double("fn", name: name, arguments: arguments)
+    end
+
+    it "passes through when no doom loop is detected" do
+      middleware = described_class.new(inner_app, threshold: 3)
+      env = build_env
+
+      result = middleware.call(env)
+
+      expect(result).to eq(response)
+      expect(env[:metadata][:doom_loop_detected]).to be_nil
+    end
+
+    it "detects consecutive identical tool calls" do
+      provider = MockProvider.new
+      ctx = LLM::Context.new(provider, tools: [])
+
+      fn = fake_function(name: "fs_read", arguments: '{"path":"x.rb"}')
+      messages = 4.times.map { assistant_msg_with_functions([fn]) }
+
+      allow(ctx).to receive(:messages).and_return(double("buffer", to_a: messages))
+      allow(ctx).to receive(:talk)
+
+      middleware = described_class.new(inner_app, threshold: 3)
+      env = build_env(context: ctx, provider: provider)
+
+      middleware.call(env)
+
+      expect(env[:metadata][:doom_loop_detected]).not_to be_nil
+    end
+
+    it "detects repeating sequences [A,B,A,B,A,B]" do
+      provider = MockProvider.new
+      ctx = LLM::Context.new(provider, tools: [])
+
+      fn_a = fake_function(name: "fs_read", arguments: '{"path":"a.rb"}')
+      fn_b = fake_function(name: "shell", arguments: '{"cmd":"ls"}')
+      messages = 3.times.flat_map do
+        [assistant_msg_with_functions([fn_a]), assistant_msg_with_functions([fn_b])]
+      end
+
+      allow(ctx).to receive(:messages).and_return(double("buffer", to_a: messages))
+      allow(ctx).to receive(:talk)
+
+      middleware = described_class.new(inner_app, threshold: 3)
+      env = build_env(context: ctx, provider: provider)
+
+      middleware.call(env)
+
+      expect(env[:metadata][:doom_loop_detected]).not_to be_nil
+    end
+
+    it "does not trigger below the threshold" do
+      provider = MockProvider.new
+      ctx = LLM::Context.new(provider, tools: [])
+
+      fn = fake_function(name: "fs_read", arguments: '{"path":"x.rb"}')
+      messages = 2.times.map { assistant_msg_with_functions([fn]) }
+
+      allow(ctx).to receive(:messages).and_return(double("buffer", to_a: messages))
+
+      middleware = described_class.new(inner_app, threshold: 3)
+      env = build_env(context: ctx, provider: provider)
+
+      middleware.call(env)
+
+      expect(env[:metadata][:doom_loop_detected]).to be_nil
+    end
+
+    describe Brute::DoomLoopDetector do
+      it "generates a warning message with repetition count" do
+        detector = described_class.new(threshold: 3)
+        msg = detector.warning_message(5)
+        expect(msg).to include("Doom loop detected")
+        expect(msg).to include("5 times")
+      end
+    end
+  end
+end