From 45f4b423ae7269dab01cbea6a31cab6d330058fe Mon Sep 17 00:00:00 2001 From: Aakash Bakhle <48802744+aakashb95@users.noreply.github.com> Date: Mon, 8 Jun 2026 13:37:21 +0530 Subject: [PATCH] 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 --- .../captain/llm/translate_query_service.rb | 4 ++ .../conversation_completion_service.rb | 4 ++ lib/captain/base_task_service.rb | 21 ++++--- lib/captain/csat_utility_analysis_service.rb | 4 ++ lib/captain/follow_up_service.rb | 4 ++ lib/captain/label_suggestion_service.rb | 4 ++ lib/captain/reply_suggestion_service.rb | 4 ++ lib/captain/rewrite_service.rb | 4 ++ lib/captain/summary_service.rb | 4 ++ .../lib/captain/base_task_service_spec.rb | 63 +++++++++++++++++++ spec/lib/captain/base_task_service_spec.rb | 56 ++++++++++++++++- .../csat_utility_analysis_service_spec.rb | 23 +++++++ 12 files changed, 185 insertions(+), 10 deletions(-) diff --git a/enterprise/app/services/captain/llm/translate_query_service.rb b/enterprise/app/services/captain/llm/translate_query_service.rb index 93f68b05b..3e05244d3 100644 --- a/enterprise/app/services/captain/llm/translate_query_service.rb +++ b/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) diff --git a/enterprise/lib/captain/conversation_completion_service.rb b/enterprise/lib/captain/conversation_completion_service.rb index aa40e8000..c45559165 100644 --- a/enterprise/lib/captain/conversation_completion_service.rb +++ b/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 diff --git a/lib/captain/base_task_service.rb b/lib/captain/base_task_service.rb index a043d38e2..d382204a5 100644 --- a/lib/captain/base_task_service.rb +++ b/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 diff --git a/lib/captain/csat_utility_analysis_service.rb b/lib/captain/csat_utility_analysis_service.rb index e04a98a7f..7aab18e6c 100644 --- a/lib/captain/csat_utility_analysis_service.rb +++ b/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 diff --git a/lib/captain/follow_up_service.rb b/lib/captain/follow_up_service.rb index f02ba9408..c4c1225be 100644 --- a/lib/captain/follow_up_service.rb +++ b/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 diff --git a/lib/captain/label_suggestion_service.rb b/lib/captain/label_suggestion_service.rb index 02f8bd89a..a0e030963 100644 --- a/lib/captain/label_suggestion_service.rb +++ b/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 diff --git a/lib/captain/reply_suggestion_service.rb b/lib/captain/reply_suggestion_service.rb index 2daf0615c..039bdcf26 100644 --- a/lib/captain/reply_suggestion_service.rb +++ b/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') diff --git a/lib/captain/rewrite_service.rb b/lib/captain/rewrite_service.rb index 3a217d3c6..6f880e775 100644 --- a/lib/captain/rewrite_service.rb +++ b/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 diff --git a/lib/captain/summary_service.rb b/lib/captain/summary_service.rb index 030c0e510..f06aa42ca 100644 --- a/lib/captain/summary_service.rb +++ b/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 diff --git a/spec/enterprise/lib/captain/base_task_service_spec.rb b/spec/enterprise/lib/captain/base_task_service_spec.rb index b3dc473eb..fb970f726 100644 --- a/spec/enterprise/lib/captain/base_task_service_spec.rb +++ b/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) diff --git a/spec/lib/captain/base_task_service_spec.rb b/spec/lib/captain/base_task_service_spec.rb index 5112e47ed..34c889967 100644 --- a/spec/lib/captain/base_task_service_spec.rb +++ b/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') diff --git a/spec/lib/captain/csat_utility_analysis_service_spec.rb b/spec/lib/captain/csat_utility_analysis_service_spec.rb index e4e980e01..34e0c9ece 100644 --- a/spec/lib/captain/csat_utility_analysis_service_spec.rb +++ b/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