@bian-womp/spark-workbench 0.3.11 → 0.3.12

package/lib/cjs/index.cjs

@@ -1387,116 +1387,145 @@ let remoteRunnerCounter = 0;
 class RemoteGraphRunner extends AbstractGraphRunner {
     /**
      * Fetch full registry description from remote and register it locally.
-     * Called automatically on first connection with retry mechanism.
+     * Simplified with straightforward retry loop.
+     * Ensures only one fetch happens at a time, even with concurrent calls.
      */
-    async fetchRegistry(client, attempt = 1) {
-        if (this.registryFetching) {
-            // Already fetching, don't start another fetch
+    async fetchRegistry(client) {
+        if (this.registryFetched) {
             return;
         }
-        this.registryFetching = true;
-        try {
-            const desc = await client.api.describeRegistry();
-            // Register types
-            for (const t of desc.types) {
-                if (t.options) {
-                    this.registry.registerEnum({
-                        id: t.id,
-                        options: t.options,
-                        bakeTarget: t.bakeTarget,
-                    });
-                }
-                else {
-                    if (!this.registry.types.has(t.id)) {
-                        this.registry.registerType({
+        // If already fetching, wait for that fetch to complete
+        if (this.registryFetchPromise) {
+            return this.registryFetchPromise;
+        }
+        // Create promise wrapper and assign atomically to prevent race conditions
+        // This ensures concurrent calls will see the promise and wait for it
+        let promise;
+        this.registryFetchPromise = promise = (async () => {
+            try {
+                await this._doFetchRegistry(client);
+            }
+            finally {
+                // Clear the promise after completion (success or failure)
+                this.registryFetchPromise = undefined;
+            }
+        })();
+        return promise;
+    }
+    /**
+     * Internal method that performs the actual registry fetch.
+     */
+    async _doFetchRegistry(client) {
+        let lastError;
+        for (let attempt = 1; attempt <= this.MAX_REGISTRY_FETCH_ATTEMPTS; attempt++) {
+            try {
+                // Add timeout to registry fetch - if it exceeds 3s, retry
+                let timeoutId;
+                const timeoutPromise = new Promise((_, reject) => {
+                    timeoutId = setTimeout(() => reject(new Error("Registry fetch timeout (3s exceeded)")), this.REGISTRY_FETCH_TIMEOUT_MS);
+                });
+                const fetchPromise = client.api.describeRegistry().finally(() => {
+                    // Clear timeout if request completes first
+                    if (timeoutId) {
+                        clearTimeout(timeoutId);
+                        timeoutId = undefined;
+                    }
+                });
+                const desc = await Promise.race([fetchPromise, timeoutPromise]);
+                // Register types
+                for (const t of desc.types) {
+                    if (t.options) {
+                        this.registry.registerEnum({
                             id: t.id,
-                            displayName: t.displayName,
-                            validate: (_v) => true,
+                            options: t.options,
                             bakeTarget: t.bakeTarget,
                         });
                     }
+                    else {
+                        if (!this.registry.types.has(t.id)) {
+                            this.registry.registerType({
+                                id: t.id,
+                                displayName: t.displayName,
+                                validate: (_v) => true,
+                                bakeTarget: t.bakeTarget,
+                            });
+                        }
+                    }
                 }
-            }
-            // Register categories
-            for (const c of desc.categories || []) {
-                if (!this.registry.categories.has(c.id)) {
-                    // Create placeholder category descriptor
-                    const category = {
-                        id: c.id,
-                        displayName: c.displayName,
-                        createRuntime: () => ({
-                            async onInputsChanged() { },
-                        }),
-                        policy: { asyncConcurrency: "switch" },
-                    };
-                    this.registry.categories.register(category);
+                // Register categories
+                for (const c of desc.categories || []) {
+                    if (!this.registry.categories.has(c.id)) {
+                        // Create placeholder category descriptor
+                        const category = {
+                            id: c.id,
+                            displayName: c.displayName,
+                            createRuntime: () => ({
+                                async onInputsChanged() { },
+                            }),
+                            policy: { asyncConcurrency: "switch" },
+                        };
+                        this.registry.categories.register(category);
+                    }
                 }
-            }
-            // Register coercions
-            for (const c of desc.coercions) {
-                if (c.async) {
-                    this.registry.registerAsyncCoercion(c.from, c.to, async (v) => v, {
-                        nonTransitive: c.nonTransitive,
-                    });
+                // Register coercions
+                for (const c of desc.coercions) {
+                    if (c.async) {
+                        this.registry.registerAsyncCoercion(c.from, c.to, async (v) => v, {
+                            nonTransitive: c.nonTransitive,
+                        });
+                    }
+                    else {
+                        this.registry.registerCoercion(c.from, c.to, (v) => v, {
+                            nonTransitive: c.nonTransitive,
+                        });
+                    }
                 }
-                else {
-                    this.registry.registerCoercion(c.from, c.to, (v) => v, {
-                        nonTransitive: c.nonTransitive,
-                    });
+                // Register nodes
+                for (const n of desc.nodes) {
+                    if (!this.registry.nodes.has(n.id)) {
+                        this.registry.registerNode({
+                            id: n.id,
+                            categoryId: n.categoryId,
+                            displayName: n.displayName,
+                            inputs: n.inputs || {},
+                            outputs: n.outputs || {},
+                            policy: n.policy || {},
+                            impl: () => { },
+                        });
+                    }
                 }
+                this.registryFetched = true;
+                this.emit("registry", this.registry);
+                return;
             }
-            // Register nodes
-            for (const n of desc.nodes) {
-                if (!this.registry.nodes.has(n.id)) {
-                    this.registry.registerNode({
-                        id: n.id,
-                        categoryId: n.categoryId,
-                        displayName: n.displayName,
-                        inputs: n.inputs || {},
-                        outputs: n.outputs || {},
-                        policy: n.policy || {},
-                        impl: () => { },
+            catch (err) {
+                lastError = err instanceof Error ? err : new Error(String(err));
+                if (attempt < this.MAX_REGISTRY_FETCH_ATTEMPTS) {
+                    const delayMs = this.INITIAL_RETRY_DELAY_MS * Math.pow(2, attempt - 1);
+                    console.warn(`Failed to fetch registry (attempt ${attempt}/${this.MAX_REGISTRY_FETCH_ATTEMPTS}), retrying in ${delayMs}ms...`, lastError);
+                    // Emit error event for UI feedback
+                    this.emit("error", {
+                        kind: "registry",
+                        message: `Registry fetch failed (attempt ${attempt}/${this.MAX_REGISTRY_FETCH_ATTEMPTS}), retrying...`,
+                        err: lastError,
+                        attempt,
+                        maxAttempts: this.MAX_REGISTRY_FETCH_ATTEMPTS,
                     });
+                    // Wait before retrying
+                    await new Promise((resolve) => setTimeout(resolve, delayMs));
                 }
             }
-            this.registryFetched = true;
-            this.registryFetching = false;
-            this.emit("registry", this.registry);
-        }
-        catch (err) {
-            this.registryFetching = false;
-            const error = err instanceof Error ? err : new Error(String(err));
-            // Retry with exponential backoff if attempts remaining
-            if (attempt < this.MAX_REGISTRY_FETCH_ATTEMPTS) {
-                const delayMs = this.INITIAL_RETRY_DELAY_MS * Math.pow(2, attempt - 1);
-                console.warn(`Failed to fetch registry (attempt ${attempt}/${this.MAX_REGISTRY_FETCH_ATTEMPTS}), retrying in ${delayMs}ms...`, error);
-                // Emit error event for UI feedback
-                this.emit("error", {
-                    kind: "registry",
-                    message: `Registry fetch failed (attempt ${attempt}/${this.MAX_REGISTRY_FETCH_ATTEMPTS}), retrying...`,
-                    err: error,
-                    attempt,
-                    maxAttempts: this.MAX_REGISTRY_FETCH_ATTEMPTS,
-                });
-                // Retry after delay
-                setTimeout(() => {
-                    this.fetchRegistry(client, attempt + 1).catch(() => {
-                        // Final failure handled below
-                    });
-                }, delayMs);
-            }
-            else {
-                // Max attempts reached, emit final error
-                console.error(`Failed to fetch registry after ${this.MAX_REGISTRY_FETCH_ATTEMPTS} attempts:`, error);
-                this.emit("error", {
-                    kind: "registry",
-                    message: `Failed to fetch registry after ${this.MAX_REGISTRY_FETCH_ATTEMPTS} attempts. Please check your connection and try refreshing.`,
-                    err: error,
-                    attempt: this.MAX_REGISTRY_FETCH_ATTEMPTS,
-                    maxAttempts: this.MAX_REGISTRY_FETCH_ATTEMPTS,
-                });
-            }
         }
+        // Max attempts reached, emit final error and throw
+        console.error(`Failed to fetch registry after ${this.MAX_REGISTRY_FETCH_ATTEMPTS} attempts:`, lastError);
+        this.emit("error", {
+            kind: "registry",
+            message: `Failed to fetch registry after ${this.MAX_REGISTRY_FETCH_ATTEMPTS} attempts. Please check your connection and try refreshing.`,
+            err: lastError,
+            attempt: this.MAX_REGISTRY_FETCH_ATTEMPTS,
+            maxAttempts: this.MAX_REGISTRY_FETCH_ATTEMPTS,
+        });
+        throw lastError;
     }
     /**
      * Build RemoteRuntimeClient config from RemoteExecutionBackend config.
@@ -1551,7 +1580,9 @@ class RemoteGraphRunner extends AbstractGraphRunner {
             return this.clientPromise;
         const backend = this.backend;
         // Create connection promise to prevent concurrent connections
-        this.clientPromise = (async () => {
+        // Create promise wrapper first, then assign immediately before async work starts
+        let promise;
+        this.clientPromise = promise = (async () => {
             // Build client config from backend config
             const clientConfig = this.buildClientConfig(backend);
             // Wrap custom event handler to intercept flow-viewport events and emit viewport event
@@ -1583,23 +1614,18 @@ class RemoteGraphRunner extends AbstractGraphRunner {
             this.client = client;
             this.valueCache.clear();
             this.listenersBound = false;
-            // Auto-fetch registry on first connection (only once)
-            if (!this.registryFetched && !this.registryFetching) {
-                // Log loading state (UI can listen to transport status for loading indication)
+            // Fetch registry before returning (wait for it to complete)
+            // Only fetch if not already fetched (handles concurrent calls)
+            if (!this.registryFetched) {
                 console.info("[RemoteGraphRunner] Loading registry from remote...");
-                this.fetchRegistry(client)
-                    .then(() => {
-                    console.info("[RemoteGraphRunner] Loaded registry from remote");
-                })
-                    .catch((err) => {
-                    console.error("[RemoteGraphRunner] Failed to fetch registry:", err);
-                });
+                await this.fetchRegistry(client);
+                console.info("[RemoteGraphRunner] Loaded registry from remote");
             }
             // Clear promise on success
             this.clientPromise = undefined;
             return client;
         })();
-        return this.clientPromise;
+        return promise;
     }
     constructor(backend) {
         super(backend);
@@ -1607,9 +1633,9 @@ class RemoteGraphRunner extends AbstractGraphRunner {
         this.valueCache = new Map();
         this.listenersBound = false;
         this.registryFetched = false;
-        this.registryFetching = false;
         this.MAX_REGISTRY_FETCH_ATTEMPTS = 3;
         this.INITIAL_RETRY_DELAY_MS = 1000; // 1 second
+        this.REGISTRY_FETCH_TIMEOUT_MS = 3000; // 3 seconds
         // Generate readable ID for this runner instance (e.g., remote-001, remote-002)
         remoteRunnerCounter++;
         this.runnerId = `remote-${String(remoteRunnerCounter).padStart(3, "0")}`;
@@ -2034,7 +2060,6 @@ class RemoteGraphRunner extends AbstractGraphRunner {
         const clientToDispose = this.client;
         this.client = undefined;
         this.registryFetched = false; // Reset so registry is fetched again on reconnect
-        this.registryFetching = false; // Reset fetching state
         if (clientToDispose) {
             try {
                 await clientToDispose.dispose();
@@ -4389,7 +4414,7 @@ function WorkbenchProvider({ wb, runner, overrides, uiVersion, children, }) {
                 console.error("[WorkbenchContext] Error updating graph:", err);
             }
         });
-        const offWbdSetValidation = wb.on("validationChanged", (r) => setValidation(r));
+        const offWbSetValidation = wb.on("validationChanged", (r) => setValidation(r));
         const offWbSelectionChanged = wb.on("selectionChanged", async (sel) => {
             setSelectedNodeId(sel.nodes?.[0]);
             setSelectedEdgeId(sel.edges?.[0]);
@@ -4462,20 +4487,8 @@ function WorkbenchProvider({ wb, runner, overrides, uiVersion, children, }) {
         });
         const offWbError = wb.on("error", add("workbench", "error"));
         // Registry updates: swap registry and refresh graph validation/UI
-        const offRunnerRegistry = runner.on("registry", async (newReg) => {
-            try {
-                wb.setRegistry(newReg);
-                // Trigger a graph update so the UI revalidates with new types/enums/nodes
-                try {
-                    await runner.update(wb.def);
-                }
-                catch {
-                    console.error("Failed to update graph definition after registry changed");
-                }
-            }
-            catch {
-                console.error("Failed to handle registry changed event");
-            }
+        const offRunnerRegistry = runner.on("registry", async (registry) => {
+            wb.setRegistry(registry);
         });
         const offFlowViewport = runner.on("viewport", (event) => {
             const viewport = event.viewport;
@@ -4542,7 +4555,7 @@ function WorkbenchProvider({ wb, runner, overrides, uiVersion, children, }) {
             offWbValidationChanged();
             offWbError();
             offWbGraphChangedForUpdate();
-            offWbdSetValidation();
+            offWbSetValidation();
             offWbSelectionChanged();
             offRunnerRegistry();
             offRunnerTransport();
@@ -5847,14 +5860,9 @@ const WorkbenchCanvasComponent = React.forwardRef((props, ref) => {
     const { nodeTypes, resolveNodeType } = React.useMemo(() => {
         // Build nodeTypes map using UI extension registry
         const custom = new Map(); // Include all types present in registry AND current graph to avoid timing issues
-        const ids = new Set([
-            ...Array.from(wb.registry.nodes.keys()),
-            ...wb.def.nodes.map((n) => n.typeId),
-        ]);
-        for (const typeId of ids) {
-            const renderer = ui.getNodeRenderer(typeId);
-            if (renderer)
-                custom.set(typeId, renderer);
+        const allNodeRenderers = ui.getAllNodeRenderers();
+        for (const typeId of Object.keys(allNodeRenderers)) {
+            custom.set(typeId, allNodeRenderers[typeId]);
         }
         const types = {
             "spark-default": DefaultNode,
@@ -5867,7 +5875,7 @@ const WorkbenchCanvasComponent = React.forwardRef((props, ref) => {
         return { nodeTypes: types, resolveNodeType: resolver };
         // Include uiVersion to recompute when custom renderers are registered
         // Include registryVersion to recompute when registry enums/types change
-    }, [wb, wb.registry, registryVersion, uiVersion, ui]);
+    }, [wb, wb.registry, registryVersion, ui, uiVersion]);
     const edgeTypes = React.useMemo(() => {
         // Use default edge renderer override if registered, otherwise use DefaultEdge
         const customEdgeRenderer = ui.getEdgeRenderer();
@@ -6526,10 +6534,7 @@ function WorkbenchStudioCanvas({ autoScroll, onAutoScrollChange, example, onExam
             return;
         const setInitialGraph = async (d, inputs) => {
             await wb.load(d);
-            try {
-                runner.build(wb.def);
-            }
-            catch { }
+            runner.build(wb.def);
             if (inputs) {
                 for (const [nodeId, map] of Object.entries(inputs)) {
                     runner.setInputs(nodeId, map, { dry: true });
@@ -6573,13 +6578,13 @@ function WorkbenchStudioCanvas({ autoScroll, onAutoScrollChange, example, onExam
         const ex = examples.find((e) => e.id === key) ?? examples[0];
         if (!ex)
             return;
-        const { registry: r, def, inputs } = await ex.load();
+        const { registry, def, inputs } = await ex.load();
         // Keep registry consistent with backend:
         // - For local backend, allow example to provide its own registry
         // - For remote backend, registry is automatically managed by RemoteGraphRunner
         if (backendKind === "local") {
-            if (r) {
-                wb.setRegistry(r);
+            if (registry) {
+                wb.setRegistry(registry);
             }
         }
         await wb.load(def);