From a656930e82b22b93a6bf50c797ca48ddc3c5ec80 Mon Sep 17 00:00:00 2001 From: Dragon <52599892+DragonnZhang@users.noreply.github.com> Date: Fri, 15 May 2026 13:58:58 +0800 Subject: [PATCH] fix(vscode): preserve thinking state and recover missing edit snapshots (#4147) * fix(vscode): allow editing sessions without local snapshots * fix(vscode): keep thinking after edited user message * fix(vscode): sync conversation id alignment through router --- .../src/services/conversationStore.test.ts | 99 ++++++ .../src/services/conversationStore.ts | 69 ++++ .../webview/handlers/BaseMessageHandler.ts | 13 + .../src/webview/handlers/MessageRouter.ts | 1 + .../handlers/SessionMessageHandler.test.ts | 310 +++++++++++++++++- .../webview/handlers/SessionMessageHandler.ts | 121 ++++--- .../hooks/message/useMessageHandling.test.tsx | 35 ++ .../hooks/message/useMessageHandling.ts | 11 +- 8 files changed, 613 insertions(+), 46 deletions(-) diff --git a/packages/vscode-ide-companion/src/services/conversationStore.test.ts b/packages/vscode-ide-companion/src/services/conversationStore.test.ts index 79bc899f7..ecd08218b 100644 --- a/packages/vscode-ide-companion/src/services/conversationStore.test.ts +++ b/packages/vscode-ide-companion/src/services/conversationStore.test.ts @@ -72,6 +72,105 @@ describe('ConversationStore', () => { expect(update).not.toHaveBeenCalled(); }); + it('renameConversationId updates a conversation id and current id', async () => { + const { store, update, conversations } = createStore([ + { + id: 'conversation-1', + title: 'Conversation', + messages: [{ role: 'user' as const, content: 'first', timestamp: 1 }], + createdAt: 1, + updatedAt: 1, + }, + ]); + store.setCurrentConversationId('conversation-1'); + + await expect( + store.renameConversationId('conversation-1', 'session-1'), + ).resolves.toBe(true); + + expect(conversations[0]?.id).toBe('session-1'); + expect(store.getCurrentConversationId()).toBe('session-1'); + expect(update).toHaveBeenCalledWith('conversations', conversations); + }); + + it('renameConversationId returns false when the target id already exists', async () => { + const { store, update, conversations } = createStore([ + { + id: 'conversation-1', + title: 'Conversation', + messages: [], + createdAt: 1, + updatedAt: 1, + }, + { + id: 'session-1', + title: 'Existing Session', + messages: [], + createdAt: 2, + updatedAt: 2, + }, + ]); + + await expect( + store.renameConversationId('conversation-1', 'session-1'), + ).resolves.toBe(false); + + expect(conversations.map((conversation) => conversation.id)).toEqual([ + 'conversation-1', + 'session-1', + ]); + expect(update).not.toHaveBeenCalled(); + }); + + it('upsertConversation inserts a missing conversation with cloned messages', async () => { + const messages = [ + { role: 'user' as const, content: 'first', timestamp: 1 }, + ]; + const { store, update, conversations } = createStore([]); + + await store.upsertConversation({ + id: 'session-1', + title: 'Session', + messages, + createdAt: 1, + updatedAt: 1, + }); + + expect(conversations).toHaveLength(1); + expect(conversations[0]?.id).toBe('session-1'); + expect(conversations[0]?.messages).toEqual(messages); + expect(conversations[0]?.messages[0]).not.toBe(messages[0]); + expect(store.getCurrentConversationId()).toBe('session-1'); + expect(update).toHaveBeenCalledWith('conversations', conversations); + }); + + it('upsertConversation replaces an existing conversation', async () => { + const { store, update, conversations } = createStore([ + { + id: 'session-1', + title: 'Old', + messages: [{ role: 'user' as const, content: 'old', timestamp: 1 }], + createdAt: 1, + updatedAt: 1, + }, + ]); + + await store.upsertConversation({ + id: 'session-1', + title: 'New', + messages: [{ role: 'assistant' as const, content: 'new', timestamp: 2 }], + createdAt: 1, + updatedAt: 2, + }); + + expect(conversations).toHaveLength(1); + expect(conversations[0]?.title).toBe('New'); + expect(conversations[0]?.messages).toEqual([ + { role: 'assistant', content: 'new', timestamp: 2 }, + ]); + expect(update).toHaveBeenCalledWith('conversations', conversations); + }); + it('truncateFromUserTurn truncates from the matching user turn', async () => { const { store, update, conversations } = createStore([ { diff --git a/packages/vscode-ide-companion/src/services/conversationStore.ts b/packages/vscode-ide-companion/src/services/conversationStore.ts index dc15a50cb..a99d417ab 100644 --- a/packages/vscode-ide-companion/src/services/conversationStore.ts +++ b/packages/vscode-ide-companion/src/services/conversationStore.ts @@ -84,6 +84,75 @@ export class ConversationStore { return true; } + async renameConversationId( + fromConversationId: string, + toConversationId: string, + ): Promise { + if (fromConversationId === toConversationId) { + return true; + } + + const conversations = await this.getAllConversations(); + const sourceIndex = conversations.findIndex( + (c) => c.id === fromConversationId, + ); + + if (sourceIndex < 0) { + console.warn( + '[ConversationStore] renameConversationId: source conversation not found:', + fromConversationId, + ); + return false; + } + + if (conversations.some((c) => c.id === toConversationId)) { + console.warn( + '[ConversationStore] renameConversationId: target conversation already exists:', + toConversationId, + ); + return false; + } + + const source = conversations[sourceIndex]; + if (!source) { + return false; + } + + conversations[sourceIndex] = { + ...source, + id: toConversationId, + updatedAt: Date.now(), + }; + + await this.context.globalState.update('conversations', conversations); + + if (this.currentConversationId === fromConversationId) { + this.currentConversationId = toConversationId; + } + + return true; + } + + async upsertConversation(conversation: Conversation): Promise { + const conversations = await this.getAllConversations(); + const storedConversation: Conversation = { + ...conversation, + messages: conversation.messages.map((message) => ({ ...message })), + }; + const existingIndex = conversations.findIndex( + (c) => c.id === conversation.id, + ); + + if (existingIndex >= 0) { + conversations[existingIndex] = storedConversation; + } else { + conversations.push(storedConversation); + } + + await this.context.globalState.update('conversations', conversations); + this.currentConversationId = conversation.id; + } + async truncateFromUserTurn( conversationId: string, targetTurnIndex: number, diff --git a/packages/vscode-ide-companion/src/webview/handlers/BaseMessageHandler.ts b/packages/vscode-ide-companion/src/webview/handlers/BaseMessageHandler.ts index 4d01fd021..90b97a25c 100644 --- a/packages/vscode-ide-companion/src/webview/handlers/BaseMessageHandler.ts +++ b/packages/vscode-ide-companion/src/webview/handlers/BaseMessageHandler.ts @@ -37,6 +37,7 @@ export abstract class BaseMessageHandler implements IMessageHandler { protected conversationStore: ConversationStore, protected currentConversationId: string | null, protected sendToWebView: (message: unknown) => void, + private readonly syncCurrentConversationId?: (id: string | null) => void, ) {} abstract handle(message: { type: string; data?: unknown }): Promise; @@ -49,6 +50,18 @@ export abstract class BaseMessageHandler implements IMessageHandler { this.currentConversationId = id; } + /** + * Update current conversation ID through the owning router when available. + */ + protected updateCurrentConversationId(id: string | null): void { + if (this.syncCurrentConversationId) { + this.syncCurrentConversationId(id); + return; + } + + this.currentConversationId = id; + } + /** * Get current conversation ID */ diff --git a/packages/vscode-ide-companion/src/webview/handlers/MessageRouter.ts b/packages/vscode-ide-companion/src/webview/handlers/MessageRouter.ts index e7d40fd47..0b307847f 100644 --- a/packages/vscode-ide-companion/src/webview/handlers/MessageRouter.ts +++ b/packages/vscode-ide-companion/src/webview/handlers/MessageRouter.ts @@ -48,6 +48,7 @@ export class MessageRouter { conversationStore, currentConversationId, sendToWebView, + (id) => this.setCurrentConversationId(id), ); this.fileHandler = new FileMessageHandler( diff --git a/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.test.ts b/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.test.ts index 415a02f13..a8400c7ef 100644 --- a/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.test.ts +++ b/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.test.ts @@ -118,6 +118,7 @@ describe('SessionMessageHandler', () => { createConversation: vi.fn().mockResolvedValue({ id: 'conversation-1' }), getConversation: vi.fn().mockResolvedValue(null), addMessage: vi.fn(), + renameConversationId: vi.fn().mockResolvedValue(true), }; const sendToWebView = vi.fn(); @@ -181,6 +182,7 @@ describe('SessionMessageHandler', () => { createConversation: vi.fn().mockResolvedValue({ id: 'conversation-1' }), getConversation: vi.fn().mockResolvedValue(null), addMessage: vi.fn(), + renameConversationId: vi.fn().mockResolvedValue(true), }; const sendToWebView = vi.fn(); @@ -222,6 +224,215 @@ describe('SessionMessageHandler', () => { ]); }); + it('keeps the conversation store aligned with the ACP session id before editing', async () => { + mockProcessImageAttachments.mockImplementation( + async (promptText: string) => ({ + formattedText: promptText, + displayText: promptText, + savedImageCount: 0, + promptImages: [], + }), + ); + + const agentManager = { + isConnected: true, + currentSessionId: 'session-1', + rewindSession: vi.fn().mockResolvedValue({ + historyBeforeRewind: [{ role: 'user', parts: [{ text: 'first' }] }], + }), + restoreSessionHistory: vi.fn().mockResolvedValue(undefined), + sendMessage: vi.fn().mockResolvedValue(undefined), + }; + let conversation = { + id: 'conversation-1', + title: 'Conversation', + messages: [] as Array<{ + role: 'user' | 'assistant' | 'thinking'; + content: string; + timestamp: number; + }>, + createdAt: 1, + updatedAt: 1, + }; + const conversationStore = { + createConversation: vi.fn().mockResolvedValue(conversation), + getConversation: vi.fn(async (id: string) => + conversation.id === id ? conversation : null, + ), + addMessage: vi.fn(async (id: string, message) => { + if (conversation.id === id) { + conversation.messages.push(message); + } + }), + renameConversationId: vi.fn(async (fromId: string, toId: string) => { + if (conversation.id !== fromId) { + return false; + } + conversation = { ...conversation, id: toId }; + return true; + }), + replaceMessages: vi.fn().mockResolvedValue(true), + truncateFromUserTurn: vi.fn().mockResolvedValue(true), + }; + const sendToWebView = vi.fn(); + + const handler = new SessionMessageHandler( + agentManager as never, + conversationStore as never, + null, + sendToWebView, + ); + + await handler.handle({ + type: 'sendMessage', + data: { text: 'first prompt' }, + }); + + await handler.handle({ + type: 'editMessage', + data: { + text: 'edited prompt', + targetTurnIndex: 0, + }, + }); + + expect(conversationStore.renameConversationId).toHaveBeenCalledWith( + 'conversation-1', + 'session-1', + ); + expect(conversationStore.getConversation).toHaveBeenCalledWith('session-1'); + expect(conversationStore.truncateFromUserTurn).toHaveBeenCalledWith( + 'session-1', + 0, + ); + expect(agentManager.rewindSession).toHaveBeenCalledWith(0); + expect(sendToWebView).not.toHaveBeenCalledWith({ + type: 'error', + data: { message: 'Failed to capture conversation state before editing.' }, + }); + }); + + it('does not switch to a colliding ACP session id when rename fails', async () => { + mockProcessImageAttachments.mockImplementation( + async (promptText: string) => ({ + formattedText: promptText, + displayText: promptText, + savedImageCount: 0, + promptImages: [], + }), + ); + + const agentManager = { + isConnected: true, + currentSessionId: 'session-1', + sendMessage: vi.fn().mockResolvedValue(undefined), + }; + const conversation = { + id: 'conversation-1', + title: 'Conversation', + messages: [], + createdAt: 1, + updatedAt: 1, + }; + const conversationStore = { + createConversation: vi.fn().mockResolvedValue(conversation), + getConversation: vi.fn().mockResolvedValue(conversation), + addMessage: vi.fn().mockResolvedValue(undefined), + renameConversationId: vi.fn().mockResolvedValue(false), + }; + const sendToWebView = vi.fn(); + const handlerRef: { current: SessionMessageHandler | null } = { + current: null, + }; + const syncCurrentConversationId = vi.fn((id: string | null) => { + handlerRef.current?.setCurrentConversationId(id); + }); + + const handler = new SessionMessageHandler( + agentManager as never, + conversationStore as never, + null, + sendToWebView, + syncCurrentConversationId, + ); + handlerRef.current = handler; + + await handler.handle({ + type: 'sendMessage', + data: { text: 'first prompt' }, + }); + + expect(conversationStore.renameConversationId).toHaveBeenCalledWith( + 'conversation-1', + 'session-1', + ); + expect(syncCurrentConversationId).toHaveBeenCalledWith('conversation-1'); + expect(syncCurrentConversationId).not.toHaveBeenCalledWith('session-1'); + expect(handler.getCurrentConversationId()).toBe('conversation-1'); + expect(sendToWebView).not.toHaveBeenCalledWith({ + type: 'sessionTitleUpdated', + data: { + sessionId: 'session-1', + title: 'first prompt', + }, + }); + }); + + it('syncs ACP session id alignment through the owning router setter', async () => { + mockProcessImageAttachments.mockImplementation( + async (promptText: string) => ({ + formattedText: promptText, + displayText: promptText, + savedImageCount: 0, + promptImages: [], + }), + ); + + const agentManager = { + isConnected: true, + currentSessionId: 'session-1', + sendMessage: vi.fn().mockResolvedValue(undefined), + }; + const conversation = { + id: 'conversation-1', + title: 'Conversation', + messages: [], + createdAt: 1, + updatedAt: 1, + }; + const conversationStore = { + createConversation: vi.fn().mockResolvedValue(conversation), + getConversation: vi.fn().mockResolvedValue(conversation), + addMessage: vi.fn().mockResolvedValue(undefined), + renameConversationId: vi.fn().mockResolvedValue(true), + }; + const sendToWebView = vi.fn(); + const handlerRef: { current: SessionMessageHandler | null } = { + current: null, + }; + const syncCurrentConversationId = vi.fn((id: string | null) => { + handlerRef.current?.setCurrentConversationId(id); + }); + + const handler = new SessionMessageHandler( + agentManager as never, + conversationStore as never, + null, + sendToWebView, + syncCurrentConversationId, + ); + handlerRef.current = handler; + + await handler.handle({ + type: 'sendMessage', + data: { text: 'first prompt' }, + }); + + expect(syncCurrentConversationId).toHaveBeenCalledWith('conversation-1'); + expect(syncCurrentConversationId).toHaveBeenCalledWith('session-1'); + expect(handler.getCurrentConversationId()).toBe('session-1'); + }); + it('rewinds the active ACP session before sending an edited message', async () => { mockProcessImageAttachments.mockResolvedValue({ formattedText: 'edited prompt', @@ -363,7 +574,7 @@ describe('SessionMessageHandler', () => { }); }); - it('aborts edits when the restore snapshot cannot be captured', async () => { + it('continues edits with ACP-only rewind when no local snapshot exists', async () => { mockProcessImageAttachments.mockResolvedValue({ formattedText: 'edited prompt', displayText: 'edited prompt', @@ -374,9 +585,12 @@ describe('SessionMessageHandler', () => { const agentManager = { isConnected: true, currentSessionId: 'session-1', - rewindSession: vi.fn(), + rewindSession: vi.fn().mockResolvedValue({ + historyBeforeRewind: [{ role: 'user', parts: [{ text: 'first' }] }], + }), restoreSessionHistory: vi.fn(), - sendMessage: vi.fn(), + sendMessage: vi.fn().mockResolvedValue(undefined), + getSessionMessages: vi.fn().mockResolvedValue([]), }; const conversationStore = { createConversation: vi.fn(), @@ -384,6 +598,7 @@ describe('SessionMessageHandler', () => { addMessage: vi.fn(), replaceMessages: vi.fn(), truncateFromUserTurn: vi.fn(), + upsertConversation: vi.fn(), }; const sendToWebView = vi.fn(); @@ -403,9 +618,94 @@ describe('SessionMessageHandler', () => { }); expect(conversationStore.truncateFromUserTurn).not.toHaveBeenCalled(); - expect(agentManager.rewindSession).not.toHaveBeenCalled(); - expect(agentManager.sendMessage).not.toHaveBeenCalled(); + expect(agentManager.rewindSession).toHaveBeenCalledWith(1); expect(sendToWebView).toHaveBeenCalledWith({ + type: 'conversationRewound', + data: { targetTurnIndex: 1 }, + }); + expect(agentManager.sendMessage).toHaveBeenCalledWith([ + { type: 'text', text: 'edited prompt' }, + ]); + expect(sendToWebView).not.toHaveBeenCalledWith({ + type: 'sessionTitleUpdated', + data: { + sessionId: 'session-1', + title: 'edited prompt', + }, + }); + expect(sendToWebView).not.toHaveBeenCalledWith({ + type: 'error', + data: expect.objectContaining({ + message: 'Failed to capture conversation state before editing.', + }), + }); + }); + + it('recovers a missing edit snapshot from persisted session messages', async () => { + mockProcessImageAttachments.mockResolvedValue({ + formattedText: 'edited prompt', + displayText: 'edited prompt', + savedImageCount: 0, + promptImages: [], + }); + + const persistedMessages = [ + { role: 'user' as const, content: 'first', timestamp: 1 }, + { role: 'assistant' as const, content: 'first reply', timestamp: 2 }, + { role: 'user' as const, content: 'second', timestamp: 3 }, + ]; + const agentManager = { + isConnected: true, + currentSessionId: 'session-1', + rewindSession: vi.fn().mockResolvedValue({ + historyBeforeRewind: [{ role: 'user', parts: [{ text: 'first' }] }], + }), + restoreSessionHistory: vi.fn().mockResolvedValue(undefined), + sendMessage: vi.fn().mockResolvedValue(undefined), + getSessionMessages: vi.fn().mockResolvedValue(persistedMessages), + }; + const conversationStore = { + createConversation: vi.fn(), + getConversation: vi.fn().mockResolvedValue(null), + addMessage: vi.fn(), + replaceMessages: vi.fn(), + truncateFromUserTurn: vi.fn().mockResolvedValue(true), + upsertConversation: vi.fn().mockResolvedValue(undefined), + }; + const sendToWebView = vi.fn(); + + const handler = new SessionMessageHandler( + agentManager as never, + conversationStore as never, + 'session-1', + sendToWebView, + ); + + await handler.handle({ + type: 'editMessage', + data: { + text: 'edited prompt', + targetTurnIndex: 1, + }, + }); + + expect(agentManager.getSessionMessages).toHaveBeenCalledWith('session-1'); + expect(conversationStore.upsertConversation).toHaveBeenCalledWith({ + id: 'session-1', + title: 'first', + messages: persistedMessages, + createdAt: 1, + updatedAt: 3, + }); + expect(conversationStore.truncateFromUserTurn).toHaveBeenCalledWith( + 'session-1', + 1, + ); + expect(agentManager.rewindSession).toHaveBeenCalledWith(1); + expect(agentManager.sendMessage).toHaveBeenCalledWith([ + { type: 'text', text: 'edited prompt' }, + ]); + expect(sendToWebView).not.toHaveBeenCalledWith({ type: 'error', data: { message: 'Failed to capture conversation state before editing.' }, }); diff --git a/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.ts b/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.ts index 04273f902..c60b866ab 100644 --- a/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.ts +++ b/packages/vscode-ide-companion/src/webview/handlers/SessionMessageHandler.ts @@ -241,14 +241,41 @@ export class SessionMessageHandler extends BaseMessageHandler { const conversation = await this.conversationStore.getConversation(conversationId); - if (!conversation) { + if (conversation) { + return { + ...conversation, + messages: conversation.messages.map((message) => ({ ...message })), + }; + } + + const getSessionMessages = ( + this.agentManager as { + getSessionMessages?: (sessionId: string) => Promise; + } + ).getSessionMessages; + if (!getSessionMessages) { return null; } - return { - ...conversation, - messages: conversation.messages.map((message) => ({ ...message })), + const messages = await getSessionMessages.call( + this.agentManager, + conversationId, + ); + if (messages.length === 0) { + return null; + } + + const timestamps = messages.map((message) => message.timestamp); + const recoveredConversation: Conversation = { + id: conversationId, + title: messages.find((message) => message.role === 'user')?.content ?? '', + messages: messages.map((message) => ({ ...message })), + createdAt: Math.min(...timestamps), + updatedAt: Math.max(...timestamps), }; + await this.conversationStore.upsertConversation(recoveredConversation); + + return recoveredConversation; } private async restoreConversationSnapshot( @@ -268,7 +295,7 @@ export class SessionMessageHandler extends BaseMessageHandler { snapshot.id, ); } - this.currentConversationId = snapshot.id; + this.updateCurrentConversationId(snapshot.id); this.sendToWebView({ type: 'conversationLoaded', data: snapshot, @@ -535,7 +562,7 @@ export class SessionMessageHandler extends BaseMessageHandler { ); try { const newConv = await this.conversationStore.createConversation(); - this.currentConversationId = newConv.id; + this.updateCurrentConversationId(newConv.id); this.sendToWebView({ type: 'conversationLoaded', data: newConv, @@ -630,25 +657,22 @@ export class SessionMessageHandler extends BaseMessageHandler { } if (!editRestoreSnapshot) { - const errorMsg = 'Failed to capture conversation state before editing.'; - console.error('[SessionMessageHandler]', errorMsg); - vscode.window.showErrorMessage(`Failed to edit message: ${errorMsg}`); - this.sendToWebView({ - type: 'error', - data: { message: errorMsg }, - }); - return; + console.warn( + '[SessionMessageHandler] Local conversation snapshot missing before edit; continuing with ACP rewind only.', + ); } try { - const truncated = await this.conversationStore.truncateFromUserTurn( - this.currentConversationId, - editTargetTurnIndex, - ); - if (!truncated) { - throw new Error('Conversation not found for edit target.'); + if (editRestoreSnapshot) { + const truncated = await this.conversationStore.truncateFromUserTurn( + this.currentConversationId, + editTargetTurnIndex, + ); + if (!truncated) { + throw new Error('Conversation not found for edit target.'); + } + editStoreMutationApplied = true; } - editStoreMutationApplied = true; const rewindResult = await this.agentManager.rewindSession(editTargetTurnIndex); @@ -691,16 +715,20 @@ export class SessionMessageHandler extends BaseMessageHandler { // Check if this is the first message let isFirstMessage = false; - try { - const conversation = await this.conversationStore.getConversation( - this.currentConversationId, - ); - isFirstMessage = !conversation || conversation.messages.length === 0; - } catch (error) { - console.error( - '[SessionMessageHandler] Failed to check conversation:', - error, - ); + if (editTargetTurnIndex !== undefined) { + isFirstMessage = editTargetTurnIndex === 0; + } else { + try { + const conversation = await this.conversationStore.getConversation( + this.currentConversationId, + ); + isFirstMessage = !conversation || conversation.messages.length === 0; + } catch (error) { + console.error( + '[SessionMessageHandler] Failed to check conversation:', + error, + ); + } } // Generate title for first message, but only if it hasn't been set yet @@ -843,7 +871,22 @@ export class SessionMessageHandler extends BaseMessageHandler { // After first message, sync ACP session ID to webview for session list highlighting const acpSessionId = this.agentManager.currentSessionId; if (acpSessionId && acpSessionId !== this.currentConversationId) { - this.currentConversationId = acpSessionId; + const previousConversationId = this.currentConversationId; + if (previousConversationId) { + const renamed = await this.conversationStore.renameConversationId( + previousConversationId, + acpSessionId, + ); + if (!renamed) { + console.warn( + '[SessionMessageHandler] Failed to align conversation store with ACP session id:', + previousConversationId, + acpSessionId, + ); + return; + } + } + this.updateCurrentConversationId(acpSessionId); this.sendToWebView({ type: 'sessionTitleUpdated', data: { @@ -966,7 +1009,7 @@ export class SessionMessageHandler extends BaseMessageHandler { const workingDir = workspaceFolder?.uri.fsPath || process.cwd(); await this.agentManager.createNewSession(workingDir, { forceNew: true }); - this.currentConversationId = null; + this.updateCurrentConversationId(null); this.sendToWebView({ type: 'conversationCleared', @@ -1021,7 +1064,7 @@ export class SessionMessageHandler extends BaseMessageHandler { // Show messages from local cache only const messages = await this.agentManager.getSessionMessages(sessionId); - this.currentConversationId = sessionId; + this.updateCurrentConversationId(sessionId); this.sendToWebView({ type: 'qwenSessionSwitched', data: { sessionId, messages }, @@ -1066,7 +1109,7 @@ export class SessionMessageHandler extends BaseMessageHandler { // Try to load session via ACP (now we should be connected) try { // Set current id and clear UI first so replayed updates append afterwards - this.currentConversationId = sessionId; + this.updateCurrentConversationId(sessionId); this.sendToWebView({ type: 'qwenSessionSwitched', data: { sessionId, messages: [], session: sessionDetails }, @@ -1131,7 +1174,7 @@ export class SessionMessageHandler extends BaseMessageHandler { // to the new ACP id here would desync the backend from the webview // and cause rename/delete/title-update flows to target the wrong // session during the fallback window. - this.currentConversationId = sessionId; + this.updateCurrentConversationId(sessionId); this.sendToWebView({ type: 'qwenSessionSwitched', @@ -1182,7 +1225,7 @@ export class SessionMessageHandler extends BaseMessageHandler { } } else { // Offline view only - this.currentConversationId = sessionId; + this.updateCurrentConversationId(sessionId); this.sendToWebView({ type: 'qwenSessionSwitched', data: { sessionId, messages, session: sessionDetails }, @@ -1307,7 +1350,7 @@ export class SessionMessageHandler extends BaseMessageHandler { if (choice === 'offline') { const messages = await this.agentManager.getSessionMessages(sessionId); - this.currentConversationId = sessionId; + this.updateCurrentConversationId(sessionId); this.sendToWebView({ type: 'qwenSessionSwitched', data: { sessionId, messages }, @@ -1324,7 +1367,7 @@ export class SessionMessageHandler extends BaseMessageHandler { // Try ACP load first try { // Pre-clear UI so replayed updates append afterwards - this.currentConversationId = sessionId; + this.updateCurrentConversationId(sessionId); this.sendToWebView({ type: 'qwenSessionSwitched', data: { sessionId, messages: [] }, diff --git a/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.test.tsx b/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.test.tsx index cf1921407..6e2182c84 100644 --- a/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.test.tsx +++ b/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.test.tsx @@ -123,6 +123,41 @@ describe('useMessageHandling', () => { expect(assistantMessages[1].timestamp).toBeGreaterThan(toolCallTimestamp); }); + it('keeps thinking after the user message when streamStart shares the same timestamp', () => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + + const rendered = renderHookHarness(); + root = rendered.root; + container = rendered.container; + + act(() => { + rendered.api.addMessage({ + role: 'user', + content: 'edited prompt', + timestamp: 1_000, + }); + }); + + act(() => { + rendered.api.startStreaming(1_000); + }); + + act(() => { + rendered.api.appendThinkingChunk('thinking'); + }); + + const sorted = [...rendered.api.messages].sort( + (a, b) => a.timestamp - b.timestamp, + ); + + expect(sorted.map((message) => message.role)).toEqual([ + 'user', + 'thinking', + 'assistant', + ]); + }); + it.fails( 'keeps every assistant segment of a turn before a user message that was sent between segments (#3273)', () => { diff --git a/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.ts b/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.ts index 2ddb8460b..028caeb75 100644 --- a/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.ts +++ b/packages/vscode-ide-companion/src/webview/hooks/message/useMessageHandling.ts @@ -59,13 +59,20 @@ export const useMessageHandling = () => { setMessages((prev) => { // Record index of the placeholder to update on chunks streamingMessageIndexRef.current = prev.length; + const maxExistingTimestamp = prev.reduce( + (max, message) => Math.max(max, message.timestamp || 0), + 0, + ); + const placeholderTimestamp = Math.max( + typeof timestamp === 'number' ? timestamp : Date.now(), + maxExistingTimestamp + 2, + ); return [ ...prev, { role: 'assistant', content: '', - // Use provided timestamp (from extension) to keep ordering stable - timestamp: typeof timestamp === 'number' ? timestamp : Date.now(), + timestamp: placeholderTimestamp, }, ]; });