diff --git a/frontend/e2e/chat.spec.ts b/frontend/e2e/chat.spec.ts index 87879faeb4..cdf7d8c358 100644 --- a/frontend/e2e/chat.spec.ts +++ b/frontend/e2e/chat.spec.ts @@ -1,5 +1,5 @@ import { readFileSync } from "node:fs"; -import { test, expect, type Page } from "@playwright/test"; +import { test, expect, type Locator, type Page } from "@playwright/test"; import { makeTarget } from "./_targets"; // --------------------------------------------------------------------------- @@ -10,6 +10,10 @@ const MOCK_CONVERSATION_ID = "e2e-conv-001"; const WIDE_IMAGE_DATA_URI = "data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='800' height='600' viewBox='0 0 800 600'%3E%3Crect width='800' height='600' fill='%230078d4'/%3E%3C/svg%3E"; +function getMessageByText(page: Page, text: string): Locator { + return page.getByTestId("message-list").getByText(text, { exact: true }); +} + /** Intercept targets & attacks APIs so the chat flow can run without real keys. */ async function mockBackendAPIs(page: Page) { // Accumulate messages so multi-turn tests get full history back @@ -236,7 +240,7 @@ test.describe("Chat Functionality", () => { const input = page.getByRole("textbox"); await input.fill("Start a mobile conversation"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("Start a mobile conversation", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "Start a mobile conversation")).toBeVisible(); await page.setViewportSize({ width: 390, height: 844 }); const chatArea = page.getByTestId("chat-area"); @@ -277,7 +281,7 @@ test.describe("Chat Functionality", () => { await page.getByRole("button", { name: /send/i }).click(); // User message appears - await expect(page.getByText("Hello, this is a test message", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "Hello, this is a test message")).toBeVisible(); // Backend response appears await expect( @@ -313,7 +317,7 @@ test.describe("Chat Functionality", () => { await input.fill("First message"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("First message", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "First message")).toBeVisible(); await expect( page.getByText("Mock response for: First message"), ).toBeVisible({ timeout: 10000 }); @@ -339,7 +343,7 @@ test.describe("Multiple Messages", () => { // Send first message await input.fill("First message"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("First message", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "First message")).toBeVisible(); await expect( page.getByText("Mock response for: First message"), ).toBeVisible({ timeout: 10000 }); @@ -347,14 +351,14 @@ test.describe("Multiple Messages", () => { // Send second message await input.fill("Second message"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("Second message", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "Second message")).toBeVisible(); await expect( page.getByText("Mock response for: Second message"), ).toBeVisible({ timeout: 10000 }); // Both user messages should still be visible - await expect(page.getByText("First message", { exact: true })).toBeVisible(); - await expect(page.getByText("Second message", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "First message")).toBeVisible(); + await expect(getMessageByText(page, "Second message")).toBeVisible(); }); }); @@ -505,7 +509,7 @@ test.describe("Multi-modal: Image response", () => { await page.getByRole("button", { name: /send/i }).click(); // User message visible - await expect(page.getByText("Generate an image", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "Generate an image")).toBeVisible(); // Image element should appear (exclude logo) const img = page.locator('img:not([alt="Co-PyRIT Logo"])'); @@ -603,7 +607,7 @@ test.describe("Multi-modal: Audio response", () => { await input.fill("Speak this out loud"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("Speak this out loud", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "Speak this out loud")).toBeVisible(); // Audio element should appear const audio = page.locator("audio"); @@ -666,7 +670,7 @@ test.describe("Multi-modal: Video response", () => { await input.fill("Create a video clip"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("Create a video clip", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "Create a video clip")).toBeVisible(); // Video element should appear const video = page.locator("video"); @@ -707,7 +711,7 @@ test.describe("Multi-modal: Mixed text + image response", () => { await page.getByRole("button", { name: /send/i }).click(); // Both text and image should be visible - await expect(page.getByText("Here is the analysis:", { exact: true })).toBeVisible({ timeout: 10000 }); + await expect(getMessageByText(page, "Here is the analysis:")).toBeVisible({ timeout: 10000 }); const img = page.locator('img:not([alt="Co-PyRIT Logo"])'); await expect(img).toBeVisible({ timeout: 10000 }); }); @@ -736,7 +740,7 @@ test.describe("Multi-modal: Error response from target", () => { await input.fill("unsafe prompt"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("unsafe prompt", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "unsafe prompt")).toBeVisible(); // Error should be displayed await expect( @@ -758,7 +762,7 @@ test.describe("Multi-turn conversation flow", () => { // Turn 1 await input.fill("First turn"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("First turn", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "First turn")).toBeVisible(); await expect( page.getByText("Mock response for: First turn"), ).toBeVisible({ timeout: 10000 }); @@ -766,7 +770,7 @@ test.describe("Multi-turn conversation flow", () => { // Turn 2 await input.fill("Second turn"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("Second turn", { exact: true })).toBeVisible({ timeout: 10000 }); + await expect(getMessageByText(page, "Second turn")).toBeVisible({ timeout: 10000 }); await expect( page.getByText("Mock response for: Second turn"), ).toBeVisible({ timeout: 10000 }); @@ -774,15 +778,15 @@ test.describe("Multi-turn conversation flow", () => { // Turn 3 await input.fill("Third turn"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("Third turn", { exact: true })).toBeVisible({ timeout: 10000 }); + await expect(getMessageByText(page, "Third turn")).toBeVisible({ timeout: 10000 }); await expect( page.getByText("Mock response for: Third turn"), ).toBeVisible({ timeout: 10000 }); // All previous messages still visible - await expect(page.getByText("First turn", { exact: true })).toBeVisible(); - await expect(page.getByText("Second turn", { exact: true })).toBeVisible(); - await expect(page.getByText("Third turn", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "First turn")).toBeVisible(); + await expect(getMessageByText(page, "Second turn")).toBeVisible(); + await expect(getMessageByText(page, "Third turn")).toBeVisible(); }); test("should reset conversation on New Chat and send again", async ({ page }) => { @@ -791,19 +795,19 @@ test.describe("Multi-turn conversation flow", () => { // Send a message await input.fill("Before reset"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("Before reset", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "Before reset")).toBeVisible(); await expect( page.getByText("Mock response for: Before reset"), ).toBeVisible({ timeout: 10000 }); // New Attack await page.getByTestId("new-attack-btn").click(); - await expect(page.getByText("Before reset", { exact: true })).not.toBeVisible(); + await expect(getMessageByText(page, "Before reset")).not.toBeVisible(); // Send new message in fresh conversation await input.fill("After reset"); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.getByText("After reset", { exact: true })).toBeVisible(); + await expect(getMessageByText(page, "After reset")).toBeVisible(); await expect( page.getByText("Mock response for: After reset"), ).toBeVisible({ timeout: 10000 }); diff --git a/frontend/e2e/errors.spec.ts b/frontend/e2e/errors.spec.ts index 84afa9953c..625a35597b 100644 --- a/frontend/e2e/errors.spec.ts +++ b/frontend/e2e/errors.spec.ts @@ -13,7 +13,7 @@ function buildSuccessMessageMock(userText: string) { messages: { messages: [ { - turn_number: 1, + turn_number: 0, role: "user", created_at: new Date().toISOString(), message_pieces: [ @@ -49,6 +49,48 @@ function buildSuccessMessageMock(userText: string) { }; } +function buildProcessingFailureMock(userText: string) { + return { + messages: { + messages: [ + { + turn_number: 2, + role: "user", + created_at: new Date().toISOString(), + message_pieces: [ + { + id: "p-processing-user", + original_value_data_type: "text", + converted_value_data_type: "text", + original_value: userText, + converted_value: userText, + scores: [], + response_error: "none", + }, + ], + }, + { + turn_number: 3, + role: "assistant", + created_at: new Date().toISOString(), + message_pieces: [ + { + id: "p-processing-error", + original_value_data_type: "text", + converted_value_data_type: "error", + original_value: "", + converted_value: + "RuntimeError: target failed\nTraceback (most recent call last): ...", + scores: [], + response_error: "processing", + }, + ], + }, + ], + }, + }; +} + /** * Set up all the mocks needed for a full chat flow. * @@ -206,6 +248,107 @@ async function triggerVisibilityChange(page: Page) { }); } +// --------------------------------------------------------------------------- +// Error scenario: persisted target processing failure returned with HTTP 200 +// --------------------------------------------------------------------------- + +test.describe("Error: target processing failure returned with HTTP 200", () => { + test("should preserve the draft and offer edit recovery", async ({ page }) => { + let callCount = 0; + let recoveryRequest: Record | undefined; + let recoveryCreated = false; + await mockAllAPIs(page, async (route) => { + const body = JSON.parse(route.request().postData() ?? "{}"); + const userText = + body?.pieces?.find( + (piece: Record) => piece.data_type === "text", + )?.original_value || "message"; + callCount++; + const response = callCount === 1 + ? buildSuccessMessageMock(userText) + : buildProcessingFailureMock(userText); + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify(response), + }); + }); + await page.route(/\/api\/attacks\/[^/]+\/conversations/, async (route) => { + if (route.request().method() === "GET" && recoveryCreated) { + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + attack_result_id: "err-ar-001", + main_conversation_id: MOCK_CONV_ID, + conversations: [ + { + conversation_id: MOCK_CONV_ID, + message_count: 4, + last_message_preview: "Target response error", + created_at: "2026-01-01T00:00:00.000Z", + }, + { + conversation_id: "err-conv-recovery", + message_count: 2, + last_message_preview: "Reply to: Setup message", + created_at: "2026-01-01T00:00:01.000Z", + }, + ], + }), + }); + return; + } + if (route.request().method() !== "POST") { + await route.fallback(); + return; + } + recoveryRequest = JSON.parse(route.request().postData() ?? "{}"); + recoveryCreated = true; + await route.fulfill({ + status: 200, + contentType: "application/json", + body: JSON.stringify({ + conversation_id: "err-conv-recovery", + created_at: "2026-01-01T00:00:01.000Z", + }), + }); + }); + + await page.goto("/"); + await activateMockTarget(page); + await sendAndWait(page, "Setup message", "Reply to: Setup message"); + + const input = page.getByRole("textbox"); + await input.fill("Preserve this draft"); + await page.getByRole("button", { name: /send/i }).click(); + + await expect( + page.getByText(/target could not process this message/i), + ).toBeVisible(); + const recoveryButton = page.getByRole("button", { + name: /edit in clean conversation/i, + }); + await expect(recoveryButton).toHaveCount(1); + await expect(input).toHaveValue("Preserve this draft"); + await expect(input).toBeDisabled(); + await expect(page.getByText(/Traceback \(most recent call last\)/i)).toHaveCount(0); + + await recoveryButton.click(); + await expect.poll(() => recoveryRequest).toEqual({ + source_conversation_id: MOCK_CONV_ID, + cutoff_index: 1, + }); + await expect(input).toBeFocused(); + await expect(input).toHaveValue("Preserve this draft"); + await expect(page.getByTestId(`conversation-item-${MOCK_CONV_ID}`)).toBeVisible(); + await expect( + page.getByTestId("conversation-item-err-conv-recovery"), + ).toBeVisible(); + expect(callCount).toBe(2); + }); +}); + // --------------------------------------------------------------------------- // Error scenario: backend returns 500 on send message // --------------------------------------------------------------------------- diff --git a/frontend/src/components/Chat/ChatInputArea.test.tsx b/frontend/src/components/Chat/ChatInputArea.test.tsx index 1c5c601b2a..0136ffee22 100644 --- a/frontend/src/components/Chat/ChatInputArea.test.tsx +++ b/frontend/src/components/Chat/ChatInputArea.test.tsx @@ -30,8 +30,13 @@ const buildCapabilities = ( const getSendButton = () => screen.getByRole("button", { name: /send/i }); describe("ChatInputArea", () => { + const sentOutcome = { status: "sent", clearDraft: true } as const; + const retryableOutcome = { + status: "retryable_failure", + clearDraft: false, + } as const; const defaultProps = { - onSend: jest.fn(), + onSend: jest.fn().mockResolvedValue(sentOutcome), disabled: false, onNewConversation: jest.fn(), onUseAsTemplate: jest.fn(), @@ -41,6 +46,7 @@ describe("ChatInputArea", () => { onInputChange: jest.fn(), onAttachmentsChange: jest.fn(), onClearConversion: jest.fn(), + onClearAllConversions: jest.fn(), onConvertedValueChange: jest.fn(), onClearMediaConversion: jest.fn(), }; @@ -159,7 +165,7 @@ describe("ChatInputArea", () => { it("should call onSend with input value when send button clicked", async () => { const user = userEvent.setup(); - const onSend = jest.fn(); + const onSend = jest.fn().mockResolvedValue(sentOutcome); render( @@ -212,9 +218,9 @@ describe("ChatInputArea", () => { expect(sendButton).toBeEnabled(); }); - it("should clear input after sending", async () => { + it("should clear input after a successful send", async () => { const user = userEvent.setup(); - const onSend = jest.fn(); + const onSend = jest.fn().mockResolvedValue(sentOutcome); render( @@ -231,9 +237,174 @@ describe("ChatInputArea", () => { }); }); + it("should preserve the draft when the send finishes in another conversation", async () => { + const user = userEvent.setup(); + const onClearAllConversions = jest.fn(); + const onSend = jest.fn().mockResolvedValue({ + status: "sent", + clearDraft: false, + }); + + render( + + + + ); + + const input = screen.getByRole("textbox"); + await user.type(input, "keep this draft"); + await user.click(getSendButton()); + + await waitFor(() => expect(onSend).toHaveBeenCalledTimes(1)); + expect(input).toHaveValue("keep this draft"); + expect(onClearAllConversions).not.toHaveBeenCalled(); + }); + + it("should not clear a newer draft when an earlier send completes", async () => { + const user = userEvent.setup(); + const onClearAllConversions = jest.fn(); + let resolveSend: ((outcome: typeof sentOutcome) => void) | undefined; + const onSend = jest.fn( + () => new Promise((resolve) => { + resolveSend = resolve; + }) + ); + + render( + + + + ); + + const input = screen.getByRole("textbox"); + await user.type(input, "conversation A draft"); + await user.click(getSendButton()); + await waitFor(() => expect(onSend).toHaveBeenCalledTimes(1)); + + await user.clear(input); + await user.type(input, "conversation B draft"); + resolveSend?.(sentOutcome); + + await waitFor(() => { + expect(input).toHaveValue("conversation B draft"); + }); + expect(onClearAllConversions).not.toHaveBeenCalled(); + }); + + it("should not clear a draft whose converter selection changed during a send", async () => { + const user = userEvent.setup(); + const onClearAllConversions = jest.fn(); + let resolveSend: ((outcome: typeof sentOutcome) => void) | undefined; + const onSend = jest.fn( + () => new Promise((resolve) => { + resolveSend = resolve; + }) + ); + + const rendered = render( + + + + ); + + const input = screen.getByRole("textbox"); + await user.type(input, "draft with converter"); + await user.click(getSendButton()); + await waitFor(() => expect(onSend).toHaveBeenCalledTimes(1)); + + rendered.rerender( + + + + ); + resolveSend?.(sentOutcome); + + await waitFor(() => { + expect(input).toHaveValue("draft with converter"); + }); + expect(onClearAllConversions).not.toHaveBeenCalled(); + }); + + it("should preserve the complete draft after a retryable send failure", async () => { + const user = userEvent.setup(); + const onSend = jest.fn().mockResolvedValue(retryableOutcome); + const onClearConversion = jest.fn(); + const onClearAllConversions = jest.fn(); + const onClearMediaConversion = jest.fn(); + const file = new File(["image"], "photo.png", { type: "image/png" }); + + render( + + + + ); + + const input = screen.getByPlaceholderText("Type prompt here"); + await user.type(input, "original prompt"); + const fileInput = document.querySelector('input[type="file"]') as HTMLInputElement; + await user.upload(fileInput, file); + await screen.findByText("photo.png", { exact: false }); + + await user.click(getSendButton()); + + await waitFor(() => { + expect(onSend).toHaveBeenCalledTimes(1); + expect(input).toHaveValue("original prompt"); + }); + expect(screen.getByText("photo.png", { exact: false })).toBeInTheDocument(); + expect(screen.getByTestId("converted-value-input")).toHaveValue("converted prompt"); + expect(screen.getByText("converted.png")).toBeInTheDocument(); + expect(onClearConversion).not.toHaveBeenCalled(); + expect(onClearAllConversions).not.toHaveBeenCalled(); + expect(onClearMediaConversion).not.toHaveBeenCalled(); + }); + it("should send message on Enter key press", async () => { const user = userEvent.setup(); - const onSend = jest.fn(); + const onSend = jest.fn().mockResolvedValue(sentOutcome); render( @@ -250,7 +421,7 @@ describe("ChatInputArea", () => { it("should not send on Shift+Enter (allows multiline)", async () => { const user = userEvent.setup(); - const onSend = jest.fn(); + const onSend = jest.fn().mockResolvedValue(sentOutcome); render( @@ -267,7 +438,7 @@ describe("ChatInputArea", () => { it("should allow sending whitespace-only messages", async () => { const user = userEvent.setup(); - const onSend = jest.fn(); + const onSend = jest.fn().mockResolvedValue(sentOutcome); render( @@ -283,7 +454,7 @@ describe("ChatInputArea", () => { }); it("should not send when input is completely empty", () => { - const onSend = jest.fn(); + const onSend = jest.fn().mockResolvedValue(sentOutcome); render( @@ -400,7 +571,7 @@ describe("ChatInputArea", () => { it("should send with attachments even without text", async () => { const user = userEvent.setup(); - const onSend = jest.fn(); + const onSend = jest.fn().mockResolvedValue(sentOutcome); render( @@ -764,8 +935,8 @@ describe("ChatInputArea", () => { }); it("should pass convertedValue to onSend when sending with conversion", async () => { - const onSend = jest.fn(); - const onClearConversion = jest.fn(); + const onSend = jest.fn().mockResolvedValue(sentOutcome); + const onClearAllConversions = jest.fn(); const user = userEvent.setup(); render( @@ -776,7 +947,7 @@ describe("ChatInputArea", () => { activeTarget={makeTarget({ target_registry_name: "t", target_type: "T", endpoint: "e", model_name: "m" })} convertedValue="convertedHello" originalValue="hello" - onClearConversion={onClearConversion} + onClearAllConversions={onClearAllConversions} /> ); @@ -786,7 +957,7 @@ describe("ChatInputArea", () => { await user.click(getSendButton()); expect(onSend).toHaveBeenCalledWith("hello", "convertedHello", []); - expect(onClearConversion).toHaveBeenCalled(); + expect(onClearAllConversions).toHaveBeenCalled(); }); it("should render converted file chip with Open link for text→file conversion", async () => { diff --git a/frontend/src/components/Chat/ChatInputArea.tsx b/frontend/src/components/Chat/ChatInputArea.tsx index db5eba03cd..8d67f65ca4 100644 --- a/frontend/src/components/Chat/ChatInputArea.tsx +++ b/frontend/src/components/Chat/ChatInputArea.tsx @@ -8,7 +8,7 @@ import { mergeClasses, } from '@fluentui/react-components' import { SendRegular, AttachRegular, DismissRegular, InfoRegular, AddRegular, CopyRegular, WarningRegular, SettingsRegular, ArrowShuffleRegular, OpenRegular, ArrowSyncRegular } from '@fluentui/react-icons' -import type { AttackTargetResolutionStatus, MessageAttachment, TargetInstance } from '../../types' +import type { AttackTargetResolutionStatus, ChatSendOutcome, MessageAttachment, TargetInstance } from '../../types' import { isTargetResolutionBlocking } from '../../utils/targetIdentity' import { useChatInputAreaStyles } from './ChatInputArea.styles' import SystemPromptSetup from './SystemPromptSetup' @@ -402,10 +402,18 @@ const formatModalityLabel = (modality: string): string => modality.replace('_pat export interface ChatInputAreaHandle { addAttachment: (att: MessageAttachment) => void setText: (text: string) => void + restoreDraft: (text: string, attachments: MessageAttachment[]) => void + focus: () => void + getDraftRevision: () => number } interface ChatInputAreaProps { - onSend: (originalValue: string, convertedValue: string | undefined, attachments: MessageAttachment[]) => void + onSend: ( + originalValue: string, + convertedValue: string | undefined, + attachments: MessageAttachment[], + ) => Promise + conversionRevisionKey?: string disabled?: boolean activeTarget?: TargetInstance | null singleTurnLimitReached?: boolean @@ -425,6 +433,7 @@ interface ChatInputAreaProps { convertedValue?: string | null originalValue?: string | null onClearConversion: () => void + onClearAllConversions?: () => void onConvertedValueChange: (value: string) => void converterOutputDataTypes?: string[] mediaConversions?: Array<{ pieceType: string; convertedValue: string; convertedDataType: string }> @@ -440,13 +449,24 @@ interface ChatInputAreaProps { onSystemPromptChange?: (value: string) => void } -const ChatInputArea = forwardRef(function ChatInputArea({ onSend, disabled = false, activeTarget, singleTurnLimitReached = false, onNewConversation, operatorLocked = false, crossTargetLocked = false, targetResolutionStatus = 'idle', onRetryTargetResolution, onUseAsTemplate, attackOperator, noTargetSelected = false, onConfigureTarget, onToggleConverterPanel, isConverterPanelOpen = false, onInputChange, onAttachmentsChange, convertedValue, originalValue: _originalValue, onClearConversion, onConvertedValueChange, converterOutputDataTypes = [], mediaConversions = [], onClearMediaConversion, convertedFileChip, onClearConvertedFileChip, showSystemPrompt = false, supportsSystemPrompt = false, systemPrompt = '', onSystemPromptChange }, ref) { +const ChatInputArea = forwardRef(function ChatInputArea({ onSend, conversionRevisionKey = '', disabled = false, activeTarget, singleTurnLimitReached = false, onNewConversation, operatorLocked = false, crossTargetLocked = false, targetResolutionStatus = 'idle', onRetryTargetResolution, onUseAsTemplate, attackOperator, noTargetSelected = false, onConfigureTarget, onToggleConverterPanel, isConverterPanelOpen = false, onInputChange, onAttachmentsChange, convertedValue, originalValue: _originalValue, onClearConversion, onClearAllConversions = () => {}, onConvertedValueChange, converterOutputDataTypes = [], mediaConversions = [], onClearMediaConversion, convertedFileChip, onClearConvertedFileChip, showSystemPrompt = false, supportsSystemPrompt = false, systemPrompt = '', onSystemPromptChange }, ref) { const styles = useChatInputAreaStyles() const [input, setInput] = useState('') const [attachments, setAttachments] = useState([]) const fileInputRef = useRef(null) const textareaRef = useRef(null) const convertedRef = useRef(null) + const inputRef = useRef('') + const attachmentsRef = useRef([]) + const draftRevisionRef = useRef(0) + const previousConversionRevisionKeyRef = useRef(conversionRevisionKey) + + useLayoutEffect(() => { + if (previousConversionRevisionKeyRef.current !== conversionRevisionKey) { + previousConversionRevisionKeyRef.current = conversionRevisionKey + draftRevisionRef.current += 1 + } + }, [conversionRevisionKey]) // Derive unsupported types from attachments AND converter outputs const unsupportedAttachmentTypes = getUnsupportedAttachmentTypes(attachments, activeTarget) @@ -461,11 +481,28 @@ const ChatInputArea = forwardRef(functi useImperativeHandle(ref, () => ({ addAttachment: (att: MessageAttachment) => { - setAttachments(prev => [...prev, att]) + const nextAttachments = [...attachmentsRef.current, att] + attachmentsRef.current = nextAttachments + draftRevisionRef.current += 1 + setAttachments(nextAttachments) }, setText: (text: string) => { + inputRef.current = text + draftRevisionRef.current += 1 setInput(text) }, + restoreDraft: (text: string, draftAttachments: MessageAttachment[]) => { + const restoredAttachments = draftAttachments.map((attachment) => ({ ...attachment })) + inputRef.current = text + attachmentsRef.current = restoredAttachments + draftRevisionRef.current += 1 + setInput(text) + setAttachments(restoredAttachments) + }, + focus: () => { + textareaRef.current?.focus() + }, + getDraftRevision: () => draftRevisionRef.current, })) const handleFileSelect = async (e: React.ChangeEvent) => { @@ -493,27 +530,47 @@ const ChatInputArea = forwardRef(functi }) } - setAttachments([...attachments, ...newAttachments]) + const nextAttachments = [...attachmentsRef.current, ...newAttachments] + attachmentsRef.current = nextAttachments + draftRevisionRef.current += 1 + setAttachments(nextAttachments) if (fileInputRef.current) { fileInputRef.current.value = '' } } const removeAttachment = (index: number) => { - const newAttachments = [...attachments] + const newAttachments = [...attachmentsRef.current] URL.revokeObjectURL(newAttachments[index].url) newAttachments.splice(index, 1) + attachmentsRef.current = newAttachments + draftRevisionRef.current += 1 setAttachments(newAttachments) } - const handleSend = () => { - if ((input || attachments.length > 0) && !disabled && !hasUnsupportedModalities) { - onSend(input, convertedValue ?? undefined, attachments) - setInput('') - setAttachments([]) - onClearConversion() - if (textareaRef.current) { - textareaRef.current.style.height = 'auto' + const handleSend = async (): Promise => { + if ( + (input || attachments.length > 0) + && !disabled + && !hasUnsupportedModalities + ) { + const submittedInput = inputRef.current + const submittedAttachments = attachmentsRef.current + const submittedRevision = draftRevisionRef.current + const outcome = await onSend(submittedInput, convertedValue ?? undefined, submittedAttachments) + if ( + outcome.clearDraft + && draftRevisionRef.current === submittedRevision + ) { + inputRef.current = '' + attachmentsRef.current = [] + draftRevisionRef.current += 1 + setInput('') + setAttachments([]) + onClearAllConversions() + if (textareaRef.current) { + textareaRef.current.style.height = 'auto' + } } } } @@ -528,7 +585,7 @@ const ChatInputArea = forwardRef(functi const handleKeyDown = (e: KeyboardEvent) => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault() - handleSend() + void handleSend() } } @@ -583,6 +640,8 @@ const ChatInputArea = forwardRef(functi }, [attachments, onAttachmentsChange]) const handleInput = (e: React.ChangeEvent) => { + inputRef.current = e.target.value + draftRevisionRef.current += 1 setInput(e.target.value) } @@ -761,7 +820,7 @@ const ChatInputArea = forwardRef(functi className={styles.sendButton} appearance="primary" icon={} - onClick={handleSend} + onClick={() => { void handleSend() }} disabled={disabled || (!input && attachments.length === 0) || hasUnsupportedModalities} aria-label="Send message" data-testid="send-message-btn" diff --git a/frontend/src/components/Chat/ChatWindow.test.tsx b/frontend/src/components/Chat/ChatWindow.test.tsx index e8db3fcc27..1d0a6deaea 100644 --- a/frontend/src/components/Chat/ChatWindow.test.tsx +++ b/frontend/src/components/Chat/ChatWindow.test.tsx @@ -1,9 +1,16 @@ -import { render, screen, waitFor } from "@testing-library/react"; +import { act, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { FluentProvider, webLightTheme } from "@fluentui/react-components"; import ChatWindow from "./ChatWindow"; import { makeTarget } from "@/test-utils/targetFixtures"; -import { Message, MessageAttachment, TargetCapabilities, TargetInfo, TargetInstance } from "../../types"; +import { + BackendMessage, + Message, + MessageAttachment, + TargetCapabilities, + TargetInfo, + TargetInstance, +} from "../../types"; import { attacksApi, convertersApi } from "../../services/api"; import * as messageMapper from "../../utils/messageMapper"; @@ -48,12 +55,17 @@ jest.mock("../../services/api", () => ({ jest.mock("../../utils/messageMapper", () => ({ buildMessagePieces: jest.fn(), + backendMessageToFrontend: jest.fn(), + backendMessageToOriginalDraft: jest.fn(), backendMessagesToFrontend: jest.fn(), })); const mockedAttacksApi = attacksApi as jest.Mocked; const mockedConvertersApi = convertersApi as jest.Mocked; const mockedMapper = messageMapper as jest.Mocked; +const actualMessageMapper = jest.requireActual( + "../../utils/messageMapper" +); const MARKDOWN_PREFERENCE_STORAGE_KEY = "pyrit.chatMarkdownMode"; const TestWrapper: React.FC<{ children: React.ReactNode }> = ({ @@ -222,12 +234,36 @@ function makeMultiModalResponse() { }; } -function makeErrorResponse(errorType: string, description: string) { +function makeErrorResponse( + errorType: string, + description: string, + failedRequestTurnNumber = 0, + hasConverters = false +) { return { messages: { messages: [ { - turn_number: 1, + turn_number: failedRequestTurnNumber, + role: "user", + message_pieces: [ + { + id: "p-failed-request", + original_value_data_type: "text", + converted_value_data_type: "text", + original_value: "failed request", + converted_value: "failed request", + converter_identifiers: hasConverters + ? [{ type: "MockConverter" }] + : [], + scores: [], + response_error: "none", + }, + ], + created_at: "2026-01-01T00:00:00Z", + }, + { + turn_number: failedRequestTurnNumber + 1, role: "assistant", message_pieces: [ { @@ -276,6 +312,11 @@ describe("ChatWindow Integration", () => { beforeEach(() => { jest.clearAllMocks(); + mockedMapper.backendMessageToFrontend.mockReset(); + mockedMapper.backendMessageToOriginalDraft.mockReset(); + mockedMapper.backendMessageToOriginalDraft.mockImplementation( + actualMessageMapper.backendMessageToOriginalDraft + ); window.localStorage.clear(); mockMatchMedia(false); // Default: panel API returns empty conversations @@ -664,6 +705,7 @@ describe("ChatWindow Integration", () => { // Messages should appear in the DOM await waitFor(() => { expect(screen.getByText("Hello back!")).toBeInTheDocument(); + expect(input).toHaveValue(""); }); }); @@ -1046,6 +1088,7 @@ describe("ChatWindow Integration", () => { await waitFor(() => { expect(screen.getByText(/Request failed with status code 404/)).toBeInTheDocument(); + expect(input).toHaveValue("test"); }); }); @@ -1495,6 +1538,378 @@ describe("ChatWindow Integration", () => { // Backend error in response piece (blocked, processing, etc.) // ----------------------------------------------------------------------- + it("should preserve the draft and expose recovery for an HTTP 200 processing error", async () => { + const user = userEvent.setup(); + const onSelectConversation = jest.fn(); + + mockedMapper.buildMessagePieces.mockResolvedValue([ + { data_type: "text", original_value: "retry this prompt" }, + ]); + mockedAttacksApi.addMessage.mockResolvedValue( + makeErrorResponse("processing", "The target could not process this message.", 2) as never + ); + mockedMapper.backendMessagesToFrontend.mockReturnValue([ + { + role: "user", + content: "retry this prompt", + timestamp: "2026-01-01T00:00:00Z", + }, + { + role: "assistant", + content: "", + timestamp: "2026-01-01T00:00:01Z", + error: { + type: "processing", + description: "The target could not process this message.", + }, + }, + ]); + mockedAttacksApi.createConversation.mockResolvedValue({ + conversation_id: "conv-processing-recovery", + } as never); + + const rendered = render( + + + + ); + + const input = screen.getByRole("textbox"); + await user.type(input, "retry this prompt"); + await user.click(screen.getByRole("button", { name: /send/i })); + + const recoveryButton = await screen.findByRole( + "button", + { name: /edit in clean conversation/i } + ); + expect(input).toHaveValue("retry this prompt"); + expect(input).toBeDisabled(); + expect(screen.queryByTestId("message-actions-1")).not.toBeInTheDocument(); + + await user.click(recoveryButton); + await waitFor(() => { + expect(mockedAttacksApi.createConversation).toHaveBeenCalledWith( + "ar-conv-processing", + { + source_conversation_id: "conv-processing", + cutoff_index: 1, + } + ); + expect(onSelectConversation).toHaveBeenCalledWith("conv-processing-recovery"); + }); + + rendered.rerender( + + + + ); + + await waitFor(() => { + expect(input).toHaveFocus(); + expect(input).toHaveValue("retry this prompt"); + }); + }); + + it("should reconstruct recovery when loading a persisted processing error", async () => { + const user = userEvent.setup(); + const onSelectConversation = jest.fn(); + const persistedMessages: BackendMessage[] = [ + { + turn_number: 2, + role: "user", + message_pieces: [ + { + id: "p-text-to-pdf", + original_value_data_type: "text", + converted_value_data_type: "binary_path", + original_value: "original persisted prompt", + converted_value: "/converted/report.pdf", + converted_value_mime_type: "application/pdf", + converted_filename: "converted.pdf", + converter_identifiers: [{ type: "PDFConverter" }], + scores: [], + response_error: "none", + }, + { + id: "p-image-to-text", + original_value_data_type: "image_path", + converted_value_data_type: "text", + original_value: "/original/evidence.png", + original_value_url: "/api/media?path=%2Foriginal%2Fevidence.png", + original_value_mime_type: "image/png", + original_filename: "evidence.png", + converted_value: "converted image description", + converter_identifiers: [{ type: "ImageToTextConverter" }], + scores: [], + response_error: "none", + }, + ], + created_at: "2026-01-01T00:00:00Z", + }, + { + turn_number: 3, + role: "assistant", + message_pieces: [ + { + id: "p-processing-error", + original_value_data_type: "text", + converted_value_data_type: "text", + original_value: "", + converted_value: "", + scores: [], + response_error: "processing", + response_error_description: "The target could not process this message.", + }, + ], + created_at: "2026-01-01T00:00:01Z", + }, + ]; + + mockedAttacksApi.getMessages.mockResolvedValue({ + messages: persistedMessages, + } as never); + mockedMapper.backendMessagesToFrontend.mockReturnValue([ + { + role: "user", + content: "", + attachments: [ + { + type: "file", + name: "converted.pdf", + url: "/converted/report.pdf", + mimeType: "application/pdf", + }, + ], + timestamp: "2026-01-01T00:00:00Z", + }, + { + role: "assistant", + content: "", + timestamp: "2026-01-01T00:00:01Z", + error: { + type: "processing", + description: "The target could not process this message.", + }, + }, + ]); + mockedAttacksApi.createConversation.mockResolvedValue({ + conversation_id: "conv-persisted-recovery", + } as never); + + const rendered = render( + + + + ); + + const recoveryButton = await screen.findByRole( + "button", + { name: /edit in clean conversation/i } + ); + expect(screen.getByText(/converter choices could not be restored/i)).toBeInTheDocument(); + expect(screen.getByRole("textbox")).toBeDisabled(); + + await user.click(recoveryButton); + await waitFor(() => { + expect(mockedAttacksApi.createConversation).toHaveBeenCalledWith( + "ar-persisted-processing", + { + source_conversation_id: "conv-persisted-processing", + cutoff_index: 1, + } + ); + expect(onSelectConversation).toHaveBeenCalledWith("conv-persisted-recovery"); + }); + + mockedAttacksApi.getMessages.mockResolvedValue({ messages: [] } as never); + mockedMapper.backendMessagesToFrontend.mockReturnValue([]); + rendered.rerender( + + + + ); + + const restoredInput = await screen.findByRole("textbox"); + expect(restoredInput).toHaveValue("original persisted prompt"); + expect(screen.getAllByText("evidence.png", { exact: false })).toHaveLength(1); + expect(screen.queryByText(/converted\.pdf/i)).not.toBeInTheDocument(); + }); + + it("should clear an unchanged submitted draft after switching conversations", async () => { + const user = userEvent.setup(); + const response = makeTextResponse("Reply from conversation A"); + let resolveMessage: ((value: typeof response) => void) | undefined; + + mockedAttacksApi.getMessages.mockResolvedValue({ messages: [] } as never); + mockedMapper.backendMessagesToFrontend.mockReturnValue([]); + mockedMapper.buildMessagePieces.mockResolvedValue([ + { data_type: "text", original_value: "conversation A draft" }, + ]); + mockedAttacksApi.addMessage.mockImplementation( + () => new Promise((resolve) => { + resolveMessage = resolve; + }) as never + ); + + const rendered = render( + + + + ); + + await waitFor(() => { + expect(mockedAttacksApi.getMessages).toHaveBeenCalledWith( + "ar-conversation-switch", + "conv-a" + ); + }); + + const input = screen.getByRole("textbox"); + await user.type(input, "conversation A draft"); + await user.click(screen.getByRole("button", { name: /send/i })); + await waitFor(() => expect(mockedAttacksApi.addMessage).toHaveBeenCalledTimes(1)); + + rendered.rerender( + + + + ); + await waitFor(() => { + expect(mockedAttacksApi.getMessages).toHaveBeenCalledWith( + "ar-conversation-switch", + "conv-b" + ); + }); + + await act(async () => { + resolveMessage?.(response); + await Promise.resolve(); + }); + + expect(input).toHaveValue(""); + }); + + it("should not overwrite another conversation when recovery completes after navigation", async () => { + const user = userEvent.setup(); + const onSelectConversation = jest.fn(); + let resolveConversation: ((value: { conversation_id: string }) => void) | undefined; + + mockedMapper.buildMessagePieces.mockResolvedValue([ + { data_type: "text", original_value: "failed draft" }, + ]); + mockedAttacksApi.addMessage.mockResolvedValue( + makeErrorResponse( + "processing", + "The target could not process this message.", + 2 + ) as never + ); + mockedMapper.backendMessagesToFrontend.mockReturnValue([ + { + role: "user", + content: "failed draft", + timestamp: "2026-01-01T00:00:00Z", + }, + { + role: "assistant", + content: "", + timestamp: "2026-01-01T00:00:01Z", + error: { + type: "processing", + description: "The target could not process this message.", + }, + }, + ]); + mockedAttacksApi.createConversation.mockImplementation( + () => new Promise<{ conversation_id: string }>((resolve) => { + resolveConversation = resolve; + }) as never + ); + + const rendered = render( + + + + ); + + const input = screen.getByRole("textbox"); + await user.type(input, "failed draft"); + await user.click(screen.getByRole("button", { name: /send/i })); + const recoveryButton = await screen.findByRole( + "button", + { name: /edit in clean conversation/i } + ); + + await user.click(recoveryButton); + await waitFor(() => { + expect(recoveryButton).toBeDisabled(); + expect(mockedAttacksApi.createConversation).toHaveBeenCalledTimes(1); + }); + await user.click(recoveryButton); + expect(mockedAttacksApi.createConversation).toHaveBeenCalledTimes(1); + + rendered.rerender( + + + + ); + await waitFor(() => expect(input).toBeEnabled()); + await user.clear(input); + await user.type(input, "newer conversation draft"); + + await act(async () => { + resolveConversation?.({ conversation_id: "conv-unused-recovery" }); + await Promise.resolve(); + }); + + expect(onSelectConversation).not.toHaveBeenCalled(); + expect(input).toHaveValue("newer conversation draft"); + }); + it("should handle blocked response from target", async () => { const user = userEvent.setup(); @@ -1532,7 +1947,94 @@ describe("ChatWindow Integration", () => { await waitFor(() => { expect(screen.getByText(/Content was filtered by safety system/)).toBeInTheDocument(); + expect(input).toHaveValue(""); + }); + expect(screen.queryByTestId(/^recover-processing-error-btn-/)).not.toBeInTheDocument(); + }); + + it("should restore a single-turn draft in a new conversation after a processing error", async () => { + const user = userEvent.setup(); + const onSelectConversation = jest.fn(); + const singleTurnTarget: TargetInstance = makeTarget({ + target_registry_name: "single-turn-target", + target_type: "OpenAIImageTarget", + capabilities: buildCapabilities({ supports_multi_turn: false }), + }); + + mockedMapper.buildMessagePieces.mockResolvedValue([ + { data_type: "text", original_value: "generate this image" }, + ]); + mockedAttacksApi.addMessage.mockResolvedValue( + makeErrorResponse("processing", "The target could not process this message.") as never + ); + mockedMapper.backendMessagesToFrontend.mockReturnValue([ + { + role: "user", + content: "generate this image", + timestamp: "2026-01-01T00:00:00Z", + }, + { + role: "assistant", + content: "", + timestamp: "2026-01-01T00:00:01Z", + error: { + type: "processing", + description: "The target could not process this message.", + }, + }, + ]); + mockedAttacksApi.createConversation.mockResolvedValue({ + conversation_id: "conv-single-recovery", + } as never); + + const rendered = render( + + + + ); + + const input = screen.getByRole("textbox"); + await user.type(input, "generate this image"); + await user.click(screen.getByRole("button", { name: /send/i })); + + const recoveryButton = await screen.findByRole( + "button", + { name: /edit in new conversation/i } + ); + expect(screen.getByTestId("single-turn-banner")).toBeInTheDocument(); + + await user.click(recoveryButton); + await waitFor(() => { + expect(mockedAttacksApi.createConversation).toHaveBeenCalledWith( + "ar-single-processing", + {} + ); + expect(onSelectConversation).toHaveBeenCalledWith("conv-single-recovery"); }); + + mockedAttacksApi.getMessages.mockResolvedValue({ messages: [] } as never); + mockedMapper.backendMessagesToFrontend.mockReturnValue([]); + rendered.rerender( + + + + ); + + const restoredInput = await screen.findByRole("textbox"); + expect(restoredInput).toHaveValue("generate this image"); }); // ----------------------------------------------------------------------- diff --git a/frontend/src/components/Chat/ChatWindow.tsx b/frontend/src/components/Chat/ChatWindow.tsx index 65369c07a7..54e1db6847 100644 --- a/frontend/src/components/Chat/ChatWindow.tsx +++ b/frontend/src/components/Chat/ChatWindow.tsx @@ -1,4 +1,4 @@ -import { useState, useRef, useEffect, useCallback, useMemo } from 'react' +import { useState, useRef, useEffect, useLayoutEffect, useCallback, useMemo } from 'react' import type { ChangeEvent } from 'react' import { Button, @@ -29,11 +29,18 @@ import LabelsBar from '../Labels/LabelsBar' import type { ChatInputAreaHandle } from './ChatInputArea' import { attacksApi } from '../../services/api' import { toApiError } from '../../services/errors' -import { buildMessagePieces, backendMessagesToFrontend } from '../../utils/messageMapper' +import { + buildMessagePieces, + backendMessageToOriginalDraft, + backendMessagesToFrontend, +} from '../../utils/messageMapper' import { exportConversation } from '../../utils/conversationExport' import type { ExportFormat } from '../../utils/conversationExport' import type { AttackTargetResolutionStatus, + BackendMessage, + ChatSendOutcome, + CreateConversationRequest, Message, MessageAttachment, TargetInstance, @@ -45,6 +52,121 @@ import { useChatWindowStyles } from './ChatWindow.styles' const NARROW_SCREEN_QUERY = '(max-width: 600px)' const MARKDOWN_PREFERENCE_STORAGE_KEY = 'pyrit.chatMarkdownMode' +const RETRYABLE_TARGET_RESPONSE_ERROR = 'processing' +const CLEAN_CONVERSATION_MESSAGE = + 'Continue in a clean conversation so the stored error is not sent back to the target.' + +interface RecoverableSendDraft { + conversationId: string + failedRequestTurnNumber: number + originalValue: string + attachments: MessageAttachment[] + conversions: Record + source: 'live' | 'persisted' + missingConverterSelections: boolean +} + +interface TargetResponseFailure { + type: string + errorTurnNumber: number + failedRequestTurnNumber?: number +} + +function getRecoveryDescription(draft: RecoverableSendDraft): string { + if (draft.source === 'live') { + return `${CLEAN_CONVERSATION_MESSAGE} Your prompt, attachments, and converter choices are preserved for editing.` + } + + const restored = 'Your prompt and attachments were restored from conversation history.' + + if (draft.missingConverterSelections) { + return `${CLEAN_CONVERSATION_MESSAGE} ${restored} Converter choices could not be restored, so review them before sending.` + } + + return `${CLEAN_CONVERSATION_MESSAGE} ${restored} Review them before sending.` +} + +function findPrecedingUserMessage( + messages: BackendMessage[], + beforeIndex: number, +): BackendMessage | undefined { + for (let index = beforeIndex - 1; index >= 0; index -= 1) { + if (messages[index].role === 'user') { + return messages[index] + } + } + return undefined +} + +function getLatestTargetResponseFailure(messages: BackendMessage[]): TargetResponseFailure | undefined { + const latestMessageIndex = messages.length - 1 + const latestMessage = messages[messages.length - 1] + if ( + !latestMessage + || (latestMessage.role !== 'assistant' && latestMessage.role !== 'simulated_assistant') + ) { + return undefined + } + + const errorType = latestMessage.message_pieces.find( + (piece) => piece.response_error && piece.response_error !== 'none', + )?.response_error + if (!errorType) { + return undefined + } + + return { + type: errorType, + errorTurnNumber: latestMessage.turn_number, + failedRequestTurnNumber: findPrecedingUserMessage(messages, latestMessageIndex)?.turn_number, + } +} + +function getPersistedProcessingRecovery( + conversationId: string, + messages: BackendMessage[], +): RecoverableSendDraft | undefined { + for (let errorIndex = messages.length - 1; errorIndex >= 0; errorIndex -= 1) { + const errorMessage = messages[errorIndex] + const isProcessingError = ( + errorMessage.role === 'assistant' + || errorMessage.role === 'simulated_assistant' + ) && errorMessage.message_pieces.some( + (piece) => piece.response_error === RETRYABLE_TARGET_RESPONSE_ERROR, + ) + if (!isProcessingError) { + continue + } + + const failedRequest = findPrecedingUserMessage(messages, errorIndex) + if (!failedRequest) { + return undefined + } + + const originalDraft = backendMessageToOriginalDraft(failedRequest) + return { + conversationId, + failedRequestTurnNumber: failedRequest.turn_number, + originalValue: originalDraft.content, + attachments: (originalDraft.attachments ?? []).map((attachment) => ({ ...attachment })), + conversions: {}, + source: 'persisted', + missingConverterSelections: failedRequest.message_pieces.some( + (piece) => Boolean(piece.converter_identifiers?.length), + ), + } + } + return undefined +} + +function findLastProcessingErrorIndex(messages: Message[]): number | undefined { + for (let index = messages.length - 1; index >= 0; index -= 1) { + if (messages[index].error?.type === RETRYABLE_TARGET_RESPONSE_ERROR) { + return index + } + } + return undefined +} function readStoredMarkdownPreference(): boolean { if (typeof window === 'undefined') return false @@ -135,8 +257,15 @@ export default function ChatWindow({ const [attachmentTypes, setAttachmentTypes] = useState([]) const [attachmentData, setAttachmentData] = useState>({}) const [pieceConversions, setPieceConversions] = useState>({}) + const [recoverableSends, setRecoverableSends] = useState>({}) + const [isRecoveringProcessingError, setIsRecoveringProcessingError] = useState(false) const [panelRefreshKey, setPanelRefreshKey] = useState(0) const inputBoxRef = useRef(null) + const recoveryInFlightRef = useRef(false) + const viewedConversationId = activeConversationId ?? conversationId + const recoverableSend = viewedConversationId + ? recoverableSends[viewedConversationId] + : undefined const handleMarkdownChange = useCallback(( _event: ChangeEvent, @@ -187,6 +316,13 @@ export default function ChatWindow({ } return hasStale ? next : pieceConversions }, [pieceConversions, chatInputText, attachmentData]) + const conversionRevisionKey = useMemo( + () => JSON.stringify( + Object.entries(activePieceConversions) + .sort(([left], [right]) => left.localeCompare(right)), + ), + [activePieceConversions], + ) // Auto-open conversation sidebar when loading a historical attack with multiple // conversations. Uses the "adjust state during render" pattern to avoid @@ -210,7 +346,9 @@ export default function ChatWindow({ // Always-current ref of the conversation being viewed so async callbacks can // check whether the user navigated away while a request was in-flight. const viewedConvRef = useRef(activeConversationId ?? conversationId) - useEffect(() => { viewedConvRef.current = activeConversationId ?? conversationId }, [activeConversationId, conversationId]) + useLayoutEffect(() => { + viewedConvRef.current = activeConversationId ?? conversationId + }, [activeConversationId, conversationId]) // Synchronous ref tracking which conversations have an in-flight send. const sendingConvIdsRef = useRef>(new Set()) // Pending user messages per conversation that may not be stored server-side yet. @@ -246,6 +384,7 @@ export default function ChatWindow({ if (attackResultId !== prevAttackResultId) { setPrevAttackResultId(attackResultId) if (!attackResultId) { + setRecoverableSends({}) setMessages([]) setLoadedConversationId(null) setSystemPrompt('') @@ -271,6 +410,22 @@ export default function ChatWindow({ // Discard stale response if user navigated away while loading if (viewedConvRef.current !== convId) { return } const frontendMessages = backendMessagesToFrontend(response.messages) + const persistedRecovery = getPersistedProcessingRecovery(convId, response.messages) + setRecoverableSends((currentRecoveries) => { + const currentRecovery = currentRecoveries[convId] + if (persistedRecovery) { + if (currentRecovery?.source === 'live') { + return currentRecoveries + } + return { ...currentRecoveries, [convId]: persistedRecovery } + } + if (!currentRecovery || currentRecovery.source === 'live') { + return currentRecoveries + } + const nextRecoveries = { ...currentRecoveries } + delete nextRecoveries[convId] + return nextRecoveries + }) // If this conversation has an in-flight send, append any pending user // messages (that the server may not have stored yet) and a loading indicator. if (sendingConvIdsRef.current.has(convId)) { @@ -329,15 +484,33 @@ export default function ChatWindow({ } }, [attackResultId, activeConversationId, isNarrowScreen, onSelectConversation, loadConversation]) - const handleSend = async (originalValue: string, convertedValue: string | undefined, attachments: MessageAttachment[]) => { + const handleSend = async ( + originalValue: string, + convertedValue: string | undefined, + attachments: MessageAttachment[], + ): Promise => { if ( !activeTarget || isLoadingAttack || isMutationLocked ) { - return + return { status: 'retryable_failure', clearDraft: false } + } + + const initialSendConvId = activeConversationId ?? conversationId ?? '__pending__' + if (sendingConvIdsRef.current.has(initialSendConvId)) { + return { status: 'retryable_failure', clearDraft: false } } + setRecoverableSends((currentRecoveries) => { + if (!currentRecoveries[initialSendConvId]) { + return currentRecoveries + } + const nextRecoveries = { ...currentRecoveries } + delete nextRecoveries[initialSendConvId] + return nextRecoveries + }) + // Capture all piece conversions upfront before any async work or state clears const conversions = { ...activePieceConversions } const textConversion = conversions['text'] @@ -345,7 +518,7 @@ export default function ChatWindow({ const isTextFileConversion = Boolean(textConversion) && !isTextTextConversion // Track which conversation this send belongs to (may be updated after attack creation) - let sendConvId = activeConversationId || '__pending__' + let sendConvId = initialSendConvId // Mark synchronously so the useEffect guard sees it immediately sendingConvIdsRef.current.add(sendConvId) @@ -450,18 +623,42 @@ export default function ChatWindow({ // Send message to target const converterIds = allConverterIds.length > 0 ? allConverterIds : undefined - const response = await attacksApi.addMessage(currentAttackResultId!, { + if (!currentAttackResultId || !effectiveConvId) { + throw new Error('Message send is missing an attack or conversation ID.') + } + const response = await attacksApi.addMessage(currentAttackResultId, { role: 'user', pieces, send: true, target_registry_name: activeTarget.target_registry_name, - target_conversation_id: effectiveConvId!, + target_conversation_id: effectiveConvId, labels: labels ?? undefined, converter_ids: converterIds, }) - // Clear converter state after successful send - setPieceConversions({}) + const targetResponseFailure = getLatestTargetResponseFailure(response.messages.messages) + const status: ChatSendOutcome['status'] = targetResponseFailure?.type === RETRYABLE_TARGET_RESPONSE_ERROR + ? 'retryable_failure' + : targetResponseFailure + ? 'non_retryable_failure' + : 'sent' + const backendMessages = backendMessagesToFrontend(response.messages.messages) + + if (targetResponseFailure?.type === RETRYABLE_TARGET_RESPONSE_ERROR) { + setRecoverableSends((currentRecoveries) => ({ + ...currentRecoveries, + [effectiveConvId]: { + conversationId: effectiveConvId, + failedRequestTurnNumber: targetResponseFailure.failedRequestTurnNumber + ?? targetResponseFailure.errorTurnNumber - 1, + originalValue, + attachments: attachments.map((attachment) => ({ ...attachment })), + conversions, + source: 'live', + missingConverterSelections: false, + }, + })) + } // Only update displayed messages if the user is still viewing this conversation. // If they switched away the response is persisted server-side and will appear @@ -470,9 +667,12 @@ export default function ChatWindow({ // Replace the entire message list with authoritative server data. // This correctly handles the case where the user switched away and // back during the request — the full conversation is restored. - const backendMessages = backendMessagesToFrontend(response.messages.messages) setMessages(backendMessages) - setLoadedConversationId(effectiveConvId!) + setLoadedConversationId(effectiveConvId) + } + return { + status, + clearDraft: status !== 'retryable_failure' || viewedConvRef.current !== effectiveConvId, } } catch (err) { const viewedConversationId = viewedConvRef.current @@ -516,10 +716,10 @@ export default function ChatWindow({ return [...prev, errorMessage] }) - // Preserve the failed message text in the input box for easy re-send - if (originalValue && inputBoxRef.current) { - inputBoxRef.current.setText(originalValue) - } + } + return { + status: 'retryable_failure', + clearDraft: viewedConversationId != null && viewedConversationId !== sendConvId, } } finally { sendingConvIdsRef.current.delete(sendConvId) @@ -533,23 +733,116 @@ export default function ChatWindow({ } } - const handleNewConversation = useCallback(async () => { - if (!attackResultId || isMutationLocked) { return } + const appendConversationCreationError = useCallback((error: unknown): void => { + const apiError = toApiError(error) + setMessages((previousMessages) => [ + ...previousMessages, + { + role: 'assistant', + content: '', + timestamp: new Date().toISOString(), + error: { + type: 'unknown', + description: `Could not create a new conversation. ${apiError.detail}`, + }, + }, + ]) + }, []) + + const createAndSelectConversation = useCallback(async ( + request: CreateConversationRequest, + ): Promise => { + if (!attackResultId || isMutationLocked) { return false } try { - const response = await attacksApi.createConversation(attackResultId, {}) + const response = await attacksApi.createConversation(attackResultId, request) onSelectConversation(response.conversation_id) setIsPanelOpen(!isNarrowScreen) - } catch { - // Silently fail + return true + } catch (err) { + appendConversationCreationError(err) + return false } }, [ + appendConversationCreationError, attackResultId, isNarrowScreen, isMutationLocked, onSelectConversation, ]) + const handleNewConversation = useCallback( + (): Promise => createAndSelectConversation({}), + [createAndSelectConversation], + ) + + const restoreRecoverableDraft = useCallback((): void => { + if (!recoverableSend) { return } + setPieceConversions(recoverableSend.conversions) + inputBoxRef.current?.restoreDraft( + recoverableSend.originalValue, + recoverableSend.attachments, + ) + inputBoxRef.current?.focus() + }, [recoverableSend]) + + const handleRecoverProcessingError = useCallback(async (): Promise => { + if ( + !attackResultId + || !recoverableSend + || isMutationLocked + || recoveryInFlightRef.current + ) { + return + } + + const supportsMultiTurn = Boolean( + activeTarget && activeTarget.capabilities?.supports_multi_turn !== false, + ) + const cutoffIndex = recoverableSend.failedRequestTurnNumber - 1 + const recoveryRequest: CreateConversationRequest = supportsMultiTurn && cutoffIndex >= 0 + ? { + source_conversation_id: recoverableSend.conversationId, + cutoff_index: cutoffIndex, + } + : {} + const sourceConversationId = recoverableSend.conversationId + const draftRevision = inputBoxRef.current?.getDraftRevision() + + recoveryInFlightRef.current = true + setIsRecoveringProcessingError(true) + try { + const response = await attacksApi.createConversation(attackResultId, recoveryRequest) + setPanelRefreshKey((currentKey) => currentKey + 1) + + const isStillViewingSource = viewedConvRef.current === sourceConversationId + const isDraftUnchanged = inputBoxRef.current?.getDraftRevision() === draftRevision + if (!isStillViewingSource || !isDraftUnchanged) { + return + } + + onSelectConversation(response.conversation_id) + setIsPanelOpen(!isNarrowScreen) + restoreRecoverableDraft() + } catch (err) { + if (viewedConvRef.current === sourceConversationId) { + appendConversationCreationError(err) + } + } finally { + recoveryInFlightRef.current = false + setIsRecoveringProcessingError(false) + } + }, [ + activeTarget, + appendConversationCreationError, + attackResultId, + isMutationLocked, + isNarrowScreen, + onSelectConversation, + recoverableSend, + restoreRecoverableDraft, + ]) + // ------------------------------------------------------------------- // Message action handlers (4 buttons on each assistant message) // ------------------------------------------------------------------- @@ -676,6 +969,12 @@ export default function ChatWindow({ ]) const singleTurnLimitReached = activeTarget?.capabilities?.supports_multi_turn === false && messages.some(m => m.role === 'user') + const recoverableProcessingErrorIndex = recoverableSend?.conversationId === viewedConversationId + ? findLastProcessingErrorIndex(messages) + : undefined + const processingRecoveryDescription = recoverableSend + ? getRecoveryDescription(recoverableSend) + : undefined // "Continue with your target" — clone the current conversation into a new attack const handleUseAsTemplate = useCallback(async () => { @@ -827,10 +1126,23 @@ export default function ChatWindow({ isCrossTarget={isCrossTargetLocked || isTargetResolutionLocked} noTargetSelected={!activeTarget} globalMarkdown={globalMarkdown} + processingErrorRecovery={recoverableProcessingErrorIndex === undefined + || processingRecoveryDescription === undefined + ? undefined + : { + messageIndex: recoverableProcessingErrorIndex, + actionLabel: activeTarget?.capabilities?.supports_multi_turn === false + ? 'Edit in new conversation' + : 'Edit in clean conversation', + description: processingRecoveryDescription, + disabled: isRecoveringProcessingError || isMutationLocked, + onRecover: handleRecoverProcessingError, + }} /> setPieceConversions((prev) => { const next = { ...prev }; delete next['text']; return next })} + onClearAllConversions={() => setPieceConversions((current) => ( + current === pieceConversions ? {} : current + ))} onConvertedValueChange={(val) => setPieceConversions((prev) => { const existing = prev['text'] if (!existing) return prev diff --git a/frontend/src/components/Chat/MessageList.styles.ts b/frontend/src/components/Chat/MessageList.styles.ts index 29603f7514..fe1c9084d4 100644 --- a/frontend/src/components/Chat/MessageList.styles.ts +++ b/frontend/src/components/Chat/MessageList.styles.ts @@ -187,6 +187,17 @@ export const useMessageListStyles = makeStyles({ errorContainer: { marginTop: tokens.spacingVerticalS, }, + errorRecovery: { + display: 'flex', + flexDirection: 'column', + alignItems: 'flex-start', + gap: tokens.spacingVerticalS, + marginTop: tokens.spacingVerticalS, + paddingBottom: tokens.spacingVerticalS, + }, + errorRecoveryButton: { + ...mobileTouchTarget, + }, mediaActions: { display: 'flex', gap: tokens.spacingHorizontalXXS, diff --git a/frontend/src/components/Chat/MessageList.test.tsx b/frontend/src/components/Chat/MessageList.test.tsx index d5e0e32e57..ed9cfa0b06 100644 --- a/frontend/src/components/Chat/MessageList.test.tsx +++ b/frontend/src/components/Chat/MessageList.test.tsx @@ -389,6 +389,7 @@ describe("MessageList", () => { }); it("should render error messages", () => { + const onRecover = jest.fn(); const errorMessages: Message[] = [ { role: "assistant", @@ -403,13 +404,65 @@ describe("MessageList", () => { render( - + ); expect( screen.getByText(/Content was filtered by safety system/) ).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /edit in clean conversation/i })).not.toBeInTheDocument(); + expect(onRecover).not.toHaveBeenCalled(); + }); + + it("should render a direct recovery action only for the current processing error", async () => { + const user = userEvent.setup(); + const onRecover = jest.fn(); + const messages: Message[] = [ + { + role: "assistant", + content: "", + timestamp: new Date().toISOString(), + error: { + type: "processing", + description: "The target could not process this message.", + }, + }, + ]; + + render( + + + + ); + + expect( + screen.getByText(/prompt, attachments, and converter choices are preserved/i) + ).toBeInTheDocument(); + expect(screen.getByText(/stored error is not sent back to the target/i)).toBeInTheDocument(); + expect(screen.queryByTestId("message-actions-0")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: /edit in clean conversation/i })); + expect(onRecover).toHaveBeenCalledTimes(1); }); it("should render multiple messages in order", () => { diff --git a/frontend/src/components/Chat/MessageList.tsx b/frontend/src/components/Chat/MessageList.tsx index b072b866f3..c976488ad1 100644 --- a/frontend/src/components/Chat/MessageList.tsx +++ b/frontend/src/components/Chat/MessageList.tsx @@ -10,11 +10,19 @@ import { Spinner, mergeClasses, } from '@fluentui/react-components' -import { ArrowDownloadRegular, ArrowReplyRegular, ArrowForwardRegular, ChatAddRegular, BranchForkRegular, OpenRegular } from '@fluentui/react-icons' +import { ArrowDownloadRegular, ArrowReplyRegular, ArrowForwardRegular, ChatAddRegular, BranchForkRegular, EditRegular, OpenRegular } from '@fluentui/react-icons' import { Message, MessageAttachment } from '../../types' import MarkdownContent from './MarkdownContent' import { useMessageListStyles } from './MessageList.styles' +interface ProcessingErrorRecovery { + messageIndex: number + actionLabel: string + description: string + disabled?: boolean + onRecover: () => void | Promise +} + interface MessageListProps { messages: Message[] /** Copy this message to the input box of the current conversation */ @@ -37,6 +45,8 @@ interface MessageListProps { noTargetSelected?: boolean /** Conversation-wide default: render message text as Markdown. */ globalMarkdown?: boolean + /** Recovery action for the processing error caused by the most recent send. */ + processingErrorRecovery?: ProcessingErrorRecovery } /** Image that shows a spinner while loading. */ @@ -108,7 +118,7 @@ function tryFormatJson(text: string): string | null { } } -export default function MessageList({ messages, onCopyToInput, onCopyToNewConversation, onBranchConversation, onBranchAttack, isLoading, isSingleTurn, isOperatorLocked, isCrossTarget, noTargetSelected, globalMarkdown = false }: MessageListProps) { +export default function MessageList({ messages, onCopyToInput, onCopyToNewConversation, onBranchConversation, onBranchAttack, isLoading, isSingleTurn, isOperatorLocked, isCrossTarget, noTargetSelected, globalMarkdown = false, processingErrorRecovery }: MessageListProps) { const styles = useMessageListStyles() const messagesEndRef = useRef(null) @@ -162,6 +172,8 @@ export default function MessageList({ messages, onCopyToInput, onCopyToNewConver const isSimulated = message.role === 'simulated_assistant' const timestamp = new Date(message.timestamp).toLocaleTimeString() const avatarName = isUser ? 'User' : isSimulated ? 'Simulated' : 'Assistant' + const canRecoverProcessingError = message.error?.type === 'processing' + && processingErrorRecovery?.messageIndex === index return (
: {message.error.description} )} + {canRecoverProcessingError && processingErrorRecovery && ( +
+ + {processingErrorRecovery.description} + + +
+ )}
@@ -318,7 +348,7 @@ export default function MessageList({ messages, onCopyToInput, onCopyToNewConver )} {/* Unified action buttons – shown on all non-user, non-loading messages */} - {!isUser && !message.isLoading && ( + {!isUser && !message.isLoading && !message.error && (
{/* 1. Copy to input box in this conversation */} {onCopyToInput && (() => { diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 0852de5827..f8d39afefd 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -14,6 +14,8 @@ export interface MessageAttachment { */ size?: number file?: File + /** Raw backend value used when reconstructing a persisted attachment for resubmission. */ + sourceValue?: string /** Backend piece ID — preserved so remix/copy can trace back to the original piece */ pieceId?: string /** Backend prompt_metadata — preserved so video_id etc. carry over on remix/copy */ @@ -45,6 +47,11 @@ export interface MessageError { description?: string } +export interface ChatSendOutcome { + status: 'sent' | 'retryable_failure' | 'non_retryable_failure' + clearDraft: boolean +} + // ============================================================================ // Backend DTO Types (mirror pyrit/backend/models) // ============================================================================ @@ -307,6 +314,7 @@ export interface BackendMessagePiece { original_filename?: string | null converted_filename?: string | null prompt_metadata?: Record | null + converter_identifiers?: Array> scores: BackendScore[] response_error: string // 'none' | 'blocked' | 'processing' | 'empty' | 'unknown' response_error_description?: string | null diff --git a/frontend/src/utils/conversationExport.test.ts b/frontend/src/utils/conversationExport.test.ts index aae2eaf3f1..e765dab607 100644 --- a/frontend/src/utils/conversationExport.test.ts +++ b/frontend/src/utils/conversationExport.test.ts @@ -270,7 +270,7 @@ describe("conversationExport", () => { expect(JSON.parse(json).messages).toHaveLength(1); }); - it("omits the in-memory File handle but keeps the other attachment fields", () => { + it("omits internal attachment values but keeps export-safe fields", () => { const file = new File(["x"], "local.png", { type: "image/png" }); const json = conversationToJson( [ @@ -282,6 +282,7 @@ describe("conversationExport", () => { url: "blob:local", mimeType: "image/png", pieceId: "piece-9", + sourceValue: "C:\\private\\raw-image.png", file, }, ], @@ -291,6 +292,7 @@ describe("conversationExport", () => { ); const attachment = JSON.parse(json).messages[0].attachments[0]; expect(attachment.file).toBeUndefined(); + expect(attachment.sourceValue).toBeUndefined(); expect(attachment.name).toBe("local.png"); expect(attachment.pieceId).toBe("piece-9"); }); @@ -320,13 +322,20 @@ describe("conversationExport", () => { expect(attachment.metadata.video_id).toBe("v1"); }); - it("strips the File handle from original attachments too", () => { + it("strips all internal values from original attachments too", () => { const file = new File(["x"], "orig.png", { type: "image/png" }); const json = conversationToJson( [ message({ originalAttachments: [ - { type: "image", name: "orig.png", url: "blob:orig", mimeType: "image/png", file }, + { + type: "image", + name: "orig.png", + url: "blob:orig", + mimeType: "image/png", + sourceValue: "C:\\private\\orig.png", + file, + }, ], }), ], @@ -334,6 +343,7 @@ describe("conversationExport", () => { ); const attachment = JSON.parse(json).messages[0].originalAttachments[0]; expect(attachment.file).toBeUndefined(); + expect(attachment.sourceValue).toBeUndefined(); expect(attachment.name).toBe("orig.png"); }); diff --git a/frontend/src/utils/conversationExport.ts b/frontend/src/utils/conversationExport.ts index 9310b155cf..25d6122afb 100644 --- a/frontend/src/utils/conversationExport.ts +++ b/frontend/src/utils/conversationExport.ts @@ -87,8 +87,8 @@ export function conversationToMarkdown( * Serialize the in-state conversation to pretty-printed JSON, exporting exactly * what the GUI holds (WYSIWYG). The envelope records the conversation id, the * export timestamp, and the messages. Loading placeholders are dropped and the - * non-serializable `File` handle is removed from each attachment; every other - * field (including attachment metadata) is preserved as-is. + * internal `File` handles and raw resubmission values are removed from each + * attachment; export-safe fields (including attachment metadata) are preserved. */ export function conversationToJson( messages: Message[], @@ -154,17 +154,18 @@ function messageForExport(message: Message): Message { } const next: Message = { ...message } if (message.attachments) { - next.attachments = message.attachments.map(attachmentWithoutFile) + next.attachments = message.attachments.map(attachmentForExport) } if (message.originalAttachments) { - next.originalAttachments = message.originalAttachments.map(attachmentWithoutFile) + next.originalAttachments = message.originalAttachments.map(attachmentForExport) } return next } -function attachmentWithoutFile(attachment: MessageAttachment): MessageAttachment { +function attachmentForExport(attachment: MessageAttachment): MessageAttachment { const next = { ...attachment } delete next.file + delete next.sourceValue return next } diff --git a/frontend/src/utils/messageMapper.test.ts b/frontend/src/utils/messageMapper.test.ts index 93d480ecba..4e12f9b3c1 100644 --- a/frontend/src/utils/messageMapper.test.ts +++ b/frontend/src/utils/messageMapper.test.ts @@ -4,6 +4,7 @@ import { dataTypeToAttachmentType, buildDataUri, backendMessageToFrontend, + backendMessageToOriginalDraft, backendMessagesToFrontend, attachmentToMessagePieceRequest, buildMessagePieces, @@ -94,6 +95,57 @@ describe("messageMapper", () => { }); }); + describe("backendMessageToOriginalDraft", () => { + it("restores only original text and media without converted output", () => { + const msg: BackendMessage = { + turn_number: 2, + role: "user", + message_pieces: [ + { + id: "text-to-pdf", + original_value_data_type: "text", + converted_value_data_type: "binary_path", + original_value: "original persisted prompt", + converted_value: "/converted/report.pdf", + converted_value_mime_type: "application/pdf", + converted_filename: "converted.pdf", + scores: [], + response_error: "none", + }, + { + id: "image-to-text", + original_value_data_type: "image_path", + converted_value_data_type: "text", + original_value: "/original/evidence.png", + original_value_url: "/api/media?path=%2Foriginal%2Fevidence.png", + original_value_mime_type: "image/png", + original_filename: "evidence.png", + converted_value: "converted image description", + scores: [], + response_error: "none", + }, + ], + created_at: "2026-01-01T00:00:00Z", + }; + + const result = backendMessageToOriginalDraft(msg); + + expect(result.content).toBe("original persisted prompt"); + expect(result.attachments).toEqual([ + expect.objectContaining({ + name: "evidence.png", + url: "/api/media?path=%2Foriginal%2Fevidence.png", + sourceValue: "/original/evidence.png", + }), + ]); + expect(result.attachments).not.toEqual( + expect.arrayContaining([ + expect.objectContaining({ name: "converted.pdf" }), + ]) + ); + }); + }); + describe("backendMessageToFrontend", () => { it("should convert a text message", () => { const msg: BackendMessage = { @@ -288,7 +340,7 @@ describe("messageMapper", () => { expect(result.attachments![0].url).toBe("/api/media?path=output%2Fdoc.pdf"); }); - it("should handle error response", () => { + it("should preserve user-facing content for a blocked response", () => { const msg: BackendMessage = { turn_number: 1, role: "assistant", @@ -296,8 +348,8 @@ describe("messageMapper", () => { { id: "p1", original_value_data_type: "text", - converted_value_data_type: "text", - converted_value: "", + converted_value_data_type: "error", + converted_value: "I cannot help with that request.", scores: [], response_error: "blocked", response_error_description: "Content was filtered", @@ -311,6 +363,33 @@ describe("messageMapper", () => { expect(result.error).toBeDefined(); expect(result.error!.type).toBe("blocked"); expect(result.error!.description).toBe("Content was filtered"); + expect(result.content).toBe("I cannot help with that request."); + }); + + it("should hide stored processing diagnostics and provide a safe description", () => { + const msg: BackendMessage = { + turn_number: 1, + role: "assistant", + message_pieces: [ + { + id: "p-processing", + original_value_data_type: "text", + converted_value_data_type: "error", + converted_value: "RuntimeError: target failed\nTraceback (most recent call last): ...", + scores: [], + response_error: "processing", + }, + ], + created_at: "2026-02-15T00:00:00Z", + }; + + const result = backendMessageToFrontend(msg); + + expect(result.content).toBe(""); + expect(result.error).toEqual({ + type: "processing", + description: "The target could not process this message.", + }); }); it("should handle multi-piece message with text + image", () => { @@ -808,6 +887,20 @@ describe("messageMapper", () => { expect(result.prompt_metadata).toEqual({ video_id: "sora-vid-456" }); }); + + it("should resubmit the raw backend value for a reconstructed attachment", async () => { + const att: MessageAttachment = { + type: "file", + name: "evidence.txt", + url: "/api/media?path=evidence.txt", + sourceValue: "stored attachment value", + mimeType: "text/plain", + }; + + const result = await attachmentToMessagePieceRequest(att); + + expect(result.original_value).toBe("stored attachment value"); + }); }); describe("reasoning summary extraction", () => { diff --git a/frontend/src/utils/messageMapper.ts b/frontend/src/utils/messageMapper.ts index 1868ca4ad0..f256f9581d 100644 --- a/frontend/src/utils/messageMapper.ts +++ b/frontend/src/utils/messageMapper.ts @@ -165,19 +165,52 @@ function pieceToAttachment( url, mimeType: mime, size, + sourceValue: value, pieceId: piece.id, metadata: piece.prompt_metadata || undefined, } } +/** + * Rebuild editable input using only a persisted message's original values. + */ +export function backendMessageToOriginalDraft( + msg: BackendMessage, +): Pick { + const textParts: string[] = [] + const attachments: MessageAttachment[] = [] + + for (const piece of msg.message_pieces) { + if (piece.original_value && !isMediaDataType(piece.original_value_data_type)) { + textParts.push(piece.original_value) + } + + const attachment = pieceToAttachment(piece, 'original') + if (attachment) { + attachments.push(attachment) + } + } + + return { + content: textParts.join('\n'), + attachments: attachments.length > 0 ? attachments : undefined, + } +} + /** * Extract an error from a backend message piece, if any. */ function pieceToError(piece: BackendMessagePiece): MessageError | undefined { if (piece.response_error && piece.response_error !== 'none') { + const fallbackDescriptions: Record = { + blocked: 'The target blocked this message.', + processing: 'The target could not process this message.', + empty: 'The target returned an empty response.', + unknown: 'The target returned an unknown error.', + } return { type: piece.response_error, - description: piece.response_error_description || undefined, + description: piece.response_error_description || fallbackDescriptions[piece.response_error], } } return undefined @@ -200,6 +233,11 @@ export function backendMessageToFrontend(msg: BackendMessage): Message { if (pieceError && !error) { error = pieceError } + if (pieceError?.type === 'processing') { + // Stored processing errors can contain exception tracebacks. Keep those + // diagnostics out of the chat transcript and recovery actions. + continue + } // Extract reasoning summaries from reasoning-type pieces if (isReasoningDataType(piece.converted_value_data_type)) { @@ -271,6 +309,8 @@ export async function attachmentToMessagePieceRequest(att: MessageAttachment): P let base64Value: string if (att.file) { base64Value = await fileToBase64(att.file) + } else if (att.sourceValue != null) { + base64Value = att.sourceValue } else if (att.url.startsWith('data:')) { base64Value = att.url.split(',')[1] || '' } else { diff --git a/pyrit/backend/mappers/_preview.py b/pyrit/backend/mappers/_preview.py index 33e5f1ea73..83a94d648e 100644 --- a/pyrit/backend/mappers/_preview.py +++ b/pyrit/backend/mappers/_preview.py @@ -66,9 +66,10 @@ def format_last_message_preview( Media-path data types are rendered as ``[Image: ]`` (and variants) so the absolute filesystem path of memory artifacts is never - exposed through API responses or UI previews. Text-like data types pass - through with truncation and an ellipsis suffix when they exceed - *max_len*. + exposed through API responses or UI previews. Error values are replaced + with a generic status so persisted exception tracebacks are not exposed. + Text-like data types pass through with truncation and an ellipsis suffix + when they exceed *max_len*. Args: value: Raw ``converted_value`` for the last piece (or ``None``). @@ -87,6 +88,9 @@ def format_last_message_preview( basename = _derive_basename(value or "") return f"[{label}: {basename}]" if basename else f"[{label}]" + if data_type == "error": + return "Target response error" + if not value: return None diff --git a/tests/unit/backend/test_preview.py b/tests/unit/backend/test_preview.py index 29e1dc498b..5b7b071d37 100644 --- a/tests/unit/backend/test_preview.py +++ b/tests/unit/backend/test_preview.py @@ -32,6 +32,14 @@ def test_unknown_data_type_treated_as_text(self) -> None: result = format_last_message_preview(value="hello", data_type=None, max_len=100) assert result == "hello" + def test_error_value_hides_persisted_diagnostics(self) -> None: + traceback = "RuntimeError: target failed\nTraceback (most recent call last):\nsecret details" + + result = format_last_message_preview(value=traceback, data_type="error", max_len=100) + + assert result == "Target response error" + assert "Traceback" not in result + def test_default_max_len_matches_conversation_stats_contract(self) -> None: # The formatter's default truncation length should track the model # constant so callers don't have to plumb it through manually.