[core] Retain workflow VMs across attributes (#3609)

* Retain workflow VMs across attributes

* Restore runtime logger spy automatically

* Explain retained attribute race ordering

* Update retained attribute telemetry expectation

* Clarify retained attribute replay coverage

* Clarify retained attribute execution comment

Author: NathanColosimo

.changeset/retain-vms-across-attributes.md

@@ -0,0 +1,5 @@
+---
+'@workflow/core': patch
+---
+
+Retain workflow VMs across workflow attribute writes.

packages/core/src/private.ts

@@ -137,11 +137,11 @@ export interface WorkflowOrchestratorContext {
   globalThis: typeof globalThis;
   /**
    * Increments when a suspension is accepted and on every retained-session
-   * resume. Step and hook suspension signals capture it when scheduled and
-   * no-op if it moved, which drops same-boundary sibling signals and timers
-   * queued at boundary N that would fire after the session resumed into
-   * boundary N+1. Sleep and attribute signals are intentionally unguarded:
-   * their presence makes the boundary unretainable, so a late signal correctly
+   * resume. Step, hook, and attribute suspension signals capture it when
+   * scheduled and no-op if it moved, which drops same-boundary sibling signals
+   * and timers queued at boundary N that would fire after the session resumed
+   * into boundary N+1. Sleep signals are intentionally unguarded: their
+   * presence makes the boundary unretainable, so a late signal correctly
    * demotes the session (workflow.ts `onWorkflowError`).
    */
   suspensionGeneration: number;

packages/core/src/retained-vm-loop.test.ts

@@ -144,6 +144,14 @@ const concurrentHookWakeWorkflow = `const createHook = globalThis[Symbol.for("WO
   }
   globalThis.__private_workflows = new Map([["workflow", workflow]]);`;
 
+const attributeThenStepWorkflow = `const setAttributes = globalThis[Symbol.for("WORKFLOW_SET_ATTRIBUTES")];
+  const s1 = globalThis[Symbol.for("WORKFLOW_USE_STEP")]("r_s1");
+  async function workflow() {
+    await setAttributes([{ key: "phase", value: "ready" }]);
+    return await s1();
+  }
+  globalThis.__private_workflows = new Map([["workflow", workflow]]);`;
+
 // Hook metadata is serialized at the same suspension boundary as step input.
 // A getter mutating workflow state must demote the retained VM just like an
 // unsafe step argument, because a cold replay skips serialization after the
@@ -781,6 +789,24 @@ describe('retained VM through the inline replay loop', () => {
     expect(on.durableLog).toEqual(off.durableLog);
   });
 
+  it('retains one VM across an attribute write and the following step', async () => {
+    process.env.WORKFLOW_RETAINED_VM = '0';
+    const off = await drive(
+      'wrun_retained_attribute_then_step',
+      attributeThenStepWorkflow
+    );
+    createContextSpy.mockClear();
+    delete process.env.WORKFLOW_RETAINED_VM;
+
+    const on = await drive(
+      'wrun_retained_attribute_then_step',
+      attributeThenStepWorkflow
+    );
+    expect(on.result).toBe(10);
+    expect(on.vmBuilds).toBe(1);
+    expect(on.durableLog).toEqual(off.durableLog);
+  });
+
   it('matches cold replay when hook metadata serialization mutates workflow state', async () => {
     process.env.WORKFLOW_RETAINED_VM = '0';
     const off = await drive(

packages/core/src/runtime.test.ts

@@ -1060,8 +1060,8 @@ describe('workflowEntrypoint replay guards', () => {
     ).toEqual([]);
   });
 
-  it('replays attribute events before executing a step that loses the same race', async () => {
-    const debug = vi
+  it('retains when an attribute write beats a step that signals first', async () => {
+    using debug = vi
       .spyOn(runtimeLogger, 'debug')
       .mockImplementation(() => undefined);
     const ops: Promise<any>[] = [];
@@ -1080,14 +1080,17 @@ describe('workflowEntrypoint replay guards', () => {
       startedAt: new Date('2024-01-01T00:00:00.000Z'),
       deploymentId: 'test-deployment',
     };
+    // Register the step consumer first so it schedules the accepted
+    // suspension. The later attribute signal must be dropped by its new
+    // generation guard after the retained session resumes.
     const workflowCode = `
       const setAttributes = globalThis[Symbol.for("WORKFLOW_SET_ATTRIBUTES")];
       const useStep = globalThis[Symbol.for("WORKFLOW_USE_STEP")];
       const slowStep = useStep("slowStep");
       async function workflow() {
         await Promise.race([
-          setAttributes([{ key: "winner", value: "attribute" }]),
           slowStep(),
+          setAttributes([{ key: "winner", value: "attribute" }]),
         ]);
         return "attribute won";
       }${getWorkflowTransformCode('workflow')}`;
@@ -1103,19 +1106,17 @@ describe('workflowEntrypoint replay guards', () => {
     expect(createdEvents).toContainEqual(
       expect.objectContaining({ eventType: 'attr_set' })
     );
-    // The attr_set suspension skips step processing and resolves through an
-    // in-process replay, where the durable attribute event wins the race —
-    // so the run completes within this same delivery, with no queue
-    // interaction.
+    // The durable attribute event wins the retained resume. The later
+    // attribute suspension signal must not demote the session.
     expect(createdEvents).toContainEqual(
       expect.objectContaining({ eventType: 'run_completed' })
     );
     expect(queueCalls).toEqual([]);
     // Under lazy inline start the step that loses the attribute race is NOT
     // eagerly created: its step_created is deferred for a lazy step_started
-    // that never fires, because the attribute-resolving replay decides the
-    // race before any step executes. So the losing step leaves no events at
-    // all — strictly less event-log garbage than the eager model.
+    // that never fires, because the attribute-resolving retained execution
+    // decides the race before any step executes. So the losing step leaves no
+    // events at all — strictly less event-log garbage than the eager model.
     expect(createdEvents).not.toContainEqual(
       expect.objectContaining({ eventType: 'step_created' })
     );
@@ -1125,8 +1126,7 @@ describe('workflowEntrypoint replay guards', () => {
     const executionModes = debug.mock.calls
       .filter(([message]) => message === 'Starting workflow execution')
       .map(([, context]) => context?.executionMode);
-    expect(executionModes).toEqual(['replay', 'replay']);
-    debug.mockRestore();
+    expect(executionModes).toEqual(['replay', 'retained']);
   });
 
   it('fails the run when the World rejects an attr_set event as invalid', async () => {
@@ -3088,11 +3088,11 @@ describe('workflowEntrypoint latency telemetry (ttfs / stso)', () => {
     }${latXform('workflow')}`;
 
   it('reports ttfs across a pre-step setAttributes, resolved in-process without a re-invoke', async () => {
-    // A workflow-body setAttributes before the first step resolves through
-    // an in-process replay: the same delivery commits attr_set, replays, and
-    // runs the step — no queue interaction. TTFS ends at the attr write, so
-    // the setAttributes call's duration is subtracted; here that detour is
-    // milliseconds, keeping ttfs ≈ the run's ULID backdate.
+    // A workflow-body setAttributes before the first step resolves in-process:
+    // the same delivery commits attr_set, reloads it, and advances the workflow
+    // by resuming or replaying — no queue interaction. TTFS ends at the attr
+    // write, so the setAttributes call's duration is subtracted. That detour is
+    // only milliseconds here, keeping ttfs ≈ the run's ULID backdate.
     const backdateMs = 10_000;
     const before = Date.now();
     const first = await driveLatency({
@@ -3115,11 +3115,13 @@ describe('workflowEntrypoint latency telemetry (ttfs / stso)', () => {
     expect(ttfs).toBeLessThanOrEqual(backdateMs + elapsed);
     expect(stso).toBeUndefined();
     // The attr detour does not end turbo: no resume invocation source
-    // exists, so the forced-optimistic fast path stays engaged.
+    // exists, so the forced-optimistic fast path stays engaged. The live VM
+    // also resumes after the attr write instead of replaying in a fresh VM.
     expect(optimizations).toEqual([
       'turbo',
       'lazyStepStart',
       'optimisticStart',
+      'retained',
     ]);
   });
 

packages/core/src/runtime.ts

@@ -524,7 +524,7 @@ type RetentionDecision =
       reason:
         | 'disabled'
         | 'serialization_executed_workflow_code'
-        | 'no_step_driver'
+        | 'no_replay_driver'
         | 'unsupported_suspension_item'
         | 'open_wait';
     };
@@ -534,9 +534,9 @@ type RetentionDecision =
  *
  * Every item type accepted here must use the suspension-generation guard when
  * signaling. Otherwise a signal scheduled at boundary N could suspend the VM
- * after it has already resumed into boundary N+1. Waits and attributes are not
- * guarded, so they remain unretainable. A step is required to drive the next
- * inline iteration; hook-only suspensions park normally.
+ * after it has already resumed into boundary N+1. Waits are not guarded, so
+ * they remain unretainable. A step or attribute write is required to drive the
+ * next inline iteration; hook-only suspensions park normally.
  *
  * Retaining across an open hook also permits a hook-woken cold replay to race
  * this invocation. That is safe only because each loaded log is a monotone,
@@ -575,17 +575,17 @@ function getRetentionDecision({
       reason: 'serialization_executed_workflow_code',
     };
   }
-  if (suspension.stepCount === 0) {
-    return { retain: false, reason: 'no_step_driver' };
+  if (suspension.stepCount === 0 && suspension.attributeCount === 0) {
+    return { retain: false, reason: 'no_replay_driver' };
   }
   if (
     !suspension.items.every((item) => {
       switch (item.type) {
         case 'step':
         case 'hook':
+        case 'attribute':
           return true;
         case 'wait':
-        case 'attribute':
           return false;
         default:
           item satisfies never;
@@ -3423,15 +3423,16 @@ export function workflowEntrypoint(
                           return await reinvoke(0);
                         }
 
-                        // Native workflow attribute events are resolved
-                        // through replay: the next loop iteration reloads the
-                        // log (now holding the just-committed attr_set) and
-                        // replays, resolving the setAttributes promise. Skip
-                        // step processing for this pass so that replay decides
-                        // races first: in Promise.race([setAttributes(),
-                        // step()]), the durable attribute event must be able
-                        // to win without executing the losing step. The replay
-                        // happens in-process rather than via a queue
+                        // Native workflow attribute events are resolved by the
+                        // next execution pass: it reloads the log (now holding
+                        // the just-committed attr_set), then either resumes the
+                        // retained VM or performs a cold replay to resolve the
+                        // setAttributes promise. Skip step processing for this
+                        // pass so that the durable event decides races first:
+                        // in Promise.race([setAttributes(), step()]), the
+                        // attribute must be able to win without executing the
+                        // losing step. The next pass happens in-process rather
+                        // than via a queue
                         // re-invocation: unlike hooks and waits, an attr_set
                         // introduces no out-of-band invocation source that the
                         // handler would need to yield the message for, so

packages/core/src/workflow.ts

@@ -367,17 +367,17 @@ async function createWorkflowSession({
         state = WorkflowSuspension.is(error)
           ? { type: 'suspended', suspension: error }
           : { type: 'replay' };
-        // Step and hook consumers can schedule the same suspension. The first
-        // signal advances the generation so the rest no-op.
+        // Step, hook, and attribute consumers can schedule the same suspension.
+        // The first signal advances the generation so the rest no-op.
         workflowContext.suspensionGeneration++;
         interruption.reject(error);
         return;
       }
       case 'suspended':
         // Same-boundary duplicates were staled by the generation bump above,
-        // so anything landing here is out-of-band — an unguarded sleep/
-        // attribute signal or a divergence. Those boundaries are unretainable
-        // (the runtime demotes them too), so fall back to replay.
+        // so anything landing here is out-of-band — an unguarded sleep signal
+        // or a divergence. Those boundaries are unretainable (the runtime
+        // demotes them too), so fall back to replay.
         state = { type: 'replay' };
         return;
       case 'replay':

packages/core/src/workflow/attribute-dispatcher.ts

@@ -2,12 +2,9 @@ import { ReplayDivergenceError } from '@workflow/errors';
 import { withResolvers } from '@workflow/utils';
 import type { AttributeChange } from '@workflow/world';
 import { EventConsumerResult } from '../events-consumer.js';
+import type { AttributeInvocationQueueItem } from '../global.js';
 import {
-  type AttributeInvocationQueueItem,
-  WorkflowSuspension,
-} from '../global.js';
-import {
-  scheduleWhenIdle,
+  scheduleWorkflowSuspension,
   type WorkflowOrchestratorContext,
 } from '../private.js';
 
@@ -30,11 +27,7 @@ export function createSetAttributes(ctx: WorkflowOrchestratorContext) {
 
     ctx.eventsConsumer.subscribe((event) => {
       if (!event) {
-        scheduleWhenIdle(ctx, () => {
-          ctx.onWorkflowError(
-            new WorkflowSuspension(ctx.invocationsQueue, ctx.globalThis)
-          );
-        });
+        scheduleWorkflowSuspension(ctx);
         return EventConsumerResult.NotConsumed;
       }