brute 0.4.0 → 0.4.1

data/lib/brute/middleware/llm_call.rb

@@ -39,3 +39,92 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../../spec/spec_helper"
+
+  RSpec.describe Brute::Middleware::LLMCall do
+    let(:provider) { MockProvider.new }
+    let(:middleware) { described_class.new }
+
+    it "calls ctx.talk with env[:input] and returns the response" do
+      ctx = LLM::Context.new(provider, tools: [])
+      prompt = ctx.prompt { |p| p.system("sys"); p.user("hello") }
+      env = build_env(context: ctx, provider: provider, input: prompt, streaming: false)
+
+      response = middleware.call(env)
+
+      expect(response).not_to be_nil
+      expect(provider.calls.size).to eq(1)
+    end
+
+    context "when not streaming" do
+      it "fires on_content callback with the response text" do
+        received_content = nil
+        callback = ->(text) { received_content = text }
+
+        response = MockResponse.new(content: "Hello world")
+        allow(provider).to receive(:complete).and_return(response)
+
+        ctx = LLM::Context.new(provider, tools: [])
+        prompt = ctx.prompt { |p| p.system("sys"); p.user("hi") }
+        env = build_env(
+          context: ctx,
+          provider: provider,
+          input: prompt,
+          streaming: false,
+          callbacks: { on_content: callback }
+        )
+
+        middleware.call(env)
+
+        expect(received_content).to eq("Hello world")
+      end
+    end
+
+    context "when streaming" do
+      it "does not fire on_content callback" do
+        callback_called = false
+        callback = ->(_text) { callback_called = true }
+
+        ctx = LLM::Context.new(provider, tools: [])
+        prompt = ctx.prompt { |p| p.system("sys"); p.user("hi") }
+        env = build_env(
+          context: ctx,
+          provider: provider,
+          input: prompt,
+          streaming: true,
+          callbacks: { on_content: callback }
+        )
+
+        middleware.call(env)
+
+        expect(callback_called).to be false
+      end
+    end
+
+    context "when response content raises NoMethodError (tool-only response)" do
+      it "does not crash and does not fire on_content" do
+        received_content = :not_called
+        callback = ->(text) { received_content = text }
+
+        bad_response = MockResponse.new(content: "")
+        allow(bad_response).to receive(:content).and_raise(NoMethodError)
+        allow(provider).to receive(:complete).and_return(bad_response)
+
+        ctx = LLM::Context.new(provider, tools: [])
+        prompt = ctx.prompt { |p| p.system("sys"); p.user("hi") }
+        env = build_env(
+          context: ctx,
+          provider: provider,
+          input: prompt,
+          streaming: false,
+          callbacks: { on_content: callback }
+        )
+
+        expect { middleware.call(env) }.not_to raise_error
+        expect(received_content).to eq(:not_called)
+      end
+    end
+  end
+end

data/lib/brute/middleware/message_tracking.rb

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     # Records every LLM exchange into a MessageStore in the OpenCode
@@ -207,3 +212,175 @@ module Brute
     end
   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
+
+      it "does not record a user message on subsequent calls (tool results)" do
+        env = build_env(input: "Hello", tool_results: nil)
+        middleware.call(env)
+
+        env[:tool_results] = [["read", "file contents"]]
+        middleware.call(env)
+
+        user_msgs = store.messages.select { |m| m[:info][:role] == "user" }
+        expect(user_msgs.size).to eq(1)
+      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
+    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
+    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
+    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)
+
+        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
+    end
+
+    describe "middleware passthrough" do
+      it "stores itself in env[:message_tracking]" do
+        env = build_env(input: "Hello", tool_results: nil)
+        middleware.call(env)
+
+        expect(env[:message_tracking]).to eq(middleware)
+      end
+
+      it "returns the inner app response unchanged" do
+        env = build_env(input: "Hello", tool_results: nil)
+        result = middleware.call(env)
+
+        expect(result).to eq(response)
+      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)
+
+        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
+    end
+  end
+end

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

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     module OTel
@@ -73,3 +78,109 @@ module Brute
     end
   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
