From d73a8957a9c3647271be008d29cab0a1d78b192e Mon Sep 17 00:00:00 2001
From: Emir Karabeg <emirkarabeg@berkeley.edu>
Date: Thu, 30 Jan 2025 11:50:52 -0800
Subject: [PATCH] Fixed and improved version history

---
 .../components/history-dropdown-item.tsx      | 41 ++++++++++++++++---
 app/w/components/control-bar/control-bar.tsx  | 33 ++++++++++++---
 stores/workflow/history-middleware.ts         | 21 ++++++++++
 stores/workflow/history-types.ts              |  1 +
 4 files changed, 86 insertions(+), 10 deletions(-)

diff --git a/app/w/components/control-bar/components/history-dropdown-item.tsx b/app/w/components/control-bar/components/history-dropdown-item.tsx
index 2c76066f1..33d2640a3 100644
--- a/app/w/components/control-bar/components/history-dropdown-item.tsx
+++ b/app/w/components/control-bar/components/history-dropdown-item.tsx
@@ -8,6 +8,7 @@ interface HistoryDropdownItemProps {
   timestamp: number
   onClick?: () => void
   isCurrent?: boolean
+  isFuture?: boolean
   id?: string
 }
 
@@ -16,25 +17,55 @@ export function HistoryDropdownItem({
   timestamp,
   onClick,
   isCurrent = false,
+  isFuture = false,
   id,
 }: HistoryDropdownItemProps) {
   const timeAgo = formatDistanceToNow(timestamp, { addSuffix: true })
 
   return (
     <DropdownMenuItem
-      className="flex items-start gap-2 p-3 cursor-pointer"
+      className={cn(
+        'flex items-start gap-2 p-3 cursor-pointer',
+        isFuture && 'text-muted-foreground/50'
+      )}
       onClick={onClick}
     >
-      <Clock className="h-4 w-4 text-muted-foreground" />
+      <Clock
+        className={cn(
+          'h-4 w-4',
+          isFuture ? 'text-muted-foreground/50' : 'text-muted-foreground'
+        )}
+      />
       <div className="flex flex-col gap-1">
         <div className="flex items-center gap-2">
           {isCurrent ? (
-            <span className="text-xs text-muted-foreground">Current</span>
+            <span
+              className={cn(
+                'text-xs',
+                isFuture ? 'text-muted-foreground/50' : 'text-muted-foreground'
+              )}
+            >
+              Current
+            </span>
           ) : (
-            <span className="text-xs text-muted-foreground">{timeAgo}</span>
+            <span
+              className={cn(
+                'text-xs',
+                isFuture ? 'text-muted-foreground/50' : 'text-muted-foreground'
+              )}
+            >
+              {timeAgo}
+            </span>
           )}
         </div>
-        <p className="text-sm text-foreground">{action}</p>
+        <p
+          className={cn(
+            'text-sm',
+            isFuture ? 'text-muted-foreground/50' : 'text-foreground'
+          )}
+        >
+          {action}
+        </p>
       </div>
     </DropdownMenuItem>
   )
diff --git a/app/w/components/control-bar/control-bar.tsx b/app/w/components/control-bar/control-bar.tsx
index 287865b8d..7f5279628 100644
--- a/app/w/components/control-bar/control-bar.tsx
+++ b/app/w/components/control-bar/control-bar.tsx
@@ -25,7 +25,7 @@ import {
 
 export function ControlBar() {
   const { notifications, getWorkflowNotifications } = useNotificationStore()
-  const { history, undo, redo } = useWorkflowStore()
+  const { history, undo, redo, revertToHistoryState } = useWorkflowStore()
   const [, forceUpdate] = useState({})
   const { isExecuting, handleRunWorkflow } = useWorkflowExecution()
   const { workflows, removeWorkflow, activeWorkflowId } = useWorkflowRegistry()
@@ -109,15 +109,33 @@ export function ControlBar() {
             </Button>
           </DropdownMenuTrigger>
 
-          {history.past.length === 0 ? (
+          {history.past.length === 0 && history.future.length === 0 ? (
             <DropdownMenuContent align="end" className="w-40">
               <DropdownMenuItem className="text-sm text-muted-foreground">
                 No history available
               </DropdownMenuItem>
             </DropdownMenuContent>
           ) : (
-            <DropdownMenuContent align="end" className="w-60">
+            <DropdownMenuContent
+              align="end"
+              className="w-60 max-h-[300px] overflow-y-auto"
+            >
               <>
+                {[...history.future].reverse().map((entry, index) => (
+                  <HistoryDropdownItem
+                    key={`future-${entry.timestamp}-${index}`}
+                    action={entry.action}
+                    timestamp={entry.timestamp}
+                    onClick={() =>
+                      revertToHistoryState(
+                        history.past.length +
+                          1 +
+                          (history.future.length - 1 - index)
+                      )
+                    }
+                    isFuture={true}
+                  />
+                ))}
                 <HistoryDropdownItem
                   key={`current-${history.present.timestamp}`}
                   action={history.present.action}
@@ -130,7 +148,9 @@ export function ControlBar() {
                     key={`past-${entry.timestamp}-${index}`}
                     action={entry.action}
                     timestamp={entry.timestamp}
-                    onClick={undo}
+                    onClick={() =>
+                      revertToHistoryState(history.past.length - 1 - index)
+                    }
                   />
                 ))}
               </>
@@ -153,7 +173,10 @@ export function ControlBar() {
               </DropdownMenuItem>
             </DropdownMenuContent>
           ) : (
-            <DropdownMenuContent align="end" className="w-60">
+            <DropdownMenuContent
+              align="end"
+              className="w-60 max-h-[300px] overflow-y-auto"
+            >
               {[...workflowNotifications]
                 .sort((a, b) => b.timestamp - a.timestamp)
                 .map((notification) => (
diff --git a/stores/workflow/history-middleware.ts b/stores/workflow/history-middleware.ts
index cce9033d5..939c7f0cc 100644
--- a/stores/workflow/history-middleware.ts
+++ b/stores/workflow/history-middleware.ts
@@ -86,6 +86,27 @@ export const withHistory = (
         set(newState)
         return newState
       },
+
+      revertToHistoryState: (index: number) => {
+        const { history, ...state } = get()
+        const allStates = [...history.past, history.present, ...history.future]
+        const targetState = allStates[index]
+        
+        if (!targetState) return
+
+        const newPast = allStates.slice(0, index)
+        const newFuture = allStates.slice(index + 1)
+
+        set({
+          ...state,
+          ...targetState.state,
+          history: {
+            past: newPast,
+            present: targetState,
+            future: newFuture,
+          },
+        })
+      },
     }
   }
 }
diff --git a/stores/workflow/history-types.ts b/stores/workflow/history-types.ts
index 94d292d33..fdb115b2b 100644
--- a/stores/workflow/history-types.ts
+++ b/stores/workflow/history-types.ts
@@ -17,4 +17,5 @@ export interface HistoryActions {
   redo: () => void
   canUndo: () => boolean
   canRedo: () => boolean
+  revertToHistoryState: (index: number) => void
 } 
\ No newline at end of file