brute 0.1.7 → 0.1.8

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4f0f32487b029541fdb462f5f4958a95e4727150911b51f6d8ab457b875187d8
-  data.tar.gz: f162d75e227b4270e4a56dba42fe3cdddd23f9492adeff5cd1f345cbbb811961
+  metadata.gz: 07be8ba156b49a76de7dda83633348bc4f9bc87552b29846a5bbf11be83c13c2
+  data.tar.gz: a3643e4bae2399cf0b92cd52f9b8083f2cc0bfa78da64228e33cb3fd6504a54f
 SHA512:
-  metadata.gz: 9aa172f042960dc5c9ec3250cf27781614077be7db114edc65d4ce72e178dfcab9ed9a9918d6583ed190d36cc54e865d8118bea9ac9814064a12a3f2b7ac8627
-  data.tar.gz: c0cc5addf257b161cc06385417a24833bf9f99a73038040abc67c3caef75ab0a0b9e3d357364b0fa0d194ca3b259555b80f09bf805fc12d17be1e4386d8b061b
+  metadata.gz: 1c6f739d3d488d5c2dd728bd7144ccb27793a592dcda4628edb228c0c3c77d6f961d38877a249b6c6d9a009ee88bc8d2d7e0ab4651d6319073e18d3edb29c35c
+  data.tar.gz: 1da1558c7b9e04027003b5bdff82d0dd894a1643e1e6f15b8d57b4123d44367e8ad0abbd7780a7927b71710f94343ad32186dd4e68ea2e74f80c067f3dd70d6e

data/lib/brute/agent_stream.rb

@@ -8,11 +8,22 @@ module Brute
   # response is still streaming. on_tool_result fires as each thread finishes.
   #
   class AgentStream < LLM::Stream
+    # 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_content = on_content
       @on_reasoning = on_reasoning
       @on_tool_call = on_tool_call
       @on_tool_result = on_tool_result
+      @pending_tool_calls = []
     end
 
     def on_content(text)
@@ -24,6 +35,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

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,195 @@
+# 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 = env[:provider]&.respond_to?(:default_model) ? env[:provider].default_model.to_s : nil
+
+        @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 ────────────────────────────────────────────────────
+
+      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

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

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Brute
+  module Middleware
+    module OTel
+      # Records tool calls the LLM requested as span events.
+      #
+      # Runs POST-call: after the LLM responds, inspects ctx.functions
+      # for any tool calls the model wants to make, and adds a span event
+      # for each one with the tool name, call ID, and arguments.
+      #
+      class ToolCalls < Base
+        def call(env)
+          response = @app.call(env)
+
+          span = env[:span]
+          if span
+            functions = env[:context].functions
+            if functions && !functions.empty?
+              span.set_attribute("brute.tool_calls.count", functions.size)
+
+              functions.each do |fn|
+                attrs = {
+                  "tool.name" => fn.name.to_s,
+                  "tool.id" => fn.id.to_s,
+                }
+                args = fn.arguments
+                attrs["tool.arguments"] = args.to_json if args
+                span.add_event("tool_call", attributes: attrs)
+              end
+            end
+          end
+
+          response
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Brute
+  module Middleware
+    module OTel
+      # Records tool results being sent back to the LLM as span events.
+      #
+      # Runs PRE-call: when env[:tool_results] is present, the orchestrator
+      # is sending tool execution results back to the LLM. Each result gets
+      # a span event with the tool name and success/error status.
+      #
+      class ToolResults < Base
+        def call(env)
+          span = env[:span]
+
+          if span && (results = env[:tool_results])
+            span.set_attribute("brute.tool_results.count", results.size)
+
+            results.each do |name, value|
+              error = value.is_a?(Hash) && value[:error]
+              attrs = { "tool.name" => name.to_s }
+              if error
+                attrs["tool.status"] = "error"
+                attrs["tool.error"] = value[:error].to_s
+              else
+                attrs["tool.status"] = "ok"
+              end
+              span.add_event("tool_result", attributes: attrs)
+            end
+          end
+
+          @app.call(env)
+        end
+      end
+    end
+  end
+end

