From 2c9ca8434e1b7a6196889a78d681ed598f4ce0e2 Mon Sep 17 00:00:00 2001 From: kev1n77 Date: Mon, 21 Sep 2026 10:53:48 +0800 Subject: [PATCH] fix(web-ui): keep image attachments when editing a user message An edit-rerun is a rollback plus a fresh submission. The rerun only carried the text the edit composer owns, so the attachments stayed out of the new turn while the rollback had already deleted the turn that held them, and the images were gone for good. Rebuild the message's image contexts the way rollback and the failed-message refill already do, and pass them to the rerun as imageContexts/imageDisplayData. Drop canEditUserMessage and message.editDisabledImages with it. That gate forbade editing any message that carried images and had no caller; images stay uneditable, but they are carried through an edit. --- .../modern/UserMessageItem.test.tsx | 155 ++++++++++++++++++ .../components/modern/UserMessageItem.tsx | 20 ++- .../services/UserMessageEditService.ts | 20 --- src/web-ui/src/locales/en-US/flow-chat.json | 1 - src/web-ui/src/locales/zh-CN/flow-chat.json | 1 - src/web-ui/src/locales/zh-TW/flow-chat.json | 1 - 6 files changed, 173 insertions(+), 25 deletions(-) diff --git a/src/web-ui/src/flow_chat/components/modern/UserMessageItem.test.tsx b/src/web-ui/src/flow_chat/components/modern/UserMessageItem.test.tsx index 867006915c..96e06f34d8 100644 --- a/src/web-ui/src/flow_chat/components/modern/UserMessageItem.test.tsx +++ b/src/web-ui/src/flow_chat/components/modern/UserMessageItem.test.tsx @@ -25,6 +25,14 @@ const rollbackServiceMock = vi.hoisted(() => ({ composerText: 'restored prompt', })), })); +const imagePayloadMock = vi.hoisted(() => ({ + buildImagePayload: vi.fn(async (_contexts: unknown) => undefined as + | { imageContexts: unknown[]; imageDisplayData: unknown[] } + | undefined), +})); +const flowChatManagerMock = vi.hoisted(() => ({ + sendMessage: vi.fn(async () => undefined), +})); const componentLibraryMock = vi.hoisted(() => ({ confirmDanger: vi.fn(async () => true), })); @@ -151,6 +159,16 @@ vi.mock('../../store/FlowChatStore', () => ({ vi.mock('../../services/SessionRollbackService', () => rollbackServiceMock); +vi.mock('../../utils/imagePayload', async importOriginal => ({ + ...await importOriginal(), + buildImagePayload: imagePayloadMock.buildImagePayload, +})); + +vi.mock('../../services/FlowChatManager', async importOriginal => ({ + ...await importOriginal(), + flowChatManager: flowChatManagerMock, +})); + vi.mock('@/shared/notification-system', () => ({ notificationService: { success: vi.fn(), @@ -211,6 +229,8 @@ describe('UserMessageItem steering tag', () => { composerText: 'restored prompt', }); editServiceMock.editAndRerunUserMessage.mockResolvedValue(undefined); + imagePayloadMock.buildImagePayload.mockResolvedValue(undefined); + flowChatManagerMock.sendMessage.mockResolvedValue(undefined); stateMachineManager.clear(); useMessageEditStore.getState().cancelEdit(); dom = new JSDOM('
', { @@ -1059,4 +1079,139 @@ describe('UserMessageItem steering tag', () => { })); expect(flowChatStoreMock.loadSessionHistory).not.toHaveBeenCalled(); }); + + it('resubmits the original image attachments when editing a message', async () => { + activeSessionRef.current = { + sessionId: 'main-session', + sessionKind: 'normal', + workspaceId: workspaceRecords.local.id, + dialogTurns: [{ id: 'turn-1', status: 'completed' }], + }; + imagePayloadMock.buildImagePayload.mockResolvedValueOnce({ + imageContexts: [{ id: 'image-1', image_path: 'E:/uploads/preview.png', mime_type: 'image/png' }], + imageDisplayData: [{ id: 'image-1', name: 'preview.png', imagePath: 'E:/uploads/preview.png' }], + }); + editServiceMock.editAndRerunUserMessage.mockImplementationOnce(async (request: any) => { + await request.rerun(request.editedContent, request.agentType, 'lease-1'); + return undefined; + }); + + act(() => { + root.render( + + + , + ); + }); + + await act(async () => { + container.querySelector('.user-message-item__edit-btn')?.click(); + }); + await act(async () => { + useMessageEditStore.getState().setDraft('Describe this image in detail'); + }); + await act(async () => { + container + .querySelector('.user-message-edit-composer__icon-button--confirm') + ?.click(); + await Promise.resolve(); + await Promise.resolve(); + }); + + expect(imagePayloadMock.buildImagePayload).toHaveBeenCalledWith([ + expect.objectContaining({ + id: 'image-1', + type: 'image', + imageName: 'preview.png', + imagePath: 'E:/uploads/preview.png', + mimeType: 'image/png', + isLocal: true, + }), + ]); + expect(flowChatManagerMock.sendMessage).toHaveBeenCalledWith( + 'Describe this image in detail', + 'main-session', + undefined, + undefined, + undefined, + expect.objectContaining({ + imageContexts: [{ id: 'image-1', image_path: 'E:/uploads/preview.png', mime_type: 'image/png' }], + imageDisplayData: [{ id: 'image-1', name: 'preview.png', imagePath: 'E:/uploads/preview.png' }], + sessionMutationLeaseId: 'lease-1', + }), + ); + }); + + it('reruns an edit without attachments when the message has no images', async () => { + activeSessionRef.current = { + sessionId: 'main-session', + sessionKind: 'normal', + workspaceId: workspaceRecords.local.id, + dialogTurns: [{ id: 'turn-1', status: 'completed' }], + }; + editServiceMock.editAndRerunUserMessage.mockImplementationOnce(async (request: any) => { + await request.rerun(request.editedContent, request.agentType, 'lease-2'); + return undefined; + }); + + act(() => { + root.render( + + + , + ); + }); + + await act(async () => { + container.querySelector('.user-message-item__edit-btn')?.click(); + }); + await act(async () => { + useMessageEditStore.getState().setDraft('Describe this more precisely'); + }); + await act(async () => { + container + .querySelector('.user-message-edit-composer__icon-button--confirm') + ?.click(); + await Promise.resolve(); + await Promise.resolve(); + }); + + expect(imagePayloadMock.buildImagePayload).toHaveBeenCalledWith([]); + expect(flowChatManagerMock.sendMessage).toHaveBeenCalledWith( + 'Describe this more precisely', + 'main-session', + undefined, + undefined, + undefined, + { sessionMutationLeaseId: 'lease-2' }, + ); + }); }); diff --git a/src/web-ui/src/flow_chat/components/modern/UserMessageItem.tsx b/src/web-ui/src/flow_chat/components/modern/UserMessageItem.tsx index 8f3c746363..cd22235689 100644 --- a/src/web-ui/src/flow_chat/components/modern/UserMessageItem.tsx +++ b/src/web-ui/src/flow_chat/components/modern/UserMessageItem.tsx @@ -8,6 +8,7 @@ import { withConversationExcerptFallback } from '@/shared/utils/conversationExce import { getAppearanceOverlayHost } from '@/infrastructure/appearance/runtime/AppearanceOverlayHost'; import { RotateCcw, Loader2 } from 'lucide-react'; import type { DialogTurn, FlowUserSteeringItem } from '../../types/flow-chat'; +import type { ImageContext } from '@/shared/types/context'; import { flowChatManager } from '../../services/FlowChatManager'; import { useFlowChatContext } from './FlowChatContext'; import { useActiveSession } from '../../store/modernFlowChatStore'; @@ -55,6 +56,7 @@ import { type ComposerPresentation, } from '../../utils/composerPresentation'; import { restoreImageContextsFromPayload } from '../../utils/imageContextRestoration'; +import { buildImagePayload } from '../../utils/imagePayload'; import { UserMessagePresentationContent, UserMessageTextContent } from './UserMessagePresentationContent'; import { UserMessageImage } from './UserMessageImage'; import { useSubmittedMessageMotion } from './useSubmittedMessageMotion'; @@ -161,6 +163,15 @@ export const UserMessageItem = React.memo( imageDisplayData: messageImages, }), ], [composerPresentation, message?.id, message?.timestamp, messageImages, turnId]); + // The edit composer only owns text and keeps its attachments outside the + // editor, so an edit-rerun resubmits the original images instead of losing + // them with the rolled-back Turn. + const restoredImageContexts = useMemo( + () => restoredComposerContexts.filter( + (context): context is ImageContext => context.type === 'image', + ), + [restoredComposerContexts], + ); const isUsageReportMessage = message?.metadata?.localCommandKind === 'usage_report'; const isGoalLoadingMessage = Boolean(message?.metadata?.threadGoalKickoff); const isThreadGoalContinuationCheck = Boolean(message?.metadata?.threadGoalContinuation); @@ -405,7 +416,10 @@ export const UserMessageItem = React.memo( originalContent: messageContent, editedContent, agentType: currentSession?.mode, - rerun: (content, agentType, sessionMutationLeaseId) => { + rerun: async (content, agentType, sessionMutationLeaseId) => { + const imagePayload = await buildImagePayload(restoredImageContexts); + const attachments = imagePayload ?? {}; + if (!editedPresentation) { return flowChatManager.sendMessage( content, @@ -413,7 +427,7 @@ export const UserMessageItem = React.memo( undefined, agentType, undefined, - { sessionMutationLeaseId }, + { ...attachments, sessionMutationLeaseId }, ); } @@ -426,6 +440,7 @@ export const UserMessageItem = React.memo( undefined, { userMessageMetadata: payload.userMessageMetadata, + ...attachments, sessionMutationLeaseId, }, ); @@ -447,6 +462,7 @@ export const UserMessageItem = React.memo( editDraft, isEditSubmitting, messageContent, + restoredImageContexts, resolvedSessionId, setEditSubmitting, t, diff --git a/src/web-ui/src/flow_chat/services/UserMessageEditService.ts b/src/web-ui/src/flow_chat/services/UserMessageEditService.ts index 253008581d..78393470d6 100644 --- a/src/web-ui/src/flow_chat/services/UserMessageEditService.ts +++ b/src/web-ui/src/flow_chat/services/UserMessageEditService.ts @@ -34,26 +34,6 @@ export function describeUserMessageEditImpact(_sessionId: string): UserMessageEd }; } -export function canEditUserMessage(request: { - sessionId?: string | null; - turnIndex: number; - hasImages?: boolean; - isUsageReportMessage?: boolean; - steeringStatus?: string; - isRemoteSession?: boolean; - isSubmitting?: boolean; -}): boolean { - return Boolean( - request.sessionId && - request.turnIndex >= 0 && - !request.hasImages && - !request.isUsageReportMessage && - !request.steeringStatus && - !request.isRemoteSession && - !request.isSubmitting, - ); -} - export async function editAndRerunUserMessage( request: EditAndRerunUserMessageRequest, ): Promise { diff --git a/src/web-ui/src/locales/en-US/flow-chat.json b/src/web-ui/src/locales/en-US/flow-chat.json index 99f83a969a..732ab41653 100644 --- a/src/web-ui/src/locales/en-US/flow-chat.json +++ b/src/web-ui/src/locales/en-US/flow-chat.json @@ -528,7 +528,6 @@ "editPlaceholder": "Edit your message...", "cannotEdit": "Cannot edit this message", "editDisabledBusy": "Wait until the session is idle before editing and rerunning.", - "editDisabledImages": "Image messages cannot be edited yet", "editDisabledRemote": "Editing and rerunning requires complete rollback, which is unavailable because remote file snapshots are not recorded.", "editDisabledDispatch": "Editing and rerunning requires history rollback, which is not yet available for remote tasks.", "editDisabledLocalCommand": "Local command messages cannot be edited", diff --git a/src/web-ui/src/locales/zh-CN/flow-chat.json b/src/web-ui/src/locales/zh-CN/flow-chat.json index eb738033fc..692e0f4bff 100644 --- a/src/web-ui/src/locales/zh-CN/flow-chat.json +++ b/src/web-ui/src/locales/zh-CN/flow-chat.json @@ -528,7 +528,6 @@ "editPlaceholder": "编辑你的消息...", "cannotEdit": "无法编辑该消息", "editDisabledBusy": "请等待会话空闲后再编辑并重跑。", - "editDisabledImages": "暂不支持编辑图片消息", "editDisabledRemote": "编辑并重跑需要完整回滚；远程工作区未记录文件快照，因此暂不支持。", "editDisabledDispatch": "编辑并重跑需要先回滚历史，远程任务暂不支持此操作。", "editDisabledLocalCommand": "无法编辑本地命令消息", diff --git a/src/web-ui/src/locales/zh-TW/flow-chat.json b/src/web-ui/src/locales/zh-TW/flow-chat.json index b8b6ec2a4c..6fe0e752a6 100644 --- a/src/web-ui/src/locales/zh-TW/flow-chat.json +++ b/src/web-ui/src/locales/zh-TW/flow-chat.json @@ -528,7 +528,6 @@ "editPlaceholder": "編輯你的消息...", "cannotEdit": "無法編輯該消息", "editDisabledBusy": "請等待會話閒置後再編輯並重跑。", - "editDisabledImages": "暫不支援編輯圖片消息", "editDisabledRemote": "編輯並重新執行需要完整回滾；遠端工作區未記錄檔案快照，因此暫不支援。", "editDisabledDispatch": "編輯並重新執行需要先回滾歷史，遠端任務暫不支援此操作。", "editDisabledLocalCommand": "無法編輯本機命令消息",