From 13e30eb0a3ac65b1b2025e302c86922588dbdbb8 Mon Sep 17 00:00:00 2001
From: Emir Karabeg <emirkarabeg@berkeley.edu>
Date: Wed, 29 Jan 2025 13:26:36 -0800
Subject: [PATCH] Added ability to add/delete workflows and persist state and
 history

---
 app/w/[id]/workflow.tsx                      |  11 ++
 app/w/components/control-bar/control-bar.tsx |  31 +++-
 app/w/components/sidebar/sidebar.tsx         |  80 ++++++--
 stores/workflow/history-middleware.ts        |  18 ++
 stores/workflow/registry-types.ts            |  23 +++
 stores/workflow/types.ts                     |   3 +
 stores/workflow/workflow-registry.ts         | 182 +++++++++++++++++++
 stores/workflow/workflow-store.ts            |  27 ++-
 8 files changed, 358 insertions(+), 17 deletions(-)
 create mode 100644 stores/workflow/registry-types.ts
 create mode 100644 stores/workflow/workflow-registry.ts

diff --git a/app/w/[id]/workflow.tsx b/app/w/[id]/workflow.tsx
index b603a192e..cb40fa053 100644
--- a/app/w/[id]/workflow.tsx
+++ b/app/w/[id]/workflow.tsx
@@ -28,6 +28,8 @@ import { NotificationList } from '@/app/w/components/notifications/notifications
 import { useNotificationStore } from '@/stores/notifications/notifications-store'
 import { executeWorkflow } from '@/lib/workflow'
 import { useWorkflowExecution } from '../hooks/use-workflow-execution'
