@bastani/atomic 0.9.45 → 0.9.46-alpha.1

package/dist/builtin/workflows/builtin/{chunk-jvna28mq.js → chunk-zhyqnf5g.js}

@@ -80,7 +80,7 @@ function reciprocalWorkflowRootRunId(runById, runId) {
       return declaredRoots.every((declared) => declared === currentId) ? currentId : undefined;
     }
     const parent = runById.get(parentRunId);
-    const boundary = parent?.stages.find((stage) => stage.id === parentStageId);
+    const boundary = parent === undefined ? undefined : findSnapshotStage(parent.stages, parentStageId);
     if (parent === undefined || authoritativeWorkflowChildRunId(boundary) !== currentId) {
       return;
     }
@@ -88,6 +88,26 @@ function reciprocalWorkflowRootRunId(runById, runId) {
   }
   return;
 }
+var runIndexes = new WeakMap;
+var stageIndexes = new WeakMap;
+function index(items, cache) {
+  let byId = cache.get(items);
+  if (byId === undefined) {
+    const built = new Map;
+    for (const item of items)
+      if (!built.has(item.id))
+        built.set(item.id, item);
+    byId = built;
+    cache.set(items, byId);
+  }
+  return byId;
+}
+function findSnapshotRun(runs, runId) {
+  return Object.isFrozen(runs) ? index(runs, runIndexes).get(runId) : runs.find((run) => run.id === runId);
+}
+function findSnapshotStage(stages, stageId) {
+  return Object.isFrozen(stages) ? index(stages, stageIndexes).get(stageId) : stages.find((stage) => stage.id === stageId);
+}
 
 // dist/builtin/workflows/src/durable/prompt-reservations.ts
 function durablePromptScope(backend, workflowId) {
@@ -1061,26 +1081,172 @@ function compactToolNode(node) {
     ...node.source !== undefined ? { source: boundedToolText(node.source) } : {}
   };
 }
