fix: preserve generated images during pending tasks and refactor image part styles (#1228)

harness/specs/tasks/acp-image-generation-compatibility.md

@@ -16,6 +16,7 @@ touchedAreas:
   - src/stores/acp-chat-session.ts
   - src/pages/Chat/index.tsx
   - src/pages/Chat/ChatInput.tsx
+  - src/pages/Chat/AcpImagePart.tsx
   - tests/unit/acp-image-generation-compat.test.ts
   - tests/unit/acp-reducer.test.ts
   - tests/unit/acp-chat-store.test.ts
@@ -64,6 +65,8 @@ acceptance:
   - Duplicate completion records do not create duplicate assistant image replies.
   - Live background image generation shows its dedicated generating label until its success or failure completion is projected, without changing the existing sending/thinking behavior.
   - Switching conversations preserves each live image-generation pending state and restores its indicator on return.
+  - Previously rendered generated images remain visible while a later image-generation task is pending, including across session reloads and navigation.
+  - Completion evidence received during a session reload is projected after the new load generation is active instead of being overwritten or dropped as stale.
   - Completion evidence received while the image conversation is inactive is deferred to that conversation, and a second prompt cannot be sent until the image task settles.
   - Stale preview resolution does not append to a different active session or generation.
 docs:

src/pages/Chat/AcpImagePart.tsx

@@ -101,7 +101,7 @@ export function AcpImagePart({ part, className }: { part: ImageRenderPart; class
     <figure
       data-testid="acp-image-part"
       className={cn(
-        'group/acp-image relative inline-flex max-w-full overflow-hidden rounded-xl border border-black/10 bg-black/5 dark:border-white/10 dark:bg-white/10',
+        'group/acp-image relative inline-flex w-fit max-w-full overflow-hidden rounded-xl border border-black/10 bg-black/5 dark:border-white/10 dark:bg-white/10',
         className,
       )}
     >

src/stores/acp-chat-session.ts

@@ -224,6 +224,34 @@ function captureLiveSession(state: AcpChatSessionState): void {
   });
 }
 
+function restoreImageGenerationProjections(
+  timeline: AcpTimelineSnapshot,
+  snapshot: AcpTimelineSnapshot,
+): AcpTimelineSnapshot {
+  let restored = timeline;
+  const restoredIds = new Set(timeline.itemOrder);
+  for (const [index, itemId] of snapshot.itemOrder.entries()) {
+    const item = snapshot.itemsById[itemId];
+    if (item?.kind !== 'message-segment' || item.compat?.source !== 'image-generation') continue;
+    let afterItemId: string | undefined;
+    for (let priorIndex = index - 1; priorIndex >= 0; priorIndex -= 1) {
+      const priorItemId = snapshot.itemOrder[priorIndex];
+      if (priorItemId && restoredIds.has(priorItemId)) {
+        afterItemId = priorItemId;
+        break;
+      }
+    }
+    restored = appendSyntheticAssistantMessage(restored, {
+      messageId: item.messageId,
+      evidenceId: item.compat.evidenceId,
+      parts: item.parts,
+      afterItemId,
+    });
+    restoredIds.add(itemId);
+  }
+  return restored;
+}
+
 function compatSession(sessionKey: string): ImageGenerationCompatSession {
   const existing = imageGenerationCompatSessions.get(sessionKey);
   if (existing) return existing;
@@ -1126,6 +1154,9 @@ export const useAcpChatSessionStore = create<AcpChatSessionState>((set, get) =>
           timeline = applyAcpSessionUpdate(timeline, notification, { historical: true });
         }
       }
+      if (!currentResumedSnapshot && restorableBackgroundSnapshot) {
+        timeline = restoreImageGenerationProjections(timeline, restorableBackgroundSnapshot.timeline);
+      }
       const pendingAttachments = newPendingAttachments(
         createEmptyAcpTimeline(input.sessionKey, generation),
         timeline,
@@ -1796,7 +1827,7 @@ export function ensureAcpChatSubscriptions(): void {
     const state = useAcpChatSessionStore.getState();
     if (
       evidence
-      && !deferInactiveImageGenerationCompletion(state.activeSessionKey, evidence)
+      && !deferInactiveImageGenerationCompletion(state.loading ? null : state.activeSessionKey, evidence)
     ) void state.projectImageGenerationCompletion(evidence);
   });
   hostEvents.onChatRuntimeEvent((event) => {
@@ -1804,7 +1835,7 @@ export function ensureAcpChatSubscriptions(): void {
     const state = useAcpChatSessionStore.getState();
     if (
       evidence
-      && !deferInactiveImageGenerationCompletion(state.activeSessionKey, evidence)
+      && !deferInactiveImageGenerationCompletion(state.loading ? null : state.activeSessionKey, evidence)
     ) void state.projectImageGenerationCompletion(evidence);
   });
 }

tests/e2e/chat-run-state-events.spec.ts

@@ -5,6 +5,7 @@ const MAIN_SESSION_KEY = 'agent:main:main';
 const MAIN_WORKSPACE = '/workspace';
 const DEFAULT_WORKSPACE = '~/.openclaw/workspace';
 const IMAGE_GENERATION_TASK_ID = '32aa3a12-a05b-4074-af4e-246cc4a9a303';
+const SECOND_IMAGE_GENERATION_TASK_ID = '1c939d2e-e7ea-480e-9dcf-4080df479fa3';
 const ONE_PIXEL_PNG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=';
 const ONE_PIXEL_PNG_DATA_URL = `data:image/png;base64,${ONE_PIXEL_PNG_BASE64}`;
 const ONE_PIXEL_SVG_BASE64 = Buffer.from('<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 16 16"><rect width="16" height="16" fill="black"/></svg>').toString('base64');
@@ -278,12 +279,16 @@ test.describe('ClawX chat run state events', () => {
   test('projects OpenClaw image-generation structured media into ACP Chat previews', async ({ launchElectronApp }) => {
     const app = await launchElectronApp({ skipSetup: true });
     const generatedPath = '/tmp/openclaw-generated-sky.png';
+    const secondGeneratedPath = '/tmp/openclaw-generated-dog.png';
 
     try {
       await installAcpChatMocks(
         app,
         { success: true, generation: 1 },
-        generatedImageHostApiMocks(generatedPath, 'e2e-live-generated-image'),
+        {
+          ...generatedImageHostApiMocks(generatedPath, 'e2e-live-generated-image'),
+          ...generatedImageHostApiMocks(secondGeneratedPath, 'e2e-second-live-generated-image'),
+        },
       );
       const page = await openChat(app);
       await expect(page.getByTestId('acp-chat-empty-state')).toBeVisible({ timeout: 30_000 });
@@ -336,6 +341,46 @@ test.describe('ClawX chat run state events', () => {
       await expect(image).toHaveAttribute('src', ONE_PIXEL_PNG_DATA_URL);
       await expect(page.getByTestId('chat-composer-image-generation-indicator')).toHaveCount(0);
       await expect(page.getByTestId('image-preview-unavailable')).toHaveCount(0);
+
+      await emitAcpSessionUpdates(app, [{
+        sessionUpdate: 'tool_call_update',
+        toolCallId: 'second-image-tool',
+        status: 'completed',
+        content: [{
+          type: 'content',
+          content: {
+            type: 'text',
+            text: `Background task started for image generation (${SECOND_IMAGE_GENERATION_TASK_ID}).`,
+          },
+        }],
+        locations: [],
+      }]);
+      await expect(page.getByTestId('chat-composer-image-generation-indicator')).toBeVisible();
+      await expect(page.getByText('Here is the exact sky scene you requested.')).toBeVisible();
+
+      await page.getByTestId('sidebar-new-chat').click();
+      await page.getByTestId(`sidebar-session-${MAIN_SESSION_KEY}`).click();
+      await expect(page.getByTestId('chat-composer-image-generation-indicator')).toBeVisible();
+      await expect(page.getByText('Here is the exact sky scene you requested.')).toBeVisible();
+      await expect(timeline.getByRole('img', { name: 'Image' })).toHaveCount(1);
+
+      await emitGatewayChatMessage(app, {
+        message: {
+          runId: `image_generate:${SECOND_IMAGE_GENERATION_TASK_ID}:ok`,
+          sessionKey: MAIN_SESSION_KEY,
+          state: 'final',
+          message: {
+            role: 'assistant',
+            content: [{
+              type: 'text',
+              text: `Here is the dog image you requested.\n\nMEDIA:${secondGeneratedPath}`,
+            }],
+          },
+        },
+      });
+      await expect(page.getByText('Here is the dog image you requested.')).toBeVisible();
+      await expect(timeline.getByRole('img', { name: 'Image' })).toHaveCount(2);
+      await expect(page.getByTestId('chat-composer-image-generation-indicator')).toHaveCount(0);
     } finally {
       await closeElectronApp(app);
     }

tests/unit/acp-chat-store.test.ts

@@ -1773,6 +1773,143 @@ describe('ACP Chat store', () => {
     expect(useAcpChatSessionStore.getState().pendingImageGenerationTaskIds).toEqual([]);
   });
 
+  it('keeps an earlier generated image while a later image task survives reload and navigation', async () => {
+    const firstTaskId = '32aa3a12-a05b-4074-af4e-246cc4a9a303';
+    const secondTaskId = '1c939d2e-e7ea-480e-9dcf-4080df479fa3';
+    const replayResult = (generation: number) => ({
+      success: true,
+      generation,
+      sessionUpdates: [{
+        sessionKey: 'agent:pi:s1',
+        generation,
+        historical: true,
+        notification: {
+          sessionId: 'agent:pi:s1',
+          update: {
+            sessionUpdate: 'user_message',
+            messageId: 'replayed-user',
+            content: [{ type: 'text', text: 'Generate an image' }],
+          },
+        },
+      }],
+    });
+    hostApiMock.loadAcpSession
+      .mockResolvedValueOnce({ success: true, generation: 1 })
+      .mockResolvedValueOnce(replayResult(2))
+      .mockResolvedValueOnce({ success: true, generation: 3 })
+      .mockResolvedValueOnce(replayResult(4));
+    hostApiMock.mediaThumbnails.mockResolvedValueOnce({
+      '/tmp/cat.png': { preview: 'data:image/png;base64,cat123', fileSize: 67 },
+    });
+    const { ensureAcpChatSubscriptions, useAcpChatSessionStore } = await importStore();
+    ensureAcpChatSubscriptions();
+    await useAcpChatSessionStore.getState().loadSession({
+      sessionKey: 'agent:pi:s1', workspaceRoot: '/repo', cwd: '/repo',
+    });
+
+    const recordStart = (taskId: string, toolCallId: string) => hostEventsMock.updateListener?.({
+      sessionKey: 'agent:pi:s1',
+      generation: 1,
+      notification: {
+        sessionId: 'agent:pi:s1',
+        update: {
+          sessionUpdate: 'tool_call_update',
+          toolCallId,
+          status: 'completed',
+          content: [{
+            type: 'content',
+            content: {
+              type: 'text',
+              text: `Background task started for image generation (${taskId}).`,
+            },
+          }],
+        },
+      },
+    });
+    recordStart(firstTaskId, 'image-tool-cat');
+    hostEventsMock.gatewayChatMessageListener?.({
+      message: {
+        sessionKey: 'agent:pi:s1',
+        runId: `image_generate:${firstTaskId}:ok`,
+        message: {
+          role: 'toolresult',
+          toolName: 'message',
+          details: { mediaUrls: ['/tmp/cat.png'] },
+        },
+      },
+    });
+    await vi.waitFor(() => expect(
+      useAcpChatSessionStore.getState().timeline.itemOrder.filter(
+        (id) => id.startsWith('compat:image-generation:'),
+      ),
+    ).toHaveLength(1));
+    const catImageItemId = useAcpChatSessionStore.getState().timeline.itemOrder.find(
+      (id) => id.startsWith('compat:image-generation:'),
+    )!;
+
+    recordStart(secondTaskId, 'image-tool-dog');
+    hostEventsMock.permissionListener?.({
+      sessionKey: 'agent:pi:s1',
+      generation: 1,
+      requestId: 'stale-permission',
+      request: {
+        sessionId: 'agent:pi:s1',
+        toolCall: { toolCallId: 'permission-tool', title: 'Old permission', status: 'pending' },
+        options: [{ optionId: 'allow-once', name: 'Allow once', kind: 'allow_once' }],
+      },
+    });
+    expect(useAcpChatSessionStore.getState().pendingImageGenerationTaskIds).toEqual([secondTaskId]);
+
+    await useAcpChatSessionStore.getState().loadSession({
+      sessionKey: 'agent:pi:s1', workspaceRoot: '/repo', cwd: '/repo',
+    });
+    expect(useAcpChatSessionStore.getState().timeline.itemOrder).toContain(catImageItemId);
+    expect(useAcpChatSessionStore.getState().timeline.itemOrder).toContain('replayed-user:0');
+    expect(useAcpChatSessionStore.getState().timeline.itemOrder).not.toContain('permission:stale-permission');
+
+    await useAcpChatSessionStore.getState().loadSession({
+      sessionKey: 'agent:pi:s2', workspaceRoot: '/repo-2', cwd: '/repo-2',
+    });
+    await useAcpChatSessionStore.getState().loadSession({
+      sessionKey: 'agent:pi:s1', workspaceRoot: '/repo', cwd: '/repo',
+    });
+    expect(useAcpChatSessionStore.getState().timeline.itemOrder).toContain(catImageItemId);
+    expect(useAcpChatSessionStore.getState().pendingImageGenerationTaskIds).toEqual([secondTaskId]);
+
+    const reload = createDeferred<ReturnType<typeof replayResult>>();
+    hostApiMock.loadAcpSession.mockReturnValueOnce(reload.promise);
+    hostApiMock.mediaThumbnails.mockResolvedValueOnce({
+      '/tmp/dog.png': { preview: 'data:image/png;base64,dog123', fileSize: 68 },
+    });
+    const reloadPromise = useAcpChatSessionStore.getState().loadSession({
+      sessionKey: 'agent:pi:s1', workspaceRoot: '/repo', cwd: '/repo',
+    });
+    await vi.waitFor(() => expect(useAcpChatSessionStore.getState().loading).toBe(true));
+    const attachmentCallCount = hostApiMock.resolveAttachment.mock.calls.length;
+    hostEventsMock.gatewayChatMessageListener?.({
+      message: {
+        sessionKey: 'agent:pi:s1',
+        runId: `image_generate:${secondTaskId}:ok`,
+        message: {
+          role: 'toolresult',
+          toolName: 'message',
+          details: { mediaUrls: ['/tmp/dog.png'] },
+        },
+      },
+    });
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(hostApiMock.resolveAttachment).toHaveBeenCalledTimes(attachmentCallCount);
+
+    reload.resolve(replayResult(5));
+    await expect(reloadPromise).resolves.toBe(true);
+    await vi.waitFor(() => expect(
+      useAcpChatSessionStore.getState().timeline.itemOrder.filter(
+        (id) => id.startsWith('compat:image-generation:'),
+      ),
+    ).toHaveLength(2));
+    expect(useAcpChatSessionStore.getState().pendingImageGenerationTaskIds).toEqual([]);
+  });
+
   it('records image-generation start detection trace entries', async () => {
     const taskId = '0d2ee919-2dfd-4b72-9da3-d87e6ee56747';
     const { ensureAcpChatSubscriptions, useAcpChatSessionStore } = await importStore();