diff --git a/src/renderer/src/store/thunk/messageThunk.ts b/src/renderer/src/store/thunk/messageThunk.ts
index b652b70463..996044bb24 100644
--- a/src/renderer/src/store/thunk/messageThunk.ts
+++ b/src/renderer/src/store/thunk/messageThunk.ts
@@ -429,6 +429,9 @@ const fetchAndProcessAssistantResponseImpl = async (
         await handleBlockTransition(baseBlock as PlaceholderMessageBlock, MessageBlockType.UNKNOWN)
       },
       onTextChunk: async (text) => {
+        const citationBlockSource = citationBlockId
+          ? (getState().messageBlocks.entities[citationBlockId] as CitationMessageBlock).response?.source
+          : WebSearchSource.WEBSEARCH
         accumulatedContent += text
         if (mainTextBlockId) {
           const blockChanges: Partial<MessageBlock> = {
@@ -441,7 +444,8 @@ const fetchAndProcessAssistantResponseImpl = async (
           const initialChanges: Partial<MessageBlock> = {
             type: MessageBlockType.MAIN_TEXT,
             content: accumulatedContent,
-            status: MessageBlockStatus.STREAMING
+            status: MessageBlockStatus.STREAMING,
+            citationReferences: citationBlockId ? [{ citationBlockId, citationBlockSource }] : []
           }
           mainTextBlockId = initialPlaceholderBlockId
           // 清理占位块
@@ -451,7 +455,8 @@ const fetchAndProcessAssistantResponseImpl = async (
           saveUpdatedBlockToDB(mainTextBlockId, assistantMsgId, topicId, getState)
         } else {
           const newBlock = createMainTextBlock(assistantMsgId, accumulatedContent, {
-            status: MessageBlockStatus.STREAMING
+            status: MessageBlockStatus.STREAMING,
+            citationReferences: citationBlockId ? [{ citationBlockId, citationBlockSource }] : []
           })
           mainTextBlockId = newBlock.id // 立即设置ID，防止竞态条件
           await handleBlockTransition(newBlock, MessageBlockType.MAIN_TEXT)
@@ -459,15 +464,9 @@ const fetchAndProcessAssistantResponseImpl = async (
       },
       onTextComplete: async (finalText) => {
         if (mainTextBlockId) {
-          let citationBlockSource: WebSearchSource | undefined
-          if (citationBlockId) {
-            const citationBlock = getState().messageBlocks.entities[citationBlockId] as CitationMessageBlock
-            citationBlockSource = citationBlock.response?.source
-          }
           const changes = {
             content: finalText,
-            status: MessageBlockStatus.SUCCESS,
-            citationReferences: citationBlockSource ? [{ citationBlockId, citationBlockSource }] : []
+            status: MessageBlockStatus.SUCCESS
           }
           cancelThrottledBlockUpdate(mainTextBlockId)
           dispatch(updateOneBlock({ id: mainTextBlockId, changes }))