From b91415dd7231ebcf3d8f55365f6ce0c694397cc1 Mon Sep 17 00:00:00 2001
From: Laila Los <44241786+ElectronicBlueberry@users.noreply.github.com>
Date: Fri, 14 Feb 2025 16:02:36 +0100
Subject: [PATCH] safely dispose store by reference counting

---
 client/src/composables/timeoutStoreDispose.ts | 30 +++++++++++++++++++
 client/src/composables/workflowStores.ts      | 21 +++++++++----
 2 files changed, 45 insertions(+), 6 deletions(-)
 create mode 100644 client/src/composables/timeoutStoreDispose.ts

diff --git a/client/src/composables/timeoutStoreDispose.ts b/client/src/composables/timeoutStoreDispose.ts
new file mode 100644
index 00000000000..c9adcacdd57
--- /dev/null
+++ b/client/src/composables/timeoutStoreDispose.ts
@@ -0,0 +1,30 @@
+import type { Store } from "pinia";
+
+const referenceCountByStore = new WeakMap<Store, number>();
+
+/**
+ * Reference counts a store, and provides a function to safely dispose it after a timeout.
+ * @param store store to reference count and dispose
+ * @param timeout how long to wait before reference counting and attempting a dispose.
+ * @returns dispose function
+ */
+export function useTimeoutStoreDispose(store: Store, timeout = 1000) {
+    const currentReferenceCount = referenceCountByStore.get(store) ?? 0;
+    referenceCountByStore.set(store, currentReferenceCount + 1);
+
+    const disposeIfReferenceFree = () => {
+        const referenceCount = referenceCountByStore.get(store) ?? 0;
+
+        if (referenceCount <= 0) {
+            store.$dispose();
+        }
+    };
+
+    const dispose = () => {
+        const referenceCount = referenceCountByStore.get(store) ?? 1;
+        referenceCountByStore.set(store, referenceCount - 1);
+        setTimeout(disposeIfReferenceFree, timeout);
+    };
+
+    return dispose;
+}
diff --git a/client/src/composables/workflowStores.ts b/client/src/composables/workflowStores.ts
index 63704068e3c..1d68e516254 100644
--- a/client/src/composables/workflowStores.ts
+++ b/client/src/composables/workflowStores.ts
@@ -7,6 +7,8 @@ import { useWorkflowStateStore } from "@/stores/workflowEditorStateStore";
 import { useWorkflowEditorToolbarStore } from "@/stores/workflowEditorToolbarStore";
 import { useWorkflowStepStore } from "@/stores/workflowStepStore";
 
+import { useTimeoutStoreDispose } from "./timeoutStoreDispose";
+
 /**
  * Creates stores scoped to a specific workflowId, and manages their lifetime.
  * In child components, use `useWorkflowStores` instead.
@@ -30,13 +32,20 @@ export function provideScopedWorkflowStores(workflowId: Ref<string> | string) {
     const toolbarStore = useWorkflowEditorToolbarStore(workflowId.value);
     const undoRedoStore = useUndoRedoStore(workflowId.value);
 
+    const disposeConnectionStore = useTimeoutStoreDispose(connectionStore);
+    const disposeStateStore = useTimeoutStoreDispose(stateStore);
+    const disposeStepStore = useTimeoutStoreDispose(stepStore);
+    const disposeCommentStore = useTimeoutStoreDispose(commentStore);
+    const disposeToolbarStore = useTimeoutStoreDispose(toolbarStore);
+    const disposeUndoRedoStore = useTimeoutStoreDispose(undoRedoStore);
+
     onScopeDispose(() => {
-        connectionStore.$dispose();
-        stateStore.$dispose();
-        stepStore.$dispose();
-        commentStore.$dispose();
-        toolbarStore.$dispose();
-        undoRedoStore.$dispose();
+        disposeConnectionStore();
+        disposeStateStore();
+        disposeStepStore();
+        disposeCommentStore();
+        disposeToolbarStore();
+        disposeUndoRedoStore();
     });
 
     return {