brute 0.4.1 → 1.0.0

data/lib/brute/loop/step.rb

@@ -0,0 +1,332 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "brute"
+
+require "securerandom"
+require "async"
+
+module Brute
+  module Loop
+  # A first-class work object with identity, state, result/error capture,
+  # optional sub-queue, and cancellation.
+  #
+  # Users subclass Step and override #perform(task). The framework calls
+  # #call(task) which owns the state machine — subclasses never touch
+  # state transitions directly.
+  #
+  # State machine:
+  #
+  #       ┌──> completed
+  #       │
+  #   pending ──> running ──┤
+  #       │                 │
+  #       │                 ├──> failed
+  #       │                 │
+  #       └──> cancelled    └──> cancelled
+  #
+  # Three terminal states. Two non-terminal. Once terminal, stays terminal.
+  #
+  class Step
+    STATES = %i[pending running completed failed cancelled].freeze
+
+    attr_reader :id
+
+    def initialize(id: nil, **attributes)
+      @id         = id || self.class.generate_id
+      @attributes = attributes
+      @state      = :pending
+      @result     = nil
+      @error      = nil
+      @task       = nil
+      @jobs       = nil
+      @mutex      = Mutex.new
+    end
+
+    def self.generate_id
+      "#{name}-#{Process.pid}-#{Thread.current.object_id}-#{SecureRandom.hex(4)}"
+    end
+
+    # Called by the queue's worker. Subclasses override #perform instead.
+    def call(task)
+      return unless transition_to_running(task)
+
+      begin
+        result = perform(task)
+        @mutex.synchronize do
+          @result = result
+          @state  = :completed
+          @task   = nil
+        end
+
+      rescue Async::Cancel
+        # Cascade to sub-queue before we lose the reference:
+        @jobs&.cancel
+        @mutex.synchronize do
+          @state = :cancelled
+          @task  = nil
+        end
+        raise
+
+      rescue => error
+        # Continue-on-failure: record the error, do NOT re-raise.
+        @mutex.synchronize do
+          @error = error
+          @state = :failed
+          @task  = nil
+        end
+      end
+    end
+
+    # Subclasses override this.
+    def perform(task)
+      raise "#{self.class}#perform not implemented"
+    end
+
+    # Lazy accessor — creates the sub-queue parented to our running task.
+    # Only valid while the step is running (inside #perform).
+    def jobs(type: Brute::Queue::SequentialQueue)
+      @mutex.synchronize do
+        raise "Step not running; sub-queue has nothing to parent to" unless @task
+        @jobs ||= type.new(parent: @task).start
+      end
+    end
+
+    def state
+      @mutex.synchronize { @state }
+    end
+
+    def result
+      @mutex.synchronize { @result }
+    end
+
+    def error
+      @mutex.synchronize { @error }
+    end
+
+    def status
+      @mutex.synchronize do
+        { id: @id, state: @state, result: @result, error: @error }
+      end
+    end
+
+    def cancel
+      task = @mutex.synchronize do
+        case @state
+        when :pending
+          @state = :cancelled
+          nil
+        when :running
+          @task
+        else
+          return false # already finished
+        end
+      end
+
+      task&.cancel
+      true
+    end
+
+    private
+
+    def transition_to_running(task)
+      @mutex.synchronize do
+        return false if @state == :cancelled
+        @state = :running
+        @task  = task
+        true
+      end
+    end
+  end
+  end
+end
+
+test do
+  class HelloStep < Brute::Loop::Step
+    def perform(task)
+      "hello"
+    end
+  end
+
+  class FailStep < Brute::Loop::Step
+    def perform(task)
+      raise "boom"
+    end
+  end
+
+  class SlowStep < Brute::Loop::Step
+    def perform(task)
+      sleep 10
+      "done"
+    end
+  end
+
+  # -- identity --
+
+  it "generates a unique id" do
+    HelloStep.new.id.should.be.kind_of String
+  end
+
+  it "accepts a custom id" do
+    HelloStep.new(id: "custom-1").id.should == "custom-1"
+  end
+
+  # -- initial state --
+
+  it "starts in pending state" do
+    HelloStep.new.state.should == :pending
+  end
+
+  it "starts with nil result" do
+    HelloStep.new.result.should.be.nil
+  end
+
+  it "starts with nil error" do
+    HelloStep.new.error.should.be.nil
+  end
+
+  # -- successful execution --
+
+  it "transitions to completed on success" do
+    Sync do
+      step = HelloStep.new
+      step.call(Async::Task.current)
+      step.state.should == :completed
+    end
+  end
+
+  it "captures the return value as result" do
+    Sync do
+      step = HelloStep.new
+      step.call(Async::Task.current)
+      step.result.should == "hello"
+    end
+  end
+
+  # -- failed execution --
+
+  it "transitions to failed on error" do
+    Sync do
+      step = FailStep.new
+      step.call(Async::Task.current)
+      step.state.should == :failed
+    end
+  end
+
+  it "captures the exception as error" do
+    Sync do
+      step = FailStep.new
+      step.call(Async::Task.current)
+      step.error.message.should == "boom"
+    end
+  end
+
+  it "does not re-raise on failure" do
+    Sync do
+      step = FailStep.new
+      lambda { step.call(Async::Task.current) }.should.not.raise
+    end
+  end
+
+  # -- cancellation of pending step --
+
+  it "cancel returns true for pending step" do
+    HelloStep.new.cancel.should.be.true
+  end
+
+  it "transitions pending step to cancelled" do
+    step = HelloStep.new
+    step.cancel
+    step.state.should == :cancelled
+  end
+
+  it "skips perform when cancelled before call" do
+    Sync do
+      step = HelloStep.new
+      step.cancel
+      step.call(Async::Task.current)
+      step.result.should.be.nil
+    end
+  end
+
+  # -- cancellation of finished step --
+
+  it "cancel returns false for completed step" do
+    Sync do
+      step = HelloStep.new
+      step.call(Async::Task.current)
+      step.cancel.should.be.false
+    end
+  end
+
+  it "cancel returns false for failed step" do
+    Sync do
+      step = FailStep.new
+      step.call(Async::Task.current)
+      step.cancel.should.be.false
+    end
+  end
+
+  # -- status --
+
+  it "status includes id" do
+    step = HelloStep.new(id: "s1")
+    step.status[:id].should == "s1"
+  end
+
+  it "status includes state" do
+    step = HelloStep.new
+    step.status[:state].should == :pending
+  end
+
+  # -- perform not implemented --
+
+  it "raises NotImplementedError for base Step" do
+    Sync do
+      step = Brute::Loop::Step.new
+      step.call(Async::Task.current)
+      step.state.should == :failed
+    end
+  end
+
+  # -- jobs raises when not running --
+
+  it "raises when accessing jobs outside perform" do
+    lambda { HelloStep.new.jobs(type: Array) }.should.raise(RuntimeError)
+  end
+
+  # -- attributes stored --
+
+  it "stores attributes" do
+    step = HelloStep.new(url: "https://example.com")
+    step.instance_variable_get(:@attributes)[:url].should == "https://example.com"
+  end
+
+  # -- nested sub-queue --
+
+  describe "nesting" do
+    class ParentStep < Brute::Loop::Step
+      def perform(task)
+        3.times { |i| jobs(type: Brute::Queue::SequentialQueue) << HelloStep.new(id: "child-#{i}") }
+        jobs.drain
+        jobs.steps.map(&:result)
+      end
+    end
+
+    it "creates a sub-queue inside perform" do
+      Sync do
+        step = ParentStep.new
+        step.call(Async::Task.current)
+        step.result.should == ["hello", "hello", "hello"]
+      end
+    end
+
+    it "sub-steps all complete" do
+      Sync do
+        step = ParentStep.new
+        step.call(Async::Task.current)
+        step.instance_variable_get(:@jobs).steps.all? { |s| s.state == :completed }.should.be.true
+      end
+    end
+  end
+end

