Merge pull request #929 from basecamp/fizzy-ask-spend-limits

Conversation cost limits

app/assets/stylesheets/conversation.css

@@ -20,6 +20,7 @@
 
   .conversation__message {
     display: flex;
+    flex-direction: column;
 
     .action-text-content {
       border-radius: 0.5rem;
@@ -51,7 +52,7 @@
   }
 
   .conversation__message--assistant {
-    justify-content: flex-start;
+    align-items: flex-start;
 
     .action-text-content {
       margin-inline-end: 15%;
@@ -67,8 +68,15 @@
     color: var(--color-negative);
   }
 
+  .conversation__message--failed[data-error]::after {
+    color: var(--color-negative);
+    content: attr(data-error);
+    display: block;
+    padding: var(--inline-space);
+  }
+
   .conversation__message--user {
-    justify-content: flex-end;
+    align-items: flex-end;
 
     .action-text-content {
       background-color: var(--color-selected);

app/controllers/conversations/messages_controller.rb

@@ -7,6 +7,10 @@ class Conversations::MessagesController < ApplicationController
 
   def create
     @conversation.ask(question, **message_params)
+  rescue Ai::Quota::UsageExceedsQuotaError
+    render json: { error: "You've depleted your quota" }, status: :too_many_requests
+  rescue Conversation::InvalidStateError
+    render json: { error: "Fizzy is still working on an answer to your last question" }, status: :conflict
   end
 
   private

app/javascript/controllers/conversation/composer_controller.js

@@ -27,7 +27,10 @@ export default class extends Controller {
     if (event.detail.success) {
       this.#reset()
     } else {
-      this.conversationMessagesOutlet.failPendingMessage(this.clientMessageIdInputTarget.value)
+      this.#failPendingMessage(
+        this.clientMessageIdInputTarget.value, 
+        event.detail.fetchResponse
+      )
     }
   }
 
@@ -60,4 +63,18 @@ export default class extends Controller {
     this.textInputTarget.value = ""
     this.clientMessageIdInputTarget.value = ""
   }
+
+  async #failPendingMessage(clientMessageId, response) {
+    let errorMessage = null
+
+    if (response?.contentType?.includes('application/json')) {
+      const error = JSON.parse(await response.responseText)
+      errorMessage = error.error
+    }
+
+    this.conversationMessagesOutlet.failPendingMessage(
+      this.clientMessageIdInputTarget.value,
+      errorMessage
+    )
+  }
 }

app/javascript/controllers/conversation/messages_controller.js

@@ -12,11 +12,15 @@ export default class extends Controller {
     this.scrollToBottom()
   }
 
-  failPendingMessage(clientMessageId) {
+  failPendingMessage(clientMessageId, errorMessage) {
     const message = this.element.querySelector(`#message_${clientMessageId}`)
 
     if (message) {
       message.classList.add(this.failedMessageClass)
+
+      if (errorMessage) {
+        message.dataset.error = errorMessage
+      }
     }
   }
 

app/models/ai.rb

@@ -0,0 +1,5 @@
+module Ai
+  def self.table_name_prefix
+    "ai_"
+  end
+end

app/models/ai/quota.rb

@@ -0,0 +1,50 @@
+class Ai::Quota < ApplicationRecord
+  class UsageExceedsQuotaError < StandardError; end
+
+  belongs_to :user
+
+  before_create -> { reset }
+
+  validates :limit, presence: true, numericality: { only_integer: true, greater_than_or_equal_to: 0 }
+  validates :used, presence: true, numericality: { only_integer: true, greater_than_or_equal_to: 0 }
+
+  def spend(cost)
+    cost = Money.wrap(cost)
+
+    transaction do
+      reset_if_due
+      increment!(:used, cost.in_microcents)
+    end
+  end
+
+  def ensure_not_depleted
+    reset_if_due
+
+    if depleted?
+      raise UsageExceedsQuotaError
+    end
+  end
+
+  private
+    def reset_if_due
+      reset if due_for_reset?
+    end
+
+    def reset
+      attributes = { used: 0, reset_at: 7.days.from_now }
+
+      if persisted?
+        update(**attributes)
+      else
+        assign_attributes(**attributes)
+      end
+    end
+
+    def due_for_reset?
+      reset_at.before?(Time.current)
+    end
+
+    def depleted?
+      used >= limit
+    end
+end

app/models/ai/quota/money.rb

@@ -0,0 +1,42 @@
+class Ai::Quota::Money < Data.define(:value)
+  CENTS_PER_DOLLAR = 100
+  MICROCENTS_PER_CENT = 1_000_000
+  MICROCENTS_PER_DOLLAR = CENTS_PER_DOLLAR * MICROCENTS_PER_CENT
+  NUMBER_REGEX = /\d+(\.\d+)?/
+
+  class << self
+    def wrap(value)
+      microcents = case value
+      when nil
+        raise ArgumentError, "#{self} can't wrap nil"
+      when self
+        value.value
+      when Integer
+        value
+      when String
+        convert_dollars_to_microcents(BigDecimal(value[NUMBER_REGEX]))
+      else
+        convert_dollars_to_microcents(value)
+      end
+
+      new(microcents)
+    end
+
+    private
+      def convert_dollars_to_microcents(dollars)
+        (dollars.to_d * MICROCENTS_PER_DOLLAR).round.to_i
+      end
+  end
+
+  def to_i
+    in_microcents
+  end
+
+  def in_microcents
+    value
+  end
+
+  def in_dollars
+    in_microcents.to_d / MICROCENTS_PER_DOLLAR
+  end
+end

app/models/conversation.rb

@@ -8,16 +8,9 @@ class Conversation < ApplicationRecord
 
   enum :state, %w[ ready thinking ].index_by(&:itself), default: :ready
 
-  def cost
-    messages.where.not(cost_microcents: nil).sum(:cost_microcents).to_d / 100_000
-  end
-
-  def clear
-    messages.delete_all
-    touch
-  end
-
   def ask(question, **attributes)
+    user.ensure_ai_quota_not_depleted
+
     create_message_with_state_change(**attributes, role: :user, content: question) do
       raise(InvalidStateError, "Can't ask questions while thinking") if thinking?
       thinking!
@@ -25,10 +18,14 @@ class Conversation < ApplicationRecord
   end
 
   def respond(answer, **attributes)
-    create_message_with_state_change(**attributes, role: :assistant, content: answer) do
+    message = create_message_with_state_change(**attributes, role: :assistant, content: answer) do
       raise(InvalidStateError, "Can't respond when not thinking") unless thinking?
       ready!
     end
+
+    user.spend_ai_quota(message.cost) if message.cost
+
+    message
   end
 
   private

app/models/conversation/message.rb

@@ -13,6 +13,10 @@ class Conversation::Message < ApplicationRecord
 
   scope :ordered, -> { order(created_at: :asc, id: :asc) }
 
+  def cost
+    cost_microcents && Ai::Quota::Money.new(cost_microcents)
+  end
+
   def all_emoji?
     content.to_plain_text.all_emoji?
   end

app/models/conversation/message/response_generator/response.rb

@@ -1,6 +1,4 @@
 class Conversation::Message::ResponseGenerator::Response
-  MICROCENTS_PER_DOLLAR = 100_000
-
   attr_reader :answer, :input_tokens, :output_tokens, :model_id, :tool_calls, :tool_call_id
 
   def initialize(answer:, input_tokens:, output_tokens:, model_id:)
@@ -10,10 +8,6 @@ class Conversation::Message::ResponseGenerator::Response
     @model_id = model_id
   end
 
-  def cost
-    cost_microcents.to_d / MICROCENTS_PER_DOLLAR
-  end
-
   def cost_microcents
     input_cost_microcents + output_cost_microcents
   end
@@ -27,7 +21,7 @@ class Conversation::Message::ResponseGenerator::Response
   def input_token_price_microcents
     return unless model_info
 
-    price_per_million_tokens_to_microcents(model_info.input_price_per_million)
+    price_per_million_tokens_in_microcents(model_info.input_price_per_million)
   end
 
   def output_cost_microcents
@@ -39,7 +33,7 @@ class Conversation::Message::ResponseGenerator::Response
   def output_token_price_microcents
     return unless model_info
 
-    price_per_million_tokens_to_microcents(model_info.output_price_per_million)
+    price_per_million_tokens_in_microcents(model_info.output_price_per_million)
   end
 
   def model_info
@@ -47,8 +41,8 @@ class Conversation::Message::ResponseGenerator::Response
   end
 
   private
-    def price_per_million_tokens_to_microcents(price)
+    def price_per_million_tokens_in_microcents(price)
       single_token_price = price.to_d / 1_000_000
-      single_token_price * MICROCENTS_PER_DOLLAR
+      Ai::Quota::Money.wrap(single_token_price).in_microcents
     end
 end

app/models/user.rb

@@ -1,6 +1,6 @@
 class User < ApplicationRecord
   include Accessor, Attachable, Assignee, Mentionable, Named, Role, Searcher,
-    SignalUser, Staff, Transferable, Conversational
+    SignalUser, Staff, Transferable, Conversational, AiQuota
   include Timelined # Depends on Accessor
 
   has_one_attached :avatar

app/models/user/ai_quota.rb

@@ -0,0 +1,22 @@
+module User::AiQuota
+  extend ActiveSupport::Concern
+
+  DEFAULT_QUOTA = Ai::Quota::Money.wrap("$100").in_microcents
+
+  included do
+    has_one :ai_quota, class_name: "Ai::Quota"
+  end
+
+  def spend_ai_quota(cost)
+    fetch_or_create_ai_quota.spend(cost)
+  end
+
+  def ensure_ai_quota_not_depleted
+    fetch_or_create_ai_quota.ensure_not_depleted
+  end
+
+  private
+    def fetch_or_create_ai_quota
+      ai_quota || create_ai_quota!(limit: DEFAULT_QUOTA)
+    end
+end

config/initializers/inflections.rb

@@ -13,4 +13,7 @@
 # These inflection rules are supported but not enabled by default:
 ActiveSupport::Inflector.inflections(:en) do |inflect|
   inflect.acronym "SQLite"
+
+  inflect.singular "quotas", "quota"
+  inflect.plural "quota", "quotas"
 end

db/migrate/20250819101032_create_ai_quotas.rb

@@ -0,0 +1,14 @@
+class CreateAiQuotas < ActiveRecord::Migration[8.1]
+  def change
+    create_table :ai_quotas do |t|
+      t.belongs_to :user, null: false, foreign_key: true
+      t.integer :limit, null: false
+      t.integer :used, null: false, default: 0
+      t.datetime :reset_at, null: false, index: true
+
+      t.timestamps
+    end
+
+    add_index :ai_quotas, :user_id, unique: true
+  end
+end

db/schema.rb

@@ -73,6 +73,17 @@ ActiveRecord::Schema[8.1].define(version: 2025_08_19_105245) do
     t.index ["blob_id", "variation_digest"], name: "index_active_storage_variant_records_uniqueness", unique: true
   end
 
+  create_table "ai_quotas", force: :cascade do |t|
+    t.datetime "created_at", null: false
+    t.integer "limit", null: false
+    t.datetime "reset_at", null: false
+    t.datetime "updated_at", null: false
+    t.integer "used", default: 0, null: false
+    t.integer "user_id", null: false
+    t.index ["reset_at"], name: "index_ai_quotas_on_reset_at"
+    t.index ["user_id"], name: "index_ai_quotas_on_user_id"
+  end
+
   create_table "assignees_filters", id: false, force: :cascade do |t|
     t.integer "assignee_id", null: false
     t.integer "filter_id", null: false
@@ -438,6 +449,7 @@ ActiveRecord::Schema[8.1].define(version: 2025_08_19_105245) do
 
   add_foreign_key "active_storage_attachments", "active_storage_blobs", column: "blob_id"
   add_foreign_key "active_storage_variant_records", "active_storage_blobs", column: "blob_id"
+  add_foreign_key "ai_quotas", "users"
   add_foreign_key "card_activity_spikes", "cards"
   add_foreign_key "card_goldnesses", "cards"
   add_foreign_key "cards", "workflow_stages", column: "stage_id"

db/schema_cache.yml

@@ -20,7 +20,7 @@ columns:
     default_function:
     collation:
     comment:
-  - &23 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+  - &24 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: collection_id
     cast_type: &3 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3Adapter::SQLite3Integer
@@ -90,7 +90,7 @@ columns:
     default_function:
     collation:
     comment:
-  - &25 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+  - &18 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: user_id
     cast_type: *3
@@ -258,7 +258,7 @@ columns:
     collation:
     comment:
   - *6
-  - &18 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+  - &19 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: key
     cast_type: *7
@@ -301,9 +301,44 @@ columns:
     default_function:
     collation:
     comment:
+  ai_quotas:
+  - *5
+  - *6
+  - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+    auto_increment:
+    name: limit
+    cast_type: *3
+    sql_type_metadata: *4
+    'null': false
+    default:
+    default_function:
+    collation:
+    comment:
+  - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+    auto_increment:
+    name: reset_at
+    cast_type: *1
+    sql_type_metadata: *2
+    'null': false
+    default:
+    default_function:
+    collation:
+    comment:
+  - *9
+  - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+    auto_increment:
+    name: used
+    cast_type: *3
+    sql_type_metadata: *4
+    'null': false
+    default: 0
+    default_function:
+    collation:
+    comment:
+  - *18
   ar_internal_metadata:
   - *5
-  - *18
+  - *19
   - *9
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
@@ -316,7 +351,7 @@ columns:
     collation:
     comment:
   assignees_filters:
-  - &20 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+  - &21 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: assignee_id
     cast_type: *3
@@ -326,7 +361,7 @@ columns:
     default_function:
     collation:
     comment:
-  - &19 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+  - &20 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: filter_id
     cast_type: *3
@@ -337,7 +372,7 @@ columns:
     collation:
     comment:
   assigners_filters:
-  - &21 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+  - &22 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: assigner_id
     cast_type: *3
@@ -347,11 +382,11 @@ columns:
     default_function:
     collation:
     comment:
-  - *19
-  assignments:
   - *20
+  assignments:
   - *21
-  - &22 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+  - *22
+  - &23 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: card_id
     cast_type: *3
@@ -365,7 +400,7 @@ columns:
   - *6
   - *9
   card_activity_spikes:
-  - *22
+  - *23
   - *5
   - *6
   - *9
@@ -394,14 +429,14 @@ columns:
     comment:
   - *9
   card_goldnesses:
-  - *22
+  - *23
   - *5
   - *6
   - *9
   cards:
-  - *23
+  - *24
   - *5
-  - &24 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
+  - &25 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: creator_id
     cast_type: *3
@@ -483,7 +518,7 @@ columns:
     default_function:
     collation:
     comment:
-  - *19
+  - *20
   closure_reasons:
   - *5
   - *6
@@ -499,7 +534,7 @@ columns:
     comment:
   - *9
   closures:
-  - *22
+  - *23
   - *5
   - *6
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
@@ -524,7 +559,7 @@ columns:
     collation:
     comment:
   collection_publications:
-  - *23
+  - *24
   - *5
   - *6
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
@@ -558,7 +593,7 @@ columns:
     collation:
     comment:
   - *5
-  - *24
+  - *25
   - *6
   - *10
   - *9
@@ -573,12 +608,12 @@ columns:
     collation:
     comment:
   collections_filters:
-  - *23
-  - *19
-  comments:
-  - *22
-  - *5
   - *24
+  - *20
+  comments:
+  - *23
+  - *5
+  - *25
   - *6
   - *9
   conversation_messages:
@@ -699,10 +734,10 @@ columns:
     collation:
     comment:
   - *9
-  - *25
+  - *18
   creators_filters:
-  - *24
-  - *19
+  - *25
+  - *20
   entropy_configurations:
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
@@ -760,7 +795,7 @@ columns:
     comment:
   - *5
   - *6
-  - *18
+  - *19
   - *9
   events:
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
@@ -773,9 +808,9 @@ columns:
     default_function:
     collation:
     comment:
-  - *23
-  - *5
   - *24
+  - *5
+  - *25
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: eventable_id
@@ -818,7 +853,7 @@ columns:
   - *9
   filters:
   - *5
-  - *24
+  - *25
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: fields
@@ -842,7 +877,7 @@ columns:
     comment:
   - *9
   filters_stages:
-  - *19
+  - *20
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: stage_id
@@ -854,7 +889,7 @@ columns:
     collation:
     comment:
   filters_tags:
-  - *19
+  - *20
   - &34 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: tag_id
@@ -935,13 +970,13 @@ columns:
   - *28
   - *29
   - *9
-  - *25
+  - *18
   pins:
-  - *22
+  - *23
   - *5
   - *6
   - *9
-  - *25
+  - *18
   push_subscriptions:
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
@@ -986,7 +1021,7 @@ columns:
     default_function:
     collation:
     comment:
-  - *25
+  - *18
   reactions:
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
@@ -1103,7 +1138,7 @@ columns:
     collation:
     comment:
   - *9
-  - *25
+  - *18
   search_results:
   - *5
   - *6
@@ -1123,9 +1158,9 @@ columns:
     comment:
   - *9
   - *30
-  - *25
+  - *18
   steps:
-  - *22
+  - *23
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: completed
@@ -1141,7 +1176,7 @@ columns:
   - *6
   - *9
   taggings:
-  - *22
+  - *23
   - *5
   - *6
   - *34
@@ -1207,11 +1242,11 @@ columns:
     comment:
   - *9
   watches:
-  - *22
+  - *23
   - *5
   - *6
   - *9
-  - *25
+  - *18
   - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column
     auto_increment:
     name: watching
@@ -1259,6 +1294,7 @@ primary_keys:
   active_storage_attachments: id
   active_storage_blobs: id
   active_storage_variant_records: id
+  ai_quotas: id
   ar_internal_metadata: key
   assignees_filters:
   assigners_filters:
@@ -1307,6 +1343,7 @@ data_sources:
   active_storage_attachments: true
   active_storage_blobs: true
   active_storage_variant_records: true
+  ai_quotas: true
   ar_internal_metadata: true
   assignees_filters: true
   assigners_filters: true
@@ -1538,6 +1575,39 @@ indexes:
     nulls_not_distinct:
     comment:
     valid: true
+  ai_quotas:
+  - !ruby/object:ActiveRecord::ConnectionAdapters::IndexDefinition
+    table: ai_quotas
+    name: index_ai_quotas_on_reset_at
+    unique: false
+    columns:
+    - reset_at
+    lengths: {}
+    orders: {}
+    opclasses: {}
+    where:
+    type:
+    using:
+    include:
+    nulls_not_distinct:
+    comment:
+    valid: true
+  - !ruby/object:ActiveRecord::ConnectionAdapters::IndexDefinition
+    table: ai_quotas
+    name: index_ai_quotas_on_user_id
+    unique: false
+    columns:
+    - user_id
+    lengths: {}
+    orders: {}
+    opclasses: {}
+    where:
+    type:
+    using:
+    include:
+    nulls_not_distinct:
+    comment:
+    valid: true
   ar_internal_metadata: []
   assignees_filters:
   - !ruby/object:ActiveRecord::ConnectionAdapters::IndexDefinition

test/fixtures/ai/quotas.yml

@@ -0,0 +1,11 @@
+kevin:
+  user: kevin
+  limit: 100_00_000_000
+  used: 1053
+  reset_at: <%= 7.days.from_now %>
+
+david:
+  user: david
+  limit: 20_00_000_000
+  used: 0
+  reset_at: <%= 3.days.from_now %>

test/models/ai/quota/money_test.rb

@@ -0,0 +1,40 @@
+require "test_helper"
+
+class Ai::Quota::MoneyTest < ActiveSupport::TestCase
+  test "wrapping" do
+    money = Ai::Quota::Money.wrap("$5.42")
+    assert_equal 5_42_000_000, money.in_microcents, "Strings with numbers are treated as dollars"
+
+    assert_raises TypeError do
+      Ai::Quota::Money.wrap("foobar")
+    end
+
+    money = Ai::Quota::Money.wrap(5.42)
+    assert_equal 5_42_000_000, money.in_microcents, "Decimals are treated as dollars"
+
+    money = Ai::Quota::Money.wrap(5)
+    assert_equal 5, money.in_microcents, "Integers are treated as microcents"
+
+    money1 = Ai::Quota::Money.wrap("$5")
+    money2 = Ai::Quota::Money.wrap(money1)
+    assert_equal money1, money2, "Money can wrap itself"
+
+    assert_raises ArgumentError do
+      Ai::Quota::Money.wrap(nil)
+    end
+  end
+
+  test "conversions" do
+    money = Ai::Quota::Money.wrap("$0")
+    assert_equal 0.0, money.in_dollars
+    assert_equal 0, money.in_microcents
+
+    money = Ai::Quota::Money.wrap("$1")
+    assert_equal 1, money.in_dollars
+    assert_equal 1_00_000_000, money.in_microcents
+
+    money = Ai::Quota::Money.wrap("$5.42")
+    assert_equal 5.42, money.in_dollars
+    assert_equal 5_42_000_000, money.in_microcents
+  end
+end

test/models/ai/quota_test.rb

@@ -0,0 +1,63 @@
+require "test_helper"
+
+class Ai::QuotaTest < ActiveSupport::TestCase
+  setup do
+    @quota = Ai::Quota.new(user: users(:jz), limit: Ai::Quota::Money.wrap("$100").in_microcents)
+  end
+
+  test "create" do
+    assert @quota.save
+
+    assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute
+    assert_equal 0, @quota.used
+    assert_equal Ai::Quota::Money.wrap("$100").in_microcents, @quota.limit
+  end
+
+  test "increment usage" do
+    @quota.save
+
+    @quota.spend("$100")
+    assert_equal Ai::Quota::Money.wrap("$100").in_microcents, @quota.used
+    @quota.spend("$500")
+    assert_equal Ai::Quota::Money.wrap("$600").in_microcents, @quota.used
+    @quota.spend("$1000")
+    assert_equal Ai::Quota::Money.wrap("$1600").in_microcents, @quota.used
+    @quota.spend("$5000")
+    assert_equal Ai::Quota::Money.wrap("$6600").in_microcents, @quota.used
+    @quota.spend("$10000")
+    assert_equal Ai::Quota::Money.wrap("$16600").in_microcents, @quota.used
+
+    @quota.used = 0
+
+    @quota.spend("$10")
+    assert_equal Ai::Quota::Money.wrap("$10").in_microcents, @quota.used
+    assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute
+
+    travel 2.days
+
+    @quota.spend("$5")
+    assert_equal Ai::Quota::Money.wrap("$15").in_microcents, @quota.used
+    assert_in_delta 5.days.from_now, @quota.reset_at, 1.minute
+
+    travel 8.days
+
+    @quota.spend("$5")
+    assert_equal Ai::Quota::Money.wrap("$5").in_microcents, @quota.used
+    assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute
+  end
+
+  test "limit checks" do
+    @quota.save
+
+    @quota.used = 0
+    @quota.ensure_not_depleted
+
+    @quota.used = Ai::Quota::Money.wrap("$300").in_microcents
+    assert_raises Ai::Quota::UsageExceedsQuotaError do
+      @quota.ensure_not_depleted
+    end
+
+    travel 10.days
+    @quota.ensure_not_depleted
+  end
+end

test/models/conversation/response_generator/response_test.rb

@@ -13,23 +13,20 @@ class Conversation::Message::ResponseGenerator::ResponseTest < ActiveSupport::Te
     # and 60 USD per million output tokens
     # That's 0.00003 cents per input token and 0.00006 cents
     # per output token
-    # Which is 3 micro-cents per input token and 6 micro-cents
+    # Which is 3000 micro-cents per input token and 6000 micro-cents
     # per output token
-    assert_equal "3.0".to_d, response.input_token_price_microcents
-    assert_equal "6.0".to_d, response.output_token_price_microcents
+    assert_equal "3000.0".to_d, response.input_token_price_microcents
+    assert_equal "6000.0".to_d, response.output_token_price_microcents
 
     # We've got 198 input tokens, so that's
-    # 193 * 3 = 594
-    assert_equal 594, response.input_cost_microcents
+    # 198 * 3000 = 594000
+    assert_equal 594000, response.input_cost_microcents
 
     # We've got 2 output tokens, so that's
-    # 2 * 6 = 12
-    assert_equal 12, response.output_cost_microcents
+    # 2 * 6000 = 12
+    assert_equal 12000, response.output_cost_microcents
 
-    # So the total is 594 + 12 micro-cents
-    assert_equal 606, response.cost_microcents
-
-    # If we convert that to a decimal value we get 0.00606 cents
-    assert_equal "0.00606".to_d, response.cost
+    # So the total is 594000 + 12000 micro-cents
+    assert_equal 606000, response.cost_microcents
   end
 end

test/models/conversation_test.rb

@@ -61,21 +61,25 @@ class ConversationTest < ActiveSupport::TestCase
     assert conversation.ready?, "The conversation should switch back to ready after a response is made"
   end
 
-  test "clearing conversation messages" do
+  test "cost limits" do
     conversation = conversations(:kevin)
 
-    assert conversation.messages.any?, "The conversation should have messages before clearing"
+    conversation.ask("Where does the planning office keep demolition notices?")
+    conversation.respond(
+      "In a locked filing cabinet in a disused lavatory",
+      cost_microcents: Ai::Quota::Money.wrap("$3").in_microcents
+    )
 
-    original_updated_at = conversation.updated_at
-    conversation.clear
-    assert conversation.updated_at > original_updated_at, "The conversation's updated_at timestamp should change after clearing messages"
+    conversation.ask("What's the meaning of life?")
+    conversation.respond("42", cost_microcents: Ai::Quota::Money.wrap("$120").in_microcents)
 
-    assert conversation.messages.empty?, "All messages should be deleted when clearing the conversation"
-  end
+    assert_raises Ai::Quota::UsageExceedsQuotaError do
+      conversation.ask("Should you leave a house without a towel?")
+    end
 
-  test "cost calculation" do
-    conversation = conversations(:kevin)
+    travel 1.month
 
-    assert_equal "0.01053".to_d, conversation.cost
+    conversation.ask("Should you leave a house without a towel?")
+    conversation.respond("Never", cost_microcents: Ai::Quota::Money.wrap("$0.01").in_microcents)
   end
 end