From 792b98bfbf1a24d1e1e0893c9a36b6cf7328d674 Mon Sep 17 00:00:00 2001 From: lewis617 Date: Tue, 30 Jun 2026 18:06:33 +0800 Subject: [PATCH] feat: replace full-message polling with incremental streaming updates Upgrade SDK to v0.18.1 and use new incremental callbacks (onAssistantMessageAdded, onAssistantContentUpdated, onAssistantReasoningUpdated, onToolBlockUpdated) instead of onMessagesChange for streaming updates. - Extract chatReducer to standalone module with 4 new incremental actions - Add 16ms throttle (~60fps) for content/reasoning, stage='end' fires immediately - React.memo on Message component with custom comparator - MessageList key changed to message.id for stable identity - onMessagesChange retained only for clear/restore/rewind scenarios - Add 23 unit tests for chatReducer incremental actions - Fix tests/tsconfig.json deprecation warnings --- package-lock.json | 8 +- package.json | 2 +- src/chatProvider.ts | 15 +- src/session/chatSession.ts | 108 +++++++- tests/reducers/chatReducer.test.ts | 334 +++++++++++++++++++++++++ tests/tsconfig.json | 4 +- webview/src/components/ChatApp.tsx | 223 ++--------------- webview/src/components/Message.tsx | 10 +- webview/src/components/MessageList.tsx | 2 +- webview/src/reducers/chatReducer.ts | 312 +++++++++++++++++++++++ webview/src/types/index.ts | 12 +- 11 files changed, 811 insertions(+), 219 deletions(-) create mode 100644 tests/reducers/chatReducer.test.ts create mode 100644 webview/src/reducers/chatReducer.ts diff --git a/package-lock.json b/package-lock.json index c667985..2ff0732 100644 --- a/package-lock.json +++ b/package-lock.json @@ -13,7 +13,7 @@ "dompurify": "^3.3.0", "marked": "^9.1.6", "mermaid": "^11.12.2", - "wave-agent-sdk": "^0.17.12" + "wave-agent-sdk": "^0.18.1" }, "devDependencies": { "@playwright/test": "^1.58.2", @@ -12485,9 +12485,9 @@ } }, "node_modules/wave-agent-sdk": { - "version": "0.17.12", - "resolved": "https://registry.npmjs.org/wave-agent-sdk/-/wave-agent-sdk-0.17.12.tgz", - "integrity": "sha512-QRL/jCOHQSmlWRQRQUFXAbTOFMU3dGU8ziGpqtT7oo67rrUWMGVOQMaXxsWPZDBvjYfOEwlP4xiXHvfMje1bFw==", + "version": "0.18.1", + "resolved": "https://registry.npmjs.org/wave-agent-sdk/-/wave-agent-sdk-0.18.1.tgz", + "integrity": "sha512-AJsR9Y0VMxkT0HWCuyW/lt6ovq8Zu7/T4VPgZLqVAVVBKM/NTD0fGcqYZCkGvNyFdP6c5Wt1c0EqIE/zDXGSGA==", "hasInstallScript": true, "license": "MIT", "dependencies": { diff --git a/package.json b/package.json index 74fdab1..4a13d3d 100644 --- a/package.json +++ b/package.json @@ -132,6 +132,6 @@ "dompurify": "^3.3.0", "marked": "^9.1.6", "mermaid": "^11.12.2", - "wave-agent-sdk": "^0.17.12" + "wave-agent-sdk": "^0.18.1" } } diff --git a/src/chatProvider.ts b/src/chatProvider.ts index d67ce14..ce3e830 100644 --- a/src/chatProvider.ts +++ b/src/chatProvider.ts @@ -150,7 +150,7 @@ export class ChatProvider implements vscode.WebviewViewProvider { this.webviewManager.postMessage({ command: 'updateTasks', tasks }, viewType, windowId); }, onSessionIdChange: (sessionId) => { - this.handleSessionIdChange(sessionId, viewType, windowId).catch(err => + this.handleSessionIdChange(sessionId, viewType, windowId).catch(err => console.error(`Error handling session ID change for ${viewType}:`, err) ); }, @@ -179,6 +179,19 @@ export class ChatProvider implements vscode.WebviewViewProvider { }, onMcpServersChange: (servers) => { this.webviewManager.postMessage({ command: 'mcpServersUpdate', servers }, viewType, windowId); + }, + // Incremental update callbacks for streaming optimization + onAssistantMessageAdded: (message) => { + this.webviewManager.postMessage({ command: 'appendMessage', message }, viewType, windowId); + }, + onStreamingContentUpdate: (params) => { + this.webviewManager.postMessage({ command: 'updateStreamingContent', ...params }, viewType, windowId); + }, + onStreamingReasoningUpdate: (params) => { + this.webviewManager.postMessage({ command: 'updateStreamingReasoning', ...params }, viewType, windowId); + }, + onToolBlockUpdate: (params) => { + this.webviewManager.postMessage({ command: 'updateToolBlock', params }, viewType, windowId); } }); } diff --git a/src/session/chatSession.ts b/src/session/chatSession.ts index 59539cd..a2039a2 100644 --- a/src/session/chatSession.ts +++ b/src/session/chatSession.ts @@ -1,5 +1,5 @@ import * as vscode from 'vscode'; -import { Agent, Message, PermissionDecision, ToolPermissionContext, AgentCallbacks, PermissionMode, Task, PromptHistoryManager, TextBlock, QueuedMessage, McpServerStatus } from 'wave-agent-sdk'; +import { Agent, Message, PermissionDecision, ToolPermissionContext, AgentCallbacks, PermissionMode, Task, PromptHistoryManager, TextBlock, QueuedMessage, McpServerStatus, ToolBlockUpdateCallbackParams } from 'wave-agent-sdk'; import { ConfigurationData } from '../services/configurationService'; import { VscodeLspAdapter } from '../services/lspAdapter'; @@ -14,6 +14,11 @@ export interface ChatSessionCallbacks { onToolPermissionRequest: (context: ToolPermissionContext) => Promise; onError: (error: any) => void; onMcpServersChange?: (servers: McpServerStatus[]) => void; + // Incremental update callbacks for streaming optimization + onAssistantMessageAdded?: (message: Message) => void; + onStreamingContentUpdate?: (params: { messageId: string; accumulated: string; stage: 'streaming' | 'end' }) => void; + onStreamingReasoningUpdate?: (params: { messageId: string; accumulated: string; stage: 'streaming' | 'end' }) => void; + onToolBlockUpdate?: (params: ToolBlockUpdateCallbackParams) => void; } export class ChatSession { @@ -40,6 +45,12 @@ export class ChatSession { private pendingUpdate: boolean = false; private forceNextUpdateImmediate: boolean = false; + // Throttled incremental update fields + private streamingContentUpdateTimer: NodeJS.Timeout | undefined; + private pendingStreamingContentUpdate: { messageId: string; accumulated: string; stage: 'streaming' | 'end' } | undefined; + private streamingReasoningUpdateTimer: NodeJS.Timeout | undefined; + private pendingStreamingReasoningUpdate: { messageId: string; accumulated: string; stage: 'streaming' | 'end' } | undefined; + constructor( public readonly viewType: 'sidebar' | 'tab' | 'window', public readonly windowId: string | undefined, @@ -68,8 +79,35 @@ export class ChatSession { const isBaseURLValid = !!config.baseURL || !!process.env.WAVE_BASE_URL || !!config.serverUrl || !!process.env.WAVE_SERVER_URL; const agentCallbacks: AgentCallbacks = { + // Keep onMessagesChange for clear/restore/rewind, but not triggered by SDK callbacks during streaming onMessagesChange: (messages: Message[]) => { - this.throttledUpdateChatMessages(messages); + this.messages = messages; + // Only trigger full update for non-streaming scenarios (clear/restore/rewind) + // During streaming, incremental callbacks handle updates + }, + onUserMessageAdded: () => { + // Find the newly added user message from agent.messages (last user message) + const userMessages = this.agent?.messages.filter(m => m.role === 'user') || []; + const newUserMessage = userMessages[userMessages.length - 1]; + if (newUserMessage) { + this.callbacks.onAssistantMessageAdded?.(newUserMessage); + } + }, + onAssistantMessageAdded: (messageId: string) => { + // Find the newly added message from agent.messages + const newMessage = this.agent?.messages.find(m => m.id === messageId); + if (newMessage) { + this.callbacks.onAssistantMessageAdded?.(newMessage); + } + }, + onAssistantContentUpdated: (params) => { + this.throttledStreamingContentUpdate(params.messageId, params.accumulated, params.stage); + }, + onAssistantReasoningUpdated: (params) => { + this.throttledStreamingReasoningUpdate(params.messageId, params.accumulated, params.stage); + }, + onToolBlockUpdated: (params) => { + this.callbacks.onToolBlockUpdate?.(params); }, onTasksChange: (tasks: Task[]) => { this.tasks = tasks; @@ -327,6 +365,62 @@ export class ChatSession { } } + private throttledStreamingContentUpdate(messageId: string, accumulated: string, stage: 'streaming' | 'end') { + // If stage is 'end', fire immediately to ensure finalization is not delayed + if (stage === 'end') { + if (this.streamingContentUpdateTimer) { + clearTimeout(this.streamingContentUpdateTimer); + this.streamingContentUpdateTimer = undefined; + } + this.pendingStreamingContentUpdate = undefined; + this.callbacks.onStreamingContentUpdate?.({ messageId, accumulated, stage }); + return; + } + + this.pendingStreamingContentUpdate = { messageId, accumulated, stage }; + + // leading edge: first call fires immediately + if (!this.streamingContentUpdateTimer) { + this.callbacks.onStreamingContentUpdate?.(this.pendingStreamingContentUpdate); + // trailing edge: fire the last update after 16ms cooldown (~60fps) + this.streamingContentUpdateTimer = setTimeout(() => { + if (this.pendingStreamingContentUpdate) { + this.callbacks.onStreamingContentUpdate?.(this.pendingStreamingContentUpdate); + this.pendingStreamingContentUpdate = undefined; + } + this.streamingContentUpdateTimer = undefined; + }, 16); + } + } + + private throttledStreamingReasoningUpdate(messageId: string, accumulated: string, stage: 'streaming' | 'end') { + // If stage is 'end', fire immediately to ensure finalization is not delayed + if (stage === 'end') { + if (this.streamingReasoningUpdateTimer) { + clearTimeout(this.streamingReasoningUpdateTimer); + this.streamingReasoningUpdateTimer = undefined; + } + this.pendingStreamingReasoningUpdate = undefined; + this.callbacks.onStreamingReasoningUpdate?.({ messageId, accumulated, stage }); + return; + } + + this.pendingStreamingReasoningUpdate = { messageId, accumulated, stage }; + + // leading edge: first call fires immediately + if (!this.streamingReasoningUpdateTimer) { + this.callbacks.onStreamingReasoningUpdate?.(this.pendingStreamingReasoningUpdate); + // trailing edge: fire the last update after 16ms cooldown (~60fps) + this.streamingReasoningUpdateTimer = setTimeout(() => { + if (this.pendingStreamingReasoningUpdate) { + this.callbacks.onStreamingReasoningUpdate?.(this.pendingStreamingReasoningUpdate); + this.pendingStreamingReasoningUpdate = undefined; + } + this.streamingReasoningUpdateTimer = undefined; + }, 16); + } + } + public async setPermissionMode(mode: PermissionMode) { if (this.agent) { await this.agent.setPermissionMode(mode); @@ -362,6 +456,14 @@ export class ChatSession { clearTimeout(this.updateTimer); this.updateTimer = undefined; } + if (this.streamingContentUpdateTimer) { + clearTimeout(this.streamingContentUpdateTimer); + this.streamingContentUpdateTimer = undefined; + } + if (this.streamingReasoningUpdateTimer) { + clearTimeout(this.streamingReasoningUpdateTimer); + this.streamingReasoningUpdateTimer = undefined; + } if (this.agent) { try { @@ -380,6 +482,8 @@ export class ChatSession { this.isStreaming = false; this.isCommandRunning = false; this.pendingUpdate = false; + this.pendingStreamingContentUpdate = undefined; + this.pendingStreamingReasoningUpdate = undefined; this.messageQueue = []; } diff --git a/tests/reducers/chatReducer.test.ts b/tests/reducers/chatReducer.test.ts new file mode 100644 index 0000000..73d12ca --- /dev/null +++ b/tests/reducers/chatReducer.test.ts @@ -0,0 +1,334 @@ +import { describe, it, expect } from 'vitest'; +import { chatReducer, initialState } from '../../webview/src/reducers/chatReducer'; +import type { Message, TextBlock, ToolBlock, ReasoningBlock } from '../../webview/src/types'; + +describe('chatReducer', () => { + describe('APPEND_MESSAGE', () => { + it('should append a new message to the end of the list', () => { + const existingMessage: Message = { + id: 'msg-1', + role: 'user', + timestamp: '0', + blocks: [] + }; + const state = { ...initialState, messages: [existingMessage] }; + + const newMessage: Message = { + id: 'msg-2', + role: 'assistant', + timestamp: '0', + blocks: [] + }; + + const newState = chatReducer(state, { type: 'APPEND_MESSAGE', payload: newMessage }); + + expect(newState.messages).toHaveLength(2); + expect(newState.messages[0]).toBe(existingMessage); + expect(newState.messages[1]).toBe(newMessage); + }); + + it('should append to empty message list', () => { + const state = { ...initialState, messages: [] }; + const newMessage: Message = { + id: 'msg-1', + role: 'assistant', + timestamp: '0', + blocks: [] + }; + + const newState = chatReducer(state, { type: 'APPEND_MESSAGE', payload: newMessage }); + + expect(newState.messages).toHaveLength(1); + expect(newState.messages[0]).toBe(newMessage); + }); + }); + + describe('UPDATE_STREAMING_CONTENT', () => { + it('should update existing text block content and stage', () => { + const textBlock: TextBlock = { + type: 'text', + content: 'Hello', + stage: 'streaming' + }; + const message: Message = { + id: 'msg-1', + role: 'assistant', + timestamp: '0', + blocks: [textBlock] + }; + const state = { ...initialState, messages: [message] }; + + const newState = chatReducer(state, { + type: 'UPDATE_STREAMING_CONTENT', + payload: { messageId: 'msg-1', accumulated: 'Hello World', stage: 'streaming' } + }); + + expect(newState.messages).toHaveLength(1); + expect(newState.messages[0].blocks).toHaveLength(1); + const updatedBlock = newState.messages[0].blocks[0] as TextBlock; + expect(updatedBlock.content).toBe('Hello World'); + expect(updatedBlock.stage).toBe('streaming'); + }); + + it('should update stage to end', () => { + const textBlock: TextBlock = { + type: 'text', + content: 'Hello', + stage: 'streaming' + }; + const message: Message = { + id: 'msg-1', + role: 'assistant', + timestamp: '0', + blocks: [textBlock] + }; + const state = { ...initialState, messages: [message] }; + + const newState = chatReducer(state, { + type: 'UPDATE_STREAMING_CONTENT', + payload: { messageId: 'msg-1', accumulated: 'Hello', stage: 'end' } + }); + + const updatedBlock = newState.messages[0].blocks[0] as TextBlock; + expect(updatedBlock.stage).toBe('end'); + }); + + it('should append text block if none exists', () => { + const message: Message = { + id: 'msg-1', + role: 'assistant', + timestamp: '0', + blocks: [] + }; + const state = { ...initialState, messages: [message] }; + + const newState = chatReducer(state, { + type: 'UPDATE_STREAMING_CONTENT', + payload: { messageId: 'msg-1', accumulated: 'Hello', stage: 'streaming' } + }); + + expect(newState.messages[0].blocks).toHaveLength(1); + const newBlock = newState.messages[0].blocks[0] as TextBlock; + expect(newBlock.type).toBe('text'); + expect(newBlock.content).toBe('Hello'); + expect(newBlock.stage).toBe('streaming'); + }); + + it('should return original state if messageId not found', () => { + const state = { ...initialState, messages: [] }; + + const newState = chatReducer(state, { + type: 'UPDATE_STREAMING_CONTENT', + payload: { messageId: 'non-existent', accumulated: 'Hello', stage: 'streaming' } + }); + + expect(newState).toBe(state); + }); + }); + + describe('UPDATE_STREAMING_REASONING', () => { + it('should update existing reasoning block content and stage', () => { + const reasoningBlock: ReasoningBlock = { + type: 'reasoning', + content: 'Thinking...', + stage: 'streaming' + }; + const message: Message = { + id: 'msg-1', + role: 'assistant', + timestamp: '0', + blocks: [reasoningBlock] + }; + const state = { ...initialState, messages: [message] }; + + const newState = chatReducer(state, { + type: 'UPDATE_STREAMING_REASONING', + payload: { messageId: 'msg-1', accumulated: 'Thinking... about this', stage: 'streaming' } + }); + + expect(newState.messages[0].blocks).toHaveLength(1); + const updatedBlock = newState.messages[0].blocks[0] as ReasoningBlock; + expect(updatedBlock.content).toBe('Thinking... about this'); + expect(updatedBlock.stage).toBe('streaming'); + }); + + it('should append reasoning block if none exists', () => { + const message: Message = { + id: 'msg-1', + role: 'assistant', + timestamp: '0', + blocks: [] + }; + const state = { ...initialState, messages: [message] }; + + const newState = chatReducer(state, { + type: 'UPDATE_STREAMING_REASONING', + payload: { messageId: 'msg-1', accumulated: 'Thinking...', stage: 'streaming' } + }); + + expect(newState.messages[0].blocks).toHaveLength(1); + const newBlock = newState.messages[0].blocks[0] as ReasoningBlock; + expect(newBlock.type).toBe('reasoning'); + expect(newBlock.content).toBe('Thinking...'); + }); + + it('should return original state if messageId not found', () => { + const state = { ...initialState, messages: [] }; + + const newState = chatReducer(state, { + type: 'UPDATE_STREAMING_REASONING', + payload: { messageId: 'non-existent', accumulated: 'Thinking...', stage: 'streaming' } + }); + + expect(newState).toBe(state); + }); + }); + + describe('UPDATE_TOOL_BLOCK', () => { + it('should update existing tool block', () => { + const toolBlock: ToolBlock = { + type: 'tool', + id: 'tool-1', + name: 'Bash', + stage: 'running', + parameters: '{"command":"ls"}', + result: '', + success: false + }; + const message: Message = { + id: 'msg-1', + role: 'assistant', + timestamp: '0', + blocks: [toolBlock] + }; + const state = { ...initialState, messages: [message] }; + + const newState = chatReducer(state, { + type: 'UPDATE_TOOL_BLOCK', + payload: { + messageId: 'msg-1', + id: 'tool-1', + name: 'Bash', + stage: 'end', + result: 'file1.txt\nfile2.txt', + success: true, + parameters: '{"command":"ls"}' + } + }); + + expect(newState.messages[0].blocks).toHaveLength(1); + const updatedBlock = newState.messages[0].blocks[0] as ToolBlock; + expect(updatedBlock.stage).toBe('end'); + expect(updatedBlock.result).toBe('file1.txt\nfile2.txt'); + expect(updatedBlock.success).toBe(true); + }); + + it('should return original state if messageId not found', () => { + const state = { ...initialState, messages: [] }; + + const newState = chatReducer(state, { + type: 'UPDATE_TOOL_BLOCK', + payload: { + messageId: 'non-existent', + id: 'tool-1', + name: 'Bash', + stage: 'end', + result: '', + success: false, + parameters: '{}' + } + }); + + expect(newState).toBe(state); + }); + + it('should return original state if tool block not found', () => { + const toolBlock: ToolBlock = { + type: 'tool', + id: 'tool-1', + name: 'Bash', + stage: 'running', + parameters: '{"command":"ls"}', + result: '', + success: false + }; + const message: Message = { + id: 'msg-1', + role: 'assistant', + timestamp: '0', + blocks: [toolBlock] + }; + const state = { ...initialState, messages: [message] }; + + const newState = chatReducer(state, { + type: 'UPDATE_TOOL_BLOCK', + payload: { + messageId: 'msg-1', + id: 'non-existent-tool', + name: 'Bash', + stage: 'end', + result: '', + success: false, + parameters: '{}' + } + }); + + expect(newState).toBe(state); + }); + }); + + describe('SET_MESSAGES', () => { + it('should replace entire message list', () => { + const oldMessage: Message = { id: 'old', role: 'user', timestamp: '0', blocks: [] }; + const state = { ...initialState, messages: [oldMessage] }; + const newMessages: Message[] = [ + { id: 'new-1', role: 'user', timestamp: '0', blocks: [] }, + { id: 'new-2', role: 'assistant', timestamp: '0', blocks: [] } + ]; + + const newState = chatReducer(state, { type: 'SET_MESSAGES', payload: newMessages }); + + expect(newState.messages).toHaveLength(2); + expect(newState.messages[0].id).toBe('new-1'); + expect(newState.messages[1].id).toBe('new-2'); + }); + }); + + describe('User message flow', () => { + it('should handle user message followed by assistant message', () => { + const state = { ...initialState, messages: [] }; + + // User sends a message + const userMessage: Message = { + id: 'user-1', + role: 'user', + timestamp: '0', + blocks: [{ type: 'text', content: 'Hello', stage: 'end' }] + }; + let newState = chatReducer(state, { type: 'APPEND_MESSAGE', payload: userMessage }); + expect(newState.messages).toHaveLength(1); + expect(newState.messages[0].role).toBe('user'); + + // Assistant starts responding + const assistantMessage: Message = { + id: 'assistant-1', + role: 'assistant', + timestamp: '0', + blocks: [] + }; + newState = chatReducer(newState, { type: 'APPEND_MESSAGE', payload: assistantMessage }); + expect(newState.messages).toHaveLength(2); + expect(newState.messages[0].role).toBe('user'); + expect(newState.messages[1].role).toBe('assistant'); + + // Assistant content streams in + newState = chatReducer(newState, { + type: 'UPDATE_STREAMING_CONTENT', + payload: { messageId: 'assistant-1', accumulated: 'Hi there', stage: 'streaming' } + }); + expect(newState.messages).toHaveLength(2); + expect(newState.messages[1].blocks).toHaveLength(1); + expect((newState.messages[1].blocks[0] as TextBlock).content).toBe('Hi there'); + }); + }); +}); diff --git a/tests/tsconfig.json b/tests/tsconfig.json index e78fdb4..d4106c0 100644 --- a/tests/tsconfig.json +++ b/tests/tsconfig.json @@ -1,10 +1,10 @@ { "compilerOptions": { - "module": "commonjs", + "module": "ES2020", "target": "ES2020", "esModuleInterop": true, "allowSyntheticDefaultImports": true, "skipLibCheck": true, - "moduleResolution": "node" + "moduleResolution": "bundler" } } diff --git a/webview/src/components/ChatApp.tsx b/webview/src/components/ChatApp.tsx index 9da2c42..f1c12d2 100644 --- a/webview/src/components/ChatApp.tsx +++ b/webview/src/components/ChatApp.tsx @@ -13,213 +13,11 @@ import StatusDialog from './StatusDialog'; import LoginDialog from './LoginDialog'; import type { ChatAppProps, - ChatState, - ChatAction, WebviewMessage, - Message, - Task } from '../types'; +import { chatReducer, initialState } from '../reducers/chatReducer'; import '../styles/ChatApp.css'; -const initialState: ChatState = { - messages: [], - tasks: [], - isTaskListVisible: false, - isTaskListCollapsed: false, - isQueueCollapsed: false, - isStreaming: false, - isCommandRunning: false, - inputDisabled: false, - shouldClearInput: false, - sessions: [], - currentSession: undefined, - sessionsLoading: false, - pendingConfirmations: [], - queuedMessages: [], - // Dialog state - activeDialog: null, - configurationData: undefined, - configurationLoading: false, - configurationError: undefined, - configuredModels: [], - currentModel: '', - currentFastModel: '', - // Permission mode state - permissionMode: 'default', - // Attached images state - attachedImages: [] -}; - -function chatReducer(state: ChatState, action: ChatAction): ChatState { - switch (action.type) { - case 'SET_MESSAGES': - return { - ...state, - messages: action.payload - }; - case 'SET_TASKS': - return { - ...state, - tasks: action.payload, - // Show task list if there are tasks - isTaskListVisible: action.payload.length > 0, - // Auto-expand task list when tasks are first created - isTaskListCollapsed: state.tasks.length === 0 && action.payload.length > 0 ? false : state.isTaskListCollapsed - }; - case 'TOGGLE_TASK_LIST_COLLAPSE': - return { - ...state, - isTaskListCollapsed: !state.isTaskListCollapsed - }; - case 'SET_TASK_LIST_COLLAPSED': - return { - ...state, - isTaskListCollapsed: action.payload - }; - case 'TOGGLE_QUEUE_COLLAPSE': - return { - ...state, - isQueueCollapsed: !state.isQueueCollapsed - }; - case 'START_STREAMING': - return { - ...state, - isStreaming: true - }; - case 'END_STREAMING': - return { - ...state, - isStreaming: false - }; - case 'SET_INPUT_DISABLED': - return { - ...state, - inputDisabled: action.payload - }; - case 'INPUT_CLEARED': - return { - ...state, - shouldClearInput: false - }; - case 'SET_SESSIONS': - return { - ...state, - sessions: action.payload, - sessionsLoading: false - }; - case 'SET_CURRENT_SESSION': - return { - ...state, - currentSession: action.payload - }; - case 'SET_SESSIONS_LOADING': - return { - ...state, - sessionsLoading: action.payload - }; - case 'SHOW_CONFIRMATION': - return { - ...state, - pendingConfirmations: [...state.pendingConfirmations, action.payload] - }; - case 'HIDE_CONFIRMATION': - return { - ...state, - pendingConfirmations: state.pendingConfirmations.filter(c => c.confirmationId !== action.payload) - }; - case 'SHOW_DIALOG': - return { - ...state, - activeDialog: action.payload.type, - configurationData: action.payload.data ?? state.configurationData, - configurationLoading: false, - configurationError: action.payload.error - }; - case 'HIDE_DIALOG': - return { - ...state, - activeDialog: null, - configurationError: undefined - }; - case 'SET_CONFIGURATION_LOADING': - return { - ...state, - configurationLoading: action.payload - }; - case 'SET_CONFIGURATION_ERROR': - return { - ...state, - configurationError: action.payload, - configurationLoading: false - }; - case 'SET_CONFIGURATION_DATA': - return { - ...state, - configurationData: action.payload, - configurationLoading: false - }; - case 'SET_CONFIGURED_MODELS': - return { - ...state, - configuredModels: action.payload - }; - case 'SET_CURRENT_MODELS': - return { - ...state, - currentModel: action.payload.model, - currentFastModel: action.payload.fastModel, - configurationData: { - ...state.configurationData, - model: action.payload.model || state.configurationData?.model, - fastModel: action.payload.fastModel || state.configurationData?.fastModel - } - }; - case 'SET_INITIAL_STATE': - return { - ...state, - messages: action.payload.messages, - tasks: action.payload.tasks || [], - isTaskListVisible: (action.payload.tasks && action.payload.tasks.length > 0) ? true : false, - isTaskListCollapsed: action.payload.isTaskListCollapsed !== undefined ? action.payload.isTaskListCollapsed : state.isTaskListCollapsed, - isStreaming: action.payload.isStreaming !== undefined ? action.payload.isStreaming : state.isStreaming, - isCommandRunning: action.payload.isCommandRunning !== undefined ? action.payload.isCommandRunning : state.isCommandRunning, - sessions: action.payload.sessions || state.sessions || [], - currentSession: action.payload.currentSession || state.currentSession, - configurationData: action.payload.configurationData || state.configurationData, - pendingConfirmations: action.payload.pendingConfirmations || [], - queuedMessages: action.payload.queuedMessages || [], - inputContent: action.payload.inputContent, - selection: action.payload.selection, - permissionMode: action.payload.permissionMode || state.permissionMode, - attachedImages: action.payload.attachedImages || [], - sessionsLoading: false, - configurationLoading: false - }; - case 'UPDATE_SELECTION': - return { - ...state, - selection: action.payload - }; - case 'SET_QUEUED_MESSAGES': - return { - ...state, - queuedMessages: action.payload - }; - case 'SET_COMMAND_RUNNING': - return { - ...state, - isCommandRunning: action.payload - }; - case 'SET_PERMISSION_MODE': - return { - ...state, - permissionMode: action.payload - }; - default: - return state; - } -} - export const ChatApp: React.FC = ({ vscode }) => { const [state, dispatch] = useReducer(chatReducer, initialState); const messageInputRef = useRef<{ focus: () => void }>(null); @@ -376,6 +174,25 @@ export const ChatApp: React.FC = ({ vscode }) => { messageListRef.current.scrollToBottom('smooth'); } break; + // Incremental update commands for streaming optimization + case 'appendMessage': + dispatch({ type: 'APPEND_MESSAGE', payload: message.message }); + break; + case 'updateStreamingContent': + dispatch({ + type: 'UPDATE_STREAMING_CONTENT', + payload: { messageId: message.messageId, accumulated: message.accumulated, stage: message.stage } + }); + break; + case 'updateStreamingReasoning': + dispatch({ + type: 'UPDATE_STREAMING_REASONING', + payload: { messageId: message.messageId, accumulated: message.accumulated, stage: message.stage } + }); + break; + case 'updateToolBlock': + dispatch({ type: 'UPDATE_TOOL_BLOCK', payload: message.params }); + break; } }; diff --git a/webview/src/components/Message.tsx b/webview/src/components/Message.tsx index 83e03d0..a7f7cfe 100644 --- a/webview/src/components/Message.tsx +++ b/webview/src/components/Message.tsx @@ -1,3 +1,4 @@ +import React from 'react'; import { ContextTag } from './ContextTag'; import { parseMentions } from '../utils/messageUtils'; import { marked } from 'marked'; @@ -100,7 +101,7 @@ const parseMarkdownWithMermaid = (content: string): ParsedMarkdownContent => { }; -export const Message: React.FC = (props) => { +export const Message: React.FC = React.memo((props) => { const { message, isStreaming = false, isQueued = false, onRewindToMessage } = props; const getMessageClassName = () => { const classes = ['message']; @@ -596,4 +597,9 @@ export const Message: React.FC = (props) => { )} ); -}; +}, (prev, next) => { + // Custom comparison for React.memo + return prev.message === next.message && + prev.isStreaming === next.isStreaming && + prev.isQueued === next.isQueued; +}); diff --git a/webview/src/components/MessageList.tsx b/webview/src/components/MessageList.tsx index 48d97dd..4a97f0e 100644 --- a/webview/src/components/MessageList.tsx +++ b/webview/src/components/MessageList.tsx @@ -122,7 +122,7 @@ export const MessageList = forwardRef<{ scrollToBottom: (behavior?: ScrollBehavi return ( 0, + // Auto-expand task list when tasks are first created + isTaskListCollapsed: state.tasks.length === 0 && action.payload.length > 0 ? false : state.isTaskListCollapsed + }; + case 'TOGGLE_TASK_LIST_COLLAPSE': + return { + ...state, + isTaskListCollapsed: !state.isTaskListCollapsed + }; + case 'SET_TASK_LIST_COLLAPSED': + return { + ...state, + isTaskListCollapsed: action.payload + }; + case 'TOGGLE_QUEUE_COLLAPSE': + return { + ...state, + isQueueCollapsed: !state.isQueueCollapsed + }; + case 'START_STREAMING': + return { + ...state, + isStreaming: true + }; + case 'END_STREAMING': + return { + ...state, + isStreaming: false + }; + case 'SET_INPUT_DISABLED': + return { + ...state, + inputDisabled: action.payload + }; + case 'INPUT_CLEARED': + return { + ...state, + shouldClearInput: false + }; + case 'SET_SESSIONS': + return { + ...state, + sessions: action.payload, + sessionsLoading: false + }; + case 'SET_CURRENT_SESSION': + return { + ...state, + currentSession: action.payload + }; + case 'SET_SESSIONS_LOADING': + return { + ...state, + sessionsLoading: action.payload + }; + case 'SHOW_CONFIRMATION': + return { + ...state, + pendingConfirmations: [...state.pendingConfirmations, action.payload] + }; + case 'HIDE_CONFIRMATION': + return { + ...state, + pendingConfirmations: state.pendingConfirmations.filter(c => c.confirmationId !== action.payload) + }; + case 'SHOW_DIALOG': + return { + ...state, + activeDialog: action.payload.type, + configurationData: action.payload.data ?? state.configurationData, + configurationLoading: false, + configurationError: action.payload.error + }; + case 'HIDE_DIALOG': + return { + ...state, + activeDialog: null, + configurationError: undefined + }; + case 'SET_CONFIGURATION_LOADING': + return { + ...state, + configurationLoading: action.payload + }; + case 'SET_CONFIGURATION_ERROR': + return { + ...state, + configurationError: action.payload, + configurationLoading: false + }; + case 'SET_CONFIGURATION_DATA': + return { + ...state, + configurationData: action.payload, + configurationLoading: false + }; + case 'SET_CONFIGURED_MODELS': + return { + ...state, + configuredModels: action.payload + }; + case 'SET_CURRENT_MODELS': + return { + ...state, + currentModel: action.payload.model, + currentFastModel: action.payload.fastModel, + configurationData: { + ...state.configurationData, + model: action.payload.model || state.configurationData?.model, + fastModel: action.payload.fastModel || state.configurationData?.fastModel + } + }; + case 'SET_INITIAL_STATE': + return { + ...state, + messages: action.payload.messages, + tasks: action.payload.tasks || [], + isTaskListVisible: (action.payload.tasks && action.payload.tasks.length > 0) ? true : false, + isTaskListCollapsed: action.payload.isTaskListCollapsed !== undefined ? action.payload.isTaskListCollapsed : state.isTaskListCollapsed, + isStreaming: action.payload.isStreaming !== undefined ? action.payload.isStreaming : state.isStreaming, + isCommandRunning: action.payload.isCommandRunning !== undefined ? action.payload.isCommandRunning : state.isCommandRunning, + sessions: action.payload.sessions || state.sessions || [], + currentSession: action.payload.currentSession || state.currentSession, + configurationData: action.payload.configurationData || state.configurationData, + pendingConfirmations: action.payload.pendingConfirmations || [], + queuedMessages: action.payload.queuedMessages || [], + inputContent: action.payload.inputContent, + selection: action.payload.selection, + permissionMode: action.payload.permissionMode || state.permissionMode, + attachedImages: action.payload.attachedImages || [], + sessionsLoading: false, + configurationLoading: false + }; + case 'UPDATE_SELECTION': + return { + ...state, + selection: action.payload + }; + case 'SET_QUEUED_MESSAGES': + return { + ...state, + queuedMessages: action.payload + }; + case 'SET_COMMAND_RUNNING': + return { + ...state, + isCommandRunning: action.payload + }; + case 'SET_PERMISSION_MODE': + return { + ...state, + permissionMode: action.payload + }; + // Incremental update actions for streaming optimization + case 'APPEND_MESSAGE': + return { + ...state, + messages: [...state.messages, action.payload] + }; + case 'UPDATE_STREAMING_CONTENT': { + const { messageId, accumulated, stage } = action.payload; + const messageIndex = state.messages.findIndex(m => m.id === messageId); + if (messageIndex === -1) return state; + + const message = state.messages[messageIndex]; + const textBlockIndex = message.blocks.findIndex(b => b.type === 'text'); + + let newBlocks: MessageBlock[]; + if (textBlockIndex === -1) { + // No text block yet, append one + const newTextBlock: TextBlock = { + type: 'text', + content: accumulated, + stage + }; + newBlocks = [...message.blocks, newTextBlock]; + } else { + // Update existing text block + newBlocks = message.blocks.map((block, idx) => { + if (idx === textBlockIndex && block.type === 'text') { + return { ...block, content: accumulated, stage } as TextBlock; + } + return block; + }); + } + + const newMessages = state.messages.map((m, idx) => { + if (idx === messageIndex) { + return { ...m, blocks: newBlocks }; + } + return m; + }); + + return { + ...state, + messages: newMessages + }; + } + case 'UPDATE_STREAMING_REASONING': { + const { messageId, accumulated, stage } = action.payload; + const messageIndex = state.messages.findIndex(m => m.id === messageId); + if (messageIndex === -1) return state; + + const message = state.messages[messageIndex]; + const reasoningBlockIndex = message.blocks.findIndex(b => b.type === 'reasoning'); + + let newBlocks: MessageBlock[]; + if (reasoningBlockIndex === -1) { + // No reasoning block yet, append one + const newReasoningBlock = { + type: 'reasoning' as const, + content: accumulated, + stage + }; + newBlocks = [...message.blocks, newReasoningBlock]; + } else { + // Update existing reasoning block + newBlocks = message.blocks.map((block, idx) => { + if (idx === reasoningBlockIndex && block.type === 'reasoning') { + return { ...block, content: accumulated, stage }; + } + return block; + }); + } + + const newMessages = state.messages.map((m, idx) => { + if (idx === messageIndex) { + return { ...m, blocks: newBlocks }; + } + return m; + }); + + return { + ...state, + messages: newMessages + }; + } + case 'UPDATE_TOOL_BLOCK': { + const { messageId, id: toolBlockId, ...updates } = action.payload; + const messageIndex = state.messages.findIndex(m => m.id === messageId); + if (messageIndex === -1) return state; + + const message = state.messages[messageIndex]; + const toolBlockIndex = message.blocks.findIndex(b => b.type === 'tool' && b.id === toolBlockId); + if (toolBlockIndex === -1) return state; + + const newBlocks = message.blocks.map((block, idx) => { + if (idx === toolBlockIndex && block.type === 'tool') { + return { ...block, ...updates }; + } + return block; + }); + + const newMessages = state.messages.map((m, idx) => { + if (idx === messageIndex) { + return { ...m, blocks: newBlocks }; + } + return m; + }); + + return { + ...state, + messages: newMessages + }; + } + default: + return state; + } +} diff --git a/webview/src/types/index.ts b/webview/src/types/index.ts index 2e832d5..69f4288 100644 --- a/webview/src/types/index.ts +++ b/webview/src/types/index.ts @@ -6,10 +6,11 @@ */ // Import message structures and session types from wave-agent-sdk -import type { Message, MessageBlock, TextBlock, ErrorBlock, ToolBlock, ImageBlock, BangBlock, CompactBlock, PermissionMode, AskUserQuestion, AskUserQuestionInput, AskUserQuestionOption, Task, TaskStatus, TaskNotificationBlock, McpServerStatus, McpServerConfig } from 'wave-agent-sdk/dist/types/index.js'; +import type { Message, MessageBlock, TextBlock, ErrorBlock, ToolBlock, ImageBlock, BangBlock, CompactBlock, ReasoningBlock, PermissionMode, AskUserQuestion, AskUserQuestionInput, AskUserQuestionOption, Task, TaskStatus, TaskNotificationBlock, McpServerStatus, McpServerConfig } from 'wave-agent-sdk/dist/types/index.js'; import type { SessionMetadata, SessionData } from 'wave-agent-sdk/dist/services/session.js'; +import type { ToolBlockUpdateCallbackParams } from 'wave-agent-sdk/dist/utils/messageOperations.js'; -export type { Message, MessageBlock, TextBlock, ErrorBlock, ToolBlock, ImageBlock, BangBlock, CompactBlock, TaskNotificationBlock, SessionData, SessionMetadata, PermissionMode, AskUserQuestion, AskUserQuestionInput, AskUserQuestionOption, Task, TaskStatus, McpServerStatus, McpServerConfig }; +export type { Message, MessageBlock, TextBlock, ErrorBlock, ToolBlock, ImageBlock, BangBlock, CompactBlock, ReasoningBlock, TaskNotificationBlock, SessionData, SessionMetadata, PermissionMode, AskUserQuestion, AskUserQuestionInput, AskUserQuestionOption, Task, TaskStatus, McpServerStatus, McpServerConfig }; // Slash command types export interface SlashCommand { @@ -423,4 +424,9 @@ export type ChatAction = permissionMode?: PermissionMode; attachedImages?: AttachedImage[]; queuedMessages?: QueuedMessage[]; - } }; + } } + // Incremental update actions for streaming optimization + | { type: 'APPEND_MESSAGE'; payload: Message } + | { type: 'UPDATE_STREAMING_CONTENT'; payload: { messageId: string; accumulated: string; stage: 'streaming' | 'end' } } + | { type: 'UPDATE_STREAMING_REASONING'; payload: { messageId: string; accumulated: string; stage: 'streaming' | 'end' } } + | { type: 'UPDATE_TOOL_BLOCK'; payload: ToolBlockUpdateCallbackParams };