feat(ui): revert indicator with undo/redo, message list, and attachment restore (#1279)

* feat(ui): revert indicator with undo/redo, message list, and attachment restore Add bidirectional undo/redo with redo stack navigation and expandable revert indicator in StatusRow. List reverted messages with inline revert/fork actions. Restore file attachments on revert. - Add revert indicator with count, expandable popover, per-button spinner - Restore file/image attachments when reverting to a message - Clear previous attachments on revert when target has none - Restore-all bypasses redo stack for direct unrevert - Survive popover close/reopen without losing loading state - Fix flash when sending message after revert - Fix count disappearing on browser refresh - Remove dead code (undoStack, getRevertHistory, fork-from-here) - Fix toast grammar (Undid -> Reverted, Redid -> Redone) - Add i18n keys for revert popover across all locales * fix(ui): Greptile review fixes and i18n for revert toasts - Fix handleSlashUndo toast always showing [No text] (moved getSyncParts before revertToMessage) - Add inputStore rollback in revertToMessage catch (restore attachments + text on API failure) - Change portal ID to per-session (prevent multi-session collisions) - Add i18n keys for undo/redo/restored toasts across all 7 locales - Use formatMessage in store for localized toast strings * fix(ui): add missing sessionId to click-outside effect deps * fix(ui): remove unused sessionActions import in ChatMessage * fix(ui): derive revert dock from session state --------- Co-authored-by: Bohdan Triapitsyn <artmore@protonmail.com>
2026-05-16 16:44:37 +03:00
parent 947f725976
commit 9c71119835
13 changed files with 371 additions and 47 deletions
@@ -86,6 +86,8 @@ export type InputState = {
  addVSCodeFileAttachment: (path: string, name: string, fileSize: number | null) => void
  addVSCodeSelectionAttachment: (path: string, file: File) => Promise<void>
  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<InputState>()((set, get) => ({
@@ -210,4 +212,23 @@ export const useInputStore = create<InputState>()((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] }))
+  },
 }))
@@ -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<Record<string, unknown>> = []
  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<string, unknown>) => (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<Record<string, unknown>>
  }

-  // 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<string, unknown> = {
-    message: { ...state.message, [sessionId]: keptMessages },
-    part: prevPart,
-  }
+  const patch: Record<string, unknown> = {}

  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<string, unknown>
+      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<void> {
 export async function unrevertSession(sessionId: string): Promise<void> {
  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<void> {
      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
+  }
 }

 /**
@@ -249,10 +249,10 @@ export type SessionUIState = {
  updateSessionTitle: (sessionId: string, title: string) => Promise<void>
  shareSession: (sessionId: string) => Promise<Session | null>
  unshareSession: (sessionId: string) => Promise<Session | null>
-  revertToMessage: (sessionId: string, messageId: string) => Promise<void>
+  revertToMessage: (sessionId: string, messageId: string, options?: { skipRedoPush?: boolean }) => Promise<void>
  forkFromMessage: (sessionId: string, messageId: string) => Promise<void>
  handleSlashUndo: (sessionId: string) => Promise<void>
-  handleSlashRedo: (sessionId: string) => Promise<void>
+  handleSlashRedo: (sessionId: string, options?: { fullUnrevert?: boolean }) => Promise<void>
  createSessionFromAssistantMessage: (sourceMessageId: string) => Promise<void>

  // Data access helpers (read from sync)
@@ -955,12 +955,15 @@ export const useSessionUIStore = create<SessionUIState>()((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<SessionUIState>()((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"))
  },

  // ---------------------------------------------------------------------------