Merge branch 'codex/cw-7431-document-usage' into codex/ai-138-v2-faq-citations

db/migrate/20260721100000_create_captain_message_sources.rb

@@ -0,0 +1,17 @@
+class CreateCaptainMessageSources < ActiveRecord::Migration[7.1]
+  def change
+    create_table :captain_message_sources do |t|
+      t.references :account, null: false, index: true
+      t.references :assistant, null: false, index: true
+      t.references :conversation, null: false, index: true
+      t.references :message, null: false, index: true
+      t.references :document, null: false, index: true
+      t.bigint :assistant_response_id, null: false
+
+      t.timestamps
+    end
+
+    add_index :captain_message_sources, [:message_id, :assistant_response_id],
+              unique: true, name: 'idx_captain_message_sources_on_message_and_response'
+  end
+end

db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[7.1].define(version: 2026_07_13_184351) do
+ActiveRecord::Schema[7.1].define(version: 2026_07_21_100000) do
   # These extensions should be enabled to support this database
   enable_extension "pg_stat_statements"
   enable_extension "pg_trgm"
@@ -477,6 +477,23 @@ ActiveRecord::Schema[7.1].define(version: 2026_07_13_184351) do
     t.index ["user_id"], name: "index_captain_message_reports_on_user_id"
   end
 
+  create_table "captain_message_sources", force: :cascade do |t|
+    t.bigint "account_id", null: false
+    t.bigint "assistant_id", null: false
+    t.bigint "conversation_id", null: false
+    t.bigint "message_id", null: false
+    t.bigint "document_id", null: false
+    t.bigint "assistant_response_id", null: false
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+    t.index ["account_id"], name: "index_captain_message_sources_on_account_id"
+    t.index ["assistant_id"], name: "index_captain_message_sources_on_assistant_id"
+    t.index ["conversation_id"], name: "index_captain_message_sources_on_conversation_id"
+    t.index ["document_id"], name: "index_captain_message_sources_on_document_id"
+    t.index ["message_id", "assistant_response_id"], name: "idx_captain_message_sources_on_message_and_response", unique: true
+    t.index ["message_id"], name: "index_captain_message_sources_on_message_id"
+  end
+
   create_table "captain_scenarios", force: :cascade do |t|
     t.string "title"
     t.text "description"

enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb

@@ -11,7 +11,7 @@ class Api::V1::Accounts::Captain::DocumentsController < Api::V1::Accounts::BaseC
     @documents = filtered_documents
     @documents_count = @documents.count
     @sync_interval_hours = current_sync_interval&.in_hours&.to_i
-    @documents = with_responses_count(@documents).page(@current_page).per(RESULTS_PER_PAGE)
+    @documents = with_document_usage(@documents).page(@current_page).per(RESULTS_PER_PAGE)
   end
 
   def show; end
@@ -61,14 +61,23 @@ class Api::V1::Accounts::Captain::DocumentsController < Api::V1::Accounts::BaseC
     apply_sort(documents, permitted_params[:sort])
   end
 
-  def with_responses_count(scope)
-    scope.left_joins(:responses)
-         .select('captain_documents.*, COUNT(captain_assistant_responses.id) AS responses_count')
-         .group('captain_documents.id')
+  def with_document_usage(scope)
+    response_counts = Captain::AssistantResponse.where(documentable_type: 'Captain::Document')
+                                                .group(:documentable_id)
+                                                .select('documentable_id AS document_id, COUNT(*) AS responses_count')
+    source_counts = Captain::MessageSource.group(:document_id).select(
+      'document_id, COUNT(DISTINCT message_id) AS used_in_answers_count, COUNT(DISTINCT conversation_id) AS used_in_conversations_count'
+    )
+
+    scope.joins("LEFT JOIN (#{response_counts.to_sql}) response_counts ON response_counts.document_id = captain_documents.id")
+         .joins("LEFT JOIN (#{source_counts.to_sql}) source_counts ON source_counts.document_id = captain_documents.id")
+         .select('captain_documents.*, COALESCE(response_counts.responses_count, 0) AS responses_count, ' \
+                 'COALESCE(source_counts.used_in_answers_count, 0) AS used_in_answers_count, ' \
+                 'COALESCE(source_counts.used_in_conversations_count, 0) AS used_in_conversations_count')
   end
 
   def set_document
