fix: captain usage for BYOK OpenAI tasks (#14587)

# Pull Request Template

## Description

Fixes:
https://linear.app/chatwoot/issue/CW-7167/label-suggestions-bad-ux

## Type of change

Please delete options that are not relevant.

- [x] Bug fix (non-breaking change which fixes an issue)

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide
instructions so we can reproduce. Please also list any relevant details
for your test configuration.


## Checklist:

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my code
- [ ] I have commented on my code, particularly in hard-to-understand
areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged and published in downstream
modules

enterprise/app/services/captain/llm/translate_query_service.rb

@@ -32,6 +32,10 @@ class Captain::Llm::TranslateQueryService < Captain::BaseTaskService
     @llm_credential ||= system_llm_credential
   end
 
+  def counts_toward_usage?
+    false
+  end
+
   def query_in_target_language?(query)
     detector = CLD3::NNetLanguageIdentifier.new(0, 1000)
     result = detector.find_language(query)

enterprise/lib/captain/conversation_completion_service.rb

@@ -62,6 +62,10 @@ class Captain::ConversationCompletionService < Captain::BaseTaskService
     @llm_credential ||= system_llm_credential
   end
 
+  def counts_toward_usage?
+    false
+  end
+
   def event_name
     'captain.conversation_completion'
   end

lib/captain/base_task_service.rb

@@ -150,13 +150,12 @@ class Captain::BaseTaskService
   end
 
   # Extension point consulted by the Enterprise quota wrapper. Subclasses
-  # whose calls run on the operator's key (e.g. internal/onboarding tasks)
-  # should override this to return false. When false, the wrapper neither
-  # blocks the call on an exhausted captain_responses quota nor decrements
-  # it on success — the call participates in the quota system in neither
-  # direction.
+  # whose calls should not consume captain_responses should override this to
+  # return false. When false, the wrapper neither blocks the call on an
+  # exhausted captain_responses quota nor decrements it on success — the call
+  # participates in the quota system in neither direction.
   def counts_toward_usage?
-    true
+    llm_credential&.dig(:source) != :hook
   end
 
   def api_key_configured?
@@ -168,7 +167,15 @@ class Captain::BaseTaskService
   end
 
   def llm_credential
-    @llm_credential ||= hook_llm_credential || system_llm_credential
+    @llm_credential ||= if use_account_openai_hook?
+                          hook_llm_credential || system_llm_credential
+                        else
+                          system_llm_credential
+                        end
+  end
+
+  def use_account_openai_hook?
+    false
   end
 
   def hook_llm_credential

lib/captain/csat_utility_analysis_service.rb

@@ -63,4 +63,8 @@ class Captain::CsatUtilityAnalysisService < Captain::BaseTaskService
   def event_name
     'csat_utility_analysis'
   end
+
+  def use_account_openai_hook?
+    true
+  end
 end

lib/captain/follow_up_service.rb

@@ -103,4 +103,8 @@ class Captain::FollowUpService < Captain::BaseTaskService
   def event_name
     'follow_up'
   end
+
+  def use_account_openai_hook?
+    true
+  end
 end

lib/captain/label_suggestion_service.rb

@@ -87,6 +87,10 @@ class Captain::LabelSuggestionService < Captain::BaseTaskService
     'label_suggestion'
   end
 
+  def use_account_openai_hook?
+    true
+  end
+
   def build_follow_up_context?
     false
   end

lib/captain/reply_suggestion_service.rb

@@ -37,6 +37,10 @@ class Captain::ReplySuggestionService < Captain::BaseTaskService
   def event_name
     'reply_suggestion'
   end
+
+  def use_account_openai_hook?
+    true
+  end
 end
 
 Captain::ReplySuggestionService.prepend_mod_with('Captain::ReplySuggestionService')

lib/captain/rewrite_service.rb

@@ -56,4 +56,8 @@ class Captain::RewriteService < Captain::BaseTaskService
   def event_name
     operation
   end
+
+  def use_account_openai_hook?
+    true
+  end
 end

lib/captain/summary_service.rb

@@ -24,4 +24,8 @@ class Captain::SummaryService < Captain::BaseTaskService
   def event_name
     'summarize'
   end
+
+  def use_account_openai_hook?
+    true
+  end
 end

spec/enterprise/lib/captain/base_task_service_spec.rb

@@ -32,6 +32,7 @@ RSpec.describe Captain::BaseTaskService, type: :model do
     before do
       allow(account).to receive(:feature_enabled?).and_call_original
       allow(account).to receive(:feature_enabled?).with('captain_tasks').and_return(true)
+      allow(Integrations::Openai::KeyValidator).to receive(:valid?).and_return(true)
     end
 
     context 'when usage limit is exceeded' do
@@ -111,6 +112,68 @@ RSpec.describe Captain::BaseTaskService, type: :model do
       end.to change { account.custom_attributes['captain_responses_usage'].to_i }.by(1)
     end
 
+    context 'when account has its own OpenAI hook key' do
+      before do
+        create(:integrations_hook, :openai, account: account, settings: { 'api_key' => 'customer-own-key' })
+      end
+
+      it 'still increments usage for services that do not opt into BYOK' do
+        expect(account).to receive(:increment_response_usage)
+        service.perform
+      end
+
+      context 'when the captain_responses quota is exhausted on Cloud' do
+        before do
+          allow(ChatwootApp).to receive(:chatwoot_cloud?).and_return(true)
+          allow(account).to receive(:usage_limits).and_return({
+                                                                captain: { responses: { current_available: 0 } }
+                                                              })
+        end
+
+        it 'returns usage limit exceeded error for services that do not opt into BYOK' do
+          result = service.perform
+          expect(result[:error]).to eq(I18n.t('captain.copilot_limit'))
+          expect(result[:error_code]).to eq(429)
+        end
+      end
+    end
+
+    context 'when subclass opts into account OpenAI hook usage' do
+      let(:test_service_class) do
+        result = perform_result
+        klass = Class.new(described_class) do
+          define_method(:perform) { result }
+          define_method(:event_name) { 'test_event' }
+          define_method(:use_account_openai_hook?) { true }
+        end
+        klass.prepend(Enterprise::Captain::BaseTaskService)
+        klass
+      end
+
+      before do
+        create(:integrations_hook, :openai, account: account, settings: { 'api_key' => 'customer-own-key' })
+      end
+
+      it 'does not increment usage on a successful result' do
+        expect(account).not_to receive(:increment_response_usage)
+        service.perform
+      end
+
+      context 'when the captain_responses quota is exhausted on Cloud' do
+        before do
+          allow(ChatwootApp).to receive(:chatwoot_cloud?).and_return(true)
+          allow(account).to receive(:usage_limits).and_return({
+                                                                captain: { responses: { current_available: 0 } }
+                                                              })
+        end
+
+        it 'bypasses the 429 gate and returns the underlying result' do
+          result = service.perform
+          expect(result).to eq(perform_result)
+        end
+      end
+    end
+
     context 'when captain is disabled' do
       before do
         allow(account).to receive(:feature_enabled?).with('captain_tasks').and_return(false)

spec/lib/captain/base_task_service_spec.rb

@@ -260,11 +260,12 @@ RSpec.describe Captain::BaseTaskService do
       expect(result[:request_messages]).to eq(messages)
     end
 
-    it 'does not track exceptions for account hook failures' do
+    it 'tracks exceptions against the system key when an account hook exists' do
       create(:integrations_hook, :openai, account: account, settings: { 'api_key' => 'hook-key' })
 
-      expect(Llm::Config).to receive(:with_api_key).with('hook-key', api_base: anything).and_raise(error)
-      expect(ChatwootExceptionTracker).not_to receive(:new)
+      expect(Llm::Config).to receive(:with_api_key).with('test-key', api_base: anything).and_raise(error)
+      expect(ChatwootExceptionTracker).to receive(:new).with(error, account: account).and_return(exception_tracker)
+      expect(exception_tracker).to receive(:capture_exception)
 
       result = service.send(:make_api_call, model: model, messages: messages)
 
@@ -279,11 +280,60 @@ RSpec.describe Captain::BaseTaskService do
 
       before { hook }
 
+      it 'uses system api key by default' do
+        expect(service.send(:api_key)).to eq('test-key')
+      end
+    end
+
+    context 'when subclass opts into account OpenAI hook usage' do
+      let(:test_service_class) do
+        Class.new(described_class) do
+          def event_name
+            'test_event'
+          end
+
+          def use_account_openai_hook?
+            true
+          end
+        end
+      end
+
+      before do
+        create(:integrations_hook, account: account, app_id: 'openai', status: 'enabled', settings: { 'api_key' => 'hook-key' })
+      end
+
       it 'uses api key from hook' do
         expect(service.send(:api_key)).to eq('hook-key')
       end
     end
 
+    it 'uses account OpenAI hook for editor task services' do
+      create(:integrations_hook, account: account, app_id: 'openai', status: 'enabled', settings: { 'api_key' => 'hook-key' })
+      user = create(:user, account: account)
+      follow_up_context = {
+        'event_name' => 'professional',
+        'original_context' => 'Original text',
+        'last_response' => 'Last response'
+      }
+
+      editor_services = [
+        Captain::RewriteService.new(account: account, content: 'Text', operation: 'improve', conversation_display_id: conversation.display_id),
+        Captain::SummaryService.new(account: account, conversation_display_id: conversation.display_id),
+        Captain::ReplySuggestionService.new(account: account, conversation_display_id: conversation.display_id, user: user),
+        Captain::LabelSuggestionService.new(account: account, conversation_display_id: conversation.display_id),
+        Captain::FollowUpService.new(
+          account: account,
+          follow_up_context: follow_up_context,
+          user_message: 'Make it shorter',
+          conversation_display_id: conversation.display_id
+        )
+      ]
+
+      editor_services.each do |editor_service|
+        expect(editor_service.send(:api_key)).to eq('hook-key')
+      end
+    end
+
     context 'when openai hook is not configured' do
       it 'uses system api key' do
         expect(service.send(:api_key)).to eq('test-key')

spec/lib/captain/csat_utility_analysis_service_spec.rb

@@ -4,6 +4,11 @@ RSpec.describe Captain::CsatUtilityAnalysisService do
   let(:account) { create(:account) }
   let(:service) { described_class.new(account: account, message: 'Test message', language: 'en', baseline: {}) }
 
+  before do
+    create(:installation_config, name: 'CAPTAIN_OPEN_AI_API_KEY', value: 'test-key')
+    allow(Integrations::Openai::KeyValidator).to receive(:valid?).and_return(true)
+  end
+
   describe '#perform' do
     before do
       allow(account).to receive(:feature_enabled?).and_call_original
@@ -21,4 +26,22 @@ RSpec.describe Captain::CsatUtilityAnalysisService do
       expect(result[:message]).to eq('{"classification":"LIKELY_UTILITY","optimized_message":"Utility-safe message"}')
     end
   end
+
+  describe '#api_key' do
+    context 'when account has an OpenAI hook key' do
+      before do
+        create(:integrations_hook, :openai, account: account, settings: { 'api_key' => 'customer-own-key' })
+      end
+
+      it 'uses the account hook key' do
+        expect(service.send(:api_key)).to eq('customer-own-key')
+      end
+    end
+
+    context 'when account does not have an OpenAI hook key' do
+      it 'uses the system key' do
+        expect(service.send(:api_key)).to eq('test-key')
+      end
+    end
+  end
 end