brute 0.1.7 → 0.1.9

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4f0f32487b029541fdb462f5f4958a95e4727150911b51f6d8ab457b875187d8
-  data.tar.gz: f162d75e227b4270e4a56dba42fe3cdddd23f9492adeff5cd1f345cbbb811961
+  metadata.gz: 5a36d054875f1465a0e9bfc380187c98ff08f837f8477f062f784652246d4256
+  data.tar.gz: a7c4df2710346a213b3ded2ee9be7d84b0bbe50d04c9e1ac9eaa679b9a35a7b2
 SHA512:
-  metadata.gz: 9aa172f042960dc5c9ec3250cf27781614077be7db114edc65d4ce72e178dfcab9ed9a9918d6583ed190d36cc54e865d8118bea9ac9814064a12a3f2b7ac8627
-  data.tar.gz: c0cc5addf257b161cc06385417a24833bf9f99a73038040abc67c3caef75ab0a0b9e3d357364b0fa0d194ca3b259555b80f09bf805fc12d17be1e4386d8b061b
+  metadata.gz: 7893f212130fc7dd94d80e3bf6b926d47ab86e37fc629866bbbec674121d4e8083791a2b66b4437af98d2cb694c0dc89655519844f5797fe1f17e8652d4ddd01
+  data.tar.gz: 1c48113815f9ad3f2d068252e851167b369e22ab8c1d5bd1d314387acab5f1e4169a49cde91f7c5bbb4847ade234779712f3f7b5d3fcc43b6c76b80478017a26

data/lib/brute/agent_stream.rb

@@ -8,11 +8,23 @@ module Brute
   # response is still streaming. on_tool_result fires as each thread finishes.
   #
   class AgentStream < LLM::Stream
-    def initialize(on_content: nil, on_reasoning: nil, on_tool_call: nil, on_tool_result: nil)
+    # Tool call metadata recorded during streaming, used by ToolUseGuard
+    # when ctx.functions is empty (nil-choice bug in llm.rb).
+    # Cleared by the guard after consumption to prevent stale data from
+    # causing duplicate synthetic assistant messages on subsequent calls.
+    attr_reader :pending_tool_calls
+
+    def clear_pending_tool_calls!
+      @pending_tool_calls.clear
+    end
+
+    def initialize(on_content: nil, on_reasoning: nil, on_tool_call: nil, on_tool_result: nil, on_question: nil)
       @on_content = on_content
       @on_reasoning = on_reasoning
       @on_tool_call = on_tool_call
       @on_tool_result = on_tool_result
+      @on_question = on_question
+      @pending_tool_calls = []
     end
 
     def on_content(text)
@@ -24,6 +36,7 @@ module Brute
     end
 
     def on_tool_call(tool, error)
+      @pending_tool_calls << { id: tool.id, name: tool.name, arguments: tool.arguments }
       @on_tool_call&.call(tool.name, tool.arguments)
 
       if error
@@ -38,8 +51,10 @@ module Brute
 
     def spawn_with_callback(tool)
       on_result = @on_tool_result
+      on_question = @on_question
       name = tool.name
       Thread.new do
+        Thread.current[:on_question] = on_question
         result = tool.call
         on_result&.call(name, result.respond_to?(:value) ? result.value : result)
         result

data/lib/brute/message_store.rb

