brute 3.0.1 → 3.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c25e73819e947f92a37f144411ba27b09f93cc48088bc0cc474247bc81149658
-  data.tar.gz: 635eff80312c14656cf2a6168f2f4ee9279413ac285e580a58e959812a9e5d4d
+  metadata.gz: 1fe7f1994a944255146721ea3030261622735e33d412f07ba6fd118287b76517
+  data.tar.gz: 0b36c9d8c38afb152c16f0bc26c305145ed94f1c53e84316399d4f32b7f8752a
 SHA512:
-  metadata.gz: c43c126ae04f4642df67df70ab780dfc11ef2070709580cc409d6b32652a2f0be41730d252a06d9029f8d156657877f046af6a06d5c2ba1d4ee9a2d5fda75d97
-  data.tar.gz: 26599b1c63616c87af1c8537dfe4e473be5a70ee408058b31fb8a67c5203f72b81d41a868e59ceb8831ac570a74a9132c6c2df42794b116ea91e67b70a4441dc
+  metadata.gz: 92f1aaba43df5f7fc0dad527c2b52a47e5336d4efabd95a6211cd9cdff97ae1a583d29ab94364d48b56760c5d7126fd3f9fe8f6e51b16518a250178daac30ee5
+  data.tar.gz: fce06871143bd87726637ece09887261b87aec3042470c0bb38ace70db42546541ef6a12f0310a3f72433a6ab65ade8f8c894a70a2f865ba21d752e94145f235

data/lib/brute/middleware/008_checkpoint.rb

