diff --git a/frontend/.changeset/calm-books-occur.md b/frontend/.changeset/calm-books-occur.md
new file mode 100644
index 000000000..ebd633e08
--- /dev/null
+++ b/frontend/.changeset/calm-books-occur.md
@@ -0,0 +1,6 @@
+---
+"@liam-hq/erd-core": patch
+"@liam-hq/cli": patch
+---
+
+refactor: Update handleLayout to accept nodes and improve hidden node handling
diff --git a/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/Toolbar/TidyUpButton/TidyUpButton.tsx b/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/Toolbar/TidyUpButton/TidyUpButton.tsx
index abb2dba2f..fd21432e1 100644
--- a/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/Toolbar/TidyUpButton/TidyUpButton.tsx
+++ b/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/Toolbar/TidyUpButton/TidyUpButton.tsx
@@ -3,10 +3,12 @@ import { useCliVersion } from '@/providers'
 import { useUserEditingStore } from '@/stores'
 import { IconButton, TidyUpIcon } from '@liam-hq/ui'
 import { ToolbarButton } from '@radix-ui/react-toolbar'
+import { useReactFlow } from '@xyflow/react'
 import { type FC, useCallback } from 'react'
 import { useAutoLayout } from '../../useAutoLayout'
 
 export const TidyUpButton: FC = () => {
+  const { getNodes } = useReactFlow()
   const { handleLayout } = useAutoLayout()
   const { showMode } = useUserEditingStore()
   const { cliVersion } = useCliVersion()
@@ -16,8 +18,8 @@ export const TidyUpButton: FC = () => {
       showMode,
       cliVer: cliVersion.version,
     })
-    handleLayout()
-  }, [handleLayout, showMode, cliVersion.version])
+    handleLayout(getNodes())
+  }, [handleLayout, showMode, getNodes, cliVersion.version])
 
   return (
     <ToolbarButton asChild>
diff --git a/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useAutoLayout/useAutoLayout.ts b/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useAutoLayout/useAutoLayout.ts
index 73c5fa947..360354751 100644
--- a/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useAutoLayout/useAutoLayout.ts
+++ b/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useAutoLayout/useAutoLayout.ts
@@ -5,25 +5,21 @@ import { useERDContentContext } from '../ERDContentContext'
 import { getElkLayout } from './getElkLayout'
 
 export const useAutoLayout = () => {
-  const { getNodes, setNodes, getEdges, fitView } = useReactFlow()
+  const { setNodes, getEdges, fitView } = useReactFlow()
   const {
     actions: { setLoading, setInitializeComplete },
   } = useERDContentContext()
 
   const handleLayout = useCallback(
-    async (
-      fitViewOptions: FitViewOptions = {},
-      initialHiddenNodeIds: string[] = [],
-    ) => {
+    async (nodes: Node[], fitViewOptions: FitViewOptions = {}) => {
       setLoading(true)
-      const nodes = getNodes()
       const edges = getEdges()
 
       const hiddenNodes: Node[] = []
       const visibleNodes: Node[] = []
       for (const node of nodes) {
-        if (initialHiddenNodeIds.includes(node.id) || node.hidden) {
-          hiddenNodes.push({ ...node, hidden: true })
+        if (node.hidden) {
+          hiddenNodes.push(node)
         } else {
           visibleNodes.push(node)
         }
@@ -47,7 +43,7 @@ export const useAutoLayout = () => {
         setInitializeComplete(true)
       }, 0)
     },
-    [getNodes, setNodes, getEdges, fitView, setLoading, setInitializeComplete],
+    [setNodes, getEdges, fitView, setLoading, setInitializeComplete],
   )
 
   return { handleLayout }
diff --git a/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useInitialAutoLayout.ts b/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useInitialAutoLayout.ts
index f91805b94..281fc90da 100644
--- a/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useInitialAutoLayout.ts
+++ b/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useInitialAutoLayout.ts
@@ -49,16 +49,20 @@ export const useInitialAutoLayout = (nodes: Node[]) => {
       updateActiveTableName(tableNameFromUrl)
       const hiddenNodeIds = await getHiddenNodeIdsFromUrl()
       addHiddenNodeIds(hiddenNodeIds)
+      const appliedNodes = nodes.map((node) => ({
+        ...node,
+        hidden: hiddenNodeIds.includes(node.id),
+      }))
 
       const fitViewOptions = tableNameFromUrl
         ? { maxZoom: 1, duration: 300, nodes: [{ id: tableNameFromUrl }] }
         : undefined
 
       if (tableNodesInitialized) {
-        handleLayout(fitViewOptions, hiddenNodeIds)
+        handleLayout(appliedNodes, fitViewOptions)
       }
     }
 
     initialize()
-  }, [tableNodesInitialized, initializeComplete, handleLayout])
+  }, [tableNodesInitialized, initializeComplete, handleLayout, nodes])
 }