ruby_llm 2.0.0 → 2.1.0

data/lib/ruby_llm/active_record/chat_methods.rb

@@ -185,7 +185,8 @@ module RubyLLM
       # +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+.
+      # explicit prompt cache boundary, and <tt>cache_until_here: { ttl: "1h" }</tt>
+      # gives that boundary its own lifetime. Returns +self+.
       #
       #   chat.with_instructions "You are a Ruby expert."
       #   chat.with_instructions "Use short bullet points.", append: true
@@ -212,20 +213,40 @@ module RubyLLM
       # 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_provider_tools with_fallbacks with_temperature
+        with_tools with_mcp with_tool_options with_provider_tools with_fallbacks with_temperature answer decline
         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
+        convert_unsupported_attachments
+        before_request before_message after_message before_tool_call after_tool_result after_tool_progress
         before_fallback after_fallback
       ].freeze
 
       ##
       # :method: with_tools
-      # :call-seq: with_tools(*tools)
+      # :call-seq: with_tools(*tools, defer: nil)
       #
       # Applies Chat#with_tools and returns this record.
 
+      ##
+      # :method: with_mcp
+      # :call-seq: with_mcp(*servers, defer: nil)
+      #
+      # Applies Chat#with_mcp and returns this record.
+
+      ##
+      # :method: answer
+      # :call-seq: answer(request, **values)
+      #
+      # Applies Chat#answer and returns this record. The answer persists
+      # on the tool call, so any process can resume the call.
+
+      ##
+      # :method: decline
+      # :call-seq: decline(request)
+      #
+      # Applies Chat#decline and returns this record.
+
       ##
       # :method: with_tool_options
       # :call-seq: with_tool_options(**options)
@@ -304,6 +325,12 @@ module RubyLLM
       #
       # Applies Chat#with_schema and returns this record.
 
+      ##
+      # :method: convert_unsupported_attachments
+      # :call-seq: convert_unsupported_attachments(&block)
+      #
+      # Applies Chat#convert_unsupported_attachments and returns this record.
+
       ##
       # :method: before_request
       # :call-seq: before_request(&block)
@@ -334,6 +361,12 @@ module RubyLLM
       #
       # Applies Chat#after_tool_result and returns this record.
 
+      ##
+      # :method: after_tool_progress
+      # :call-seq: after_tool_progress(&block)
+      #
+      # Applies Chat#after_tool_progress and returns this record.
+
       ##
       # :method: before_fallback
       # :call-seq: before_fallback(&block)
@@ -436,6 +469,51 @@ module RubyLLM
       #
       # Delegates to Chat#tools. See that method for arguments and return values.
 
+      ##
+      # :method: deferred_tools
+      # :call-seq: deferred_tools
+      #
+      # Delegates to Chat#deferred_tools. See that method for arguments and return values.
+
+      ##
+      # :method: mcp
+      # :call-seq: mcp
+      #
+      # Delegates to Chat#mcp. See that method for arguments and return values.
+
+      ##
+      # :method: waiting?
+      # :call-seq: waiting?
+      #
+      # Delegates to Chat#waiting?. Approvals, answers, and tasks persist on
+      # their tool calls, so any process sees what the chat waits on.
+
+      ##
+      # :method: awaiting_input?
+      # :call-seq: awaiting_input?
+      #
+      # Delegates to Chat#awaiting_input?. See that method for arguments and return values.
+
+      ##
+      # :method: pending_inputs
+      # :call-seq: pending_inputs
+      #
+      # Delegates to Chat#pending_inputs. Requests persist on their tool
+      # calls, so they survive restarts.
+
+      ##
+      # :method: awaiting_tasks?
+      # :call-seq: awaiting_tasks?
+      #
+      # Delegates to Chat#awaiting_tasks?. See that method for arguments and return values.
+
+      ##
+      # :method: pending_tasks
+      # :call-seq: pending_tasks
+      #
+      # Delegates to Chat#pending_tasks. Tasks persist on their tool calls,
+      # so any process can check on them and resume the chat.
+
       CHAINABLE_CHAT_DELEGATES.each do |name|
         define_method(name) do |*args, **kwargs, &block|
           to_llm.public_send(name, *args, **kwargs, &block)
