ruby_llm 1.15.0 → 2.0.0.rc1

data/lib/ruby_llm/active_record/chat_methods.rb

@@ -1,484 +1,1002 @@
 # frozen_string_literal: true
 
 require 'active_support/concern'
+require 'ruby_llm/active_record/attachment_helpers'
 
 module RubyLLM
   module ActiveRecord
-    # Methods mixed into chat models.
+    # ChatMethods provides the RubyLLM::Chat API on ActiveRecord models
+    # declared with <tt>acts_as_chat</tt>, persisting every message to the
+    # database. Configuration methods return +self+ so calls can be chained.
+    #
+    #   class Chat < ApplicationRecord
+    #     acts_as_chat
+    #   end
+    #
+    #   chat = Chat.create!(model: 'gpt-5.6-luna')
+    #   chat.ask "What is the capital of France?"
+    #   chat.messages.count # => 2
+    #
     module ChatMethods
       extend ActiveSupport::Concern
+      include Enumerable
+      include AttachmentHelpers
+
+      CANCELLATION_POLL_INTERVAL = 1.0 # :nodoc:
+      COMPLETION_ERRORS = [ # :nodoc:
+        RubyLLM::CancelledError, RubyLLM::Error, Faraday::Error, Timeout::Error, Errno::ETIMEDOUT
+      ].freeze
 
       included do
-        before_save :resolve_model_from_strings
+        before_save :resolve_model
       end
 
-      attr_accessor :assume_model_exists, :context
+      # When +true+, skips the model registry lookup so unregistered model ids
+      # are accepted. Not persisted; set it again after reloading the record.
+      attr_accessor :assume_model_exists
 
-      def model=(value)
-        @model_string = value if value.is_a?(String)
-        return if value.is_a?(String)
+      # Overrides the wire protocol the provider would pick for the model, such
+      # as +:responses+ or +:chat_completions+ for OpenAI, or +nil+ for the
+      # provider default. Not persisted; set it again after reloading the record.
+      attr_accessor :protocol
+
+      # An optional RubyLLM::Context supplying per-chat configuration, used
+      # when building the underlying chat. Not persisted; set it again after
+      # reloading the record.
+      attr_accessor :context
+
+      # Requests cancellation of the current in-flight chat operation. The
+      # request is persisted so a background job can observe it from another
+      # process.
+      def cancel
+        if persisted?
+          update_column(:cancelled, true)
+        else
+          self[:cancelled] = true
+        end
 
-        if self.class.model_association_name == :model
+        @chat&.cancel
+        self
+      end
+
+      # Returns whether this record or its memoized in-memory chat has a
+      # pending cancellation request.
+      def cancelled?
+        @chat&.cancelled? || self[:cancelled]
+      end
+
+      # Records approval for +tool_call+ (a ToolCall, a tool call record, or
+      # an id) on the persisted tool call, so the next #complete executes it
+      # from any process. Returns +self+.
+      #
+      #   chat.approve(params[:tool_call_id])
+      #   CompleteJob.perform_later(chat.id)
+      #
+      def approve(tool_call)
+        record_tool_call_decision(tool_call, 'approved')
+      end
+
+      # Records denial for +tool_call+ (a ToolCall, a tool call record, or
+      # an id) on the persisted tool call. The next #complete appends a
+      # structured denial result instead of executing the tool. Returns
+      # +self+.
+      def deny(tool_call)
+        record_tool_call_decision(tool_call, 'denied')
+      end
+
+      # Returns whether the conversation is waiting on tool calls that
+      # require approval and have no recorded decision. See
+      # RubyLLM::Chat#awaiting_approval?.
+      def awaiting_approval?
+        to_llm.awaiting_approval?
+      end
+
+      # Returns the persisted tool call records that require approval and
+      # have no recorded decision, ready to render as approval cards. Pass
+      # a record (or its tool_call_id) to #approve or #deny.
+      # A record's +remote?+ identifies a provider-executed call.
+      #
+      #   chat.pending_approvals.each { |record| render record }
+      #
+      def pending_approvals
+        ids = to_llm.pending_approvals.map(&:id)
+        RubyLLM::ActiveRecord::ToolCall.where(
+          tool_call_id: ids,
+          message_type: self.class.message_class.constantize.polymorphic_name,
+          message_id: messages_association.select(:id)
+        )
+      end
+
+      def messages_association # :nodoc:
+        send(messages_association_name)
+      end
+
+      # Sets the chat's model from an id, a RubyLLM::Model value, or the
+      # associated internal model record.
+      #
+      #   chat.model = 'gpt-5.6-luna'
+      #
+      def model=(value)
+        if value.is_a?(RubyLLM::ActiveRecord::Model)
+          @pending_model_id = nil
+          @pending_provider = nil
           super
         else
-          self.model_association = value
+          @pending_model_id = value.respond_to?(:id) ? value.id : value
+          @pending_provider = value.provider if value.respond_to?(:provider)
         end
       end
 
+      # Stores +value+ as the model id, resolved to a model record before save.
       def model_id=(value)
-        @model_string = value
+        @pending_model_id = value
       end
 
+      # Returns the model id of the associated model record, or +nil+.
+      #
+      #   chat.model_id # => "gpt-5.6-luna"
+      #
       def model_id
-        model_association&.model_id
+        model&.model_id || @pending_model_id
       end
 
+      # Stores +value+ as the provider used when resolving the model id
+      # before save.
       def provider=(value)
-        @provider_string = value
+        @pending_provider = value
       end
 
+      # Returns the provider of the associated model record, or +nil+.
       def provider
