llm_gateway 0.5.0 → 0.7.0

data/lib/llm_gateway/agents/in_memory_session_manager.rb

@@ -0,0 +1,222 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require "time"
+
+module LlmGateway
+  module Agents
+    class InMemorySessionManager
+      MESSAGE_QUEUED = :queued
+      MESSAGE_STARTED = :started
+      QUEUES = [ :steer, :follow_up, :next_turn ].freeze
+      DRAIN_MODES = [ :one_at_a_time, :all ].freeze
+
+      attr_reader :session_id, :session_start
+
+      def initialize(session_id = nil)
+        @state = :idle
+        @session_id = session_id
+        @message_queues = Hash.new { |hash, key| hash[key] = [] }
+      end
+
+      def busy!
+        @state = :busy
+      end
+
+      def idle!
+        @state = :idle
+      end
+
+      def drain_message_queue(queue = :next_turn, mode: :all)
+        messages = queued_messages(queue, mode)
+        messages.each { |message| push_message(message) }
+        messages
+      end
+
+      def queued_messages?(queue = :next_turn)
+        @message_queues[validate_queue!(queue)].any?
+      end
+
+      def push_message_to_queue(message, queue = :next_turn)
+        @message_queues[validate_queue!(queue)] << message
+      end
+
+      def busy?
+        @state == :busy
+      end
+
+      def validate_queue!(queue)
+        queue = queue.to_sym
+        raise ArgumentError, "Invalid queue mode: #{queue}" unless QUEUES.include?(queue)
+
+        queue
+      end
+
+      def validate_drain_mode!(mode)
+        mode = mode.to_sym
+        raise ArgumentError, "Invalid queue drain mode: #{mode}" unless DRAIN_MODES.include?(mode)
+
+        mode
+      end
+
+      def start_or_enqueue_user_message(payload, queue: :next_turn)
+        if busy?
+          push_message_to_queue(payload, queue)
+          MESSAGE_QUEUED
+        else
+          yield if block_given?
+          push_message(payload)
+          busy!
+          MESSAGE_STARTED
+        end
+      end
+
+      def push_message(payload)
+        payload = payload.deep_symbolize_keys
+
+        push_entry(
+          type: "message",
+          usage: message_usage(payload),
+          data: payload,
+        )
+      end
+
+      def push_entry(entry)
+        id = SecureRandom.uuid
+        new_entry = {
+          id: id,
+          parent_id: parent_id_for_new_entry,
+          timestamp: Time.now.iso8601,
+          **entry
+        }
+
+        persist_entry(new_entry)
+        new_entry
+      end
+
+      def active_messages
+        active_message_events.map { |event| event[:data] }
+      end
+
+      def last_message_id
+        message_events.last&.dig(:id)
+      end
+
+      def last_model_used
+        events.reverse.find { |event| event[:type] == "model_change" }&.dig(:model_id)
+      end
+
+      def last_reasoning_level_used
+        events.reverse.find { |event| event[:type] == "reasoning_change" }&.dig(:reasoning)
+      end
+
+      def events_until(event_id)
+        index = events.index { |event| event[:id] == event_id }
+        raise ArgumentError, "Event not found in session: #{event_id}" unless index
+
+        events[0..index]
+      end
+
+      def events
+        @events ||= [ new_session_event ]
+      end
+
+      def build_model_input_messages
+        return active_messages unless last_compaction_entry
+
+        [ last_compaction_entry[:data], *active_messages ]
+      end
+
+      def total_tokens
+        entry = active_message_events.reverse.find { |event| event.dig(:usage, :total_tokens) }
+        entry&.dig(:usage, :total_tokens) || 0
+      end
+
+      def last_assistant_message_at
+        entry = active_message_events.reverse.find { |event| event.dig(:data, :role) == "assistant" }
+        Time.parse(entry[:timestamp]) if entry
+      end
+
+      def compaction(adapter)
+        response = adapter.stream(
+          active_messages,
+          system: "Summarize the conversation so far for future context.",
+          tools: []
+        )
+        message = response.to_h
+
+        push_entry(
+          type: "compaction",
+          usage: message_usage(message),
+          data: message
+        )
+      end
+
+      private
+
+      def queued_messages(queue, mode)
+        queue = validate_queue!(queue)
+        case validate_drain_mode!(mode)
+        when :one_at_a_time
+          message = @message_queues[queue].shift
+          message ? [ message ] : []
+        when :all
+          @message_queues[queue].shift(@message_queues[queue].length)
+        end
+      end
+
+      def parent_id_for_new_entry
+        events.length.positive? ? events.last[:id] : nil
+      end
+
+      def message_events
+        events.select { |event| event[:type] == "message" }
+      end
+
+      def active_message_events
+        compaction_event = last_compaction_entry
+        return message_events unless compaction_event
+
+        compaction_index = events.index(compaction_event)
+        events[(compaction_index + 1)..].select { |event| event[:type] == "message" }
+      end
+
+      def last_compaction_entry
+        events.reverse.find { |event| event[:type] == "compaction" }
+      end
+
+      def message_usage(message)
+        usage = message[:usage] || message["usage"]
+        return {} unless usage
+
+        usage.transform_keys(&:to_sym)
+      end
+
+      def persist_entry(entry)
+        attributes = {
+          session_id: @session_id,
+          position: next_position,
+          id: entry[:id],
+          parent_id: entry[:parent_id],
+          timestamp: entry[:timestamp],
+          type: entry[:type],
+          usage: entry[:usage],
+          data: entry[:data]
+        }
+
+        events << entry
+        attributes
+      end
+
+      def next_position
+        events.length
+      end
+
+      def new_session_event
+        @session_id ||= SecureRandom.uuid
+        @session_start = Time.now.strftime("%Y%m%d_%H%M%S")
+        { type: "session", id: session_id, timestamp: session_start }
+      end
+    end
+  end
+end