data/lib/brute/middleware/otel.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Brute
+  module Middleware
+    # OpenTelemetry instrumentation for the LLM pipeline.
+    #
+    # Each middleware is independent and communicates through env[:span].
+    # OTel::Span must be outermost — it creates the span. The rest
+    # decorate it with events and attributes from their position in the
+    # pipeline.
+    #
+    # All middlewares are no-ops when opentelemetry-sdk is not loaded.
+    #
+    # Usage in pipeline:
+    #
+    #   use Brute::Middleware::OTel::Span
+    #   use Brute::Middleware::OTel::ToolResults
+    #   use Brute::Middleware::OTel::ToolCalls
+    #   use Brute::Middleware::OTel::TokenUsage
+    #
+    module OTel
+    end
+  end
+end
+
+require_relative "otel/span"
+require_relative "otel/tool_results"
+require_relative "otel/tool_calls"
+require_relative "otel/token_usage"

data/lib/brute/middleware/tool_use_guard.rb

@@ -11,8 +11,16 @@ module Brute
     # blocks is lost. This causes "unexpected tool_use_id" on the next call
     # because tool_result references a tool_use that's missing from the buffer.
     #
-    # This middleware runs post-call and injects a synthetic assistant message
-    # when tool calls exist but no assistant message was recorded.
+    # This middleware runs post-call and ensures every pending tool_use ID
+    # is covered by an assistant message in the buffer. It handles three
+    # cases:
+    #
+    #   1. ctx.functions is non-empty and the assistant message exists → no-op
+    #   2. ctx.functions is non-empty but the assistant message is missing
+    #      (or has different IDs) → inject synthetic message
+    #   3. ctx.functions is empty (nil-choice bug) but the stream recorded
+    #      tool calls → inject synthetic message using stream metadata
+    #
     class ToolUseGuard
       def initialize(app)
         @app = app
@@ -22,32 +30,67 @@ module Brute
         response = @app.call(env)
 
         ctx = env[:context]
-        functions = ctx.functions
 
-        # If there are pending tool calls, ensure the buffer has an assistant
-        # message with tool_use blocks.
+        # Collect pending tool data from ctx.functions (primary) or the
+        # stream's recorded metadata (fallback for nil-choice bug).
+        tool_data = collect_tool_data(ctx, env)
+        return response if tool_data.empty?
+
+        # Find all tool_use IDs already covered by assistant messages.
+        covered_ids = covered_tool_ids(ctx)
+
+        # Inject a synthetic assistant message for any uncovered tool calls.
+        uncovered = tool_data.reject { |td| covered_ids.include?(td[:id]) }
+        inject_synthetic!(ctx, uncovered) unless uncovered.empty?
+
+        response
+      end
+
+      private
+
+      def collect_tool_data(ctx, env)
+        functions = ctx.functions
         if functions && !functions.empty?
-          messages = ctx.messages.to_a
-          last_assistant = messages.reverse.find { |m| m.role.to_s == "assistant" }
-
-          unless last_assistant&.tool_call?
-            # Build a synthetic assistant message with the tool_use data
-            tool_calls = functions.map do |fn|
-              LLM::Object.from(id: fn.id, name: fn.name, arguments: fn.arguments)
-            end
-            original_tool_calls = functions.map do |fn|
-              { "type" => "tool_use", "id" => fn.id, "name" => fn.name, "input" => fn.arguments || {} }
-            end
-
-            synthetic = LLM::Message.new(:assistant, "", {
-              tool_calls: tool_calls,
-              original_tool_calls: original_tool_calls,
-            })
-            ctx.messages.concat([synthetic])
+          functions.map { |fn| { id: fn.id, name: fn.name, arguments: fn.arguments } }
+        elsif env[:streaming]
+          stream = resolve_stream(ctx)
+          if stream
+            data = stream.pending_tool_calls.dup
+            stream.clear_pending_tool_calls!
+            data
+          else
+            []
           end
+        else
+          []
         end
+      end
 
