Fix step-level getWritable() stream timeout on Vercel (#770) (#1518)

Signed-off-by: ceolinwill <4393133+ceolinwill@users.noreply.github.com>

Author: ceolinwill

.changeset/fix-step-getwritable-timeout.md

@@ -0,0 +1,5 @@
+---
+"@workflow/core": patch
+---
+
+Fix `getWritable()` in step functions to resolve on lock release instead of requiring stream close, preventing Vercel function timeouts

packages/core/src/step/writable-stream.test.ts

@@ -0,0 +1,111 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { LOCK_POLL_INTERVAL_MS } from '../flushable-stream.js';
+
+vi.mock('../runtime/world.js', () => ({
+  getWorld: vi.fn(),
+}));
+
+describe('step-level getWritable', () => {
+  beforeEach(async () => {
+    const mockWorld = {
+      writeToStream: vi.fn().mockResolvedValue(undefined),
+      writeToStreamMulti: vi.fn().mockResolvedValue(undefined),
+      closeStream: vi.fn().mockResolvedValue(undefined),
+    };
+
+    const { getWorld } = await import('../runtime/world.js');
+    (getWorld as ReturnType<typeof vi.fn>).mockReturnValue(mockWorld);
+  });
+
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('ops promise should resolve when writer lock is released (without closing stream)', async () => {
+    const { contextStorage } = await import('./context-storage.js');
+
+    const ops: Promise<void>[] = [];
+    const ctx = {
+      stepMetadata: {
+        stepName: 'test-step',
+        stepId: 'step_001',
+        stepStartedAt: new Date(),
+        attempt: 1,
+      },
+      workflowMetadata: {
+        workflowName: 'test-workflow',
+        workflowRunId: 'wrun_test123',
+        workflowStartedAt: new Date(),
+      },
+      ops,
+      encryptionKey: undefined,
+    };
+
+    const writable = await contextStorage.run(ctx, async () => {
+      const { getWritable } = await import('./writable-stream.js');
+      return getWritable<string>();
+    });
+
+    // Simulate user pattern: write data, then release lock
+    const writer = writable.getWriter();
+    await writer.write('hello');
+    await writer.write('world');
+    writer.releaseLock();
+
+    // Without the fix (.pipeTo()), this hangs because pipeTo only resolves on stream close.
+    // With flushablePipe + pollWritableLock, it resolves once the lock is released.
+    await expect(
+      Promise.race([
+        Promise.all(ops),
+        new Promise((_, r) =>
+          setTimeout(
+            () => r(new Error('ops did not resolve after releaseLock')),
+            LOCK_POLL_INTERVAL_MS * 5 + 200
+          )
+        ),
+      ])
+    ).resolves.not.toThrow();
+  });
+
+  it('ops promise should resolve when stream is explicitly closed', async () => {
+    const { contextStorage } = await import('./context-storage.js');
+
+    const ops: Promise<void>[] = [];
+    const ctx = {
+      stepMetadata: {
+        stepName: 'test-step',
+        stepId: 'step_001',
+        stepStartedAt: new Date(),
+        attempt: 1,
+      },
+      workflowMetadata: {
+        workflowName: 'test-workflow',
+        workflowRunId: 'wrun_test123',
+        workflowStartedAt: new Date(),
+      },
+      ops,
+      encryptionKey: undefined,
+    };
+
+    const writable = await contextStorage.run(ctx, async () => {
+      const { getWritable } = await import('./writable-stream.js');
+      return getWritable<string>();
+    });
+
+    const writer = writable.getWriter();
+    await writer.write('data');
+    await writer.close();
+
+    await expect(
+      Promise.race([
+        Promise.all(ops),
+        new Promise((_, r) =>
+          setTimeout(
+            () => r(new Error('ops did not resolve after close')),
+            LOCK_POLL_INTERVAL_MS * 5 + 200
+          )
+        ),
+      ])
+    ).resolves.not.toThrow();
+  });
+});

packages/core/src/step/writable-stream.ts

@@ -1,3 +1,8 @@
+import {
+  createFlushableState,
+  flushablePipe,
+  pollWritableLock,
+} from '../flushable-stream.js';
 import {
   getExternalReducers,
   getSerializeStream,
@@ -47,11 +52,19 @@ export function getWritable<W = any>(
     ctx.encryptionKey
   );
 
-  // Pipe the serialized data to the workflow server stream
-  // Register this async operation with the runtime's ops array so it's awaited via waitUntil
-  ctx.ops.push(
-    serialize.readable.pipeTo(new WorkflowServerWritableStream(name, runId))
-  );
+  // Use flushable pipe so the ops promise resolves when the user releases
+  // their writer lock, not only when the stream is explicitly closed.
+  // Without this, Vercel functions hang until the runtime timeout because
+  // .pipeTo() only resolves on stream close.
+  const serverWritable = new WorkflowServerWritableStream(name, runId);
+  const state = createFlushableState();
+  ctx.ops.push(state.promise);
+
+  flushablePipe(serialize.readable, serverWritable, state).catch(() => {
+    // Errors are handled via state.reject
+  });
+
+  pollWritableLock(serialize.writable, state);
 
   // Return the writable side of the transform stream
   return serialize.writable;