console_agent 0.10.0 → 0.12.0

data/lib/console_agent/safety_guards.rb

@@ -0,0 +1,207 @@
+module ConsoleAgent
+  # Raised by safety guards to block dangerous operations.
+  # Host apps should raise this error in their custom guards.
+  # ConsoleAgent will catch it and guide the user to use 'd' or /danger.
+  class SafetyError < StandardError
+    attr_reader :guard, :blocked_key
+
+    def initialize(message, guard: nil, blocked_key: nil)
+      super(message)
+      @guard = guard
+      @blocked_key = blocked_key
+    end
+  end
+
+  class SafetyGuards
+    attr_reader :guards
+
+    def initialize
+      @guards = {}
+      @enabled = true
+      @allowlist = {}  # { guard_name => [String or Regexp, ...] }
+    end
+
+    def add(name, &block)
+      @guards[name.to_sym] = block
+    end
+
+    def remove(name)
+      @guards.delete(name.to_sym)
+    end
+
+    def enabled?
+      @enabled
+    end
+
+    def enable!
+      @enabled = true
+    end
+
+    def disable!
+      @enabled = false
+    end
+
+    def empty?
+      @guards.empty?
+    end
+
+    def names
+      @guards.keys
+    end
+
+    def allow(guard_name, key)
+      guard_name = guard_name.to_sym
+      @allowlist[guard_name] ||= []
+      @allowlist[guard_name] << key unless @allowlist[guard_name].include?(key)
+    end
+
+    def allowed?(guard_name, key)
+      entries = @allowlist[guard_name.to_sym]
+      return false unless entries
+
+      entries.any? do |entry|
+        case entry
+        when Regexp then key.match?(entry)
+        else entry.to_s == key.to_s
+        end
+      end
+    end
+
+    def allowlist
+      @allowlist
+    end
+
+    # Compose all guards around a block of code.
+    # Each guard is an around-block: guard.call { inner }
+    # Result: guard_1 { guard_2 { guard_3 { yield } } }
+    def wrap(&block)
+      return yield unless @enabled && !@guards.empty?
+
+      @guards.values.reduce(block) { |inner, guard|
+        -> { guard.call(&inner) }
+      }.call
+    end
+  end
+
+  # Built-in guard: database write prevention
+  # Works on all Rails versions (5+) and all database adapters.
+  # Prepends a write-intercepting module once, controlled by a thread-local flag.
+  module BuiltinGuards
+    # Blocks INSERT, UPDATE, DELETE, DROP, CREATE, ALTER, TRUNCATE
+    module WriteBlocker
+      WRITE_PATTERN = /\A\s*(INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE)\b/i
+      TABLE_PATTERN = /\b(?:INTO|FROM|UPDATE|TABLE|TRUNCATE)\s+[`"]?(\w+)[`"]?/i
+
+      private
+
+      def console_agent_check_write!(sql)
+        return unless Thread.current[:console_agent_block_writes] && sql.match?(WRITE_PATTERN)
+
+        table = sql.match(TABLE_PATTERN)&.captures&.first
+        guards = ConsoleAgent.configuration.safety_guards
+        return if table && guards.allowed?(:database_writes, table)
+
+        raise ConsoleAgent::SafetyError.new(
+          "Database write blocked: #{sql.strip.split(/\s+/).first(3).join(' ')}...",
+          guard: :database_writes,
+          blocked_key: table
+        )
+      end
+
+      public
+
+      def execute(sql, *args, **kwargs)
+        console_agent_check_write!(sql)
+        super
+      end
+
+      def exec_delete(sql, *args, **kwargs)
+        console_agent_check_write!(sql)
+        super
+      end
+
+      def exec_update(sql, *args, **kwargs)
+        console_agent_check_write!(sql)
+        super
+      end
+    end
+
+    def self.database_writes
+      ->(& block) {
+        ensure_write_blocker_installed!
+        Thread.current[:console_agent_block_writes] = true
+        begin
+          block.call
+        ensure
+          Thread.current[:console_agent_block_writes] = false
+        end
+      }
+    end
+
+    def self.ensure_write_blocker_installed!
+      return if @write_blocker_installed
+
+      connection = ActiveRecord::Base.connection
+      unless connection.class.ancestors.include?(WriteBlocker)
+        connection.class.prepend(WriteBlocker)
+      end
+      @write_blocker_installed = true
+    end
+
+    # Blocks non-safe HTTP requests (POST, PUT, PATCH, DELETE, etc.) via Net::HTTP.
+    # Since most Ruby HTTP libraries (HTTParty, RestClient, Faraday) use Net::HTTP
+    # under the hood, this covers them all.
+    module HttpBlocker
+      SAFE_METHODS = %w[GET HEAD OPTIONS TRACE].freeze
+
+      def request(req, *args, &block)
+        if Thread.current[:console_agent_block_http] && !SAFE_METHODS.include?(req.method)
+          host = @address.to_s
+          guards = ConsoleAgent.configuration.safety_guards
+          unless guards.allowed?(:http_mutations, host)
+            raise ConsoleAgent::SafetyError.new(
+              "HTTP #{req.method} blocked (#{host}#{req.path})",
+              guard: :http_mutations,
+              blocked_key: host
+            )
+          end
+        end
+        super
+      end
+    end
+
+    def self.http_mutations
+      ->(&block) {
+        ensure_http_blocker_installed!
+        Thread.current[:console_agent_block_http] = true
+        begin
+          block.call
+        ensure
+          Thread.current[:console_agent_block_http] = false
+        end
+      }
+    end
+
+    def self.mailers
+      ->(&block) {
+        old_value = ActionMailer::Base.perform_deliveries
+        ActionMailer::Base.perform_deliveries = false
+        begin
+          block.call
+        ensure
+          ActionMailer::Base.perform_deliveries = old_value
+        end
+      }
+    end
+
+    def self.ensure_http_blocker_installed!
+      return if @http_blocker_installed
+
+      require 'net/http'
+      unless Net::HTTP.ancestors.include?(HttpBlocker)
+        Net::HTTP.prepend(HttpBlocker)
+      end
+      @http_blocker_installed = true
+    end
+  end
+end

data/lib/console_agent/session_logger.rb

@@ -5,12 +5,12 @@ module ConsoleAgent
         return unless ConsoleAgent.configuration.session_logging
         return unless table_exists?
 
-        record = session_class.create!(
+        create_attrs = {
           query:         attrs[:query],
           conversation:  Array(attrs[:conversation]).to_json,
           input_tokens:  attrs[:input_tokens] || 0,
           output_tokens: attrs[:output_tokens] || 0,
-          user_name:     current_user_name,
+          user_name:     attrs[:user_name] || current_user_name,
           mode:          attrs[:mode].to_s,
           name:          attrs[:name],
           code_executed: attrs[:code_executed],
@@ -22,7 +22,9 @@ module ConsoleAgent
           model:         ConsoleAgent.configuration.resolved_model,
           duration_ms:   attrs[:duration_ms],
           created_at:    Time.respond_to?(:current) ? Time.current : Time.now
-        )
+        }
+        create_attrs[:slack_thread_ts] = attrs[:slack_thread_ts] if attrs[:slack_thread_ts]
+        record = session_class.create!(create_attrs)
         record.id
       rescue => e
         msg = "ConsoleAgent: session logging failed: #{e.class}: #{e.message}"
@@ -31,6 +33,15 @@ module ConsoleAgent
         nil
       end
 
+      def find_by_slack_thread(thread_ts)
+        return nil unless ConsoleAgent.configuration.session_logging
+        return nil unless table_exists?
+        session_class.where(slack_thread_ts: thread_ts).order(created_at: :desc).first
+      rescue => e
+        ConsoleAgent.logger.warn("ConsoleAgent: session lookup failed: #{e.class}: #{e.message}")
+        nil
+      end
+
       def update(id, attrs)
         return unless id
         return unless ConsoleAgent.configuration.session_logging