-        response
+      def resolve_stream(ctx)
+        stream = ctx.instance_variable_get(:@params)&.dig(:stream)
+        stream if stream.respond_to?(:pending_tool_calls)
+      end
+
+      def covered_tool_ids(ctx)
+        ctx.messages.to_a
+          .select { |m| m.role.to_s == "assistant" && m.tool_call? }
+          .flat_map { |m| (m.extra.original_tool_calls || []).map { |tc| tc["id"] } }
+          .to_set
+      end
+
+      def inject_synthetic!(ctx, uncovered)
+        tool_calls = uncovered.map do |td|
+          LLM::Object.from(id: td[:id], name: td[:name], arguments: td[:arguments])
+        end
+        original_tool_calls = uncovered.map do |td|
+          { "type" => "tool_use", "id" => td[:id], "name" => td[:name], "input" => td[:arguments] || {} }
+        end
+
+        synthetic = LLM::Message.new(:assistant, "", {
+          tool_calls: tool_calls,
+          original_tool_calls: original_tool_calls,
+        })
+        ctx.messages.concat([synthetic])
       end
     end
   end

data/lib/brute/orchestrator.rb

@@ -20,7 +20,7 @@ module Brute
   class Orchestrator
     MAX_REQUESTS_PER_TURN = 100
 
-    attr_reader :context, :session, :pipeline, :env, :barrier
+    attr_reader :context, :session, :pipeline, :env, :barrier, :message_store
 
     def initialize(
       provider:,
@@ -40,6 +40,7 @@ module Brute
       @cwd = cwd
       @session = session || Session.new
       @logger = logger || Logger.new($stderr, level: Logger::INFO)
+      @message_store = @session.message_store
 
       # Build system prompt
       custom_rules = load_custom_rules
@@ -65,6 +66,7 @@ module Brute
         session: @session,
         logger: @logger,
         reasoning: reasoning,
+        message_store: @message_store,
       )
 
       # The shared env hash — passed to every pipeline.call()
@@ -115,7 +117,7 @@ module Brute
 
       # --- Agent loop ---
       loop do
-        break if @context.functions.empty?
+        break if @context.functions.empty? && (!@stream || @stream.queue.empty?)
 
         # Collect tool results.
         # Streaming: tools already spawned threads during the LLM response — just join them.
@@ -135,7 +137,7 @@ module Brute
         @request_count += 1
 
         # Check limits
-        break if @context.functions.empty?
+        break if @context.functions.empty? && (!@stream || @stream.queue.empty?)
         break if @request_count >= MAX_REQUESTS_PER_TURN
         break if @env[:metadata][:tool_error_limit_reached]
       end
@@ -149,28 +151,42 @@ module Brute
     # Pipeline construction
     # ------------------------------------------------------------------
 
-    def build_pipeline(compactor:, session:, logger:, reasoning:)
+    def build_pipeline(compactor:, session:, logger:, reasoning:, message_store:)
       sys_prompt = @system_prompt
       tools = @tool_classes
+      stream = @stream
 
       Pipeline.new do
-        # Outermost: timing and logging (sees total elapsed including retries)
+        # OTel span lifecycle (outermost — creates env[:span])
+        use Middleware::OTel::Span
+
+        # Timing and logging
         use Middleware::Tracing, logger: logger
 
+        # OTel: record tool results being sent back (pre-call)
+        use Middleware::OTel::ToolResults
+
         # Retry transient errors (wraps everything below)
         use Middleware::Retry
 
         # Save after each successful LLM call
         use Middleware::SessionPersistence, session: session
 
+        # Record structured messages in OpenCode {info, parts} format
+        use Middleware::MessageTracking, store: message_store
+
         # Track cumulative token usage
         use Middleware::TokenTracking
 
+        # OTel: record token usage from response (post-call)
+        use Middleware::OTel::TokenUsage
+
         # Check context size and compact if needed
         use Middleware::CompactionCheck,
           compactor: compactor,
           system_prompt: sys_prompt,
-          tools: tools
+          tools: tools,
+          stream: stream
 
         # Track per-tool errors
         use Middleware::ToolErrorTracking
@@ -184,6 +200,9 @@ module Brute
         # Guard against tool-only responses dropping the assistant message
         use Middleware::ToolUseGuard
 
+        # OTel: record tool calls the LLM requested (post-call, after ToolUseGuard)
+        use Middleware::OTel::ToolCalls
+
         # Innermost: the actual LLM call
         run Middleware::LLMCall.new
       end

data/lib/brute/session.rb

