diff --git a/src/lib/ComfyGraph.ts b/src/lib/ComfyGraph.ts
index 776051e..17a0382 100644
--- a/src/lib/ComfyGraph.ts
+++ b/src/lib/ComfyGraph.ts
@@ -32,9 +32,10 @@ export default class ComfyGraph extends LGraph {
     workflowID: WorkflowInstID | null = null;
 
     get workflow(): ComfyWorkflow | null {
-        if (this.workflowID == null)
+        const workflowID = (this.getRootGraph() as ComfyGraph)?.workflowID;
+        if (workflowID == null)
             return null;
-        return workflowState.getWorkflow(this.workflowID)
+        return workflowState.getWorkflow(workflowID)
     }
 
     constructor(workflowID?: WorkflowInstID) {
diff --git a/src/lib/nodes/ComfyConfigureQueuePromptButton.ts b/src/lib/nodes/ComfyConfigureQueuePromptButton.ts
index 1edd66f..c10edd5 100644
--- a/src/lib/nodes/ComfyConfigureQueuePromptButton.ts
+++ b/src/lib/nodes/ComfyConfigureQueuePromptButton.ts
@@ -30,9 +30,9 @@ export default class ComfyConfigureQueuePromptButton extends ComfyGraphNode {
             }
 
             if (typeof param === "string")
-                this.workflow.attrs.queuePromptButtonName = param || ""
+                this.workflow.setAttribute("queuePromptButtonName", param || "")
             else if (typeof param === "object" && "buttonName" in param)
-                this.workflow.attrs.queuePromptButtonName = param.buttonName || ""
+                this.workflow.setAttribute("queuePromptButtonName", param.buttonName || "")
         }
     }
 }
diff --git a/src/lib/stores/workflowState.ts b/src/lib/stores/workflowState.ts
index 0b8d1f2..b119593 100644
--- a/src/lib/stores/workflowState.ts
+++ b/src/lib/stores/workflowState.ts
@@ -99,6 +99,11 @@ export class ComfyWorkflow {
         store.set(get(store));
     }
 
+    setAttribute<K extends keyof WorkflowAttributes>(key: K, value: WorkflowAttributes[K]) {
+        this.attrs[key] = value;
+        this.notifyModified();
+    }
+
     start(key: string, canvas: ComfyGraphCanvas) {
         if (this.canvases[key] != null)
             throw new Error(`This workflow is already being displayed on canvas ${key}`)
@@ -176,6 +181,7 @@ export class ComfyWorkflow {
         }
 
         this.graph.configure(data.graph);
+        this.graph.workflowID = this.id;
 
         for (const node of this.graph._nodes) {
             const size = node.computeSize();
@@ -228,9 +234,10 @@ function getWorkflow(id: WorkflowInstID): ComfyWorkflow | null {
 }
 
 function getWorkflowByGraph(graph: LGraph): ComfyWorkflow | null {
-    if ("workflowID" in graph && graph.workflowID != null)
-        return getWorkflow((graph as ComfyGraph).workflowID);
-    return null;
+    const workflowID = (graph.getRootGraph() as ComfyGraph)?.workflowID;
+    if (workflowID == null)
+        return null;
+    return getWorkflow(workflowID);
 }
 
 function getWorkflowByNode(node: LGraphNode): ComfyWorkflow | null {