brute 0.4.1 → 1.0.0

data/lib/brute/middleware/message_tracking.rb

@@ -1,9 +1,7 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
@@ -20,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
@@ -116,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|
@@ -170,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
@@ -213,174 +207,133 @@ module Brute
   end
 end
 
-if __FILE__ == $0
-  require_relative "../../../spec/spec_helper"
-
-  RSpec.describe Brute::Middleware::MessageTracking do
-    let(:tmpdir) { Dir.mktmpdir("brute_test_") }
-    let(:store) { Brute::MessageStore.new(session_id: "test-session", dir: tmpdir) }
-    let(:response) { MockResponse.new(content: "Hello from the LLM") }
-    let(:inner_app) { ->(_env) { response } }
-    let(:middleware) { described_class.new(inner_app, store: store) }
-
-    after { FileUtils.rm_rf(tmpdir) }
-
-    describe "user message recording" do
-      it "records a user message on the first call of a turn" do
-        env = build_env(input: "What is Ruby?", tool_results: nil)
-        middleware.call(env)
-
-        msgs = store.messages
-        user_msg = msgs.find { |m| m[:info][:role] == "user" }
-        expect(user_msg).not_to be_nil
-        expect(user_msg[:parts][0][:text]).to eq("What is Ruby?")
-      end
+test do
+  require_relative "../../../spec/support/mock_provider"
+  require_relative "../../../spec/support/mock_response"
+  require "tmpdir"
+  require "fileutils"
 
-      it "does not record a user message on subsequent calls (tool results)" do
-        env = build_env(input: "Hello", tool_results: nil)
-        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
 
-        env[:tool_results] = [["read", "file contents"]]
-        middleware.call(env)
+  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
 
-        user_msgs = store.messages.select { |m| m[:info][:role] == "user" }
-        expect(user_msgs.size).to eq(1)
-      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
 
-    describe "assistant message recording" do
-      it "records an assistant message after each LLM call" do
-        env = build_env(input: "Hello", tool_results: nil)
-        middleware.call(env)
-
-        msgs = store.messages
-        asst_msg = msgs.find { |m| m[:info][:role] == "assistant" }
-        expect(asst_msg).not_to be_nil
-        expect(asst_msg[:info][:parentID]).not_to be_nil
-      end
-
-      it "captures text content as a text part" do
-        env = build_env(input: "Hello", tool_results: nil)
-        middleware.call(env)
-
-        asst_msg = store.messages.find { |m| m[:info][:role] == "assistant" }
-        text_parts = asst_msg[:parts].select { |p| p[:type] == "text" }
-        expect(text_parts.size).to eq(1)
-        expect(text_parts[0][:text]).to eq("Hello from the LLM")
-      end
-
-      it "captures token usage from response" do
-        env = build_env(input: "Hello", tool_results: nil)
-        middleware.call(env)
-
-        asst_msg = store.messages.find { |m| m[:info][:role] == "assistant" }
-        expect(asst_msg[:info][:tokens][:input]).to eq(100)
-        expect(asst_msg[:info][:tokens][:output]).to eq(50)
-      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
 
-    describe "tool call recording" do
-      it "records tool calls as tool parts in running state" do
-        fn = double("function", id: "call_001", name: "read", arguments: { file_path: "/test" })
-        provider = MockProvider.new
-        ctx = LLM::Context.new(provider, tools: [])
-        allow(ctx).to receive(:functions).and_return([fn])
-        allow(provider).to receive(:complete).and_return(response)
-
-        env = build_env(input: "Read the file", tool_results: nil, context: ctx)
-        middleware.call(env)
-
-        asst_msg = store.messages.find { |m| m[:info][:role] == "assistant" }
-        tool_parts = asst_msg[:parts].select { |p| p[:type] == "tool" }
-        expect(tool_parts.size).to eq(1)
-        expect(tool_parts[0][:tool]).to eq("read")
-        expect(tool_parts[0][:callID]).to eq("call_001")
-        expect(tool_parts[0][:state][:status]).to eq("running")
-      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
 
-    describe "tool result completion" do
-      it "updates tool parts when results arrive" do
-        fn = double("function", id: "call_001", name: "read", arguments: { file_path: "/test" })
-        provider = MockProvider.new
-        ctx = LLM::Context.new(provider, tools: [])
-        allow(ctx).to receive(:functions).and_return([fn])
-        allow(provider).to receive(:complete).and_return(response)
-
-        env = build_env(input: "Read the file", tool_results: nil, context: ctx)
-        middleware.call(env)
-
-        allow(ctx).to receive(:functions).and_return([])
-        env[:tool_results] = [["read", "file contents here"]]
-        middleware.call(env)
-
-        msgs = store.messages
-        first_asst = msgs.find { |m| m[:info][:role] == "assistant" }
-        tool_part = first_asst[:parts].find { |p| p[:type] == "tool" }
-        expect(tool_part[:state][:status]).to eq("completed")
-        expect(tool_part[:state][:output]).to eq("file contents here")
-      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
 