@@ -8,79 +8,154 @@ module Brute
   # Manages session persistence. Each session is a conversation that can be
   # saved to disk and resumed later.
   #
-  # Sessions are stored as JSON files in a configurable directory
-  # (defaults to ~/.brute/sessions/).
+  # New directory-based layout (per-session directory):
+  #
+  #   ~/.brute/sessions/{session-id}/
+  #     session.meta.json          # session metadata
+  #     context.json               # llm.rb context blob (for resumption)
+  #     msg_0001.json              # structured messages (OpenCode format)
+  #     msg_0002.json
+  #     ...
+  #
+  # Also supports the legacy flat layout for reading:
+  #
+  #   ~/.brute/sessions/{session-id}.json
+  #   ~/.brute/sessions/{session-id}.meta.json
+  #
   class Session
     attr_reader :id, :title, :path
 
     def initialize(id: nil, dir: nil)
       @id = id || SecureRandom.uuid
-      @dir = dir || File.join(Dir.home, ".brute", "sessions")
-      @path = File.join(@dir, "#{@id}.json")
+      @base_dir = dir || File.join(Dir.home, ".brute", "sessions")
+      @session_dir = File.join(@base_dir, @id)
+      @path = File.join(@session_dir, "context.json")
       @title = nil
       @metadata = {}
-      FileUtils.mkdir_p(@dir)
+      FileUtils.mkdir_p(@session_dir)
+
+      # Check for legacy flat-file layout and migrate path if present
+      @legacy_path = File.join(@base_dir, "#{@id}.json")
+      @legacy_meta = File.join(@base_dir, "#{@id}.meta.json")
+    end
+
+    # Returns a MessageStore for this session's structured messages.
+    def message_store
+      @message_store ||= MessageStore.new(session_id: @id, dir: @session_dir)
     end
 
-    # Save a context to this session file.
+    # Save a context to this session.
     def save(context, title: nil, metadata: {})
       @title = title if title
       @metadata.merge!(metadata)
 
-      data = {
-        id: @id,
-        title: @title,
-        saved_at: Time.now.iso8601,
-        metadata: @metadata,
-      }
-
-      # Use llm.rb's built-in serialization
+      # Use llm.rb's built-in serialization for context (used for resumption)
       context.save(path: @path)
 
       # Write metadata sidecar
-      meta_path = @path.sub(/\.json$/, ".meta.json")
-      File.write(meta_path, JSON.pretty_generate(data))
+      save_meta
     end
 
-    # Restore a context from this session file.
+    # Restore a context from this session.
     # Returns true if restored successfully, false if no session file found.
     def restore(context)
-      return false unless File.exist?(@path)
+      # Try new layout first, then legacy
+      ctx_path = if File.exist?(@path)
+        @path
+      elsif File.exist?(@legacy_path)
+        @legacy_path
+      end
 
-      context.restore(path: @path)
+      return false unless ctx_path
 
-      # Load metadata sidecar if present
-      meta_path = @path.sub(/\.json$/, ".meta.json")
-      if File.exist?(meta_path)
-        data = JSON.parse(File.read(meta_path), symbolize_names: true)
-        @title = data[:title]
-        @metadata = data[:metadata] || {}
-      end
+      context.restore(path: ctx_path)
+
+      # Load metadata
+      load_meta
 
       true
     end
 
     # List all saved sessions, newest first.
+    # Scans both new directory-based layout and legacy flat files.
     def self.list(dir: nil)
       dir ||= File.join(Dir.home, ".brute", "sessions")
       return [] unless File.directory?(dir)
 
