fix(web-shell): message edit fails closed while the transcript window is incomplete

packages/vscode-ide-companion/src/webview/EmbeddedApp.tsx

@@ -323,10 +323,13 @@ export function EmbeddedApp() {
   const [switchingSessionId, setSwitchingSessionId] = useState<string>();
   const [creatingSession, setCreatingSession] = useState(false);
   const [editingMessage, setEditingMessage] = useState<EditingMessage>();
-  const latestSubmittedPromptRef = useRef<{
-    sessionId: string;
-    prompt: string;
-  } | undefined>(undefined);
+  const latestSubmittedPromptRef = useRef<
+    | {
+        sessionId: string;
+        prompt: string;
+      }
+    | undefined
+  >(undefined);
   const sessionSwitchStartedAtRef = useRef(0);
   const sessionSwitchTimerRef = useRef<
     ReturnType<typeof setTimeout> | undefined

packages/web-shell/client/components/MessageList.dom.test.tsx

@@ -53,6 +53,7 @@ vi.mock('./MessageItem', async () => {
       assistantTurnFooterInfo,
       sendFailed,
       onRetrySend,
+      onEditUserMessage,
     }: {
       message: Message;
       showAssistantActions?: boolean;
@@ -63,6 +64,7 @@ vi.mock('./MessageItem', async () => {
       assistantTurnFooterInfo?: WebShellAssistantTurnFooterRenderInfo;
       sendFailed?: boolean;
       onRetrySend?: () => void;
+      onEditUserMessage?: () => void;
     }) => {
       if (message.role === 'tool_group') {
         messageItemTestState.toolArrays.push(message.tools);
@@ -107,6 +109,17 @@ vi.mock('./MessageItem', async () => {
               'data-testid': `disclosure-${message.id}`,
             })
           : null,
+        onEditUserMessage
+          ? React.createElement(
+              'button',
+              {
+                'data-testid': `edit-${message.id}`,
+                onClick: onEditUserMessage,
+                type: 'button',
+              },
+              'edit',
+            )
+          : null,
         showAssistantBranch
           ? React.createElement('button', {
               'data-testid': `branch-${message.id}`,
@@ -353,6 +366,7 @@ function mount(
     pendingApproval?: PermissionRequest | null;
     failedPromptMessageId?: string;
     onRetryFailedPrompt?: () => void;
+    onEditUserMessage?: (targetTurnIndex: number, content: string) => void;
   } = {},
 ): HTMLElement {
   const container = document.createElement('div');
@@ -392,6 +406,7 @@ function mount(
                 onCanScrollToBottomChange={opts.onCanScrollToBottomChange}
                 failedPromptMessageId={opts.failedPromptMessageId}
                 onRetryFailedPrompt={opts.onRetryFailedPrompt}
+                onEditUserMessage={opts.onEditUserMessage}
               />
             </TranscriptRenderModeProvider>
           </CompactModeContext.Provider>
@@ -416,8 +431,11 @@ function rerenderMessages(
     catchingUp?: boolean;
     isResponding?: boolean;
     hasOlderHistory?: boolean;
+    historyCapacityReached?: boolean;
+    historyPaginationError?: boolean;
     onLoadOlderHistory?: (options?: { force?: boolean }) => Promise<void>;
     sessionKey?: string;
+    onEditUserMessage?: (targetTurnIndex: number, content: string) => void;
   } = {},
 ): void {
   const entry = mounted.find((item) => item.container === container);
@@ -435,8 +453,11 @@ function rerenderMessages(
                 catchingUp={opts.catchingUp}
                 isResponding={opts.isResponding}
                 hasOlderHistory={opts.hasOlderHistory}
+                historyCapacityReached={opts.historyCapacityReached}
+                historyPaginationError={opts.historyPaginationError}
                 onLoadOlderHistory={opts.onLoadOlderHistory}
                 sessionKey={opts.sessionKey}
+                onEditUserMessage={opts.onEditUserMessage}
               />
             </TranscriptRenderModeProvider>
           </CompactModeContext.Provider>
@@ -6505,3 +6526,130 @@ describe('MessageList — turn collapse (DOM)', () => {
     expect(parallelAgentsSummary(c)).toBeNull();
   });
 });
+
+describe('user message edit affordance (issue #10385)', () => {
+  function clickEdit(container: HTMLElement, messageId: string): void {
+    const button = container.querySelector(`[data-testid="edit-${messageId}"]`);
+    expect(button).not.toBeNull();
+    act(() => {
+      (button as HTMLButtonElement).dispatchEvent(
+        new MouseEvent('click', { bubbles: true }),
+      );
+    });
+  }
+
+  it('does not offer editing while older history is unloaded', () => {
+    const onEditUserMessage = vi.fn();
+    // Session-global history has more user turns than the loaded window:
+    // only the last two user messages are rendered. The window-local ordinal
+    // of the last message (1) is not its session-global turn index, so the
+    // edit affordance must not be offered while the window is incomplete.
+    const c = mount([userMsg('u4'), asstMsg('a4'), userMsg('u5')], undefined, {
+      hasOlderHistory: true,
+      onLoadOlderHistory: async () => {},
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u5"]')).toBeNull();
+    expect(c.querySelector('[data-testid="edit-u4"]')).toBeNull();
+    expect(onEditUserMessage).not.toHaveBeenCalled();
+  });
+
+  it('does not offer editing when older history was dropped for window capacity', () => {
+    const onEditUserMessage = vi.fn();
+    const c = mount([userMsg('u4'), asstMsg('a4'), userMsg('u5')], undefined, {
+      historyCapacityReached: true,
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u5"]')).toBeNull();
+    expect(onEditUserMessage).not.toHaveBeenCalled();
+  });
+
+  it('does not offer editing while an older-history page terminally failed to load', () => {
+    const onEditUserMessage = vi.fn();
+    // A non-retryable load-older failure latches the provider at
+    // hasMore=false, capacityReached=false, paginationError=true while the
+    // window is still missing older turns, so the window-local ordinal of
+    // the last message (1) is not its session-global turn index either.
+    const c = mount([userMsg('u4'), asstMsg('a4'), userMsg('u5')], undefined, {
+      historyPaginationError: true,
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u5"]')).toBeNull();
+    expect(onEditUserMessage).not.toHaveBeenCalled();
+  });
+
+  it('offers editing only for the last user message when the window is complete', () => {
+    const onEditUserMessage = vi.fn();
+    const c = mount([userMsg('u1'), asstMsg('a1'), userMsg('u2')], undefined, {
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u1"]')).toBeNull();
+    clickEdit(c, 'u2');
+    expect(onEditUserMessage).toHaveBeenCalledTimes(1);
+    expect(onEditUserMessage).toHaveBeenCalledWith(1, 'q');
+  });
+
+  it('re-offers editing when hasOlderHistory flips without a message identity change', () => {
+    const onEditUserMessage = vi.fn();
+    const messages = [userMsg('u4'), asstMsg('a4'), userMsg('u5')];
+    const c = mount(messages, undefined, {
+      hasOlderHistory: true,
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u5"]')).toBeNull();
+    // A load-older request resolving with zero additional turns flips the
+    // flag while the messages array keeps its identity; the render callback
+    // must not hold the stale flag value.
+    rerenderMessages(c, messages, {
+      hasOlderHistory: false,
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u5"]')).not.toBeNull();
+    clickEdit(c, 'u5');
+    expect(onEditUserMessage).toHaveBeenCalledWith(1, 'q');
+  });
+
+  it('re-offers editing when historyCapacityReached flips without a message identity change', () => {
+    const onEditUserMessage = vi.fn();
+    const messages = [userMsg('u4'), asstMsg('a4'), userMsg('u5')];
+    const c = mount(messages, undefined, {
+      historyCapacityReached: true,
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u5"]')).toBeNull();
+    rerenderMessages(c, messages, {
+      historyCapacityReached: false,
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u5"]')).not.toBeNull();
+  });
+
+  it('re-offers editing when historyPaginationError flips without a message identity change', () => {
+    const onEditUserMessage = vi.fn();
+    const messages = [userMsg('u4'), asstMsg('a4'), userMsg('u5')];
+    const c = mount(messages, undefined, {
+      historyPaginationError: true,
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u5"]')).toBeNull();
+    // A successful retry of the failed load-older page clears the flag
+    // while the messages array keeps its identity; the render callback
+    // must not hold the stale flag value.
+    rerenderMessages(c, messages, {
+      historyPaginationError: false,
+      onEditUserMessage,
+    });
+    expect(c.querySelector('[data-testid="edit-u5"]')).not.toBeNull();
+  });
+
+  it('does not count user_shell echoes when numbering user turns', () => {
+    const onEditUserMessage = vi.fn();
+    const c = mount(
+      [userMsg('u1'), userShellMsg('s1'), asstMsg('a1'), userMsg('u2')],
+      undefined,
+      { onEditUserMessage },
+    );
+    clickEdit(c, 'u2');
+    expect(onEditUserMessage).toHaveBeenCalledWith(1, 'q');
+  });
+});

packages/web-shell/client/components/MessageList.tsx

@@ -5437,8 +5437,19 @@ export const MessageList = memo(
               onEditUserMessage={
                 onEditUserMessage &&
                 !isResponding &&
+                // Rewind snapshots are indexed session-globally, but the
+                // rendered transcript is a capped/paginated window. While
+                // older history is still unloaded (hasOlderHistory), was
+                // dropped for window capacity (historyCapacityReached), or
+                // the load-older page failed terminally
+                // (historyPaginationError — the provider then latches
+                // hasMore=false and merges nothing), the window-local
+                // user-turn ordinal is not the session-global turn index,
+                // so offering edit here would rewind to the wrong snapshot
+                // (or none at all). Fail closed (#10385).
                 !hasOlderHistory &&
                 !historyCapacityReached &&
+                !historyPaginationError &&
                 displayItem.message.role === 'user' &&
                 editableUserContent !== undefined &&
                 displayItem.message.id === editableUserTurn.lastId
@@ -5516,6 +5527,7 @@ export const MessageList = memo(
         editableUserTurn,
         hasOlderHistory,
         historyCapacityReached,
+        historyPaginationError,
         generateContent,
         headerOffset,
         visibleItems,