-        model_association&.provider
-      end
-
-      private
-
-      def resolve_model_from_strings # rubocop:disable Metrics/PerceivedComplexity
-        config = context&.config || RubyLLM.config
-        @model_string ||= config.default_model unless model_association
-        return unless @model_string
-
-        model_info, _provider = Models.resolve(
-          @model_string,
-          provider: @provider_string,
-          assume_exists: assume_model_exists || false,
-          config: config
-        )
-
-        model_class = self.class.model_class.constantize
-        model_record = model_class.find_or_create_by!(
-          model_id: model_info.id,
-          provider: model_info.provider
-        ) do |m|
-          m.name = model_info.name || model_info.id
-          m.family = model_info.family
-          m.context_window = model_info.context_window
-          m.max_output_tokens = model_info.max_output_tokens
-          m.capabilities = model_info.capabilities || []
-          m.modalities = model_info.modalities || {}
-          m.pricing = model_info.pricing || {}
-          m.metadata = model_info.metadata || {}
-        end
-
-        self.model_association = model_record
-        @model_string = nil
-        @provider_string = nil
+        model&.provider || @pending_provider
       end
 
-      public
-
+      # Returns the underlying RubyLLM::Chat for this record, building it on
+      # first call and memoizing it. The chat is loaded with the persisted
+      # messages and wired to persist new ones. Subsequent calls return the
+      # same chat without touching the database; use #reload to refresh its
+      # message history from the record.
       def to_llm
-        model_record = model_association
-        @chat ||= (context || RubyLLM).chat(
-          model: model_record.model_id,
-          provider: model_record.provider.to_sym,
-          assume_model_exists: assume_model_exists || false
-        )
-        @chat.reset_messages!
-
-        ordered_messages = order_messages_for_llm(messages_association.to_a)
-        ordered_messages.each do |msg|
-          @chat.add_message(msg.to_llm)
-        end
-        reapply_runtime_instructions(@chat)
-
-        setup_persistence_callbacks
+        @chat ||= build_llm_chat # rubocop:disable Naming/MemoizedInstanceVariableName
       end
 
-      def with_instructions(instructions, append: false, replace: nil)
-        append = append_instructions?(append:, replace:)
-        persist_system_instruction(instructions, append:)
+      # Reloads the record from the database, Rails-style, and refreshes the
+      # underlying chat's persisted message history to match. Runtime-only
+      # configuration such as tools, temperature, and callbacks is preserved.
+      # Returns +self+.
+      def reload(...)
+        super
+        sync_messages if @chat
+        self
+      end
 
-        to_llm.with_instructions(instructions, append:, replace:)
+      # Rebinds the underlying chat to +value+ so subsequent requests use its
+      # configuration. Pass +nil+ to return to the global RubyLLM
+      # configuration. The Context itself is runtime-only and is not
+      # persisted.
+      def with_context(value)
+        self.context = value
+        @chat&.with_context(value)
         self
       end
 
-      def with_runtime_instructions(instructions, append: false, replace: nil)
-        append = append_instructions?(append:, replace:)
-        store_runtime_instruction(instructions, append:)
+      # Sets the system instructions, persisting them as a message with the
+      # +:system+ role. Replaces any persisted system messages unless
+      # +append:+ is true. Pass <tt>persist: false</tt> to apply the
+      # instructions only to the in-memory chat for this record instance.
+      # With <tt>cache_until_here: true</tt> the instruction becomes an
+      # explicit prompt cache boundary. Returns +self+.
+      #
+      #   chat.with_instructions "You are a Ruby expert."
+      #   chat.with_instructions "Use short bullet points.", append: true
+      #   chat.with_instructions current_context, persist: false
+      #
+      def with_instructions(instructions, append: false, persist: true, cache_until_here: false)
+        to_llm
+
+        if persist
+          if instructions.nil?
+            clear_persisted_system_instructions
+          else
+            persist_system_instruction(instructions, append:, cache_until_here:)
+          end
+        else
+          store_unpersisted_instruction(instructions, append:, cache_until_here:)
+        end
 
-        to_llm.with_instructions(instructions, append:, replace:)
+        sync_messages
         self
       end
 