-    @document = @documents.find(permitted_params[:id])
+    @document = with_document_usage(@documents).find(permitted_params[:id])
   end
 
   def set_assistant

enterprise/app/jobs/captain/conversation/message_builder.rb

@@ -1,4 +1,8 @@
 module Captain::Conversation::MessageBuilder
+  SOURCE_MARKER_PATTERN = /\[\[source:(\d+)\]\]/
+  MALFORMED_SOURCE_MARKER_PATTERN = /\[\[source:[^\]\n]*\]\]?/
+  CUSTOMER_CITATION_PATTERN = /\[\[(\d+)\]\([^\)\n]*\)\]/
+
   private
 
   def collect_previous_messages
@@ -28,8 +32,67 @@ module Captain::Conversation::MessageBuilder
   end
 
   def create_messages
-    validate_message_content!(@response['response'])
-    create_outgoing_message(@response['response'], agent_name: @response['agent_name'])
+    message_content = resolve_citations(@response['response'])
+    validate_message_content!(message_content)
+    create_outgoing_message(message_content, agent_name: @response['agent_name'])
+  end
+
+  def resolve_citations(content)
+    return content if content.blank?
+
+    citation_urls = trusted_citation_urls
+    content = content.gsub(CUSTOMER_CITATION_PATTERN) do
+      citation_for(Regexp.last_match(1), citation_urls)
+    end
+    content = content.gsub(SOURCE_MARKER_PATTERN) do
+      citation_for(Regexp.last_match(1), citation_urls)
+    end
+    content.gsub(MALFORMED_SOURCE_MARKER_PATTERN, '')
+  end
+
+  def citation_for(source_number, citation_urls)
+    url = citation_urls[source_number.to_s]
+    return '' if url.blank?
+
+    "[[#{source_number}](#{url})]"
+  end
+
+  def trusted_citation_urls
+    return {} unless @assistant.config['feature_citation']
+
+    source_map = citation_source_map
+    return {} if source_map.empty?
+
+    document_links = Captain::Document.where(
+      id: source_map.values,
+      account_id: account.id,
+      assistant_id: @assistant.id
+    ).pluck(:id, :external_link).to_h
+
+    source_map.each_with_object({}) do |(source_number, document_id), urls|
+      external_link = document_links[document_id]
+      urls[source_number] = external_link if customer_safe_document_link?(external_link)
+    end
+  end
+
+  def citation_source_map
+    raw_citation_source_map.filter_map do |source_number, document_id|
+      [source_number.to_s, document_id.to_i] if valid_citation_source?(source_number, document_id)
+    end.to_h
+  end
+
+  def raw_citation_source_map
+    context = @run_result&.context || {}
+    state = context[:state] || context['state'] || {}
+    state[:captain_v2_citation_source_map] || state['captain_v2_citation_source_map'] || {}
+  end
+
+  def valid_citation_source?(source_number, document_id)
+    source_number.to_s.match?(/\A[1-9]\d*\z/) && document_id.to_s.match?(/\A\d+\z/)
+  end
+
+  def customer_safe_document_link?(external_link)
+    external_link.present? && !external_link.start_with?('PDF:')
   end
 
   def validate_message_content!(content)

enterprise/app/jobs/captain/conversation/response_builder_job.rb

@@ -68,7 +68,7 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob
         # left is the customer-facing follow-up message.
         process_v2_handoff
       end
-      capture_assistant_session(result_message: @handoff_message, credits_consumed: 0.0)
+      capture_assistant_session(result_message: @handoff_message, credits_consumed: 0.0, capture_message_sources: false)
     elsif v1_handoff_requested?
       # V1 only signals via the response string — no state has been touched yet. If
       # the conversation isn't pending anymore, a human took over mid-run; bail out
@@ -83,7 +83,7 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob
         Rails.logger.info("[CAPTAIN][ResponseBuilderJob] Incrementing response usage for #{account.id}")
         account.increment_response_usage
       end
