feat(chat): add revert to previous message

Add revert button on user messages except first message
Restore reverted message text to input field
Filter out messages after revert point

packages/ui/src/stores/messageStore.ts

@@ -290,6 +290,30 @@ const resolveSessionDirectory = async (sessionId: string | null | undefined): Pr
     }
 };
 
+const getSessionRevertMessageId = (sessionId: string | null | undefined): string | null => {
+    if (!sessionId) return null;
+    try {
+        const sessionStore = useSessionStore.getState();
+        const session = sessionStore.sessions.find((entry) => entry.id === sessionId) as { revert?: { messageID?: string } } | undefined;
+        return session?.revert?.messageID ?? null;
+    } catch {
+        return null;
+    }
+};
+
+const filterRevertedMessages = (
+    messages: { info: Message; parts: Part[] }[],
+    revertMessageId: string | null
+): { info: Message; parts: Part[] }[] => {
+    if (!revertMessageId) return messages;
+
+    const revertIndex = messages.findIndex((m) => m.info.id === revertMessageId);
+    if (revertIndex === -1) return messages;
+
+    // Keep only messages before the revert point (exclusive)
+    return messages.slice(0, revertIndex);
+};
+
 const executeWithSessionDirectory = async <T>(sessionId: string | null | undefined, operation: () => Promise<T>): Promise<T> => {
     const directoryOverride = await resolveSessionDirectory(sessionId);
     if (directoryOverride) {
@@ -358,15 +382,20 @@ export const useMessageStore = create<MessageStore>()(
 
                 loadMessages: async (sessionId: string, limit: number = MEMORY_LIMITS.VIEWPORT_MESSAGES) => {
                         const allMessages = await executeWithSessionDirectory(sessionId, () => opencodeClient.getSessionMessages(sessionId));
+
+                        // Filter out reverted messages first
+                        const revertMessageId = getSessionRevertMessageId(sessionId);
+                        const messagesWithoutReverted = filterRevertedMessages(allMessages, revertMessageId);
+
                         const watermark = get().sessionMemoryState.get(sessionId)?.trimmedHeadMaxId;
 
                         const afterWatermark = watermark
-                            ? allMessages.filter((message) => {
+                            ? messagesWithoutReverted.filter((message) => {
                                   const messageId = message?.info?.id;
                                   if (!messageId) return true;
                                   return isIdNewer(messageId, watermark);
                               })
-                            : allMessages;
+                            : messagesWithoutReverted;
                         const messagesToKeep = afterWatermark.slice(-limit);
 
                         set((state) => {
@@ -1836,15 +1865,19 @@ export const useMessageStore = create<MessageStore>()(
                 },
 
                 syncMessages: (sessionId: string, messages: { info: Message; parts: Part[] }[]) => {
+                    // Filter out reverted messages first
+                    const revertMessageId = getSessionRevertMessageId(sessionId);
+                    const messagesWithoutReverted = filterRevertedMessages(messages, revertMessageId);
+
                     const watermark = get().sessionMemoryState.get(sessionId)?.trimmedHeadMaxId;
                     const messagesFiltered = watermark
-                        ? messages.filter((message) => {
+                        ? messagesWithoutReverted.filter((message) => {
                               const messageId = message?.info?.id;
                               if (!messageId) return true;
 
                               return isIdNewer(messageId, watermark);
                           })
-                        : messages;
+                        : messagesWithoutReverted;
 
                     set((state) => {
                         const newMessages = new Map(state.messages);

packages/ui/src/stores/types/sessionTypes.ts

@@ -96,6 +96,8 @@ export interface SessionStore {
 
     userSummaryTitles: Map<string, { title: string; createdAt: number | null }>;
 
+    pendingInputText: string | null;
+
     getSessionAgentEditMode: (sessionId: string, agentName: string | undefined, defaultMode?: EditPermissionMode) => EditPermissionMode;
     toggleSessionAgentEditMode: (sessionId: string, agentName: string | undefined, defaultMode?: EditPermissionMode) => void;
     setSessionAgentEditMode: (sessionId: string, agentName: string | undefined, mode: EditPermissionMode, defaultMode?: EditPermissionMode) => void;
@@ -169,4 +171,8 @@ export interface SessionStore {
 
      pollForTokenUpdates: (sessionId: string, messageId: string, maxAttempts?: number) => void;
      updateSession: (session: Session) => void;
+
+     revertToMessage: (sessionId: string, messageId: string) => Promise<void>;
+     setPendingInputText: (text: string | null) => void;
+     consumePendingInputText: () => string | null;
  }

packages/ui/src/stores/useSessionStore.ts

@@ -94,6 +94,7 @@ export const useSessionStore = create<SessionStore>()(
             abortPromptExpiresAt: null,
             sessionActivityPhase: new Map(),
             userSummaryTitles: new Map(),
+            pendingInputText: null,
 
                 getSessionAgentEditMode: (sessionId: string, agentName: string | undefined, defaultMode?: EditPermissionMode) => {
                     return useContextStore.getState().getSessionAgentEditMode(sessionId, agentName, defaultMode);
@@ -340,6 +341,63 @@ export const useSessionStore = create<SessionStore>()(
                     return useContextStore.getState().pollForTokenUpdates(sessionId, messageId, messages, maxAttempts);
                 },
                 updateSession: (session: Session) => useSessionManagementStore.getState().updateSession(session),
+
+                revertToMessage: async (sessionId: string, messageId: string) => {
+                    // Get the message text before reverting
+                    const messages = useMessageStore.getState().messages.get(sessionId) || [];
+                    const targetMessage = messages.find((m) => m.info.id === messageId);
+                    let messageText = '';
+
+                    if (targetMessage && targetMessage.info.role === 'user') {
+                        // Extract text from user message parts
+                        const textParts = targetMessage.parts.filter((p) => p.type === 'text');
+                        messageText = textParts
+                            .map((p) => {
+                                const part = p as { text?: string; content?: string };
+                                return part.text || part.content || '';
+                            })
+                            .join('\n')
+                            .trim();
+                    }
+
+                    // Call revert API
+                    const updatedSession = await opencodeClient.revertSession(sessionId, messageId);
+
+                    // Update session in store (this stores the revert.messageID)
+                    useSessionManagementStore.getState().updateSession(updatedSession);
+
+                    // Filter out reverted messages from the store
+                    // Messages with ID >= revert.messageID should be removed
+                    const currentMessages = useMessageStore.getState().messages.get(sessionId) || [];
+                    const revertMessageId = updatedSession.revert?.messageID;
+
+                    if (revertMessageId) {
+                        // Find the index of the revert message
+                        const revertIndex = currentMessages.findIndex((m) => m.info.id === revertMessageId);
+                        if (revertIndex !== -1) {
+                            // Keep only messages before the revert point
+                            const filteredMessages = currentMessages.slice(0, revertIndex);
+                            useMessageStore.getState().syncMessages(sessionId, filteredMessages);
+                        }
+                    }
+
+                    // Set pending input text for ChatInput to consume
+                    if (messageText) {
+                        set({ pendingInputText: messageText });
+                    }
+                },
+
+                setPendingInputText: (text: string | null) => {
+                    set({ pendingInputText: text });
+                },
+
+                consumePendingInputText: () => {
+                    const text = get().pendingInputText;
+                    if (text !== null) {
+                        set({ pendingInputText: null });
+                    }
+                    return text;
+                },
             }),
         {
             name: "composed-session-store",