diff --git a/frontend/src/components/Chat/ChatWindow.tsx b/frontend/src/components/Chat/ChatWindow.tsx index a3c70f152d..97828bf4cb 100644 --- a/frontend/src/components/Chat/ChatWindow.tsx +++ b/frontend/src/components/Chat/ChatWindow.tsx @@ -43,11 +43,12 @@ import { attachmentDraft, canBranchFirstPrompt, draftDataTypes, draftMessageToComposer, editorTargetDisabledReason, historyReplayDisabledReason, newDraftPiece, toConversationDraft, } from '@/utils/conversationDraft' -import { useConversationSave } from '@/hooks/useConversationSave' import { useConversationDraft } from '@/hooks/useConversationDraft' import { useConversationMessages } from '@/hooks/useConversationMessages' import { useAttackMutations } from '@/hooks/useAttackMutations' import { useConversationExport } from '@/hooks/useConversationExport' +import { useConversationCopy } from '@/hooks/useConversationCopy' +import { conversationPieceSignature, preparePromptRestore } from '@/utils/promptRestore' import { useChatDraftHandoff } from '@/hooks/useChatDraftHandoff' import { shouldApplyPreservedDraft, @@ -107,7 +108,6 @@ import type { CreateConversationRequest, Message, MessageAttachment, - MessageDisplayPiece, MessageSendConversation, MessageSendRequest, MessageSendStatus, @@ -135,8 +135,6 @@ interface ChatInfoNotice { readonly fading: boolean } -type ConversationCopyDestination = 'same_attack' | 'new_attack' | 'technique' - interface RecoverableSendDraft { history: Array> conversationId: string @@ -191,20 +189,6 @@ interface UndoableSend { history: Array> } -function preparePromptRestore( - saved: ChatConverterSnapshot | undefined, current: ChatConverterSnapshot, -): { snapshot: ChatConverterSnapshot; pipelineChanged: boolean } { - const pipelineChanged = saved - ? JSON.stringify(current.pipelines) !== JSON.stringify(saved.pipelines) - : Object.values(current.pipelines).some((stages: ConverterPipelineStage[]) => stages.length > 0) - return { - pipelineChanged, - snapshot: saved && !pipelineChanged ? saved : { - ...current, workingInputs: {}, stageResults: {}, errors: {}, applied: {}, - }, - } -} - interface RepeatSendGroup { source: PendingSend progress: MessageSendStatus @@ -259,10 +243,6 @@ function userPieceIds(response: ConversationMessagesResponse): Set { .flatMap((message) => message.message_pieces.map((piece) => piece.id))) } -function conversationPieceSignature(response: ConversationMessagesResponse): string { - return JSON.stringify(response.messages.map((message: BackendMessage) => message.message_pieces.map((piece) => piece.id))) -} - function getRecoveryHistoryCutoff(messages: BackendMessage[], failedRequestTurnNumber: number): number { let precedingUserTurnNumber: number | undefined for (const message of messages) { @@ -501,8 +481,6 @@ export default function ChatWindow({ const editorRef = useRef(null) const [editorInitialConverters, setEditorInitialConverters] = useState() const [isLoadingEdit, setIsLoadingEdit] = useState(false) - const copyingRef = useRef(false) - const copySave = useConversationSave(newAttackContext) const [editorNotice, setEditorNoticeState] = useState(null) const setEditorNotice = useCallback((message: string | null): void => { setEditorNoticeState(message === null ? null : { id: generateClientId(), message, fading: false }) @@ -2290,155 +2268,53 @@ export default function ChatWindow({ && undoableSend.history.every((message: Pick) => message.role === 'system' || message.role === 'developer') ? undoableSend : undefined - const copyConversation = async ( - messageIndex: number, destination: ConversationCopyDestination, - ): Promise => { - if (!attackResultId || !viewedConversationId || copyingRef.current || isSending || isLoadingEdit - || isLoadingMessages || awaitingConversationLoad || isTranscriptIncomplete) return - if (destination === 'same_attack' && isMutationLocked) return - if (destination !== 'same_attack' && (!runtime.ready || !defaultsReady || !scorerSelectionReady)) return - if (destination === 'same_attack' && recoverableSend - && messageIndex === recoverableSend.failedRequestMessageIndex) { - await handleRecoverProcessingError() - return - } - const sourceId = viewedConversationId - const draftRevision = inputBoxRef.current?.getDraftRevision() - const navigationRevision = navigationRevisionRef.current - const branchConverterSnapshot = converters.snapshot() - 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 - && !convertersRef.current.isConverting - && JSON.stringify(convertersRef.current.snapshot()) === JSON.stringify(branchConverterSnapshot) - copyingRef.current = true - setIsLoadingEdit(true) - setEditorError(null) - setEditorNotice(null) - try { - const source = await attacksApi.getMessages(attackResultId, sourceId) - if (destination !== 'same_attack' && ( - launchStateRef.current.generation !== runtime.generation - || !launchStateRef.current.ready || !launchStateRef.current.defaultsReady - || !launchStateRef.current.scorerSelectionReady - )) { - throw new Error('Runtime or default labels changed while loading this conversation. Retry after default labels finish loading.') - } - const isFailedPrompt = destination === 'new_attack' && recoverableSend && messageIndex === recoverableSend.failedRequestMessageIndex - const sourceMessages = isFailedPrompt - ? 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 - if (restoresFirstPrompt && activeTarget) { - const selected = source.messages[messageIndex] - if (!selected || !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) - const restoration = preparePromptRestore(saved?.converters, branchConverterSnapshot) - promptSnapshot = restoration.snapshot - const promptMessage = copiedMessages[copiedMessages.length - 1] - if (!promptMessage) { - throw new Error('No prompt message to branch from.') - } - restoredPrompt = saved ? { - text: saved.originalValue, attachments: saved.attachments, textMetadata: saved.textMetadata, converters: promptSnapshot, - } : draftMessageToComposer(restoration.pipelineChanged ? { - ...promptMessage, - pieces: promptMessage.pieces.map((piece: ConversationDraftPiece) => ({ - ...piece, converted_value: undefined, converted_value_data_type: undefined, - applied_converter_ids: undefined, applied_converter_provenance: undefined, - })), - } : promptMessage, promptSnapshot) - copiedMessages = copiedMessages.slice(0, -1) - } - if (isFailedPrompt && activeTarget - && editorTargetDisabledReason(activeTarget, draftDataTypes(copiedMessages), copiedMessages)) { - copiedMessages = [] - } - if (destination === 'same_attack' && targetResolutionStatus !== 'unbound' && activeTarget) { - const reason = editorTargetDisabledReason(activeTarget, draftDataTypes(copiedMessages), copiedMessages) - if (reason) throw new Error(reason) - } - let target = activeTarget - if (destination === 'technique' || (destination === 'same_attack' && targetResolutionStatus === 'unbound')) { - target = null - } else if (destination === 'new_attack' && activeTarget && !isFailedPrompt - && editorTargetDisabledReason(activeTarget, draftDataTypes(copiedMessages), copiedMessages)) { - target = null - } - if (isFailedPrompt && copiedMessages.length === 0 && !executionReadOnly) { - const response = await attacksApi.createAttack({ - target_registry_name: target?.target_registry_name, - objective: objective || undefined, - labels, - ...(defaultChatScorer === null || objective.trim() && defaultChatScorer - ? { chat_scorer: defaultChatScorer } : {}), - }) - if (getViewedConversationId() !== sourceId || inputBoxRef.current?.getDraftRevision() !== draftRevision) return - onConversationCreated(response.attack_result_id, response.conversation_id, objective, target) - restoreRecoverableDraft() - setIsPanelOpen(true) - setPanelRefreshKey((key: number) => key + 1) - return - } - const response = await copySave.save({ - sourceAttackId: attackResultId, - sourceConversationId: sourceId, - initialObjective: objective, - objective, - target, - 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 }) - navigate(`/scanner/launch/technique.execute?${params.toString()}`) - return - } - if (isFailedPrompt && inputBoxRef.current?.getDraftRevision() !== draftRevision) return - if (destination === 'same_attack') { - if (restoredPrompt) { - setMessages(handleConversationLoaded( - response.messages.conversation_id, response.messages, 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() - if (restoredPrompt && promptSnapshot) { - restoreEditorPrompt(restoredPrompt, promptSnapshot) - setUndoableSend(null) - } - setIsPanelOpen(true) - } catch (error) { - if (getViewedConversationId() === sourceId) setEditorError(toApiError(error).detail) - } finally { - copyingRef.current = false - setIsLoadingEdit(false) - } - } + const { copyConversation, branchDisabledReason: branchConversationDisabledReason } = useConversationCopy({ + attackResultId, + viewedConversationId, + messages, + activeTarget, + targetResolutionStatus, + objective, + labels, + defaultChatScorer, + newAttackContext, + runtimeGeneration: runtime.generation, + runtimeReady: runtime.ready, + defaultsReady, + scorerSelectionReady, + executionReadOnly, + isMutationLocked, + isBusy: isSending || isLoadingEdit || isLoadingMessages || awaitingConversationLoad || isTranscriptIncomplete, + recoverableSend, + firstPromptSend, + converters, + promptRestoreDisabledReason, + getViewedConversationId, + getViewedAttackId: () => viewedAttackRef.current, + getNavigationRevision: () => navigationRevisionRef.current, + getDraftRevision: () => inputBoxRef.current?.getDraftRevision(), + getLaunchState: () => launchStateRef.current, + getConverters: () => convertersRef.current, + setBusy: setIsLoadingEdit, + setEditorError, + setEditorNotice, + recoverProcessingError: handleRecoverProcessingError, + restoreRecoverableDraft, + restoreEditorPrompt, + clearUndoableSend: () => { setUndoableSend(null) }, + showConversation: (response: ConversationMessagesResponse) => { + setMessages(handleConversationLoaded( + response.conversation_id, response, backendMessagesToFrontend(response.messages), + )) + markConversationLoaded(response.conversation_id) + }, + refreshPanel: () => { setPanelRefreshKey((key: number) => key + 1) }, + openPanel: () => { setIsPanelOpen(true) }, + navigate: (path: string) => { void navigate(path) }, + onSelectConversation, + onConversationCreated, + onAttackChange, + }) const sameAttackDisabledReason = !attackResultId ? 'No saved attack exists yet.' : executionReadOnly ? 'Automated results are read-only. Choose Continue manually.' @@ -2469,20 +2345,6 @@ export default function ChatWindow({ const undoHistory = undoableSend?.attackResultId === attackResultId && 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) { - return historyReplayDisabledReason( - activeTarget, activeTarget.capabilities?.supports_multi_turn ? recoverableSend.history : [], - ) - } - return historyReplayDisabledReason(activeTarget, messages.slice(0, index + 1)) - } const undoDisabledReason = !undoableSend ? 'Undo is available after a single successful send from this tab, with its converters still in memory.' diff --git a/frontend/src/hooks/useConversationCopy.test.ts b/frontend/src/hooks/useConversationCopy.test.ts new file mode 100644 index 0000000000..aab0057e3c --- /dev/null +++ b/frontend/src/hooks/useConversationCopy.test.ts @@ -0,0 +1,232 @@ +import { act, renderHook } from '@testing-library/react' + +import { attacksApi } from '@/services/api' +import { deferred } from '@/test-utils/deferred' +import type { + AddMessageResponse, AttackSummary, ChatConverterController, ConversationMessagesResponse, Message, +} from '@/types' + +import { useConversationCopy } from './useConversationCopy' +import type { UseConversationCopyOptions } from './useConversationCopy' + +jest.mock('@/services/api', () => ({ + attacksApi: { getMessages: jest.fn(), saveConversation: jest.fn(), createAttack: jest.fn() }, +})) + +const source: ConversationMessagesResponse = { + conversation_id: 'main', + messages: [0, 1, 2].map((turn: number) => ({ + turn_number: turn, role: turn % 2 ? 'assistant' : 'user', message_pieces: [], created_at: '', + })), + target_response_status: null, +} +const attack: AttackSummary = { + attack_result_id: 'copied-attack', conversation_id: 'copy', attack_type: 'ManualAttack', objective: 'Objective', + converters: [], message_count: 2, related_conversation_ids: [], labels: {}, created_at: '', updated_at: '', +} +const saved: AddMessageResponse = { + attack, messages: { conversation_id: 'copy', messages: [], target_response_status: null }, +} +const converters = { + isConverting: false, + snapshot: () => ({ generation: 'g1', pipelines: {}, workingInputs: {}, stageResults: {}, errors: {}, applied: {} }), +} as unknown as ChatConverterController +const messages = [0, 1, 2].map((turn: number): Message => ({ role: turn % 2 ? 'assistant' : 'user', content: `${turn}`, timestamp: '' })) + +function renderCopy(overrides: Partial = {}) { + const options: UseConversationCopyOptions = { + attackResultId: 'attack', + viewedConversationId: 'main', + messages, + activeTarget: null, + targetResolutionStatus: 'unbound', + objective: 'Objective', + labels: { operator: 'me' }, + runtimeGeneration: 'g1', + runtimeReady: true, + defaultsReady: true, + scorerSelectionReady: true, + executionReadOnly: false, + isMutationLocked: false, + isBusy: false, + converters, + promptRestoreDisabledReason: () => undefined, + getViewedConversationId: () => 'main', + getViewedAttackId: () => 'attack', + getNavigationRevision: () => 0, + getDraftRevision: () => 1, + getLaunchState: () => ({ generation: 'g1', ready: true, defaultsReady: true, scorerSelectionReady: true }), + getConverters: () => converters, + setBusy: jest.fn(), + setEditorError: jest.fn(), + setEditorNotice: jest.fn(), + recoverProcessingError: jest.fn().mockResolvedValue(undefined), + restoreRecoverableDraft: jest.fn(), + restoreEditorPrompt: jest.fn(), + clearUndoableSend: jest.fn(), + showConversation: jest.fn(), + refreshPanel: jest.fn(), + openPanel: jest.fn(), + navigate: jest.fn(), + onSelectConversation: jest.fn(), + onConversationCreated: jest.fn(), + onAttackChange: jest.fn(), + ...overrides, + } + return { options, ...renderHook((props: UseConversationCopyOptions) => useConversationCopy(props), { initialProps: options }) } +} + +describe('useConversationCopy', () => { + const getMessages = jest.mocked(attacksApi.getMessages) + const saveConversation = jest.mocked(attacksApi.saveConversation) + + beforeEach(() => { + jest.resetAllMocks() + getMessages.mockResolvedValue(source) + saveConversation.mockResolvedValue(saved) + }) + + it('should copy to a new conversation in the same attack and select it', async () => { + const { result, options } = renderCopy() + + await act(async () => { await result.current.copyConversation(1, 'same_attack') }) + + expect(getMessages).toHaveBeenCalledWith('attack', 'main') + expect(saveConversation).toHaveBeenCalledWith(expect.objectContaining({ + destination: 'same_attack', attack_result_id: 'attack', source_conversation_id: 'main', + })) + expect(saveConversation.mock.calls[0][0].messages).toHaveLength(2) + expect(options.onSelectConversation).toHaveBeenCalledWith('copy') + expect(options.onConversationCreated).not.toHaveBeenCalled() + expect(options.onAttackChange).toHaveBeenCalledWith(attack) + expect(options.refreshPanel).toHaveBeenCalled() + expect(options.openPanel).toHaveBeenCalled() + expect(options.setBusy).toHaveBeenLastCalledWith(false) + }) + + it('should branch into a new attack and report the created conversation', async () => { + const { result, options } = renderCopy() + + await act(async () => { await result.current.copyConversation(2, 'new_attack') }) + + expect(saveConversation).toHaveBeenCalledWith(expect.objectContaining({ destination: 'new_attack', objective: 'Objective' })) + expect(options.onConversationCreated).toHaveBeenCalledWith('copied-attack', 'copy', 'Objective', null) + expect(options.onSelectConversation).not.toHaveBeenCalled() + }) + + it('should send technique copies to the technique runner without a target', async () => { + const { result, options } = renderCopy() + + await act(async () => { await result.current.copyConversation(1, 'technique') }) + + expect(saveConversation.mock.calls[0][0].target_registry_name).toBeUndefined() + expect(options.navigate).toHaveBeenCalledWith('/scanner/launch/technique.execute?prepended_conversation_id=copy') + expect(options.onConversationCreated).not.toHaveBeenCalled() + }) + + it('should ignore a second copy while one is in flight', async () => { + const load = deferred() + getMessages.mockReturnValueOnce(load.promise) + const { result } = renderCopy() + + let first!: Promise + act(() => { first = result.current.copyConversation(1, 'same_attack') }) + await act(async () => { await result.current.copyConversation(1, 'new_attack') }) + await act(async () => { load.resolve(source); await first }) + + expect(getMessages).toHaveBeenCalledTimes(1) + expect(saveConversation).toHaveBeenCalledTimes(1) + }) + + it('should not navigate when the user left the source conversation during the copy', async () => { + const load = deferred() + getMessages.mockReturnValueOnce(load.promise) + let viewed = 'main' + const { result, options } = renderCopy({ getViewedConversationId: () => viewed }) + + let pending!: Promise + act(() => { pending = result.current.copyConversation(1, 'same_attack') }) + viewed = 'other' + await act(async () => { load.resolve(source); await pending }) + + expect(options.onSelectConversation).not.toHaveBeenCalled() + expect(options.openPanel).not.toHaveBeenCalled() + }) + + it('should not copy into the same attack while mutations are locked', async () => { + const { result } = renderCopy({ isMutationLocked: true }) + + await act(async () => { await result.current.copyConversation(1, 'same_attack') }) + + expect(getMessages).not.toHaveBeenCalled() + }) + + it('should not start a new attack before default labels are ready', async () => { + const { result } = renderCopy({ defaultsReady: false }) + + await act(async () => { await result.current.copyConversation(1, 'new_attack') }) + + expect(getMessages).not.toHaveBeenCalled() + }) + + it('should not start a new attack before the remembered chat scorer is ready', async () => { + const { result } = renderCopy({ scorerSelectionReady: false }) + + await act(async () => { await result.current.copyConversation(1, 'new_attack') }) + + expect(getMessages).not.toHaveBeenCalled() + }) + + it('should recover a failed prompt in place instead of copying it into the same attack', async () => { + const { result, options } = renderCopy({ + recoverableSend: { history: [], historyCutoffIndex: 0, failedRequestMessageIndex: 2 }, + }) + + await act(async () => { await result.current.copyConversation(2, 'same_attack') }) + + expect(options.recoverProcessingError).toHaveBeenCalled() + expect(getMessages).not.toHaveBeenCalled() + }) + + it('should report a failed copy only while the source conversation is viewed', async () => { + getMessages.mockRejectedValue(new Error('Load failed')) + const { result, options } = renderCopy() + + await act(async () => { await result.current.copyConversation(1, 'same_attack') }) + + expect(options.setEditorError).toHaveBeenLastCalledWith('Load failed') + expect(options.setBusy).toHaveBeenLastCalledWith(false) + }) + + it.each([ + { label: 'with an objective', objective: 'Keep this goal', sent: 'Keep this goal' }, + { label: 'without an objective', objective: '', sent: undefined }, + ])('should create a failed-prompt branch $label using objective and not name', async ({ objective, sent }) => { + const createAttack = jest.mocked(attacksApi.createAttack) + createAttack.mockResolvedValue({ attack_result_id: 'retry', conversation_id: 'retry-conv', created_at: '' }) + const { result, options } = renderCopy({ + objective, + recoverableSend: { history: [], historyCutoffIndex: -1, failedRequestMessageIndex: 2 }, + }) + + await act(async () => { await result.current.copyConversation(2, 'new_attack') }) + + expect(createAttack).toHaveBeenCalledWith(expect.objectContaining({ + target_registry_name: undefined, + objective: sent, + labels: { operator: 'me' }, + })) + expect(createAttack.mock.calls[0][0]).not.toHaveProperty('name') + expect(saveConversation).not.toHaveBeenCalled() + expect(options.onConversationCreated).toHaveBeenCalledWith('retry', 'retry-conv', objective, null) + expect(options.restoreRecoverableDraft).toHaveBeenCalled() + }) + + it('should explain why branching needs a configured target', () => { + const { result, rerender, options } = renderCopy({ targetResolutionStatus: 'idle' }) + + expect(result.current.branchDisabledReason(1)).toBe('Configure a target before branching this conversation.') + rerender({ ...options, targetResolutionStatus: 'unbound' }) + expect(result.current.branchDisabledReason(1)).toBeUndefined() + }) +}) diff --git a/frontend/src/hooks/useConversationCopy.ts b/frontend/src/hooks/useConversationCopy.ts new file mode 100644 index 0000000000..d9056aba39 --- /dev/null +++ b/frontend/src/hooks/useConversationCopy.ts @@ -0,0 +1,284 @@ +import { useRef } from 'react' + +import { useConversationSave } from '@/hooks/useConversationSave' +import { attacksApi } from '@/services/api' +import { toApiError } from '@/services/errors' +import type { + AttackSummary, + AttackTargetResolutionStatus, + BackendMessage, + BackendMessagePiece, + ChatConverterController, + ChatConverterSnapshot, + ConversationDraftPiece, + ConversationMessagesResponse, + Message, + MessageAttachment, + MessageDisplayPiece, + NewAttackContext, + RegisteredComponentReference, + TargetInstance, +} from '@/types' +import { + canBranchFirstPrompt, draftDataTypes, draftMessageToComposer, editorTargetDisabledReason, + historyReplayDisabledReason, toConversationDraft, +} from '@/utils/conversationDraft' +import { conversationPieceSignature, preparePromptRestore } from '@/utils/promptRestore' + +export type ConversationCopyDestination = 'same_attack' | 'new_attack' | 'technique' + +export type RestoredPrompt = ReturnType + +/** The failed send that can be recovered or branched from. */ +export interface CopyRecoverableSend { + history: Array> + historyCutoffIndex: number + failedRequestMessageIndex: number +} + +/** The first prompt sent in the viewed conversation, kept so a branch can restore it into the composer. */ +export interface CopyFirstPromptSend { + historySignature: string + cutoffIndex: number + originalValue: string + textMetadata?: Record + attachments: MessageAttachment[] + converters: ChatConverterSnapshot +} + +export interface CopyLaunchState { + generation: string + ready: boolean + defaultsReady: boolean + scorerSelectionReady: boolean +} + +export interface UseConversationCopyOptions { + attackResultId: string | null + viewedConversationId: string | null + messages: Message[] + activeTarget: TargetInstance | null + targetResolutionStatus: AttackTargetResolutionStatus | undefined + objective: string + labels?: Record + defaultChatScorer?: RegisteredComponentReference | null + newAttackContext?: NewAttackContext + runtimeGeneration: string + runtimeReady: boolean + defaultsReady: boolean + scorerSelectionReady: boolean + executionReadOnly: boolean + isMutationLocked: boolean + /** True while a send or another conversation edit is in flight. */ + isBusy: boolean + recoverableSend?: CopyRecoverableSend + firstPromptSend?: CopyFirstPromptSend + converters: ChatConverterController + promptRestoreDisabledReason: (snapshot?: ChatConverterSnapshot) => string | undefined + /** Live reads used to detect that the user moved on while the copy was loading. */ + getViewedConversationId: () => string | null + getViewedAttackId: () => string | null + getNavigationRevision: () => number + getDraftRevision: () => number | undefined + getLaunchState: () => CopyLaunchState + getConverters: () => ChatConverterController + setBusy: (busy: boolean) => void + setEditorError: (error: string | null) => void + setEditorNotice: (notice: string | null) => void + recoverProcessingError: () => Promise + restoreRecoverableDraft: () => void + restoreEditorPrompt: (prompt: RestoredPrompt | undefined, snapshot: ChatConverterSnapshot) => void + clearUndoableSend: () => void + showConversation: (response: ConversationMessagesResponse) => void + refreshPanel: () => void + openPanel: () => void + navigate: (path: string) => void + onSelectConversation: (conversationId: string) => void + onConversationCreated: ( + attackResultId: string, conversationId: string, objective?: string, target?: TargetInstance | null, + ) => void + onAttackChange?: (attack: AttackSummary) => void +} + +export interface ConversationCopy { + /** Copies the conversation up to a message into the same attack, a new attack, or a technique run. */ + copyConversation: (messageIndex: number, destination: ConversationCopyDestination) => Promise + /** Why branching from a message is unavailable, if it is. */ + branchDisabledReason: (messageIndex: number) => string | undefined +} + +/** Owns copy-to-conversation and branch workflows, including the lock that prevents overlapping copies. */ +export function useConversationCopy(options: UseConversationCopyOptions): ConversationCopy { + const copyingRef = useRef(false) + const copySave = useConversationSave(options.newAttackContext) + + const copyConversation = async ( + messageIndex: number, destination: ConversationCopyDestination, + ): Promise => { + const { + attackResultId, viewedConversationId, messages, activeTarget, targetResolutionStatus, objective, labels, + defaultChatScorer, runtimeGeneration, executionReadOnly, recoverableSend, firstPromptSend, converters, + getViewedConversationId, getViewedAttackId, getNavigationRevision, getDraftRevision, getLaunchState, getConverters, + } = options + if (!attackResultId || !viewedConversationId || copyingRef.current || options.isBusy) return + if (destination === 'same_attack' && options.isMutationLocked) return + if (destination !== 'same_attack' && (!options.runtimeReady || !options.defaultsReady || !options.scorerSelectionReady)) return + if (destination === 'same_attack' && recoverableSend + && messageIndex === recoverableSend.failedRequestMessageIndex) { + await options.recoverProcessingError() + return + } + const sourceId = viewedConversationId + const draftRevision = getDraftRevision() + const navigationRevision = getNavigationRevision() + const branchConverterSnapshot = converters.snapshot() + const selectedPieceIds = messages[messageIndex]?.displayPieces?.map((piece: MessageDisplayPiece) => piece.pieceId) ?? [] + const restoresFirstPrompt = destination === 'same_attack' && activeTarget + && canBranchFirstPrompt(activeTarget, messages, messageIndex) + const draftUnchanged = (): boolean => { + const launch = getLaunchState() + return getViewedAttackId() === attackResultId + && getViewedConversationId() === sourceId + && getNavigationRevision() === navigationRevision + && getDraftRevision() === draftRevision + && launch.generation === runtimeGeneration && launch.ready + && !getConverters().isConverting + && JSON.stringify(getConverters().snapshot()) === JSON.stringify(branchConverterSnapshot) + } + copyingRef.current = true + options.setBusy(true) + options.setEditorError(null) + options.setEditorNotice(null) + try { + const source = await attacksApi.getMessages(attackResultId, sourceId) + const launch = getLaunchState() + if (destination !== 'same_attack' && (launch.generation !== runtimeGeneration || !launch.ready || !launch.defaultsReady + || !launch.scorerSelectionReady)) { + throw new Error('Runtime or default labels changed while loading this conversation. Retry after default labels finish loading.') + } + const isFailedPrompt = destination === 'new_attack' && recoverableSend && messageIndex === recoverableSend.failedRequestMessageIndex + const sourceMessages = isFailedPrompt + ? source.messages.filter((message: BackendMessage) => message.turn_number <= recoverableSend.historyCutoffIndex) + : source.messages.slice(0, messageIndex + 1) + let copiedMessages = toConversationDraft(sourceMessages) + let restoredPrompt: RestoredPrompt | undefined + let promptSnapshot: ChatConverterSnapshot | undefined + if (restoresFirstPrompt && activeTarget) { + const selected = source.messages[messageIndex] + if (!selected || !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 = options.promptRestoreDisabledReason(saved?.converters) + if (restoreError) throw new Error(restoreError) + const restoration = preparePromptRestore(saved?.converters, branchConverterSnapshot) + promptSnapshot = restoration.snapshot + const promptMessage = copiedMessages[copiedMessages.length - 1] + if (!promptMessage) { + throw new Error('No prompt message to branch from.') + } + restoredPrompt = saved ? { + text: saved.originalValue, attachments: saved.attachments, textMetadata: saved.textMetadata, converters: promptSnapshot, + } : draftMessageToComposer(restoration.pipelineChanged ? { + ...promptMessage, + pieces: promptMessage.pieces.map((piece: ConversationDraftPiece) => ({ + ...piece, converted_value: undefined, converted_value_data_type: undefined, + applied_converter_ids: undefined, applied_converter_provenance: undefined, + })), + } : promptMessage, promptSnapshot) + copiedMessages = copiedMessages.slice(0, -1) + } + if (isFailedPrompt && activeTarget + && editorTargetDisabledReason(activeTarget, draftDataTypes(copiedMessages), copiedMessages)) { + copiedMessages = [] + } + if (destination === 'same_attack' && targetResolutionStatus !== 'unbound' && activeTarget) { + const reason = editorTargetDisabledReason(activeTarget, draftDataTypes(copiedMessages), copiedMessages) + if (reason) throw new Error(reason) + } + let target = activeTarget + if (destination === 'technique' || (destination === 'same_attack' && targetResolutionStatus === 'unbound')) { + target = null + } else if (destination === 'new_attack' && activeTarget && !isFailedPrompt + && editorTargetDisabledReason(activeTarget, draftDataTypes(copiedMessages), copiedMessages)) { + target = null + } + if (isFailedPrompt && copiedMessages.length === 0 && !executionReadOnly) { + const response = await attacksApi.createAttack({ + target_registry_name: target?.target_registry_name, + objective: objective || undefined, + labels, + ...(defaultChatScorer === null || objective.trim() && defaultChatScorer + ? { chat_scorer: defaultChatScorer } : {}), + }) + if (getViewedConversationId() !== sourceId || getDraftRevision() !== draftRevision) return + options.onConversationCreated(response.attack_result_id, response.conversation_id, objective, target) + options.restoreRecoverableDraft() + options.openPanel() + options.refreshPanel() + return + } + const response = await copySave.save({ + sourceAttackId: attackResultId, + sourceConversationId: sourceId, + initialObjective: objective, + objective, + target, + labels, + messages: copiedMessages, + }, destination === 'same_attack' ? 'same_attack' : 'new_attack') + options.refreshPanel() + if (restoredPrompt && !draftUnchanged()) { + options.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 }) + options.navigate(`/scanner/launch/technique.execute?${params.toString()}`) + return + } + if (isFailedPrompt && getDraftRevision() !== draftRevision) return + if (destination === 'same_attack') { + if (restoredPrompt) { + options.showConversation(response.messages) + } + options.onSelectConversation(response.messages.conversation_id) + } + else options.onConversationCreated(response.attack.attack_result_id, response.messages.conversation_id, response.attack.objective, target) + options.onAttackChange?.(response.attack) + if (isFailedPrompt) options.restoreRecoverableDraft() + if (restoredPrompt && promptSnapshot) { + options.restoreEditorPrompt(restoredPrompt, promptSnapshot) + options.clearUndoableSend() + } + options.openPanel() + } catch (error) { + if (getViewedConversationId() === sourceId) options.setEditorError(toApiError(error).detail) + } finally { + copyingRef.current = false + options.setBusy(false) + } + } + + const branchDisabledReason = (index: number): string | undefined => { + const { activeTarget, messages, recoverableSend, firstPromptSend, targetResolutionStatus } = options + if (activeTarget && canBranchFirstPrompt(activeTarget, messages, index) + && recoverableSend?.failedRequestMessageIndex !== index) { + return options.promptRestoreDisabledReason(firstPromptSend?.converters) + } + if (targetResolutionStatus === 'unbound') return undefined + if (!activeTarget) return 'Configure a target before branching this conversation.' + if (recoverableSend?.failedRequestMessageIndex === index) { + return historyReplayDisabledReason( + activeTarget, activeTarget.capabilities?.supports_multi_turn ? recoverableSend.history : [], + ) + } + return historyReplayDisabledReason(activeTarget, messages.slice(0, index + 1)) + } + + return { copyConversation, branchDisabledReason } +} diff --git a/frontend/src/utils/promptRestore.ts b/frontend/src/utils/promptRestore.ts new file mode 100644 index 0000000000..9c016b40ff --- /dev/null +++ b/frontend/src/utils/promptRestore.ts @@ -0,0 +1,24 @@ +import type { BackendMessage, ChatConverterSnapshot, ConversationMessagesResponse, ConverterPipelineStage } from '@/types' + +/** + * Picks the converter snapshot to restore with a sent prompt. A saved snapshot is reused only when the + * current pipeline still matches it; otherwise the current pipeline is kept and its results are cleared. + */ +export function preparePromptRestore( + saved: ChatConverterSnapshot | undefined, current: ChatConverterSnapshot, +): { snapshot: ChatConverterSnapshot; pipelineChanged: boolean } { + const pipelineChanged = saved + ? JSON.stringify(current.pipelines) !== JSON.stringify(saved.pipelines) + : Object.values(current.pipelines).some((stages: ConverterPipelineStage[]) => stages.length > 0) + return { + pipelineChanged, + snapshot: saved && !pipelineChanged ? saved : { + ...current, workingInputs: {}, stageResults: {}, errors: {}, applied: {}, + }, + } +} + +/** Identifies a conversation's persisted pieces so a stale restore can be detected. */ +export function conversationPieceSignature(response: ConversationMessagesResponse): string { + return JSON.stringify(response.messages.map((message: BackendMessage) => message.message_pieces.map((piece) => piece.id))) +}