diff --git a/packages/ui/src/components/chat/ChatInput.tsx b/packages/ui/src/components/chat/ChatInput.tsx index 1086a296..f571479c 100644 --- a/packages/ui/src/components/chat/ChatInput.tsx +++ b/packages/ui/src/components/chat/ChatInput.tsx @@ -10,7 +10,7 @@ import { useSelectionStore } from '@/sync/selection-store'; import { useInputStore } from '@/sync/input-store'; import type { AttachedFile } from '@/stores/types/sessionTypes'; import * as sessionActions from '@/sync/session-actions'; -import { useUserMessageHistory } from '@/sync/sync-context'; +import { useDirectorySync, useUserMessageHistory } from '@/sync/sync-context'; import { useInlineCommentDraftStore, type InlineCommentDraft } from '@/stores/useInlineCommentDraftStore'; import { appendInlineComments } from '@/lib/messages/inlineComments'; import { renderMagicPrompt } from '@/lib/magicPrompts'; @@ -30,6 +30,7 @@ import { MobileModelButton } from './MobileModelButton'; import { MobileSessionStatusBar } from './MobileSessionStatusBar'; import { useCurrentSessionActivity } from '@/hooks/useSessionActivity'; import { toast } from '@/components/ui'; +import { Button } from '@/components/ui/button'; // useMessageStore removed — messages now come from sync system import { isTauriShell, isVSCodeRuntime } from '@/lib/desktop'; import { isIMECompositionEvent } from '@/lib/ime'; @@ -63,9 +64,12 @@ import { fetchResponseStyleInstruction } from '@/lib/responseStyle'; import { wrapSystemReminder } from '@/lib/systemReminder'; import { getSyncMessages } from '@/sync/sync-refs'; import { eventMatchesShortcut, getEffectiveShortcutCombo, normalizeCombo } from '@/lib/shortcuts'; +import { isSyntheticPart } from '@/lib/messages/synthetic'; +import type { Message, Part } from '@opencode-ai/sdk/v2/client'; const MAX_VISIBLE_TEXTAREA_LINES = 8; const EMPTY_QUEUE: QueuedMessage[] = []; +const EMPTY_MESSAGES: Message[] = []; const FILE_MENTION_TOKEN = /^@[^\s]+$/; const CHAT_DRAFT_PERSIST_DEBOUNCE_MS = 500; const VS_CODE_DROP_DATA_TYPES = [ @@ -81,6 +85,26 @@ const hasUserMessages = (sessionId: string, directory?: string) => { return getSyncMessages(sessionId, directory).some((message) => message.role === 'user'); }; +const getRevertedPreview = (parts: Part[], fallback: string): string => { + const text = parts + .filter((part) => part.type === 'text' && !isSyntheticPart(part)) + .map((part) => { + const record = part as Record; + return typeof record.text === 'string' + ? record.text + : typeof record.content === 'string' + ? record.content + : ''; + }) + .join('\n') + .replace(/\s+/g, ' ') + .trim(); + + if (text) return text; + const filePart = parts.find((part) => part.type === 'file') as (Part & { filename?: string }) | undefined; + return filePart?.filename ? `[${filePart.filename}]` : fallback; +}; + const FILE_URI_PREFIX = 'file://'; const encodeFilePath = (filepath: string): string => { @@ -228,6 +252,144 @@ const MemoMobileAgentButton = React.memo(MobileAgentButton); const MemoMobileModelButton = React.memo(MobileModelButton); const MemoStatusRow = React.memo(StatusRow); +type RevertedMessageDockProps = { + sessionId: string | null; + directory?: string; +}; + +const RevertedMessageDock: React.FC = React.memo(({ sessionId, directory }) => { + const { t } = useI18n(); + const revertToMessage = useSessionUIStore((s) => s.revertToMessage); + const forkFromMessage = useSessionUIStore((s) => s.forkFromMessage); + const handleSlashRedo = useSessionUIStore((s) => s.handleSlashRedo); + const [restoringId, setRestoringId] = React.useState(null); + const [forkingId, setForkingId] = React.useState(null); + const [collapsed, setCollapsed] = React.useState(true); + const revertMessageID = useDirectorySync( + React.useCallback((state) => { + if (!sessionId) return undefined; + const session = state.session.find((item) => item.id === sessionId); + return (session as { revert?: { messageID?: string } } | undefined)?.revert?.messageID; + }, [sessionId]), + directory, + ); + const sessionMessages = useDirectorySync( + React.useCallback((state) => (sessionId ? state.message[sessionId] ?? EMPTY_MESSAGES : EMPTY_MESSAGES), [sessionId]), + directory, + ); + const partsByMessage = useDirectorySync(React.useCallback((state) => state.part, []), directory); + + const userMessages = React.useMemo( + () => sessionMessages.filter((message): message is Message & { role: 'user' } => message.role === 'user'), + [sessionMessages], + ); + const noTextContent = t('chat.revertPopover.noTextContent'); + const items = React.useMemo(() => { + if (!revertMessageID) return []; + return userMessages + .filter((message) => message.id >= revertMessageID) + .map((message) => ({ + id: message.id, + text: getRevertedPreview(partsByMessage[message.id] ?? [], noTextContent), + })); + }, [noTextContent, partsByMessage, revertMessageID, userMessages]); + const firstRevertedMessageId = items[0]?.id; + + React.useEffect(() => { + setCollapsed(true); + }, [revertMessageID, firstRevertedMessageId]); + + const handleRestore = React.useCallback(async (messageId: string) => { + if (!sessionId || restoringId) return; + setRestoringId(messageId); + try { + const nextMessage = userMessages.find((message) => message.id > messageId); + if (nextMessage) { + await revertToMessage(sessionId, nextMessage.id, { skipRedoPush: true }); + } else { + await handleSlashRedo(sessionId, { fullUnrevert: true }); + } + } finally { + setRestoringId(null); + } + }, [handleSlashRedo, revertToMessage, restoringId, sessionId, userMessages]); + + const handleFork = React.useCallback(async (messageId: string) => { + if (!sessionId || forkingId) return; + setForkingId(messageId); + try { + await forkFromMessage(sessionId, messageId); + } finally { + setForkingId(null); + } + }, [forkFromMessage, forkingId, sessionId]); + + if (!sessionId || items.length === 0) return null; + + return ( +
+
+ + {!collapsed && ( +
+ {items.map((item) => ( +
+ + {item.text} + + + +
+ ))} +
+ )} +
+
+ ); +}); + +RevertedMessageDock.displayName = 'RevertedMessageDock'; + type ComposerAttachmentControlsProps = { isMobile: boolean; isVSCode: boolean; @@ -754,6 +916,9 @@ const ChatInputComponent: React.FC = ({ onOpenSettings, scrollTo ).current; const currentSessionId = useSessionUIStore((s) => s.currentSessionId); const currentDirectory = useDirectoryStore((s) => s.currentDirectory); + const currentSessionDirectoryForSync = useSessionUIStore( + React.useCallback((s) => currentSessionId ? s.getDirectoryForSession(currentSessionId) : null, [currentSessionId]), + ); const newSessionDraft = useSessionUIStore((s) => s.newSessionDraft); const newSessionDraftOpen = Boolean(newSessionDraft?.open); const setNewSessionDraftTarget = useSessionUIStore((s) => s.setNewSessionDraftTarget); @@ -3523,6 +3688,10 @@ const ChatInputComponent: React.FC = ({ onOpenSettings, scrollTo )} + = ({ const getAgentModelForSession = useSelectionStore((s) => s.getAgentModelForSession); const getSessionModelSelection = useSelectionStore((s) => s.getSessionModelSelection); - const revertToMessage = sessionActions.revertToMessage; - const forkFromMessage = sessionActions.forkFromMessage; + const revertToMessage = useSessionUIStore((s) => s.revertToMessage); + const forkFromMessage = useSessionUIStore((s) => s.forkFromMessage); streamPerfCount('ui.chat_message.render'); if (isInActiveTurn) { diff --git a/packages/ui/src/components/chat/StatusRow.tsx b/packages/ui/src/components/chat/StatusRow.tsx index 4dba3e65..ddde86f2 100644 --- a/packages/ui/src/components/chat/StatusRow.tsx +++ b/packages/ui/src/components/chat/StatusRow.tsx @@ -288,8 +288,8 @@ export const StatusRow: React.FC = ({ return (
- {/* Left: Abort status or Working placeholder or leftAccessory */} -
+ {/* Left: Abort status | Working placeholder | leftAccessory */} +
{showAssistantStatus && showAbortStatus ? (
diff --git a/packages/ui/src/lib/i18n/messages/en.ts b/packages/ui/src/lib/i18n/messages/en.ts index 232cd279..06659e72 100644 --- a/packages/ui/src/lib/i18n/messages/en.ts +++ b/packages/ui/src/lib/i18n/messages/en.ts @@ -1402,6 +1402,19 @@ export const dict = { 'chat.statusRow.tasksTitle': 'Tasks', 'chat.statusRow.summary.activeLeft': '{active} active · {left} left', 'chat.statusRow.aborted': 'Aborted', + 'chat.revertIndicator.redo': 'Redo', + 'chat.revertIndicator.redoAria': 'Redo — restore reverted messages', + 'chat.revertPopover.title': 'Reverted', + 'chat.revertPopover.noTextContent': 'No text content', + 'chat.revertPopover.forkFromHere': 'Fork from here', + 'chat.revertPopover.forkFromMessage': 'Fork from this message', + 'chat.revertPopover.restoreAll': 'Restore all', + 'chat.revertPopover.restore': 'Restore', + 'chat.revertPopover.revert': 'Revert', + 'chat.revertPopover.fork': 'Fork', + 'chat.revert.toast.undo': 'Reverted to {preview}', + 'chat.revert.toast.redo': 'Redone', + 'chat.revert.toast.restored': 'Restored all messages', 'chat.errorBoundary.title': 'Chat Error', 'chat.errorBoundary.description': 'The chat interface encountered an error. This might be due to a temporary network issue or corrupted message data.', 'chat.errorBoundary.sessionLabel': 'Session', diff --git a/packages/ui/src/lib/i18n/messages/es.ts b/packages/ui/src/lib/i18n/messages/es.ts index bb28d4de..f9d0a9e4 100644 --- a/packages/ui/src/lib/i18n/messages/es.ts +++ b/packages/ui/src/lib/i18n/messages/es.ts @@ -1368,6 +1368,19 @@ export const dict: Record = { "chat.statusRow.tasksTitle": "Tareas", "chat.statusRow.summary.activeLeft": "{active} activas · {left} restantes", "chat.statusRow.aborted": "Interrumpido", + "chat.revertIndicator.redo": "Rehacer", + "chat.revertIndicator.redoAria": "Rehacer — restaurar mensajes revertidos", + "chat.revertPopover.title": "Revertidos", + "chat.revertPopover.noTextContent": "Sin contenido de texto", + "chat.revertPopover.forkFromHere": "Bifurcar desde aquí", + "chat.revertPopover.forkFromMessage": "Bifurcar desde este mensaje", + "chat.revertPopover.restoreAll": "Restaurar todo", + "chat.revertPopover.restore": "Restaurar", + "chat.revertPopover.revert": "Revertir", + "chat.revertPopover.fork": "Bifurcar", + "chat.revert.toast.undo": "Revertido a {preview}", + "chat.revert.toast.redo": "Rehecho", + "chat.revert.toast.restored": "Todos los mensajes restaurados", "chat.errorBoundary.title": "Error en la conversación", "chat.errorBoundary.description": "La interfaz de la conversación encontró un error. Esto podría deberse a un problema de red temporal o a datos de mensaje corruptos.", "chat.errorBoundary.sessionLabel": "Sesión", diff --git a/packages/ui/src/lib/i18n/messages/ko.ts b/packages/ui/src/lib/i18n/messages/ko.ts index 174530a5..3598cf33 100644 --- a/packages/ui/src/lib/i18n/messages/ko.ts +++ b/packages/ui/src/lib/i18n/messages/ko.ts @@ -1404,6 +1404,19 @@ export const dict: Record = { 'chat.statusRow.tasksTitle': '작업', 'chat.statusRow.summary.activeLeft': '{active}개 활성 · {left}개 남음', 'chat.statusRow.aborted': '중단됨', + 'chat.revertIndicator.redo': '다시 실행', + 'chat.revertIndicator.redoAria': '다시 실행 — 되돌린 메시지 복원', + 'chat.revertPopover.title': '되돌림', + 'chat.revertPopover.noTextContent': '텍스트 없음', + 'chat.revertPopover.forkFromHere': '여기서 분기', + 'chat.revertPopover.forkFromMessage': '이 메시지에서 분기', + 'chat.revertPopover.restoreAll': '모두 복원', + 'chat.revertPopover.restore': '복원', + 'chat.revertPopover.revert': '되돌리기', + 'chat.revertPopover.fork': '분기', + 'chat.revert.toast.undo': '{preview}(으)로 되돌림', + 'chat.revert.toast.redo': '다시 실행', + 'chat.revert.toast.restored': '모든 메시지 복원됨', 'chat.errorBoundary.title': '채팅 오류', 'chat.errorBoundary.description': '채팅 인터페이스에서 오류가 발생했습니다. 일시적인 네트워크 이슈 또는 손상된 메시지 데이터 때문일 수 있습니다.', 'chat.errorBoundary.sessionLabel': '세션', diff --git a/packages/ui/src/lib/i18n/messages/pl.ts b/packages/ui/src/lib/i18n/messages/pl.ts index e45f60a9..e9e9ef82 100644 --- a/packages/ui/src/lib/i18n/messages/pl.ts +++ b/packages/ui/src/lib/i18n/messages/pl.ts @@ -459,6 +459,19 @@ export const dict: Record = { 'chat.statusRow.tasksTitle': 'Zadania', 'chat.statusRow.summary.activeLeft': '{active} aktywne · {left} pozostało', 'chat.statusRow.aborted': 'Przerwane', + 'chat.revertIndicator.redo': 'Ponów', + 'chat.revertIndicator.redoAria': 'Ponów — przywróć cofnięte wiadomości', + 'chat.revertPopover.title': 'Cofnięte', + 'chat.revertPopover.noTextContent': 'Brak treści tekstowej', + 'chat.revertPopover.forkFromHere': 'Rozgałęź stąd', + 'chat.revertPopover.forkFromMessage': 'Rozgałęź od tej wiadomości', + 'chat.revertPopover.restoreAll': 'Przywróć wszystkie', + 'chat.revertPopover.restore': 'Przywróć', + 'chat.revertPopover.revert': 'Cofnij', + 'chat.revertPopover.fork': 'Rozgałęź', + 'chat.revert.toast.undo': 'Cofnięte do {preview}', + 'chat.revert.toast.redo': 'Ponowione', + 'chat.revert.toast.restored': 'Przywrócono wszystkie wiadomości', 'chat.errorBoundary.title': 'Błąd Czatu', 'chat.errorBoundary.description': 'Interfejs czatu napotkał błąd. Może to być spowodowane tymczasowym problemem sieciowym lub uszkodzonymi danymi wiadomości.', 'chat.errorBoundary.sessionLabel': 'Sesja', diff --git a/packages/ui/src/lib/i18n/messages/pt-BR.ts b/packages/ui/src/lib/i18n/messages/pt-BR.ts index 87ee6229..62aa9bab 100644 --- a/packages/ui/src/lib/i18n/messages/pt-BR.ts +++ b/packages/ui/src/lib/i18n/messages/pt-BR.ts @@ -1368,6 +1368,19 @@ export const dict: Record = { "chat.statusRow.tasksTitle": "Tarefas", "chat.statusRow.summary.activeLeft": "{active} ativas · {left} restantes", "chat.statusRow.aborted": "Interrompido", + "chat.revertIndicator.redo": "Refazer", + "chat.revertIndicator.redoAria": "Refazer — restaurar mensagens revertidas", + "chat.revertPopover.title": "Revertidas", + "chat.revertPopover.noTextContent": "Sem conteúdo de texto", + "chat.revertPopover.forkFromHere": "Criar ramificação aqui", + "chat.revertPopover.forkFromMessage": "Criar ramificação desta mensagem", + "chat.revertPopover.restoreAll": "Restaurar todas", + "chat.revertPopover.restore": "Restaurar", + "chat.revertPopover.revert": "Reverter", + "chat.revertPopover.fork": "Bifurcar", + "chat.revert.toast.undo": "Revertido para {preview}", + "chat.revert.toast.redo": "Refeito", + "chat.revert.toast.restored": "Todas as mensagens restauradas", "chat.errorBoundary.title": "Erro na conversa", "chat.errorBoundary.description": "A interface da conversa encontrou um erro. Isso pode ter sido causado por um problema temporário de rede ou por dados de mensagem corrompidos.", "chat.errorBoundary.sessionLabel": "Sessão", diff --git a/packages/ui/src/lib/i18n/messages/uk.ts b/packages/ui/src/lib/i18n/messages/uk.ts index 3bbe3127..c224ca70 100644 --- a/packages/ui/src/lib/i18n/messages/uk.ts +++ b/packages/ui/src/lib/i18n/messages/uk.ts @@ -1368,6 +1368,19 @@ export const dict: Record = { "chat.statusRow.tasksTitle": "завдання", "chat.statusRow.summary.activeLeft": "Активних: {active} · залишилось: {left}", "chat.statusRow.aborted": "Перервано", + "chat.revertIndicator.redo": "Повторити", + "chat.revertIndicator.redoAria": "Повторити — відновити відкочені повідомлення", + "chat.revertPopover.title": "Відкочено", + "chat.revertPopover.noTextContent": "Немає текстового змісту", + "chat.revertPopover.forkFromHere": "Розгалуження тут", + "chat.revertPopover.forkFromMessage": "Розгалуження від цього повідомлення", + "chat.revertPopover.restoreAll": "Відновити всі", + "chat.revertPopover.restore": "Відновити", + "chat.revertPopover.revert": "Відкотити", + "chat.revertPopover.fork": "Розгалуження", + "chat.revert.toast.undo": "Відкочено до {preview}", + "chat.revert.toast.redo": "Повторено", + "chat.revert.toast.restored": "Всі повідомлення відновлено", "chat.errorBoundary.title": "Помилка чату", "chat.errorBoundary.description": "В інтерфейсі чату сталася помилка. Причиною може бути тимчасова проблема з мережею або пошкоджені дані повідомлення.", "chat.errorBoundary.sessionLabel": "Сесія", diff --git a/packages/ui/src/lib/i18n/messages/zh-CN.ts b/packages/ui/src/lib/i18n/messages/zh-CN.ts index 1449e33c..b3924e3d 100644 --- a/packages/ui/src/lib/i18n/messages/zh-CN.ts +++ b/packages/ui/src/lib/i18n/messages/zh-CN.ts @@ -1368,6 +1368,19 @@ export const dict: Record = { 'chat.statusRow.tasksTitle': '任务', 'chat.statusRow.summary.activeLeft': '{active} 个活跃 · 剩余 {left} 个', 'chat.statusRow.aborted': '已中止', + 'chat.revertIndicator.redo': '重做', + 'chat.revertIndicator.redoAria': '重做 — 恢复已撤回的消息', + 'chat.revertPopover.title': '已撤回', + 'chat.revertPopover.noTextContent': '无文本内容', + 'chat.revertPopover.forkFromHere': '从此处分支', + 'chat.revertPopover.forkFromMessage': '从此消息分支', + 'chat.revertPopover.restoreAll': '恢复全部', + 'chat.revertPopover.restore': '恢复', + 'chat.revertPopover.revert': '撤回', + 'chat.revertPopover.fork': '分叉', + 'chat.revert.toast.undo': '已撤回至 {preview}', + 'chat.revert.toast.redo': '已重做', + 'chat.revert.toast.restored': '已恢复全部消息', 'chat.errorBoundary.title': '聊天错误', 'chat.errorBoundary.description': '聊天界面发生错误，可能是临时网络问题或消息数据损坏导致。', 'chat.errorBoundary.sessionLabel': '会话', diff --git a/packages/ui/src/sync/input-store.ts b/packages/ui/src/sync/input-store.ts index bb10bed3..62984664 100644 --- a/packages/ui/src/sync/input-store.ts +++ b/packages/ui/src/sync/input-store.ts @@ -86,6 +86,8 @@ export type InputState = { addVSCodeFileAttachment: (path: string, name: string, fileSize: number | null) => void addVSCodeSelectionAttachment: (path: string, file: File) => Promise setActiveEditorFile: (file: VSCodeActiveEditorFile | null) => void + /** Add attachments restored from a reverted message (file already on server) */ + addRestoredAttachment: (file: { url: string; mimeType: string; filename: string }) => void } export const useInputStore = create()((set, get) => ({ @@ -210,4 +212,23 @@ export const useInputStore = create()((set, get) => ({ if (isSameVSCodeActiveEditorFile(get().activeEditorFile, file)) return set({ activeEditorFile: file }) }, + + addRestoredAttachment: ({ url, mimeType, filename }) => { + const id = `restored-${Date.now()}-${Math.random().toString(36).slice(2)}` + // Use "local" source so the file renders in AttachedFilesList. + // Set serverPath to the URL so ImagePreview can use it as the img src + // when dataUrl is not a data: URL. sanitizeAttachmentsForSend leaves + // dataUrl alone for non-server sources, so the URL stays intact on send. + const attached: AttachedFile = { + id, + file: new File([], filename, { type: mimeType }), + dataUrl: url, + mimeType, + filename, + size: 0, + source: "local", + serverPath: url, + } + set((s) => ({ attachedFiles: [...s.attachedFiles, attached] })) + }, })) diff --git a/packages/ui/src/sync/session-actions.ts b/packages/ui/src/sync/session-actions.ts index a8ed080f..ae10afdf 100644 --- a/packages/ui/src/sync/session-actions.ts +++ b/packages/ui/src/sync/session-actions.ts @@ -18,6 +18,8 @@ import { stripMessageDiffSnapshots } from "./sanitize" const MESSAGE_REFETCH_LIMIT = 200 const MESSAGE_REFETCH_SKIP_PARTS = new Set(["patch", "step-start", "step-finish"]) +const UNREVERT_REFETCH_ATTEMPTS = 3 +const UNREVERT_REFETCH_RETRY_MS = 150 // Reference set by SyncProvider — allows actions to access SDK and stores let _sdk: OpencodeClient | null = null @@ -26,6 +28,8 @@ let _getDirectory: () => string = () => "" let _optimisticAdd: ((input: { sessionID: string; message: Message; parts: Part[] }) => void) | null = null let _optimisticRemove: ((input: { sessionID: string; messageID: string }) => void) | null = null +const wait = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)) + export function setActionRefs( sdk: OpencodeClient, childStores: ChildStoreManager, @@ -577,6 +581,7 @@ export async function revertToMessage(sessionId: string, messageId: string): Pro const messages = state.message[sessionId] ?? [] const targetMsg = messages.find((m) => m.id === messageId) let messageText = "" + let submittedFileParts: Array> = [] if (targetMsg && targetMsg.role === "user") { const parts = state.part[messageId] ?? [] const textParts = parts.filter((p) => p.type === "text" && !isSyntheticPart(p)) @@ -584,9 +589,16 @@ export async function revertToMessage(sessionId: string, messageId: string): Pro .map((p: Record) => (p as { text?: string }).text || (p as { content?: string }).content || "") .join("\n") .trim() + // Snapshot file parts for later restoration to the input (line ~626). + // File parts (type="file") contain url/mime/filename — the file already + // exists on the server so we record it as a "server" source attachment. + submittedFileParts = parts.filter((p) => p.type === "file") as Array> } - // Optimistically remove reverted messages + set marker + // Optimistically set only the revert marker. Keep messages and parts in the + // local store; visible-message selectors derive the displayed timeline from + // session.revert. This matches the server model and preserves reverted + // messages for the restore dock without maintaining a separate shadow copy. const prevRevert = (() => { const s = state.session.find((s) => s.id === sessionId) return (s as Session & { revert?: unknown })?.revert @@ -594,19 +606,7 @@ export async function revertToMessage(sessionId: string, messageId: string): Pro const sessions = [...state.session] const sessionIdx = sessions.findIndex((s) => s.id === sessionId) - // Remove messages at and after the revert point from the store - const prevMessages = state.message[sessionId] ?? [] - const prevPart = { ...state.part } - const keptMessages = prevMessages.filter((m) => m.id < messageId) - const removedMessages = prevMessages.filter((m) => m.id >= messageId) - for (const m of removedMessages) { - delete prevPart[m.id] - } - - const patch: Record = { - message: { ...state.message, [sessionId]: keptMessages }, - part: prevPart, - } + const patch: Record = {} if (sessionIdx >= 0) { sessions[sessionIdx] = { ...sessions[sessionIdx], revert: { messageID: messageId } } as Session @@ -615,7 +615,13 @@ export async function revertToMessage(sessionId: string, messageId: string): Pro store.setState(patch) - // Restore reverted message text to input + // Save input store state before mutations — if the API fails we need to + // roll back both text and attachments to their previous values. + const prevInputAttachments = [...useInputStore.getState().attachedFiles] + const prevInputText = useInputStore.getState().pendingInputText + const prevInputMode = useInputStore.getState().pendingInputMode + + // Restore reverted message text and file attachments to input if (messageText) { useInputStore.setState({ pendingInputText: messageText, @@ -623,6 +629,20 @@ export async function revertToMessage(sessionId: string, messageId: string): Pro }) } + // Restore file/image attachments from the target message. + // Clear existing attachments first — previous revert's attachments + // must not carry over, even when the current message has no files. + useInputStore.getState().clearAttachedFiles() + for (const fp of submittedFileParts) { + const f = fp as Record + const url = typeof f.url === "string" ? f.url : "" + const mime = typeof f.mime === "string" ? f.mime : "application/octet-stream" + const filename = typeof f.filename === "string" ? f.filename : "attachment" + if (url) { + useInputStore.getState().addRestoredAttachment({ url, mimeType: mime, filename }) + } + } + // Call SDK and merge authoritative result into store try { const result = await sdk().session.revert({ sessionID: sessionId, directory: dir(), messageID: messageId }) @@ -645,8 +665,12 @@ export async function revertToMessage(sessionId: string, messageId: string): Pro } store.setState({ session: rollback, - message: { ...current.message, [sessionId]: prevMessages }, - part: { ...current.part, ...Object.fromEntries(removedMessages.map((m) => [m.id, state.part[m.id] ?? []])) }, + }) + // Rollback input store: restore previous text and attachments + useInputStore.setState({ + pendingInputText: prevInputText, + pendingInputMode: prevInputMode, + attachedFiles: prevInputAttachments, }) throw err } @@ -679,6 +703,7 @@ export async function refetchSessionMessages(sessionId: string): Promise { export async function unrevertSession(sessionId: string): Promise { const store = dirStore() const state = store.getState() + const previousMessageCount = state.message[sessionId]?.length ?? 0 // Abort if busy const status = state.session_status[sessionId] @@ -700,7 +725,12 @@ export async function unrevertSession(sessionId: string): Promise { store.setState({ session: sessions }) } } - await refetchSessionMessages(sessionId) + for (let attempt = 0; attempt < UNREVERT_REFETCH_ATTEMPTS; attempt += 1) { + if (attempt > 0) await wait(UNREVERT_REFETCH_RETRY_MS) + await refetchSessionMessages(sessionId) + const nextMessageCount = store.getState().message[sessionId]?.length ?? 0 + if (nextMessageCount > previousMessageCount) return + } } /** diff --git a/packages/ui/src/sync/session-ui-store.ts b/packages/ui/src/sync/session-ui-store.ts index 46f8a89b..9e367fe1 100644 --- a/packages/ui/src/sync/session-ui-store.ts +++ b/packages/ui/src/sync/session-ui-store.ts @@ -249,10 +249,10 @@ export type SessionUIState = { updateSessionTitle: (sessionId: string, title: string) => Promise shareSession: (sessionId: string) => Promise unshareSession: (sessionId: string) => Promise - revertToMessage: (sessionId: string, messageId: string) => Promise + revertToMessage: (sessionId: string, messageId: string, options?: { skipRedoPush?: boolean }) => Promise forkFromMessage: (sessionId: string, messageId: string) => Promise handleSlashUndo: (sessionId: string) => Promise - handleSlashRedo: (sessionId: string) => Promise + handleSlashRedo: (sessionId: string, options?: { fullUnrevert?: boolean }) => Promise createSessionFromAssistantMessage: (sourceMessageId: string) => Promise // Data access helpers (read from sync) @@ -955,12 +955,15 @@ export const useSessionUIStore = create()((set, get) => ({ // revertToMessage — delegates to session-actions (single implementation) // --------------------------------------------------------------------------- revertToMessage: async (sessionId, messageId) => { + // Ensure the complete message range is present before applying the revert + // marker. Reverted UI is derived from session.revert + stored messages. + await refetchSessionMessages(sessionId) const { revertToMessage: revert } = await import("./session-actions") await revert(sessionId, messageId) }, // --------------------------------------------------------------------------- - // handleSlashUndo — reads from sync + // handleSlashUndo — reads from sync, records history for redo // --------------------------------------------------------------------------- handleSlashUndo: async (sessionId) => { const messages = getSyncMessages(sessionId) @@ -980,52 +983,63 @@ export const useSessionUIStore = create()((set, get) => ({ if (!targetMessage) return + // Read target message parts BEFORE calling revertToMessage. + // revertToMessage optimistically deletes messages from the sync store + // before the API call, so getSyncParts must run first. const targetParts = getSyncParts(targetMessage.id) const textPart = targetParts.find((p: Part) => p.type === "text") as TextPart | undefined const preview = textPart?.text ? String(textPart.text).slice(0, 50) + (textPart.text.length > 50 ? "..." : "") : "[No text]" + // revertToMessage handles the redo stack push internally await get().revertToMessage(sessionId, targetMessage.id) const { toast } = await import("sonner") - toast.success(`Undid to: ${preview}`) + const { useI18nStore, formatMessage } = await import("@/lib/i18n/store") + const { dictionary } = useI18nStore.getState() + toast.success(formatMessage(dictionary, "chat.revert.toast.undo", { preview })) }, // --------------------------------------------------------------------------- - // handleSlashRedo — reads from sync + // handleSlashRedo — moves the authoritative revert marker forward // --------------------------------------------------------------------------- - handleSlashRedo: async (sessionId) => { + handleSlashRedo: async (sessionId, options) => { + if (options?.fullUnrevert) { + const { unrevertSession } = await import("./session-actions") + await unrevertSession(sessionId) + const { toast } = await import("sonner") + const { useI18nStore, formatMessage } = await import("@/lib/i18n/store") + const { dictionary } = useI18nStore.getState() + toast.success(formatMessage(dictionary, "chat.revert.toast.restored")) + return + } + const sessions = getSyncSessions() const currentSession = sessions.find((s) => s.id === sessionId) const revertToId = currentSession?.revert?.messageID if (!revertToId) return await refetchSessionMessages(sessionId) - const messages = getSyncMessages(sessionId) const userMessages = messages.filter((m) => m.role === "user") const targetMessage = userMessages.find((m) => m.id > revertToId) if (targetMessage) { - const targetParts = getSyncParts(targetMessage.id) - const textPart = targetParts.find((p: Part) => p.type === "text") as TextPart | undefined - const preview = textPart?.text - ? String(textPart.text).slice(0, 50) + (textPart.text.length > 50 ? "..." : "") - : "[No text]" - - await get().revertToMessage(sessionId, targetMessage.id) - + await get().revertToMessage(sessionId, targetMessage.id, { skipRedoPush: true }) const { toast } = await import("sonner") - toast.success(`Redid to: ${preview}`) - } else { - // Full unrevert - const { unrevertSession } = await import("./session-actions") - await unrevertSession(sessionId) - - const { toast } = await import("sonner") - toast.success("Restored all messages") + const { useI18nStore, formatMessage } = await import("@/lib/i18n/store") + const { dictionary } = useI18nStore.getState() + toast.success(formatMessage(dictionary, "chat.revert.toast.redo")) + return } + + const { unrevertSession } = await import("./session-actions") + await unrevertSession(sessionId) + const { toast } = await import("sonner") + const { useI18nStore, formatMessage } = await import("@/lib/i18n/store") + const { dictionary } = useI18nStore.getState() + toast.success(formatMessage(dictionary, "chat.revert.toast.restored")) }, // ---------------------------------------------------------------------------