ruby_llm 2.1.1 → 2.2.0

data/lib/generators/ruby_llm/sandbox/sandbox_generator.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require 'rails/generators'
+
+module RubyLLM
+  module Generators
+    # Generator for sandbox classes that code execution runs model code in.
+    class SandboxGenerator < Rails::Generators::NamedBase
+      INITIALIZER = 'config/initializers/ruby_llm.rb'
+      CONFIGURE_BLOCK = /^RubyLLM\.configure do \|config\|\n/
+      MXC_REQUIREMENT = '>= 0.4.1'
+
+      source_root File.expand_path('templates', __dir__)
+
+      namespace 'ruby_llm:sandbox'
+
+      desc 'Creates a sandbox class for code execution in app/sandboxes'
+
+      def check_sandbox_class_collision
+        class_collisions sandbox_class_name
+      end
+
+      def create_sandbox_file
+        template 'sandbox.rb.tt', File.join('app/sandboxes', class_path, "#{sandbox_file_name}.rb")
+      end
+
+      def add_mxc_gem
+        gemfile = File.join(destination_root, 'Gemfile')
+        return unless File.exist?(gemfile)
+
+        existing = File.read(gemfile)[/^\s*gem\s+["']mxc["'].*$/]
+        return warn_about_mxc_constraint(existing) if existing
+
+        gem 'mxc', MXC_REQUIREMENT
+        say "\nRun bundle install" if behavior == :invoke
+      end
+
+      def default_to_new_sandbox
+        initializer = File.join(destination_root, INITIALIZER)
+        return unless File.exist?(initializer)
+        return if behavior == :invoke && File.read(initializer).match?(/^\s*config\.sandbox\s*(\|\|)?=/)
+
+        unless File.read(initializer).match?(CONFIGURE_BLOCK)
+          return say %(Set config.sandbox = "#{sandbox_class_name}" where you configure RubyLLM), :yellow
+        end
+
+        inject_into_file INITIALIZER, %(  config.sandbox = "#{sandbox_class_name}"\n), after: CONFIGURE_BLOCK
+      end
+
+      private
+
+      def warn_about_mxc_constraint(line)
+        requirement = line[/["']mxc["']\s*,\s*["']([^"']+)["']/, 1]
+        return if requirement.nil? || Gem::Requirement.new(requirement).satisfied_by?(Gem::Version.new('0.4.1'))
+
+        say %(Change the mxc requirement in your Gemfile to "#{MXC_REQUIREMENT}"), :yellow
+      end
+
+      def sandbox_class_name
+        class_name.end_with?('Sandbox') ? class_name : "#{class_name}Sandbox"
+      end
+
+      def sandbox_file_name
+        sandbox_class_name.demodulize.underscore
+      end
+    end
+  end
+end

data/lib/generators/ruby_llm/sandbox/templates/sandbox.rb.tt

@@ -0,0 +1,3 @@
+class <%= sandbox_class_name %> < MXC::Sandbox
+  timeout 120
+end

data/lib/generators/ruby_llm/skill/skill_generator.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require 'fileutils'
+require 'rails/generators'
+require 'tmpdir'
+
+module RubyLLM
+  module Generators
+    # Generator for skills in the open Agent Skills format.
+    class SkillGenerator < Rails::Generators::Base
+      namespace 'ruby_llm:skill'
+
+      def self.exit_on_failure?
+        true
+      end
+
+      argument :skill_name, type: :string, banner: 'name'
+      argument :description, type: :string, banner: '"What it does and when to use it"'
+
+      desc 'Creates a skill folder with its SKILL.md in app/skills'
+
+      def check_skill
+        Dir.mktmpdir do |root|
+          folder = File.join(root, skill_name)
+          FileUtils.mkdir_p(folder)
+          File.write(File.join(folder, 'SKILL.md'), skill_markdown)
+          RubyLLM::Skill.new(folder)
+        end
+      rescue RubyLLM::InvalidSkillError => e
+        raise Thor::Error, e.message.sub(%r{\A\S*/SKILL\.md }, '')
+      end
+
+      def create_skill_file
+        create_file File.join(skill_path, 'SKILL.md'), skill_markdown
+      end
+
+      private
+
+      def skill_path
+        File.join('app/skills', skill_name)
+      end
+
+      def skill_markdown
+        "---\nname: #{skill_name}\ndescription: #{description.to_json}\n---\n"
+      end
+    end
+  end
+end

data/lib/generators/ruby_llm/tool/templates/tailwind/tool_call.html.erb.tt

@@ -1,7 +1,7 @@
 <div id="<%= message_variable_name %>_tool_call_<%%= tool_call.id %>" class="w-full sm:w-auto my-5 space-y-3 rounded-md px-3 py-2 bg-gray-50 border border-gray-200">
   <div>
     <span class="inline-block rounded px-2 py-0.5 text-xs font-medium bg-gray-200 text-gray-700">
-      <%= tool_display_name %> Call
+      <%= tool_partial_name.humanize %> call
     </span>
   </div>
 

data/lib/generators/ruby_llm/tool/templates/tailwind/tool_result.html.erb.tt

@@ -8,7 +8,7 @@
   <div id="<%= message_variable_name %>_<%%= tool.id %>" class="w-full sm:w-auto my-5 space-y-3 rounded-md px-3 py-2 bg-gray-50 border border-gray-200">
     <div>
       <span class="inline-block rounded px-2 py-0.5 text-xs font-medium bg-gray-200 text-gray-700">
-        <%= tool_display_name %> Result
+        <%= tool_partial_name.humanize %> result
       </span>
     </div>
 

data/lib/generators/ruby_llm/tool/templates/tool.rb.tt

@@ -2,8 +2,5 @@ class <%= tool_class_name %> < RubyLLM::Tool
   description "TODO: describe what this tool does"
 
   def execute
-    # TODO: return something to the LLM.
-    # Anything you return here gets converted to a string.
-    "TODO: implement <%= tool_class_name %>"
   end
 end

data/lib/generators/ruby_llm/tool/templates/tool_call.html.erb.tt

@@ -1,6 +1,6 @@
 <div id="<%= message_variable_name %>_tool_call_<%%= tool_call.id %>" class="<%= message_variable_name %>"
      style="margin-bottom: 20px; padding: 10px; border-left: 3px solid #6b7280; background: #f9fafb;">
-  <div style="font-weight: bold; margin-bottom: 5px;"><%= tool_display_name %> Call</div>
+  <div style="font-weight: bold; margin-bottom: 5px;"><%= tool_partial_name.humanize %> call</div>
   <pre style="white-space: pre-wrap; margin: 0;"><%%= tool_call.name %>(<%%= tool_call.arguments.map { |k, v| "#{k}: #{v.inspect}" }.join(", ") %>)</pre>
   <div style="font-size: 0.85em; color: #666; margin-top: 5px;">
     <%%= message.created_at&.strftime("%I:%M %p") %>

data/lib/generators/ruby_llm/tool/templates/tool_result.html.erb.tt

@@ -7,7 +7,7 @@
 <%% else %>
   <div id="<%= message_variable_name %>_<%%= tool.id %>" class="<%= message_variable_name %>"
        style="margin-bottom: 20px; padding: 10px; border-left: 3px solid #6b7280; background: #f9fafb;">
-    <div style="font-weight: bold; margin-bottom: 5px;"><%= tool_display_name %> Result</div>
+    <div style="font-weight: bold; margin-bottom: 5px;"><%= tool_partial_name.humanize %> result</div>
     <pre style="white-space: pre-wrap; margin: 0;"><%%= tool.content.presence || "(no output)" %></pre>
     <div style="font-size: 0.85em; color: #666; margin-top: 5px;">
       <%%= tool.created_at&.strftime("%I:%M %p") %>

data/lib/generators/ruby_llm/tool/tool_generator.rb

@@ -49,13 +49,9 @@ module RubyLLM
         tool_class_name.demodulize.underscore
       end
 
-      def tool_display_name
-        tool_class_name.demodulize.delete_suffix('Tool')
-      end
-
       # Mirrors RubyLLM::Tool.tool_name, the key the runtime partial lookup uses.
       def tool_partial_name
-        tool_file_name.delete_suffix('_tool')
+        RubyLLM::Support::Utils.underscore(tool_class_name).delete_suffix('_tool').gsub('::', '__')
       end
     end
   end

data/lib/generators/ruby_llm/upgrade/templates/upgrade_ruby_llm_to_2_2.rb.tt

@@ -0,0 +1,164 @@
+class <%= 'upgrade_ruby_llm_to_2_2'.camelize %> < ActiveRecord::Migration<%= migration_version %>
+  def up
+    unless column_exists?(:ruby_llm_tool_calls, :provider_call_id)
+      if column_exists?(:ruby_llm_tool_calls, :provider_id)
+        rename_column :ruby_llm_tool_calls, :provider_id, :provider_call_id
+      else
+        add_column :ruby_llm_tool_calls, :provider_call_id, :string
+      end
+    end
+
+    add_column :ruby_llm_tool_calls, :namespace, :string unless column_exists?(:ruby_llm_tool_calls, :namespace)
+
+    unless column_exists?(:<%= chat_table_name %>, :claim_token)
+      add_column :<%= chat_table_name %>, :claim_token, :string
+    end
+
+    unless column_exists?(:<%= chat_table_name %>, :claimed_until)
+      add_column :<%= chat_table_name %>, :claimed_until, :datetime
+    end
+
+    unless index_exists?(:<%= chat_table_name %>, :claimed_until)
+      add_index :<%= chat_table_name %>, :claimed_until<%= ', where: "claim_token IS NOT NULL"' unless mysql? %>
+    end
+
+    unless column_exists?(:ruby_llm_tool_calls, :claimed_until)
+      add_column :ruby_llm_tool_calls, :claimed_until, :datetime
+    end
+
+    unless column_exists?(:<%= message_table_name %>, :phase)
+      add_column :<%= message_table_name %>, :phase, :string
+    end
+
+    unless column_exists?(:<%= message_table_name %>, :continuation)
+      add_column :<%= message_table_name %>, :continuation, :text<%= ', size: :long' if mysql? %>
+    end
+<% if mysql? -%>
+
+    message_columns = %i[content thinking_text thinking_signature].select do |column|
+      column_exists?(:<%= message_table_name %>, column)
+    end
+    if message_columns.any?
+      change_table :<%= message_table_name %>, bulk: true do |t|
+        message_columns.each { |column| t.change column, :text, size: :long }
+      end
+    end
+    change_column :ruby_llm_tool_calls, :thought_signature, :text, size: :long
+<% end -%>
+
+    %i[workflow_id workflow_name workflow_step_id workflow_step_name].each do |column|
+      add_column :ruby_llm_usages, column, :string unless column_exists?(:ruby_llm_usages, column)
+    end
+    add_index :ruby_llm_usages, :workflow_id unless index_exists?(:ruby_llm_usages, :workflow_id)
+
+    add_column :ruby_llm_usages, :characters, :integer unless column_exists?(:ruby_llm_usages, :characters)
+    unless column_exists?(:ruby_llm_usages, :audio_seconds)
+      add_column :ruby_llm_usages, :audio_seconds, :decimal, precision: 16, scale: 6
+    end
+    add_column :ruby_llm_usages, :search_units, :integer unless column_exists?(:ruby_llm_usages, :search_units)
+    %i[characters_cost audio_seconds_cost search_units_cost].each do |column|
+      unless column_exists?(:ruby_llm_usages, column)
+        add_column :ruby_llm_usages, column, :decimal, precision: 16, scale: 10
+      end
+    end
+
+    unless table_exists?(:ruby_llm_queued_messages)
+      create_table :ruby_llm_queued_messages, id: :<%= reference_type %> do |t|
+        t.references :chat, polymorphic: true, null: false, type: :<%= reference_type %>, index: false
+        t.references :sender, polymorphic: true, type: :<%= reference_type %>, index: false
+        t.text :content<%= ', size: :long' if mysql? %>
+        t.boolean :steer, default: true, null: false
+        t.timestamps
+
+        t.index [ :chat_type, :chat_id ]
+        t.index [ :sender_type, :sender_id ]
+      end
+    end
+
+    unless table_exists?(:ruby_llm_chat_files)
+      create_table :ruby_llm_chat_files, id: :<%= reference_type %> do |t|
+        t.references :chat, polymorphic: true, null: false, type: :<%= reference_type %>, index: false
+        t.references :owner, polymorphic: true, type: :<%= reference_type %>, index: false
+        t.string :path, null: false
+        t.string :content_type
+        t.bigint :byte_size
+        t.string :digest
+        t.text :text<%= ', size: :long' if mysql? %>
+        t.string :read_digest
+        t.timestamps
+
+        t.index [ :chat_type, :chat_id, :path ]
+        t.index [ :owner_type, :owner_id ]
+      end
+    end
+    unless column_exists?(:ruby_llm_chat_files, :read_digest)
+      add_column :ruby_llm_chat_files, :read_digest, :string
+    end
+
+    allow_code_execution_usage unless usage_accepts_code_execution?
+    # Asks again: a check of the application's may still refuse the row.
+    warn_about_code_execution_usage unless usage_accepts_code_execution?
+  end
+
+  def down
+    drop_table :ruby_llm_chat_files, if_exists: true
+    drop_table :ruby_llm_queued_messages, if_exists: true
+    remove_column :<%= message_table_name %>, :continuation
+    %i[search_units_cost audio_seconds_cost characters_cost search_units audio_seconds characters].each do |column|
+      remove_column :ruby_llm_usages, column
+    end
+    remove_index :ruby_llm_usages, :workflow_id, if_exists: true
+    remove_column :ruby_llm_usages, :workflow_step_name
+    remove_column :ruby_llm_usages, :workflow_step_id
+    remove_column :ruby_llm_usages, :workflow_name
+    remove_column :ruby_llm_usages, :workflow_id
+    remove_column :<%= message_table_name %>, :phase
+    remove_column :ruby_llm_tool_calls, :namespace
+    remove_column :ruby_llm_tool_calls, :claimed_until
+    remove_index :<%= chat_table_name %>, :claimed_until, if_exists: true
+    remove_column :<%= chat_table_name %>, :claimed_until
+    remove_column :<%= chat_table_name %>, :claim_token
+    remove_column :ruby_llm_tool_calls, :provider_call_id
+  end
+
+  private
+
+  # Asks the database, since each one prints a check differently: a check
+  # refusing the row leaves code runs unmetered, with only a warning at run time.
+  def usage_accepts_code_execution?
+    connection.transaction(requires_new: true) do
+      execute <<~SQL.squish
+        INSERT INTO ruby_llm_usages (operation, provider, model, status, created_at, updated_at)
+        VALUES ('code_execution', 'sandbox', 'upgrade check', 'succeeded', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+      SQL
+      raise ActiveRecord::Rollback
+    end
+    true
+  rescue ActiveRecord::StatementInvalid
+    false
+  end
+
+  def allow_code_execution_usage
+    lists = check_constraints(:ruby_llm_usages).select { |check| operation_list?(check.expression) }
+    lists.each { |check| remove_check_constraint :ruby_llm_usages, name: check.name }
+    return if lists.empty?
+
+    add_check_constraint :ruby_llm_usages, "operation<%= '::text' if postgresql? %> IN (<%= usage_operations_sql %>)"
+  end
+
+  # The list the install generator wrote, as each database prints it. A check
+  # with other conditions is the application's, and stays.
+  def operation_list?(expression)
+    names = expression.scan(/'(\w+)'/).flatten
+    expression.match?(/\A\(*[`"]?operation[`"]?\)*(?:::\w+)?\s*(?:in\b|=\s*any\b)/i) &&
+      !expression.match?(/\b(?:and|or|not)\b|<>|!=/i) && names.size > 1 && (names - %w[<%= ::RubyLLM::Accounting::Usage::Entry::OPERATIONS.join(' ') %>]).empty?
+  end
+
+  def warn_about_code_execution_usage
+    checks = check_constraints(:ruby_llm_usages).select { |check| check.expression.match?(/\boperation\b/) }
+                                                 .map(&:name).join(', ')
+    say "RubyLLM: ruby_llm_usages refuses 'code_execution' rows, so code runs will go unmetered with only a " \
+        "warning in the log. Check these constraints, or any NOT NULL column you added: #{checks}. Allow the " \
+        'value in a new migration: see "Code Runs Write Usage Rows" in the upgrading guide.'
+  end
+end

data/lib/generators/ruby_llm/upgrade/upgrade_generator.rb

@@ -6,7 +6,7 @@ require_relative '../generator_helpers'
 
 module RubyLLM
   module Generators
-    # Upgrades the schema of an application on RubyLLM 2.0 to 2.1.
+    # Upgrades the schema of an application on RubyLLM 2.1 to 2.2.
     class UpgradeGenerator < Rails::Generators::Base
       include Rails::Generators::Migration
       include RubyLLM::Generators::GeneratorHelpers
@@ -16,7 +16,7 @@ module RubyLLM
 
       argument :model_mappings, type: :array, default: [], banner: 'chat:ChatName message:MessageName'
 
-      desc 'Upgrades a RubyLLM 2.0 Rails schema to 2.1\n' \
+      desc "Upgrades a RubyLLM 2.1 Rails schema to 2.2\n" \
            'Usage: bin/rails g ruby_llm:upgrade [chat:ChatName] [message:MessageName]'
 
       def self.next_migration_number(dirname)
@@ -24,7 +24,45 @@ module RubyLLM
       end
 
       def create_migration_file
-        migration_template 'upgrade_ruby_llm_to_2_1.rb.tt', 'db/migrate/upgrade_ruby_llm_to_2_1.rb'
+        migration_template 'upgrade_ruby_llm_to_2_2.rb.tt', 'db/migrate/upgrade_ruby_llm_to_2_2.rb'
+      end
+
+      def create_application_agent_file
+        create_application_agent
+      end
+
+      def remove_generated_message_code_from_message_model
+        return if chat_ui_uses_generated_broadcasting?
+
+        enable_message_broadcasting if remove_generated_message_code
+      end
+
+      def install_active_storage
+        return if table_exists?('active_storage_blobs')
+
+        rails_command 'active_storage:install'
+      end
+
+      def show_next_step
+        say "\nRun bin/rails db:migrate" if behavior == :invoke
+      end
+
+      private
+
+      def chat_ui_uses_generated_broadcasting?
+        return false unless behavior == :invoke
+
+        model = File.join(destination_root, "app/models/#{message_model_name.underscore}.rb")
+        return false unless File.exist?(model)
+
+        source = File.read(model)
+        return false unless generated_broadcasting_blocks.any? { |block| source.include?(block) }
+
+        stream = %("#{chat_variable_name}_\#{)
+        Dir.glob(File.join(destination_root, 'app/**/*.{rb,erb}')).without(model).any? do |path|
+          source = File.read(path)
+          source.include?('broadcast_append_chunk') || source.include?(stream)
+        end
       end
     end
   end

data/lib/ruby_llm/accounting/operation.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Accounting
+    # The body the standalone operations share: resolve the model, open the
+    # instrumentation event, attribute the provider call to its owner, and
+    # report the result on the event. +run+ takes the event fields particular
+    # to the operation, in event order, and yields the provider and the
+    # resolved model.
+    #
+    # +details:+ adds fields after the response model. +report:+ replaces the
+    # reporting of an operation whose result does not read like the others.
+    # +usage: false+ leaves tokens and cost off the event. +resolved:+ takes
+    # a Model already resolved, such as a fallback given as one, and skips
+    # the lookup.
+    class Operation # :nodoc: all
+      def initialize(name, model, owner:, context: nil, provider: nil, assume_model_exists: false,
+                     default: nil, operation: nil, category: nil, usage: true, details: nil, report: nil,
+                     resolved: nil)
+        @name = name
+        @owner = owner
+        @category = category
+        @usage = usage
+        @details = details
+        @report = report
+        @config = context&.config || RubyLLM.config
+        @model, @provider = if resolved
+                              [resolved, Provider.resolve!(resolved.provider).new(@config)]
+                            else
+                              resolve(model, default, operation, provider:, assume_model_exists:)
+                            end
+      end
+
+      def run(**fields)
+        RubyLLM.instrument("#{@name}.ruby_llm", payload(fields), config: @config) do |event|
+          result = Usage.owned_by(@owner) { yield @provider, @model }
+          report(event, result)
+          result
+        end
+      end
+
+      private
+
+      def resolve(model, default, operation, **options)
+        default_model = default && @config.public_send(default)
+        options[:config] = @config
+        return Models.resolve(model, **options, operation:, default_model:) if operation
+
+        Models.resolve(model || default_model, **options)
+      end
+
+      def payload(fields)
+        common = { provider: @provider.slug, provider_class: @provider.name, model: @model&.id, model_info: @model }
+        return common.merge(fields) unless @usage
+
+        tokens = Tokens.new
+        common.merge(fields, tokens:, cost: Cost.new(tokens:, model: @model, **{ category: @category }.compact))
+      end
+
+      def report(event, result)
+        return @report.call(event, result) if @report
+
+        event[:result] = result
+        event[:response_model] = result.model
+        @details&.call(event, result)
+        return unless @usage
+
+        event[:tokens] = result.tokens
+        event[:cost] = result.cost
+      end
+    end
+  end
+end

data/lib/ruby_llm/accounting/usage.rb

@@ -31,6 +31,10 @@ module RubyLLM
           owner.nil? ? yield : with_owner(owner, &)
         end
 
+        def workflow
+          Support::Instrumentation.current_workflow
+        end
+
         def record(entry)
           ledger&.record(entry)
         end
@@ -40,20 +44,24 @@ module RubyLLM
           instrument(entry, config:)
         end
 
+        # An attempt reports the workflow it started in, even when it
+        # finishes outside that block, such as a polled video job.
         def instrument(entry, config:)
-          RubyLLM.instrument(
-            'usage.ruby_llm',
-            {
-              operation: entry.operation,
-              provider: entry.provider,
-              model: entry.model,
-              status: entry.status,
-              tokens: entry.tokens,
-              cost: entry.cost,
-              owner: entry.owner
-            },
-            config:
-          )
+          Support::Instrumentation.with_workflow(entry.workflow || workflow) do
+            RubyLLM.instrument(
+              'usage.ruby_llm',
+              {
+                operation: entry.operation,
+                provider: entry.provider,
+                model: entry.model,
+                status: entry.status,
+                tokens: entry.tokens,
+                cost: entry.cost,
+                owner: entry.owner
+              },
+              config:
+            )
+          end
         end
       end
 
@@ -82,12 +90,26 @@ module RubyLLM
         end
       end
 
+      # The registry Model of a result that names one. The usage tracker
+      # assigns it; a result built by hand looks up #model.
+      module ModelInfo
+        attr_writer :model_info # :nodoc:
+
+        def model_info # :nodoc:
+          @model_info ||= RubyLLM.models.find(model)
+        rescue ModelNotFoundError
+          nil
+        end
+      end
+
       # One physical provider request attempt and its accounting facts.
       class Entry
-        OPERATIONS = %i[chat embedding moderation image speech transcription ocr rerank judgment video research].freeze
+        OPERATIONS = %i[
+          chat embedding moderation image speech transcription ocr rerank judgment video research code_execution
+        ].freeze
         STATUSES = %i[pending succeeded failed cancelled].freeze
 
-        attr_reader :operation, :provider, :model, :status, :tokens, :owner
+        attr_reader :operation, :provider, :model, :status, :tokens, :owner, :workflow
         attr_accessor :message
 
         include Support::Inspectable
@@ -97,7 +119,7 @@ module RubyLLM
         end
 
         def initialize(operation:, provider:, model:, status: :pending, tokens: nil, cost: nil, message: nil,
-                       owner: nil)
+                       owner: nil, workflow: nil)
           @operation = operation.to_sym
           @provider = provider.to_s
           @model = model&.to_s
@@ -106,6 +128,7 @@ module RubyLLM
           @cost = cost
           @message = message
           @owner = owner
+          @workflow = workflow
           validate!
         end
 
@@ -161,6 +184,11 @@ module RubyLLM
           judgment: :text_tokens
         }.freeze
 
+        CONNECT_FAILURES = [
+          Errno::ECONNREFUSED, Errno::EHOSTUNREACH, Errno::ENETUNREACH, SocketError, Net::OpenTimeout
+        ].freeze
+        TLS_HANDSHAKE_FAILURE = /\ASSL_connect|does not match the server certificate/
+
         attr_reader :entries
 
         def initialize(operation:, provider:, model:, config:, on_finish: nil, model_id: nil)
@@ -176,13 +204,14 @@ module RubyLLM
 
         def start
           entry = Entry.new(operation: @operation, provider: @provider.slug, model: @model_id || @model_info&.id,
-                            owner: Usage.owner)
+                            owner: Usage.owner, workflow: Usage.workflow)
           @entries << entry
           @pending << entry
           entry
         end
 
         def observe(chunk)
+          @answered = @pending.last
           observe_tokens(chunk.tokens) if chunk.respond_to?(:tokens)
         end
 
@@ -207,6 +236,12 @@ module RubyLLM
           @pending.dup.each { |entry| fail_attempt(entry, error) }
         end
 
+        # Finishes the attempts an exception outside StandardError
+        # interrupted, such as a stopped task or a shutdown.
+        def cancel_pending
+          @pending.dup.each { |entry| finish(entry, status: :cancelled, tokens: entry.tokens) }
+        end
+
         def succeed(result)
           # A request that produced several results, like a multi-image
           # generation, is billed once: the first result carries the call.
@@ -246,23 +281,38 @@ module RubyLLM
         end
 
         # A request that never reached the provider, or that the provider
-        # refused before running it, cannot have been billed; possibly-billed
-        # failures keep their usage unknown.
+        # answered with an error before streaming anything, cannot have been
+        # billed; possibly-billed failures keep their usage unknown.
         def failure_tokens(entry, error)
           return entry.tokens unless @model_info
-          return entry.tokens if entry.usage_available? || (!never_sent?(error) && !refused?(error))
+          return entry.tokens if entry.usage_available? || (!never_sent?(entry, error) && !refused?(entry, error))
 
           Tokens.new(input: 0, output: 0)
         end
 
-        def never_sent?(error)
-          error.is_a?(Faraday::ConnectionFailed) || error.is_a?(Faraday::SSLError)
+        # Faraday also raises its connection errors for a connection that
+        # drops after the request was sent, so only the wrapped cause tells
+        # whether the connection was ever established.
+        def never_sent?(entry, error)
+          return false if entry.equal?(@answered)
+
+          error = transport_error(error)
+          return false unless error.is_a?(Faraday::ConnectionFailed) || error.is_a?(Faraday::SSLError)
+
+          cause = error.wrapped_exception || error.cause
+          CONNECT_FAILURES.any? { |failure| cause.is_a?(failure) } ||
+            (cause.is_a?(OpenSSL::SSL::SSLError) && cause.message.match?(TLS_HANDSHAKE_FAILURE))
+        end
+
+        def transport_error(error)
+          error.is_a?(ConnectionError) ? error.cause : error
         end
 
-        def refused?(error)
+        def refused?(entry, error)
           return false unless error.is_a?(Error) && error.response.respond_to?(:status)
 
-          (400..499).cover?(error.response.status)
+          status = error.response.status
+          (400..499).cover?(status) || ((500..599).cover?(status) && !entry.equal?(@answered))
         end
 
         def finish(entry, status:, tokens:, **details)