-      def with_tool(...)
-        to_llm.with_tool(...)
-        self
+      # Chat configuration and callback methods forwarded to the underlying
+      # RubyLLM::Chat. Each behaves exactly as documented on RubyLLM::Chat,
+      # then returns the record so calls chain.
+      CHAINABLE_CHAT_DELEGATES = %i[
+        with_tools with_tool_options with_server_tools with_fallbacks with_temperature
+        with_max_output_tokens with_thinking with_citations with_caching
+        with_end_user with_compaction
+        with_provider_options with_headers with_schema
+        before_request before_message after_message before_tool_call after_tool_result
+        before_fallback after_fallback
+      ].freeze
+
+      ##
+      # :method: with_tools
+      # :call-seq: with_tools(*tools)
+      #
+      # Applies Chat#with_tools and returns this record.
+
+      ##
+      # :method: with_tool_options
+      # :call-seq: with_tool_options(**options)
+      #
+      # Applies Chat#with_tool_options and returns this record.
+
+      ##
+      # :method: with_server_tools
+      # :call-seq: with_server_tools(*tools, **tools_with_options)
+      #
+      # Applies Chat#with_server_tools and returns this record.
+
+      ##
+      # :method: with_fallbacks
+      # :call-seq: with_fallbacks(*models, on: Fallback::DEFAULT_ERRORS)
+      #
+      # Applies Chat#with_fallbacks and returns this record.
+
+      ##
+      # :method: with_temperature
+      # :call-seq: with_temperature(temperature)
+      #
+      # Applies Chat#with_temperature and returns this record.
+
+      ##
+      # :method: with_max_output_tokens
+      # :call-seq: with_max_output_tokens(max_output_tokens)
+      #
+      # Applies Chat#with_max_output_tokens and returns this record.
+
+      ##
+      # :method: with_thinking
+      # :call-seq: with_thinking(enabled = true, **options)
+      #
+      # Applies Chat#with_thinking and returns this record.
+
+      ##
+      # :method: with_citations
+      # :call-seq: with_citations(enabled = true)
+      #
+      # Applies Chat#with_citations and returns this record.
+
+      ##
+      # :method: with_caching
+      # :call-seq: with_caching(options = {})
+      #
+      # Applies Chat#with_caching and returns this record.
+
+      ##
+      # :method: with_end_user
+      # :call-seq: with_end_user(end_user)
+      #
+      # Applies Chat#with_end_user and returns this record.
+
+      ##
+      # :method: with_compaction
+      # :call-seq: with_compaction(options = {})
+      #
+      # Applies Chat#with_compaction and returns this record.
+
+      ##
+      # :method: with_provider_options
+      # :call-seq: with_provider_options(provider_options)
+      #
+      # Applies Chat#with_provider_options and returns this record.
+
+      ##
+      # :method: with_headers
+      # :call-seq: with_headers(headers)
+      #
+      # Applies Chat#with_headers and returns this record.
+
+      ##
+      # :method: with_schema
+      # :call-seq: with_schema(schema)
+      #
+      # Applies Chat#with_schema and returns this record.
+
+      ##
+      # :method: before_request
+      # :call-seq: before_request(&block)
+      #
+      # Applies Chat#before_request and returns this record.
+
+      ##
+      # :method: before_message
+      # :call-seq: before_message(&block)
+      #
+      # Applies Chat#before_message and returns this record.
+
+      ##
+      # :method: after_message
+      # :call-seq: after_message(&block)
+      #
+      # Applies Chat#after_message and returns this record.
+
+      ##
+      # :method: before_tool_call
+      # :call-seq: before_tool_call(&block)
+      #
+      # Applies Chat#before_tool_call and returns this record.
+
+      ##
+      # :method: after_tool_result
+      # :call-seq: after_tool_result(&block)
+      #
+      # Applies Chat#after_tool_result and returns this record.
+
+      ##
+      # :method: before_fallback
+      # :call-seq: before_fallback(&block)
+      #
+      # Applies Chat#before_fallback and returns this record.
+
+      ##
+      # :method: after_fallback
+      # :call-seq: after_fallback(&block)
+      #
+      # Applies Chat#after_fallback and returns this record.
+
+      ##
+      # :method: caching
+      # :call-seq: caching
+      #
+      # Delegates to Chat#caching. See that method for arguments and return values.
+
+      ##
+      # :method: citations
+      # :call-seq: citations
+      #
+      # Delegates to Chat#citations. See that method for arguments and return values.
+
+      ##
+      # :method: compaction
+      # :call-seq: compaction
+      #
+      # Delegates to Chat#compaction. See that method for arguments and return values.
+
+      ##
+      # :method: concurrency
+      # :call-seq: concurrency
+      #
+      # Delegates to Chat#concurrency. See that method for arguments and return values.
+
+      ##
+      # :method: end_user
+      # :call-seq: end_user
+      #
+      # Delegates to Chat#end_user. See that method for arguments and return values.
+
+      ##
+      # :method: fallbacks
+      # :call-seq: fallbacks
+      #
+      # Delegates to Chat#fallbacks. See that method for arguments and return values.
+
+      ##
+      # :method: headers
+      # :call-seq: headers
+      #
+      # Delegates to Chat#headers. See that method for arguments and return values.
+
+      ##
+      # :method: max_output_tokens
+      # :call-seq: max_output_tokens
+      #
+      # Delegates to Chat#max_output_tokens. See that method for arguments and return values.
+
+      ##
+      # :method: provider_options
+      # :call-seq: provider_options
+      #
+      # Delegates to Chat#provider_options. See that method for arguments and return values.
+
+      ##
+      # :method: schema
+      # :call-seq: schema
+      #
+      # Delegates to Chat#schema. See that method for arguments and return values.
+
+      ##
+      # :method: server_tools
+      # :call-seq: server_tools
+      #
+      # Delegates to Chat#server_tools. See that method for arguments and return values.
+
+      ##
+      # :method: temperature
+      # :call-seq: temperature
+      #
+      # Delegates to Chat#temperature. See that method for arguments and return values.
+
+      ##
+      # :method: thinking
+      # :call-seq: thinking()
+      #
+      # Delegates to Chat#thinking. See that method for arguments and return values.
+
+      ##
+      # :method: tool_options
+      # :call-seq: tool_options()
+      #
+      # Delegates to Chat#tool_options. See that method for arguments and return values.
+
+      ##
+      # :method: tools
+      # :call-seq: tools
+      #
+      # Delegates to Chat#tools. See that method for arguments and return values.
+
+      CHAINABLE_CHAT_DELEGATES.each do |name|
+        define_method(name) do |*args, **kwargs, &block|
+          to_llm.public_send(name, *args, **kwargs, &block)
+          self
+        end
       end
 
-      def with_tools(...)
-        to_llm.with_tools(...)
-        self
+      # Chat values and operations whose return values pass through unchanged.
+      #
+      # The public methods behave as documented on RubyLLM::Chat.
+
+      PASSTHROUGH_CHAT_DELEGATES = %i[
+        caching citations compaction concurrency end_user fallbacks headers max_output_tokens provider_options
+        schema server_tools temperature thinking tool_options tools
+        add_completion count_tokens each render
+      ].freeze
+
+      ##
+      # :method: count_tokens
+      # :call-seq: count_tokens(message = nil)
+      #
+      # Returns the number of input tokens the next request would carry,
+      # counted by the provider over the persisted conversation.
+
+      ##
+      # :method: each
+      # :call-seq: each(&block)
+      #
+      # Yields each message in the conversation. Returns an Enumerator without a block.
+
+      ##
+      # :method: render
+      # :call-seq: render
+      #
+      # Returns the next request payload with #before_request hooks applied.
+
+      PASSTHROUGH_CHAT_DELEGATES.each do |name|
+        define_method(name) do |*args, **kwargs, &block|
+          to_llm.public_send(name, *args, **kwargs, &block)
+        end
       end
 