-      Dir.glob(File.join(dir, "*.meta.json")).map { |meta_path|
+      sessions = {}
+
+      # New layout: {id}/session.meta.json
+      Dir.glob(File.join(dir, "*", "session.meta.json")).each do |meta_path|
+        data = JSON.parse(File.read(meta_path), symbolize_names: true)
+        id = data[:id]
+        next unless id
+        sessions[id] = {
+          id: id,
+          title: data[:title],
+          saved_at: data[:saved_at],
+          path: File.join(File.dirname(meta_path), "context.json"),
+        }
+      end
+
+      # Legacy layout: {id}.meta.json (only if not already found)
+      Dir.glob(File.join(dir, "*.meta.json")).each do |meta_path|
+        # Skip files inside session subdirectories
+        next if meta_path.include?("/session.meta.json")
         data = JSON.parse(File.read(meta_path), symbolize_names: true)
-        {
-          id: data[:id],
+        id = data[:id]
+        next unless id
+        next if sessions.key?(id)  # new layout takes precedence
+        sessions[id] = {
+          id: id,
           title: data[:title],
           saved_at: data[:saved_at],
           path: meta_path.sub(/\.meta\.json$/, ".json"),
         }
-      }.sort_by { |s| s[:saved_at] || "" }.reverse
+      end
+
+      sessions.values.sort_by { |s| s[:saved_at] || "" }.reverse
     end
 
-    # Delete a session from disk.
+    # Delete a session from disk (both new and legacy layouts).
     def delete
-      File.delete(@path) if File.exist?(@path)
-      meta_path = @path.sub(/\.json$/, ".meta.json")
-      File.delete(meta_path) if File.exist?(meta_path)
+      # New layout: remove the whole directory
+      FileUtils.rm_rf(@session_dir) if File.directory?(@session_dir)
+
+      # Legacy layout: remove flat files
+      File.delete(@legacy_path) if File.exist?(@legacy_path)
+      File.delete(@legacy_meta) if File.exist?(@legacy_meta)
+    end
+
+    private
+
+    def meta_path
+      File.join(@session_dir, "session.meta.json")
+    end
+
+    def save_meta
+      data = {
+        id: @id,
+        title: @title,
+        saved_at: Time.now.iso8601,
+        metadata: @metadata,
+      }
+      FileUtils.mkdir_p(@session_dir)
+      File.write(meta_path, JSON.pretty_generate(data))
+    end
+
+    def load_meta
+      # Try new layout first
+      path = if File.exist?(meta_path)
+        meta_path
+      elsif File.exist?(@legacy_meta)
+        @legacy_meta
+      end
+
+      return unless path
+
+      data = JSON.parse(File.read(path), symbolize_names: true)
+      @title = data[:title]
+      @metadata = data[:metadata] || {}
     end
   end
 end

data/lib/brute/version.rb

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

data/lib/brute.rb

@@ -28,6 +28,7 @@ require_relative 'brute/doom_loop'
 require_relative 'brute/hooks'
 require_relative 'brute/compactor'
 require_relative 'brute/system_prompt'
+require_relative 'brute/message_store'
 require_relative 'brute/session'
 require_relative 'brute/pipeline'
 require_relative 'brute/agent_stream'
@@ -44,10 +45,12 @@ require_relative 'brute/middleware/doom_loop_detection'
 require_relative 'brute/middleware/token_tracking'
 require_relative 'brute/middleware/compaction_check'
 require_relative 'brute/middleware/session_persistence'
+require_relative 'brute/middleware/message_tracking'
 require_relative 'brute/middleware/tracing'
 require_relative 'brute/middleware/tool_error_tracking'
 require_relative 'brute/middleware/reasoning_normalizer'
 require_relative "brute/middleware/tool_use_guard"
+require_relative "brute/middleware/otel"
 
 # Tools
 require_relative 'brute/tools/fs_read'

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: brute
 version: !ruby/object:Gem::Version
-  version: 0.1.7
+  version: 0.1.8
 platform: ruby
 authors:
 - Brute Contributors
@@ -92,10 +92,17 @@ files:
 - lib/brute/doom_loop.rb
 - lib/brute/file_mutation_queue.rb
 - lib/brute/hooks.rb
+- lib/brute/message_store.rb
 - lib/brute/middleware/base.rb
 - lib/brute/middleware/compaction_check.rb
 - lib/brute/middleware/doom_loop_detection.rb
 - lib/brute/middleware/llm_call.rb
+- lib/brute/middleware/message_tracking.rb
+- lib/brute/middleware/otel.rb
+- lib/brute/middleware/otel/span.rb
+- lib/brute/middleware/otel/token_usage.rb
+- lib/brute/middleware/otel/tool_calls.rb
+- lib/brute/middleware/otel/tool_results.rb
 - lib/brute/middleware/reasoning_normalizer.rb
 - lib/brute/middleware/retry.rb
 - lib/brute/middleware/session_persistence.rb