@kiberon-labs/behave-graph-flow 2.0.0 → 3.0.0

package/src/plugin/docs/panel/styles.module.css

@@ -10,7 +10,7 @@
 
 .header {
     padding: 1rem;
-    border-bottom: 1px solid var(--vscode-input-border, var(--colors-borderSubtle));
+    border-bottom: 1px solid var(--ds-input-border, var(--colors-borderSubtle));
     background: var(--colors-bgPanel);
 }
 
@@ -28,7 +28,7 @@
     align-items: center;
     justify-content: center;
     border-radius: 0.5rem;
-    border: 1px solid var(--vscode-input-border, var(--colors-borderSubtle));
+    border: 1px solid var(--ds-input-border, var(--colors-borderSubtle));
     background: var(--colors-bgCanvas);
     flex-shrink: 0;
 }
@@ -62,8 +62,8 @@
     font-size: 0.75rem;
     padding: 0.25rem 0.5rem;
     border-radius: 0.25rem;
-    background: var(--vscode-badge-background, #4d4d4d);
-    color: var(--vscode-badge-foreground, #ffffff);
+    background: var(--ds-badge-bg, #4d4d4d);
+    color: var(--ds-badge-fg, #ffffff);
 }
 
 .content {
@@ -77,8 +77,8 @@
     line-height: 1.5;
     margin-bottom: 1.5rem;
     padding: 0.75rem;
-    background: var(--vscode-textBlockQuote-background, rgba(127, 127, 127, 0.1));
-    border-left: 4px solid var(--vscode-textBlockQuote-border, #007acc);
+    background: var(--ds-blockquote-bg, rgba(127, 127, 127, 0.1));
+    border-left: 4px solid var(--ds-blockquote-border, #007acc);
     border-radius: 0.25rem;
 }
 
@@ -98,7 +98,7 @@
     font-weight: 600;
     margin-bottom: 0.75rem;
     padding-bottom: 0.25rem;
-    border-bottom: 1px solid var(--vscode-input-border, var(--colors-borderSubtle));
+    border-bottom: 1px solid var(--ds-input-border, var(--colors-borderSubtle));
 }
 
 .socketList {
@@ -114,7 +114,7 @@
     padding: 0.5rem;
     background: var(--colors-bgPanel);
     border-radius: 0.25rem;
-    border: 1px solid var(--vscode-input-border, var(--colors-borderSubtle));
+    border: 1px solid var(--ds-input-border, var(--colors-borderSubtle));
 }
 
 .socketIcon {

package/src/plugin/graphrunner/actions.ts

@@ -1,17 +1,88 @@
-import type { System } from '../../system/system';
-import type { StoreApi } from 'zustand';
-import type { GraphRunnerClientStore } from './store';
+import type { GraphSession } from '@/system/graphSession';
 import type { GraphRunnerClient } from './client';
+import type { GraphRunner } from './runner';
 import { executing } from '@/annotations';
-import { sleep } from '@/util/sleep';
 import { type ValueJSON } from '@kiberon-labs/behave-graph';
+import type { TraceBatchEvent } from './types';
 
-// Helper to clear executing state from all nodes
-async function clearAllExecutingStates(system: System) {
-  await sleep(1); // Delay to allow any final traces to process
-  system.nodeStore.getState().setNodes((nodes) =>
-    nodes.map((node) => {
+/**
+ * Per-session batcher for the `executing` node annotation.
+ *
+ * Trace events arrive per node execution (start + end), which at display-rate
+ * ticking means hundreds of events per frame. Writing the annotation straight
+ * through did a full O(nodes) copy-map of the node array per event; instead,
+ * accumulate the net executing state here and apply it once per animation
+ * frame with a single identity-preserving setNodes pass.
+ */
+type ExecutingBatch = {
+  pending: Map<string, boolean>;
+  scheduled: boolean;
+};
+
+const executingBatches = new WeakMap<GraphSession, ExecutingBatch>();
+
+const scheduleFrame = (cb: () => void): void => {
+  if (typeof requestAnimationFrame === 'function') {
+    requestAnimationFrame(cb);
+  } else {
+    setTimeout(cb, 16);
+  }
+};
+
+function flushExecutingState(session: GraphSession, batch: ExecutingBatch) {
+  batch.scheduled = false;
+  if (batch.pending.size === 0) return;
+  const updates = batch.pending;
+  batch.pending = new Map();
+
+  session.nodeStore.getState().setNodes((nodes) => {
+    let changed = false;
+    const next = nodes.map((node) => {
+      if (!('data' in node)) return node;
+      const target = updates.get(node.id);
+      if (target === undefined) return node;
+      if (Boolean(node.data.annotations?.[executing]) === target) return node;
+      changed = true;
+      return {
+        ...node,
+        data: {
+          ...node.data,
+          annotations: {
+            ...node.data.annotations,
+            [executing]: target
+          }
+        }
+      };
+    });
+    // Keep the original array identity when nothing changed so selector-based
+    // subscribers (the React Flow canvas) skip the re-render entirely.
+    return changed ? next : nodes;
+  });
+}
+
+function markExecuting(session: GraphSession, nodeId: string, state: boolean) {
+  let batch = executingBatches.get(session);
+  if (!batch) {
+    batch = { pending: new Map(), scheduled: false };
+    executingBatches.set(session, batch);
+  }
+  batch.pending.set(nodeId, state);
+  if (!batch.scheduled) {
+    batch.scheduled = true;
+    scheduleFrame(() => flushExecutingState(session, batch));
+  }
+}
+
+// Clear executing state from all nodes of a session, dropping any queued
+// per-node updates so a pending flush can't re-highlight after the run ended.
+function clearAllExecutingStates(session: GraphSession) {
+  const batch = executingBatches.get(session);
+  if (batch) batch.pending.clear();
+  session.nodeStore.getState().setNodes((nodes) => {
+    let changed = false;
+    const next = nodes.map((node) => {
       if ('data' in node && node.data.annotations?.[executing]) {
+        changed = true;
         return {
           ...node,
           data: {
@@ -24,95 +95,92 @@ async function clearAllExecutingStates(system: System) {
         };
       }
       return node;
-    })
-  );
+    });
+    return changed ? next : nodes;
+  });
 }
 
+/** Apply one trace event to the session's trace spans + executing annotation. */
+function processTraceEvent(session: GraphSession, ev: TraceBatchEvent) {
+  const traceStore = session.traceStore.getState();
+  if (ev.event === 'start') {
+    let name = ev.nodeId;
+    if (ev.data && typeof ev.data === 'object' && 'typeName' in ev.data) {
+      const typeName = (ev.data as { typeName?: unknown }).typeName;
+      if (typeof typeName === 'string') {
+        name = typeName;
+      }
+    }
+    traceStore.addSpan({
+      nodeId: ev.nodeId,
+      name,
+      // `??` not `||`: the worker sends run-relative ms, so 0 is a valid (and
+      // common, for the first node) start , `||` fell back to the main thread's
+      // performance.now(), producing huge, wrong-clock timestamps.
+      start: ev.timestamp ?? performance.now(),
+      // Open span: NaN until the matching `end` event arrives. The store/render
+      // treat NaN as "still running"; a literal end let it render mis-sized.
+      end: Number.NaN
+      // lane omitted: let the store allocate/free lanes so concurrent spans
+      // stack instead of all piling into lane 0.
+    });
+    markExecuting(session, ev.nodeId, true);
+  } else if (ev.event === 'end') {
+    traceStore.updateSpan(ev.nodeId, {
+      end: ev.timestamp ?? performance.now()
+    });
+    markExecuting(session, ev.nodeId, false);
+  }
+}
+
+/**
+ * Clients that already have listeners attached. Guards against double-wiring when
+ * both the plugin's `runner.connect()` and a host (e.g. the webworker runner)
+ * call {@link setupClientEventListeners} on the same client , which would record
+ * every trace span, log and event twice.
+ */
+const wiredClients = new WeakSet<GraphRunnerClient>();
+
 /**
- * Setup persistent event listeners on the client for trace, logs, and run lifecycle
- * This should be called once when the client is connected
+ * Setup persistent event listeners on the shared client. Registered once when
+ * the client connects; every message is routed to the session that started its
+ * run (via {@link GraphRunner.runIndex}), so concurrent graphs stay isolated.
+ *
+ * Idempotent per client: calling it again with the same client is a no-op.
  */
 export function setupClientEventListeners(
   client: GraphRunnerClient,
-  system: System,
-  store: StoreApi<GraphRunnerClientStore>
+  runner: GraphRunner
 ) {
-  const { setIsExecuting, setCurrentRunId, setCurrentGraphId, setIsPaused } =
-    store.getState();
-
-  // Listen for trace events - these apply to all runs
-  client.on('trace', async (message) => {
-    const traceStore = system.traceStore.getState();
-    if (message.event === 'start') {
-      let name = message.nodeId;
-      if (
-        message.data &&
-        typeof message.data === 'object' &&
-        'typeName' in message.data
-      ) {
-        const typeName = (message.data as { typeName?: unknown }).typeName;
-        if (typeof typeName === 'string') {
-          name = typeName;
-        }
-      }
-      traceStore.addSpan({
-        nodeId: message.nodeId,
-        name,
-        start: message.timestamp || performance.now(),
-        end: 1,
-        lane: 0
-      });
-
-      // Mark node as executing
-      system.nodeStore.getState().setNodes((nodes) =>
-        nodes.map((node) =>
-          node.id === message.nodeId && 'data' in node
-            ? {
-                ...node,
-                data: {
-                  ...node.data,
-                  annotations: {
-                    ...node.data.annotations,
-                    [executing]: true
-                  }
-                }
-              }
-            : node
-        )
-      );
-    } else if (message.event === 'end') {
-      traceStore.updateSpan(message.nodeId, {
-        end: message.timestamp || performance.now()
-      });
-
-      //Delay to allow UI to show executing state
-      await sleep(1);
-
-      // Mark node as no longer executing
-      system.nodeStore.getState().setNodes((nodes) =>
-        nodes.map((node) =>
-          node.id === message.nodeId && 'data' in node
-            ? {
-                ...node,
-                data: {
-                  ...node.data,
-                  annotations: {
-                    ...node.data.annotations,
-                    [executing]: false
-                  }
-                }
-              }
-            : node
-        )
-      );
+  if (wiredClients.has(client)) return;
+  wiredClients.add(client);
+
+  // Resolve the session that owns a given run id, or null if unknown.
+  const sessionFor = (runId: string): GraphSession | null =>
+    runner.runIndex.get(runId)?.session ?? null;
+
+  // Batched trace events (one message per flush window, many events inside).
+  client.on('traceBatch', (message) => {
+    const session = sessionFor(message.runId);
+    if (!session) return;
+    for (const ev of message.events) {
+      processTraceEvent(session, ev);
     }
   });
 
+  // Single trace events , kept for remote servers that predate `traceBatch`.
+  client.on('trace', (message) => {
+    const session = sessionFor(message.runId);
+    if (!session) return;
+    processTraceEvent(session, message);
+  });
+
   // Listen for log messages
   client.on('log', (message) => {
-    const logStore = system.logsStore.getState();
+    const session = sessionFor(message.runId);
+    if (!session) return;
     const formattedMessage = `[${message.runId}/${message.graphId}] ${message.message}${message.data !== undefined ? ` ${JSON.stringify(message.data)}` : ''}`;
-    logStore.append({
+    session.logsStore.getState().append({
       time: new Date(),
       data: {
         message: formattedMessage
@@ -121,133 +189,138 @@ export function setupClientEventListeners(
     });
   });
 
-  // Listen for variable change events from server
+  // Listen for variable change events from server. A tick-driven graph writes
+  // variables every frame, so coalesce to the latest value per variable and
+  // apply once per animation frame instead of one store write per change.
+  const pendingVariableUpdates = new Map<GraphSession, Map<string, unknown>>();
+  let variableFlushScheduled = false;
+
+  const flushVariableUpdates = () => {
+    variableFlushScheduled = false;
+    for (const [session, updates] of pendingVariableUpdates) {
+      const variableStore = session.variableStore.getState();
+      for (const [id, newValue] of updates) {
+        const existingVariable = variableStore.variables[id];
+        if (existingVariable) {
+          variableStore.setVariable(id, {
+            ...existingVariable,
+            initialValue: newValue as ValueJSON
+          });
+        } else {
+          const inferredType = typeof newValue;
+          variableStore.setVariable(id, {
+            id,
+            name: id,
+            valueTypeName: inferredType === 'object' ? 'string' : inferredType,
+            initialValue: newValue as ValueJSON
+          });
+        }
+      }
+    }
+    pendingVariableUpdates.clear();
+  };
+
   client.on('variableChanged', (message) => {
-    const variableStore = system.variableStore.getState();
-    const id = message.variableName;
-
-    // Get existing variable or create new one
-    const existingVariable = variableStore.variables[id];
-
-    if (existingVariable) {
-      // Update existing variable
-      system.variableStore.getState().setVariable(id, {
-        ...existingVariable,
-        initialValue: message.newValue as ValueJSON
-      });
-    } else {
-      // Create new variable if it doesn't exist
-      const inferredType = typeof message.newValue;
-      const newVariable = {
-        id,
-        name: message.variableName,
-        valueTypeName: inferredType === 'object' ? 'string' : inferredType,
-        initialValue: message.newValue as ValueJSON
-      };
-      variableStore.setVariable(id, newVariable);
+    const session = sessionFor(message.runId);
+    if (!session) return;
+    let updates = pendingVariableUpdates.get(session);
+    if (!updates) {
+      updates = new Map();
+      pendingVariableUpdates.set(session, updates);
+    }
+    updates.set(message.variableName, message.newValue);
+    if (!variableFlushScheduled) {
+      variableFlushScheduled = true;
+      scheduleFrame(flushVariableUpdates);
     }
   });
 
-  // Listen for run completion events
+  // Run lifecycle events
   client.on('completed', (message) => {
-    const currentRunId = store.getState().currentRunId;
-    if (message.runId === currentRunId) {
-      system.notifications.success(`Graph completed: ${message.graphId}`);
-      setIsExecuting(false);
-      setCurrentRunId(null);
-      setCurrentGraphId(null);
-      setIsPaused(false);
-      clearAllExecutingStates(system);
-    }
+    const controller = runner.runIndex.get(message.runId);
+    if (!controller) return;
+    controller.session.editor.notifications.success(
+      `Graph completed: ${message.graphId}`
+    );
+    controller.finishRun();
+    clearAllExecutingStates(controller.session);
   });
 
   client.on('error', (message) => {
-    const currentRunId = store.getState().currentRunId;
-    if (message.runId === currentRunId) {
-      system.notifications.error(`Graph failed: ${message.graphId}`);
-      setIsExecuting(false);
-      setCurrentRunId(null);
-      setCurrentGraphId(null);
-      setIsPaused(false);
-      clearAllExecutingStates(system);
-    }
+    if (!message.runId) return;
+    const controller = runner.runIndex.get(message.runId);
+    if (!controller) return;
+    controller.session.editor.notifications.error(
+      `Graph failed: ${message.graphId}`
+    );
+    controller.finishRun();
+    clearAllExecutingStates(controller.session);
   });
 
   client.on('stopped', (message) => {
-    const currentRunId = store.getState().currentRunId;
-    if (message.runId === currentRunId) {
-      system.notifications.info(`Graph stopped: ${message.graphId}`);
-      setIsExecuting(false);
-      setCurrentRunId(null);
-      setCurrentGraphId(null);
-      setIsPaused(false);
-      clearAllExecutingStates(system);
-    }
+    const controller = runner.runIndex.get(message.runId);
+    if (!controller) return;
+    controller.session.editor.notifications.info(
+      `Graph stopped: ${message.graphId}`
+    );
+    controller.finishRun();
+    clearAllExecutingStates(controller.session);
   });
 
   // Realtime state change listeners
   client.on('nodeRemoved', (message) => {
-    const currentRunId = store.getState().currentRunId;
-    if (message.runId === currentRunId) {
-      // Update node store to reflect removal
-      system.nodeStore
-        .getState()
-        .setNodes((nodes) =>
-          nodes.filter((node) => node.id !== message.nodeId)
-        );
-      system.notifications.info(`Node removed: ${message.nodeId}`);
-    }
+    const session = sessionFor(message.runId);
+    if (!session) return;
+    session.nodeStore
+      .getState()
+      .setNodes((nodes) => nodes.filter((node) => node.id !== message.nodeId));
+    session.editor.notifications.info(`Node removed: ${message.nodeId}`);
   });
 
   client.on('linkCreated', (message) => {
-    const currentRunId = store.getState().currentRunId;
-    if (message.runId === currentRunId) {
-      system.notifications.info(
-        `Link created: ${message.fromNodeId}/${message.fromSocket} -> ${message.toNodeId}/${message.toSocket}`
-      );
-    }
+    const session = sessionFor(message.runId);
+    if (!session) return;
+    session.editor.notifications.info(
+      `Link created: ${message.fromNodeId}/${message.fromSocket} -> ${message.toNodeId}/${message.toSocket}`
+    );
   });
 
   client.on('linkRemoved', (message) => {
-    const currentRunId = store.getState().currentRunId;
-    if (message.runId === currentRunId) {
-      system.notifications.info(
-        `Link removed: ${message.fromNodeId}/${message.fromSocket} -> ${message.toNodeId}/${message.toSocket}`
-      );
-    }
+    const session = sessionFor(message.runId);
+    if (!session) return;
+    session.editor.notifications.info(
+      `Link removed: ${message.fromNodeId}/${message.fromSocket} -> ${message.toNodeId}/${message.toSocket}`
+    );
   });
 
   client.on('nodeParamUpdated', (message) => {
-    const currentRunId = store.getState().currentRunId;
-    if (message.runId === currentRunId) {
-      system.notifications.info(`Parameter updated on ${message.nodeId}`);
-    }
+    const session = sessionFor(message.runId);
+    if (!session) return;
+    session.editor.notifications.info(`Parameter updated on ${message.nodeId}`);
   });
 
   client.on('affectedNodes', (message) => {
-    const currentRunId = store.getState().currentRunId;
-    if (message.runId === currentRunId) {
-      // Highlight affected nodes
-      system.nodeStore.getState().setNodes((nodes) =>
-        nodes.map((node) => {
-          if (message.nodeIds.includes(node.id)) {
-            return {
-              ...node,
-              data: {
-                ...node.data,
-                annotations: {
-                  ...node.data?.annotations,
-                  [executing]: true
-                }
+    const session = sessionFor(message.runId);
+    if (!session) return;
+    session.nodeStore.getState().setNodes((nodes) =>
+      nodes.map((node) => {
+        if (message.nodeIds.includes(node.id)) {
+          return {
+            ...node,
+            data: {
+              ...node.data,
+              annotations: {
+                ...node.data?.annotations,
+                [executing]: true
               }
-            };
-          }
-          return node;
-        })
-      );
-      system.notifications.info(
-        `Executing ${message.reason}: ${message.nodeIds.length} node(s)`
-      );
-    }
+            }
+          };
+        }
+        return node;
+      })
+    );
+    session.editor.notifications.info(
+      `Executing ${message.reason}: ${message.nodeIds.length} node(s)`
+    );
   });
 }

package/src/plugin/graphrunner/buttons.tsx

@@ -1,29 +1,37 @@
 import { VscodeButton } from '@vscode-elements/react-elements';
-import { useStore } from 'zustand';
+import { useStore, type StoreApi } from 'zustand';
 import { Play, PauseWindow, Square, ArrowRight } from 'iconoir-react';
-import type { StoreApi } from 'zustand';
+import { useGraph } from '@/system/provider';
 import type { GraphRunnerClientStore } from './store';
+import type { GraphRunController } from './runController';
 
-interface GraphRunnerButtonsProps {
-  store: StoreApi<GraphRunnerClientStore>;
-  onPlay: () => void;
-  onPause: () => void;
-  onResume: () => void;
-  onStep: () => void;
-  onStop: () => void;
-}
+/**
+ * Execution controls for the graph in the surrounding tab. Connection state
+ * comes from the shared runner connection; run state comes from this graph's own
+ * run controller, so each open graph has independent controls.
+ */
+export const GraphRunnerButtons = () => {
+  const session = useGraph();
+  const controller = session.runController;
+  if (!controller) return null;
+  return (
+    <Buttons
+      controller={controller}
+      connectionStore={session.editor.runner.store}
+    />
+  );
+};
 
-export const GraphRunnerButtons = ({
-  store,
-  onPlay,
-  onPause,
-  onResume,
-  onStep,
-  onStop
-}: GraphRunnerButtonsProps) => {
-  const connectionState = useStore(store, (state) => state.connectionState);
-  const isExecuting = useStore(store, (state) => state.isExecuting);
-  const isPaused = useStore(store, (state) => state.isPaused);
+const Buttons = ({
+  controller,
+  connectionStore
+}: {
+  controller: GraphRunController;
+  connectionStore: StoreApi<GraphRunnerClientStore>;
+}) => {
+  const connectionState = useStore(connectionStore, (s) => s.connectionState);
+  const isExecuting = useStore(controller.store, (s) => s.isExecuting);
+  const isPaused = useStore(controller.store, (s) => s.isPaused);
 
   const isConnected = connectionState === 'connected';
 
@@ -34,7 +42,7 @@ export const GraphRunnerButtons = ({
           secondary
           iconOnly
           title="Play Graph"
-          onClick={onPlay}
+          onClick={() => controller.play()}
           disabled={!isConnected}
         >
           <Play />
@@ -45,7 +53,7 @@ export const GraphRunnerButtons = ({
           secondary
           iconOnly
           title="Pause Graph"
-          onClick={onPause}
+          onClick={() => controller.pause()}
           disabled={!isConnected}
         >
           <PauseWindow />
@@ -57,7 +65,7 @@ export const GraphRunnerButtons = ({
             secondary
             iconOnly
             title="Resume Graph"
-            onClick={onResume}
+            onClick={() => controller.resume()}
             disabled={!isConnected}
           >
             <Play />
@@ -66,7 +74,7 @@ export const GraphRunnerButtons = ({
             secondary
             iconOnly
             title="Step Forward"
-            onClick={onStep}
+            onClick={() => controller.step()}
             disabled={!isConnected}
           >
             <ArrowRight />
@@ -77,7 +85,7 @@ export const GraphRunnerButtons = ({
         secondary
         iconOnly
         title="Stop Graph"
-        onClick={onStop}
+        onClick={() => controller.stop()}
         disabled={!isExecuting}
       >
         <Square />

package/src/plugin/graphrunner/client.ts

@@ -301,7 +301,10 @@ export class GraphRunnerClient {
       graph: options?.graph as GraphJSON,
       inputs: options?.inputs,
       options: {
-        trace: options?.trace ?? true
+        // Default off , tracing is a debugging aid with a real per-node cost;
+        // callers opt in explicitly (the run controller passes the panel's
+        // "Enable execution tracing" preference).
+        trace: options?.trace ?? false
       }
     });