From c179393a0be93de0ef1443a8f7d715e333cfe10e Mon Sep 17 00:00:00 2001
From: Shivam Mishra <scm.mymail@gmail.com>
Date: Wed, 14 Jan 2026 11:31:12 +0530
Subject: [PATCH] feat: add model based usage helper

---
 .../account/plan_usage_and_limits.rb          |  6 +++++
 .../account/plan_usage_and_limits_spec.rb     | 23 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb b/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb
index 92d8410c4..b624fa10f 100644
--- a/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb
+++ b/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb
@@ -21,6 +21,12 @@ module Enterprise::Account::PlanUsageAndLimits
     save
   end
 
+  def increment_response_usage_for_model(model)
+    credits = using_openai_hook_key? ? 1 : Llm::Models.credit_multiplier_for(model)
+    Rails.logger.info("[CAPTAIN] Incrementing response usage for account #{id} by #{credits} credits (model: #{model})")
+    increment_response_usage(credits: credits)
+  end
+
   def reset_response_usage
     custom_attributes[CAPTAIN_RESPONSES_USAGE] = 0
     save
diff --git a/spec/enterprise/models/enterprise/account/plan_usage_and_limits_spec.rb b/spec/enterprise/models/enterprise/account/plan_usage_and_limits_spec.rb
index 31f8820f6..89f62fa25 100644
--- a/spec/enterprise/models/enterprise/account/plan_usage_and_limits_spec.rb
+++ b/spec/enterprise/models/enterprise/account/plan_usage_and_limits_spec.rb
@@ -28,4 +28,27 @@ RSpec.describe Enterprise::Account::PlanUsageAndLimits do
       expect(account.custom_attributes['captain_responses_usage']).to eq(2)
     end
   end
+
+  describe '#increment_response_usage_for_model' do
+    it 'increments by model credit multiplier when using system key' do
+      # gpt-5.1 has credit_multiplier of 2
+      account.increment_response_usage_for_model('gpt-5.1')
+      account.reload
+      expect(account.custom_attributes['captain_responses_usage']).to eq(2)
+    end
+
+    it 'increments by 1 when using hook key' do
+      create(:integrations_hook, account: account, app_id: 'openai', status: 'enabled', settings: { 'api_key' => 'user-key' })
+      # gpt-5.1 has credit_multiplier of 2, but hook key overrides to 1
+      account.increment_response_usage_for_model('gpt-5.1')
+      account.reload
+      expect(account.custom_attributes['captain_responses_usage']).to eq(1)
+    end
+
+    it 'defaults to 1 credit for unknown models' do
+      account.increment_response_usage_for_model('unknown-model')
+      account.reload
+      expect(account.custom_attributes['captain_responses_usage']).to eq(1)
+    end
+  end
 end