data/lib/llm_gateway/agents/tools/bash_tool.rb

@@ -0,0 +1,132 @@
+require "securerandom"
+require "tmpdir"
+require_relative "tool_utils"
+
+class BashTool < LlmGateway::Tool
+  # Pi adaptation notes:
+  # - Keep timeout schema as integer: gruv treats integer and number schemas equivalently for seconds.
+  # - Do not add pi's pluggable operations, shell/env hooks, command prefix, AbortSignal handling, partial updates, or UI render details: those are runtime/UI extension concerns outside this tool contract.
+  name "bash"
+  description "Execute a bash command in the current working directory. Returns stdout and stderr. Output is truncated to last #{ToolUtils::DEFAULT_MAX_LINES} lines or #{ToolUtils::DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first). If truncated, full output is saved to a temp file. Optionally provide a timeout in seconds."
+  input_schema({
+    type: "object",
+    properties: {
+      command: { type: "string", description: "Bash command to execute" },
+      timeout: { type: "integer", description: "Timeout in seconds (optional, no default timeout)" }
+    },
+    required: [ "command" ]
+  })
+
+  def execute(input)
+    command = input[:command]
+    timeout = input[:timeout]
+
+    result = run_command(command, timeout)
+    out = format_output(result[:output], empty_text: result[:timed_out] ? "" : "(no output)")
+
+    if result[:timed_out]
+      return append_status(out, "Command timed out after #{timeout} seconds")
+    end
+
+    if result[:exit_status] && result[:exit_status] != 0
+      return append_status(out, "Command exited with code #{result[:exit_status]}")
+    end
+
+    out
+  rescue StandardError => e
+    e.message
+  end
+
+  private
+
+  def run_command(command, timeout)
+    output = +""
+    timed_out = false
+    read_io, write_io = IO.pipe
+    pid = Process.spawn(command, chdir: Dir.pwd, in: File::NULL, out: write_io, err: write_io, pgroup: true)
+    write_io.close
+
+    deadline = timeout && timeout.positive? ? Time.now + timeout : nil
+
+    loop do
+      remaining = deadline ? deadline - Time.now : nil
+      if remaining && remaining <= 0
+        timed_out = true
+        terminate_process_group(pid)
+        break
+      end
+
+      ready = IO.select([ read_io ], nil, nil, remaining)
+      unless ready
+        timed_out = true
+        terminate_process_group(pid)
+        break
+      end
+
+      begin
+        output << read_io.readpartial(16 * 1024)
+      rescue EOFError
+        break
+      end
+    end
+
+    _, status = Process.wait2(pid)
+    drain_available_output(read_io, output)
+    read_io.close
+
+    { output: output, exit_status: status.exitstatus, timed_out: timed_out }
+  ensure
+    write_io.close if write_io && !write_io.closed?
+    read_io.close if read_io && !read_io.closed?
+  end
+
+  def drain_available_output(read_io, output)
+    loop do
+      ready = IO.select([ read_io ], nil, nil, 0.1)
+      break unless ready
+
+      begin
+        output << read_io.readpartial(16 * 1024)
+      rescue EOFError
+        break
+      end
+    end
+  end
+
+  def terminate_process_group(pid)
+    Process.kill("TERM", -pid)
+    sleep 0.1
+    Process.kill("KILL", -pid)
+  rescue Errno::ESRCH, Errno::EPERM
+    nil
+  end
+
+  def format_output(output, empty_text: "(no output)")
+    truncation = ToolUtils.truncate_tail(output)
+    out = truncation[:content]
+    out = empty_text if out.empty?
+
+    return out unless truncation[:truncated]
+
+    temp_path = File.join(Dir.tmpdir, "pi-bash-#{SecureRandom.hex(8)}.log")
+    File.write(temp_path, output)
+
+    start_line = truncation[:total_lines] - truncation[:output_lines] + 1
+    end_line = truncation[:total_lines]
+
+    notice = if truncation[:last_line_partial]
+      last_line = output.split("\n", -1).last
+      "[Showing last #{ToolUtils.format_size(truncation[:output_bytes])} of line #{end_line} (line is #{ToolUtils.format_size(last_line.bytesize)}). Full output: #{temp_path}]"
+    elsif truncation[:truncated_by] == "lines"
+      "[Showing lines #{start_line}-#{end_line} of #{truncation[:total_lines]}. Full output: #{temp_path}]"
+    else
+      "[Showing lines #{start_line}-#{end_line} of #{truncation[:total_lines]} (#{ToolUtils.format_size(ToolUtils::DEFAULT_MAX_BYTES)} limit). Full output: #{temp_path}]"
+    end
+
+    "#{out}\n\n#{notice}"
+  end
+
+  def append_status(text, status)
+    text.empty? ? status : "#{text}\n\n#{status}"
+  end
+end

