Make getWritable work inside steps (#222)

Author: pranaygp

.changeset/smart-queens-post.md

@@ -0,0 +1,5 @@
+---
+"@workflow/core": patch
+---
+
+Add support for getWritable directly in step functions

docs/content/docs/api-reference/workflow/get-writable.mdx

@@ -6,12 +6,12 @@ import { generateDefinition } from "@/lib/tsdoc"
 
 # `getWritable()`
 
-Retrieves the current workflow run's default writable stream. The writable stream is intended to be passed as an argument to steps which can write to it. Chunks written to this stream can be read outside the workflow by using the `readable` property of the [`Run` object](/docs/api-reference/workflow-api/get-run).
+Retrieves the current workflow run's default writable stream. The writable stream can be used in both workflow and step functions to write data that can be read outside the workflow by using the `readable` property of the [`Run` object](/docs/api-reference/workflow-api/get-run).
 
-Use this function in your workflows to produce streaming output that can be consumed by clients in real-time.
+Use this function in your workflows and steps to produce streaming output that can be consumed by clients in real-time.
 
 <Callout type="warn">
-This function can only be called inside a workflow function (functions with `"use workflow"` directive)
+This function can only be called inside a workflow or step function (functions with `"use workflow"` or `"use step"` directive)
 </Callout>
 
 ```typescript lineNumbers
@@ -56,7 +56,9 @@ Returns a `WritableStream<W>` where `W` is the type of data you plan to write to
 
 ## Good to Know
 
-- The stream should typically be passed to step functions for writing.
+- The stream can be obtained from either workflow or step functions using the same `getWritable()` call.
+- When called from a workflow, the stream can be passed as an argument to steps.
+- When called from a step, it retrieves the same workflow-scoped stream directly.
 - Always release the writer lock after writing to prevent resource leaks.
 - The stream can write binary data (using `TextEncoder`) or structured objects.
 - Remember to close the stream when finished to signal completion.
@@ -100,6 +102,89 @@ async function stepCloseOutputStream(writable: WritableStream) {
 }
 ```
 
+### Calling `getWritable()` Inside Steps
+
+You can also call `getWritable()` directly inside step functions without passing it as a parameter:
+
+```typescript lineNumbers
+import { sleep, getWritable } from 'workflow';
+
+export async function outputStreamFromStepWorkflow() {
+  "use workflow";
+
+  // No need to create or pass the stream - steps can get it themselves
+  await sleep("1s");
+  await stepWithOutputStreamInside();
+  await sleep("1s");
+  await stepCloseOutputStreamInside();
+
+  return 'done';
+}
+
+async function stepWithOutputStreamInside() {
+  "use step";
+
+  // Call getWritable() directly inside the step // [!code highlight]
+  const writable = getWritable(); // [!code highlight]
+  const writer = writable.getWriter();
+
+  await writer.write(new TextEncoder().encode('Hello from step!'));
+  writer.releaseLock();
+}
+
+async function stepCloseOutputStreamInside() {
+  "use step";
+
+  // Call getWritable() to get the same stream // [!code highlight]
+  const writable = getWritable(); // [!code highlight]
+  await writable.close();
+}
+```
+
+### Using Namespaced Streams in Steps
+
+You can also use namespaced streams when calling `getWritable()` from steps:
+
+```typescript lineNumbers
+import { getWritable } from 'workflow';
+
+export async function multiStreamWorkflow() {
+  "use workflow";
+
+  // Steps will access both streams by namespace
+  await writeToDefaultStream();
+  await writeToNamedStream();
+  await closeStreams();
+
+  return 'done';
+}
+
+async function writeToDefaultStream() {
+  "use step";
+
+  const writable = getWritable(); // Default stream
+  const writer = writable.getWriter();
+  await writer.write({ message: 'Default stream data' });
+  writer.releaseLock();
+}
+
+async function writeToNamedStream() {
+  "use step";
+
+  const writable = getWritable({ namespace: 'logs' }); // [!code highlight]
+  const writer = writable.getWriter();
+  await writer.write({ log: 'Named stream data' });
+  writer.releaseLock();
+}
+
+async function closeStreams() {
+  "use step";
+
+  await getWritable().close(); // Close default stream
+  await getWritable({ namespace: 'logs' }).close(); // Close named stream
+}
+```
+
 ### Advanced Chat Streaming
 
 Here's a more complex example showing how you might stream AI chat responses:

packages/core/e2e/e2e.test.ts

@@ -412,6 +412,65 @@ describe('e2e', () => {
     expect(returnValue).toEqual('done');
   });
 
+  test(
+    'outputStreamInsideStepWorkflow - getWritable() called inside step functions',
+    { timeout: 60_000 },
+    async () => {
+      const run = await triggerWorkflow('outputStreamInsideStepWorkflow', []);
+      const stream = await fetch(
+        `${deploymentUrl}/api/trigger?runId=${run.runId}&output-stream=1`
+      );
+      const namedStream = await fetch(
+        `${deploymentUrl}/api/trigger?runId=${run.runId}&output-stream=step-ns`
+      );
+      const textDecoderStream = new TextDecoderStream();
+      stream.body?.pipeThrough(textDecoderStream);
+      const reader = textDecoderStream.readable.getReader();
+
+      const namedTextDecoderStream = new TextDecoderStream();
+      namedStream.body?.pipeThrough(namedTextDecoderStream);
+      const namedReader = namedTextDecoderStream.readable.getReader();
+
+      // First message from default stream
+      const r1 = await reader.read();
+      assert(r1.value);
+      const chunk1 = JSON.parse(r1.value);
+      const binaryData1 = Buffer.from(chunk1.data, 'base64');
+      expect(binaryData1.toString()).toEqual('Hello from step!');
+
+      // First message from named stream
+      const r1Named = await namedReader.read();
+      assert(r1Named.value);
+      const chunk1Named = JSON.parse(r1Named.value);
+      expect(chunk1Named).toEqual({
+        message: 'Hello from named stream in step!',
+      });
+
+      // Second message from default stream
+      const r2 = await reader.read();
+      assert(r2.value);
+      const chunk2 = JSON.parse(r2.value);
+      const binaryData2 = Buffer.from(chunk2.data, 'base64');
+      expect(binaryData2.toString()).toEqual('Second message');
+
+      // Second message from named stream
+      const r2Named = await namedReader.read();
+      assert(r2Named.value);
+      const chunk2Named = JSON.parse(r2Named.value);
+      expect(chunk2Named).toEqual({ counter: 42 });
+
+      // Verify streams are closed
+      const r3 = await reader.read();
+      expect(r3.done).toBe(true);
+
+      const r3Named = await namedReader.read();
+      expect(r3Named.done).toBe(true);
+
+      const returnValue = await getWorkflowReturnValue(run.runId);
+      expect(returnValue).toEqual('done');
+    }
+  );
+
   test('fetchWorkflow', { timeout: 60_000 }, async () => {
     const run = await triggerWorkflow('fetchWorkflow', []);
     const returnValue = await getWorkflowReturnValue(run.runId);

packages/core/src/index.ts

@@ -34,4 +34,7 @@ export {
   type WorkflowMetadata,
 } from './step/get-workflow-metadata.js';
 export { sleep } from './sleep.js';
-export { getWritable } from './writable-stream.js';
+export {
+  getWritable,
+  type WorkflowWritableStreamOptions,
+} from './step/writable-stream.js';

packages/core/src/runtime.ts

@@ -674,6 +674,7 @@ export const stepEntrypoint =
                     ? `https://${process.env.VERCEL_URL}`
                     : `http://localhost:${process.env.PORT || 3000}`,
                 },
+                ops,
               },
               () => stepFn(...args)
             );