-      def with_model(model_name, provider: nil, assume_exists: false)
+      private_constant :CHAINABLE_CHAT_DELEGATES, :PASSTHROUGH_CHAT_DELEGATES
+
+      # Switches the chat to +model_name+, resolving and saving the model
+      # record and updating the underlying chat. Falls back to the configured
+      # default model when +model_name+ is +nil+. Pass +protocol:+ to override
+      # the wire protocol the provider would pick for the model. Returns +self+.
+      #
+      #   chat.with_model 'claude-sonnet-5'
+      #
+      def with_model(model_name, provider: nil, protocol: nil, assume_model_exists: false)
+        model_name ||= (context&.config || RubyLLM.config).default_model
         self.model = model_name
         self.provider = provider if provider
-        self.assume_model_exists = assume_exists
-        resolve_model_from_strings
+        self.protocol = protocol
+        self.assume_model_exists = assume_model_exists
+        resolve_model
         save!
-        to_llm.with_model(model_association.model_id, provider: model_association.provider.to_sym, assume_exists:)
+        to_llm.with_model(model_id, provider: provider&.to_sym, protocol:, assume_model_exists:)
         self
       end
 
-      def with_temperature(...)
-        to_llm.with_temperature(...)
-        self
-      end
+      # Persists +message_or_attributes+ as a message record, including any
+      # attachments and tool calls. Accepts a RubyLLM::Message, an attributes
+      # Hash, or a record responding to +to_llm+.
+      # Returns the message record.
+      #
+      #   chat.add_message(role: :user, content: long_context)
+      #
+      def add_message(message_or_attributes)
+        llm_message = message_or_attributes
+        llm_message = llm_message.to_llm if llm_message.respond_to?(:to_llm)
+        llm_message = RubyLLM::Message.new(llm_message) unless llm_message.is_a?(RubyLLM::Message)
 
-      def with_thinking(...)
-        to_llm.with_thinking(...)
-        self
+        message_record = messages_association.create!(message_attributes(llm_message))
+
+        if llm_message.tool_call_id && (tool_call = find_tool_call(llm_message.tool_call_id))
+          tool_call.update!(result: message_record)
+        end
+
+        persist_content(message_record, llm_message.attachments) if llm_message.attachments.any?
+        persist_tool_calls(llm_message.tool_calls, message_record:) if llm_message.tool_calls.present?
+
+        @chat&.add_message(llm_message)
+
+        message_record
       end
 
-      def with_params(...)
-        to_llm.with_params(...)
+      # Marks the latest persisted message as a prompt cache boundary, or the
+      # latest in-memory message when none is persisted yet. Returns +self+.
+      #
+      #   chat.with_instructions('Reusable analysis prompt').cache_until_here
+      #
+      # Raises ArgumentError if the chat has no messages.
+      def cache_until_here
+        message_record = messages_association.order(:id).last
+        if message_record
+          message_record.cache_until_here
+        elsif @chat&.messages&.any?
+          @chat.cache_until_here
+        else
+          raise ArgumentError, 'No messages to cache'
+        end
+
         self
       end
 
-      def with_headers(...)
-        to_llm.with_headers(...)
-        self
+      # Returns token usage aggregated across every persisted usage entry,
+      # including retries and attempts that did not produce a message.
+      #
+      #   chat.tokens.input
+      #
+      def tokens
+        RubyLLM::Tokens.aggregate(ruby_llm_usages.map(&:tokens))
       end
 
-      def with_schema(...)
-        to_llm.with_schema(...)
-        self
+      # Returns a RubyLLM::Cost aggregating every persisted usage entry,
+      # including retries and attempts that did not produce a message.
+      #
+      #   chat.cost.total
+      #
+      def cost
+        records = ruby_llm_usages.to_a
+        RubyLLM::Cost.aggregate(records.map(&:cost), complete: records.all?(&:cost_available?))
       end
 
-      def on_new_message(&)
-        to_llm.on_new_message(&)
-        self
+      # Persists +message+ as a user message, then runs the conversation loop
+      # and returns the latest assistant RubyLLM::Message. The loop pauses
+      # when #awaiting_approval? is true. Yields streaming chunks to a block.
+      #
+      #   chat.ask "What is the capital of France?"
+      #   chat.ask "What's in this file?", with: "diagram.png"
+      #
+      def ask(message = nil, with: nil, &)
+        ask_later(message, with: with)
+        complete(&)
       end
 
-      def on_end_message(&)
-        to_llm.on_end_message(&)
+      alias say ask
+
+      # Persists +message+ as a user message without calling the model, so
+      # #complete can run later. Returns +self+.
+      #
+      #   chat.ask_later "Summarize this document."
+      #   chat.complete
+      #
+      def ask_later(message = nil, with: nil)
+        to_llm.raise_if_pending_tool_calls!
+        add_message(role: :user, content: message, attachments: with)
         self
       end
 
-      def before_message(...)
-        to_llm.before_message(...)
-        self
+      # Makes a single generation attempt, persists the response, and returns it as a
+      # RubyLLM::Message. Tool calls in the response are not executed. See
+      # RubyLLM::Chat#generate.
+      def generate(...)
+        to_llm.generate(...)
+      rescue *COMPLETION_ERRORS => e
+        cleanup_after_failure(e)
+        raise
       end
 
