diff --git a/enterprise/lib/enterprise/captain/base_task_service.rb b/enterprise/lib/enterprise/captain/base_task_service.rb index d2b75f7aa..b080ff807 100644 --- a/enterprise/lib/enterprise/captain/base_task_service.rb +++ b/enterprise/lib/enterprise/captain/base_task_service.rb @@ -28,7 +28,8 @@ module Enterprise::Captain::BaseTaskService end def increment_usage - Rails.logger.info("[CAPTAIN][#{self.class.name}] Incrementing response usage for account #{account.id}") - account.increment_response_usage + credits = Llm::Models.credit_multiplier_for(configured_model) + Rails.logger.info("[CAPTAIN][#{self.class.name}] Incrementing response usage for account #{account.id} by #{credits} credits") + account.increment_response_usage(credits: credits) end end diff --git a/spec/enterprise/lib/captain/base_task_service_spec.rb b/spec/enterprise/lib/captain/base_task_service_spec.rb index c696674d5..e0dd5fbab 100644 --- a/spec/enterprise/lib/captain/base_task_service_spec.rb +++ b/spec/enterprise/lib/captain/base_task_service_spec.rb @@ -15,6 +15,10 @@ RSpec.describe Captain::BaseTaskService, type: :model do def event_name 'test_event' end + + def feature_key + 'editor' + end end # Manually prepend enterprise module to test class klass.prepend(Enterprise::Captain::BaseTaskService) @@ -53,8 +57,15 @@ RSpec.describe Captain::BaseTaskService, type: :model do end end - it 'increments response usage on successful execution' do - expect(account).to receive(:increment_response_usage) + it 'increments response usage with credit multiplier on successful execution' do + expect(account).to receive(:captain_editor_model).and_return('gpt-4.1-mini') + expect(account).to receive(:increment_response_usage).with(credits: 1) + service.perform + end + + it 'uses correct credit multiplier for higher cost models' do + expect(account).to receive(:captain_editor_model).and_return('gpt-5.1') + expect(account).to receive(:increment_response_usage).with(credits: 2) service.perform end @@ -103,13 +114,22 @@ RSpec.describe Captain::BaseTaskService, type: :model do end end - it 'actually increments the usage counter in custom_attributes' do + it 'actually increments the usage counter in custom_attributes by credit multiplier' do + allow(account).to receive(:captain_editor_model).and_return('gpt-4.1-mini') expect do service.perform account.reload end.to change { account.custom_attributes['captain_responses_usage'].to_i }.by(1) end + it 'increments by correct credits for expensive models' do + allow(account).to receive(:captain_editor_model).and_return('gpt-5.1') + expect do + service.perform + account.reload + end.to change { account.custom_attributes['captain_responses_usage'].to_i }.by(2) + end + context 'when captain is disabled' do before do allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(false)