packages/core/src/step/context-storage.ts

@@ -5,4 +5,5 @@ import type { StepMetadata } from './get-step-metadata.js';
 export const contextStorage = /* @__PURE__ */ new AsyncLocalStorage<{
   stepMetadata: StepMetadata;
   workflowMetadata: WorkflowMetadata;
+  ops: Promise<any>[];
 }>();

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

@@ -0,0 +1,51 @@
+import {
+  WorkflowServerWritableStream,
+  getSerializeStream,
+  getExternalReducers,
+} from '../serialization.js';
+import { getWorkflowRunStreamId } from '../util.js';
+import type { WorkflowWritableStreamOptions } from '../writable-stream.js';
+import { contextStorage } from './context-storage.js';
+
+export type { WorkflowWritableStreamOptions };
+
+/**
+ * Retrieves a writable stream that is associated with the current workflow.
+ *
+ * The writable stream is intended to be used within step functions to write
+ * data that can be read outside the workflow by using the readable method of getRun.
+ *
+ * @param options - Optional configuration for the writable stream
+ * @returns The writable stream associated with the current workflow run
+ * @throws Error if called outside a workflow or step function
+ */
+export function getWritable<W = any>(
+  options: WorkflowWritableStreamOptions = {}
+): WritableStream<W> {
+  const ctx = contextStorage.getStore();
+  if (!ctx) {
+    throw new Error(
+      '`getWritable()` can only be called inside a workflow or step function'
+    );
+  }
+
+  const { namespace } = options;
+  const name = getWorkflowRunStreamId(
+    ctx.workflowMetadata.workflowRunId,
+    namespace
+  );
+
+  // Create a transform stream that serializes chunks and pipes to the workflow server
+  const serialize = getSerializeStream(
+    getExternalReducers(globalThis, ctx.ops)
+  );
+
+  // 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))
+  );
+
+  // Return the writable side of the transform stream
+  return serialize.writable;
+}