-      def after_message(...)
-        to_llm.after_message(...)
-        self
+      # Compacts the model context and persists its assistant Message without
+      # deleting earlier messages. See RubyLLM::Chat#compact.
+      def compact
+        to_llm.compact
+      rescue *COMPLETION_ERRORS => e
+        cleanup_after_failure(e)
+        raise
       end
 
-      def before_tool_call(...)
-        to_llm.before_tool_call(...)
+      # Executes the pending tool calls and persists their results without
+      # calling the model. See RubyLLM::Chat#run_tools. Returns +self+.
+      def run_tools
+        to_llm.run_tools
         self
       end
 
-      def after_tool_result(...)
-        to_llm.after_tool_result(...)
-        self
+      # Advances the conversation by one move: runs the pending tool calls if
+      # there are any, otherwise generates a response. Returns +nil+ once the
+      # chat is complete or waiting for approval. See RubyLLM::Chat#step.
+      #
+      #   chat.step until chat.complete? || chat.awaiting_approval?
+      #
+      def step(...)
+        to_llm.step(...)
+      rescue *COMPLETION_ERRORS => e
+        cleanup_after_failure(e)
+        raise
       end
 
-      def on_tool_call(...)
-        to_llm.on_tool_call(...)
-        self
+      # Returns whether the conversation has no pending work, neither a
+      # response to generate nor tool calls to run. See
+      # RubyLLM::Chat#complete?.
+      def complete?
+        to_llm.complete?
       end
 
-      def on_tool_result(...)
-        to_llm.on_tool_result(...)
-        self
+      # Runs the completion loop on the underlying chat, persisting each
+      # message, and returns the latest RubyLLM::Message. Pauses when a tool
+      # requires approval. When the API call fails, destroys the empty assistant
+      # message and any orphaned tool results, then re-raises the error.
+      def complete(...)
+        to_llm.complete(...)
+      rescue *COMPLETION_ERRORS => e
+        cleanup_after_failure(e)
+        raise
       end
 
-      def add_message(message_or_attributes)
-        llm_message = message_or_attributes.is_a?(RubyLLM::Message) ? message_or_attributes : RubyLLM::Message.new(message_or_attributes)
-        content_text, attachments, content_raw = prepare_content_for_storage(llm_message.content)
-
-        attrs = { role: llm_message.role, content: content_text }
-        parent_tool_call_assoc = messages_association.klass.reflect_on_association(:parent_tool_call)
-        if parent_tool_call_assoc && llm_message.tool_call_id
-          tool_call_id = find_tool_call_id(llm_message.tool_call_id)
-          attrs[parent_tool_call_assoc.foreign_key] = tool_call_id if tool_call_id
+      private
+
+      def persist_usage_entry(entry)
+        tokens = entry.tokens
+        cost = entry.cost
+        attributes = {
+          operation: entry.operation,
+          provider: entry.provider,
+          model: entry.model,
+          status: entry.status,
+          input_tokens: tokens.input,
+          output_tokens: tokens.output,
+          cache_read_tokens: tokens.cache_read,
+          cache_write_tokens: tokens.cache_write,
+          thinking_tokens: tokens.thinking,
+          input_cost: cost.input,
+          output_cost: cost.output,
+          cache_read_cost: cost.cache_read,
+          cache_write_cost: cost.cache_write,
+          thinking_cost: cost.thinking,
+          total_cost: cost.total
+        }
+        record = ruby_llm_usages.create!(attributes)
+        usage_records_by_entry[entry] = record
+      end
+
+      def link_usage_entries(message)
+        message.ruby_llm_usage_entries.each do |entry|
+          record = usage_records_by_entry[entry]
+          record&.update!(message: @message)
         end
+      end
 
-        message_record = messages_association.create!(attrs)
-        message_record.update!(content_raw:) if message_record.respond_to?(:content_raw=)
+      def usage_records_by_entry
+        @usage_records_by_entry ||= {}.compare_by_identity
+      end
 
-        persist_content(message_record, attachments) if attachments.present?
-        persist_tool_calls(llm_message.tool_calls, message_record:) if llm_message.tool_calls.present?
+      def resolve_model
+        config = context&.config || RubyLLM.config
+        @pending_model_id ||= config.default_model unless model
+        return unless @pending_model_id
 
-        message_record
+        model_info = resolve_model_info
+
+        self.model = find_or_create_model(model_info)
+        @pending_model_id = nil
+        @pending_provider = nil
       end
 
-      def cost
-        RubyLLM::Cost.aggregate(messages_association.map(&:cost))
+      def resolve_model_info
+        return find_registered_model unless assume_model_exists
+
+        raise ArgumentError, 'Provider must be specified if assume_model_exists is true' unless @pending_provider
+
+        begin
+          find_registered_model
+        rescue RubyLLM::ModelNotFoundError
+          RubyLLM::Model.default(@pending_model_id, @pending_provider)
+        end
       end
 
-      def create_user_message(content, with: nil)
-        add_message(role: :user, content: build_content(content, with))
+      def find_registered_model
+        RubyLLM.models.find(@pending_model_id, provider: @pending_provider, config: context&.config)
       end
 
-      def ask(message = nil, with: nil, &)
-        add_message(role: :user, content: build_content(message, with))
-        complete(&)
+      # An empty store would otherwise hold only this chat's model on the next
+      # boot, since the registry reads whatever the store returns.
+      def load_model_registry_into_store
+        RubyLLM::ActiveRecord::Model.save_to_database(RubyLLM.models)
+      rescue ::ActiveRecord::RecordInvalid, ::ActiveRecord::RecordNotUnique
+        nil
       end
 