@@ -0,0 +1,268 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "brute"
+require "json"
+require "securerandom"
+require "fileutils"
+
+module Brute
+  module Middleware
+    # Durable execution for the tool loop. Where SessionLog persists the
+    # conversation once per turn (outermost), Checkpoint snapshots it after
+    # every pass through the inner stack — one checkpoint per LLM call + tool
+    # execution. Place it just inside Loop::ToolResult:
+    #
+    #   use Brute::Middleware::Loop::ToolResult
+    #   use Brute::Middleware::Checkpoint, path: "tmp/checkpoints.jsonl"
+    #   use Brute::Middleware::MaxIterations
+    #   use Brute::Middleware::ToolPipeline, tools: Brute::Tools::ALL
+    #
+    # The store is just a JSONL log of snapshots — one line per checkpoint,
+    # each carrying the full message log plus its own id and the id of the
+    # parent checkpoint it grew from. That append-only chain buys three
+    # things:
+    #
+    #   resume       pass resume: :latest — a crash mid-turn costs at most
+    #                one iteration instead of the whole turn
+    #   time travel  pass resume: "<checkpoint id>" — restart from any
+    #                snapshot in the chain
+    #   forking      checkpoints written after a time-travel resume carry the
+    #                resumed id as parent_id, branching the chain in place
+    #
+    # System messages are not persisted (SystemPrompt re-adds them each
+    # turn); restored history is inserted after any leading system message
+    # and before the current turn's input.
+    class Checkpoint
+      def initialize(app, path:, resume: nil)
+        @app    = app
+        @path   = path
+        @resume = resume
+      end
+
+      def call(env)
+        restore(env) unless env[:metadata][:checkpoint]
+        @app.call(env)
+        persist(env)
+        env
+      end
+
+      # Parsed checkpoint records (symbol keys, messages as plain hashes),
+      # oldest first.
+      def self.list(path)
+        return [] unless path && File.exist?(path)
+
+        File.foreach(path).filter_map do |line|
+          line = line.strip
+          JSON.parse(line, symbolize_names: true) unless line.empty?
+        end
+      end
+
+      private
+
+        def restore(env)
+          env[:metadata][:checkpoint] = { path: @path }
+          return unless @resume
+
+          record = find_record
+          if record.nil?
+            return if @resume == :latest
+
+            raise KeyError, "no checkpoint #{@resume.inspect} in #{@path}"
+          end
+
+          history = record[:messages].map { |h| Brute::Message.new(**h) }
+          index = env[:messages].index { |m| m.role != :system } || env[:messages].size
+          env[:messages].insert(index, *history)
+          env[:metadata][:checkpoint][:id] = record[:id]
+        end
+
+        def find_record
+          records = self.class.list(@path)
+          @resume == :latest ? records.last : records.find { |r| r[:id] == @resume.to_s }
+        end
+
+        def persist(env)
+          FileUtils.mkdir_p(File.dirname(@path))
+          record = {
+            id:        SecureRandom.hex(6),
+            parent_id: env[:metadata][:checkpoint][:id],
+            ts:        Time.now.utc.iso8601,
+            iteration: env[:current_iteration],
+            messages:  env[:messages].reject { |m| m.role == :system }.map(&:to_h),
+          }
+          File.open(@path, "a") { |f| f.puts(JSON.generate(record)) }
+          env[:metadata][:checkpoint][:id] = record[:id]
+        end
+    end
+  end
+end
+
+__END__
+
+describe "brute/middleware/008_checkpoint" do
+  require "brute/messages"
+  require "tmpdir"
+
+  it "snapshots after every pass through the loop" do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, "cp.jsonl")
+      calls = 0
+      inner = ->(env) do
+        calls += 1
+        if calls < 3
+          env[:messages] << Brute::Message.new(role: :tool, content: "r#{calls}", tool_call_id: "tc#{calls}")
+        else
+          env[:messages].assistant("done")
+        end
+      end
+
+      stack = Brute::Middleware::Loop::ToolResult.new(
+        Brute::Middleware::Checkpoint.new(inner, path: path)
+      )
+      env = { messages: Brute.log.tap { |l| l.user("go") }, metadata: {}, current_iteration: 1 }
+      stack.call(env)
+
+      records = Brute::Middleware::Checkpoint.list(path)
+      records.size.should == 3
+      records.map { |r| r[:iteration] }.should == [1, 2, 3]
+      records.last[:messages].last[:content].should == "done"
+    end
+  end
+
+  it "chains checkpoints via parent_id" do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, "cp.jsonl")
+      mw = Brute::Middleware::Checkpoint.new(->(env) { env[:messages].assistant("a") }, path: path)
+      env = { messages: Brute.log.tap { |l| l.user("go") }, metadata: {}, current_iteration: 1 }
+      mw.call(env)
+      mw.call(env)
+
+      records = Brute::Middleware::Checkpoint.list(path)
+      records[0][:parent_id].should == nil
+      records[1][:parent_id].should == records[0][:id]
+    end
+  end
+
+  it "resume: :latest restores the conversation before this turn's input" do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, "cp.jsonl")
+      Brute::Middleware::Checkpoint.new(->(env) { env[:messages].assistant("first reply") }, path: path)
+        .call({ messages: Brute.log.tap { |l| l.user("hi") }, metadata: {}, current_iteration: 1 })
+
+      env = { messages: Brute.log.tap { |l| l.user("again") }, metadata: {}, current_iteration: 1 }
+      Brute::Middleware::Checkpoint.new(->(_e) {}, path: path, resume: :latest).call(env)
+
+      env[:messages].map(&:content).should == ["hi", "first reply", "again"]
+    end
+  end
+
+  it "inserts restored history after the system message" do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, "cp.jsonl")
+      Brute::Middleware::Checkpoint.new(->(env) { env[:messages].assistant("old") }, path: path)
+        .call({ messages: Brute.log.tap { |l| l.user("hi") }, metadata: {}, current_iteration: 1 })
+
+      env = { messages: Brute.log, metadata: {}, current_iteration: 1 }
+      env[:messages].system("rules")
+      env[:messages].user("new input")
+      Brute::Middleware::Checkpoint.new(->(_e) {}, path: path, resume: :latest).call(env)
+
+      env[:messages].map(&:role).should == [:system, :user, :assistant, :user]
+      env[:messages].last.content.should == "new input"
+    end
+  end
+
+  it "resumes from a specific checkpoint id (time travel)" do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, "cp.jsonl")
+      mw = Brute::Middleware::Checkpoint.new(->(env) { env[:messages].assistant("reply #{env[:messages].size}") }, path: path)
+      env = { messages: Brute.log.tap { |l| l.user("go") }, metadata: {}, current_iteration: 1 }
+      mw.call(env)
+      mw.call(env)
+
+      first = Brute::Middleware::Checkpoint.list(path).first
+      env2 = { messages: Brute.log, metadata: {}, current_iteration: 1 }
+      Brute::Middleware::Checkpoint.new(->(_e) {}, path: path, resume: first[:id]).call(env2)
+
+      env2[:messages].size.should == first[:messages].size
+    end
+  end
+
+  it "forks: checkpoints after a time-travel resume carry the resumed id as parent" do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, "cp.jsonl")
+      mw = Brute::Middleware::Checkpoint.new(->(env) { env[:messages].assistant("a") }, path: path)
+      env = { messages: Brute.log.tap { |l| l.user("go") }, metadata: {}, current_iteration: 1 }
+      mw.call(env)
+      mw.call(env)
+
+      root = Brute::Middleware::Checkpoint.list(path).first
+      Brute::Middleware::Checkpoint.new(->(env) { env[:messages].assistant("fork") }, path: path, resume: root[:id])
+        .call({ messages: Brute.log, metadata: {}, current_iteration: 1 })
+
+      Brute::Middleware::Checkpoint.list(path).last[:parent_id].should == root[:id]
+    end
+  end
+
+  it "resume: :latest with no store yet starts fresh" do
+    Dir.mktmpdir do |dir|
+      env = { messages: Brute.log.tap { |l| l.user("go") }, metadata: {}, current_iteration: 1 }
+      Brute::Middleware::Checkpoint.new(->(_e) {}, path: File.join(dir, "cp.jsonl"), resume: :latest).call(env)
+      env[:messages].map(&:content).should == ["go"]
+    end
+  end
+
+  it "raises on an unknown checkpoint id" do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, "cp.jsonl")
+      Brute::Middleware::Checkpoint.new(->(_e) {}, path: path).call(
+        { messages: Brute.log, metadata: {}, current_iteration: 1 }
+      )
+
+      lambda do
+        Brute::Middleware::Checkpoint.new(->(_e) {}, path: path, resume: "nope").call(
+          { messages: Brute.log, metadata: {}, current_iteration: 1 }
+        )
+      end.should.raise(KeyError)
+    end
+  end
+
+  it "does not persist system messages" do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, "cp.jsonl")
+      env = { messages: Brute.log, metadata: {}, current_iteration: 1 }
+      env[:messages].system("rules")
+      env[:messages].user("go")
+      Brute::Middleware::Checkpoint.new(->(_e) {}, path: path).call(env)
+
+      roles = Brute::Middleware::Checkpoint.list(path).last[:messages].map { |m| m[:role] }
+      roles.should == ["user"]
+    end
+  end
+
+  it "restores only once even though the loop re-enters" do
+    Dir.mktmpdir do |dir|
+      path = File.join(dir, "cp.jsonl")
+      Brute::Middleware::Checkpoint.new(->(env) { env[:messages].assistant("old") }, path: path)
+        .call({ messages: Brute.log.tap { |l| l.user("hi") }, metadata: {}, current_iteration: 1 })
+
+      calls = 0
+      inner = ->(env) do
+        calls += 1
+        if calls < 2
+          env[:messages] << Brute::Message.new(role: :tool, content: "r", tool_call_id: "tc1")
+        else
+          env[:messages].assistant("done")
+        end
+      end
+      stack = Brute::Middleware::Loop::ToolResult.new(
+        Brute::Middleware::Checkpoint.new(inner, path: path, resume: :latest)
+      )
+      env = { messages: Brute.log.tap { |l| l.user("again") }, metadata: {}, current_iteration: 1 }
+      stack.call(env)
+
+      env[:messages].count { |m| m.content == "old" }.should == 1
+    end
+  end
+end

data/lib/brute/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Brute
-  VERSION = "3.0.1"
+  VERSION = "3.1.0"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: brute
 version: !ruby/object:Gem::Version
-  version: 3.0.1
+  version: 3.1.0
 platform: ruby
 authors:
 - Brute Contributors
@@ -184,6 +184,7 @@ files:
 - lib/brute/middleware/004_summarize.rb
 - lib/brute/middleware/005_tracing.rb
 - lib/brute/middleware/006_loop.rb
+- lib/brute/middleware/008_checkpoint.rb
 - lib/brute/middleware/010_max_iterations.rb
 - lib/brute/middleware/015_otel_token_usage.rb
 - lib/brute/middleware/020_system_prompt.rb