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": "無法編輯本機命令消息",