From 61425b6a3b76809bb9987bbaed7b8e60099c7499 Mon Sep 17 00:00:00 2001 From: Shivam Mishra Date: Tue, 13 Jan 2026 18:09:34 +0530 Subject: [PATCH] feat: wire up credit usage (#13260) --- app/javascript/dashboard/composables/useAI.js | 35 ++-- config/locales/en.yml | 2 + .../enterprise/captain/base_task_service.rb | 34 ++++ lib/captain/base_task_service.rb | 11 ++ lib/llm/config.rb | 3 +- .../lib/captain/base_task_service_spec.rb | 168 ++++++++++++++++++ spec/lib/captain/base_task_service_spec.rb | 4 + spec/lib/captain/follow_up_service_spec.rb | 7 + .../captain/label_suggestion_service_spec.rb | 4 + .../captain/reply_suggestion_service_spec.rb | 4 + spec/lib/captain/rewrite_service_spec.rb | 4 + spec/lib/captain/summary_service_spec.rb | 4 + 12 files changed, 261 insertions(+), 19 deletions(-) create mode 100644 enterprise/lib/enterprise/captain/base_task_service.rb create mode 100644 spec/enterprise/lib/captain/base_task_service_spec.rb diff --git a/app/javascript/dashboard/composables/useAI.js b/app/javascript/dashboard/composables/useAI.js index 22336689a..8d5b3a73f 100644 --- a/app/javascript/dashboard/composables/useAI.js +++ b/app/javascript/dashboard/composables/useAI.js @@ -109,6 +109,21 @@ export function useAI() { } }; + /** + * Handles API errors and displays appropriate error messages. + * Silently returns for aborted requests. + * @param {Error} error - The error object from the API call. + */ + const handleAPIError = error => { + if (error.name === 'AbortError' || error.name === 'CanceledError') { + return; + } + const errorMessage = + error.response?.data?.error || + t('INTEGRATION_SETTINGS.OPEN_AI.GENERATE_ERROR'); + useAlert(errorMessage); + }; + /** * Records analytics for AI-related events. * @param {string} type - The type of event. @@ -171,15 +186,7 @@ export function useAI() { } = result; return { message: generatedMessage, followUpContext }; } catch (error) { - // Don't show error for aborted requests - if (error.name === 'AbortError' || error.name === 'CanceledError') { - return { message: '' }; - } - const errorData = error.response?.data?.error; - const errorMessage = - errorData?.error?.message || - t('INTEGRATION_SETTINGS.OPEN_AI.GENERATE_ERROR'); - useAlert(errorMessage); + handleAPIError(error); return { message: '' }; } }; @@ -203,15 +210,7 @@ export function useAI() { } = result; return { message: generatedMessage, followUpContext: updatedContext }; } catch (error) { - // Don't show error for aborted requests - if (error.name === 'AbortError' || error.name === 'CanceledError') { - return { message: '', followUpContext }; - } - const errorData = error.response?.data?.error; - const errorMessage = - errorData?.error?.message || - t('INTEGRATION_SETTINGS.OPEN_AI.GENERATE_ERROR'); - useAlert(errorMessage); + handleAPIError(error); return { message: '', followUpContext }; } }; diff --git a/config/locales/en.yml b/config/locales/en.yml index d9ca2f1be..f5cc97a0a 100644 --- a/config/locales/en.yml +++ b/config/locales/en.yml @@ -336,6 +336,8 @@ en: copilot_message_required: Message is required copilot_error: 'Please connect an assistant to this inbox to use Copilot' copilot_limit: 'You are out of Copilot credits. You can buy more credits from the billing section.' + upgrade: 'Upgrade your plan to enable Captain AI' + disabled: 'Captain AI is disabled for this account.' copilot: using_tool: 'Using tool %{function_name}' completed_tool_call: 'Completed %{function_name} tool call' diff --git a/enterprise/lib/enterprise/captain/base_task_service.rb b/enterprise/lib/enterprise/captain/base_task_service.rb new file mode 100644 index 000000000..d2b75f7aa --- /dev/null +++ b/enterprise/lib/enterprise/captain/base_task_service.rb @@ -0,0 +1,34 @@ +module Enterprise::Captain::BaseTaskService + def perform + return { error: I18n.t('captain.copilot_limit'), error_code: 429 } unless responses_available? + + unless captain_enabled? + return { error: I18n.t('captain.upgrade') } if ChatwootApp.chatwoot_cloud? + + return { error: I18n.t('captain.disabled') } + end + + result = super + increment_usage if successful_result?(result) + result + end + + private + + def captain_enabled? + account.feature_enabled?('captain_integration') + end + + def responses_available? + account.usage_limits[:captain][:responses][:current_available].positive? + end + + def successful_result?(result) + result.is_a?(Hash) && result[:message].present? && !result[:error] + end + + def increment_usage + Rails.logger.info("[CAPTAIN][#{self.class.name}] Incrementing response usage for account #{account.id}") + account.increment_response_usage + end +end diff --git a/lib/captain/base_task_service.rb b/lib/captain/base_task_service.rb index c60edf686..966057886 100644 --- a/lib/captain/base_task_service.rb +++ b/lib/captain/base_task_service.rb @@ -8,6 +8,15 @@ class Captain::BaseTaskService TOKEN_LIMIT = 400_000 GPT_MODEL = Llm::Config::DEFAULT_MODEL + # Prepend enterprise module to subclasses when they're defined. + # This ensures the enterprise perform wrapper is applied even when + # subclasses define their own perform method, since prepend puts + # the module before the class in the ancestor chain. + def self.inherited(subclass) + super + subclass.prepend_mod_with('Captain::BaseTaskService') + end + pattr_initialize [:account!, { conversation_display_id: nil }] private @@ -147,3 +156,5 @@ class Captain::BaseTaskService user_msg ? user_msg[:content] : nil end end + +Captain::BaseTaskService.prepend_mod_with('Captain::BaseTaskService') diff --git a/lib/llm/config.rb b/lib/llm/config.rb index f983f5daf..48de51022 100644 --- a/lib/llm/config.rb +++ b/lib/llm/config.rb @@ -1,7 +1,8 @@ require 'ruby_llm' module Llm::Config - DEFAULT_MODEL = 'gpt-5-mini'.freeze + DEFAULT_MODEL = 'gpt-4.1-mini'.freeze + class << self def initialized? @initialized ||= false diff --git a/spec/enterprise/lib/captain/base_task_service_spec.rb b/spec/enterprise/lib/captain/base_task_service_spec.rb new file mode 100644 index 000000000..c696674d5 --- /dev/null +++ b/spec/enterprise/lib/captain/base_task_service_spec.rb @@ -0,0 +1,168 @@ +require 'rails_helper' + +RSpec.describe Captain::BaseTaskService, type: :model do + let(:account) { create(:account) } + let(:inbox) { create(:inbox, account: account) } + let(:conversation) { create(:conversation, account: account, inbox: inbox) } + let(:perform_result) { { message: 'Test response' } } + + # Create a concrete test service class with enterprise module prepended + let(:test_service_class) do + result = perform_result + klass = Class.new(described_class) do + define_method(:perform) { result } + + def event_name + 'test_event' + end + end + # Manually prepend enterprise module to test class + klass.prepend(Enterprise::Captain::BaseTaskService) + klass + end + + let(:service) { test_service_class.new(account: account, conversation_display_id: conversation.display_id) } + + before do + create(:installation_config, name: 'CAPTAIN_OPEN_AI_API_KEY', value: 'test-key') + end + + describe '#perform with enterprise usage tracking' do + # Ensure captain is enabled by default for tests unless explicitly testing disabled state + before do + allow(account).to receive(:feature_enabled?).and_call_original + allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(true) + end + + context 'when usage limit is exceeded' do + before do + allow(account).to receive(:usage_limits).and_return({ + captain: { responses: { current_available: 0 } } + }) + end + + it 'returns usage limit exceeded error' do + result = service.perform + expect(result[:error]).to eq(I18n.t('captain.copilot_limit')) + expect(result[:error_code]).to eq(429) + end + + it 'does not increment usage' do + expect(account).not_to receive(:increment_response_usage) + service.perform + end + end + + it 'increments response usage on successful execution' do + expect(account).to receive(:increment_response_usage) + service.perform + end + + context 'when result has an error' do + let(:perform_result) { { error: 'API Error' } } + + it 'does not increment usage' do + expect(account).not_to receive(:increment_response_usage) + service.perform + end + end + + context 'when result is nil' do + let(:perform_result) { nil } + + it 'does not increment usage' do + expect(account).not_to receive(:increment_response_usage) + service.perform + end + end + + context 'when result is empty hash' do + let(:perform_result) { {} } + + it 'does not increment usage' do + expect(account).not_to receive(:increment_response_usage) + service.perform + end + end + + context 'when result has blank message' do + let(:perform_result) { { message: '' } } + + it 'does not increment usage' do + expect(account).not_to receive(:increment_response_usage) + service.perform + end + end + + context 'when result has nil message' do + let(:perform_result) { { message: nil } } + + it 'does not increment usage' do + expect(account).not_to receive(:increment_response_usage) + service.perform + end + end + + it 'actually increments the usage counter in custom_attributes' do + expect do + service.perform + account.reload + end.to change { account.custom_attributes['captain_responses_usage'].to_i }.by(1) + end + + context 'when captain is disabled' do + before do + allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(false) + end + + context 'when on Chatwoot Cloud' do + before do + allow(ChatwootApp).to receive(:chatwoot_cloud?).and_return(true) + end + + it 'returns upgrade error message' do + result = service.perform + expect(result[:error]).to eq(I18n.t('captain.upgrade')) + end + + it 'does not increment usage' do + expect(account).not_to receive(:increment_response_usage) + service.perform + end + end + + context 'when self-hosted' do + before do + allow(ChatwootApp).to receive(:chatwoot_cloud?).and_return(false) + end + + it 'returns disabled error message' do + result = service.perform + expect(result[:error]).to eq(I18n.t('captain.disabled')) + end + + it 'does not increment usage' do + expect(account).not_to receive(:increment_response_usage) + service.perform + end + end + end + + context 'when captain is enabled' do + before do + allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(true) + end + + it 'proceeds with the task' do + result = service.perform + expect(result[:message]).to eq('Test response') + expect(result[:error]).to be_nil + end + + it 'increments usage' do + expect(account).to receive(:increment_response_usage) + service.perform + end + end + end +end diff --git a/spec/lib/captain/base_task_service_spec.rb b/spec/lib/captain/base_task_service_spec.rb index ff3e97105..9ab6dc01e 100644 --- a/spec/lib/captain/base_task_service_spec.rb +++ b/spec/lib/captain/base_task_service_spec.rb @@ -22,6 +22,10 @@ RSpec.describe Captain::BaseTaskService do before do create(:installation_config, name: 'CAPTAIN_OPEN_AI_API_KEY', value: 'test-key') + # Stub captain enabled check to allow OSS specs to test base functionality + # without enterprise module interference + allow(account).to receive(:feature_enabled?).and_call_original + allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(true) end describe '#perform' do diff --git a/spec/lib/captain/follow_up_service_spec.rb b/spec/lib/captain/follow_up_service_spec.rb index 77648ab37..263707a47 100644 --- a/spec/lib/captain/follow_up_service_spec.rb +++ b/spec/lib/captain/follow_up_service_spec.rb @@ -25,6 +25,13 @@ RSpec.describe Captain::FollowUpService do ) end + before do + # Stub captain enabled check to allow specs to test base functionality + # without enterprise module interference + allow(account).to receive(:feature_enabled?).and_call_original + allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(true) + end + describe '#perform' do context 'when conversation_display_id is provided' do it 'resolves conversation for instrumentation' do diff --git a/spec/lib/captain/label_suggestion_service_spec.rb b/spec/lib/captain/label_suggestion_service_spec.rb index f99864288..37062e606 100644 --- a/spec/lib/captain/label_suggestion_service_spec.rb +++ b/spec/lib/captain/label_suggestion_service_spec.rb @@ -18,6 +18,10 @@ RSpec.describe Captain::LabelSuggestionService do allow(Llm::Config).to receive(:with_api_key).and_yield(mock_context) allow(mock_chat).to receive(:with_instructions) allow(mock_chat).to receive(:ask).and_return(mock_response) + # Stub captain enabled check to allow specs to test base functionality + # without enterprise module interference + allow(account).to receive(:feature_enabled?).and_call_original + allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(true) end describe '#label_suggestion_message' do diff --git a/spec/lib/captain/reply_suggestion_service_spec.rb b/spec/lib/captain/reply_suggestion_service_spec.rb index 64b1c0f11..c5cf52284 100644 --- a/spec/lib/captain/reply_suggestion_service_spec.rb +++ b/spec/lib/captain/reply_suggestion_service_spec.rb @@ -15,6 +15,10 @@ RSpec.describe Captain::ReplySuggestionService do allow(mock_chat).to receive(:with_instructions) allow(mock_chat).to receive(:add_message) allow(mock_chat).to receive(:ask).and_return(mock_response) + # Stub captain enabled check to allow specs to test base functionality + # without enterprise module interference + allow(account).to receive(:feature_enabled?).and_call_original + allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(true) end describe '#perform' do diff --git a/spec/lib/captain/rewrite_service_spec.rb b/spec/lib/captain/rewrite_service_spec.rb index f776118c8..cf53fc8f1 100644 --- a/spec/lib/captain/rewrite_service_spec.rb +++ b/spec/lib/captain/rewrite_service_spec.rb @@ -16,6 +16,10 @@ RSpec.describe Captain::RewriteService do allow(Llm::Config).to receive(:with_api_key).and_yield(mock_context) allow(mock_chat).to receive(:with_instructions) allow(mock_chat).to receive(:ask).and_return(mock_response) + # Stub captain enabled check to allow specs to test base functionality + # without enterprise module interference + allow(account).to receive(:feature_enabled?).and_call_original + allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(true) end describe '#perform with fix_spelling_grammar operation' do diff --git a/spec/lib/captain/summary_service_spec.rb b/spec/lib/captain/summary_service_spec.rb index 88e87f52b..87e90f557 100644 --- a/spec/lib/captain/summary_service_spec.rb +++ b/spec/lib/captain/summary_service_spec.rb @@ -14,6 +14,10 @@ RSpec.describe Captain::SummaryService do allow(Llm::Config).to receive(:with_api_key).and_yield(mock_context) allow(mock_chat).to receive(:with_instructions) allow(mock_chat).to receive(:ask).and_return(mock_response) + # Stub captain enabled check to allow specs to test base functionality + # without enterprise module interference + allow(account).to receive(:feature_enabled?).and_call_original + allow(account).to receive(:feature_enabled?).with('captain_integration').and_return(true) end describe '#perform' do