-    describe "model name resolution" do
-      it "records the provider default_model when no override is set" do
-        env = build_env(input: "Hello", tool_results: nil)
-        middleware.call(env)
-
-        asst_msg = store.messages.find { |m| m[:info][:role] == "assistant" }
-        expect(asst_msg[:info][:modelID]).to eq("mock-model")
-      end
-
-      it "records the overridden model when context was created with model:" do
-        provider = MockProvider.new
-        ctx = LLM::Context.new(provider, tools: [], model: "custom-haiku-model")
-
-        env = build_env(input: "Hello", tool_results: nil, context: ctx, provider: provider)
-        middleware.call(env)
-
-        asst_msg = store.messages.find { |m| m[:info][:role] == "assistant" }
-        expect(asst_msg[:info][:modelID]).to eq("custom-haiku-model")
-      end
-
-      it "does not fall back to default_model when an override is present" do
-        provider = MockProvider.new
-        ctx = LLM::Context.new(provider, tools: [], model: "claude-3-haiku-20240307")
-
-        env = build_env(input: "Hello", tool_results: nil, context: ctx, provider: provider)
-        middleware.call(env)
+  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
 
-        asst_msg = store.messages.find { |m| m[:info][:role] == "assistant" }
-        expect(asst_msg[:info][:modelID]).not_to eq(provider.default_model)
-        expect(asst_msg[:info][:modelID]).to eq("claude-3-haiku-20240307")
-      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
 
-    describe "middleware passthrough" do
-      it "stores itself in env[:message_tracking]" do
-        env = build_env(input: "Hello", tool_results: nil)
-        middleware.call(env)
+  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
 
-        expect(env[:message_tracking]).to eq(middleware)
-      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 "returns the inner app response unchanged" do
-        env = build_env(input: "Hello", tool_results: nil)
-        result = middleware.call(env)
+  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
 
-        expect(result).to eq(response)
-      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
 
-    describe "step-finish parts" do
-      it "adds a step-finish part to each assistant message" do
-        env = build_env(input: "Hello", tool_results: nil)
-        middleware.call(env)
+  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
 
-        asst_msg = store.messages.find { |m| m[:info][:role] == "assistant" }
-        step_finish = asst_msg[:parts].find { |p| p[:type] == "step-finish" }
-        expect(step_finish).not_to be_nil
-        expect(step_finish[:reason]).to eq("stop")
-      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,9 +1,7 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
@@ -28,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)
@@ -79,108 +77,29 @@ module Brute
   end
 end
 
-if __FILE__ == $0
-  require_relative "../../../../spec/spec_helper"
-
-  RSpec.describe Brute::Middleware::OTel::Span do
-    let(:response) { MockResponse.new(content: "hello from LLM") }
-    let(:inner_app) { ->(_env) { response } }
-    let(:middleware) { described_class.new(inner_app) }
-
-    context "when OpenTelemetry::SDK is not defined" do
-      it "passes through to the inner app without touching env[:span]" do
-        hide_const("OpenTelemetry::SDK") if defined?(OpenTelemetry::SDK)
-
-        env = build_env
-        result = middleware.call(env)
-
-        expect(result).to eq(response)
-        expect(env[:span]).to be_nil
-      end
-    end
-
-    context "when OpenTelemetry::SDK is defined" do
-      let(:span) { mock_span }
-      let(:tracer) { double("tracer") }
-      let(:tracer_provider) { double("tracer_provider", tracer: tracer) }
-
-      before do
-        stub_const("OpenTelemetry::SDK", Module.new)
-        stub_const("OpenTelemetry::Trace::Status", Class.new {
-          def self.error(msg) = new(msg)
-          def initialize(msg = nil) = nil
-        })
+test do
+  require_relative "../../../../spec/support/mock_provider"
+  require_relative "../../../../spec/support/mock_response"
 