-      capture_assistant_session(result_message: message, credits_consumed: 1.0)
+      capture_assistant_session(result_message: message, credits_consumed: 1.0, capture_message_sources: captain_v2_enabled?)
     end
   end
 
@@ -142,9 +142,10 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob
   # Capture runs outside the delivery transaction and never raises (the service
   # swallows its own failures): a session-logging bug must never roll back the
   # customer reply or trigger the top-level handle_error handoff on top of it.
-  def capture_assistant_session(result_message:, credits_consumed:)
+  def capture_assistant_session(result_message:, credits_consumed:, capture_message_sources:)
     Captain::Assistant::SessionCaptureService.new(assistant: @assistant, conversation: @conversation, run_result: @run_result,
-                                                  result_message: result_message, credits_consumed: credits_consumed).capture
+                                                  result_message: result_message, credits_consumed: credits_consumed,
+                                                  capture_message_sources: capture_message_sources).capture
   end
 
   def handle_error(error)

enterprise/app/models/captain/assistant.rb

@@ -108,6 +108,7 @@ class Captain::Assistant < ApplicationRecord
       name: name,
       description: description,
       product_name: config['product_name'] || 'this product',
+      citation_enabled: config['feature_citation'],
       scenarios: scenarios.enabled.map do |scenario|
         {
           title: scenario.title,

enterprise/app/models/captain/document.rb

@@ -34,6 +34,7 @@ class Captain::Document < ApplicationRecord
 
   belongs_to :assistant, class_name: 'Captain::Assistant'
   has_many :responses, class_name: 'Captain::AssistantResponse', dependent: :destroy, as: :documentable
+  has_many :message_sources, class_name: 'Captain::MessageSource', dependent: :destroy_async
   belongs_to :account
   has_one_attached :pdf_file
   store_accessor :metadata, :content_fingerprint, :last_sync_error_code, :sync_step, :openai_file_id

enterprise/app/models/captain/message_source.rb

@@ -0,0 +1,33 @@
+# == Schema Information
+#
+# Table name: captain_message_sources
+#
+#  id                    :bigint           not null, primary key
+#  created_at            :datetime         not null
+#  updated_at            :datetime         not null
+#  account_id            :bigint           not null
+#  assistant_id          :bigint           not null
+#  assistant_response_id :bigint           not null
+#  conversation_id       :bigint           not null
+#  document_id           :bigint           not null
+#  message_id            :bigint           not null
+#
+# Indexes
+#
+#  idx_captain_message_sources_on_message_and_response  (message_id,assistant_response_id) UNIQUE
+#  index_captain_message_sources_on_account_id          (account_id)
+#  index_captain_message_sources_on_assistant_id        (assistant_id)
+#  index_captain_message_sources_on_conversation_id     (conversation_id)
+#  index_captain_message_sources_on_document_id         (document_id)
+#  index_captain_message_sources_on_message_id          (message_id)
+#
+class Captain::MessageSource < ApplicationRecord
+  self.table_name = 'captain_message_sources'
+
+  belongs_to :account
+  belongs_to :assistant, class_name: 'Captain::Assistant'
+  belongs_to :conversation, class_name: '::Conversation'
+  belongs_to :message
+  belongs_to :document, class_name: 'Captain::Document'
+  belongs_to :assistant_response, class_name: 'Captain::AssistantResponse', optional: true
+end

enterprise/app/models/enterprise/concerns/account.rb

@@ -16,6 +16,7 @@ module Enterprise::Concerns::Account
     has_many :captain_documents, dependent: :destroy_async, class_name: 'Captain::Document'
     has_many :captain_custom_tools, dependent: :destroy_async, class_name: 'Captain::CustomTool'
     has_many :captain_agent_sessions, dependent: :destroy_async, class_name: 'Captain::AgentSession'
+    has_many :captain_message_sources, dependent: :destroy_async, class_name: 'Captain::MessageSource'
 
     has_many :copilot_threads, dependent: :destroy_async
     has_many :companies, dependent: :destroy_async

enterprise/app/models/enterprise/concerns/message.rb

@@ -4,5 +4,6 @@ module Enterprise::Concerns::Message
   included do
     has_one :call, dependent: :nullify
     has_many :message_reports, class_name: 'Captain::MessageReport', dependent: :destroy_async
+    has_many :captain_message_sources, class_name: 'Captain::MessageSource', dependent: :destroy_async
   end
 end

enterprise/app/services/captain/assistant/session_capture_service.rb

@@ -1,12 +1,13 @@
 class Captain::Assistant::SessionCaptureService
   SCENARIO_AGENT_REGEX = /\A#{Captain::Scenario::HANDOFF_KEY_PREFIX}_(\d+)_/
 
-  def initialize(assistant:, conversation:, run_result:, result_message:, credits_consumed:)
+  def initialize(assistant:, conversation:, run_result:, result_message:, **options)
     @assistant = assistant
     @conversation = conversation
     @run_result = run_result
     @result_message = result_message
-    @credits_consumed = credits_consumed
+    @credits_consumed = options.fetch(:credits_consumed)
+    @capture_message_sources = options.fetch(:capture_message_sources, false)
   end
 
   def capture
@@ -23,7 +24,7 @@ class Captain::Assistant::SessionCaptureService
   def capture!
     model = @assistant.agent_model
 
-    Captain::AgentSession.create!(
+    session = Captain::AgentSession.create!(
       assistant: @assistant,
       session_type: :assistant,
       subject: @conversation,
@@ -35,6 +36,8 @@ class Captain::Assistant::SessionCaptureService
       scenario_ids: scenario_ids,
       run_context: current_turn_history
     )
+    capture_message_sources(metadata) if @capture_message_sources
+    session
   end
 
   private
@@ -70,6 +73,42 @@ class Captain::Assistant::SessionCaptureService
     ids & @assistant.scenarios.where(id: ids).pluck(:id)
   end
 
+  def capture_message_sources(metadata)
+    sources = Array(metadata[:message_sources])
+    return if sources.empty?
+
+    document_ids = @assistant.documents.where(id: sources.pluck(:document_id)).pluck(:id)
+    return if document_ids.empty?
+
+    insert_message_sources(message_source_rows(sources, document_ids))
+  end
+
+  def message_source_rows(sources, document_ids)
+    timestamp = Time.current
+    sources.filter_map do |source|
+      next unless document_ids.include?(source[:document_id])
+
+      {
+        account_id: @assistant.account_id,
+        assistant_id: @assistant.id,
+        conversation_id: @conversation.id,
+        message_id: @result_message.id,
+        document_id: source[:document_id],
+        assistant_response_id: source[:assistant_response_id],
+        created_at: timestamp,
+        updated_at: timestamp
+      }
+    end
+  end
+
+  def insert_message_sources(rows)
+    return if rows.empty?
+
+    # rubocop:disable Rails/SkipsModelValidations
+    Captain::MessageSource.insert_all(rows, unique_by: :idx_captain_message_sources_on_message_and_response)
+    # rubocop:enable Rails/SkipsModelValidations
+  end
+
   # Trim to the current turn: the last user message and everything after it
   # (assistant replies, tool calls/results, handoff hops).
   def current_turn_history

enterprise/app/views/api/v1/models/captain/_document.json.jbuilder

@@ -11,6 +11,10 @@ json.file_size resource.file_size
 json.pdf_document resource.pdf_document?
 responses_count = resource.respond_to?(:responses_count) ? resource.responses_count : resource.responses.count
 json.responses_count responses_count.to_i
+used_in_answers_count = resource.respond_to?(:used_in_answers_count) ? resource.used_in_answers_count : resource.message_sources.distinct.count(:message_id)
+json.used_in_answers_count used_in_answers_count.to_i
+used_in_conversations_count = resource.respond_to?(:used_in_conversations_count) ? resource.used_in_conversations_count : resource.message_sources.distinct.count(:conversation_id)
+json.used_in_conversations_count used_in_conversations_count.to_i
 json.id resource.id
 json.name resource.name
 json.status resource.status

enterprise/lib/captain/prompts/assistant.liquid

@@ -10,6 +10,11 @@ You are {{name}}, a helpful, friendly, and knowledgeable assistant for the produ
 
 Don't digress away from your instructions, and use all the available tools at your disposal for solving customer issues. If you are to state something factual about {{product_name}}, use the `captain--tools--faq_lookup` tool to check the available information first.
 
+{% if citation_enabled -%}
+# Citations
+When you use FAQ content that includes a source marker, add that exact marker at the end of the supported sentence or paragraph. Source markers have the form `[[source:n]]`. Do not create URLs or markdown links. Reuse the same source marker when the same document supports more than one statement. Do not cite conversation context.
+{% endif -%}
+
 {% render 'current_time', current_time: current_time %}
 
 {% render 'core_rules' %}

enterprise/lib/captain/tools/faq_lookup_tool.rb

@@ -14,7 +14,7 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool
       "No relevant FAQs found for: #{query}"
     else
       log_tool_usage('found_results', { query: query, count: responses.size })
-      format_responses(responses)
+      format_responses(tool_context, responses)
     end
   end
 
@@ -26,34 +26,46 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool
     metadata = tool_context.state[:cw_metadata] ||= {}
     metadata[:faq_ids] = Array(metadata[:faq_ids]) | responses.map(&:id)
 
-    document_ids = responses.filter_map { |response| response.documentable_id if response.documentable_type == 'Captain::Document' }
+    document_ids = document_responses(responses).map(&:documentable_id)
     metadata[:document_ids] = Array(metadata[:document_ids]) | document_ids
+    metadata[:message_sources] = Array(metadata[:message_sources]) | message_sources(document_responses(responses))
+
+    document_ids.each { |document_id| citation_source_number(tool_context, document_id) }
   end
 
-  def format_responses(responses)
-    responses.map { |response| format_response(response) }.join
+  def document_responses(responses)
+    responses.select { |response| response.documentable_type == 'Captain::Document' }
   end
 
-  def format_response(response)
+  def message_sources(responses)
+    responses.map { |response| { assistant_response_id: response.id, document_id: response.documentable_id } }
+  end
+
+  def format_responses(tool_context, responses)
+    responses.map { |response| format_response(tool_context, response) }.join
+  end
+
+  def format_response(tool_context, response)
     formatted_response = "
         Question: #{response.question}
         Answer: #{response.answer}
         "
-    if should_show_source?(response)
+    if response.documentable_type == 'Captain::Document'
       formatted_response += "
-          Source: #{response.documentable.external_link}
+          Source marker: [[source:#{citation_source_number(tool_context, response.documentable_id)}]]
           "
     end
 
     formatted_response
   end
 
-  def should_show_source?(response)
-    return false if response.documentable.blank?
-    return false unless response.documentable.try(:external_link)
+  def citation_source_number(tool_context, document_id)
+    source_map = tool_context.state[:captain_v2_citation_source_map] ||= {}
+    existing_number = source_map.find { |_number, id| id.to_i == document_id }&.first
+    return existing_number if existing_number.present?
 
-    # Don't show source if it's a PDF placeholder
-    external_link = response.documentable.external_link
-    !external_link.start_with?('PDF:')
+    source_number = (source_map.keys.map(&:to_i).max || 0) + 1
+    source_map[source_number.to_s] = document_id
+    source_number.to_s
   end
 end

spec/enterprise/lib/captain/tools/faq_lookup_tool_spec.rb

@@ -66,12 +66,13 @@ RSpec.describe Captain::Tools::FaqLookupTool, type: :model do
         expect(result).to include('Answer: Go to settings and update email')
       end
 
-      it 'includes source link when document has external_link' do
+      it 'includes a source marker instead of the document link' do
         document.update!(external_link: 'https://help.example.com/password')
 
         result = tool.perform(tool_context, query: 'password')
 
-        expect(result).to include('Source: https://help.example.com/password')
+        expect(result).to include('Source marker: [[source:1]]')
+        expect(result).not_to include('https://help.example.com/password')
       end
 
       it 'logs tool usage for search' do