From 2fb0894c2a370720205e023d1a4a81c29bdefca3 Mon Sep 17 00:00:00 2001
From: Waleed Latif <walif6@gmail.com>
Date: Wed, 28 May 2025 17:19:57 -0700
Subject: [PATCH] fix(persistence): persist collapsed state for responseFormat
 code subblock (#429)

* fix(persistence): persist collapsed state for responseFormat code subblock

* add additional type safety
---
 .../components/sub-block/components/code.tsx   | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/apps/sim/app/w/[id]/components/workflow-block/components/sub-block/components/code.tsx b/apps/sim/app/w/[id]/components/workflow-block/components/sub-block/components/code.tsx
index 4d0ca96c55..27caaf7637 100644
--- a/apps/sim/app/w/[id]/components/workflow-block/components/sub-block/components/code.tsx
+++ b/apps/sim/app/w/[id]/components/workflow-block/components/sub-block/components/code.tsx
@@ -12,6 +12,7 @@ import { checkTagTrigger, TagDropdown } from '@/components/ui/tag-dropdown'
 import { createLogger } from '@/lib/logs/console-logger'
 import { cn } from '@/lib/utils'
 import { useCodeGeneration } from '@/app/w/[id]/hooks/use-code-generation'
+import { useSubBlockStore } from '@/stores/workflows/subblock/store'
 import { CodePromptBar } from '../../../../code-prompt-bar/code-prompt-bar'
 import { useSubBlockValue } from '../hooks/use-sub-block-value'
 
@@ -68,10 +69,21 @@ export function Code({
   const [cursorPosition, setCursorPosition] = useState(0)
   const [activeSourceBlockId, setActiveSourceBlockId] = useState<string | null>(null)
   const [visualLineHeights, setVisualLineHeights] = useState<number[]>([])
-  const [isCollapsed, setIsCollapsed] = useState(false)
+
+  const collapsedStateKey = `${subBlockId}_collapsed`
+  const isCollapsed =
+    (useSubBlockStore((state) => state.getValue(blockId, collapsedStateKey)) as boolean) ?? false
+  const setCollapsedValue = useSubBlockStore((state) => state.setValue)
+
+  const showCollapseButton = subBlockId === 'responseFormat' && code.split('\n').length > 5
 
   const editorRef = useRef<HTMLDivElement>(null)
 
+  // Function to toggle collapsed state
+  const toggleCollapsed = () => {
+    setCollapsedValue(blockId, collapsedStateKey, !isCollapsed)
+  }
+
   // AI Code Generation Hook
   const handleStreamStart = () => {
     setCode('')
@@ -311,11 +323,11 @@ export function Code({
             </Button>
           )}
 
-          {code.split('\n').length > 5 && !isAiStreaming && (
+          {showCollapseButton && !isAiStreaming && (
             <Button
               variant='ghost'
               size='sm'
-              onClick={() => setIsCollapsed(!isCollapsed)}
+              onClick={toggleCollapsed}
               aria-label={isCollapsed ? 'Expand code' : 'Collapse code'}
               className='h-8 px-2 text-muted-foreground hover:text-foreground'
             >