From 0313f24a8364f29f29c145f24adda8f4432c09f0 Mon Sep 17 00:00:00 2001
From: Emir Karabeg <emirkarabeg@berkeley.edu>
Date: Thu, 16 Jan 2025 12:53:33 -0800
Subject: [PATCH] State works without subblock state

---
 .../sub-block/components/dropdown.tsx         |  9 +++++++-
 .../sub-block/components/long-input.tsx       |  8 ++++++-
 .../sub-block/components/short-input.tsx      |  9 +++++++-
 .../sub-block/components/slider-input.tsx     |  4 ++++
 .../sub-block/hooks/use-sub-block-value.ts    | 23 +++++++++++++++++++
 .../workflow-block/sub-block/sub-block.tsx    | 21 ++++++++++++++---
 .../workflow-block/workflow-block.tsx         |  2 +-
 stores/workflow/types.ts                      |  6 ++---
 stores/workflow/workflow-store.ts             | 23 ++++++++++++-------
 9 files changed, 87 insertions(+), 18 deletions(-)
 create mode 100644 app/w/components/workflow-block/sub-block/hooks/use-sub-block-value.ts

diff --git a/app/w/components/workflow-block/sub-block/components/dropdown.tsx b/app/w/components/workflow-block/sub-block/components/dropdown.tsx
index de2458867b..caf4b10200 100644
--- a/app/w/components/workflow-block/sub-block/components/dropdown.tsx
+++ b/app/w/components/workflow-block/sub-block/components/dropdown.tsx
@@ -9,9 +9,16 @@ import {
 interface DropdownProps {
   options: string[]
   defaultValue?: string
+  blockId: string
+  subBlockId: string
 }
 
-export function Dropdown({ options, defaultValue }: DropdownProps) {
+export function Dropdown({
+  options,
+  defaultValue,
+  blockId,
+  subBlockId,
+}: DropdownProps) {
   return (
     <Select defaultValue={defaultValue ?? options[0]}>
       <SelectTrigger>
diff --git a/app/w/components/workflow-block/sub-block/components/long-input.tsx b/app/w/components/workflow-block/sub-block/components/long-input.tsx
index c4f6f5eae5..4afc3e0dbd 100644
--- a/app/w/components/workflow-block/sub-block/components/long-input.tsx
+++ b/app/w/components/workflow-block/sub-block/components/long-input.tsx
@@ -3,9 +3,15 @@ import { useState } from 'react'
 
 interface LongInputProps {
   placeholder?: string
+  blockId: string
+  subBlockId: string
 }
 
-export function LongInput({ placeholder }: LongInputProps) {
+export function LongInput({
+  placeholder,
+  blockId,
+  subBlockId,
+}: LongInputProps) {
   const [value, setValue] = useState('')
 
   return (
diff --git a/app/w/components/workflow-block/sub-block/components/short-input.tsx b/app/w/components/workflow-block/sub-block/components/short-input.tsx
index e7f18751ac..013037a9a1 100644
--- a/app/w/components/workflow-block/sub-block/components/short-input.tsx
+++ b/app/w/components/workflow-block/sub-block/components/short-input.tsx
@@ -4,9 +4,16 @@ import { useState } from 'react'
 interface ShortInputProps {
   placeholder?: string
   password?: boolean
+  blockId: string
+  subBlockId: string
 }
 
-export function ShortInput({ placeholder, password }: ShortInputProps) {
+export function ShortInput({
+  blockId,
+  subBlockId,
+  placeholder,
+  password,
+}: ShortInputProps) {
   const [isFocused, setIsFocused] = useState(false)
   const [value, setValue] = useState('')
 
diff --git a/app/w/components/workflow-block/sub-block/components/slider-input.tsx b/app/w/components/workflow-block/sub-block/components/slider-input.tsx
index b747bd30da..1ebd0b99b7 100644
--- a/app/w/components/workflow-block/sub-block/components/slider-input.tsx
+++ b/app/w/components/workflow-block/sub-block/components/slider-input.tsx
@@ -5,12 +5,16 @@ interface SliderInputProps {
   min?: number
   max?: number
   defaultValue: number
+  blockId: string
+  subBlockId: string
 }
 
 export function SliderInput({
   min = 0,
   max = 100,
   defaultValue,
+  blockId,
+  subBlockId,
 }: SliderInputProps) {
   const [sliderValue, setSliderValue] = useState(defaultValue)
 
diff --git a/app/w/components/workflow-block/sub-block/hooks/use-sub-block-value.ts b/app/w/components/workflow-block/sub-block/hooks/use-sub-block-value.ts
new file mode 100644
index 0000000000..ded4766e62
--- /dev/null
+++ b/app/w/components/workflow-block/sub-block/hooks/use-sub-block-value.ts
@@ -0,0 +1,23 @@
+import { useWorkflowStore } from '@/stores/workflow/workflow-store'
+import { useCallback } from 'react'
+
+export function useSubBlockValue(blockId: string, subBlockId: string) {
+  const value = useWorkflowStore(
+    useCallback(
+      (state) =>
+        (state.blocks[blockId]?.subBlocks[subBlockId]?.value) ?? null,
+      [blockId, subBlockId]
+    )
+  )
+
+  const updateSubBlock = useWorkflowStore((state) => state.updateSubBlock)
+
+  const setValue = useCallback(
+    (newValue: any) => {
+      updateSubBlock(blockId, subBlockId, newValue)
+    },
+    [blockId, subBlockId, updateSubBlock]
+  )
+
+  return [value, setValue] as const
+} 
\ No newline at end of file
diff --git a/app/w/components/workflow-block/sub-block/sub-block.tsx b/app/w/components/workflow-block/sub-block/sub-block.tsx
index 5e89d01137..0fbdef8ad3 100644
--- a/app/w/components/workflow-block/sub-block/sub-block.tsx
+++ b/app/w/components/workflow-block/sub-block/sub-block.tsx
@@ -8,10 +8,11 @@ import { Table } from './components/table'
 import { Code } from './components/code'
 
 interface SubBlockProps {
+  blockId: string
   config: SubBlockConfig
 }
 
-export function SubBlock({ config }: SubBlockProps) {
+export function SubBlock({ blockId, config }: SubBlockProps) {
   const handleMouseDown = (e: React.MouseEvent) => {
     e.stopPropagation()
   }
@@ -21,21 +22,35 @@ export function SubBlock({ config }: SubBlockProps) {
       case 'short-input':
         return (
           <ShortInput
+            blockId={blockId}
+            subBlockId={config.id || ''}
             placeholder={config.placeholder}
             password={config.password}
           />
         )
       case 'long-input':
-        return <LongInput placeholder={config.placeholder} />
+        return (
+          <LongInput
+            blockId={blockId}
+            subBlockId={config.id || ''}
+            placeholder={config.placeholder}
+          />
+        )
       case 'dropdown':
         return (
           <div onMouseDown={handleMouseDown}>
-            <Dropdown options={config.options ?? []} />
+            <Dropdown
+              blockId={blockId}
+              subBlockId={config.id || ''}
+              options={config.options ?? []}
+            />
           </div>
         )
       case 'slider':
         return (
           <SliderInput
+            blockId={blockId}
+            subBlockId={config.id || ''}
             min={config.min}
             max={config.max}
             defaultValue={
diff --git a/app/w/components/workflow-block/workflow-block.tsx b/app/w/components/workflow-block/workflow-block.tsx
index abb9308091..cb1f912a9a 100644
--- a/app/w/components/workflow-block/workflow-block.tsx
+++ b/app/w/components/workflow-block/workflow-block.tsx
@@ -75,7 +75,7 @@ export function WorkflowBlock({ id, type, config, name }: WorkflowBlockProps) {
                   subBlock.layout === 'half' ? 'flex-1' : 'w-full'
                 }`}
               >
-                <SubBlock config={subBlock} />
+                <SubBlock blockId={id} config={subBlock} />
               </div>
             ))}
           </div>
diff --git a/stores/workflow/types.ts b/stores/workflow/types.ts
index 47ff3180a6..d12a0b56fd 100644
--- a/stores/workflow/types.ts
+++ b/stores/workflow/types.ts
@@ -11,11 +11,11 @@ export interface BlockState {
   type: BlockType
   name: string
   position: Position
-  subBlocks: Record<string, SubBlock>
+  subBlocks: Record<string, SubBlockState>
   outputType: OutputType
 }
 
-export interface SubBlock {
+export interface SubBlockState {
   id: string
   type: SubBlockType
   value: string | number | string[][] | null
@@ -35,7 +35,7 @@ export interface WorkflowActions {
     position: Position
   ) => void
   updateBlockPosition: (id: string, position: Position) => void
-  updateBlockInput: (blockId: string, inputId: string, value: any) => void
+  updateSubBlock: (blockId: string, subBlockId: string, subBlock: SubBlockState) => void
   removeBlock: (id: string) => void
   addEdge: (edge: Edge) => void
   removeEdge: (edgeId: string) => void
diff --git a/stores/workflow/workflow-store.ts b/stores/workflow/workflow-store.ts
index 4bd16bd248..64c0e716c7 100644
--- a/stores/workflow/workflow-store.ts
+++ b/stores/workflow/workflow-store.ts
@@ -2,8 +2,7 @@ import { create } from 'zustand'
 import { devtools } from 'zustand/middleware'
 import { Edge } from 'reactflow'
 import { BlockType } from '@/blocks/types/block'
-import { Position } from '@/stores/workflow/types'
-import { WorkflowStore } from './types'
+import { Position, WorkflowStore } from './types'
 import { getBlock } from '@/blocks/configs'
 
 const initialState = {
@@ -23,8 +22,9 @@ export const useWorkflowStore = create<WorkflowStore>()(
 
         const subBlocks: Record<string, any> = {}
         blockConfig.workflow.subBlocks.forEach((subBlock) => {
-          subBlocks[subBlock.id || crypto.randomUUID()] = {
-            id: subBlock.id || crypto.randomUUID(),
+          const subBlockId = subBlock.id || crypto.randomUUID()
+          subBlocks[subBlockId] = {
+            id: subBlockId,
             type: subBlock.type,
             value: null,
           }
@@ -60,7 +60,7 @@ export const useWorkflowStore = create<WorkflowStore>()(
         }))
       },
 
-      updateBlockInput: (blockId: string, inputId: string, value: any) => {
+      updateSubBlock: (blockId: string, subBlockId: string, value: any) => {
         set((state) => ({
           blocks: {
             ...state.blocks,
@@ -68,8 +68,8 @@ export const useWorkflowStore = create<WorkflowStore>()(
               ...state.blocks[blockId],
               subBlocks: {
                 ...state.blocks[blockId].subBlocks,
-                [inputId]: {
-                  ...state.blocks[blockId].subBlocks[inputId],
+                [subBlockId]: {
+                  ...state.blocks[blockId].subBlocks[subBlockId],
                   value,
                 },
               },
@@ -94,7 +94,14 @@ export const useWorkflowStore = create<WorkflowStore>()(
 
       addEdge: (edge: Edge) => {
         set((state) => ({
-          edges: [...state.edges, edge],
+          edges: [
+            ...state.edges,
+            {
+              id: edge.id || crypto.randomUUID(),
+              source: edge.source,
+              target: edge.target,
+            },
+          ],
         }))
       },