+        })
+
+        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
+
+      it "cleans up env[:span] even on error" do
+        app = ->(_env) { raise "boom" }
+        mw = described_class.new(app)
+        env = build_env
+
+        expect { mw.call(env) }.to raise_error(RuntimeError)
+        expect(env[:span]).to be_nil
+      end
+    end
+  end
+end

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

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     module OTel
@@ -28,3 +33,91 @@ module Brute
     end
   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 }
+
+      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)
+
+        expect(result).to eq(no_usage)
+        expect(span).not_to have_received(:set_attribute)
+      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)
+
+        mw.call(env)
+
+        expect(span).not_to have_received(:set_attribute)
+      end
+    end
+  end
+end

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

@@ -1,5 +1,10 @@
 # frozen_string_literal: true
 
+if __FILE__ == $0
+  require "bundler/setup"
+  require "brute"
+end
+
 module Brute
   module Middleware
     module OTel
@@ -37,3 +42,111 @@ module Brute
     end
   end
 end
+
+if __FILE__ == $0
+  require_relative "../../../../spec/spec_helper"
+
+  RSpec.describe Brute::Middleware::OTel::ToolCalls do
+    let(:response) { MockResponse.new(content: "here's my plan") }
+    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 even with pending functions" do
+        ctx = build_env[:context]
+        fn = double("function", name: "fs_read", id: "tc_001", arguments: { "path" => "/tmp" })
+        allow(ctx).to receive(:functions).and_return([fn])
+
+        env = build_env(context: ctx)
+        result = middleware.call(env)
+        expect(result).to eq(response)
+      end
+    end
+
+    context "when env[:span] is present" do
+      let(:span) { mock_span }
+
+      it "does nothing when there are no pending functions" do
+        ctx = build_env[:context]
+        allow(ctx).to receive(:functions).and_return([])
+
+        env = build_env(context: ctx, span: span)
+        middleware.call(env)
+
+        expect(span).not_to have_received(:add_event)
+        expect(span).not_to have_received(:set_attribute)
+      end
+
+      it "does nothing when functions is nil" do
+        ctx = build_env[:context]
+        allow(ctx).to receive(:functions).and_return(nil)
+
+        env = build_env(context: ctx, span: span)
+        middleware.call(env)
+
+        expect(span).not_to have_received(:add_event)
+      end
+
+      it "records a tool_call event per pending function" do
+        ctx = build_env[:context]
+        fn1 = double("function", name: "fs_read", id: "tc_001", arguments: { "path" => "/src/main.rb" })
+        fn2 = double("function", name: "shell", id: "tc_002", arguments: { "command" => "rspec" })
+        allow(ctx).to receive(:functions).and_return([fn1, fn2])
+
+        env = build_env(context: ctx, span: span)
+        middleware.call(env)
+
+        expect(span).to have_received(:set_attribute).with("brute.tool_calls.count", 2)
+        expect(span).to have_received(:add_event).with(
+          "tool_call",
+          attributes: hash_including(
+            "tool.name" => "fs_read",
+            "tool.id" => "tc_001"
+          )
+        )
+        expect(span).to have_received(:add_event).with(
+          "tool_call",
+          attributes: hash_including(
+            "tool.name" => "shell",
+            "tool.id" => "tc_002"
+          )
+        )
+      end
+
+      it "serializes arguments as JSON" do
+        ctx = build_env[:context]
+        args = { "path" => "/tmp/test.rb", "content" => "puts 'hi'" }
+        fn = double("function", name: "fs_write", id: "tc_003", arguments: args)
+        allow(ctx).to receive(:functions).and_return([fn])
+
+        env = build_env(context: ctx, span: span)
+        middleware.call(env)
+
+        expect(span).to have_received(:add_event).with(
+          "tool_call",
+          attributes: hash_including("tool.arguments" => args.to_json)
+        )
+      end
+
+      it "handles nil arguments" do
+        ctx = build_env[:context]
+        fn = double("function", name: "todo_read", id: "tc_004", arguments: nil)
+        allow(ctx).to receive(:functions).and_return([fn])
+
+        env = build_env(context: ctx, span: span)
+        middleware.call(env)
+
+        expect(span).to have_received(:add_event).with(
+          "tool_call",
+          attributes: { "tool.name" => "todo_read", "tool.id" => "tc_004" }
+        )
+      end
+    end
+  end
+end