packages/core/src/writable-stream.ts

@@ -12,18 +12,19 @@ export interface WorkflowWritableStreamOptions {
 /**
  * Retrieves a writable stream that is associated with the current workflow.
  *
- * The writable stream is intended to be passed as an argument to steps which can
- * write to it. Chunks written to this stream can be read outside the workflow
- * by using readable method of getRun.
+ * The writable stream can be used in both workflow and step functions.
+ * In workflows, it can be passed as an argument to steps. In steps, it can
+ * be called directly. Chunks written to this stream can be read outside the
+ * workflow by using the readable method of getRun.
  *
- * @note This function can only be called inside a workflow function.
+ * @note This function can only be called inside a workflow or step function.
  * @returns The writable stream.
  */
 export function getWritable<W = any>(
   // @ts-expect-error `options` is here for types/docs
   options: WorkflowWritableStreamOptions = {}
 ): WritableStream<W> {
   throw new Error(
-    '`getWritable()` can only be called inside a workflow function'
+    '`getWritable()` can only be called inside a workflow or step function'
   );
 }

workbench/example/workflows/99_e2e.ts

@@ -289,6 +289,54 @@ export async function outputStreamWorkflow() {
 
 //////////////////////////////////////////////////////////
 
+async function stepWithOutputStreamInsideStep(text: string) {
+  'use step';
+  // Call getWritable directly inside the step function
+  const writable = getWritable();
+  const writer = writable.getWriter();
+  await writer.write(new TextEncoder().encode(text));
+  writer.releaseLock();
+}
+
+async function stepWithNamedOutputStreamInsideStep(
+  namespace: string,
+  obj: any
+) {
+  'use step';
+  // Call getWritable with namespace directly inside the step function
+  const writable = getWritable({ namespace });
+  const writer = writable.getWriter();
+  await writer.write(obj);
+  writer.releaseLock();
+}
+
+async function stepCloseOutputStreamInsideStep(namespace?: string) {
+  'use step';
+  // Call getWritable directly inside the step function and close it
+  const writable = getWritable({ namespace });
+  await writable.close();
+}
+
+export async function outputStreamInsideStepWorkflow() {
+  'use workflow';
+  await sleep('1s');
+  await stepWithOutputStreamInsideStep('Hello from step!');
+  await sleep('1s');
+  await stepWithNamedOutputStreamInsideStep('step-ns', {
+    message: 'Hello from named stream in step!',
+  });
+  await sleep('1s');
+  await stepWithOutputStreamInsideStep('Second message');
+  await sleep('1s');
+  await stepWithNamedOutputStreamInsideStep('step-ns', { counter: 42 });
+  await sleep('1s');
+  await stepCloseOutputStreamInsideStep();
+  await stepCloseOutputStreamInsideStep('step-ns');
+  return 'done';
+}
+
+//////////////////////////////////////////////////////////
+
 export async function fetchWorkflow() {
   'use workflow';
   const response = await fetch('https://jsonplaceholder.typicode.com/todos/1');

workbench/nextjs-turbopack/tsconfig.json

@@ -11,7 +11,7 @@
     "moduleResolution": "bundler",
     "resolveJsonModule": true,
     "isolatedModules": true,
-    "jsx": "preserve",
+    "jsx": "react-jsx",
     "incremental": true,
     "plugins": [
       {
@@ -30,7 +30,8 @@
     "**/*.tsx",
     ".next/types/**/*.ts",
     "next-env.d.ts",
-    "**/*.mts"
+    "**/*.mts",
+    ".next/dev/types/**/*.ts"
   ],
   "exclude": ["node_modules"]
 }