From 4445dcfbaab8293802239bef24b5c721d752ecc7 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 23 Oct 2026 17:41:52 -0700 Subject: [PATCH 1/2] FIX Improve single-turn chat controls Keep the composer visible, explain single-turn restrictions, support first-prompt branching, and save objectives with Enter. Protect converter edits during prompt restoration. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 04870b03-21bf-46bb-8b0a-0ec0e2865965 --- .../components/Chat/ChatInputArea.test.tsx | 74 ++--- .../src/components/Chat/ChatInputArea.tsx | 17 +- .../src/components/Chat/ChatWindow.test.tsx | 277 +++++++++++++++++- frontend/src/components/Chat/ChatWindow.tsx | 122 ++++++-- .../src/components/Chat/MessageList.test.tsx | 19 ++ frontend/src/components/Chat/MessageList.tsx | 13 +- .../components/Chat/ObjectiveHeader.test.tsx | 69 ++++- .../src/components/Chat/ObjectiveHeader.tsx | 13 +- frontend/src/utils/conversationDraft.test.ts | 38 ++- frontend/src/utils/conversationDraft.ts | 11 + 10 files changed, 541 insertions(+), 112 deletions(-) diff --git a/frontend/src/components/Chat/ChatInputArea.test.tsx b/frontend/src/components/Chat/ChatInputArea.test.tsx index 828976c607..0436d1b104 100644 --- a/frontend/src/components/Chat/ChatInputArea.test.tsx +++ b/frontend/src/components/Chat/ChatInputArea.test.tsx @@ -38,7 +38,6 @@ describe("ChatInputArea", () => { const defaultProps = { onSend: jest.fn().mockResolvedValue(sentOutcome), disabled: false, - onNewConversation: jest.fn(), onUseAsTemplate: jest.fn(), onConfigureTarget: jest.fn(), onToggleConverterPanel: jest.fn(), @@ -105,14 +104,15 @@ describe("ChatInputArea", () => { it("should keep Undo reachable for an eligible first-turn undo when a single-turn prompt is disabled", async () => { const user = userEvent.setup(); const onUndo = jest.fn(); - render(); expect(screen.getByRole("textbox")).toBeDisabled(); expect(screen.getByRole("button", { name: "Toggle converter panel" })).toBeDisabled(); expect(screen.getByRole("button", { name: "Undo last send" })).toBeEnabled(); await user.click(screen.getByRole("button", { name: "Undo last send" })); expect(onUndo).toHaveBeenCalledTimes(1); - expect(screen.getByTestId("single-turn-banner")).toBeInTheDocument(); + expect(screen.queryByTestId("single-turn-banner")).not.toBeInTheDocument(); }); it("should retain the disabled Undo explanation in the prompt box", async () => { @@ -934,7 +934,7 @@ describe("ChatInputArea", () => { expect( screen.getByText( - /does not track conversation history/ + /only supports single-turn conversations/ ) ).toBeInTheDocument(); }); @@ -1200,62 +1200,30 @@ describe("ChatInputArea", () => { expect(screen.queryByText(/\(\s*\d+(?:\.\d+)?\s*[KM]?B\s*\)/)).toBeNull(); }); - it("should show single-turn banner when singleTurnLimitReached is true", () => { - render( - - - - ); - - expect(screen.getByTestId("single-turn-banner")).toBeInTheDocument(); - expect(screen.getByText(/only supports single-turn/)).toBeInTheDocument(); - expect(screen.getByTestId("new-conversation-btn")).toBeInTheDocument(); - // Input area should not be rendered - expect(screen.queryByRole("textbox")).not.toBeInTheDocument(); - }); - - it("should call onNewConversation when New Conversation button clicked", async () => { + it.each([false, true])("keeps the disabled single-turn composer visible with Undo present=%s", async (withUndo: boolean) => { const user = userEvent.setup(); - const onNewConversation = jest.fn(); - - render( - - - - ); - - await user.click(screen.getByTestId("new-conversation-btn")); - expect(onNewConversation).toHaveBeenCalledTimes(1); - }); - - it("should show New Conversation button when singleTurnLimitReached", () => { - render( - - - - ); - - expect(screen.getByTestId("single-turn-banner")).toBeInTheDocument(); - expect(screen.getByTestId("new-conversation-btn")).toBeInTheDocument(); + render(); + expect(screen.getByRole("textbox")).toBeDisabled(); + expect(screen.queryByTestId("single-turn-banner")).not.toBeInTheDocument(); + expect(screen.queryByTestId("new-conversation-btn")).not.toBeInTheDocument(); + await user.click(screen.getByLabelText("Why the prompt box is disabled")); + expect(screen.getByRole("listitem")).toHaveTextContent("This target only supports single-turn conversations."); + expect(screen.queryByText(/Wait for the current operation/)).not.toBeInTheDocument(); + await user.keyboard("{Escape}"); + await user.hover(getSendButton()); + expect(await screen.findByRole("tooltip")).toHaveTextContent("This target only supports single-turn conversations."); + await user.click(getSendButton()); + expect(defaultProps.onSend).not.toHaveBeenCalled(); }); - it("should show normal input when singleTurnLimitReached is false", () => { + it("should show normal input before a single-turn send", () => { render( ); diff --git a/frontend/src/components/Chat/ChatInputArea.tsx b/frontend/src/components/Chat/ChatInputArea.tsx index a11d75e380..0a0086c7e3 100644 --- a/frontend/src/components/Chat/ChatInputArea.tsx +++ b/frontend/src/components/Chat/ChatInputArea.tsx @@ -1,7 +1,7 @@ import { useState, useCallback, useEffect, useLayoutEffect, useRef, forwardRef, useImperativeHandle, type KeyboardEvent } from 'react' import { Button, Caption1, Popover, PopoverSurface, PopoverTrigger, Tooltip, Text } from '@fluentui/react-components' -import { SendRegular, DismissRegular, InfoRegular, AddRegular, CopyRegular, WarningRegular, SettingsRegular, ArrowSyncRegular } from '@fluentui/react-icons' +import { SendRegular, DismissRegular, InfoRegular, CopyRegular, WarningRegular, SettingsRegular, ArrowSyncRegular } from '@fluentui/react-icons' import type { AttackTargetResolutionStatus, ChatSendOutcome, ConvertedFileChip, MessageAttachment, MultiSendOptions, PieceConversion, TargetInstance, UndoLastSendAction } from '@/types' import { isTargetResolutionBlocking } from '@/utils/targetIdentity' @@ -120,8 +120,6 @@ interface ChatInputAreaProps { sendDisabled?: boolean sendDisabledReasons?: string[] activeTarget?: TargetInstance | null - singleTurnLimitReached?: boolean - onNewConversation?: () => void operatorLocked?: boolean crossTargetLocked?: boolean targetResolutionStatus?: AttackTargetResolutionStatus @@ -151,7 +149,7 @@ interface ChatInputAreaProps { const ChatInputArea = forwardRef(function ChatInputArea({ onSend, conversionRevisionKey = '', disabled = false, disabledReasons = [], sendDisabled = false, sendDisabledReasons = [], activeTarget, - singleTurnLimitReached = false, onNewConversation, operatorLocked = false, crossTargetLocked = false, + operatorLocked = false, crossTargetLocked = false, targetResolutionStatus = 'idle', onRetryTargetResolution, onUseAsTemplate, attackOperator, onConfigureTarget, onToggleConverterPanel, isConverterPanelOpen = false, onInputChange, onAttachmentsChange, convertedValue, onClearConversion, onClearAllConversions, onConvertedValueChange, @@ -288,17 +286,8 @@ const ChatInputArea = forwardRef(functi text="This attack uses a different target. Continue with your target to keep the conversation." buttonText="Continue with your target" buttonIcon={} onButtonClick={onUseAsTemplate} testId="cross-target-banner" buttonTestId="use-as-template-btn" buttonClassName={styles.touchTarget} /> - ) : singleTurnLimitReached && !undoLastSend ? ( - } - text="This target only supports single-turn conversations." buttonText="New Conversation" buttonIcon={} - onButtonClick={onNewConversation} testId="single-turn-banner" - buttonTestId="new-conversation-btn" buttonClassName={styles.touchTarget} /> ) : ( <> - {singleTurnLimitReached && } text="This target only supports single-turn conversations." - buttonText="New Conversation" buttonIcon={} onButtonClick={onNewConversation} - testId="single-turn-banner" buttonTestId="new-conversation-btn" buttonClassName={styles.touchTarget} />}
(functi {activeTarget?.capabilities?.supports_multi_turn === false && + content="This target only supports single-turn conversations." relationship="description"> } diff --git a/frontend/src/components/Chat/ChatWindow.test.tsx b/frontend/src/components/Chat/ChatWindow.test.tsx index c3717b2a54..feceb95aa0 100644 --- a/frontend/src/components/Chat/ChatWindow.test.tsx +++ b/frontend/src/components/Chat/ChatWindow.test.tsx @@ -1461,6 +1461,52 @@ describe("ChatWindow Integration", () => { expect(await screen.findByPlaceholderText("Type prompt here")).toHaveValue("Original prompt"); }); + it("restores matching in-memory converter stages when branching the first prompt", async () => { + const { user, selected } = await prepareUndoChat({ immutable: true, emptyHistory: true }); + const attack = await mockedAttacksApi.getAttack("undo-attack"); + mockedAttacksApi.saveConversation.mockResolvedValue({ + attack, messages: { conversation_id: "undo-branch", messages: [], target_response_status: null }, + }); + await user.click(screen.getAllByRole("button", { name: "Copy or Branch" })[0]); + await user.click(screen.getByRole("menuitem", { name: "Branch conversation", exact: true })); + await waitFor(() => expect(selected).toHaveBeenCalledWith("undo-branch")); + expect(screen.getByPlaceholderText("Type prompt here")).toHaveValue("Original prompt"); + expect(screen.getByRole("textbox", { name: "Stage 1 output - Text" })).toHaveValue("Custom middle text"); + expect(screen.getByRole("textbox", { name: "Stage 2 output - Text" })).toHaveValue("Custom middle text|second"); + expect(screen.getByRole("textbox", { name: "Converted prompt" })).toHaveValue("Custom middle text|second"); + expect(mockedConvertersApi.previewConversion).toHaveBeenCalledTimes(2); + expect(mockedAttacksApi.submitMessageSend).toHaveBeenCalledTimes(1); + expect(mockedAttacksApi.saveConversation).toHaveBeenCalledWith(expect.objectContaining({ + destination: "same_attack", messages: [], + })); + }); + + it.each(["pipeline", "working input"])("protects %s edits before first-prompt restore", async (edit: string) => { + const { user, selected } = await prepareUndoChat({ immutable: true, emptyHistory: true }); + if (edit === "pipeline") { + await user.click(screen.getByRole("button", { name: "Remove converter second" })); + } else { + await user.type(screen.getByRole("textbox", { name: "Working input - Text" }), "Keep my working input"); + } + expect(screen.getByRole("button", { name: "Undo last send" })).toBeDisabled(); + await user.click(screen.getAllByRole("button", { name: "Copy or Branch" })[0]); + const branch = screen.getByRole("menuitem", { name: "Branch conversation", exact: true }); + expect(branch).toHaveAttribute("aria-disabled", "true"); + expect(branch).toHaveAccessibleDescription(expect.stringMatching( + edit === "pipeline" ? /converter pipeline changed/i : /clear the current prompt and conversion edits/i, + )); + await user.click(branch); + expect(mockedAttacksApi.saveConversation).not.toHaveBeenCalled(); + expect(mockedAttacksApi.createConversation).not.toHaveBeenCalled(); + expect(selected).not.toHaveBeenCalled(); + if (edit === "pipeline") { + expect(screen.queryByRole("button", { name: "Remove converter second" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Remove converter first" })).toBeInTheDocument(); + } else { + expect(screen.getByRole("textbox", { name: "Working input - Text" })).toHaveValue("Keep my working input"); + } + }); + it("disables undo rather than overwriting a new prompt", async () => { const { user } = await prepareUndoChat(); await user.type(screen.getByPlaceholderText("Type prompt here"), "Keep my new draft"); @@ -5078,7 +5124,8 @@ describe("ChatWindow Integration", () => { await user.click(screen.getByRole("button", { name: "Send message" })); const recoveryButton = await findRecoveryAction(user); - expect(screen.getByTestId("single-turn-banner")).toBeInTheDocument(); + expect(screen.queryByTestId("single-turn-banner")).not.toBeInTheDocument(); + expect(screen.getByRole("textbox")).toBeDisabled(); await user.click(recoveryButton); await waitFor(() => { @@ -5323,7 +5370,209 @@ describe("ChatWindow Integration", () => { // Single-turn target UX // ----------------------------------------------------------------------- - it("should show single-turn banner for single-turn target with existing user messages", async () => { + describe("first-prompt branching", () => { + async function setupFirstPromptChat(options: { + sent?: boolean; imageResponse?: boolean; convertedPrompt?: boolean; withAttachment?: boolean; restorable?: boolean; + } = {}) { + const user = userEvent.setup(); + const target = makeTarget({ + target_registry_name: "http-target", target_type: "HTTPTarget", supports_history_replay: options.restorable ?? false, + capabilities: buildCapabilities({ + supports_multi_turn: false, supported_input_modalities: ["text", "image_path"], + supports_multi_message_pieces: true, + }), + }); + const attack: AttackSummary = { + attack_result_id: "first-attack", conversation_id: "first-source", attack_type: "ManualAttack", + objective: "Keep this objective", converters: [], message_count: 2, related_conversation_ids: [], + labels: {}, created_at: "", updated_at: "", + }; + const source = makeErrorResponse("none", "").messages; + source.messages[0].message_pieces[0] = { + ...source.messages[0].message_pieces[0], original_value: "Original request", + converted_value: options.convertedPrompt ? "Converted request" : "Original request", + prompt_metadata: { custom: "preserved" }, + converter_identifiers: options.convertedPrompt ? [{ type: "SavedConverter" }] : [], + }; + if (options.withAttachment) source.messages[0].message_pieces.push({ + id: "original-image", original_value_data_type: "image_path", original_value: "original.png", + converted_value_data_type: "image_path", converted_value: "original.png", + original_filename: "evidence.png", original_value_mime_type: "image/png", + scores: [], response_error: "none", + }); + source.messages[1].message_pieces[0] = { + ...source.messages[1].message_pieces[0], + original_value: options.imageResponse ? "response.png" : "Original response", + converted_value: options.imageResponse ? "response.png" : "Original response", + converted_value_data_type: options.imageResponse ? "image_path" : "text", + }; + const empty = (id: string) => ({ conversation_id: id, messages: [], target_response_status: null }); + const stored = new Map([ + ["first-source", options.sent === false ? empty("first-source") : { ...source, conversation_id: "first-source" }], + ["other", empty("other")], + ]); + mockedAttacksApi.getMessages.mockImplementation(async (_attackId, conversationId) => { + const messages = stored.get(conversationId ?? "first-source"); + if (!messages) throw new Error("Unexpected conversation"); + return messages; + }); + mockedAttacksApi.getAttack.mockResolvedValue(attack); + mockedAttacksApi.saveConversation.mockImplementation(async () => { + const messages = empty("first-branch"); + stored.set("first-branch", messages); + return { attack, messages }; + }); + mockSendResult.mockImplementation(async (_attackId, request) => { + const messages = { ...source, conversation_id: request.target_conversation_id }; + stored.set(request.target_conversation_id, messages); + return { attack, messages }; + }); + const selected = jest.fn(); + function BranchHarness() { + const [conversation, setConversation] = useState("first-source"); + return <> + + { selected(id); setConversation(id); }} /> + ; + } + render(); + if (options.sent !== false) await screen.findByText(options.convertedPrompt ? "Converted request" : "Original request"); + else await waitFor(() => expect(screen.getByPlaceholderText("Type prompt here")).toBeEnabled()); + return { user, source, stored, selected, attack }; + } + + async function branchFirstPrompt(user: ReturnType): Promise { + await user.click(screen.getAllByRole("button", { name: "Copy or Branch" })[0]); + const branch = screen.getByRole("menuitem", { name: "Branch conversation", exact: true }); + expect(branch).not.toHaveAttribute("aria-disabled", "true"); + expect(branch).toHaveAccessibleDescription( + "Restore this prompt as an unsent draft in a new conversation in the same attack. The response is not copied.", + ); + await user.click(branch); + } + + it.each([ + { imageResponse: false, restorable: false }, + { imageResponse: true, restorable: false }, + { imageResponse: false, restorable: true }, + ])("keeps the composer mounted after an HTTP send with image response=$imageResponse and restorable=$restorable and can branch and resend", async ({ + imageResponse, restorable, + }: { imageResponse: boolean; restorable: boolean }) => { + const { user, stored, selected } = await setupFirstPromptChat({ sent: false, imageResponse, restorable }); + const composer = screen.getByPlaceholderText("Type prompt here"); + await user.type(composer, "Original request"); + await user.click(screen.getByRole("button", { name: "Send message" })); + await screen.findByText("Original request"); + await waitFor(() => expect(composer).toBeDisabled()); + expect(screen.getByPlaceholderText("Type prompt here")).toBe(composer); + expect(screen.queryByTestId("single-turn-banner")).not.toBeInTheDocument(); + if (imageResponse) expect(screen.queryByRole("button", { name: "Undo last send" })).not.toBeInTheDocument(); + else await waitFor(() => expect(screen.getByRole("button", { name: "Undo last send" })).toBeEnabled()); + await user.click(screen.getByLabelText("Why the prompt box is disabled")); + expect(screen.getByRole("listitem")).toHaveTextContent("This target only supports single-turn conversations."); + await user.keyboard("{Escape}"); + for (const button of screen.getAllByRole("button", { name: "Copy or Branch" })) { + await user.click(button); + expect(screen.getByRole("menuitem", { name: "Copy to next message" })).toHaveAttribute("aria-disabled", "true"); + await user.keyboard("{Escape}"); + } + const sourceBeforeBranch = stored.get("first-source"); + await branchFirstPrompt(user); + await waitFor(() => expect(selected).toHaveBeenCalledWith("first-branch")); + expect(mockedAttacksApi.saveConversation).toHaveBeenCalledWith(expect.objectContaining({ + destination: "same_attack", attack_result_id: "first-attack", messages: [], + target_registry_name: "http-target", source_conversation_id: "first-source", + })); + expect(stored.get("first-source")).toBe(sourceBeforeBranch); + expect(composer).toBeEnabled(); + expect(composer).toHaveValue("Original request"); + expect(screen.getByRole("button", { name: "Edit objective" })).toHaveTextContent("Keep this objective"); + expect(mockedAttacksApi.createAttack).not.toHaveBeenCalled(); + expect(mockedAttacksApi.submitMessageSend).toHaveBeenCalledTimes(1); + await user.click(screen.getByRole("button", { name: "Send message" })); + await waitFor(() => expect(mockedAttacksApi.submitMessageSend).toHaveBeenCalledTimes(2)); + expect(mockedAttacksApi.submitMessageSend).toHaveBeenLastCalledWith("first-attack", expect.objectContaining({ + target_conversation_id: "first-branch", pieces: [expect.objectContaining({ original_value: "Original request" })], + })); + expect(mockedConvertersApi.previewConversion).not.toHaveBeenCalled(); + }); + + it("restores a saved first prompt with original values, conversions, attachments and metadata", async () => { + const { user } = await setupFirstPromptChat({ convertedPrompt: true, withAttachment: true }); + await branchFirstPrompt(user); + await waitFor(() => expect(screen.getByPlaceholderText("Type prompt here")).toBeEnabled()); + expect(screen.getByPlaceholderText("Type prompt here")).toHaveValue("Original request"); + expect(screen.getByRole("button", { name: "Remove evidence.png" })).toBeInTheDocument(); + expect(screen.getByText(/Converter choices could not be restored/)).toBeInTheDocument(); + expect(screen.queryByText("Original response")).not.toBeInTheDocument(); + expect(mockedAttacksApi.submitMessageSend).not.toHaveBeenCalled(); + await user.click(screen.getByRole("button", { name: "Send message" })); + await waitFor(() => expect(mockedAttacksApi.submitMessageSend).toHaveBeenCalledTimes(1)); + const request = mockedAttacksApi.submitMessageSend.mock.calls[0][1]; + expect(request.pieces[0]).toMatchObject({ + original_value: "Original request", converted_value: "Converted request", prompt_metadata: { custom: "preserved" }, + }); + expect(request.pieces[1]).toMatchObject({ original_value: "original.png", data_type: "image_path" }); + expect(mockedConvertersApi.previewConversion).not.toHaveBeenCalled(); + }); + + it("keeps assistant history blocked for a single-turn target", async () => { + const { user } = await setupFirstPromptChat(); + await user.click(screen.getAllByRole("button", { name: "Copy or Branch" })[1]); + expect(screen.getByRole("menuitem", { name: "Branch conversation", exact: true })).toHaveAttribute("aria-disabled", "true"); + expect(mockedAttacksApi.saveConversation).not.toHaveBeenCalled(); + }); + + it("rechecks the selected prompt before saving a branch", async () => { + const { user, source, stored } = await setupFirstPromptChat(); + stored.set("first-source", { + ...source, conversation_id: "first-source", messages: [{ + ...source.messages[0], message_pieces: [{ ...source.messages[0].message_pieces[0], id: "changed-piece" }], + }, source.messages[1]], + }); + await branchFirstPrompt(user); + expect(await screen.findByText(/selected prompt changed/)).toBeInTheDocument(); + expect(mockedAttacksApi.saveConversation).not.toHaveBeenCalled(); + expect(screen.getByPlaceholderText("Type prompt here")).toHaveValue(""); + }); + + it("does not replace the destination draft if navigation changes during a branch save", async () => { + const { user, stored, selected, attack } = await setupFirstPromptChat(); + let finishSave: ((response: AddMessageResponse) => void) | undefined; + mockedAttacksApi.saveConversation.mockImplementationOnce(() => new Promise((resolve) => { finishSave = resolve; })); + await branchFirstPrompt(user); + await waitFor(() => expect(mockedAttacksApi.saveConversation).toHaveBeenCalledTimes(1)); + await user.click(screen.getByRole("button", { name: "Switch conversation" })); + await waitFor(() => expect(screen.getByPlaceholderText("Type prompt here")).toBeEnabled()); + await user.type(screen.getByPlaceholderText("Type prompt here"), "Keep this new draft"); + const other = stored.get("other"); + if (!other || !finishSave) throw new Error("Expected the pending save and destination conversation"); + const completeSave = finishSave; + await act(async () => { completeSave({ attack, messages: { ...other, conversation_id: "first-branch" } }); }); + expect(await screen.findByText(/without replacing your current draft/)).toBeInTheDocument(); + expect(selected).not.toHaveBeenCalled(); + expect(screen.getByPlaceholderText("Type prompt here")).toHaveValue("Keep this new draft"); + }); + + it("preserves the sent conversation on a failed branch and reuses the save ID on retry", async () => { + const { user } = await setupFirstPromptChat(); + mockedAttacksApi.saveConversation.mockRejectedValueOnce(new Error("Connection lost")); + await branchFirstPrompt(user); + expect(await screen.findByText("Connection lost")).toBeInTheDocument(); + expect(screen.getByText("Original response")).toBeInTheDocument(); + expect(screen.getByPlaceholderText("Type prompt here")).toHaveValue(""); + const saveId = mockedAttacksApi.saveConversation.mock.calls[0][0].save_id; + await branchFirstPrompt(user); + await waitFor(() => expect(screen.getByPlaceholderText("Type prompt here")).toBeEnabled()); + expect(mockedAttacksApi.saveConversation.mock.calls[1][0].save_id).toBe(saveId); + }); + }); + + it("keeps the composer visible and explains the single-turn restriction on saved chats", async () => { + const user = userEvent.setup(); const singleTurnTarget: TargetInstance = makeTarget({ target_registry_name: "openai_image_1", target_type: "OpenAIImageTarget", @@ -5351,10 +5600,12 @@ describe("ChatWindow Integration", () => { ); await waitFor(() => { - expect(screen.getByTestId("single-turn-banner")).toBeInTheDocument(); - expect(screen.getByText(/only supports single-turn/)).toBeInTheDocument(); + expect(screen.getByRole("textbox")).toBeDisabled(); }); - expect(screen.queryByRole("textbox")).not.toBeInTheDocument(); + expect(screen.queryByTestId("single-turn-banner")).not.toBeInTheDocument(); + await user.click(screen.getByLabelText("Why the prompt box is disabled")); + expect(screen.getByRole("listitem")).toHaveTextContent("This target only supports single-turn conversations."); + expect(screen.queryByText(/Wait for the current operation/)).not.toBeInTheDocument(); }); it("should not show single-turn banner for single-turn target with no messages", () => { @@ -5406,7 +5657,7 @@ describe("ChatWindow Integration", () => { expect(screen.getByRole("textbox")).toBeInTheDocument(); }); - it("should show New Conversation button in single-turn banner when conversation exists", async () => { + it("keeps the composer visible without a single-turn banner when Undo is unavailable", async () => { const singleTurnTarget: TargetInstance = makeTarget({ target_registry_name: "openai_tts_1", target_type: "OpenAITTSTarget", @@ -5434,8 +5685,10 @@ describe("ChatWindow Integration", () => { ); await waitFor(() => { - expect(screen.getByTestId("new-conversation-btn")).toBeInTheDocument(); + expect(screen.getByRole("textbox")).toBeDisabled(); }); + expect(screen.queryByTestId("single-turn-banner")).not.toBeInTheDocument(); + expect(screen.queryByTestId("new-conversation-btn")).not.toBeInTheDocument(); }); it("should show cross-target banner when attackTarget differs from activeTarget", () => { @@ -5696,6 +5949,7 @@ describe("ChatWindow Integration", () => { // ----------------------------------------------------------------------- it("should create a new conversation and select it via handleNewConversation", async () => { + const user = userEvent.setup(); const onSelectConversation = jest.fn(); mockedAttacksApi.createConversation.mockResolvedValue({ conversation_id: "new-conv-from-new", @@ -5728,12 +5982,9 @@ describe("ChatWindow Integration", () => { ); - // For single-turn targets with existing messages, there's a New Conversation button - await waitFor(() => { - expect(screen.getByTestId("new-conversation-btn")).toBeInTheDocument(); - }); - - await userEvent.click(screen.getByTestId("new-conversation-btn")); + await screen.findByText("Generate an image"); + await user.click(screen.getByRole("button", { name: "Toggle conversations panel" })); + await user.click(await screen.findByTestId("new-conversation-btn")); await waitFor(() => { expect(mockedAttacksApi.createConversation).toHaveBeenCalledWith("ar-new-conv", {}); diff --git a/frontend/src/components/Chat/ChatWindow.tsx b/frontend/src/components/Chat/ChatWindow.tsx index e438282b96..5b74911c19 100644 --- a/frontend/src/components/Chat/ChatWindow.tsx +++ b/frontend/src/components/Chat/ChatWindow.tsx @@ -41,7 +41,7 @@ import ObjectiveHeader from './ObjectiveHeader' import ConversationEditor from './ConversationEditor' import type { ConversationEditorHandle } from './ConversationEditor' import { - attachmentDraft, draftDataTypes, draftMessageToComposer, editorTargetDisabledReason, historyReplayDisabledReason, + attachmentDraft, canBranchFirstPrompt, draftDataTypes, draftMessageToComposer, editorTargetDisabledReason, historyReplayDisabledReason, newDraftPiece, toConversationDraft, } from '@/utils/conversationDraft' import { useConversationSave } from '@/hooks/useConversationSave' @@ -79,16 +79,19 @@ import type { AttackSummary, AttackTargetResolutionStatus, BackendMessage, + BackendMessagePiece, BackendScore, ChatSendOutcome, ChatConverterSnapshot, ConversationMessagesResponse, ConversationDraftPiece, + ConverterInputPiece, ConverterPipelineStage, CreateAttackRequest, CreateConversationRequest, Message, MessageAttachment, + MessageDisplayPiece, MessageSendConversation, MessageSendRequest, MessageSendStatus, @@ -1435,10 +1438,11 @@ export default function ChatWindow({ /** 1. Copy the clicked message's content/attachments into the current conversation's input box */ const handleCopyToInput = useCallback((messageIndex: number) => { + if (activeTarget?.capabilities?.supports_multi_turn === false && messages.some((message: Message) => message.role === 'user')) return const msg = messages[messageIndex] if (!msg) { return } copyMessageToInput(msg) - }, [copyMessageToInput, messages]) + }, [activeTarget, copyMessageToInput, messages]) const handleChangeMainConversation = useCallback(async (convId: string) => { if ( @@ -1590,6 +1594,32 @@ export default function ChatWindow({ return pending } + const promptRestoreDisabledReason = (snapshot?: ChatConverterSnapshot): string | undefined => { + if (converters.isConverting) return 'Wait for the current conversion to finish before restoring a sent prompt.' + if (chatInputText || draftAttachments.length || Object.entries(converters.workingInputs).some( + ([id, value]: [string, string]) => value !== converters.inputs.find((input: ConverterInputPiece) => input.id === id)?.value, + ) || Object.keys(converters.applied).length) { + return 'Clear the current prompt and conversion edits before restoring a sent prompt.' + } + if (snapshot + ? JSON.stringify(converters.pipelines) !== JSON.stringify(snapshot.pipelines) + : Object.values(converters.pipelines).some((stages: ConverterPipelineStage[]) => stages.length > 0)) { + return snapshot + ? 'The converter pipeline changed after sending. Restore is unavailable for this prompt.' + : 'Clear the current converter pipeline before restoring a stored prompt.' + } + if (snapshot && Object.keys(snapshot.errors).length > 0) { + return 'The previous conversion was incomplete. Restore is unavailable for this prompt.' + } + return undefined + } + + const firstPromptSend = undoableSend?.attackResultId === attackResultId + && undoableSend.conversationId === viewedConversationId + && undoableSend.converters.generation === runtime.generation + && undoableSend.history.every((message: Pick) => message.role === 'system' || message.role === 'developer') + ? undoableSend : undefined + const copyConversation = async ( messageIndex: number, destination: ConversationCopyDestination, ): Promise => { @@ -1603,6 +1633,15 @@ export default function ChatWindow({ } const sourceId = viewedConversationId const draftRevision = inputBoxRef.current?.getDraftRevision() + const navigationRevision = navigationRevisionRef.current + const selectedPieceIds = messages[messageIndex]?.displayPieces?.map((piece: MessageDisplayPiece) => piece.pieceId) ?? [] + const restoresFirstPrompt = destination === 'same_attack' && activeTarget + && canBranchFirstPrompt(activeTarget, messages, messageIndex) + const draftUnchanged = (): boolean => viewedAttackRef.current === attackResultId + && getViewedConversationId() === sourceId + && navigationRevisionRef.current === navigationRevision + && inputBoxRef.current?.getDraftRevision() === draftRevision + && launchStateRef.current.generation === runtime.generation && launchStateRef.current.ready copyingRef.current = true setIsLoadingEdit(true) setEditorError(null) @@ -1620,6 +1659,29 @@ export default function ChatWindow({ ? source.messages.filter((message: BackendMessage) => message.turn_number <= recoverableSend.historyCutoffIndex) : source.messages.slice(0, messageIndex + 1) let copiedMessages = toConversationDraft(sourceMessages) + let restoredPrompt: ReturnType | undefined + let promptSnapshot: ChatConverterSnapshot | undefined + let missingConverterSelections = false + if (restoresFirstPrompt && activeTarget) { + const selected = source.messages[messageIndex] + if (!canBranchFirstPrompt(activeTarget, source.messages, messageIndex) + || selectedPieceIds.some((id: string) => !selected.message_pieces.some((piece: BackendMessagePiece) => piece.id === id))) { + throw new Error('The selected prompt changed. Reload the conversation before branching.') + } + if (!draftUnchanged()) throw new Error('The view, prompt, or runtime changed. Your current draft was not replaced.') + const saved = firstPromptSend?.cutoffIndex === selected.turn_number - 1 + && firstPromptSend.historySignature === conversationPieceSignature(source) ? firstPromptSend : undefined + const restoreError = promptRestoreDisabledReason(saved?.converters) + if (restoreError) throw new Error(restoreError) + promptSnapshot = saved?.converters ?? { + generation: runtime.generation, pipelines: {}, workingInputs: {}, stageResults: {}, errors: {}, applied: {}, + } + restoredPrompt = saved ? { + text: saved.originalValue, attachments: saved.attachments, textMetadata: saved.textMetadata, converters: saved.converters, + } : draftMessageToComposer(copiedMessages[copiedMessages.length - 1], promptSnapshot) + missingConverterSelections = !saved && selected.message_pieces.some((piece: BackendMessagePiece) => Boolean(piece.converter_identifiers?.length)) + copiedMessages = copiedMessages.slice(0, -1) + } if (isFailedPrompt && activeTarget && editorTargetDisabledReason(activeTarget, draftDataTypes(copiedMessages), copiedMessages)) { copiedMessages = [] @@ -1657,6 +1719,11 @@ export default function ChatWindow({ labels, messages: copiedMessages, }, destination === 'same_attack' ? 'same_attack' : 'new_attack') + setPanelRefreshKey((key: number) => key + 1) + if (restoredPrompt && !draftUnchanged()) { + setEditorNotice(`Branch created conversation ${response.messages.conversation_id} without replacing your current draft.`) + return + } if (getViewedConversationId() !== sourceId) return if (destination === 'technique') { const params = new URLSearchParams({ prepended_conversation_id: response.messages.conversation_id }) @@ -1664,12 +1731,22 @@ export default function ChatWindow({ return } if (isFailedPrompt && inputBoxRef.current?.getDraftRevision() !== draftRevision) return - if (destination === 'same_attack') onSelectConversation(response.messages.conversation_id) + if (destination === 'same_attack') { + if (restoredPrompt) { + setMessages(backendMessagesToFrontend(response.messages.messages)) + markConversationLoaded(response.messages.conversation_id) + } + onSelectConversation(response.messages.conversation_id) + } else onConversationCreated(response.attack.attack_result_id, response.messages.conversation_id, response.attack.objective, target) onAttackChange?.(response.attack) if (isFailedPrompt) restoreRecoverableDraft(response.messages.conversation_id) + if (restoredPrompt && promptSnapshot) { + restoreEditorPrompt(restoredPrompt, promptSnapshot) + setUndoableSend(null) + setRecoveredDraftWarningConversationId(missingConverterSelections ? response.messages.conversation_id : null) + } setIsPanelOpen(true) - setPanelRefreshKey((key: number) => key + 1) } catch (error) { if (getViewedConversationId() === sourceId) setEditorError(toApiError(error).detail) } finally { @@ -1706,6 +1783,10 @@ export default function ChatWindow({ && undoableSend.conversationId === viewedConversationId ? undoableSend.history : messages const showUndoLastSend = !activeTarget || !historyReplayDisabledReason(activeTarget, undoHistory) const branchConversationDisabledReason = (index: number): string | undefined => { + if (activeTarget && canBranchFirstPrompt(activeTarget, messages, index) + && recoverableSend?.failedRequestMessageIndex !== index) { + return promptRestoreDisabledReason(firstPromptSend?.converters) + } if (targetResolutionStatus === 'unbound') return undefined if (!activeTarget) return 'Configure a target before branching this conversation.' if (recoverableSend?.failedRequestMessageIndex === index) { @@ -1730,15 +1811,7 @@ export default function ChatWindow({ ? 'Open the conversation containing the last send to undo it.' : activeTarget && historyReplayDisabledReason(activeTarget, undoableSend.history) ? 'This target cannot replay edited conversation history.' - : (chatInputText || draftAttachments.length || Object.entries(converters.workingInputs).some( - ([id, value]: [string, string]) => value !== converters.inputs.find((input) => input.id === id)?.value, - ) || Object.keys(converters.applied).length) - ? 'Clear the current prompt and conversion edits before undoing the previous send.' - : JSON.stringify(converters.pipelines) !== JSON.stringify(undoableSend.converters.pipelines) - ? 'The converter pipeline changed after sending. Undo is unavailable for this turn.' - : Object.keys(undoableSend.converters.errors).length > 0 - ? 'The previous conversion was incomplete. Undo is unavailable for this turn.' - : undefined + : promptRestoreDisabledReason(undoableSend.converters) const handleUndoLastSend = async (): Promise => { if (undoDisabledReason || !undoableSend || undoInFlight.current) return @@ -1800,12 +1873,18 @@ export default function ChatWindow({ } const singleTurnLimitReached = activeTarget?.capabilities?.supports_multi_turn === false && messages.some(m => m.role === 'user') + const singleTurnDisabledReason = singleTurnLimitReached ? 'This target only supports single-turn conversations.' : undefined const hasProcessingError = Boolean(recoverableSend) const inputDisabledReasons = [ + ...(singleTurnDisabledReason ? [singleTurnDisabledReason] : []), ...(isUndoingSend ? ['The previous send is being restored. Wait for undo to finish.'] : []), ...(!runtime.ready ? ['The runtime is not ready. Wait for initialization to finish.'] : []), ...(isSending ? ['A message is being sent. Wait for the response.'] : []), ...(isLoadingAttack ? ['The attack is loading. Wait for it to finish.'] : []), + ...(editDraft !== null ? ['Finish editing the conversation before sending.'] : []), + ...(isOperatorLocked ? ['This attack belongs to another operator.'] : []), + ...(isCrossTargetLocked ? ['This attack uses a different target.'] : []), + ...(isTargetResolutionLocked ? ['The target used by this attack could not be resolved.'] : []), ...(isEvidenceUnavailable ? ['This finding evidence is unavailable. Sending is disabled.'] : []), ...(hasProcessingError ? ['This conversation contains a target error. Use Copy or Branch on the failed prompt to branch into a new conversation or attack.'] @@ -2092,12 +2171,16 @@ export default function ChatWindow({ {editDraft === null && { void copyConversation(index, 'same_attack') }} onCopyToNewAttack={newAttackDisabledReason ? undefined : (index: number) => { void copyConversation(index, 'new_attack') }} onUseInTechnique={newAttackDisabledReason ? undefined : (index: number) => { void copyConversation(index, 'technique') }} copyConversationDisabled={isSending || isLoadingEdit || isRecoveringProcessingError || Boolean(sendIssue?.blocking)} newConversationDisabledReason={isMutationLocked ? "This attack is read-only. Copy to a new attack instead." : undefined} getBranchConversationDisabledReason={branchConversationDisabledReason} + getBranchConversationDescription={(index: number) => activeTarget && canBranchFirstPrompt(activeTarget, messages, index) + ? 'Restore this prompt as an unsent draft in a new conversation in the same attack. The response is not copied.' + : undefined} isLoading={isLoadingAttack || isLoadingMessages || awaitingConversationLoad} globalMarkdown={globalMarkdown} onEditConversation={() => { void beginEdit() }} @@ -2163,19 +2246,8 @@ export default function ChatWindow({ supportsSystemPrompt={supportsSystemPrompt} systemPrompt={systemPrompt} onSystemPromptChange={setSystemPrompt} - disabled={ - !runtime.ready - || isSending - || editDraft !== null - || isLoadingAttack - || singleTurnLimitReached - || isMutationLocked - || hasProcessingError - || isUndoingSend - } + disabled={inputDisabledReasons.length > 0} activeTarget={activeTarget} - singleTurnLimitReached={singleTurnLimitReached} - onNewConversation={handleNewConversation} operatorLocked={isOperatorLocked} crossTargetLocked={isCrossTargetLocked} targetResolutionStatus={targetResolutionStatus} diff --git a/frontend/src/components/Chat/MessageList.test.tsx b/frontend/src/components/Chat/MessageList.test.tsx index d9e862207a..ff45c8b989 100644 --- a/frontend/src/components/Chat/MessageList.test.tsx +++ b/frontend/src/components/Chat/MessageList.test.tsx @@ -90,6 +90,25 @@ describe("MessageList", () => { expect(screen.getByRole("menuitem", { name: "Branch attack" })).not.toHaveAttribute("aria-disabled", "true"); }); + it("disables next-message copying without disabling branching", async () => { + const user = userEvent.setup(); + const onCopyToInput = jest.fn(); + const reason = "This target only supports single-turn conversations."; + render(); + for (const button of screen.getAllByRole("button", { name: "Copy or Branch" })) { + await user.click(button); + const copy = screen.getByRole("menuitem", { name: "Copy to next message" }); + expect(copy).toHaveAttribute("aria-disabled", "true"); + expect(copy).toHaveAccessibleDescription(reason); + await user.click(copy); + expect(onCopyToInput).not.toHaveBeenCalled(); + expect(screen.getByRole("menuitem", { name: "Branch conversation" })).not.toHaveAttribute("aria-disabled", "true"); + await user.keyboard("{Escape}"); + } + }); + it("should show an accessible description for every copy or branch action and support keyboard selection", async () => { const user = userEvent.setup(); const onUseInTechnique = jest.fn(); diff --git a/frontend/src/components/Chat/MessageList.tsx b/frontend/src/components/Chat/MessageList.tsx index 23b672684c..5281361791 100644 --- a/frontend/src/components/Chat/MessageList.tsx +++ b/frontend/src/components/Chat/MessageList.tsx @@ -40,12 +40,14 @@ interface MessageListProps { messages: Message[] /** Copy this message to the input box of the current conversation */ onCopyToInput?: (messageIndex: number) => void + copyToInputDisabledReason?: string onCopyToNewConversation?: (messageIndex: number) => void onCopyToNewAttack?: (messageIndex: number) => void onUseInTechnique?: (messageIndex: number) => void copyConversationDisabled?: boolean newConversationDisabledReason?: string getBranchConversationDisabledReason?: (messageIndex: number) => string | undefined + getBranchConversationDescription?: (messageIndex: number) => string | undefined /** True while loading a historical attack's messages */ isLoading?: boolean /** Conversation-wide default: render message text as Markdown. */ @@ -436,7 +438,8 @@ function MessageScores({ scores, groupId }: { scores: DisplayScore[]; groupId: s */ export default function MessageList({ messages, onCopyToInput, onCopyToNewConversation, onCopyToNewAttack, onUseInTechnique, copyConversationDisabled = false, - newConversationDisabledReason, getBranchConversationDisabledReason, isLoading, globalMarkdown = false, onEditConversation, + newConversationDisabledReason, getBranchConversationDisabledReason, getBranchConversationDescription, copyToInputDisabledReason, + isLoading, globalMarkdown = false, onEditConversation, }: MessageListProps) { const styles = useMessageListStyles() const messagesEndRef = useRef(null) @@ -679,15 +682,17 @@ export default function MessageList({ >Copy or Branch - onCopyToInput(index)} + { if (!copyToInputDisabledReason) onCopyToInput(index) }} aria-label="Copy to next message" aria-describedby={`${copyMenuId}-${index}-prompt`} - subText={{ id: `${copyMenuId}-${index}-prompt`, children: 'Copy only this message to the prompt box.' }}> + subText={{ id: `${copyMenuId}-${index}-prompt`, + children: copyToInputDisabledReason ?? 'Copy only this message to the prompt box.' }}> Copy to next message onCopyToNewConversation?.(index)}>Branch conversation { + beforeEach(() => jest.clearAllMocks()) + + it.each([ + { objective: '', next: ' New goal ', draftMode: false }, + { objective: 'Existing goal', next: 'Changed goal', draftMode: false }, + { objective: 'Existing goal', next: '', draftMode: false }, + { objective: 'Existing goal', next: 'Local goal', draftMode: true }, + ])('saves objective "$next" with Enter in draft mode $draftMode', async ({ objective, next, draftMode }) => { + const user = userEvent.setup() + const onAdd = jest.fn().mockResolvedValue(undefined) + const onKeyDown = jest.fn() + render(
+ +
) + await user.click(screen.getByRole('button', { name: objective ? 'Edit objective' : 'Add objective' })) + const input = screen.getByRole('textbox', { name: 'Attack objective' }) + await user.clear(input) + if (next) await user.type(input, next) + onKeyDown.mockClear() + await user.keyboard('{Enter}') + await waitFor(() => expect(onAdd).toHaveBeenCalledWith(next.trim(), objective)) + expect(onAdd).toHaveBeenCalledTimes(1) + expect(onKeyDown).not.toHaveBeenCalled() + expect(screen.queryByRole('textbox', { name: 'Attack objective' })).not.toBeInTheDocument() + }) + + it('does not save an empty new objective with Enter', async () => { + const user = userEvent.setup() + const onAdd = jest.fn() + render() + await user.click(screen.getByRole('button', { name: 'Add objective' })) + await user.keyboard('{Enter}') + expect(onAdd).not.toHaveBeenCalled() + expect(screen.getByRole('button', { name: 'Save' })).toBeDisabled() + }) + + it('does not save during input-method composition', async () => { + const user = userEvent.setup() + const onAdd = jest.fn() + render() + await user.click(screen.getByRole('button', { name: 'Edit objective' })) + fireEvent.keyDown(screen.getByRole('textbox', { name: 'Attack objective' }), { key: 'Enter', isComposing: true }) + expect(onAdd).not.toHaveBeenCalled() + }) + + it('prevents duplicate objective saves and keeps a failed draft available for retry', async () => { + const user = userEvent.setup() + let rejectSave: (error: Error) => void = () => {} + const onAdd = jest.fn().mockImplementationOnce(() => new Promise((_resolve, reject) => { + rejectSave = reject + })).mockResolvedValue(undefined) + render() + await user.click(screen.getByRole('button', { name: 'Edit objective' })) + const input = screen.getByRole('textbox', { name: 'Attack objective' }) + await user.type(input, ' changed{Enter}{Enter}') + await user.click(screen.getByRole('button', { name: 'Saving...' })) + expect(onAdd).toHaveBeenCalledTimes(1) + expect(input).toBeDisabled() + await act(async () => { rejectSave(new Error('Save failed')) }) + expect(await screen.findByRole('alert')).toHaveTextContent('Unable to save the objective.') + expect(input).toHaveValue('Goal changed') + await user.click(input) + await user.keyboard('{Enter}') + await waitFor(() => expect(onAdd).toHaveBeenCalledTimes(2)) + expect(onAdd).toHaveBeenLastCalledWith('Goal changed', 'Goal') + }) + afterEach(() => { delete (HTMLElement.prototype as { scrollWidth?: number }).scrollWidth delete (HTMLElement.prototype as { clientWidth?: number }).clientWidth diff --git a/frontend/src/components/Chat/ObjectiveHeader.tsx b/frontend/src/components/Chat/ObjectiveHeader.tsx index de96dab4c4..790a9cd5e3 100644 --- a/frontend/src/components/Chat/ObjectiveHeader.tsx +++ b/frontend/src/components/Chat/ObjectiveHeader.tsx @@ -1,4 +1,5 @@ import { useId, useLayoutEffect, useRef, useState } from 'react' +import type { KeyboardEvent } from 'react' import { Badge, @@ -82,6 +83,7 @@ export default function ObjectiveHeader({ const [draft, setDraft] = useState('') const [initialObjective, setInitialObjective] = useState('') const [isSaving, setIsSaving] = useState(false) + const savePending = useRef(false) const [error, setError] = useState('') const initialVerdict = scoreVerdict(humanScore ?? automatedScore) const [humanVerdict, setHumanVerdict] = useState<'success' | 'failure'>( @@ -122,8 +124,9 @@ export default function ObjectiveHeader({ const handleSave = async (): Promise => { const trimmedObjective = draft.trim() - if ((!trimmedObjective && missingObjective) || !onAdd) return + if (savePending.current || (!trimmedObjective && missingObjective) || !onAdd) return + savePending.current = true setIsSaving(true) setError('') try { @@ -133,6 +136,7 @@ export default function ObjectiveHeader({ } catch { setError('Unable to save the objective.') } finally { + savePending.current = false setIsSaving(false) } } @@ -322,6 +326,13 @@ export default function ObjectiveHeader({ className={styles.input} value={draft} onChange={(_event, data) => setDraft(data.value)} + onKeyDown={(event: KeyboardEvent) => { + if (event.key !== 'Enter' || event.nativeEvent.isComposing) return + event.preventDefault() + event.stopPropagation() + void handleSave() + }} + disabled={isSaving} placeholder="Enter an objective" aria-label="Attack objective" autoFocus diff --git a/frontend/src/utils/conversationDraft.test.ts b/frontend/src/utils/conversationDraft.test.ts index 0bfa897e7f..08ae2df371 100644 --- a/frontend/src/utils/conversationDraft.test.ts +++ b/frontend/src/utils/conversationDraft.test.ts @@ -4,10 +4,46 @@ import { makeTarget } from '@/test-utils/targetFixtures' import { buildMessagePieces } from './messageMapper' import { draftConverterInputs, draftDataTypes, draftMessageToComposer, draftMessageToInput, editorTargetDisabledReason, - historyReplayDisabledReason, serializeDraft, toConversationDraft, unansweredToolCallId, validateDraft, + canBranchFirstPrompt, historyReplayDisabledReason, serializeDraft, toConversationDraft, unansweredToolCallId, validateDraft, } from './conversationDraft' describe('conversation drafts', () => { + it.each([ + { roles: ['user', 'assistant'], index: 0, multiTurn: false, system: false, allowed: true }, + { roles: ['user', 'assistant'], index: 0, multiTurn: true, system: false, allowed: true }, + { roles: ['user', 'assistant'], index: 1, multiTurn: false, system: false, allowed: false }, + { roles: ['user', 'assistant', 'user'], index: 2, multiTurn: true, system: false, allowed: false }, + { roles: ['system', 'developer', 'user'], index: 2, multiTurn: true, system: true, allowed: true }, + { roles: ['system', 'user'], index: 1, multiTurn: false, system: true, allowed: false }, + { roles: ['system', 'user'], index: 1, multiTurn: true, system: false, allowed: false }, + { roles: ['simulated_assistant', 'user'], index: 1, multiTurn: true, system: true, allowed: false }, + { roles: ['tool', 'user'], index: 1, multiTurn: true, system: true, allowed: false }, + { roles: [], index: 0, multiTurn: false, system: false, allowed: false }, + ])('allows first-prompt branching=$allowed for $roles at $index', ({ roles, index, multiTurn, system, allowed }) => { + const target = makeTarget({ + target_registry_name: 'http', capabilities: { + supports_multi_turn: multiTurn, supports_editable_history: false, supports_system_prompt: system, + }, + }) + expect(canBranchFirstPrompt(target, roles.map((role: string) => ({ role })), index)).toBe(allowed) + }) + + it('keeps normal history branching for targets with effective replay support', () => { + const target = makeTarget({ target_registry_name: 'restorable', supports_history_replay: true }) + expect(canBranchFirstPrompt(target, [{ role: 'user' }], 0)).toBe(false) + }) + + it('restores the first prompt for an adapted single-turn target', () => { + const target = makeTarget({ + target_registry_name: 'adapted', supports_history_replay: true, + capabilities: { supports_multi_turn: false, supports_editable_history: false }, + }) + const messages = [{ role: 'user' }, { role: 'assistant' }] + expect(canBranchFirstPrompt(target, messages, 0)).toBe(true) + expect(canBranchFirstPrompt(target, messages, 1)).toBe(false) + expect(canBranchFirstPrompt(target, [...messages, { role: 'user' }], 2)).toBe(false) + }) + it('returns original text and media with separate conversions and remapped converter stages', () => { const stage = { id: 'stage', converterId: 'base64' } const snapshot: ChatConverterSnapshot = { diff --git a/frontend/src/utils/conversationDraft.ts b/frontend/src/utils/conversationDraft.ts index a91067314a..668fe90e09 100644 --- a/frontend/src/utils/conversationDraft.ts +++ b/frontend/src/utils/conversationDraft.ts @@ -39,6 +39,17 @@ export function editorTargetDisabledReason( return undefined } +export function canBranchFirstPrompt( + target: TargetInstance, messages: ReadonlyArray<{ readonly role: string }>, index: number, +): boolean { + const history = messages.slice(0, index) + return messages[index]?.role === 'user' + && history.every((message: { readonly role: string }) => message.role === 'system' || message.role === 'developer') + && !historyReplayDisabledReason(target, history) + && (target.capabilities?.supports_multi_turn === false + || Boolean(historyReplayDisabledReason(target, messages.slice(0, index + 1)))) +} + export function unansweredToolCallId(messages: ConversationDraftMessage[], messageId: string): string | undefined { const responses = new Set() const calls: string[] = [] From 5f93c22924a1ab9d074e771d8d87e0ef68a4547e Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 27 Oct 2026 08:33:19 -0700 Subject: [PATCH 2/2] Fix prompt restoration after multimodal responses Keep the sent prompt snapshot for every successful response modality so Undo and first-prompt branching preserve converter stages for image, audio, video, and mixed results. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- frontend/README.md | 7 +++ .../src/components/Chat/ChatWindow.test.tsx | 52 ++++++++++++++----- frontend/src/components/Chat/ChatWindow.tsx | 1 - 3 files changed, 46 insertions(+), 14 deletions(-) diff --git a/frontend/README.md b/frontend/README.md index 6895cc1c59..38038a9cf3 100644 --- a/frontend/README.md +++ b/frontend/README.md @@ -29,6 +29,13 @@ conversations), open the conversation panel, and select the copy at its full `/attacks/{attack_id}/conversations/{conversation_id}` URL. Copying within a target-free attack keeps it target-free until you send a prompt. +For single-turn targets or targets that cannot replay the first prompt as history, +**Branch conversation** on that first user turn restores it to the prompt box +without copying the response. After a successful single send, **Undo last send** +also restores the prompt in a new conversation, if the target can replay the +preceding history. Both actions restore the saved prompt and converter results +for text, image, audio, video, and mixed responses. They do not overwrite new +prompt edits or changed converter pipelines. When you save an edited conversation, its final user turn moves to the prompt box instead of saved history. Text and attachments retain their metadata, including text-only video remix IDs. You can enter and edit prompts without a target. diff --git a/frontend/src/components/Chat/ChatWindow.test.tsx b/frontend/src/components/Chat/ChatWindow.test.tsx index 230cdf3d5a..23dd85e56d 100644 --- a/frontend/src/components/Chat/ChatWindow.test.tsx +++ b/frontend/src/components/Chat/ChatWindow.test.tsx @@ -1461,6 +1461,7 @@ describe("ChatWindow Integration", () => { async function prepareUndoChat(options: { immutable?: boolean; emptyHistory?: boolean; assistantHistory?: boolean; restorable?: boolean; + responseDataTypes?: string[]; } = {}) { const user = userEvent.setup(); const prefix: BackendMessage[] = [{ @@ -1539,10 +1540,13 @@ describe("ChatWindow Integration", () => { }], }, { role: "assistant", turn_number: prefix.length + 1, created_at: "", - message_pieces: [{ - id: "last-assistant", original_value_data_type: "text", original_value: "Accepted result", - converted_value_data_type: "text", converted_value: "Accepted result", scores: [], response_error: "none", - }], + message_pieces: (options.responseDataTypes ?? ["text"]).map((dataType: string, index: number) => ({ + id: `last-assistant-${index}`, original_value_data_type: dataType, + original_value: dataType === "text" ? "Accepted result" : `response-${index}`, + converted_value_data_type: dataType, + converted_value: dataType === "text" ? "Accepted result" : `response-${index}`, + scores: [], response_error: "none", + })), }], }; return { attack, messages: source }; @@ -1568,7 +1572,7 @@ describe("ChatWindow Integration", () => { await user.type(middle, "Custom middle text"); await user.click(screen.getByRole("button", { name: "Convert Text from stage 2 to end" })); await user.click(screen.getByRole("button", { name: "Send message" })); - await screen.findByText("Accepted result"); + await waitFor(() => expect(screen.getByPlaceholderText("Type prompt here")).toHaveValue("")); if (options.immutable && !options.emptyHistory && !options.restorable) { await waitFor(() => expect(screen.queryByRole("button", { name: "Undo last send" })).not.toBeInTheDocument()); } else { @@ -1588,8 +1592,16 @@ describe("ChatWindow Integration", () => { }; } - it("restores the last prompt and exact middle-stage converter edits with a single click", async () => { - const { user, selected } = await prepareUndoChat(); + it.each([ + { responseDataTypes: ["text"] }, + { responseDataTypes: ["image_path"] }, + { responseDataTypes: ["audio_path"] }, + { responseDataTypes: ["video_path"] }, + { responseDataTypes: ["text", "image_path", "audio_path"] }, + ])("restores the last prompt and exact middle-stage converter edits after a $responseDataTypes response", async ({ + responseDataTypes, + }: { responseDataTypes: string[] }) => { + const { user, selected } = await prepareUndoChat({ responseDataTypes }); expect(screen.getByPlaceholderText("Type prompt here")).toHaveValue(""); await user.click(screen.getByRole("button", { name: "Undo last send" })); await waitFor(() => expect(selected).toHaveBeenCalledWith("undo-branch")); @@ -1640,8 +1652,16 @@ describe("ChatWindow Integration", () => { expect(await screen.findByPlaceholderText("Type prompt here")).toHaveValue("Original prompt"); }); - it("restores matching in-memory converter stages when branching the first prompt", async () => { - const { user, selected } = await prepareUndoChat({ immutable: true, emptyHistory: true }); + it.each([ + { responseDataTypes: ["text"] }, + { responseDataTypes: ["image_path"] }, + { responseDataTypes: ["audio_path"] }, + { responseDataTypes: ["video_path"] }, + { responseDataTypes: ["text", "image_path", "audio_path"] }, + ])("restores matching in-memory converter stages when branching the first prompt after a $responseDataTypes response", async ({ + responseDataTypes, + }: { responseDataTypes: string[] }) => { + const { user, selected } = await prepareUndoChat({ immutable: true, emptyHistory: true, responseDataTypes }); const attack = await mockedAttacksApi.getAttack("undo-attack"); mockedAttacksApi.saveConversation.mockResolvedValue({ attack, messages: { conversation_id: "undo-branch", messages: [], target_response_status: null }, @@ -1660,8 +1680,15 @@ describe("ChatWindow Integration", () => { })); }); - it.each(["pipeline", "working input"])("protects %s edits before first-prompt restore", async (edit: string) => { - const { user, selected } = await prepareUndoChat({ immutable: true, emptyHistory: true }); + it.each([ + { edit: "pipeline", responseDataTypes: ["text"] }, + { edit: "working input", responseDataTypes: ["text"] }, + { edit: "pipeline", responseDataTypes: ["image_path"] }, + { edit: "working input", responseDataTypes: ["audio_path"] }, + ])("protects $edit edits before first-prompt restore after a $responseDataTypes response", async ({ + edit, responseDataTypes, + }: { edit: string; responseDataTypes: string[] }) => { + const { user, selected } = await prepareUndoChat({ immutable: true, emptyHistory: true, responseDataTypes }); if (edit === "pipeline") { await user.click(screen.getByRole("button", { name: "Remove converter second" })); } else { @@ -5648,8 +5675,7 @@ describe("ChatWindow Integration", () => { await waitFor(() => expect(composer).toBeDisabled()); expect(screen.getByPlaceholderText("Type prompt here")).toBe(composer); expect(screen.queryByTestId("single-turn-banner")).not.toBeInTheDocument(); - if (imageResponse) expect(screen.queryByRole("button", { name: "Undo last send" })).not.toBeInTheDocument(); - else await waitFor(() => expect(screen.getByRole("button", { name: "Undo last send" })).toBeEnabled()); + await waitFor(() => expect(screen.getByRole("button", { name: "Undo last send" })).toBeEnabled()); await user.click(screen.getByLabelText("Why the prompt box is disabled")); expect(screen.getByRole("listitem")).toHaveTextContent("This target only supports single-turn conversations."); await user.keyboard("{Escape}"); diff --git a/frontend/src/components/Chat/ChatWindow.tsx b/frontend/src/components/Chat/ChatWindow.tsx index 506beadcb2..24727b822f 100644 --- a/frontend/src/components/Chat/ChatWindow.tsx +++ b/frontend/src/components/Chat/ChatWindow.tsx @@ -993,7 +993,6 @@ export default function ChatWindow({ && typeof targetResponseStatus?.request_turn_number === 'number' && targetResponseStatus.request_turn_number === operation.progress.request_turn_number && lastMessage?.role === 'assistant' && lastMessage.message_pieces.length > 0 - && lastMessage.message_pieces.every((piece) => piece.converted_value_data_type === 'text') ) { setUndoableSend({ attackResultId: operation.attackResultId,