-        allow(tracer).to receive(:in_span) do |_name, **_opts, &block|
-          block.call(span)
-        end
-
-        allow(::OpenTelemetry).to receive(:tracer_provider).and_return(tracer_provider)
-
-        middleware.instance_variable_set(:@tracer, nil)
-      end
-
-      it "creates a span and sets env[:span] during the call" do
-        captured_span = nil
-        app = ->(env) { captured_span = env[:span]; response }
-        mw = described_class.new(app)
-        env = build_env
-
-        mw.call(env)
-
-        expect(captured_span).to eq(span)
-      end
-
-      it "cleans up env[:span] after the call" do
-        env = build_env
-        middleware.call(env)
-
-        expect(env[:span]).to be_nil
-      end
-
-      it "passes the response through" do
-        env = build_env
-        result = middleware.call(env)
-
-        expect(result).to eq(response)
-      end
-
-      it "creates a span with the tracer" do
-        env = build_env
-        middleware.call(env)
-
-        expect(tracer).to have_received(:in_span).with(
-          anything,
-          hash_including(
-            attributes: hash_including(
-              "brute.provider" => anything,
-              "brute.streaming" => false,
-              "brute.context_messages" => anything
-            ),
-            kind: :internal
-          )
-        )
-      end
-
-      it "records exceptions on the span and re-raises" do
-        error = RuntimeError.new("LLM exploded")
-        app = ->(_env) { raise error }
-        mw = described_class.new(app)
-        env = build_env
-
-        expect { mw.call(env) }.to raise_error(RuntimeError, "LLM exploded")
-        expect(span).to have_received(:record_exception).with(error)
-        expect(span).to have_received(:status=)
-      end
+  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 "cleans up env[:span] even on error" do
-        app = ->(_env) { raise "boom" }
-        mw = described_class.new(app)
-        env = build_env
+  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
 
-        expect { mw.call(env) }.to raise_error(RuntimeError)
-        expect(env[:span]).to be_nil
-      end
-    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,9 +1,7 @@
 # frozen_string_literal: true
 
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
+require "bundler/setup"
+require "brute"
 
 module Brute
   module Middleware
@@ -34,90 +32,37 @@ module Brute
   end
 end
 
-if __FILE__ == $0
-  require_relative "../../../../spec/spec_helper"
-
-  RSpec.describe Brute::Middleware::OTel::TokenUsage do
-    let(:response) do
-      MockResponse.new(
-        content: "hello",
-        usage: LLM::Usage.new(input_tokens: 100, output_tokens: 50, reasoning_tokens: 10, total_tokens: 160)
-      )
-    end
-    let(:inner_app) { ->(_env) { response } }
-    let(:middleware) { described_class.new(inner_app) }
-
-    it "passes the response through unchanged" do
-      env = build_env
-      result = middleware.call(env)
-      expect(result).to eq(response)
-    end
-
-    context "when env[:span] is nil" do
-      it "passes through without error" do
-        env = build_env
-        result = middleware.call(env)
-        expect(result).to eq(response)
-      end
-    end
-
-    context "when env[:span] is present" do
-      let(:span) { mock_span }
+test do
+  require_relative "../../../../spec/support/mock_provider"
+  require_relative "../../../../spec/support/mock_response"
 
-      it "sets token usage attributes on the span" do
-        env = build_env(span: span)
-        middleware.call(env)
-
-        expect(span).to have_received(:set_attribute).with("gen_ai.usage.input_tokens", 100)
-        expect(span).to have_received(:set_attribute).with("gen_ai.usage.output_tokens", 50)
-        expect(span).to have_received(:set_attribute).with("gen_ai.usage.total_tokens", 160)
-      end
-
-      it "sets reasoning_tokens when greater than zero" do
-        env = build_env(span: span)
-        middleware.call(env)
-
-        expect(span).to have_received(:set_attribute).with("gen_ai.usage.reasoning_tokens", 10)
-      end
-
-      it "omits reasoning_tokens when zero" do
-        zero_reasoning = MockResponse.new(
-          content: "hello",
-          usage: LLM::Usage.new(input_tokens: 100, output_tokens: 50, reasoning_tokens: 0, total_tokens: 150)
-        )
-        app = ->(_env) { zero_reasoning }
-        mw = described_class.new(app)
-        env = build_env(span: span)
-
-        mw.call(env)
-
-        expect(span).not_to have_received(:set_attribute).with("gen_ai.usage.reasoning_tokens", anything)
-      end
-
-      it "handles a response without usage gracefully" do
-        no_usage = double("response")
-        allow(no_usage).to receive(:respond_to?).with(:usage).and_return(false)
-        app = ->(_env) { no_usage }
-        mw = described_class.new(app)
-        env = build_env(span: span)
-
-        result = mw.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(result).to eq(no_usage)
-        expect(span).not_to have_received(:set_attribute)
-      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 "handles a response where usage returns nil" do
-        nil_usage = double("response", usage: nil)
-        allow(nil_usage).to receive(:respond_to?).with(:usage).and_return(true)
-        app = ->(_env) { nil_usage }
-        mw = described_class.new(app)
-        env = build_env(span: span)
+  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
 
-        mw.call(env)
+  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
 
-        expect(span).not_to have_received(:set_attribute)
-      end
-    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