@@ -0,0 +1,269 @@
+# frozen_string_literal: true
+
+require "json"
+require "fileutils"
+require "securerandom"
+
+module Brute
+  # Stores session messages as individual JSON files in the OpenCode
+  # {info, parts} format. Each session gets a directory; each message
+  # is a numbered JSON file inside it.
+  #
+  # Storage layout:
+  #
+  #   ~/.brute/sessions/{session-id}/
+  #     session.meta.json
+  #     msg_0001.json
+  #     msg_0002.json
+  #     ...
+  #
+  # Message format matches OpenCode's MessageV2.WithParts:
+  #
+  #   { info: { id:, sessionID:, role:, time:, ... },
+  #     parts: [{ id:, type:, ... }, ...] }
+  #
+  class MessageStore
+    attr_reader :session_id, :dir
+
+    def initialize(session_id:, dir: nil)
+      @session_id = session_id
+      @dir = dir || File.join(Dir.home, ".brute", "sessions", session_id)
+      @messages = {}   # id => { info:, parts: }
+      @seq = 0
+      @part_seq = 0
+      @mutex = Mutex.new
+      load_existing
+    end
+
+    # ── Append messages ──────────────────────────────────────────────
+
+    # Record a user message.
+    def append_user(text:, message_id: nil)
+      id = message_id || next_message_id
+      msg = {
+        info: {
+          id: id,
+          sessionID: @session_id,
+          role: "user",
+          time: { created: now_ms },
+        },
+        parts: [
+          { id: next_part_id, sessionID: @session_id, messageID: id,
+            type: "text", text: text },
+        ],
+      }
+      save_message(id, msg)
+      id
+    end
+
+    # Record the start of an assistant message. Returns the message ID.
+    # Call complete_assistant later to fill in tokens/timing.
+    def append_assistant(message_id: nil, parent_id: nil, model_id: nil, provider_id: nil)
+      id = message_id || next_message_id
+      msg = {
+        info: {
+          id: id,
+          sessionID: @session_id,
+          role: "assistant",
+          parentID: parent_id,
+          time: { created: now_ms },
+          modelID: model_id,
+          providerID: provider_id,
+          tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+          cost: 0.0,
+        },
+        parts: [],
+      }
+      save_message(id, msg)
+      id
+    end
+
+    # ── Parts ────────────────────────────────────────────────────────
+
+    # Add a text part to an existing message.
+    def add_text_part(message_id:, text:)
+      @mutex.synchronize do
+        msg = @messages[message_id]
+        return unless msg
+
+        part = { id: next_part_id, sessionID: @session_id, messageID: message_id,
+                 type: "text", text: text }
+        msg[:parts] << part
+        persist(message_id)
+        part[:id]
+      end
+    end
+
+    # Add a tool part in "running" state. Returns the part ID.
+    def add_tool_part(message_id:, tool:, call_id:, input:)
+      @mutex.synchronize do
+        msg = @messages[message_id]
+        return unless msg
+
+        part = {
+          id: next_part_id, sessionID: @session_id, messageID: message_id,
+          type: "tool", callID: call_id, tool: tool,
+          state: {
+            status: "running",
+            input: input,
+            time: { start: now_ms },
+          },
+        }
+        msg[:parts] << part
+        persist(message_id)
+        part[:id]
+      end
+    end
+
+    # Mark a tool part as completed with output.
+    def complete_tool_part(message_id:, call_id:, output:)
+      @mutex.synchronize do
+        msg = @messages[message_id]
+        return unless msg
+
+        part = msg[:parts].find { |p| p[:type] == "tool" && p[:callID] == call_id }
+        return unless part
+
+        part[:state][:status] = "completed"
+        part[:state][:output] = output
+        part[:state][:time][:end] = now_ms
+        persist(message_id)
+      end
+    end
+
+    # Mark a tool part as errored.
+    def error_tool_part(message_id:, call_id:, error:)
+      @mutex.synchronize do
+        msg = @messages[message_id]
+        return unless msg
+
+        part = msg[:parts].find { |p| p[:type] == "tool" && p[:callID] == call_id }
+        return unless part
+
+        part[:state][:status] = "error"
+        part[:state][:error] = error.to_s
+        part[:state][:time][:end] = now_ms
+        persist(message_id)
+      end
+    end
+
+    # Add a step-finish part to an assistant message.
+    def add_step_finish(message_id:, tokens: nil)
+      @mutex.synchronize do
+        msg = @messages[message_id]
+        return unless msg
+
+        part = {
+          id: next_part_id, sessionID: @session_id, messageID: message_id,
+          type: "step-finish",
+          reason: "stop",
+          tokens: tokens || { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+        }
+        msg[:parts] << part
+        persist(message_id)
+      end
+    end
+
+    # ── Complete / update ────────────────────────────────────────────
+
+    # Finalize an assistant message with token counts and completion time.
+    def complete_assistant(message_id:, tokens: nil)
+      @mutex.synchronize do
+        msg = @messages[message_id]
+        return unless msg
+
+        msg[:info][:time][:completed] = now_ms
+        if tokens
+          msg[:info][:tokens] = {
+            input: tokens[:input] || tokens[:total_input] || 0,
+            output: tokens[:output] || tokens[:total_output] || 0,
+            reasoning: tokens[:reasoning] || tokens[:total_reasoning] || 0,
+            cache: tokens[:cache] || { read: 0, write: 0 },
+          }
+        end
+        persist(message_id)
+      end
+    end
+
+    # ── Queries ──────────────────────────────────────────────────────
+
+    # All messages in order.
+    def messages
+      @mutex.synchronize { @messages.values }
+    end
+
+    # Single message by ID.
+    def message(id)
+      @mutex.synchronize { @messages[id] }
+    end
+
+    # Number of stored messages.
+    def count
+      @mutex.synchronize { @messages.size }
+    end
+
+    private
+
+    # ── ID generation ────────────────────────────────────────────────
+
+    def next_message_id
+      @seq += 1
+      format("msg_%04d", @seq)
+    end
+
+    def next_part_id
+      @part_seq += 1
+      format("prt_%04d", @part_seq)
+    end
+
+    def now_ms
+      (Time.now.to_f * 1000).to_i
+    end
+
+    # ── Persistence ──────────────────────────────────────────────────
+
+    def save_message(id, msg)
+      @mutex.synchronize do
+        @messages[id] = msg
+        persist(id)
+      end
+    end
+
+    def persist(id)
+      FileUtils.mkdir_p(@dir)
+      msg = @messages[id]
+      return unless msg
+
+      path = File.join(@dir, "#{id}.json")
+      File.write(path, JSON.pretty_generate(msg))
+    end
+
+    # Load any existing message files from disk on init.
+    def load_existing
+      return unless File.directory?(@dir)
+
+      Dir.glob(File.join(@dir, "msg_*.json")).sort.each do |path|
+        data = JSON.parse(File.read(path), symbolize_names: true)
+        id = data.dig(:info, :id)
+        next unless id
+
+        @messages[id] = data
+
+        # Track sequence numbers so new IDs don't collide
+        if (m = id.match(/\Amsg_(\d+)\z/))
+          n = m[1].to_i
+          @seq = n if n > @seq
+        end
+
+        # Track part sequences too
+        (data[:parts] || []).each do |part|
+          pid = part[:id]
+          if pid.is_a?(String) && (m = pid.match(/\Aprt_(\d+)\z/))
+            n = m[1].to_i
+            @part_seq = n if n > @part_seq
+          end
+        end
+      end
+    end
+  end
+end

data/lib/brute/middleware/compaction_check.rb

@@ -10,11 +10,12 @@ module Brute
     # rebuilds the context with the summary + recent messages.
     #
     class CompactionCheck < Base
-      def initialize(app, compactor:, system_prompt:, tools:)
+      def initialize(app, compactor:, system_prompt:, tools:, stream: nil)
         super(app)
         @compactor = compactor
         @system_prompt = system_prompt
         @tools = tools
+        @stream = stream
       end
 
       def call(env)
@@ -43,7 +44,9 @@ module Brute
 
       def rebuild_context!(env, summary_text)
         provider = env[:provider]
-        new_ctx = LLM::Context.new(provider, tools: @tools)
+        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}"

