From 4f2f7003c84da955974adce0685b9859124d25cd Mon Sep 17 00:00:00 2001 From: Syngnat Date: Sun, 28 Jun 2026 16:03:22 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20perf(ai-chat):=20=E9=99=8D?= =?UTF-8?q?=E4=BD=8E=E6=B5=81=E5=BC=8F=E6=80=9D=E8=80=83=E8=BE=93=E5=87=BA?= =?UTF-8?q?=E6=B8=B2=E6=9F=93=E5=BC=80=E9=94=80?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 合并 AI 流式 token 刷新，减少高频状态写入 - 避免纯流式更新重排会话列表，并收窄当前会话订阅 - 补充 thinking 合并刷新和会话列表稳定性回归测试 --- .../messageBubble/AIMessageMarkdown.test.tsx | 10 +++ .../components/ai/useAIChatSessionState.ts | 9 +-- .../ai/useAIChatStreamSubscription.test.tsx | 42 +++++++++++++ .../ai/useAIChatStreamSubscription.ts | 63 +++++++++++++++++-- frontend/src/store.test.ts | 49 +++++++++++++++ frontend/src/store.ts | 21 ++++++- 6 files changed, 183 insertions(+), 11 deletions(-) diff --git a/frontend/src/components/ai/messageBubble/AIMessageMarkdown.test.tsx b/frontend/src/components/ai/messageBubble/AIMessageMarkdown.test.tsx index 3b116af..7f71bd2 100644 --- a/frontend/src/components/ai/messageBubble/AIMessageMarkdown.test.tsx +++ b/frontend/src/components/ai/messageBubble/AIMessageMarkdown.test.tsx @@ -10,6 +10,16 @@ import { buildOverlayWorkbenchTheme } from '../../../utils/overlayWorkbenchTheme vi.mock('antd', () => ({ Tooltip: ({ children }: { children: React.ReactNode }) => children, + Modal: Object.assign( + ({ children, open }: { children?: React.ReactNode; open?: boolean }) => (open ?
{children}
: null), + { + info: vi.fn(), + success: vi.fn(), + error: vi.fn(), + warning: vi.fn(), + confirm: vi.fn(), + }, + ), message: { error: vi.fn() }, })); diff --git a/frontend/src/components/ai/useAIChatSessionState.ts b/frontend/src/components/ai/useAIChatSessionState.ts index fa514a5..d3884fe 100644 --- a/frontend/src/components/ai/useAIChatSessionState.ts +++ b/frontend/src/components/ai/useAIChatSessionState.ts @@ -5,6 +5,7 @@ import { loadAISessionsFromBackend, useStore, } from '../../store'; +import type { AIChatMessage } from '../../types'; interface UseAIChatSessionStateOptions { aiActiveSessionId: string | null; @@ -12,13 +13,16 @@ interface UseAIChatSessionStateOptions { createNewAISession: () => void; } +const EMPTY_AI_CHAT_MESSAGES: AIChatMessage[] = []; + export const useAIChatSessionState = ({ aiActiveSessionId, aiPanelVisible, createNewAISession, }: UseAIChatSessionStateOptions) => { - const aiChatHistory = useStore((state) => state.aiChatHistory); const aiChatSessions = useStore((state) => state.aiChatSessions); + const sid = aiActiveSessionId || 'session-fallback'; + const messages = useStore((state) => state.aiChatHistory[sid] || EMPTY_AI_CHAT_MESSAGES); useEffect(() => { if (!aiActiveSessionId) { @@ -26,9 +30,6 @@ export const useAIChatSessionState = ({ } }, [aiActiveSessionId, createNewAISession]); - const sid = aiActiveSessionId || 'session-fallback'; - const messages = aiChatHistory[sid] || []; - const sessionsLoadedRef = useRef(false); useEffect(() => { if (!aiPanelVisible || sessionsLoadedRef.current) { diff --git a/frontend/src/components/ai/useAIChatStreamSubscription.test.tsx b/frontend/src/components/ai/useAIChatStreamSubscription.test.tsx index 0c4cc3c..adb90b6 100644 --- a/frontend/src/components/ai/useAIChatStreamSubscription.test.tsx +++ b/frontend/src/components/ai/useAIChatStreamSubscription.test.tsx @@ -41,6 +41,7 @@ const translate = ( params?: Record, ) => (translatedCopy[key] || key).replace(/\{\{(\w+)\}\}/g, (_match, name) => String(params?.[name] ?? '')); let nextId = 0; +let patchMessageCalls = 0; const emitStreamChunk = async (data: any) => { const handler = runtimeMock.handlers.get(`ai:stream:${SESSION_ID}`); @@ -71,6 +72,7 @@ const patchMessage = ( messageId: string, patch: Parameters['updateAIChatMessage']>[2], ) => { + patchMessageCalls += 1; useStore.setState((state) => { const messages = state.aiChatHistory[sessionId]; if (!messages) { @@ -132,6 +134,7 @@ describe('useAIChatStreamSubscription', () => { beforeEach(() => { nextId = 0; + patchMessageCalls = 0; aiChatStreamMock.mockClear(); generateTitleForSessionMock.mockClear(); runtimeMock.handlers.clear(); @@ -185,6 +188,7 @@ describe('useAIChatStreamSubscription', () => { }); it('keeps streamed chunks in the same assistant message after a parent rerender', async () => { + vi.useFakeTimers(); let renderer: ReactTestRenderer | undefined; await act(async () => { @@ -193,6 +197,9 @@ describe('useAIChatStreamSubscription', () => { await emitStreamChunk({ content: 'Hello' }); await emitStreamChunk({ content: ' world' }); + await act(async () => { + await vi.advanceTimersByTimeAsync(90); + }); const messages = useStore.getState().aiChatHistory[SESSION_ID] || []; const assistantMessages = messages.filter((message) => message.role === 'assistant'); @@ -210,6 +217,41 @@ describe('useAIChatStreamSubscription', () => { }); }); + it('coalesces high-frequency thinking chunks before writing them to the store', async () => { + vi.useFakeTimers(); + let renderer: ReactTestRenderer | undefined; + + await act(async () => { + renderer = create(); + }); + + await emitStreamChunk({ thinking: 'A' }); + const callsAfterScheduling = patchMessageCalls; + await emitStreamChunk({ thinking: 'B' }); + await emitStreamChunk({ thinking: 'C' }); + + expect(patchMessageCalls).toBe(callsAfterScheduling); + expect( + (useStore.getState().aiChatHistory[SESSION_ID] || []).find((message) => message.id === 'assistant-connecting'), + ).not.toHaveProperty('thinking'); + + await act(async () => { + await vi.advanceTimersByTimeAsync(90); + }); + + expect( + (useStore.getState().aiChatHistory[SESSION_ID] || []).find((message) => message.id === 'assistant-connecting'), + ).toMatchObject({ + thinking: 'ABC', + phase: 'thinking', + }); + expect(patchMessageCalls).toBe(callsAfterScheduling + 1); + + await act(async () => { + renderer?.unmount(); + }); + }); + it('resends a localized force-tool-call nudge when the model only describes the next action', async () => { vi.useFakeTimers(); let renderer: ReactTestRenderer | undefined; diff --git a/frontend/src/components/ai/useAIChatStreamSubscription.ts b/frontend/src/components/ai/useAIChatStreamSubscription.ts index b9d8d02..01739a3 100644 --- a/frontend/src/components/ai/useAIChatStreamSubscription.ts +++ b/frontend/src/components/ai/useAIChatStreamSubscription.ts @@ -57,14 +57,18 @@ interface AIChatStreamState { content: string; }; flushPending: boolean; + lastFlushAt: number | null; } +const AI_CHAT_STREAM_FLUSH_INTERVAL_MS = 80; + const createAIChatStreamState = (sid: string): AIChatStreamState => ({ sid, assistantMsgId: '', isFirstCompletion: false, streamBuffer: { thinking: '', reasoningContent: '', content: '' }, flushPending: false, + lastFlushAt: null, }); const resetAIChatStreamProgress = (state: AIChatStreamState) => { @@ -74,6 +78,7 @@ const resetAIChatStreamProgress = (state: AIChatStreamState) => { state.streamBuffer.reasoningContent = ''; state.streamBuffer.content = ''; state.flushPending = false; + state.lastFlushAt = null; }; const translatePanelCopy = ( @@ -131,9 +136,25 @@ export const useAIChatStreamSubscription = ({ // 缓冲高频 token，避免把流式吞吐直接转成同步重绘风暴 const streamBuffer = streamState.streamBuffer; + let flushTimerId: ReturnType | null = null; + let flushFrameId: number | null = null; + + const cancelScheduledFlush = () => { + if (flushTimerId !== null) { + clearTimeout(flushTimerId); + flushTimerId = null; + } + if (flushFrameId !== null && typeof cancelAnimationFrame === 'function') { + cancelAnimationFrame(flushFrameId); + } + flushFrameId = null; + streamState.flushPending = false; + }; const flushStreamBuffer = () => { streamState.flushPending = false; + flushTimerId = null; + flushFrameId = null; if (!streamState.assistantMsgId) return; const current = useStore.getState().aiChatHistory[sid]; const existing = current?.find((message) => message.id === streamState.assistantMsgId); @@ -157,9 +178,43 @@ export const useAIChatStreamSubscription = ({ if (Object.keys(updates).length > 0) { updateAIChatMessage(sid, streamState.assistantMsgId, updates); + streamState.lastFlushAt = Date.now(); } }; + const requestFlushFrame = () => { + if (typeof requestAnimationFrame !== 'function') { + flushStreamBuffer(); + return; + } + + let completedSynchronously = false; + const frameId = requestAnimationFrame(() => { + completedSynchronously = true; + flushFrameId = null; + flushStreamBuffer(); + }); + flushFrameId = completedSynchronously ? null : frameId; + }; + + const scheduleStreamFlush = () => { + if (streamState.flushPending) return; + streamState.flushPending = true; + + const lastFlushAt = streamState.lastFlushAt; + const delay = + lastFlushAt === null + ? 0 + : Math.max(0, AI_CHAT_STREAM_FLUSH_INTERVAL_MS - (Date.now() - lastFlushAt)); + + if (delay > 0) { + flushTimerId = setTimeout(requestFlushFrame, delay); + return; + } + + requestFlushFrame(); + }; + const handler = (data: AIChatStreamChunk) => { if (!streamState.assistantMsgId) { const history = useStore.getState().aiChatHistory[sid] || []; @@ -202,6 +257,7 @@ export const useAIChatStreamSubscription = ({ jvmDiagnosticPlanContext: pendingJVMDiagnosticPlanContextRef.current, }); } + cancelScheduledFlush(); resetAIChatStreamProgress(streamState); setSending(false); return; @@ -275,14 +331,12 @@ export const useAIChatStreamSubscription = ({ } if (streamBuffer.thinking || streamBuffer.reasoningContent || streamBuffer.content) { - if (!streamState.flushPending) { - streamState.flushPending = true; - requestAnimationFrame(flushStreamBuffer); - } + scheduleStreamFlush(); } if (data.done) { if (streamBuffer.thinking || streamBuffer.reasoningContent || streamBuffer.content) { + cancelScheduledFlush(); flushStreamBuffer(); } const doneAssistantId = streamState.assistantMsgId; @@ -380,6 +434,7 @@ export const useAIChatStreamSubscription = ({ EventsOn(eventName, handler); return () => { + cancelScheduledFlush(); EventsOff(eventName); }; }, [ diff --git a/frontend/src/store.test.ts b/frontend/src/store.test.ts index 3297998..4d47ca5 100644 --- a/frontend/src/store.test.ts +++ b/frontend/src/store.test.ts @@ -1117,6 +1117,55 @@ describe('store appearance persistence', () => { } }); + it('keeps streaming-only AI message patches from reordering the session list', async () => { + vi.useFakeTimers(); + try { + const { useStore } = await importStore(); + useStore.setState({ + aiChatSessions: [ + { id: 'session-other', title: 'other', updatedAt: 20 }, + { id: 'session-stream', title: 'stream', updatedAt: 10 }, + ], + aiChatHistory: { + 'session-stream': [ + { + id: 'assistant-1', + role: 'assistant', + phase: 'connecting', + content: '', + timestamp: 1, + loading: true, + }, + ], + }, + }); + + const sessionsBeforeStreamingPatch = useStore.getState().aiChatSessions; + useStore.getState().updateAIChatMessage('session-stream', 'assistant-1', { + thinking: 'planning', + phase: 'thinking', + }); + + expect(useStore.getState().aiChatSessions).toBe(sessionsBeforeStreamingPatch); + expect(useStore.getState().aiChatSessions.map((session) => session.id)).toEqual([ + 'session-other', + 'session-stream', + ]); + + useStore.getState().updateAIChatMessage('session-stream', 'assistant-1', { + loading: false, + phase: 'idle', + }); + + expect(useStore.getState().aiChatSessions.map((session) => session.id)).toEqual([ + 'session-stream', + 'session-other', + ]); + } finally { + vi.useRealTimers(); + } + }); + it('keeps store fallback titles out of production source literals', async () => { const { readFileSync } = await import('node:fs'); const source = readFileSync(new URL('./store.ts', import.meta.url), 'utf8'); diff --git a/frontend/src/store.ts b/frontend/src/store.ts index 60be8f1..38ce15b 100644 --- a/frontend/src/store.ts +++ b/frontend/src/store.ts @@ -1495,6 +1495,23 @@ interface AppState { setAIActiveSessionId: (sessionId: string | null) => void; } +const AI_STREAMING_MESSAGE_UPDATE_KEYS = new Set([ + "content", + "thinking", + "reasoning_content", + "phase", +]); + +const isAIStreamingOnlyMessageUpdate = ( + updates: Partial, +): boolean => { + const updateKeys = Object.keys(updates) as Array; + return ( + updateKeys.length > 0 && + updateKeys.every((key) => AI_STREAMING_MESSAGE_UPDATE_KEYS.has(key)) + ); +}; + const sanitizeSqlSnippets = (value: unknown): SqlSnippet[] => { if (!Array.isArray(value)) return DEFAULT_SQL_SNIPPETS; const result: SqlSnippet[] = []; @@ -3455,9 +3472,7 @@ export const useStore = create()( const newMessages = [...messages]; newMessages[idx] = { ...newMessages[idx], ...updates }; const history = { ...state.aiChatHistory, [sessionId]: newMessages }; - const isContentOnlyUpdate = - Object.keys(updates).length === 1 && "content" in updates; - if (!isContentOnlyUpdate) { + if (!isAIStreamingOnlyMessageUpdate(updates)) { let newSessions = [...state.aiChatSessions]; const existingSession = newSessions.find((s) => s.id === sessionId); if (existingSession) {