-function compactRun(run) {
-  const { inputs: _inputs, result: sourceResult, stages, toolNodes, pendingPrompt, ...metadata } = run;
+function compactRun(run, stages = run.stages.map(compactStage), toolNodes = run.toolNodes.map(compactToolNode)) {
+  const {
+    inputs: _inputs,
+    result: sourceResult,
+    stages: _stages,
+    toolNodes: _toolNodes,
+    pendingPrompt,
+    ...metadata
+  } = run;
   const result = compactRunResult(sourceResult, run.exited === true && run.status === "failed");
   return {
     ...metadata,
     inputs: {},
-    stages: stages.map(compactStage),
-    toolNodes: toolNodes.map(compactToolNode),
+    stages,
+    toolNodes,
     ...pendingPrompt !== undefined ? { pendingPrompt: clonePrompt(pendingPrompt) } : {},
     ...result !== undefined ? { result } : {}
   };
 }
-function createGraphStoreSnapshot(runs, notices, version) {
-  const snapshot = {
-    runs: runs.map(compactRun),
-    notices: notices.map((notice) => ({ ...notice })),
-    version
+function sourceShape(value) {
+  const keys = Object.keys(value);
+  return { keys, values: keys.map((key) => value[key]) };
+}
+function sameSourceShape(value, shape) {
+  let index = 0;
+  for (const key in value) {
+    if (index >= shape.keys.length || key !== shape.keys[index] || value[key] !== shape.values[index])
+      return false;
+    index++;
+  }
+  return index === shape.keys.length;
+}
+function projectStage(stage) {
+  return {
+    shape: sourceShape(stage),
+    parentCount: stage.parentIds.length,
+    toolEventCount: stage.toolEvents.length,
+    lastToolEvent: stage.toolEvents.at(-1),
+    noticeCount: stage.notices?.length,
+    compact: compactStage(stage)
   };
-  deepFreezeGraphValue(snapshot);
-  return snapshot;
+}
+function stageUnchanged(stage, projected) {
+  return sameSourceShape(stage, projected.shape) && stage.parentIds.length === projected.parentCount && stage.toolEvents.length === projected.toolEventCount && stage.toolEvents.at(-1) === projected.lastToolEvent && stage.notices?.length === projected.noticeCount;
+}
+function projectToolNode(node) {
+  return { shape: sourceShape(node), parentCount: node.parentIds.length, compact: compactToolNode(node) };
+}
+function toolNodeUnchanged(node, projected) {
+  return sameSourceShape(node, projected.shape) && node.parentIds.length === projected.parentCount;
+}
+
+class GraphStoreSnapshotBuilder {
+  projected = new Map;
+  changedRuns = new Set;
+  changedStages = new Map;
+  changedToolNodes = new Map;
+  markChanged(change) {
+    if (change === undefined) {
+      this.projected.clear();
+      this.changedRuns.clear();
+      this.changedStages.clear();
+      this.changedToolNodes.clear();
+      return;
+    }
+    if (change.stageId === undefined && change.toolNodeId === undefined) {
+      this.changedRuns.add(change.runId);
+      return;
+    }
+    if (change.stageId !== undefined)
+      addScoped(this.changedStages, change.runId, change.stageId);
+    if (change.toolNodeId !== undefined)
+      addScoped(this.changedToolNodes, change.runId, change.toolNodeId);
+  }
+  build(runs, notices, version) {
+    const previous = this.projected;
+    const next = new Map;
+    const compactRuns = [];
+    for (const run of runs) {
+      const known = previous.get(run.id);
+      const reusable = known?.source === run ? known : undefined;
+      const projected = reusable !== undefined && !this.changedRuns.has(run.id) && this.runUnchanged(run, reusable) ? reusable : this.projectRun(run, reusable);
+      next.set(run.id, projected);
+      compactRuns.push(projected.compact);
+    }
+    this.projected = next;
+    this.changedRuns.clear();
+    this.changedStages.clear();
+    this.changedToolNodes.clear();
+    const snapshot = {
+      runs: compactRuns,
+      notices: notices.map((notice) => ({ ...notice })),
+      version
+    };
+    deepFreezeGraphValue(snapshot);
+    return snapshot;
+  }
+  buildRun(run) {
+    const known = this.projected.get(run.id);
+    const reusable = known?.source === run ? known : undefined;
+    const projected = reusable !== undefined && !this.changedRuns.has(run.id) && this.runUnchanged(run, reusable) ? reusable : this.projectRun(run, reusable);
+    deepFreezeGraphValue(projected.compact);
+    this.projected.set(run.id, projected);
+    this.changedRuns.delete(run.id);
+    this.changedStages.delete(run.id);
+    this.changedToolNodes.delete(run.id);
+    return projected.compact;
+  }
+  runUnchanged(run, projected) {
+    if (!sameSourceShape(run, projected.shape))
+      return false;
+    const { stages, toolNodes } = projected.compact;
+    if (run.stages.length !== stages.length || run.toolNodes.length !== toolNodes.length)
+      return false;
+    const changedStages = this.changedStages.get(run.id);
+    for (let index = 0;index < run.stages.length; index++) {
+      const stage = run.stages[index];
+      const known = projected.stages.get(stage);
+      if (known === undefined || known.compact !== stages[index] || changedStages?.has(stage.id) === true || !stageUnchanged(stage, known))
+        return false;
+    }
+    const changedToolNodes = this.changedToolNodes.get(run.id);
+    for (let index = 0;index < run.toolNodes.length; index++) {
+      const node = run.toolNodes[index];
+      const known = projected.toolNodes.get(node);
+      if (known === undefined || known.compact !== toolNodes[index] || changedToolNodes?.has(node.id) === true || !toolNodeUnchanged(node, known))
+        return false;
+    }
+    return true;
+  }
+  projectRun(run, previous) {
+    const changedStages = this.changedStages.get(run.id);
+    const stages = new Map;
+    const compactStages = [];
+    for (const stage of run.stages) {
+      let projected = stages.get(stage) ?? previous?.stages.get(stage);
+      if (projected === undefined || changedStages?.has(stage.id) === true || !stageUnchanged(stage, projected))
+        projected = projectStage(stage);
+      stages.set(stage, projected);
+      compactStages.push(projected.compact);
+    }
+    const changedToolNodes = this.changedToolNodes.get(run.id);
+    const toolNodes = new Map;
+    const compactToolNodes = [];
+    for (const node of run.toolNodes) {
+      let projected = toolNodes.get(node) ?? previous?.toolNodes.get(node);
+      if (projected === undefined || changedToolNodes?.has(node.id) === true || !toolNodeUnchanged(node, projected))
+        projected = projectToolNode(node);
+      toolNodes.set(node, projected);
+      compactToolNodes.push(projected.compact);
+    }
+    return {
+      source: run,
+      shape: sourceShape(run),
+      stages,
+      toolNodes,
+      compact: compactRun(run, compactStages, compactToolNodes)
+    };
+  }
+}
+function addScoped(scopes, runId, nodeId) {
+  const ids = scopes.get(runId);
+  if (ids === undefined)
+    scopes.set(runId, new Set([nodeId]));
+  else
+    ids.add(nodeId);
 }
 function deepFreezeGraphValue(value) {
   if (value === null || typeof value !== "object" || Object.isFrozen(value))
@@ -1099,11 +1265,16 @@ function deepFreezeGraphValue(value) {
 // dist/builtin/workflows/src/shared/workflow-observation-runtime.ts
 import { randomUUID } from "node:crypto";
 
+// dist/builtin/workflows/src/shared/expanded-workflow-graph.ts
+var leafExpansions = new WeakMap;
+var snapshotExpansions = new WeakMap;
+
 // dist/builtin/workflows/src/shared/workflow-activity.ts
 function workflowActivityNodeKey(runId, nodeId) {
   return `${runId}:${nodeId}`;
 }
 var endedStatuses = new Set(["completed", "cached", "failed", "cancelled"]);
+var projectedGraphNodes = new WeakMap;
 
 // dist/builtin/workflows/src/shared/workflow-observation-runtime.ts
 class WorkflowObservationRuntime {
@@ -1118,6 +1289,7 @@ class WorkflowObservationRuntime {
   listeners = new Set;
   replayStageIds = new Set;
   previous = new Map;
+  owners = new Map;
   recoveryDepth = 0;
   constructor(runs, invalidate) {
     this.runs = runs;
@@ -1143,7 +1315,7 @@ class WorkflowObservationRuntime {
         if (id.startsWith(`${runId}:`))
           ids.delete(id);
     }
-    this.invalidate();
+    this.invalidate({ runId });
   }
   emit(event) {
     for (const listener of this.listeners) {
@@ -1192,7 +1364,7 @@ class WorkflowObservationRuntime {
     if (action === "resume")
       this.stoppingRunIds.delete(runId);
     this.lifecycle({ kind: "run", runId, status: run.status, action }, attribution);
-    this.invalidate();
+    this.invalidate({ runId });
   }
   heartbeat(runId, scheduledAt, intervalMinutes) {
     const run = this.runs().find((candidate) => candidate.id === runId);
@@ -1217,35 +1389,103 @@ class WorkflowObservationRuntime {
         this.emit({ kind: "availability" });
     };
   }
-  capture() {
-    const next = new Map;
-    const visit = (key, target, historical = false) => {
+  capture(changes) {
+    if (changes === undefined) {
+      const next = new Map;
+      const nextOwners = new Map;
+      for (const run of this.runs()) {
+        const owners = nextOwners.get(run.id) ?? new Map;
+        nextOwners.set(run.id, owners);
+        this.captureRun(run, this.previous.has(`run:${run.id}`), next, owners);
+      }
+      this.previous = next;
+      this.owners = nextOwners;
+      return;
+    }
+    if (changes.length === 0)
+      return;
+    const changed = new Map;
+    for (const change of changes) {
+      const current = changed.get(change.runId);
+      if (current === "run")
+        continue;
+      if (change.stageId === undefined && change.toolNodeId === undefined) {
+        changed.set(change.runId, "run");
+        continue;
+      }
+      const nodes = current ?? { stages: new Set, tools: new Set };
+      if (change.stageId !== undefined)
+        nodes.stages.add(change.stageId);
+      if (change.toolNodeId !== undefined)
+        nodes.tools.add(change.toolNodeId);
+      changed.set(change.runId, nodes);
+    }
+    for (const run of this.runs()) {
+      const change = changed.get(run.id);
+      if (change === undefined)
+        continue;
+      const existing = this.previous.has(`run:${run.id}`);
+      const nodes = change === "run" || !existing ? undefined : change;
+      const next = new Map;
+      const nextOwners = new Map;
+      this.captureRun(run, existing, next, nextOwners, nodes);
+      const owners = this.owners.get(run.id) ?? new Map;
+      const reread = nodes === undefined ? [...owners.keys()] : [...[...nodes.stages].map((id) => `stage:${id}`), ...[...nodes.tools].map((id) => `tool:${id}`)];
+      for (const owner of reread)
+        for (const key of owners.get(owner) ?? [])
+          if (!next.has(key))
+            this.previous.delete(key);
+      for (const [key, target] of next)
+        this.previous.set(key, target);
+      if (nodes === undefined) {
+        this.owners.set(run.id, nextOwners);
+        continue;
+      }
+      for (const owner of reread) {
+        const keys = nextOwners.get(owner);
+        if (keys === undefined)
+          owners.delete(owner);
+        else
+          owners.set(owner, keys);
+      }
+      this.owners.set(run.id, owners);
+    }
+  }
+  captureRun(run, existing, next, owners, nodes) {
+    const visit = (owner, key, target, historical = false) => {
       next.set(key, target);
+      const keys = owners.get(owner);
+      if (keys === undefined)
+        owners.set(owner, [key]);
+      else
+        keys.push(key);
       const previous = this.previous.get(key);
       if (!historical && this.liveRunIds.has(target.runId) && previous?.status !== target.status) {
-        this.lifecycle(previous && previous.kind !== "prompt" && target.kind !== "prompt" ? { ...target, previousStatus: previous.status } : target, this.runs().find((run) => run.id === target.runId)?.origin);
+        this.lifecycle(previous && previous.kind !== "prompt" && target.kind !== "prompt" ? { ...target, previousStatus: previous.status } : target, this.runs().find((candidate) => candidate.id === target.runId)?.origin);
       }
     };
-    for (const run of this.runs()) {
-      const existing = this.previous.has(`run:${run.id}`);
-      visit(`run:${run.id}`, { kind: "run", runId: run.id, status: run.status });
-      for (const stage of run.stages) {
-        visit(`stage:${workflowActivityNodeKey(run.id, stage.id)}`, { kind: "stage", runId: run.id, stageId: stage.id, stageName: stage.name, status: stage.status }, !existing);
-        if (stage.pendingPrompt)
-          visit(`prompt:${stage.pendingPrompt.id}`, {
-            kind: "prompt",
-            runId: run.id,
-            stageId: stage.id,
-            promptId: stage.pendingPrompt.id,
-            status: "opened"
-          }, !existing);
-      }
-      for (const tool of run.toolNodes)
-        visit(`tool:${workflowActivityNodeKey(run.id, tool.id)}`, { kind: "tool", runId: run.id, toolNodeId: tool.id, toolName: tool.name, status: tool.status }, !existing);
-      if (run.pendingPrompt)
-        visit(`prompt:${run.pendingPrompt.id}`, { kind: "prompt", runId: run.id, promptId: run.pendingPrompt.id, status: "opened" }, !existing);
+    visit("run", `run:${run.id}`, { kind: "run", runId: run.id, status: run.status });
+    for (const stage of run.stages) {
+      if (nodes !== undefined && !nodes.stages.has(stage.id))
+        continue;
+      const owner = `stage:${stage.id}`;
+      visit(owner, `stage:${workflowActivityNodeKey(run.id, stage.id)}`, { kind: "stage", runId: run.id, stageId: stage.id, stageName: stage.name, status: stage.status }, !existing);
+      if (stage.pendingPrompt)
+        visit(owner, `prompt:${stage.pendingPrompt.id}`, {
+          kind: "prompt",
+          runId: run.id,
+          stageId: stage.id,
+          promptId: stage.pendingPrompt.id,
+          status: "opened"
+        }, !existing);
+    }
+    for (const tool of run.toolNodes) {
+      if (nodes !== undefined && !nodes.tools.has(tool.id))
+        continue;
+      visit(`tool:${tool.id}`, `tool:${workflowActivityNodeKey(run.id, tool.id)}`, { kind: "tool", runId: run.id, toolNodeId: tool.id, toolName: tool.name, status: tool.status }, !existing);
     }
-    this.previous = next;
+    if (nodes === undefined && run.pendingPrompt)
+      visit("run", `prompt:${run.pendingPrompt.id}`, { kind: "prompt", runId: run.id, promptId: run.pendingPrompt.id, status: "opened" }, !existing);
   }
 }
 var workflowObservationRuntimeKey = Symbol.for("atomic-workflows/store-observation-runtime@1");
@@ -1363,25 +1603,41 @@ function createStoreState() {
 function createStoreContext(state = createStoreState()) {
   let cachedGraphVersion = -1;
   let cachedGraphSnapshot;
+  const graphBuilder = new GraphStoreSnapshotBuilder;
   const observation = new WorkflowObservationRuntime(() => state.runs, bumpAndNotify);
   function snapshot() {
     return JSON.parse(JSON.stringify({ runs: state.runs, notices: state.notices, version: state.version }));
   }
   function graphSnapshot() {
     if (cachedGraphSnapshot === undefined || cachedGraphVersion !== state.version) {
-      cachedGraphSnapshot = createGraphStoreSnapshot(state.runs, state.notices, state.version);
+      cachedGraphSnapshot = graphBuilder.build(state.runs, state.notices, state.version);
       cachedGraphVersion = state.version;
     }
     return cachedGraphSnapshot;
   }
+  function graphRun(runId) {
+    if (cachedGraphSnapshot !== undefined && cachedGraphVersion === state.version)
+      return findSnapshotRun(cachedGraphSnapshot.runs, runId);
+    const run = state.runs.find((candidate) => candidate.id === runId);
+    return run === undefined ? undefined : graphBuilder.buildRun(run);
+  }
   let batchDepth = 0;
   let batchedNotification = false;
-  function notify() {
+  let batchedChanges = [];
+  function notify(change) {
+    graphBuilder.markChanged(change);
     if (batchDepth > 0) {
       batchedNotification = true;
+      if (change === undefined)
+        batchedChanges = undefined;
+      else
+        batchedChanges?.push(change);
       return;
     }
-    observation.capture();
+    observation.capture(change === undefined ? undefined : [change]);
+    publish();
+  }
+  function publish() {
     for (const fn of state.invalidationListeners) {
       try {
         fn();
@@ -1396,9 +1652,9 @@ function createStoreContext(state = createStoreState()) {
       }
     }
   }
-  function bumpAndNotify() {
+  function bumpAndNotify(change) {
     state.version++;
-    notify();
+    notify(change);
   }
   function batch(update) {
     batchDepth++;
@@ -1407,8 +1663,11 @@ function createStoreContext(state = createStoreState()) {
     } finally {
       batchDepth--;
       if (batchDepth === 0 && batchedNotification) {
+        const changes = batchedChanges;
         batchedNotification = false;
-        notify();
+        batchedChanges = [];
+        observation.capture(changes);
+        publish();
       }
     }
   }
@@ -1464,6 +1723,7 @@ function createStoreContext(state = createStoreState()) {
     observation,
     snapshot,
     graphSnapshot,
+    graphRun,
     notify,
     batch,
     bumpAndNotify,
@@ -2252,7 +2512,7 @@ function createRunStoreMethods(context) {
         state.runs.push(run);
       else
         state.runs[index] = run;
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId: run.id });
     },
     recordRunSummary(runId, summary, summarizedThrough) {
       const run = context.findRun(runId);
@@ -2268,7 +2528,7 @@ function createRunStoreMethods(context) {
       if (update.phase !== undefined && update.phase !== run.phase)
         run.phaseStartedAt = Date.now();
       Object.assign(run, update);
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId });
     },
     reconcileRunParentStage(runId, expectedParentStageId, parentStageId) {
       const run = context.findRun(runId);
@@ -2330,7 +2590,7 @@ function createRunStoreMethods(context) {
         context.rejectPrompt(pending.id, `atomic-workflows: run ${runId} ended before prompt resolved`);
       }
       context.rejectAllStagePrompts(runId, run, `atomic-workflows: run ${runId} ended before prompt resolved`);
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId });
       return true;
     },
     recordRunBlocked(runId, error, metadata) {
@@ -2522,6 +2782,9 @@ function createRunStoreMethods(context) {
     graphSnapshot() {
       return context.graphSnapshot();
     },
+    graphRun(runId) {
+      return context.graphRun(runId);
+    },
     subscribe(fn) {
       state.listeners.add(fn);
       return () => {
@@ -2559,7 +2822,7 @@ function createToolNodeStoreMethods(context) {
       if (node.executionOrder === undefined)
         node.executionOrder = nextExecutionOrder(run);
       nodes.push(node);
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, toolNodeId: node.id });
       return true;
     },
     recordToolNodeRunning(runId, nodeId, startedAt) {
@@ -2570,7 +2833,7 @@ function createToolNodeStoreMethods(context) {
       node.startedAt = startedAt;
       if (context.observation.liveRunIds.has(runId))
         context.observation.executingToolNodeIds.add(workflowActivityNodeKey(runId, nodeId));
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, toolNodeId: nodeId });
       return true;
     },
     recordToolNodeEnd(runId, nodeId, update) {
@@ -2589,7 +2852,7 @@ function createToolNodeStoreMethods(context) {
         node.resultSummary = update.resultSummary;
       if (update.error !== undefined)
         node.error = update.error;
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, toolNodeId: nodeId });
       return true;
     }
   };
@@ -2610,7 +2873,7 @@ function createStageStoreMethods(context) {
       if (stage.status === "running" && context.observation.liveRunIds.has(runId)) {
         context.observation.executingStageIds.add(workflowActivityNodeKey(runId, stage.id));
       }
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId: stage.id });
     },
     recordStageWorkflowChildRun(runId, stageId, ref) {
       const run = context.findRun(runId);
@@ -2622,11 +2885,11 @@ function createStageStoreMethods(context) {
       const released = context.observation.executingStageIds.delete(workflowActivityNodeKey(runId, stageId));
       if (stage.workflowChildRun?.runId === ref.runId && stage.workflowChildRun.alias === ref.alias && stage.workflowChildRun.workflow === ref.workflow) {
         if (released)
-          context.bumpAndNotify();
+          context.bumpAndNotify({ runId, stageId });
         return false;
       }
       stage.workflowChildRun = { ...ref };
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordToolStart(runId, stageId, evt) {
@@ -2640,7 +2903,7 @@ function createStageStoreMethods(context) {
       if (!exists) {
         stage.toolEvents.push(evt);
       }
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
     },
     recordToolEnd(runId, stageId, evt) {
       const run = context.findRun(runId);
@@ -2654,7 +2917,7 @@ function createStageStoreMethods(context) {
         existing.endedAt = evt.endedAt;
         existing.output = evt.output;
       }
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
     },
     recordStageEnd(runId, stage) {
       const run = context.findRun(runId);
@@ -2710,7 +2973,7 @@ function createStageStoreMethods(context) {
       delete existing.awaitingInputSince;
       delete existing.inputRequest;
       context.rejectStagePrompt(runId, existing, `atomic-workflows: stage ${stage.id} ended before prompt resolved`);
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId: stage.id });
     },
     recordStageSession(runId, stageId, session) {
       const run = context.findRun(runId);
@@ -2720,7 +2983,7 @@ function createStageStoreMethods(context) {
       if (!stage || stage.sessionId === session.sessionId)
         return false;
       stage.sessionId = session.sessionId;
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordStageAttachable(runId, stageId, attachable) {
@@ -2738,7 +3001,7 @@ function createStageStoreMethods(context) {
       } else {
         stage.attachable = next;
       }
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordStageAttached(runId, stageId, attached) {
@@ -2756,7 +3019,7 @@ function createStageStoreMethods(context) {
       } else {
         stage.attached = next;
       }
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordStageAwaitingInput(runId, stageId, awaiting, ts) {
@@ -2783,7 +3046,7 @@ function createStageStoreMethods(context) {
         stage.status = "running";
         delete stage.awaitingInputSince;
       }
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordStageInputRequest(runId, stageId, request) {
@@ -2800,7 +3063,7 @@ function createStageStoreMethods(context) {
       if (stage.inputRequest?.id === request.id)
         return false;
       stage.inputRequest = { ...request };
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     clearStageInputRequest(runId, stageId) {
@@ -2811,7 +3074,7 @@ function createStageStoreMethods(context) {
       if (!stage || stage.inputRequest === undefined)
         return false;
       delete stage.inputRequest;
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordStageBlocked(runId, stageId, blockedBy) {
@@ -2834,7 +3097,7 @@ function createStageStoreMethods(context) {
         stage.blockedByStageId = blockedBy;
         delete stage.awaitingInputSince;
       }
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordStageUnblocked(runId, stageId) {
@@ -2847,7 +3110,7 @@ function createStageStoreMethods(context) {
       stage.status = "pending";
       delete stage.blockedByStageId;
       delete stage.awaitingInputSince;
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordStageNotice(runId, stageId, notice) {
@@ -2860,7 +3123,7 @@ function createStageStoreMethods(context) {
       if (!stage.notices)
         stage.notices = [];
       stage.notices.push(notice);
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordStagePaused(runId, stageId, pausedAt, metadata) {
@@ -2878,7 +3141,7 @@ function createStageStoreMethods(context) {
         if (stage.pauseActor === metadata.actor)
           return false;
         stage.pauseActor = metadata.actor;
-        context.bumpAndNotify();
+        context.bumpAndNotify({ runId, stageId });
         return false;
       }
       context.observation.executingStageIds.delete(workflowActivityNodeKey(runId, stageId));
@@ -2891,7 +3154,7 @@ function createStageStoreMethods(context) {
         delete stage.pauseActor;
       else
         stage.pauseActor = metadata.actor;
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     },
     recordStageResumed(runId, stageId, resumedAt, metadata) {
@@ -2910,7 +3173,7 @@ function createStageStoreMethods(context) {
         if (stage.resumeActor === metadata.actor)
           return false;
         stage.resumeActor = metadata.actor;
-        context.bumpAndNotify();
+        context.bumpAndNotify({ runId, stageId });
         return false;
       }
       const resumedTs = nextControlTimestamp(resumedAt, stage.pausedAt);
@@ -2929,7 +3192,7 @@ function createStageStoreMethods(context) {
         delete stage.resumeActor;
       else
         stage.resumeActor = metadata.actor;
-      context.bumpAndNotify();
+      context.bumpAndNotify({ runId, stageId });
       return true;
     }
   };