data/lib/brute/middleware/message_tracking.rb

@@ -0,0 +1,209 @@
+# frozen_string_literal: true
+
+module Brute
+  module Middleware
+    # Records every LLM exchange into a MessageStore in the OpenCode
+    # {info, parts} format so sessions can be viewed later.
+    #
+    # Lifecycle per pipeline call:
+    #
+    #   1. PRE-CALL  — if this is the first call of a turn (env[:tool_results]
+    #      is nil), record the user message.
+    #   2. POST-CALL — record the assistant message: text content as a "text"
+    #      part, each tool call as a "tool" part in "running" state.
+    #   3. When the pipeline is called again with tool results, update the
+    #      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.
+    #
+    class MessageTracking < Base
+      attr_reader :store
+
+      def initialize(app, store:)
+        super(app)
+        @store = store
+        @current_user_id = nil
+        @current_assistant_id = nil
+      end
+
+      def call(env)
+        env[:message_tracking] = self
+
+        # ── Pre-call: record user message or update tool results ──
+        if env[:tool_results].nil?
+          # New turn — record the user message
+          record_user_message(env)
+        else
+          # Tool results coming back — complete the tool parts
+          complete_tool_parts(env)
+        end
+
+        # ── LLM call ──
+        response = @app.call(env)
+
+        # ── Post-call: record assistant message ──
+        record_assistant_message(env, response)
+
+        response
+      end
+
+      # The current assistant message ID (used by external callbacks).
+      def current_assistant_id
+        @current_assistant_id
+      end
+
+      private
+
+      # ── User message ───────────────────────────────────────────────
+
+      def record_user_message(env)
+        text = extract_user_text(env)
+        return unless text
+
+        @current_user_id = @store.append_user(text: text)
+      end
+
+      def extract_user_text(env)
+        input = env[:input]
+        case input
+        when String
+          input
+        when Array
+          # llm.rb prompt format: array of message hashes
+          user_msg = input.reverse_each.find { |m| m.respond_to?(:role) && m.role.to_s == "user" }
+          user_msg&.content.to_s if user_msg
+        else
+          # Could be a prompt object — try to extract user content
+          if input.respond_to?(:messages)
+            msgs = input.messages.to_a
+            user_msg = msgs.reverse_each.find { |m| m.role.to_s == "user" }
+            user_msg&.content.to_s if user_msg
+          end
+        end
+      end
+
+      # ── Assistant message ──────────────────────────────────────────
+
+      def record_assistant_message(env, response)
+        provider_name = env[:provider]&.class&.name&.split("::")&.last&.downcase
+        model_name = resolve_model_name(env)
+
+        @current_assistant_id = @store.append_assistant(
+          parent_id: @current_user_id,
+          model_id: model_name,
+          provider_id: provider_name,
+        )
+
+        # Text content
+        text = safe_content(response)
+        @store.add_text_part(message_id: @current_assistant_id, text: text) if text && !text.empty?
+
+        # Tool calls
+        record_tool_calls(env)
+
+        # Token usage
+        tokens = extract_tokens(env, response)
+        @store.complete_assistant(message_id: @current_assistant_id, tokens: tokens) if tokens
+
+        # Step finish
+        @store.add_step_finish(message_id: @current_assistant_id, tokens: tokens)
+      end
+
+      def record_tool_calls(env)
+        ctx = env[:context]
+        functions = ctx.functions
+        return if functions.nil? || functions.empty?
+
+        functions.each do |fn|
+          @store.add_tool_part(
+            message_id: @current_assistant_id,
+            tool: fn.name,
+            call_id: fn.id,
+            input: fn.arguments,
+          )
+        end
+      end
+
+      # ── Tool results ───────────────────────────────────────────────
+
+      def complete_tool_parts(env)
+        return unless @current_assistant_id
+
+        results = env[:tool_results]
+        return unless results.is_a?(Array)
+
+        results.each do |name, value|
+          # Find the tool part by name (tool results come as [name, value] pairs)
+          msg = @store.message(@current_assistant_id)
+          next unless msg
+
+          # Match by tool name — find the first running tool part with this name
+          part = msg[:parts]&.find do |p|
+            p[:type] == "tool" && p[:tool] == name && p.dig(:state, :status) == "running"
+          end
+          next unless part
+
+          call_id = part[:callID]
+          if value.is_a?(Hash) && value[:error]
+            @store.error_tool_part(
+              message_id: @current_assistant_id,
+              call_id: call_id,
+              error: value[:error],
+            )
+          else
+            output = value.is_a?(String) ? value : value.to_s
+            @store.complete_tool_part(
+              message_id: @current_assistant_id,
+              call_id: call_id,
+              output: output,
+            )
+          end
+        end
+      end
+
+      # ── 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.
+      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
+
+        # Fall back to provider default
+        env[:provider]&.respond_to?(:default_model) ? env[:provider].default_model.to_s : nil
+      end
+
+      def safe_content(response)
+        return nil unless response.respond_to?(:content)
+        response.content
+      rescue NoMethodError
+        nil
+      end
+
+      def extract_tokens(env, response)
+        # Prefer the metadata accumulated by TokenTracking middleware
+        meta_tokens = env.dig(:metadata, :tokens, :last_call)
+        if meta_tokens
+          {
+            input: meta_tokens[:input] || 0,
+            output: meta_tokens[:output] || 0,
+            reasoning: 0,
+            cache: { read: 0, write: 0 },
+          }
+        elsif response.respond_to?(:usage) && (usage = response.usage)
+          {
+            input: usage.input_tokens.to_i,
+            output: usage.output_tokens.to_i,
+            reasoning: usage.reasoning_tokens.to_i,
+            cache: { read: 0, write: 0 },
+          }
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module Brute
+  module Middleware
+    module OTel
+      # Outermost OTel middleware. Creates a span per LLM pipeline call
+      # and passes it through env[:span] for inner OTel middlewares to
+      # decorate with events and attributes.
+      #
+      # When opentelemetry-sdk is not loaded, this is a pure pass-through.
+      #
+      # Pipeline position: outermost (wraps everything including retries).
+      #
+      #   use Brute::Middleware::OTel::Span
+      #   use Brute::Middleware::OTel::ToolResults
+      #   use Brute::Middleware::OTel::ToolCalls
+      #   use Brute::Middleware::OTel::TokenUsage
+      #   # ... existing middleware ...
+      #   run Brute::Middleware::LLMCall.new
+      #
+      class Span < Base
+        def call(env)
+          return @app.call(env) unless defined?(::OpenTelemetry::SDK)
+
+          provider_name = provider_type(env[:provider])
+          model = begin; env[:context].model; rescue; nil; end
+          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,
+          }
+          attributes["brute.model"] = model.to_s if model
+          attributes["brute.session_id"] = env[:metadata][:session_id].to_s if env.dig(:metadata, :session_id)
+
+          tracer.in_span(span_name, attributes: attributes, kind: :internal) do |span|
+            env[:span] = span
+            response = @app.call(env)
+
+            # Record response model if it differs from request model
+            resp_model = begin; response.model; rescue; nil; end
+            span.set_attribute("brute.response_model", resp_model.to_s) if resp_model && resp_model != model
+
+            response
+          rescue ::StandardError => e
+            span.record_exception(e)
+            span.status = ::OpenTelemetry::Trace::Status.error(e.message)
+            raise
+          ensure
+            env.delete(:span)
+          end
+        end
+
+        private
+
+        def tracer
+          @tracer ||= ::OpenTelemetry.tracer_provider.tracer("brute", Brute::VERSION)
+        end
+
+        def provider_type(provider)
+          name = provider.class.name.to_s.downcase
+          if name.include?("anthropic") then "anthropic"
+          elsif name.include?("openai") then "openai"
+          elsif name.include?("google") || name.include?("gemini") then "google"
+          elsif name.include?("deepseek") then "deepseek"
+          elsif name.include?("ollama") then "ollama"
+          elsif name.include?("xai") then "xai"
+          else "unknown"
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module Brute
+  module Middleware
+    module OTel
+      # Records token usage from the LLM response as span attributes.
+      #
+      # Runs POST-call: reads token counts from the response usage object
+      # and sets them as attributes on the span.
+      #
+      class TokenUsage < Base
+        def call(env)
+          response = @app.call(env)
+
+          span = env[:span]
+          if span && response.respond_to?(:usage) && (usage = response.usage)
+            span.set_attribute("gen_ai.usage.input_tokens", usage.input_tokens.to_i)
+            span.set_attribute("gen_ai.usage.output_tokens", usage.output_tokens.to_i)
+            span.set_attribute("gen_ai.usage.total_tokens", usage.total_tokens.to_i)
+
+            reasoning = usage.reasoning_tokens.to_i
+            span.set_attribute("gen_ai.usage.reasoning_tokens", reasoning) if reasoning > 0
+          end
+
+          response
+        end
+      end
+    end
+  end
+end