fix(captain): resolve V2 FAQ citations from trusted sources (#15159)

Captain V2 now adds FAQ citations from a structured model response. The
model returns ordered response parts with citation indexes, and Chatwoot
turns only trusted indexes into customer links.

## Before

Captain V2 asked the model to copy text markers such as `[[faq:1]]`.
Chatwoot used one regular expression to replace those markers with links
in the outgoing message and another regular expression to remove the
rendered links before the next model turn. Long conversations depended
on parsing the customer message to recover plain model context.

## After

The FAQ lookup tool now gives each eligible source document a numeric
index and never gives the model a URL. FAQ results from the same
document reuse the same index. The model returns `response_parts`, where
each part contains customer text and the supporting citation indexes.
Chatwoot checks every index against the document IDs registered during
the current run.

Only stored HTTP or HTTPS web-document links without embedded
credentials can appear in the customer reply. Blank links, PDF sources,
attachments, non-HTTP storage links, and unknown indexes do not create
links. Sources receive display numbers in the order they first appear,
and repeated sources keep the same display number.

Chatwoot saves the structured response parts with each newly generated
Captain message. Later Captain V2 turns use the saved plain text for
those messages, so they never need to parse rendered links. Existing
messages remain unchanged and continue to use their stored content. When
citations are disabled, Chatwoot clears citation indexes before it
returns or saves the response.

Captain V1, Copilot, legacy prompts, legacy tools, and the playground
response contract are unchanged. The playground continues to show the
plain `response` field.

## Closes

[AI-138](https://linear.app/chatwoot/issue/AI-138/faq-citation-fix)

## How to test

1. Open a conversation handled by a Captain V2 assistant and turn
citations off. Ask a greeting, an FAQ question, a code question, and a
follow up question. Confirm that the assistant answers normally and
shows no source links.
2. Turn citations on and ask a question that matches one public web
document. Confirm that the reply shows the stored public link after the
supported text.
3. Ask a question that needs two public web documents. Confirm that the
response order stays correct, each link appears after the supported
text, and repeated sources keep the same display number.
4. Ask a question that retrieves several FAQ results from one document.
Confirm that the reply shows the document once at each supported
response part rather than exposing separate FAQ sources.
5. Ask a question supported by a PDF, attachment, blank link, or
non-HTTP storage link. Confirm that Captain can use the information but
does not show a customer link.
6. Ask for a fenced code example with a citation. Confirm that the code
block stays complete and the citation appears after the closing fence.
7. Continue the conversation with a follow up question. Confirm that
Captain uses the earlier plain response text and does not receive or
repeat rendered citation links.
8. Test a scenario handoff in a conversation. Confirm that the handoff
and final response still work.

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

@@ -6,9 +6,11 @@ class Api::V1::Accounts::Captain::AgentSessionsController < Api::V1::Accounts::B
     @agent_session = Current.account.captain_agent_sessions.find_by(result_type: 'Message', result_id: @message.id)
     return head :not_found if @agent_session.blank?
 
-    @citations = Current.account.captain_assistant_responses
-                        .where(id: @agent_session.faq_ids)
-                        .includes(:documentable)
+    @citations = Current.account.captain_documents.where(id: @agent_session.cited_document_ids)
+    @used_faqs = Current.account.captain_assistant_responses.approved.where(
+      id: @agent_session.used_faq_ids,
+      documentable_type: 'User'
+    )
     @scenario_titles = Captain::Scenario.where(account_id: Current.account.id, id: @agent_session.scenario_ids)
                                         .pluck(:id, :title).to_h
   end

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

@@ -28,6 +28,16 @@ module Captain::Conversation::MessageBuilder
   end
 
   def create_messages
+    return create_v1_message unless captain_v2_enabled?
+
+    response_parts = Captain::Assistant::ResponseParts.from_response(@response)
+    citation_urls = @assistant.trusted_citation_urls(@run_result)
+    message_content = response_parts.customer_message_content(citation_urls: citation_urls)
+    validate_message_content!(message_content)
+    create_outgoing_message(message_content, agent_name: @response['agent_name'], response_parts: response_parts.to_a)
+  end
+
+  def create_v1_message
     validate_message_content!(@response['response'])
     create_outgoing_message(@response['response'], agent_name: @response['agent_name'])
   end
@@ -36,9 +46,10 @@ module Captain::Conversation::MessageBuilder
     raise ArgumentError, 'Message content cannot be blank' if content.blank?
   end
 
-  def create_outgoing_message(message_content, agent_name: nil, preserve_waiting_since: false)
+  def create_outgoing_message(message_content, agent_name: nil, response_parts: nil, preserve_waiting_since: false)
     additional_attrs = {}
     additional_attrs[:agent_name] = agent_name if agent_name.present?
+    additional_attrs[Captain::Assistant::ResponseParts::MESSAGE_ATTRIBUTE_KEY] = response_parts unless response_parts.nil?
 
     @conversation.messages.create!(
       message_type: :outgoing,

enterprise/app/models/captain/agent_session.rb

@@ -2,23 +2,25 @@
 #
 # Table name: agent_sessions
 #
-#  id               :bigint           not null, primary key
-#  credits_consumed :float
-#  document_ids     :jsonb
-#  faq_ids          :jsonb
-#  llm_model        :string
-#  result_type      :string
-#  run_context      :jsonb
-#  scenario_ids     :jsonb
-#  session_type     :integer          not null
-#  subject_type     :string           not null
-#  created_at       :datetime         not null
-#  updated_at       :datetime         not null
-#  account_id       :bigint           not null
-#  assistant_id     :bigint           not null
-#  result_id        :bigint
-#  subject_id       :bigint           not null
-#  user_id          :bigint
+#  id                 :bigint           not null, primary key
+#  cited_document_ids :jsonb            not null
+#  credits_consumed   :float
+#  document_ids       :jsonb
+#  faq_ids            :jsonb
+#  llm_model          :string
+#  result_type        :string
+#  run_context        :jsonb
+#  scenario_ids       :jsonb
+#  session_type       :integer          not null
+#  subject_type       :string           not null
+#  used_faq_ids       :jsonb            not null
+#  created_at         :datetime         not null
+#  updated_at         :datetime         not null
+#  account_id         :bigint           not null
+#  assistant_id       :bigint           not null
+#  result_id          :bigint
+#  subject_id         :bigint           not null
+#  user_id            :bigint
 #
 # Indexes
 #
@@ -27,6 +29,8 @@
 #  idx_on_account_id_subject_type_subject_id_6d60963b3d  (account_id,subject_type,subject_id)
 #  index_agent_sessions_on_account_id                    (account_id)
 #  index_agent_sessions_on_assistant_id                  (assistant_id)
+#  index_agent_sessions_on_cited_document_ids            (cited_document_ids) USING gin
+#  index_agent_sessions_on_used_faq_ids                  (used_faq_ids) USING gin
 #  index_agent_sessions_on_user_id                       (user_id)
 #
 class Captain::AgentSession < ApplicationRecord

enterprise/app/models/captain/assistant.rb

@@ -18,6 +18,7 @@
 #
 class Captain::Assistant < ApplicationRecord
   DESCRIPTION_LENGTH_LIMIT = 500
+  CITATION_SOURCES_STATE_KEY = :captain_v2_citation_sources
   AUTO_RESOLVE_MODES = %w[disabled legacy evaluated].freeze
 
   include Avatarable
@@ -107,6 +108,25 @@ class Captain::Assistant < ApplicationRecord
     }
   end
 
+  def customer_visible_citation_urls(citation_document_ids)
+    citation_documents = documents.where(id: citation_document_ids.values).index_by(&:id)
+    citation_urls = citation_document_ids.transform_values do |document_id|
+      citation_documents[document_id.to_i]&.customer_visible_source_url
+    end
+    citation_urls.compact.transform_keys(&:to_i)
+  end
+
+  def citations_enabled?
+    config['feature_citation']
+  end
+
+  def trusted_citation_urls(run_result)
+    return {} unless citations_enabled?
+
+    citation_document_ids = run_result&.context&.dig(:state, CITATION_SOURCES_STATE_KEY) || {}
+    customer_visible_citation_urls(citation_document_ids)
+  end
+
   private
 
   def set_default_auto_resolve_mode
@@ -132,6 +152,7 @@ class Captain::Assistant < ApplicationRecord
       name: name,
       description: description,
       product_name: config['product_name'] || 'this product',
+      citation_enabled: citations_enabled?,
       scenarios: scenarios.enabled.map do |scenario|
         {
           title: scenario.title,

enterprise/app/models/captain/assistant_response.rb

@@ -51,6 +51,10 @@ class Captain::AssistantResponse < ApplicationRecord
     nearest_neighbors(:embedding, embedding, distance: 'cosine').limit(5)
   end
 
+  def customer_visible_source_url
+    documentable.customer_visible_source_url if documentable.is_a?(Captain::Document)
+  end
+
   private
 
   def ensure_status

enterprise/app/models/captain/document.rb

@@ -103,6 +103,21 @@ class Captain::Document < ApplicationRecord
     end
   end
 
+  def customer_visible_source_url
+    return unless customer_visible_source?
+
+    url = external_link.presence
+    return unless url
+
+    uri = URI.parse(url)
+    return unless customer_visible_uri?(uri)
+    return if File.extname(uri.path).casecmp('.pdf').zero?
+
+    uri.to_s
+  rescue URI::InvalidURIError
+    nil
+  end
+
   def to_llm_metadata
     { document_id: id, assistant_id: assistant_id, external_link: external_link }
   end
@@ -121,6 +136,26 @@ class Captain::Document < ApplicationRecord
 
   private
 
+  def customer_visible_source?
+    !pdf_document? && !pdf_file.attached?
+  end
+
+  def customer_visible_uri?(uri)
+    return false unless uri.is_a?(URI::HTTP) && uri.host.present? && uri.userinfo.blank?
+
+    addresses = SsrfFilter::DEFAULT_RESOLVER.call(uri.host)
+    addresses.present? && addresses.all? { |ip| publicly_routable_address?(ip) }
+  rescue Resolv::ResolvError, Resolv::ResolvTimeout, IPAddr::InvalidAddressError
+    false
+  end
+
+  def publicly_routable_address?(ip)
+    return false if ip.ipv6? && SsrfFilter::NAT64_LOCAL_PREFIX.dup.include?(ip)
+
+    blocked_ranges = ip.ipv4? ? SsrfFilter::IPV4_BLACKLIST : SsrfFilter::IPV6_BLACKLIST
+    blocked_ranges.none? { |range| range.include?(ip) }
+  end
+
   def enqueue_crawl_job
     return if status != 'in_progress'
 

enterprise/app/models/captain/scenario.rb

@@ -67,6 +67,7 @@ class Captain::Scenario < ApplicationRecord
       instructions: resolved_instructions,
       tools: resolved_tools,
       assistant_name: assistant.name.downcase.gsub(/\s+/, '_'),
+      citation_enabled: assistant.citations_enabled?,
       response_guidelines: response_guidelines || [],
       guardrails: guardrails || []
     }

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

@@ -0,0 +1,67 @@
+module Captain::Assistant::AgentRunResponse
+  private
+
+  def process_agent_result(run_result)
+    Rails.logger.info "[Captain V2] Agent result: #{run_result.inspect}"
+    model_output = run_result.output
+    structured_response = if model_output.is_a?(Hash)
+                            model_output.with_indifferent_access
+                          else
+                            { 'response' => model_output.to_s, 'reasoning' => 'Processed by agent' }
+                          end
+    response_parts = Captain::Assistant::ResponseParts.from_response(structured_response)
+    response_parts = response_parts.without_citations unless @assistant.citations_enabled?
+    structured_response['response_parts'] = response_parts.to_a
+    structured_response['response'] = response_parts.plain_text
+    structured_response['agent_name'] = run_result.context&.dig(:current_agent)
+    structured_response['handoff_tool_called'] = run_result.context&.dig(:captain_v2_handoff_tool_called) || false
+    structured_response
+  end
+
+  def rewrite_oversized_response(run_result)
+    response_parts = Captain::Assistant::ResponseParts.from_response(run_result.output)
+    rendered_customer_message = customer_message_content(run_result)
+    citation_markup_length = rendered_customer_message.length - response_parts.plain_text.length
+    response_text_limit = message_length_limit - citation_markup_length
+    raise 'Captain citation links exceed the channel limit' unless response_text_limit.positive?
+
+    response_rewriter.rewrite(run_result, response_parts: response_parts, response_text_limit: response_text_limit)
+  end
+
+  def response_rewriter
+    @response_rewriter ||= Captain::Assistant::ResponseRewriter.new(
+      assistant: @assistant,
+      attribute_provider: Captain::Assistant::InstrumentationAttributeProvider.new(self)
+    )
+  end
+
+  def record_turn_start(run_result)
+    history = Array(run_result.context&.dig(:conversation_history))
+    turn_start_index = history.rindex { |message| message[:role].to_s == 'user' }
+    run_result.context[:captain_v2_turn_start_index] = turn_start_index if turn_start_index
+  end
+
+  def response_too_long?(run_result)
+    message_length_limit && customer_message_content(run_result).length > message_length_limit
+  end
+
+  def customer_message_content(run_result)
+    response_parts = Captain::Assistant::ResponseParts.from_response(run_result.output)
+    response_parts.customer_message_content(citation_urls: @assistant.trusted_citation_urls(run_result))
+  end
+
+  def message_length_limit
+    @message_length_limit ||= Captain::MessageLengthLimit.for(@conversation)
+  end
+
+  def error_response(error)
+    {
+      'response' => 'conversation_handoff',
+      'response_parts' => [{ 'text' => 'conversation_handoff', 'citation_indexes' => [] }],
+      'reasoning' => "Error occurred: #{error.message}",
+      'error' => true,
+      'error_reason' => error.class.name.underscore.tr('/', '_'),
+      'handoff_tool_called' => @handoff_tool_called
+    }
+  end
+end

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

@@ -3,6 +3,7 @@ require 'agents/instrumentation'
 
 class Captain::Assistant::AgentRunnerService
   include Captain::Assistant::RunnerCallbacksHelper
+  include Captain::Assistant::AgentRunResponse
   include Captain::Assistant::RunnerInstrumentationHelper
   include Captain::Assistant::TracePayloadHelper
   include Captain::Assistant::RunnerStateHelper
@@ -86,7 +87,10 @@ class Captain::Assistant::AgentRunnerService
 
   def extract_text_from_content(content)
     # Handle structured output from agents
-    return content[:response] || content['response'] || content.to_s if content.is_a?(Hash)
+    if content.is_a?(Hash)
+      response_text = Captain::Assistant::ResponseParts.from_response(content).plain_text
+      return response_text.presence || content.to_s
+    end
 
     return content unless content.is_a?(Array)
 
@@ -94,54 +98,6 @@ class Captain::Assistant::AgentRunnerService
     text_parts.join(' ')
   end
 
-  def process_agent_result(result)
-    Rails.logger.info "[Captain V2] Agent result: #{result.inspect}"
-    output = result.output
-    response = output.is_a?(Hash) ? output.with_indifferent_access : { 'response' => output.to_s, 'reasoning' => 'Processed by agent' }
-    response['agent_name'] = result.context&.dig(:current_agent)
-    response['handoff_tool_called'] = result.context&.dig(:captain_v2_handoff_tool_called) || false
-    response
-  end
-
-  def rewrite_oversized_response(result)
-    response_rewriter.rewrite(result, response: response_text(result), limit: message_length_limit)
-  end
-
-  def response_rewriter
-    @response_rewriter ||= Captain::Assistant::ResponseRewriter.new(
-      assistant: @assistant,
-      attribute_provider: Captain::Assistant::InstrumentationAttributeProvider.new(self)
-    )
-  end
-
-  def record_turn_start(result)
-    history = Array(result.context&.dig(:conversation_history))
-    turn_start_index = history.rindex { |message| message[:role].to_s == 'user' }
-    result.context[:captain_v2_turn_start_index] = turn_start_index if turn_start_index
-  end
-
-  def response_too_long?(result)
-    message_length_limit && response_text(result).length > message_length_limit
-  end
-
-  def response_text(result)
-    extract_text_from_content(result.output).to_s
-  end
-
-  def message_length_limit
-    @message_length_limit ||= Captain::MessageLengthLimit.for(@conversation)
-  end
-
-  def error_response(error)
-    {
-      'response' => 'conversation_handoff',
-      'reasoning' => "Error occurred: #{error.message}",
-      'error' => true,
-      'error_reason' => error.class.name.underscore.tr('/', '_'),
-      'handoff_tool_called' => @handoff_tool_called
-    }
-  end
-
   def build_and_wire_agents
     assistant_agent = @assistant.agent
     scenario_agents = @assistant.scenarios.enabled.map(&:agent)

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

@@ -0,0 +1,70 @@
+class Captain::Assistant::ResponseParts
+  MESSAGE_ATTRIBUTE_KEY = 'captain_v2_response_parts'.freeze
+  CLOSING_CODE_FENCE_LINE = /\A(?:`{3,}|~{3,})\z/
+
+  attr_reader :parts
+
+  def self.from_response(response)
+    return new([{ text: response.to_s, citation_indexes: [] }]) unless response.is_a?(Hash)
+
+    response = response.with_indifferent_access
+    return new(response[:response_parts]) if response.key?(:response_parts)
+
+    new([{ text: response[:response], citation_indexes: [] }])
+  end
+
+  def initialize(response_parts)
+    @parts = Array(response_parts).filter_map { |part| normalize_part(part) }
+  end
+
+  def plain_text
+    parts.pluck('text').join("\n\n")
+  end
+
+  def without_citations
+    self.class.new(parts.map { |part| part.merge('citation_indexes' => []) })
+  end
+
+  def customer_message_content(citation_urls: {})
+    display_numbers = {}
+
+    parts.map do |part|
+      links = part['citation_indexes'].filter_map do |citation_index|
+        url = citation_urls[citation_index]
+        next if url.blank?
+
+        # Number trusted sources by first appearance and reuse that number for later references.
+        display_number = display_numbers[url] ||= display_numbers.size + 1
+        markdown_safe_url = url.gsub('(', '%28').gsub(')', '%29')
+        "[[#{display_number}](#{markdown_safe_url})]"
+      end.uniq
+
+      final_text_line = part['text'].lines.last.to_s.strip
+      citation_separator = final_text_line.match?(CLOSING_CODE_FENCE_LINE) ? "\n" : ' '
+      [part['text'], links.join(' ')].compact_blank.join(citation_separator)
+    end.join("\n\n")
+  end
+
+  def to_a
+    parts
+  end
+
+  private
+
+  def normalize_part(part)
+    return unless part.is_a?(Hash)
+
+    part = part.with_indifferent_access
+    return unless part[:text].is_a?(String)
+
+    text = part[:text].strip
+    return if text.blank?
+
+    {
+      'text' => text,
+      'citation_indexes' => Array(part[:citation_indexes]).select do |citation_index|
+        citation_index.is_a?(Integer) && citation_index.positive?
+      end.uniq
+    }
+  end
+end

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

@@ -12,26 +12,55 @@ class Captain::Assistant::ResponseRewriter
     @attribute_provider = attribute_provider
   end
 
-  def rewrite(result, response:, limit:)
-    prompt = "The response below is #{response.length} characters, but this channel allows a maximum of #{limit}. " \
-             'Shorten it while preserving names, numbers, dates, links, warnings, and completed actions. ' \
-             "Do not add facts.\n\nResponse:\n#{response}"
-    context = {
-      session_id: result.context[:session_id],
-      state: result.context[:state],
-      captain_v2_trace_input: prompt
+  def rewrite(run_result, response_parts:, response_text_limit:)
+    rewrite_prompt = build_rewrite_prompt(response_parts, response_text_limit)
+    rewrite_context = {
+      session_id: run_result.context[:session_id],
+      state: run_result.context[:state],
+      captain_v2_trace_input: rewrite_prompt
     }
-    rewritten_result = runner.run(prompt, context: context, max_turns: 1)
-    raise rewritten_result.error || 'Captain response rewrite failed' if rewritten_result.failed?
+    rewrite_run_result = runner.run(rewrite_prompt, context: rewrite_context, max_turns: 1)
+    raise rewrite_run_result.error || 'Captain response rewrite failed' if rewrite_run_result.failed?
 
-    replace_final_assistant_output(result.context[:conversation_history], rewritten_result.output)
-    replace_final_assistant_output(result.messages, rewritten_result.output)
-    result.output = rewritten_result.output
-    result
+    rewritten_model_output = rewritten_response_with_original_citations(rewrite_run_result.output, response_parts)
+    replace_final_assistant_output(run_result.context[:conversation_history], rewritten_model_output)
+    replace_final_assistant_output(run_result.messages, rewritten_model_output)
+    run_result.output = rewritten_model_output
+    run_result
   end
 
   private
 
+  def build_rewrite_prompt(response_parts, response_text_limit)
+    response_text = response_parts.plain_text
+
+    "The response text below is #{response_text.length} characters, but it must be at most #{response_text_limit} characters " \
+      'so the complete customer message fits this channel. ' \
+      'Shorten the text in each response part so the joined customer message fits the limit. ' \
+      'Keep the same number and order of response parts. Change only each text field. ' \
+      'Preserve names, numbers, dates, links, warnings, and completed actions. Keep Markdown valid within each text field. Do not add facts. ' \
+      "Return the complete structured response.\n\nResponse parts:\n#{response_parts.to_a.to_json}"
+  end
+
+  def rewritten_response_with_original_citations(rewritten_model_output, original_response_parts)
+    rewritten_model_output = rewritten_model_output.with_indifferent_access
+    rewritten_response_parts = Captain::Assistant::ResponseParts.from_response(rewritten_model_output)
+    original_parts = original_response_parts.to_a
+    rewritten_parts = rewritten_response_parts.to_a
+    raise 'Captain response rewrite changed the number of response parts' unless rewritten_parts.size == original_parts.size
+
+    original_citation_indexes = original_parts.pluck('citation_indexes')
+    rewritten_citation_indexes = rewritten_parts.pluck('citation_indexes')
+    if @assistant.citations_enabled? && rewritten_citation_indexes != original_citation_indexes
+      raise 'Captain response rewrite changed the response part citation order'
+    end
+
+    rewritten_model_output['response_parts'] = rewritten_parts.zip(original_parts).map do |rewritten_part, original_part|
+      rewritten_part.merge('citation_indexes' => original_part['citation_indexes'])
+    end
+    rewritten_model_output
+  end
+
   def runner
     @runner ||= begin
       agent = Agents::Agent.new(

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

@@ -31,6 +31,8 @@ class Captain::Assistant::SessionCaptureService
       llm_model: "#{Llm::Models.provider_for(model)}-#{model}",
       credits_consumed: @credits_consumed,
       faq_ids: metadata[:faq_ids] || [],
+      used_faq_ids: metadata[:used_faq_ids] || [],
+      cited_document_ids: cited_document_ids,
       document_ids: metadata[:document_ids] || [],
       scenario_ids: scenario_ids,
       run_context: current_turn_history
@@ -47,6 +49,18 @@ class Captain::Assistant::SessionCaptureService
     @metadata ||= context.dig(:state, :cw_metadata) || {}
   end
 
+  def cited_document_ids
+    return [] unless @assistant.config['feature_citation']
+
+    citation_document_ids = (context.dig(:state, Captain::Assistant::CITATION_SOURCES_STATE_KEY) || {}).transform_keys(&:to_i)
+    visible_citation_indexes = @assistant.customer_visible_citation_urls(citation_document_ids).keys
+    stored_response_parts = result_message.additional_attributes.to_h[Captain::Assistant::ResponseParts::MESSAGE_ATTRIBUTE_KEY]
+    response_parts = Captain::Assistant::ResponseParts.new(stored_response_parts)
+    selected_citation_indexes = response_parts.to_a.flat_map { |part| part['citation_indexes'] }.uniq
+
+    (selected_citation_indexes & visible_citation_indexes).filter_map { |index| citation_document_ids[index] }.uniq
+  end
+
   # On handoff, HandoffTool records the private reason note it created; the session
   # attaches there so agents can inspect the generation path on the note itself.
   def result_message

enterprise/app/services/captain/conversation/message_history_builder_service.rb

@@ -24,8 +24,11 @@ class Captain::Conversation::MessageHistoryBuilderService
   def message_hash_for_context(message)
     return activity_message_hash(message) if message.message_type == 'activity'
 
+    content = prepare_multimodal_message_content(message)
+    return if content.blank?
+
     {
-      content: prepare_multimodal_message_content(message),
+      content: content,
       role: determine_role(message)
     }
   end
@@ -45,6 +48,14 @@ class Captain::Conversation::MessageHistoryBuilderService
   end
 
   def prepare_multimodal_message_content(message)
+    if message.outgoing? && message.sender_type == 'Captain::Assistant' && !message.deleted
+      message_attributes = message.additional_attributes.to_h
+      response_parts_attribute = Captain::Assistant::ResponseParts::MESSAGE_ATTRIBUTE_KEY
+      if message_attributes.key?(response_parts_attribute)
+        return Captain::Assistant::ResponseParts.new(message_attributes[response_parts_attribute]).plain_text.presence
+      end
+    end
+
     Captain::OpenAiMessageBuilderService.new(message: message).generate_content
   end
 end

enterprise/app/views/api/v1/accounts/captain/agent_sessions/show.json.jbuilder

@@ -5,13 +5,17 @@ json.credits_consumed @agent_session.credits_consumed
 json.run_context @agent_session.run_context.is_a?(Array) ? @agent_session.run_context : []
 json.citations @citations do |citation|
   json.id citation.id
-  json.title citation.question
+  json.title citation.name
   # display_url resolves uploaded PDFs to their blob URL; external_link holds a
   # "PDF: ..." placeholder for those. Guard on scheme so placeholders render as
   # plain text instead of dead anchors.
-  link = citation.documentable.is_a?(Captain::Document) ? citation.documentable.display_url : nil
+  link = citation.display_url
   json.link link&.match?(%r{\Ahttps?://}) ? link : nil
 end
+json.used_faqs @used_faqs do |faq|
+  json.id faq.id
+  json.title faq.question
+end
 json.scenarios @scenario_titles do |id, title|
   json.id id
   json.title title

enterprise/lib/captain/prompts/assistant.liquid

@@ -10,6 +10,8 @@ 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.
 
+{% render 'citations', citation_enabled: citation_enabled %}
+
 {% render 'current_time', current_time: current_time %}
 
 {% render 'core_rules' %}

enterprise/lib/captain/prompts/scenario.liquid

@@ -8,6 +8,8 @@ You are a specialized agent called "{{ title }}", your task is to handle the fol
 
 If you believe the user's request is not within the scope of your role, you can assign this conversation back to the orchestrator agent using the `handoff_to_{{ assistant_name }}` tool
 
+{% render 'citations', citation_enabled: citation_enabled %}
+
 {% render 'current_time', current_time: current_time %}
 
 {% render 'core_rules' %}

enterprise/lib/captain/prompts/snippets/citations.liquid

@@ -0,0 +1,4 @@
+{% if citation_enabled -%}
+# Citations
+Each eligible FAQ result can include a numeric citation index. Put the supporting result's index in the `citation_indexes` array for each response part that uses it. Use only indexes shown in FAQ results. Reuse the same index for repeated references. Leave the array empty when the response part has no eligible source. Never include a source URL in response text or a citation field.
+{% endif -%}

enterprise/lib/captain/response_schema.rb

@@ -1,6 +1,16 @@
 # TODO: Wrap the schema lib under ai-agents
 # So we can extend it as Agents::Schema
 class Captain::ResponseSchema < RubyLLM::Schema
-  string :response, description: 'The message to send to the user'
+  array :response_parts,
+        description: 'Ordered parts of the message to send to the user. Keep all customer-visible text within each part text field.',
+        min_items: 1 do
+    object do
+      string :text, description: 'Customer-visible response text without citation markers or source URLs.', min_length: 1
+      array :citation_indexes,
+            description: 'Numeric citation indexes from FAQ results that support this text. Use an empty array when none apply.' do
+        integer minimum: 1
+      end
+    end
+  end
   string :reasoning, description: "Agent's thought process"
 end

enterprise/lib/captain/tools/faq_lookup_tool.rb

@@ -6,7 +6,7 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool
     log_tool_usage('searching', { query: query })
 
     # Use existing vector search on approved responses
-    responses = @assistant.responses.approved.search(query).to_a
+    responses = @assistant.responses.approved.search(query).includes(:documentable).to_a
     record_retrieved_sources(tool_context, responses)
 
     if responses.empty?
@@ -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
 
@@ -30,34 +30,42 @@ 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' }
+    responses_by_type = responses.group_by(&:documentable_type)
+    document_ids = Array(responses_by_type['Captain::Document']).map(&:documentable_id)
     metadata[:document_ids] = Array(metadata[:document_ids]) | document_ids
+
+    used_faq_ids = Array(responses_by_type['User']).map(&:id)
+    metadata[:used_faq_ids] = Array(metadata[:used_faq_ids]) | used_faq_ids
   end
 
-  def format_responses(responses)
-    responses.map { |response| format_response(response) }.join
+  def format_responses(tool_context, responses)
+    responses.map { |response| format_response(tool_context, response) }.join
   end
 
-  def format_response(response)
+  def format_response(tool_context, response)
     formatted_response = "
+        FAQ result:
+        "
+    if @assistant.citations_enabled? && response.customer_visible_source_url.present?
+      formatted_response += "
+          Citation index: #{citation_index(tool_context, response)}
+          "
+    end
+    formatted_response += "
         Question: #{response.question}
         Answer: #{response.answer}
         "
-    if should_show_source?(response)
-      formatted_response += "
-          Source: #{response.documentable.external_link}
-          "
-    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_index(tool_context, response)
+    citation_document_ids = tool_context.state[Captain::Assistant::CITATION_SOURCES_STATE_KEY] ||= {}
+    existing_index = citation_document_ids.find { |_index, document_id| document_id == response.documentable_id }&.first
+    return existing_index if existing_index.present?
 
-    # Don't show source if it's a PDF placeholder
-    external_link = response.documentable.external_link
-    !external_link.start_with?('PDF:')
+    next_citation_index = citation_document_ids.size + 1
+    citation_document_ids[next_citation_index] = response.documentable_id
+    next_citation_index
   end
 end