feat: enhance citation handling and add metadata support in citation blocks

src/renderer/src/pages/home/Messages/Blocks/CitationBlock.tsx

@@ -12,14 +12,14 @@ import CitationsList from '../CitationsList'
 
 function CitationBlock({ block }: { block: CitationMessageBlock }) {
   const formattedCitations = useSelector((state: RootState) => selectFormattedCitationsByBlockId(state, block.id))
-  const hasCitations = useMemo(() => {
   const hasGeminiBlock = block.response?.source === WebSearchSource.GEMINI
+  const hasCitations = useMemo(() => {
     return (
       (formattedCitations && formattedCitations.length > 0) ||
       hasGeminiBlock ||
       (block.knowledge && block.knowledge.length > 0)
     )
-  }, [formattedCitations, block.response, block.knowledge])
+  }, [formattedCitations, block.knowledge, hasGeminiBlock])
 
   if (block.status === MessageBlockStatus.PROCESSING) {
     return <Spinner text="message.searching" />
@@ -29,12 +29,10 @@ function CitationBlock({ block }: { block: CitationMessageBlock }) {
     return null
   }
 
-  const isGemini = block.response?.source === WebSearchSource.GEMINI
-
   return (
     <>
       {block.status === MessageBlockStatus.SUCCESS &&
-        (isGemini ? (
+        (hasGeminiBlock ? (
           <>
             <CitationsList citations={formattedCitations} />
             <SearchEntryPoint

src/renderer/src/pages/home/Messages/Blocks/MainTextBlock.tsx

@@ -1,8 +1,9 @@
+import { GroundingSupport } from '@google/genai'
 import { useSettings } from '@renderer/hooks/useSettings'
 import { getModelUniqId } from '@renderer/services/ModelService'
 import type { RootState } from '@renderer/store'
 import { selectFormattedCitationsByBlockId } from '@renderer/store/messageBlock'
-import type { Model } from '@renderer/types'
+import { type Model, WebSearchSource } from '@renderer/types'
 import type { MainTextMessageBlock, Message } from '@renderer/types/newMessage'
 import { Flex } from 'antd'
 import React, { useMemo } from 'react'
@@ -47,8 +48,9 @@ const MainTextBlock: React.FC<Props> = ({ block, citationBlockId, role, mentions
       return content
     }
 
-    // FIXME：性能问题，需要优化
+    switch (block.citationReferences[0].citationBlockSource) {
-    // Replace all citation numbers in the content with formatted citations
+      case WebSearchSource.OPENAI_COMPATIBLE:
+      case WebSearchSource.OPENAI: {
         formattedCitations.forEach((citation) => {
           const citationNum = citation.number
           const supData = {
@@ -58,12 +60,83 @@ const MainTextBlock: React.FC<Props> = ({ block, citationBlockId, role, mentions
             content: citation.content?.substring(0, 200)
           }
           const citationJson = encodeHTML(JSON.stringify(supData))
+
+          // Handle[<sup>N</sup>](url)
+          const preFormattedRegex = new RegExp(`\\[<sup>${citationNum}</sup>\\]\\(.*?\\)`, 'g')
+
           const citationTag = `[<sup data-citation='${citationJson}'>${citationNum}</sup>](${citation.url})`
 
-      // Replace all occurrences of [citationNum] with the formatted citation
+          content = content.replace(preFormattedRegex, citationTag)
-      const regex = new RegExp(`\\[${citationNum}\\]`, 'g')
-      content = content.replace(regex, citationTag)
         })
+        break
+      }
+      case WebSearchSource.GEMINI: {
+        // First pass: Add basic citation marks using metadata
+        let processedContent = content
+        const firstCitation = formattedCitations[0]
+        if (firstCitation?.metadata) {
+          console.log('groundingSupport, ', firstCitation.metadata)
+          firstCitation.metadata.forEach((support: GroundingSupport) => {
+            const citationNums = support.groundingChunkIndices!
+
+            if (support.segment) {
+              const text = support.segment.text!
+              // 生成引用标记
+              const basicTag = citationNums
+                .map((citationNum) => {
+                  const citation = formattedCitations.find((c) => c.number === citationNum + 1)
+                  return citation ? `[<sup>${citationNum + 1}</sup>](${citation.url})` : ''
+                })
+                .join('')
+
+              // 在文本后面添加引用标记，而不是替换
+              if (text && basicTag) {
+                processedContent = processedContent.replace(text, `${text}${basicTag}`)
+              }
+            }
+          })
+          content = processedContent
+        }
+        // Second pass: Replace basic citations with full citation data
+        formattedCitations.forEach((citation) => {
+          const citationNum = citation.number
+          const supData = {
+            id: citationNum,
+            url: citation.url,
+            title: citation.title || citation.hostname || '',
+            content: citation.content?.substring(0, 200)
+          }
+          const citationJson = encodeHTML(JSON.stringify(supData))
+
+          // Replace basic citation with full citation including data
+          const basicCitationRegex = new RegExp(`\\[<sup>${citationNum}</sup>\\]\\(${citation.url}\\)`, 'g')
+          const fullCitationTag = `[<sup data-citation='${citationJson}'>${citationNum}</sup>](${citation.url})`
+          content = content.replace(basicCitationRegex, fullCitationTag)
+        })
+        break
+      }
+      default: {
+        // FIXME：性能问题，需要优化
+        // Replace all citation numbers and pre-formatted links with formatted citations
+        formattedCitations.forEach((citation) => {
+          const citationNum = citation.number
+          const supData = {
+            id: citationNum,
+            url: citation.url,
+            title: citation.title || citation.hostname || '',
+            content: citation.content?.substring(0, 200)
+          }
+          const citationJson = encodeHTML(JSON.stringify(supData))
+
+          // Handle both plain references [N] and pre-formatted links [<sup>N</sup>](url)
+          const plainRefRegex = new RegExp(`\\[${citationNum}\\]`, 'g')
+
+          const citationTag = `[<sup data-citation='${citationJson}'>${citationNum}</sup>](${citation.url})`
+
+          content = content.replace(plainRefRegex, citationTag)
+        })
+      }
+    }
 
     return content
   }, [block.content, block.citationReferences, citationBlockId, formattedCitations])

src/renderer/src/pages/home/Messages/CitationsList.tsx

@@ -16,6 +16,7 @@ export interface Citation {
   content?: string
   showFavicon?: boolean
   type?: string
+  metadata?: Record<string, any>
 }
 
 interface CitationsListProps {

src/renderer/src/providers/AiProvider/OpenAICompatibleProvider.ts

@@ -331,7 +331,7 @@ export default class OpenAICompatibleProvider extends OpenAIProvider {
     const model = assistant.model || defaultModel
 
     const { contextCount, maxTokens, streamOutput } = getAssistantSettings(assistant)
-    const isEnabledWebSearch = assistant.enableWebSearch || !!assistant.webSearchProviderId
+    const isEnabledBultinWebSearch = assistant.enableWebSearch
     messages = addImageFileToContents(messages)
     const enableReasoning =
       ((isSupportedThinkingTokenModel(model) || isSupportedReasoningEffortModel(model)) &&
@@ -597,7 +597,7 @@ export default class OpenAICompatibleProvider extends OpenAIProvider {
                 }
               }
               if (
-                isEnabledWebSearch &&
+                isEnabledBultinWebSearch &&
                 isZhipuModel(model) &&
                 finishReason === 'stop' &&
                 originalFinishRawChunk?.web_search
@@ -611,7 +611,7 @@ export default class OpenAICompatibleProvider extends OpenAIProvider {
                 } as LLMWebSearchCompleteChunk)
               }
               if (
-                isEnabledWebSearch &&
+                isEnabledBultinWebSearch &&
                 isHunyuanSearchModel(model) &&
                 originalFinishRawChunk?.search_info?.search_results
               ) {

src/renderer/src/providers/AiProvider/OpenAIProvider.ts

@@ -287,7 +287,7 @@ export default class OpenAIProvider extends BaseProvider {
     const model = assistant.model || defaultModel
     const { contextCount, maxTokens, streamOutput } = getAssistantSettings(assistant)
 
-    const isEnabledWebSearch = assistant.enableWebSearch || !!assistant.webSearchProviderId
+    const isEnabledBuiltinWebSearch = assistant.enableWebSearch
     onChunk({ type: ChunkType.LLM_RESPONSE_CREATED })
     // 退回到 OpenAI 兼容模式
     if (isOpenAIWebSearch(model)) {
@@ -342,7 +342,7 @@ export default class OpenAIProvider extends BaseProvider {
           const delta = chunk.choices[0]?.delta
           const finishReason = chunk.choices[0]?.finish_reason
           if (delta?.content) {
-            if (delta?.annotations) {
+            if (isOpenAIWebSearch(model)) {
               delta.content = convertLinks(delta.content || '', isFirstChunk)
             }
             if (isFirstChunk) {
@@ -388,7 +388,10 @@ export default class OpenAIProvider extends BaseProvider {
       return
     }
     const tools: OpenAI.Responses.Tool[] = []
-    if (isEnabledWebSearch) {
+    const toolChoices: OpenAI.Responses.ToolChoiceTypes = {
+      type: 'web_search_preview'
+    }
+    if (isEnabledBuiltinWebSearch) {
       tools.push({
         type: 'web_search_preview'
       })
@@ -558,17 +561,22 @@ export default class OpenAIProvider extends BaseProvider {
               thinking_millsec: new Date().getTime() - time_first_token_millsec
             })
             break
-          case 'response.output_text.delta':
+          case 'response.output_text.delta': {
+            let delta = chunk.delta
+            if (isEnabledBuiltinWebSearch) {
+              delta = convertLinks(delta)
+            }
             onChunk({
               type: ChunkType.TEXT_DELTA,
-              text: chunk.delta
+              text: delta
             })
-            content += chunk.delta
+            content += delta
             break
+          }
           case 'response.output_text.done':
             onChunk({
               type: ChunkType.TEXT_COMPLETE,
-              text: chunk.text
+              text: content
             })
             break
           case 'response.content_part.done':
@@ -633,6 +641,7 @@ export default class OpenAIProvider extends BaseProvider {
         max_output_tokens: maxTokens,
         stream: streamOutput,
         tools: tools.length > 0 ? tools : undefined,
+        tool_choice: isEnabledBuiltinWebSearch ? toolChoices : undefined,
         service_tier: this.getServiceTier(model),
         ...this.getResponseReasoningEffort(assistant, model),
         ...this.getCustomParameters(assistant)

src/renderer/src/store/messageBlock.ts

@@ -85,19 +85,22 @@ const formatCitationsFromBlock = (block: CitationMessageBlock | undefined): Cita
   if (!block) return []
 
   let formattedCitations: Citation[] = []
-  // 1. Handle Web Search Responses (Non-Gemini)
+  // 1. Handle Web Search Responses
   if (block.response) {
     switch (block.response.source) {
-      case WebSearchSource.GEMINI:
+      case WebSearchSource.GEMINI: {
+        const groundingMetadata = block.response.results as GroundingMetadata
         formattedCitations =
-          (block.response?.results as GroundingMetadata)?.groundingChunks?.map((chunk, index) => ({
+          groundingMetadata?.groundingChunks?.map((chunk, index) => ({
             number: index + 1,
             url: chunk?.web?.uri || '',
             title: chunk?.web?.title,
-            showFavicon: false,
+            showFavicon: true,
+            metadata: groundingMetadata.groundingSupports,
             type: 'websearch'
           })) || []
         break
+      }
       case WebSearchSource.OPENAI:
         formattedCitations =
           (block.response.results as OpenAI.Responses.ResponseOutputText.URLCitation[])?.map((result, index) => {

src/renderer/src/store/thunk/messageThunk.ts

@@ -470,7 +470,6 @@ const fetchAndProcessAssistantResponseImpl = async (
         saveUpdatedBlockToDB(citationBlock.id, assistantMsgId, topicId, getState)
       },
       onExternalToolComplete: (externalToolResult: ExternalToolResult) => {
-        console.warn('onExternalToolComplete received.', externalToolResult)
         if (citationBlockId) {
           const changes: Partial<CitationMessageBlock> = {
             response: externalToolResult.webSearch,
@@ -505,19 +504,24 @@ const fetchAndProcessAssistantResponseImpl = async (
           )
           citationBlockId = citationBlock.id
           handleBlockTransition(citationBlock, MessageBlockType.CITATION)
+        }
         if (mainTextBlockId) {
           const state = getState()
           const existingMainTextBlock = state.messageBlocks.entities[mainTextBlockId]
           if (existingMainTextBlock && existingMainTextBlock.type === MessageBlockType.MAIN_TEXT) {
             const currentRefs = existingMainTextBlock.citationReferences || []
             if (!currentRefs.some((ref) => ref.citationBlockId === citationBlockId)) {
-                const mainTextChanges = { citationReferences: [...currentRefs, { citationBlockId }] }
+              const mainTextChanges = {
+                citationReferences: [
+                  ...currentRefs,
+                  { citationBlockId, citationBlockSource: llmWebSearchResult.source }
+                ]
+              }
               dispatch(updateOneBlock({ id: mainTextBlockId, changes: mainTextChanges }))
               saveUpdatedBlockToDB(mainTextBlockId, assistantMsgId, topicId, getState)
             }
           }
         }
-        }
       },
       onImageCreated: () => {
         const imageBlock = createImageBlock(assistantMsgId, {

src/renderer/src/types/newMessage.ts

@@ -11,7 +11,8 @@ import type {
   Model,
   Topic,
   Usage,
-  WebSearchResponse
+  WebSearchResponse,
+  WebSearchSource
 } from '.'
 
 // MessageBlock 类型枚举 - 根据实际API返回特性优化
@@ -63,6 +64,7 @@ export interface MainTextMessageBlock extends BaseMessageBlock {
   // Citation references
   citationReferences?: {
     citationBlockId?: string
+    citationBlockSource?: WebSearchSource
   }[]
 }
 

src/renderer/src/utils/__tests__/linkConverter.test.ts

@@ -99,12 +99,6 @@ describe('linkConverter', () => {
       expect(result).toBe('这里有链接 [<sup>1</sup>](https://example.com)')
     })
 
-    it('should preserve non-domain link text', () => {
-      const input = '点击[这里](https://example.com)查看更多'
-      const result = convertLinks(input, true)
-      expect(result).toBe('点击这里[<sup>1</sup>](https://example.com)查看更多')
-    })
-
     it('should use the same counter for duplicate URLs', () => {
       const input =
         '第一个链接 [example.com](https://example.com) 和第二个相同链接 [subdomain.example.com](https://example.com)'
@@ -113,24 +107,6 @@ describe('linkConverter', () => {
         '第一个链接 [<sup>1</sup>](https://example.com) 和第二个相同链接 [<sup>1</sup>](https://example.com)'
       )
     })
-
-    it('should correctly convert links in Zhipu mode', () => {
-      const input = '这里是引用 [ref_1]'
-      const result = convertLinks(input, true, true)
-      expect(result).toBe('这里是引用 [<sup>1</sup>]()')
-    })
-
-    it('should handle incomplete links in chunked input', () => {
-      // 第一个块包含未完成的链接
-      const chunk1 = '这是链接 ['
-      const result1 = convertLinks(chunk1, true)
-      expect(result1).toBe('这是链接 ')
-
-      // 第二个块完成链接
-      const chunk2 = 'example.com](https://example.com)'
-      const result2 = convertLinks(chunk2, false)
-      expect(result2).toBe('[<sup>1</sup>](https://example.com)')
-    })
   })
 
   describe('convertLinksToOpenRouter', () => {

src/renderer/src/utils/fetch.ts

@@ -126,3 +126,20 @@ export async function fetchWebContent(
     }
   }
 }
+
+export async function fetchRedirectUrl(url: string) {
+  try {
+    const response = await fetch(url, {
+      method: 'HEAD',
+      redirect: 'follow',
+      headers: {
+        'User-Agent':
+          'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
+      }
+    })
+    return response.url
+  } catch (e) {
+    console.error(`Failed to fetch redirect url: ${e}`)
+    return url
+  }
+}

src/renderer/src/utils/linkConverter.ts

@@ -113,6 +113,7 @@ export function convertLinksToHunyuan(text: string, webSearch: any[], resetCount
  * Converts Markdown links in the text to numbered links based on the rules:
  * 1. ([host](url)) -> [cnt](url)
  * 2. [host](url) -> [cnt](url)
+ * 3. [any text except url](url)-> any text [cnt](url)
  *
  * @param text The current chunk of text to process
  * @param resetCounter Whether to reset the counter and buffer
@@ -131,7 +132,6 @@ export function convertLinks(text: string, resetCounter = false): string {
   // Find the safe point - the position after which we might have incomplete patterns
   let safePoint = buffer.length
 
-
   // Check for potentially incomplete patterns from the end
   for (let i = buffer.length - 1; i >= 0; i--) {
     if (buffer[i] === '(') {
@@ -198,6 +198,7 @@ export function convertLinks(text: string, resetCounter = false): string {
 
       if (match) {
         // Found complete regular link
+        const linkText = match[1]
         const url = match[2]
 
         // Check if this URL has been seen before
@@ -209,7 +210,13 @@ export function convertLinks(text: string, resetCounter = false): string {
           urlToCounterMap.set(url, counter)
         }
 
+        // Rule 3: If the link text is not a URL/host, keep the text and add the numbered link
+        if (!isHost(linkText)) {
+          result += `${linkText} [<sup>${counter}</sup>](${url})`
+        } else {
+          // Rule 2: If the link text is a URL/host, replace with numbered link
           result += `[<sup>${counter}</sup>](${url})`
+        }
 
         position += match[0].length
         continue
@@ -317,7 +324,7 @@ export function extractUrlsFromMarkdown(text: string): string[] {
 
   // 匹配所有Markdown链接格式
   const linkPattern = /\[(?:[^[\]]*)\]\(([^()]+)\)/g
-  let match
+  let match: RegExpExecArray | null
 
   while ((match = linkPattern.exec(text)) !== null) {
     const url = match[1].trim()