diff --git a/src/renderer/src/components/flow/BlocksLibrary.tsx b/src/renderer/src/components/flow/BlocksLibrary.tsx
index 29be26b..e989fec 100644
--- a/src/renderer/src/components/flow/BlocksLibrary.tsx
+++ b/src/renderer/src/components/flow/BlocksLibrary.tsx
@@ -6,7 +6,7 @@ import { X } from 'lucide-react';
 
 const BlocksLibrary = () => {
   const { functionDefinitions, removeDefinition } = useFlowchartStore((state) => ({
-    functionDefinitions: state.functionDefinitions,
+    functionDefinitions: state.functionDefinitionBlocks,
     removeDefinition: state.removeDefinition
   }));
 
diff --git a/src/renderer/src/components/flow/FlowCanvas.tsx b/src/renderer/src/components/flow/FlowCanvas.tsx
index e126318..4989f20 100644
--- a/src/renderer/src/components/flow/FlowCanvas.tsx
+++ b/src/renderer/src/components/flow/FlowCanvas.tsx
@@ -34,20 +34,17 @@ const FlowCanvas = () => {
   const { menu, onPaneClick, onNodeContextMenu } = useContextMenu(reactFlowRef);
 
   const { takeSnapshot } = useUndoRedo();
-  const { edges, onEdgesChange, nodes, onNodesChange, onConnect, addNode, functionDefinitions } =
-    useFlowchartStore(
-      useShallow((state) => ({
-        edges: state.edges,
-        onEdgesChange: state.onEdgesChange,
-        nodes: state.nodes,
-        onNodesChange: state.onNodesChange,
-        onConnect: state.onConnect,
-        addNode: state.addNode,
-        functionDefinitions: state.functionDefinitions
-      }))
-    );
-
-  console.log(functionDefinitions);
+  const { edges, onEdgesChange, nodes, onNodesChange, onConnect, addNode } = useFlowchartStore(
+    useShallow((state) => ({
+      edges: state.edges,
+      onEdgesChange: state.onEdgesChange,
+      nodes: state.nodes,
+      onNodesChange: state.onNodesChange,
+      onConnect: state.onConnect,
+      addNode: state.addNode,
+      saveDefinition: state.saveDefinition
+    }))
+  );
 
   const onNodeDragStart: NodeDragHandler = useCallback(() => {
     // 👇 make dragging a node undoable
diff --git a/src/renderer/src/components/flow/FlowControlsTopRight.tsx b/src/renderer/src/components/flow/FlowControlsTopRight.tsx
index 2e5ba8c..fa60119 100644
--- a/src/renderer/src/components/flow/FlowControlsTopRight.tsx
+++ b/src/renderer/src/components/flow/FlowControlsTopRight.tsx
@@ -4,13 +4,14 @@ import { useFlowchartStore } from '@/stores/flowchart';
 import { useShallow } from 'zustand/react/shallow';
 import { useLifecycleStore } from '@/stores/lifecycle';
 import { trpcClient } from '@/main';
+import { FunctionDefinition, Name } from '@/types/block';
 
 const FlowControlsTopRight = (): JSX.Element => {
-  const { edges, nodes, functionDefinitions } = useFlowchartStore(
+  const { edges, nodes, functionDefinitionBlocks } = useFlowchartStore(
     useShallow((state) => ({
       edges: state.edges,
       nodes: state.nodes,
-      functionDefinitions: state.functionDefinitions
+      functionDefinitionBlocks: state.functionDefinitionBlocks
     }))
   );
 
@@ -33,6 +34,20 @@ const FlowControlsTopRight = (): JSX.Element => {
   };
 
   const onStart = async () => {
+    const functionDefinitions: Record<Name, FunctionDefinition> = {};
+    for (const [name, block] of Object.entries(functionDefinitionBlocks)) {
+      const bodyNodes = nodes.filter((n) => n.parentNode === block.id);
+      const bodyNodeIds = new Set(bodyNodes.map((n) => n.id));
+      bodyNodeIds.add(block.id);
+      const bodyEdges = edges.filter((e) => bodyNodeIds.has(e.target) && bodyNodeIds.has(e.source));
+
+      functionDefinitions[name] = {
+        block,
+        nodes: bodyNodes,
+        edges: bodyEdges
+      };
+    }
+
     await trpcClient.startFlowchart.mutate(
       JSON.stringify({
         event: {
diff --git a/src/renderer/src/stores/flowchart.ts b/src/renderer/src/stores/flowchart.ts
index a8de545..c011d0b 100644
--- a/src/renderer/src/stores/flowchart.ts
+++ b/src/renderer/src/stores/flowchart.ts
@@ -13,7 +13,7 @@ import {
   applyEdgeChanges,
   XYPosition
 } from 'reactflow';
-import { BlockData, FunctionDefinition, Name } from '@/types/block';
+import { BlockData, Name } from '@/types/block';
 
 import { v4 as uuidv4 } from 'uuid';
 import { createJSONStorage, persist } from 'zustand/middleware';
@@ -31,7 +31,7 @@ interface FlowchartState {
   edges: Edge[];
   controls: Node[];
 
-  functionDefinitions: Record<Name, FunctionDefinition>;
+  functionDefinitionBlocks: Record<Name, Node<BlockData>>;
 
   setNodes: (nodes: Node<BlockData>[]) => void;
   setEdges: (edges: Edge[]) => void;
@@ -61,7 +61,7 @@ export const useFlowchartStore = create<FlowchartState>()(
         edges: [] as Edge[],
         controls: [] as Node[],
 
-        functionDefinitions: {},
+        functionDefinitionBlocks: {},
 
         setNodes: (nodes: Node<BlockData>[]) => set({ nodes }),
         setEdges: (edges: Edge[]) => set({ edges }),
@@ -80,34 +80,25 @@ export const useFlowchartStore = create<FlowchartState>()(
         saveDefinition: (definitionBlockId: string) => {
           set((state) => {
             const node = state.nodes.find((n) => n.id === definitionBlockId);
-            if (node === undefined || node.data.block_type !== 'flojoy.intrinsics.function') {
+            if (node?.data.block_type !== 'flojoy.intrinsics.function') {
               return;
             }
-            const bodyNodes = state.nodes.filter((n) => n.parentNode === definitionBlockId);
-            const bodyNodeIds = new Set(bodyNodes.map((n) => n.id));
-            bodyNodeIds.add(definitionBlockId);
-            const bodyEdges = state.edges.filter(
-              (e) => bodyNodeIds.has(e.target) && bodyNodeIds.has(e.source)
-            );
 
-            state.functionDefinitions[node.data.label] = {
-              block: node,
-              nodes: bodyNodes,
-              edges: bodyEdges
-            };
+            state.functionDefinitionBlocks[node.data.label] = node;
           });
         },
 
         removeDefinition: (name: string) => {
           set((state) => {
-            delete state.functionDefinitions[name];
+            delete state.functionDefinitionBlocks[name];
           });
         },
 
-        onNodesChange: (changes: NodeChange[]) =>
+        onNodesChange: (changes: NodeChange[]) => {
           set({
             nodes: applyNodeChanges(changes, get().nodes)
-          }),
+          });
+        },
         onEdgesChange: (changes: EdgeChange[]) => {
           set({
             edges: applyEdgeChanges(changes, get().edges)
@@ -168,7 +159,7 @@ export const useFlowchartStore = create<FlowchartState>()(
               };
               break;
             case 'function_instance': {
-              const definitions = get().functionDefinitions;
+              const definitions = get().functionDefinitionBlocks;
               if (!(payload.name in definitions)) {
                 throw new Error(`Undefined function block ${payload.name}`);
               }
@@ -179,8 +170,8 @@ export const useFlowchartStore = create<FlowchartState>()(
                 block_type: 'function_instance',
                 label: payload.name,
                 intrinsic_parameters: {},
-                inputs: definition.block.data.inputs,
-                outputs: definition.block.data.outputs
+                inputs: definition.data.inputs,
+                outputs: definition.data.outputs
               };
               break;
             }