-      alias say ask
+      def find_or_create_model(model_info)
+        attributes = { model_id: model_info.id, provider: model_info.provider }
+        load_model_registry_into_store if RubyLLM::ActiveRecord::Model.none?
+
+        RubyLLM::ActiveRecord::Model.find_or_create_by!(attributes) do |record|
+          record.name = model_info.name || model_info.id
+          record.family = model_info.family
+          record.model_created_at = model_info.created_at
+          record.context_window = model_info.context_window
+          record.max_output_tokens = model_info.max_output_tokens
+          record.knowledge_cutoff = model_info.knowledge_cutoff
+          record.capabilities = model_info.capabilities || []
+          record.modalities = model_info.modalities.to_h
+          record.pricing = model_info.pricing.to_h
+          record.metadata = model_info.metadata || {}
+        end
+      rescue ::ActiveRecord::RecordInvalid, ::ActiveRecord::RecordNotUnique
+        # Another process can insert the row between the lookup and the insert.
+        RubyLLM::ActiveRecord::Model.find_by(attributes) || raise
+      end
 
-      def complete(...)
-        to_llm.complete(...)
-      rescue RubyLLM::Error => e
-        cleanup_failed_messages if @message&.persisted? && @message.content.blank?
+      def cleanup_after_failure(error)
+        reason = error.is_a?(RubyLLM::CancelledError) ? 'chat cancelled' : 'API call failed'
+        cleanup_failed_messages(reason:) if blank_placeholder?
         cleanup_orphaned_tool_results
-        raise e
       end
 
-      private
-
-      def cleanup_failed_messages
-        RubyLLM.logger.warn "RubyLLM: API call failed, destroying message: #{@message.id}"
+      def cleanup_failed_messages(reason:)
+        RubyLLM.logger.warn "RubyLLM: #{reason}, destroying message: #{@message.id}"
         @message.destroy
       end
 
       def cleanup_orphaned_tool_results # rubocop:disable Metrics/PerceivedComplexity
         messages_association.reload
-        last = messages_association.order(:id).last
+        last = eager_load_messages.last
 
         return unless last&.tool_call? || last&.tool_result?
 
         if last.tool_call?
           last.destroy
         elsif last.tool_result?
-          tool_call_message = last.parent_tool_call.message_association
-          expected_results = tool_call_message.tool_calls_association.pluck(:id)
-          fk_column = tool_call_message.class.reflections['tool_results'].foreign_key
-          actual_results = tool_call_message.tool_results.pluck(fk_column)
+          parent = last.ruby_llm_parent_tool_call
+          tool_call_message = parent.message
+          calls = tool_call_message.ruby_llm_tool_calls
 
-          if expected_results.sort != actual_results.sort
-            tool_call_message.tool_results.each(&:destroy)
+          if calls.any? { |call| call.result.nil? }
+            calls.filter_map(&:result).each(&:destroy)
             tool_call_message.destroy
           end
         end
       end
 
-      def setup_persistence_callbacks
-        return @chat if @chat.instance_variable_get(:@_persistence_callbacks_setup)
+      def eager_load_messages
+        assoc = messages_association
+        messages = assoc.to_a
+        return messages unless assoc.respond_to?(:klass)
+
+        associations = %i[ruby_llm_tool_calls ruby_llm_parent_tool_call ruby_llm_usages]
+
+        ::ActiveRecord::Associations::Preloader.new(records: messages, associations: associations).call
+        messages
+      end
+
+      def build_llm_chat
+        chat = (context || RubyLLM).chat(
+          model: model_id,
+          provider: provider&.to_sym,
+          protocol: protocol,
+          assume_model_exists: assume_model_exists || false
+        )
+        sync_messages(chat)
+        chat.cancellation_checker = proc { consume_persisted_cancellation_request }
+        chat.approval_checker = proc { |tool_call| persisted_tool_call_approval(tool_call) }
+        install_persistence_callbacks(chat)
+      end
 
-        @chat.before_message { persist_new_message }
-        @chat.after_message { |msg| persist_message_completion(msg) }
+      def record_tool_call_decision(tool_call, decision)
+        id = tool_call.respond_to?(:tool_call_id) ? tool_call.tool_call_id : tool_call
+        id = id.id if id.respond_to?(:id) && !id.is_a?(String)
+        record = find_tool_call(id)
+        raise ArgumentError, "Unknown tool call: #{id.inspect}" unless record
 
-        @chat.instance_variable_set(:@_persistence_callbacks_setup, true)
-        @chat
+        record.update!(approval: decision)
+        self
       end
 
-      def replace_persisted_system_instructions(instructions)
-        system_messages = messages_association.where(role: :system).order(:id).to_a
+      def persisted_tool_call_approval(tool_call)
+        record = RubyLLM::ActiveRecord::ToolCall.uncached { find_tool_call(tool_call.id) }
+        return unless record&.has_attribute?(:approval)
 
-        if system_messages.empty?
-          messages_association.create!(role: :system, content: instructions)
-          return
+        case record.approval
+        when 'approved' then true
+        when 'denied' then false
         end
+      end
 
-        primary_message = system_messages.shift
-        primary_message.update!(content: instructions) if primary_message.content != instructions
-        system_messages.each(&:destroy!)
+      def sync_messages(chat = @chat)
+        message_records = eager_load_messages
+        chat.messages = message_records
+        linked_entry_pairs = message_records.zip(chat.messages).flat_map do |record, message|
+          record.ruby_llm_usages.zip(message.ruby_llm_usage_entries)
+        end
+        linked_entries = linked_entry_pairs.to_h { |record, entry| [record.id, entry] }
+        chat.usage_entries = ruby_llm_usages.map do |record|
+          linked_entries.fetch(record.id) { record.to_entry }
+        end
+        reapply_runtime_instructions(chat)
+        chat
       end
 