@@ -449,8 +527,8 @@ module RubyLLM
 
       PASSTHROUGH_CHAT_DELEGATES = %i[
         caching citations compaction concurrency end_user fallbacks headers max_output_tokens provider_options
-        schema provider_tools temperature thinking tool_options tools
-        add_completion count_tokens each render
+        schema provider_tools temperature thinking tool_options tools deferred_tools mcp waiting? awaiting_input?
+        pending_inputs awaiting_tasks? pending_tasks add_completion count_tokens each render
       ].freeze
 
       ##
@@ -526,17 +604,19 @@ module RubyLLM
       end
 
       # Marks the latest persisted message as a prompt cache boundary, or the
-      # latest in-memory message when none is persisted yet. Returns +self+.
+      # latest in-memory message when none is persisted yet. Pass +ttl:+ to
+      # give the boundary its own cache lifetime. Returns +self+.
       #
       #   chat.with_instructions('Reusable analysis prompt').cache_until_here
+      #   chat.with_instructions('Reusable analysis prompt').cache_until_here(ttl: "1h")
       #
       # Raises ArgumentError if the chat has no messages.
-      def cache_until_here
+      def cache_until_here(ttl: nil)
         message_record = messages_association.order(:id).last
         if message_record
-          message_record.cache_until_here
+          message_record.cache_until_here(ttl:)
         elsif @chat&.messages&.any?
-          @chat.cache_until_here
+          @chat.cache_until_here(ttl:)
         else
           raise ArgumentError, 'No messages to cache'
         end
@@ -585,7 +665,11 @@ module RubyLLM
       #
       def ask_later(message = nil, with: nil)
         to_llm.raise_if_pending_tool_calls!
-        add_message(role: :user, content: message, attachments: with)
+        if message.is_a?(RubyLLM::MCP::Prompt)
+          message.messages.each { |prompt_message| add_message(prompt_message) }
+        else
+          add_message(role: :user, content: message, attachments: with)
+        end
         self
       end
 
@@ -617,9 +701,9 @@ module RubyLLM
 
       # 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 is complete or #waiting?. See RubyLLM::Chat#step.
       #
-      #   chat.step until chat.complete? || chat.awaiting_approval?
+      #   chat.step until chat.complete? || chat.waiting?
       #
       def step(...)
         to_llm.step(...)
@@ -649,27 +733,11 @@ module RubyLLM
       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
+        usage_records_by_entry[entry] = Usage.transaction(requires_new: true) do
+          ruby_llm_usages.create!(Usage.attributes_for(entry))
+        end
+      rescue StandardError => e
+        RubyLLM.logger.warn("RubyLLM could not record #{entry.operation} usage: #{e.class}: #{e.message}")
       end
 
       def link_usage_entries(message)
@@ -800,6 +868,8 @@ module RubyLLM
         sync_messages(chat)
         chat.cancellation_checker = proc { consume_persisted_cancellation_request }
         chat.approval_checker = proc { |tool_call| persisted_tool_call_approval(tool_call) }
+        chat.input_checker = proc { |tool_call| persisted_tool_call_input(tool_call) }
+        chat.input_recorder = proc { |tool_call, input| persist_tool_call_input(tool_call, input) }
         install_persistence_callbacks(chat)
       end
 
@@ -813,6 +883,16 @@ module RubyLLM
         self
       end
 
+      def persisted_tool_call_input(tool_call)
+        record = RubyLLM::ActiveRecord::ToolCall.uncached { find_tool_call(tool_call.id) }
+        record.mcp_state if record&.has_attribute?(:mcp_state)
+      end
+
+      def persist_tool_call_input(tool_call, input)
+        record = find_tool_call(tool_call.id)
+        record.update!(mcp_state: input) if record&.has_attribute?(:mcp_state)
+      end
+
       def persisted_tool_call_approval(tool_call)
         record = RubyLLM::ActiveRecord::ToolCall.uncached { find_tool_call(tool_call.id) }
         return unless record&.has_attribute?(:approval)
@@ -827,7 +907,7 @@ module RubyLLM
         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)
+          record.ruby_llm_usage_records.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|
@@ -856,15 +936,17 @@ module RubyLLM
           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)
+          messages_association.create!(role: :system, content: instructions,
+                                       **cache_boundary_attributes(cache_until_here))
         end
       end
 
       # 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)