data/lib/brute/loop/tool_call_step.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "brute"
+
+module Brute
+  module Loop
+  # A Step that wraps an LLM::Function tool call.
+  #
+  # Identity comes from the function's call ID so tool results
+  # can be correlated back to the LLM's request.
+  #
+  class ToolCallStep < Step
+    attr_reader :function
+
+    def initialize(function:, **rest)
+      super(id: function.id, **rest)
+      @function = function
+    end
+
+    def perform(task)
+      @function.call
+    end
+  end
+  end
+end
+
+test do
+  FakeFunction = Struct.new(:id, :name, :arguments, :return_value) do
+    def call
+      return_value
+    end
+  end
+
+  class FailFunction
+    attr_reader :id, :name, :arguments
+    def initialize
+      @id = "fail_1"
+      @name = "fail"
+      @arguments = {}
+    end
+    def call
+      raise "tool exploded"
+    end
+  end
+
+  it "uses function id as step id" do
+    fn = FakeFunction.new("call_123", "read", {}, "content")
+    Brute::Loop::ToolCallStep.new(function: fn).id.should == "call_123"
+  end
+
+  it "calls the function in perform" do
+    Sync do
+      fn = FakeFunction.new("call_1", "read", {}, "file contents")
+      step = Brute::Loop::ToolCallStep.new(function: fn)
+      step.call(Async::Task.current)
+      step.result.should == "file contents"
+    end
+  end
+
+  it "transitions to completed on success" do
+    Sync do
+      fn = FakeFunction.new("call_2", "write", {}, "ok")
+      step = Brute::Loop::ToolCallStep.new(function: fn)
+      step.call(Async::Task.current)
+      step.state.should == :completed
+    end
+  end
+
+  it "captures function as accessor" do
+    fn = FakeFunction.new("call_3", "shell", {}, nil)
+    Brute::Loop::ToolCallStep.new(function: fn).function.should.be.identical_to fn
+  end
+
+  it "transitions to failed when function raises" do
+    Sync do
+      step = Brute::Loop::ToolCallStep.new(function: FailFunction.new)
+      step.call(Async::Task.current)
+      step.state.should == :failed
+    end
+  end
+
+  it "captures function error" do
+    Sync do
+      step = Brute::Loop::ToolCallStep.new(function: FailFunction.new)
+      step.call(Async::Task.current)
+      step.error.message.should == "tool exploded"
+    end
+  end
+end