-      def append_instructions?(append:, replace:)
-        return append if replace.nil?
+      def install_persistence_callbacks(chat)
+        chat.before_message { persist_new_message }
+        chat.usage_recorder = method(:persist_usage_entry)
+        chat.after_message { |msg| persist_message_completion(msg) }
+        chat
+      end
+
+      def clear_persisted_system_instructions
+        association = messages_association
+        association.where(role: :system).destroy_all
+        association.reset
+      end
 
-        append || (replace == false)
+      def replace_persisted_system_instructions(instructions, cache_until_here:)
+        existing = messages_association.where(role: :system).order(:id).to_a
+        if existing.one?
+          update_persisted_system_instruction(existing.first, instructions, cache_until_here:)
+        else
+          clear_persisted_system_instructions
+          messages_association.create!(role: :system, content: instructions, cache_until_here: cache_until_here)
+        end
       end
 
-      def persist_system_instruction(instructions, append:)
+      # Rewriting the same instructions every turn would move the system row
+      # behind the user messages and rebroadcast it each time.
+      def update_persisted_system_instruction(record, instructions, cache_until_here:)
+        attributes = { content: instructions }
+        attributes[:cache_until_here] = cache_until_here if record.has_attribute?(:cache_until_here)
+        changed = attributes.any? { |column, value| record[column] != value }
+        record.update!(attributes) if changed
+        messages_association.reset
+        record
+      end
+
+      def persist_system_instruction(instructions, append:, cache_until_here:)
         transaction do
           if append
-            messages_association.create!(role: :system, content: instructions)
+            messages_association.create!(
+              role: :system,
+              content: instructions,
+              cache_until_here: cache_until_here
+            )
           else
-            replace_persisted_system_instructions(instructions)
+            replace_persisted_system_instructions(instructions, cache_until_here:)
           end
         end
       end
 
-      def order_messages_for_llm(messages)
-        system_messages, non_system_messages = messages.partition { |msg| msg.role.to_s == 'system' }
-        system_messages + non_system_messages
+      def unpersisted_instructions
+        @unpersisted_instructions ||= []
       end
 
-      def runtime_instructions
-        @runtime_instructions ||= []
-      end
+      def store_unpersisted_instruction(instructions, append:, cache_until_here:)
+        if instructions.nil?
+          @unpersisted_instructions = []
+          return
+        end
 
-      def store_runtime_instruction(instructions, append:)
+        entry = [instructions, append, cache_until_here]
         if append
-          runtime_instructions << instructions
+          unpersisted_instructions << entry
         else
-          @runtime_instructions = [instructions]
+          @unpersisted_instructions = [entry]
         end
       end
 
       def reapply_runtime_instructions(chat)
-        return if runtime_instructions.empty?
-
-        first, *rest = runtime_instructions
-        chat.with_instructions(first)
-        rest.each { |instruction| chat.with_instructions(instruction, append: true) }
+        unpersisted_instructions.each do |instructions, append, cache_until_here|
+          chat.with_instructions(instructions, append:, cache_until_here:)
+        end
       end
 
       def persist_new_message
-        @message = messages_association.create!(role: :assistant, content: '')
+        @message.destroy if blank_placeholder?
+
+        attrs = { role: :assistant, content: '' }
+        @message = messages_association.create!(attrs)
+      end
+
+      def blank_placeholder?
+        return false unless @message&.persisted? && @message.content.blank?
+        return false if @message.respond_to?(:attachments) && @message.attachments.attached?
+
+        !@message.ruby_llm_tool_calls.exists?
       end
 
-      # rubocop:disable Metrics/PerceivedComplexity
       def persist_message_completion(message)
         return unless message
 
-        tool_call_id = find_tool_call_id(message.tool_call_id) if message.tool_call_id
+        tool_call = find_tool_call(message.tool_call_id) if message.tool_call_id
+        attrs = message_attributes(message)
 
         transaction do
-          content_text, attachments_to_persist, content_raw = prepare_content_for_storage(message.content)
-
-          attrs = {
-            role: message.role,
-            content: content_text,
-            input_tokens: message.input_tokens,
-            output_tokens: message.output_tokens
-          }
-          attrs[:cached_tokens] = message.cached_tokens if @message.has_attribute?(:cached_tokens)
-          if @message.has_attribute?(:cache_creation_tokens)
-            attrs[:cache_creation_tokens] = message.cache_creation_tokens
-          end
-          attrs[:thinking_text] = message.thinking&.text if @message.has_attribute?(:thinking_text)
-          attrs[:thinking_signature] = message.thinking&.signature if @message.has_attribute?(:thinking_signature)
-          attrs[:thinking_tokens] = message.thinking_tokens if @message.has_attribute?(:thinking_tokens)
-
-          # Add model association dynamically
-          attrs[self.class.model_association_name] = model_association
-
-          if tool_call_id
-            parent_tool_call_assoc = @message.class.reflect_on_association(:parent_tool_call)
-            attrs[parent_tool_call_assoc.foreign_key] = tool_call_id
-          end
-
           @message.assign_attributes(attrs)
-          @message.content_raw = content_raw if @message.respond_to?(:content_raw=)
           @message.save!
+          tool_call&.update!(result: @message)
 
-          persist_content(@message, attachments_to_persist) if attachments_to_persist
+          persist_content(@message, message.attachments) if message.attachments.any?
           persist_tool_calls(message.tool_calls) if message.tool_calls.present?
+          link_usage_entries(message)
         end
       end
-      # rubocop:enable Metrics/PerceivedComplexity
 
