From 22af6af9c706752edc740f83f04adc3d98f91fae Mon Sep 17 00:00:00 2001
From: "Apple\\Apple" <zeraturing@foxmail.com>
Date: Tue, 27 May 2025 02:07:42 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=9B=A0=EF=B8=8F=20fix(chat):=20enhance=20?=
 =?UTF-8?q?message=20generation=20stop=20behavior?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This commit improves the handling of message generation interruption in the
playground chat interface, ensuring both content and thinking process are
properly preserved.

Key changes:
- Preserve existing content when stopping message generation
- Handle both direct reasoningContent and <think> tag formats
- Extract and merge thinking process from unclosed <think> tags
- Maintain consistent thinking chain format with separators
- Auto-collapse reasoning panel after stopping for cleaner UI

Technical details:
- Modified onStopGenerator to properly handle SSE connection closure
- Added regex pattern to extract thinking content from <think> tags
- Implemented proper state management for incomplete messages
- Ensured all content types are preserved in their respective fields

This fix resolves the issue where thinking chain content would be lost
when stopping message generation mid-stream.
---
 web/src/pages/Playground/Playground.js | 55 +++++++++++++++++++++++++-
 1 file changed, 54 insertions(+), 1 deletion(-)

diff --git a/web/src/pages/Playground/Playground.js b/web/src/pages/Playground/Playground.js
index c5e695e3..3021e106 100644
--- a/web/src/pages/Playground/Playground.js
+++ b/web/src/pages/Playground/Playground.js
@@ -1,4 +1,4 @@
-import React, { useCallback, useContext, useEffect, useState } from 'react';
+import React, { useCallback, useContext, useEffect, useState, useRef } from 'react';
 import { useSearchParams } from 'react-router-dom';
 import { UserContext } from '../../context/User/index.js';
 import {
@@ -100,6 +100,7 @@ const Playground = () => {
   const [groups, setGroups] = useState([]);
   const [showSettings, setShowSettings] = useState(true);
   const [styleState, styleDispatch] = useContext(StyleContext);
+  const sseSourceRef = useRef(null);
 
   const handleInputChange = (name, value) => {
     setInputs((inputs) => ({ ...inputs, [name]: value }));
@@ -210,9 +211,13 @@ const Playground = () => {
       payload: JSON.stringify(payload),
     });
 
+    // 保存 source 引用以便后续停止生成
+    sseSourceRef.current = source;
+
     source.addEventListener('message', (e) => {
       if (e.data === '[DONE]') {
         source.close();
+        sseSourceRef.current = null;
         completeMessage();
         return;
       }
@@ -240,6 +245,7 @@ const Playground = () => {
       const errorMessage = e.data || t('请求发生错误');
       streamMessageUpdate(errorMessage, 'content');
       completeMessage('error');
+      sseSourceRef.current = null;
       source.close();
     });
 
@@ -352,6 +358,51 @@ const Playground = () => {
     });
   }, [setMessage]);
 
+  const onStopGenerator = useCallback(() => {
+    if (sseSourceRef.current) {
+      sseSourceRef.current.close();
+      sseSourceRef.current = null;
+      setMessage((prevMessage) => {
+        const lastMessage = prevMessage[prevMessage.length - 1];
+        if (lastMessage.status === 'loading' || lastMessage.status === 'incomplete') {
+          let content = lastMessage.content || '';
+          let reasoningContent = lastMessage.reasoningContent || '';
+
+          // 处理 <think> 标签格式的思维链
+          if (content.includes('<think>')) {
+            const thinkTagRegex = /<think>([\s\S]*?)(?:<\/think>|$)/g;
+            let thoughts = [];
+            let replyParts = [];
+            let lastIndex = 0;
+            let match;
+
+            while ((match = thinkTagRegex.exec(content)) !== null) {
+              replyParts.push(content.substring(lastIndex, match.index));
+              thoughts.push(match[1]);
+              lastIndex = match.index + match[0].length;
+            }
+            replyParts.push(content.substring(lastIndex));
+
+            // 更新内容和思维链
+            content = replyParts.join('').trim();
+            if (thoughts.length > 0) {
+              reasoningContent = thoughts.join('\n\n---\n\n');
+            }
+          }
+
+          return [...prevMessage.slice(0, -1), {
+            ...lastMessage,
+            status: 'complete',
+            reasoningContent: reasoningContent,
+            content: content,
+            isReasoningExpanded: false  // 停止时折叠思维链面板
+          }];
+        }
+        return prevMessage;
+      });
+    }
+  }, [setMessage]);
+
   const SettingsToggle = () => {
     if (!styleState.isMobile) return null;
     return (
@@ -664,6 +715,8 @@ const Playground = () => {
             chats={message}
             onMessageSend={onMessageSend}
             showClearContext
+            showStopGenerate
+            onStopGenerator={onStopGenerator}
             onClear={() => {
               setMessage([]);
             }}