data/lib/llm_gateway/agents/tools/edit_tool.rb

@@ -0,0 +1,215 @@
+require "json"
+require_relative "tool_utils"
+
+class EditTool < LlmGateway::Tool
+  # Pi adaptation notes:
+  # - Legacy single-edit input: pi accepts oldText/newText and converts to edits,
+  #   but this tool intentionally exposes/supports only the edits array to keep the LLM contract unambiguous.
+  # - Do not add pi's diff/patch details, preview rendering, pluggable operations, or AbortSignal handling: those are UI/runtime extension concerns outside this tool contract.
+  name "edit"
+
+  class EditError < StandardError; end
+  description "Edit a single file using exact text replacement. Every edits[].oldText must match a unique, non-overlapping region of the original file. If two changes affect the same block or nearby lines, merge them into one edit instead of emitting overlapping edits."
+  input_schema({
+    type: "object",
+    properties: {
+      path: { type: "string", description: "Path to the file to edit (relative or absolute)" },
+      edits: {
+        type: "array",
+        description: "One or more targeted replacements. Each edit is matched against the original file, not incrementally. Do not include overlapping or nested edits.",
+        items: {
+          type: "object",
+          properties: {
+            oldText: { type: "string", description: "Exact text for one targeted replacement. It must be unique in the original file and must not overlap with any other edits[].oldText in the same call." },
+            newText: { type: "string", description: "Replacement text for this targeted edit." }
+          },
+          required: [ "oldText", "newText" ]
+        }
+      }
+    },
+    required: [ "path", "edits" ]
+  })
+
+  def execute(input)
+    path = input[:path]
+    edits = prepare_edits(input[:edits])
+
+    return "Edit tool input is invalid. edits must contain at least one replacement." if !edits.is_a?(Array) || edits.empty?
+
+    absolute_path = ToolUtils.resolve_to_cwd(path)
+
+    ToolUtils.with_file_mutation_lock(absolute_path) do
+      begin
+        File.open(absolute_path, File::RDWR) { }
+      rescue SystemCallError => e
+        return "Could not edit file: #{path}. Error code: #{e.class.name.split("::").last}."
+      end
+
+      raw_content = File.binread(absolute_path)
+
+      bom = raw_content.start_with?("\xEF\xBB\xBF".b) ? "\xEF\xBB\xBF".b : "".b
+      content_without_bom = bom.empty? ? raw_content : raw_content.byteslice(3..)
+      content_utf8 = content_without_bom.force_encoding("UTF-8")
+
+      original_ending = detect_line_ending(content_utf8)
+      normalized_content = normalize_to_lf(content_utf8)
+      base_content, new_content = apply_edits_to_normalized_content(normalized_content, edits, path)
+
+      restored = restore_line_endings(new_content, original_ending)
+      final_bytes = bom + restored.encode("UTF-8").b
+
+      File.binwrite(absolute_path, final_bytes)
+      "Successfully replaced #{edits.length} block(s) in #{path}."
+    end
+  rescue EditError => e
+    e.message
+  rescue StandardError => e
+    "Error editing file: #{e.message}"
+  end
+
+  private
+
+  def prepare_edits(edits)
+    return JSON.parse(edits, symbolize_names: true) if edits.is_a?(String)
+
+    edits
+  end
+
+  def detect_line_ending(content)
+    crlf_index = content.index("\r\n")
+    lf_index = content.index("\n")
+    return "\n" unless lf_index
+    return "\n" unless crlf_index
+
+    crlf_index < lf_index ? "\r\n" : "\n"
+  end
+
+  def normalize_to_lf(text)
+    text.gsub("\r\n", "\n").gsub("\r", "\n")
+  end
+
+  def restore_line_endings(text, ending)
+    ending == "\r\n" ? text.gsub("\n", "\r\n") : text
+  end
+
+  def normalize_for_fuzzy_match(text)
+    text
+      .unicode_normalize(:nfkc)
+      .split("\n", -1)
+      .map(&:rstrip)
+      .join("\n")
+      .gsub(/[\u2018\u2019\u201A\u201B]/, "'")
+      .gsub(/[\u201C\u201D\u201E\u201F]/, '"')
+      .gsub(/[\u2010\u2011\u2012\u2013\u2014\u2015\u2212]/, "-")
+      .gsub(/[\u00A0\u2002-\u200A\u202F\u205F\u3000]/, " ")
+  end
+
+  def fuzzy_find_text(content, old_text)
+    exact_index = content.index(old_text)
+    if exact_index
+      return {
+        found: true,
+        index: exact_index,
+        match_length: old_text.length,
+        used_fuzzy_match: false,
+        content_for_replacement: content
+      }
+    end
+
+    fuzzy_content = normalize_for_fuzzy_match(content)
+    fuzzy_old_text = normalize_for_fuzzy_match(old_text)
+    fuzzy_index = fuzzy_content.index(fuzzy_old_text)
+
+    return { found: false, index: -1, match_length: 0, used_fuzzy_match: false, content_for_replacement: content } unless fuzzy_index
+
+    {
+      found: true,
+      index: fuzzy_index,
+      match_length: fuzzy_old_text.length,
+      used_fuzzy_match: true,
+      content_for_replacement: fuzzy_content
+    }
+  end
+
+  def count_occurrences(content, old_text)
+    fuzzy_content = normalize_for_fuzzy_match(content)
+    fuzzy_old_text = normalize_for_fuzzy_match(old_text)
+    fuzzy_content.split(fuzzy_old_text, -1).length - 1
+  end
+
+  def apply_edits_to_normalized_content(normalized_content, edits, path)
+    normalized_edits = edits.map do |edit|
+      {
+        oldText: normalize_to_lf(edit[:oldText]),
+        newText: normalize_to_lf(edit[:newText])
+      }
+    end
+
+    normalized_edits.each_with_index do |edit, index|
+      raise EditError, empty_old_text_error(path, index, normalized_edits.length) if edit[:oldText].empty?
+    end
+
+    initial_matches = normalized_edits.map { |edit| fuzzy_find_text(normalized_content, edit[:oldText]) }
+    base_content = initial_matches.any? { |match| match[:used_fuzzy_match] } ? normalize_for_fuzzy_match(normalized_content) : normalized_content
+
+    matched_edits = []
+    normalized_edits.each_with_index do |edit, index|
+      match = fuzzy_find_text(base_content, edit[:oldText])
+      raise EditError, not_found_error(path, index, normalized_edits.length) unless match[:found]
+
+      occurrences = count_occurrences(base_content, edit[:oldText])
+      raise EditError, duplicate_error(path, index, normalized_edits.length, occurrences) if occurrences > 1
+
+      matched_edits << {
+        edit_index: index,
+        match_index: match[:index],
+        match_length: match[:match_length],
+        new_text: edit[:newText]
+      }
+    end
+
+    matched_edits.sort_by! { |edit| edit[:match_index] }
+    matched_edits.each_cons(2) do |previous, current|
+      next unless previous[:match_index] + previous[:match_length] > current[:match_index]
+
+      raise EditError, "edits[#{previous[:edit_index]}] and edits[#{current[:edit_index]}] overlap in #{path}. Merge them into one edit or target disjoint regions."
+    end
+
+    new_content = base_content.dup
+    matched_edits.reverse_each do |edit|
+      new_content = new_content[0...edit[:match_index]] + edit[:new_text] + new_content[(edit[:match_index] + edit[:match_length])..]
+    end
+
+    raise EditError, no_change_error(path, normalized_edits.length) if base_content == new_content
+
+    [ base_content, new_content ]
+  end
+
+  def not_found_error(path, edit_index, total_edits)
+    return "Could not find the exact text in #{path}. The old text must match exactly including all whitespace and newlines." if total_edits == 1
+
+    "Could not find edits[#{edit_index}] in #{path}. The oldText must match exactly including all whitespace and newlines."
+  end
+
+  def duplicate_error(path, edit_index, total_edits, occurrences)
+    if total_edits == 1
+      return "Found #{occurrences} occurrences of the text in #{path}. The text must be unique. Please provide more context to make it unique."
+    end
+
+    "Found #{occurrences} occurrences of edits[#{edit_index}] in #{path}. Each oldText must be unique. Please provide more context to make it unique."
+  end
+
+  def empty_old_text_error(path, edit_index, total_edits)
+    return "oldText must not be empty in #{path}." if total_edits == 1
+
+    "edits[#{edit_index}].oldText must not be empty in #{path}."
+  end
+
+  def no_change_error(path, total_edits)
+    if total_edits == 1
+      return "No changes made to #{path}. The replacement produced identical content. This might indicate an issue with special characters or the text not existing as expected."
+    end
+
+    "No changes made to #{path}. The replacements produced identical content."
+  end
+end