-      def persist_tool_calls(tool_calls, message_record: @message)
-        tool_call_klass = message_record.tool_calls_association.klass
-        supports_thought_signature = tool_call_klass.column_names.include?('thought_signature')
+      def message_attributes(message)
+        attrs = { role: message.role, content: message.content }
+        assign_supported_attribute(attrs, :thinking_text, message.thinking&.text)
+        assign_supported_attribute(attrs, :thinking_signature, message.thinking&.signature)
+        assign_supported_attribute(attrs, :citations, message.citations.map(&:to_h).presence)
+        assign_supported_attribute(attrs, :server_tool_calls, message.server_tool_calls.map(&:to_h).presence)
+        assign_supported_attribute(attrs, :raw_content, message.raw_content)
+        assign_supported_attribute(attrs, :raw_reasoning, message.raw_reasoning)
+        assign_supported_attribute(attrs, :finish_reason, message.finish_reason)
+        assign_supported_attribute(attrs, :cache_until_here, message.cache_until_here?)
+        attrs
+      end
 
+      def assign_supported_attribute(attributes, name, value)
+        attributes[name] = value if messages_association.klass.column_names.include?(name.to_s)
+      end
+
+      def persist_tool_calls(tool_calls, message_record: @message)
         tool_calls.each_value do |tool_call|
           attributes = tool_call.to_h
-          attributes.delete(:thought_signature) unless supports_thought_signature
           attributes[:tool_call_id] = attributes.delete(:id)
-          message_record.tool_calls_association.create!(**attributes)
+          message_record.ruby_llm_tool_calls.create!(**attributes)
         end
       end
 
-      def find_tool_call_id(tool_call_id)
-        messages = messages_association
-        message_class = messages.klass
-        tool_calls_assoc = message_class.tool_calls_association_name
-        tool_call_table_name = message_class.reflect_on_association(tool_calls_assoc).table_name
-
-        message_with_tool_call = messages.joins(tool_calls_assoc)
-                                         .find_by(tool_call_table_name => { tool_call_id: tool_call_id })
-        return nil unless message_with_tool_call
-
-        tool_call = message_with_tool_call.tool_calls_association.find_by(tool_call_id: tool_call_id)
-        tool_call&.id
-      end
-
-      def persist_content(message_record, attachments)
-        return unless message_record.respond_to?(:attachments)
-
-        attachables = prepare_for_active_storage(attachments)
-        message_record.attachments.attach(attachables) if attachables.any?
-      end
-
-      def prepare_for_active_storage(attachments)
-        Utils.to_safe_array(attachments).filter_map do |attachment|
-          case attachment
-          when ActionDispatch::Http::UploadedFile, ActiveStorage::Blob
-            attachment
-          when ActiveStorage::Attachment, ActiveStorage::Attached::One, ActiveStorage::Attached::Many
-            active_storage_blobs(attachment)
-          when Hash
-            attachment.values.map { |v| prepare_for_active_storage(v) }
-          else
-            convert_to_active_storage_format(attachment)
-          end
-        end.flatten.compact
+      def find_tool_call(tool_call_id)
+        RubyLLM::ActiveRecord::ToolCall.find_by(
+          tool_call_id: tool_call_id,
+          message_type: self.class.message_class.constantize.polymorphic_name,
+          message_id: messages_association.select(:id)
+        )
       end
 
-      def convert_to_active_storage_format(source)
-        return if source.blank?
-
-        attachment = source.is_a?(RubyLLM::Attachment) ? source : RubyLLM::Attachment.new(source)
-
-        if attachment.active_storage?
-          active_storage_blobs(attachment.source)
-        else
-          {
-            io: StringIO.new(attachment.content),
-            filename: attachment.filename,
-            content_type: attachment.mime_type
-          }
+      def consume_persisted_cancellation_request
+        if self[:cancelled]
+          clear_cancellation_request
+          return :cancelled
         end
-      rescue StandardError => e
-        RubyLLM.logger.warn "Failed to process attachment #{source}: #{e.message}"
-        nil
-      end
+        return unless persisted?
+        return unless cancellation_poll_due?
+        return unless persisted_cancellation_request?
 
-      def active_storage_blobs(attachment)
-        case attachment
-        when ActiveStorage::Blob then attachment
-        when ActiveStorage::Attachment, ActiveStorage::Attached::One then attachment.blob
-        when ActiveStorage::Attached::Many then attachment.blobs
-        end
+        clear_cancellation_request
+        :cancelled
       end
 
-      def build_content(message, attachments)
-        return message if content_like?(message)
-
-        RubyLLM::Content.new(message, attachments)
+      # Jobs and requests run with the query cache on, which would replay the
+      # first poll's answer for the rest of the run.
+      def persisted_cancellation_request?
+        self.class.uncached do
+          self.class.unscoped.where(self.class.primary_key => id).pick(:cancelled)
+        end
       end
 
-      def content_like?(object)
-        object.is_a?(RubyLLM::Content) || object.is_a?(RubyLLM::Content::Raw)
+      def clear_cancellation_request
+        self[:cancelled] = false
+        self.class.unscoped.where(self.class.primary_key => id).update_all(cancelled: false) if persisted?
       end
 
-      def prepare_content_for_storage(content)
-        attachments = nil
-        content_raw = nil
-        content_text = content
-
-        case content
-        when RubyLLM::Content::Raw
-          content_raw = content.value
-          content_text = nil
-        when RubyLLM::Content
-          attachments = content.attachments if content.attachments.any?
-          content_text = content.text
-        when Hash, Array
-          content_raw = content
-          content_text = nil
-        end
+      def cancellation_poll_due?
+        now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        return false if @last_cancellation_poll_at && now - @last_cancellation_poll_at < CANCELLATION_POLL_INTERVAL
 
-        [content_text, attachments, content_raw]
+        @last_cancellation_poll_at = now
+        true
       end
     end
   end