@@ -12,6 +12,7 @@ import { checkTagTrigger, TagDropdown } from '@/components/ui/tag-dropdown'
1212import { createLogger } from '@/lib/logs/console-logger'
1313import { cn } from '@/lib/utils'
1414import { useCodeGeneration } from '@/app/w/[id]/hooks/use-code-generation'
15+ import { useSubBlockStore } from '@/stores/workflows/subblock/store'
1516import { CodePromptBar } from '../../../../code-prompt-bar/code-prompt-bar'
1617import { useSubBlockValue } from '../hooks/use-sub-block-value'
1718
@@ -68,10 +69,21 @@ export function Code({
6869 const [ cursorPosition , setCursorPosition ] = useState ( 0 )
6970 const [ activeSourceBlockId , setActiveSourceBlockId ] = useState < string | null > ( null )
7071 const [ visualLineHeights , setVisualLineHeights ] = useState < number [ ] > ( [ ] )
71- const [ isCollapsed , setIsCollapsed ] = useState ( false )
72+
73+ const collapsedStateKey = `${ subBlockId } _collapsed`
74+ const isCollapsed =
75+ ( useSubBlockStore ( ( state ) => state . getValue ( blockId , collapsedStateKey ) ) as boolean ) ?? false
76+ const setCollapsedValue = useSubBlockStore ( ( state ) => state . setValue )
77+
78+ const showCollapseButton = subBlockId === 'responseFormat' && code . split ( '\n' ) . length > 5
7279
7380 const editorRef = useRef < HTMLDivElement > ( null )
7481
82+ // Function to toggle collapsed state
83+ const toggleCollapsed = ( ) => {
84+ setCollapsedValue ( blockId , collapsedStateKey , ! isCollapsed )
85+ }
86+
7587 // AI Code Generation Hook
7688 const handleStreamStart = ( ) => {
7789 setCode ( '' )
@@ -311,11 +323,11 @@ export function Code({
311323 </ Button >
312324 ) }
313325
314- { code . split ( '\n' ) . length > 5 && ! isAiStreaming && (
326+ { showCollapseButton && ! isAiStreaming && (
315327 < Button
316328 variant = 'ghost'
317329 size = 'sm'
318- onClick = { ( ) => setIsCollapsed ( ! isCollapsed ) }
330+ onClick = { toggleCollapsed }
319331 aria-label = { isCollapsed ? 'Expand code' : 'Collapse code' }
320332 className = 'h-8 px-2 text-muted-foreground hover:text-foreground'
321333 >
0 commit comments