+        boundary = cache_boundary_attributes(cache_until_here)
+        boundary.delete(:cache_until_here) unless record.has_attribute?(:cache_until_here)
+        attributes = { content: instructions, **boundary }
         changed = attributes.any? { |column, value| record[column] != value }
         record.update!(attributes) if changed
         messages_association.reset
@@ -877,7 +959,7 @@ module RubyLLM
             messages_association.create!(
               role: :system,
               content: instructions,
-              cache_until_here: cache_until_here
+              **cache_boundary_attributes(cache_until_here)
             )
           else
             replace_persisted_system_instructions(instructions, cache_until_here:)
@@ -885,6 +967,13 @@ module RubyLLM
         end
       end
 
+      def cache_boundary_attributes(cache_until_here)
+        ttl = RubyLLM::Message.cache_boundary_options(cache_until_here)[:ttl] if cache_until_here
+        attributes = { cache_until_here: cache_until_here ? true : false }
+        attributes[:cache_ttl] = ttl if ttl || messages_association.klass.column_names.include?('cache_ttl')
+        attributes
+      end
+
       def unpersisted_instructions
         @unpersisted_instructions ||= []
       end
@@ -932,7 +1021,7 @@ module RubyLLM
         transaction do
           @message.assign_attributes(attrs)
           @message.save!
-          tool_call&.update!(result: @message)
+          tool_call&.update!(tool_result_attributes(tool_call, message))
 
           persist_content(@message, message.attachments) if message.attachments.any?
           persist_tool_calls(message.tool_calls) if message.tool_calls.present?
@@ -950,9 +1039,16 @@ module RubyLLM
         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?)
+        assign_supported_attribute(attrs, :cache_ttl, message.cache_ttl)
         attrs
       end
 
+      def tool_result_attributes(tool_call, message)
+        attributes = { result: @message }
+        attributes[:mcp_result] = message.mcp_result&.dump if tool_call.has_attribute?(:mcp_result)
+        attributes
+      end
+
       def assign_supported_attribute(attributes, name, value)
         attributes[name] = value if messages_association.klass.column_names.include?(name.to_s)
       end

data/lib/ruby_llm/active_record/mcp_credential.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module ActiveRecord
+    # RubyLLM's persistence for MCP OAuth credentials, encrypted with Active
+    # Record encryption.
+    class MCPCredential < Record # :nodoc:
+      self.table_name = 'ruby_llm_mcp_credentials'
+
+      belongs_to :owner, polymorphic: true, optional: true
+
+      serialize :data, coder: JSON
+      encrypts :data
+
+      validates :key, presence: true
+
+      class << self
+        def read(key)
+          find_by(key:)&.data
+        end
+
+        def write(key, data, owner: nil)
+          record = find_or_initialize_by(key:)
+          record.update!(data:, owner: owner.is_a?(::ActiveRecord::Base) ? owner : nil)
+        end
+
+        def delete(key)
+          where(key:).delete_all
+        end
+
+        # Updating the row locks it on every adapter; SQLite ignores FOR UPDATE.
+        def synchronize(key)
+          transaction do
+            where(key:).touch_all
+            yield
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/active_record/message_methods.rb

@@ -32,14 +32,20 @@ module RubyLLM
           usage_entries: entries,
           tool_calls: tool_calls,
           tool_call_id: parent_tool_call&.id,
+          mcp_result: mcp_result,
           finish_reason: optional_column(:finish_reason),
           model: entries.reverse.find(&:succeeded?)&.model,
-          cache_until_here: cache_until_here?
+          cache_until_here: cache_until_here?,
+          cache_ttl: cache_ttl
         )
       end
 
       def ruby_llm_usage_entries # :nodoc:
-        ruby_llm_usages.map(&:to_entry)
+        ruby_llm_usage_records.map(&:to_entry)
+      end
+
+      def ruby_llm_usage_records # :nodoc:
+        preloaded_records(:ruby_llm_usages) || ruby_llm_usages.to_a
       end
 
       # Returns the model ID from the last successful attempt, or +nil+.
@@ -52,9 +58,13 @@ module RubyLLM
         to_llm.model_info
       end
 
