diff --git a/web/app/components/sub-graph/components/sub-graph-children.tsx b/web/app/components/sub-graph/components/sub-graph-children.tsx
index 618d0c9e20..b17472ff7d 100644
--- a/web/app/components/sub-graph/components/sub-graph-children.tsx
+++ b/web/app/components/sub-graph/components/sub-graph-children.tsx
@@ -27,29 +27,13 @@ const SubGraphChildren: FC<SubGraphChildrenProps> = ({
   const isChatMode = useIsChatMode()
   const nodePanelWidth = useStore(s => s.nodePanelWidth)
 
-  const { selectedNode, nodes } = useReactFlowStore(useShallow((s) => {
-    const nodes = s.getNodes()
-    const currentNode = nodes.find(node => node.data.selected)
-
-    if (currentNode?.data.type === BlockEnum.LLM) {
-      return {
-        selectedNode: {
-          id: currentNode.id,
-          type: currentNode.type,
-          data: currentNode.data,
-        },
-        nodes,
-      }
-    }
-    return {
-      selectedNode: null,
-      nodes,
-    }
+  const selectedNode = useReactFlowStore(useShallow((s) => {
+    return s.getNodes().find(node => node.data.selected)
   }))
 
-  const extractorNode = useMemo(() => {
-    return nodes.find(node => node.data.type === BlockEnum.LLM)
-  }, [nodes])
+  const extractorNode = useReactFlowStore(useShallow((s) => {
+    return s.getNodes().find(node => node.data.type === BlockEnum.LLM)
+  }))
 
   const availableNodes = useMemo(() => {
     return extractorNode ? [extractorNode] : []
diff --git a/web/app/components/sub-graph/index.tsx b/web/app/components/sub-graph/index.tsx
index 79f13f0161..6a4aef5537 100644
--- a/web/app/components/sub-graph/index.tsx
+++ b/web/app/components/sub-graph/index.tsx
@@ -12,7 +12,7 @@ import SubGraphMain from './components/sub-graph-main'
 import { useSubGraphNodes } from './hooks'
 import { createSubGraphSlice } from './store'
 
-const SUB_GRAPH_EDGE_GAP = 180
+const SUB_GRAPH_EDGE_GAP = 160
 const SUB_GRAPH_ENTRY_POSITION = {
   x: START_INITIAL_POSITION.x,
   y: 150,
@@ -28,7 +28,7 @@ const defaultViewport: Viewport = {
   zoom: 1.3,
 }
 
-const SubGraph: FC<SubGraphProps> = (props) => {
+const SubGraphContent: FC<SubGraphProps> = (props) => {
   const {
     toolNodeId,
     paramKey,
@@ -186,14 +186,14 @@ const SubGraph: FC<SubGraphProps> = (props) => {
   )
 }
 
-const SubGraphWrapper: FC<SubGraphProps> = (props) => {
+const SubGraph: FC<SubGraphProps> = (props) => {
   return (
     <WorkflowContextProvider
-      injectWorkflowStoreSliceFn={createSubGraphSlice as unknown as InjectWorkflowStoreSliceFn}
+      injectWorkflowStoreSliceFn={createSubGraphSlice as InjectWorkflowStoreSliceFn}
     >
-      <SubGraph {...props} />
+      <SubGraphContent {...props} />
     </WorkflowContextProvider>
   )
 }
 
-export default memo(SubGraphWrapper)
+export default memo(SubGraph)
diff --git a/web/app/components/workflow/store/workflow/index.ts b/web/app/components/workflow/store/workflow/index.ts
index c2c0c00201..ff67ff29a5 100644
--- a/web/app/components/workflow/store/workflow/index.ts
+++ b/web/app/components/workflow/store/workflow/index.ts
@@ -15,6 +15,7 @@ import type { VersionSliceShape } from './version-slice'
 import type { WorkflowDraftSliceShape } from './workflow-draft-slice'
 import type { WorkflowSliceShape } from './workflow-slice'
 import type { RagPipelineSliceShape } from '@/app/components/rag-pipeline/store'
+import type { SubGraphSliceShape } from '@/app/components/sub-graph/types'
 import type { WorkflowSliceShape as WorkflowAppSliceShape } from '@/app/components/workflow-app/store/workflow/workflow-slice'
 import { useContext } from 'react'
 import {
@@ -30,7 +31,6 @@ import { createHelpLineSlice } from './help-line-slice'
 import { createHistorySlice } from './history-slice'
 import { createLayoutSlice } from './layout-slice'
 import { createNodeSlice } from './node-slice'
-
 import { createPanelSlice } from './panel-slice'
 import { createToolSlice } from './tool-slice'
 import { createVersionSlice } from './version-slice'
@@ -40,6 +40,7 @@ import { createWorkflowSlice } from './workflow-slice'
 export type SliceFromInjection
   = Partial<WorkflowAppSliceShape>
     & Partial<RagPipelineSliceShape>
+    & Partial<SubGraphSliceShape>
 
 export type Shape
   = ChatVariableSliceShape