From eee6ccc6f56cbca5b02e46b150842ceaca750cd6 Mon Sep 17 00:00:00 2001
From: Shivam Mishra <scm.mymail@gmail.com>
Date: Tue, 13 Jan 2026 19:31:36 +0530
Subject: [PATCH] feat: add feature key for preference lookup

---
 lib/captain/base_task_service.rb           |  8 +++++++
 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 ++++
 spec/lib/captain/base_task_service_spec.rb | 27 ++++++++++++++++++++++
 7 files changed, 55 insertions(+)

diff --git a/lib/captain/base_task_service.rb b/lib/captain/base_task_service.rb
index 966057886..739239cd2 100644
--- a/lib/captain/base_task_service.rb
+++ b/lib/captain/base_task_service.rb
@@ -25,6 +25,14 @@ class Captain::BaseTaskService
     raise NotImplementedError, "#{self.class} must implement #event_name"
   end
 
+  def feature_key
+    raise NotImplementedError, "#{self.class} must implement #feature_key"
+  end
+
+  def configured_model
+    account.public_send("captain_#{feature_key}_model")
+  end
+
   def conversation
     @conversation ||= account.conversations.find_by(display_id: conversation_display_id)
   end
diff --git a/lib/captain/follow_up_service.rb b/lib/captain/follow_up_service.rb
index 4b925f396..5d87ef99c 100644
--- a/lib/captain/follow_up_service.rb
+++ b/lib/captain/follow_up_service.rb
@@ -95,4 +95,8 @@ class Captain::FollowUpService < Captain::BaseTaskService
   def event_name
     'follow_up'
   end
+
+  def feature_key
+    'editor'
+  end
 end
diff --git a/lib/captain/label_suggestion_service.rb b/lib/captain/label_suggestion_service.rb
index 02f8bd89a..2a1068a8b 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 feature_key
+    'label_suggestion'
+  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 6b80947a9..f4f6bbaf6 100644
--- a/lib/captain/reply_suggestion_service.rb
+++ b/lib/captain/reply_suggestion_service.rb
@@ -15,4 +15,8 @@ class Captain::ReplySuggestionService < Captain::BaseTaskService
   def event_name
     'reply_suggestion'
   end
+
+  def feature_key
+    'editor'
+  end
 end
diff --git a/lib/captain/rewrite_service.rb b/lib/captain/rewrite_service.rb
index 3a217d3c6..3340ec75b 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 feature_key
+    'editor'
+  end
 end
diff --git a/lib/captain/summary_service.rb b/lib/captain/summary_service.rb
index 16ee57b51..c609575f4 100644
--- a/lib/captain/summary_service.rb
+++ b/lib/captain/summary_service.rb
@@ -16,4 +16,8 @@ class Captain::SummaryService < Captain::BaseTaskService
   def event_name
     'summarize'
   end
+
+  def feature_key
+    'editor'
+  end
 end
diff --git a/spec/lib/captain/base_task_service_spec.rb b/spec/lib/captain/base_task_service_spec.rb
index 9ab6dc01e..1007796b2 100644
--- a/spec/lib/captain/base_task_service_spec.rb
+++ b/spec/lib/captain/base_task_service_spec.rb
@@ -15,6 +15,10 @@ RSpec.describe Captain::BaseTaskService do
       def event_name
         'test_event'
       end
+
+      def feature_key
+        'editor'
+      end
     end
   end
 
@@ -46,6 +50,29 @@ RSpec.describe Captain::BaseTaskService do
     end
   end
 
+  describe '#feature_key' do
+    it 'raises NotImplementedError for base class' do
+      base_service = described_class.new(account: account, conversation_display_id: conversation.display_id)
+      expect { base_service.send(:feature_key) }.to raise_error(NotImplementedError, /must implement #feature_key/)
+    end
+
+    it 'returns custom feature key in subclass' do
+      expect(service.send(:feature_key)).to eq('editor')
+    end
+  end
+
+  describe '#configured_model' do
+    it 'returns the configured model for the feature_key' do
+      expect(account).to receive(:captain_editor_model).and_return('gpt-5.1')
+      expect(service.send(:configured_model)).to eq('gpt-5.1')
+    end
+
+    it 'uses the default model when account model is not set' do
+      allow(account).to receive(:captain_editor_model).and_return('gpt-4.1-mini')
+      expect(service.send(:configured_model)).to eq('gpt-4.1-mini')
+    end
+  end
+
   describe '#conversation' do
     it 'finds conversation by display_id' do
       expect(service.send(:conversation)).to eq(conversation)