-      # Marks the message as a prompt cache boundary and returns it.
-      def cache_until_here
-        update!(cache_until_here: true)
+      # Marks the message as a prompt cache boundary and returns it. Pass
+      # +ttl:+ to give the boundary its own cache lifetime, as
+      # RubyLLM::Message#cache_until_here describes.
+      def cache_until_here(ttl: nil)
+        attributes = { cache_until_here: true }
+        attributes[:cache_ttl] = ttl if ttl || has_attribute?(:cache_ttl)
+        update!(attributes)
         self
       end
 
@@ -63,6 +73,12 @@ module RubyLLM
         optional_column(:cache_until_here) || false
       end
 
+      # Returns the cache lifetime given to the message's boundary, or +nil+
+      # when the chat's caching policy applies.
+      def cache_ttl
+        optional_column(:cache_ttl) if cache_until_here?
+      end
+
       # Returns the reasoning the model returned as a RubyLLM::Thinking, or +nil+.
       def thinking
         RubyLLM::Thinking.build(
@@ -94,7 +110,8 @@ module RubyLLM
 
       # Returns the tool calls as RubyLLM::ToolCall values keyed by call id.
       def tool_calls
-        ruby_llm_tool_calls.to_h { |record| [record.tool_call_id, record.to_llm] }
+        records = preloaded_records(:ruby_llm_tool_calls) || ruby_llm_tool_calls
+        records.to_h { |record| [record.tool_call_id, record.to_llm] }
       end
 
       # Returns the tool call this message answers, or +nil+.
@@ -102,6 +119,15 @@ module RubyLLM
         ruby_llm_parent_tool_call&.to_llm
       end
 
+      # Returns the RubyLLM::MCP::Result kept for an MCP tool with a UI, as
+      # RubyLLM::Message#mcp_result does, or +nil+. It persists on the tool
+      # call, so the UI renders again after a reload.
+      def mcp_result
+        record = ruby_llm_parent_tool_call
+        data = record[:mcp_result] if record&.has_attribute?(:mcp_result)
+        RubyLLM::MCP::Result.load(data) if data
+      end
+
       # Returns the message records that answer this message's tool calls.
       def tool_results
         ruby_llm_tool_calls.filter_map(&:result)
@@ -183,7 +209,6 @@ module RubyLLM
         action_text_attachments = action_text_attachment_sources(content)
         return [] unless content_attachments?(action_text_attachments)
 
-        @_tempfiles = []
         collect_attachments(action_text_attachments)
       end
     end

data/lib/ruby_llm/active_record/model.rb

@@ -8,7 +8,7 @@ module RubyLLM
     # point at it through +belongs_to :model+. The +listed+ scope returns the
     # models the provider still serves; +unlisted+ returns the ones kept only
     # because a chat references them.
-    class Model < ::ActiveRecord::Base
+    class Model < Record
       self.table_name = 'ruby_llm_models'
 
       validates :model_id, presence: true, uniqueness: { scope: :provider }

data/lib/ruby_llm/active_record/provider_file.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module ActiveRecord
+    # RubyLLM's record of the provider files each Active Storage blob was
+    # uploaded to, so a chat loaded in another process sends the file it
+    # already uploaded instead of uploading the blob again. Rows name the
+    # blob by its key, which Active Storage never hands out twice, so a row
+    # left by a deleted blob can never match a newer one.
+    class ProviderFile < Record # :nodoc:
+      self.table_name = 'ruby_llm_provider_files'
+
+      validates :blob_key, :provider, :account, :file_id, presence: true
+
+      # Returns a store for the blob whose key the block returns, or +nil+
+      # while the table has not been migrated.
+      def self.store(&)
+        Store.new(&) if table_exists?
+      end
+
+      # Deletes the uploads recorded for +blob+ as Active Storage destroys it.
+      def self.forget_blob(blob)
+        where(blob_key: blob.key).delete_all if table_exists?
+      end
+
+      # The uploads of one blob. The block finds the blob the first time a
+      # provider asks, so attachments that are never uploaded cost nothing.
+      class Store # :nodoc:
+        def initialize(&blob_key)
+          @find_blob_key = blob_key
+        end
+
+        def fetch(provider:, account:)
+          record = blob_key && ProviderFile.find_by(blob_key:, provider:, account:)
+          RubyLLM::UploadedFile.new(id: record.file_id, provider:, expires_at: record.expires_at) if record
+        end
+
+        # Another process can record an upload of the same blob first; its
+        # file serves as well as this one.
+        def store(upload, provider:, account:)
+          return unless blob_key
+
+          ProviderFile.find_or_initialize_by(blob_key:, provider:, account:)
+                      .update!(file_id: upload.id, expires_at: upload.expires_at)
+        rescue ::ActiveRecord::RecordNotUnique
+          nil
+        end
+
+        # Only a row that still names the missing file goes: another process
+        # may have recorded a newer upload meanwhile.
+        def forget(id, provider:, account:)
+          ProviderFile.where(blob_key:, provider:, account:, file_id: id).delete_all if blob_key
+        end
+
+        private
+
+        def blob_key
+          @blob_key = @find_blob_key.call unless defined?(@blob_key)
+          @blob_key
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/active_record/record.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module ActiveRecord
+    # Abstract base for RubyLLM's model, tool-call, usage, and batch records.
+    # Inherits the application's default connection unless configured otherwise.
+    # Keep these records on the same database and connection pool as chats and
+    # messages to preserve foreign keys and shared transactions.
+    #
+    #   Rails.application.config.to_prepare do
+    #     RubyLLM::ActiveRecord::Record.connection_specification_name =
+    #       LlmRecord.connection_specification_name
+    #   end
+    #
+    class Record < ::ActiveRecord::Base
+      self.abstract_class = true
+    end
+  end
+end

data/lib/ruby_llm/active_record/tool_call.rb

@@ -5,7 +5,7 @@ require 'ruby_llm/active_record/payload_helpers'
 module RubyLLM
   module ActiveRecord
     # RubyLLM's private persistence for provider tool calls.
-    class ToolCall < ::ActiveRecord::Base # :nodoc:
+    class ToolCall < Record # :nodoc:
       self.table_name = 'ruby_llm_tool_calls'
 
       include PayloadHelpers

data/lib/ruby_llm/active_record/usage.rb

@@ -3,11 +3,12 @@
 module RubyLLM
   module ActiveRecord
     # RubyLLM's private normalized persistence for one provider attempt.
-    class Usage < ::ActiveRecord::Base # :nodoc:
+    class Usage < Record # :nodoc:
       self.table_name = 'ruby_llm_usages'
 
-      belongs_to :chat, polymorphic: true
+      belongs_to :chat, polymorphic: true, optional: true
       belongs_to :message, polymorphic: true, optional: true
+      belongs_to :owner, polymorphic: true, optional: true
 
       validates :operation, inclusion: { in: ::RubyLLM::Accounting::Usage::Entry::OPERATIONS.map(&:to_s) }
       validates :status, inclusion: { in: ::RubyLLM::Accounting::Usage::Entry::STATUSES.map(&:to_s) }
@@ -15,13 +16,62 @@ module RubyLLM
 
       scope :chronological, -> { order(created_at: :asc, id: :asc) }
 
+      # Writes an entry that belongs to no chat record, attributed to its
+      # owner. A ledger that cannot store such rows yet is skipped, and a
+      # failed write is logged so it never breaks the operation that billed
+      # it. A savepoint keeps a failed insert from aborting the caller's
+      # transaction, and a thread that held no connection gives back the
+      # one it borrowed.
+      def self.record(entry)
+        return unless entry.model
+
+        connection_pool.with_connection do
+          next unless ledger_available?
+
+          transaction(requires_new: true) { create!(attributes_for(entry)) }
+        end
+      rescue StandardError => e
+        RubyLLM.logger.warn("RubyLLM could not record #{entry.operation} usage: #{e.class}: #{e.message}")
+        nil
+      end
+
+      def self.ledger_available?
+        table_exists? && columns_hash['chat_id']&.null
+      end
+
+      def self.attributes_for(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
+        }
+        attributes[:server_tool_use] = tokens.server_tool_use if column_names.include?('server_tool_use')
+        attributes[:owner] = entry.owner if entry.owner.is_a?(::ActiveRecord::Base) && column_names.include?('owner_id')
+        attributes
+      end
+
       def tokens
         RubyLLM::Tokens.new(
           input: input_tokens,
           output: output_tokens,
           cache_read: cache_read_tokens,
           cache_write: cache_write_tokens,
-          thinking: thinking_tokens
+          thinking: thinking_tokens,
+          server_tool_use: (self[:server_tool_use] if has_attribute?(:server_tool_use))
         )
       end