+import { useWorkflowRegistry } from '@/stores/workflow/workflow-registry'
+import { useParams } from 'next/navigation'
 
 /**
  * Represents the data structure for a workflow node
@@ -304,6 +306,15 @@ function WorkflowCanvas() {
  * Root workflow component that provides the ReactFlow context to the canvas
  */
 export default function Workflow() {
+  const params = useParams()
+  const { setActiveWorkflow } = useWorkflowRegistry()
+
+  useEffect(() => {
+    if (params.id) {
+      setActiveWorkflow(params.id as string)
+    }
+  }, [params.id, setActiveWorkflow])
+
   return (
     <ReactFlowProvider>
       <WorkflowCanvas />
diff --git a/app/w/components/control-bar/control-bar.tsx b/app/w/components/control-bar/control-bar.tsx
index d346835fe..88ac787a3 100644
--- a/app/w/components/control-bar/control-bar.tsx
+++ b/app/w/components/control-bar/control-bar.tsx
@@ -7,7 +7,7 @@ import {
   DropdownMenuItem,
   DropdownMenuTrigger,
 } from '@/components/ui/dropdown-menu'
-import { History, Bell, Play } from 'lucide-react'
+import { History, Bell, Play, Trash2 } from 'lucide-react'
 import { useNotificationStore } from '@/stores/notifications/notifications-store'
 import { NotificationDropdownItem } from './components/notification-dropdown-item'
 import { useWorkflowStore } from '@/stores/workflow/workflow-store'
@@ -15,12 +15,31 @@ import { HistoryDropdownItem } from './components/history-dropdown-item'
 import { formatDistanceToNow } from 'date-fns'
 import { useEffect, useState } from 'react'
 import { useWorkflowExecution } from '../../hooks/use-workflow-execution'
+import { useWorkflowRegistry } from '@/stores/workflow/workflow-registry'
+import { useRouter } from 'next/navigation'
 
 export function ControlBar() {
   const { notifications } = useNotificationStore()
   const { history, undo, redo } = useWorkflowStore()
   const [, forceUpdate] = useState({})
   const { isExecuting, handleRunWorkflow } = useWorkflowExecution()
+  const { workflows, removeWorkflow, activeWorkflowId } = useWorkflowRegistry()
+  const router = useRouter()
+
+  const handleDeleteWorkflow = () => {
+    if (!activeWorkflowId) return
+    const newWorkflows = { ...workflows }
+    delete newWorkflows[activeWorkflowId]
+    const remainingIds = Object.keys(newWorkflows)
+
+    removeWorkflow(activeWorkflowId)
+
+    if (remainingIds.length > 0) {
+      router.push(`/w/${remainingIds[0]}`)
+    } else {
+      router.push('/')
+    }
+  }
 
   // Update the time display every minute
   useEffect(() => {
@@ -44,6 +63,16 @@ export function ControlBar() {
 
       {/* Right Section - Actions */}
       <div className="flex items-center gap-3">
+        <Button
+          variant="ghost"
+          size="icon"
+          onClick={handleDeleteWorkflow}
+          disabled={Object.keys(workflows).length <= 1}
+        >
+          <Trash2 className="h-5 w-5" />
+          <span className="sr-only">Delete Workflow</span>
+        </Button>
+
         <DropdownMenu>
           <DropdownMenuTrigger asChild>
             <Button variant="ghost" size="icon">
diff --git a/app/w/components/sidebar/sidebar.tsx b/app/w/components/sidebar/sidebar.tsx
index 6aae63a2b..fd1d56e70 100644
--- a/app/w/components/sidebar/sidebar.tsx
+++ b/app/w/components/sidebar/sidebar.tsx
@@ -1,3 +1,5 @@
+'use client'
+
 import Link from 'next/link'
 import { NavItem } from './components/nav-item'
 import { Settings, Plus } from 'lucide-react'
@@ -7,32 +9,82 @@ import {
   TooltipTrigger,
 } from '@/components/ui/tooltip'
 import { AgentIcon } from '@/components/icons'
+import { useWorkflowRegistry } from '@/stores/workflow/workflow-registry'
+import { useRouter } from 'next/navigation'
+import { Button } from '@/components/ui/button'
 
 export function Sidebar() {
+  const WORKFLOW_COLORS = [
+    '#3972F6',
+    '#F639DD',
+    '#F6B539',
+    '#8139F6',
+    '#F64439',
+  ]
+  const { workflows, addWorkflow } = useWorkflowRegistry()
+  const router = useRouter()
+
+  const handleCreateWorkflow = () => {
+    const id = crypto.randomUUID()
+    const colorIndex = Object.keys(workflows).length % WORKFLOW_COLORS.length
+    const newWorkflow = {
+      id,
+      name: `Workflow ${Object.keys(workflows).length + 1}`,
+      lastModified: new Date(),
+      description: 'New workflow',
+      color: WORKFLOW_COLORS[colorIndex],
+    }
+
+    addWorkflow(newWorkflow)
+    router.push(`/w/${id}`)
+  }
+
   return (
     <aside className="fixed inset-y-0 left-0 z-10 hidden w-14 flex-col border-r bg-background sm:flex">
-      <nav className="flex flex-col items-center gap-4 px-2 sm:py-5">
+      <nav className="flex flex-col items-center gap-4 px-2 py-5">
         <Link
-          href="#"
+          href="/"
           className="group flex h-8 w-8 items-center justify-center rounded-lg bg-[#7F2FFF]"
         >
           <AgentIcon className="text-white transition-all group-hover:scale-110 -translate-y-[0.5px] w-5 h-5" />
           <span className="sr-only">Sim Studio</span>
         </Link>
 
-        <NavItem href="#" label="Add Workflow">
-          <Plus className="h-5 w-5" />
-        </NavItem>
-
-        <NavItem href="#" label="Workflow 1">
-          <div className="h-4 w-4 rounded-full bg-[#3972F6]" />
-        </NavItem>
-
-        <NavItem href="#" label="Workflow 2">
-          <div className="h-4 w-4 rounded-full bg-[#F639DD]" />
-        </NavItem>
+        <Tooltip>
+          <TooltipTrigger asChild>
+            <Button
+              variant="ghost"
+              size="icon"
+              onClick={handleCreateWorkflow}
+              className="h-9 w-9 md:h-8 md:w-8"
+            >
+              <Plus className="h-5 w-5" />
+              <span className="sr-only">Add Workflow</span>
+            </Button>
+          </TooltipTrigger>
+          <TooltipContent side="right">Add Workflow</TooltipContent>
+        </Tooltip>
       </nav>
-      <nav className="mt-auto flex flex-col items-center gap-4 px-2 sm:py-5">
+
+      {/* Scrollable workflows section */}
+      <nav className="flex-1 overflow-y-auto px-2 [&::-webkit-scrollbar]:hidden [-ms-overflow-style:'none'] [scrollbar-width:'none']">
+        <div className="flex flex-col items-center gap-4">
+          {Object.values(workflows).map((workflow) => (
+            <NavItem
+              key={workflow.id}
+              href={`/w/${workflow.id}`}
+              label={workflow.name}
+            >
+              <div
+                className="h-4 w-4 rounded-full"
+                style={{ backgroundColor: workflow.color || '#3972F6' }}
+              />
+            </NavItem>
+          ))}
+        </div>
+      </nav>
+
+      <nav className="flex flex-col items-center gap-4 px-2 py-5">
         <Tooltip>
           <TooltipTrigger asChild>
             <Link
diff --git a/stores/workflow/history-middleware.ts b/stores/workflow/history-middleware.ts
index 8ecc91e63..cce9033d5 100644
--- a/stores/workflow/history-middleware.ts
+++ b/stores/workflow/history-middleware.ts
@@ -68,6 +68,24 @@ export const withHistory = (
           },
         })
       },
+
+      clear: () => {
+        const newState = {
+          blocks: {},
+          edges: [],
+          history: {
+            past: [],
+            present: {
+              state: { blocks: {}, edges: [] },
+              timestamp: Date.now(),
+              action: 'Clear workflow'
+            },
+            future: []
+          }
+        }
+        set(newState)
+        return newState
+      },
     }
   }
 }
diff --git a/stores/workflow/registry-types.ts b/stores/workflow/registry-types.ts
new file mode 100644
index 000000000..a77573bb1
--- /dev/null
+++ b/stores/workflow/registry-types.ts
@@ -0,0 +1,23 @@
+export interface WorkflowMetadata {
+  id: string
+  name: string
+  lastModified: Date
+  description?: string
+  color?: string
+}
+
+export interface WorkflowRegistryState {
+  workflows: Record<string, WorkflowMetadata>
+  activeWorkflowId: string | null
+  isLoading: boolean
+  error: string | null
+}
+
+export interface WorkflowRegistryActions {
+  setActiveWorkflow: (id: string) => Promise<void>
+  addWorkflow: (metadata: WorkflowMetadata) => void
+  removeWorkflow: (id: string) => void
+  updateWorkflow: (id: string, metadata: Partial<WorkflowMetadata>) => void
+}
+
+export type WorkflowRegistry = WorkflowRegistryState & WorkflowRegistryActions 
\ No newline at end of file
diff --git a/stores/workflow/types.ts b/stores/workflow/types.ts
index 9cf2b2833..792f3fbfe 100644
--- a/stores/workflow/types.ts
+++ b/stores/workflow/types.ts
@@ -1,5 +1,6 @@
 import { Node, Edge } from 'reactflow'
 import { OutputType, SubBlockType } from '@/blocks/types'
+import { WorkflowHistory } from './history-types'
 
 export interface Position {
   x: number
@@ -24,6 +25,7 @@ export interface SubBlockState {
 export interface WorkflowState {
   blocks: Record<string, BlockState>
   edges: Edge[]
+  lastSaved?: number
 }
 
 export interface WorkflowActions {
@@ -39,6 +41,7 @@ export interface WorkflowActions {
   addEdge: (edge: Edge) => void
   removeEdge: (edgeId: string) => void
   clear: () => void
+  updateLastSaved: () => void
 }
 
 export type WorkflowStore = WorkflowState & WorkflowActions 
\ No newline at end of file
diff --git a/stores/workflow/workflow-registry.ts b/stores/workflow/workflow-registry.ts
new file mode 100644
index 000000000..45aa3226f
--- /dev/null
+++ b/stores/workflow/workflow-registry.ts
@@ -0,0 +1,182 @@
+import { create } from 'zustand'
+import { devtools } from 'zustand/middleware'
+import { WorkflowRegistry, WorkflowMetadata } from './registry-types'
+import { useWorkflowStore } from './workflow-store'
+
+export const useWorkflowRegistry = create<WorkflowRegistry>()(
+  devtools(
+    (set, get) => ({
+      workflows: {},
+      activeWorkflowId: null,
+      isLoading: false,
+      error: null,
+
+      setActiveWorkflow: async (id: string) => {
+        const { workflows } = get()
+        if (!workflows[id]) {
+          set({ error: `Workflow ${id} not found` })
+          return
+        }
+
+        // Save current workflow state before switching
+        const currentId = get().activeWorkflowId
+        if (currentId) {
+          const currentState = useWorkflowStore.getState()
+          localStorage.setItem(`workflow-${currentId}`, JSON.stringify({
+            blocks: currentState.blocks,
+            edges: currentState.edges,
+            history: currentState.history // Save history state
+          }))
+        }
+
+        // Load new workflow state
+        const savedState = localStorage.getItem(`workflow-${id}`)
+        if (savedState) {
+          const { blocks, edges, history } = JSON.parse(savedState)
+          useWorkflowStore.setState({ 
+            blocks, 
+            edges,
+            history: history || {
+              past: [],
+              present: {
+                state: { blocks, edges },
+                timestamp: Date.now(),
+                action: 'Initial state'
+              },
+              future: []
+            }
+          })
+        } else {
+          useWorkflowStore.setState({
+            blocks: {},
+            edges: [],
+            history: {
+              past: [],
+              present: {
+                state: { blocks: {}, edges: [] },
+                timestamp: Date.now(),
+                action: 'Initial state'
+              },
+              future: []
+            },
+            lastSaved: Date.now()
+          })
+        }
+
+        set({ activeWorkflowId: id, error: null })
+      },
+
+      addWorkflow: (metadata: WorkflowMetadata) => {
+        set((state) => ({
+          workflows: {
+            ...state.workflows,
+            [metadata.id]: metadata
+          },
+          error: null
+        }))
+
+        // Save workflow list to localStorage
+        const workflows = get().workflows
+        localStorage.setItem('workflow-registry', JSON.stringify(workflows))
+      },
+
+      removeWorkflow: (id: string) => {
+        set((state) => {
+          const newWorkflows = { ...state.workflows }
+          delete newWorkflows[id]
+
+          // Remove workflow state from localStorage
+          localStorage.removeItem(`workflow-${id}`)
+          
+          // Update registry in localStorage
+          localStorage.setItem('workflow-registry', JSON.stringify(newWorkflows))
+
+          // If deleting active workflow, switch to another one
+          let newActiveWorkflowId = state.activeWorkflowId
+          if (state.activeWorkflowId === id) {
+            const remainingIds = Object.keys(newWorkflows)
+            // Switch to first available workflow
+            newActiveWorkflowId = remainingIds[0]
+            const savedState = localStorage.getItem(`workflow-${newActiveWorkflowId}`)
+            if (savedState) {
+              const { blocks, edges, history } = JSON.parse(savedState)
+              useWorkflowStore.setState({ 
+                blocks, 
+                edges,
+                history: history || {
+                  past: [],
+                  present: {
+                    state: { blocks, edges },
+                    timestamp: Date.now(),
+                    action: 'Initial state'
+                  },
+                  future: []
+                }
+              })
+            } else {
+              useWorkflowStore.setState({
+                blocks: {},
+                edges: [],
+                history: {
+                  past: [],
+                  present: {
+                    state: { blocks: {}, edges: [] },
+                    timestamp: Date.now(),
+                    action: 'Initial state'
+                  },
+                  future: []
+                },
+                lastSaved: Date.now()
+              })
+            }
+          }
+
+          return {
+            workflows: newWorkflows,
+            activeWorkflowId: newActiveWorkflowId,
+            error: null
+          }
+        })
+      },
+
+      updateWorkflow: (id: string, metadata: Partial<WorkflowMetadata>) => {
+        set((state) => {
+          const workflow = state.workflows[id]
+          if (!workflow) return state
+
+          const updatedWorkflows = {
+            ...state.workflows,
+            [id]: {
+              ...workflow,
+              ...metadata,
+              lastModified: new Date()
+            }
+          }
+
+          // Update registry in localStorage
+          localStorage.setItem('workflow-registry', JSON.stringify(updatedWorkflows))
+
+          return {
+            workflows: updatedWorkflows,
+            error: null
+          }
+        })
+      }
+    }),
+    { name: 'workflow-registry' }
+  )
+)
+
+// Initialize registry from localStorage
+const initializeRegistry = () => {
+  const savedRegistry = localStorage.getItem('workflow-registry')
+  if (savedRegistry) {
+    const workflows = JSON.parse(savedRegistry)
+    useWorkflowRegistry.setState({ workflows })
+  }
+}
+
+// Call this in your app's entry point
+if (typeof window !== 'undefined') {
+  initializeRegistry()
+} 
\ No newline at end of file
diff --git a/stores/workflow/workflow-store.ts b/stores/workflow/workflow-store.ts
index 2f8edf92f..7dbe17cd7 100644
--- a/stores/workflow/workflow-store.ts
+++ b/stores/workflow/workflow-store.ts
@@ -9,6 +9,7 @@ import { resolveOutputType } from '@/blocks/utils'
 const initialState = {
   blocks: {},
   edges: [],
+  lastSaved: undefined,
   history: {
     past: [],
     present: {
@@ -98,6 +99,7 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
 
         set(newState)
         pushHistory(set, get, newState, `Add ${type} block`)
+        get().updateLastSaved()
       },
 
       updateBlockPosition: (id: string, position: Position) => {
@@ -111,6 +113,7 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
           },
           edges: [...state.edges],
         }))
+        get().updateLastSaved()
       },
 
       removeBlock: (id: string) => {
@@ -124,6 +127,7 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
         
         set(newState)
         pushHistory(set, get, newState, 'Remove block')
+        get().updateLastSaved()
       },
 
       addEdge: (edge: Edge) => {
@@ -143,6 +147,7 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
         
         set(newState)
         pushHistory(set, get, newState, 'Add connection')
+        get().updateLastSaved()
       },
 
       removeEdge: (edgeId: string) => {
@@ -153,12 +158,30 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
         
         set(newState)
         pushHistory(set, get, newState, 'Remove connection')
+        get().updateLastSaved()
       },
 
       clear: () => {
-        const newState = initialState
+        const newState = {
+          blocks: {},
+          edges: [],
+          history: {
+            past: [],
+            present: {
+              state: { blocks: {}, edges: [] },
+              timestamp: Date.now(),
+              action: 'Initial state'
+            },
+            future: []
+          },
+          lastSaved: Date.now(),
+        }
         set(newState)
-        pushHistory(set, get, newState, 'Clear workflow')
+        return newState
+      },
+
+      updateLastSaved: () => {
+        set({ lastSaved: Date.now() })
       },
     })),
     { name: 'workflow-store' }