From 677110ff9162254d0fe222e8a89a2bac1132fb3f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=9E=BF=E4=BD=90=E9=B9=8F?= <806072719@qq.com>
Date: Tue, 23 Dec 2025 16:43:12 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E7=9F=A5=E8=AF=86=E5=BA=93?=
 =?UTF-8?q?=E5=92=8C=E8=81=94=E7=BD=91=E6=90=9C=E7=B4=A2=E4=BB=85=E5=88=9D?=
 =?UTF-8?q?=E5=A7=8B=E5=8C=96=E4=B8=80=E6=AC=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../chat/AiChatMessageServiceImpl.java        | 30 +++++++++++--------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/yudao-module-ai/yudao-module-ai-server/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java b/yudao-module-ai/yudao-module-ai-server/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
index f29a5c3f9..bdd2ba2a1 100644
--- a/yudao-module-ai/yudao-module-ai-server/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
+++ b/yudao-module-ai/yudao-module-ai-server/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
@@ -59,6 +59,8 @@ import reactor.core.publisher.Flux;
 
 import java.time.LocalDateTime;
 import java.util.*;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
@@ -231,20 +233,24 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
         // 4.3 流式返回
         StringBuffer contentBuffer = new StringBuffer();
         StringBuffer reasoningContentBuffer = new StringBuffer();
+
+        // 防止执行多次知识库和联网搜索
+        AtomicBoolean firstExecuteFlag = new AtomicBoolean(true);
+        AtomicReference<List<AiChatMessageRespVO.KnowledgeSegment>> cacheSegments = new AtomicReference<>();
+        AtomicReference<List<AiWebSearchResponse.WebPage>> cacheWebSearchPages = new AtomicReference<>();
         return streamResponse.map(chunk -> {
             // 仅首次：返回知识库、联网搜索
-            List<AiChatMessageRespVO.KnowledgeSegment> segments = null;
-            List<AiWebSearchResponse.WebPage> webSearchPages = null;
             if (StrUtil.isEmpty(contentBuffer)) {
-                Map<Long, AiKnowledgeDocumentDO> documentMap = TenantUtils.executeIgnore(() ->
-                        knowledgeDocumentService.getKnowledgeDocumentMap(
-                                convertSet(knowledgeSegments, AiKnowledgeSegmentSearchRespBO::getDocumentId)));
-                segments = BeanUtils.toBean(knowledgeSegments, AiChatMessageRespVO.KnowledgeSegment.class, segment ->  {
-                    AiKnowledgeDocumentDO document = documentMap.get(segment.getDocumentId());
-                    segment.setDocumentName(document != null ? document.getName() : null);
-                });
-                if (webSearchResponse != null) {
-                    webSearchPages = webSearchResponse.getLists();
+                if (firstExecuteFlag.compareAndSet(true, false)) { // CAS 操作，确保仅执行一次
+                    Map<Long, AiKnowledgeDocumentDO> documentMap = TenantUtils.executeIgnore(() -> knowledgeDocumentService.getKnowledgeDocumentMap(
+                            convertSet(knowledgeSegments, AiKnowledgeSegmentSearchRespBO::getDocumentId)));
+                    cacheSegments.set(BeanUtils.toBean(knowledgeSegments, AiChatMessageRespVO.KnowledgeSegment.class, segment -> {
+                        AiKnowledgeDocumentDO document = documentMap.get(segment.getDocumentId());
+                        segment.setDocumentName(document != null ? document.getName() : null);
+                    }));
+                    if (webSearchResponse != null) {
+                        cacheWebSearchPages.set(webSearchResponse.getLists());
+                    }
                 }
             }
             // 响应结果
@@ -261,7 +267,7 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
                     .setReceive(BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class)
                             .setContent(StrUtil.nullToDefault(newContent, "")) // 避免 null 的 情况
                             .setReasoningContent(StrUtil.nullToDefault(newReasoningContent, "")) // 避免 null 的 情况
-                            .setSegments(segments).setWebSearchPages(webSearchPages))); // 知识库 + 联网搜索
+                            .setSegments(cacheSegments.get()).setWebSearchPages(cacheWebSearchPages.get()))); // 知识库 + 联网搜索
         }).doOnComplete(() -> {
             // 忽略租户，因为 Flux 异步无法透传租户
             TenantUtils.executeIgnore(() -> chatMessageMapper.updateById(