data/lib/brute/middleware/compaction_check.rb

@@ -1,192 +1,106 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
     # Checks context size after each LLM call and triggers compaction
     # when thresholds are exceeded.
     #
-    # Runs POST-call: inspects message count and token usage from the
-    # response. If compaction is needed, summarizes older messages and
-    # rebuilds the context with the summary + recent messages.
+    # Runs POST-call: inspects message count and token usage. If compaction
+    # is needed, summarizes older messages and replaces env[:messages] with
+    # the summary so the next LLM call starts with a compact history.
     #
     class CompactionCheck < Base
-      def initialize(app, compactor:, system_prompt:, tools:, stream: nil)
+      def initialize(app, compactor:, system_prompt:)
         super(app)
         @compactor = compactor
         @system_prompt = system_prompt
-        @tools = tools
-        @stream = stream
       end
 
       def call(env)
         response = @app.call(env)
 
-        ctx = env[:context]
-        messages = ctx.messages.to_a.compact
-        usage = ctx.usage rescue nil
+        messages = env[:messages]
+        usage = env[:metadata].dig(:tokens, :last_call)
 
         if @compactor.should_compact?(messages, usage: usage)
           result = @compactor.compact(messages)
           if result
             summary_text, _recent = result
-            rebuild_context!(env, summary_text)
             env[:metadata][:compaction] = {
               messages_before: messages.size,
               timestamp: Time.now.iso8601,
             }
+            # Replace the message history with the summary
+            env[:messages] = [
+              LLM::Message.new(:system, @system_prompt),
+              LLM::Message.new(:user, "[Previous conversation summary]\n\n#{summary_text}"),
+            ]
           end
         end
 
         response
       end
-
-      private
-
-      def rebuild_context!(env, summary_text)
-        provider = env[:provider]
-        ctx_opts = { tools: @tools }
-        ctx_opts[:stream] = @stream if @stream
-        new_ctx = LLM::Context.new(provider, **ctx_opts)
-        prompt = new_ctx.prompt do |p|
-          p.system @system_prompt
-          p.user "[Previous conversation summary]\n\n#{summary_text}"
-        end
-        new_ctx.talk(prompt)
-        env[:context] = new_ctx
-      end
     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]
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
 
-      middleware.call(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(env[:context]).to equal(original_ctx)
-      expect(env[:metadata][:compaction]).to be_nil
+  def make_compactor(should: false, result: nil)
+    Object.new.tap do |c|
+      c.define_singleton_method(:should_compact?) { |_msgs, **_| should }
+      c.define_singleton_method(:compact) { |_msgs| result }
     end
+  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", []])
+  it "passes the response through when compaction is not needed" do
+    response = MockResponse.new(content: "compaction response")
+    compactor = make_compactor(should: false)
+    middleware = Brute::Middleware::CompactionCheck.new(->(_env) { response }, compactor: compactor, system_prompt: "sys")
+    result = middleware.call(build_env)
+    result.should == response
+  end
 
-        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)
+  it "does not set compaction metadata when not needed" do
+    compactor = make_compactor(should: false)
+    middleware = Brute::Middleware::CompactionCheck.new(->(_env) { MockResponse.new }, compactor: compactor, system_prompt: "sys")
+    env = build_env
+    middleware.call(env)
+    env[:metadata][:compaction].should.be.nil
+  end
 
-        env = build_env(
-          context: original_ctx,
-          provider: provider,
-          streaming: true,
-          callbacks: { on_content: callback },
-        )
-        middleware_with_stream.call(env)
+  it "replaces messages with summary when compaction triggers" do
+    compactor = make_compactor(should: true, result: ["Summary of conversation", []])
+    middleware = Brute::Middleware::CompactionCheck.new(->(_env) { MockResponse.new }, compactor: compactor, system_prompt: "sys")
+    env = build_env(messages: [LLM::Message.new(:user, "hello"), LLM::Message.new(:assistant, "hi"), LLM::Message.new(:user, "how")])
+    middleware.call(env)
+    env[:metadata][:compaction][:messages_before].should == 3
+  end
 
-        new_ctx = env[:context]
+  it "creates two messages after compaction" do
+    compactor = make_compactor(should: true, result: ["Summary", []])
+    middleware = Brute::Middleware::CompactionCheck.new(->(_env) { MockResponse.new }, compactor: compactor, system_prompt: "sys")
+    env = build_env(messages: [LLM::Message.new(:user, "hello")])
+    middleware.call(env)
+    env[:messages].size.should == 2
+  end
 
-        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
+  it "handles compactor returning nil gracefully" do
+    compactor = make_compactor(should: true, result: nil)
+    middleware = Brute::Middleware::CompactionCheck.new(->(_env) { MockResponse.new }, compactor: compactor, system_prompt: "sys")
+    env = build_env(messages: [LLM::Message.new(:user, "hello")])
+    middleware.call(env)
+    env[:metadata][:compaction].should.be.nil
   end
 end