From 483fa354bfe0b6a9df12c07ffab0442b10f8d89b Mon Sep 17 00:00:00 2001
From: Emir Karabeg <emirkarabeg@berkeley.edu>
Date: Thu, 16 Jan 2025 17:26:48 -0800
Subject: [PATCH] Created code subblock state

---
 .../sub-block/components/code.tsx             | 30 +++++++++++++++----
 .../workflow-block/sub-block/sub-block.tsx    |  2 +-
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/app/w/components/workflow-block/sub-block/components/code.tsx b/app/w/components/workflow-block/sub-block/components/code.tsx
index 8dfada74c1..5454d22d01 100644
--- a/app/w/components/workflow-block/sub-block/components/code.tsx
+++ b/app/w/components/workflow-block/sub-block/components/code.tsx
@@ -1,5 +1,6 @@
 import { useState, useRef, useEffect, useCallback, useMemo } from 'react'
 import { cn } from '@/lib/utils'
+import { useSubBlockValue } from '../hooks/use-sub-block-value'
 
 interface CodeLine {
   id: string
@@ -25,15 +26,32 @@ const MATCHING_PAIRS: Record<string, string> = {
   '`': '`',
 }
 
-function useCodeLines() {
-  const [lines, setLines] = useState<CodeLine[]>(() =>
-    Array(INITIAL_LINES)
+interface CodeProps {
+  blockId: string
+  subBlockId: string
+}
+
+function useCodeLines(blockId: string, subBlockId: string) {
+  const [value, setValue] = useSubBlockValue(blockId, subBlockId)
+
+  // Initialize lines from store value or create initial state
+  const [lines, setLines] = useState<any[]>(() => {
+    if (Array.isArray(value)) {
+      return value
+    }
+    return Array(INITIAL_LINES)
       .fill(null)
       .map(() => ({
         id: crypto.randomUUID(),
         content: '',
       }))
-  )
+  })
+
+  // Update store whenever lines change
+  useEffect(() => {
+    setValue(lines)
+  }, [lines, setValue])
+
   const [currentLine, setCurrentLine] = useState(0)
   const [selection, setSelection] = useState<SelectionState | null>(null)
 
@@ -65,7 +83,7 @@ function useCodeLines() {
   }
 }
 
-export function Code() {
+export function Code({ blockId, subBlockId }: CodeProps) {
   const {
     lines,
     setLines,
@@ -74,7 +92,7 @@ export function Code() {
     handleChange,
     selection,
     setSelection,
-  } = useCodeLines()
+  } = useCodeLines(blockId, subBlockId)
   const textareaRefs = useRef<(HTMLTextAreaElement | null)[]>([])
   const scrollContainerRef = useRef<HTMLDivElement>(null)
   const contentWrapperRef = useRef<HTMLDivElement>(null)
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 2828777eb0..aaf97f2660 100644
--- a/app/w/components/workflow-block/sub-block/sub-block.tsx
+++ b/app/w/components/workflow-block/sub-block/sub-block.tsx
@@ -67,7 +67,7 @@ export function SubBlock({ blockId, config }: SubBlockProps) {
           />
         )
       case 'code':
-        return <Code />
+        return <Code blockId={blockId} subBlockId={config.id} />
       default:
         return null
     }