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

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([
       {

packages/vscode-ide-companion/src/services/conversationStore.ts

@@ -84,6 +84,75 @@ export class ConversationStore {
     return true;
   }
 
+  async renameConversationId(
+    fromConversationId: string,
+    toConversationId: string,
+  ): Promise<boolean> {
+    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<void> {
+    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,

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<void>;
@@ -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
    */

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(

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.' },
     });

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<ChatMessage[]>;
+      }
+    ).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: [] },

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)',
     () => {

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,
         },
       ];
     });