Add result type and .unwrap() helper to inputStream.once().

ericallam · ericallam · commit f6f2b772566e · 2026-02-27T19:46:19.000Z
diff --git a/packages/core/src/v3/inputStreams/index.ts b/packages/core/src/v3/inputStreams/index.ts
@@ -1,6 +1,6 @@
 import { getGlobal, registerGlobal } from "../utils/globals.js";
 import { NoopInputStreamManager } from "./noopManager.js";
-import { InputStreamManager } from "./types.js";
+import { InputStreamManager, InputStreamOncePromise } from "./types.js";
 import { InputStreamOnceOptions } from "../realtimeStreams/types.js";
 
 const API_NAME = "input-streams";
@@ -39,7 +39,7 @@ export class InputStreamsAPI implements InputStreamManager {
     return this.#getManager().on(streamId, handler);
   }
 
-  public once(streamId: string, options?: InputStreamOnceOptions): Promise<unknown> {
+  public once(streamId: string, options?: InputStreamOnceOptions): InputStreamOncePromise<unknown> {
     return this.#getManager().once(streamId, options);
   }
 
diff --git a/packages/core/src/v3/inputStreams/manager.ts b/packages/core/src/v3/inputStreams/manager.ts
@@ -1,11 +1,16 @@
 import { ApiClient } from "../apiClient/index.js";
-import { InputStreamManager } from "./types.js";
+import {
+  InputStreamManager,
+  InputStreamOncePromise,
+  InputStreamOnceResult,
+  InputStreamTimeoutError,
+} from "./types.js";
 import { InputStreamOnceOptions } from "../realtimeStreams/types.js";
 
 type InputStreamHandler = (data: unknown) => void | Promise<void>;
 
 type OnceWaiter = {
-  resolve: (data: unknown) => void;
+  resolve: (result: InputStreamOnceResult<unknown>) => void;
   reject: (error: Error) => void;
   timeoutHandle?: ReturnType<typeof setTimeout>;
 };
@@ -72,7 +77,7 @@ export class StandardInputStreamManager implements InputStreamManager {
     };
   }
 
-  once(streamId: string, options?: InputStreamOnceOptions): Promise<unknown> {
+  once(streamId: string, options?: InputStreamOnceOptions): InputStreamOncePromise<unknown> {
     this.#requireV2Streams();
 
     // Lazily connect a tail for this stream
@@ -85,10 +90,12 @@ export class StandardInputStreamManager implements InputStreamManager {
       if (buffered.length === 0) {
         this.buffer.delete(streamId);
       }
-      return Promise.resolve(data);
+      return new InputStreamOncePromise((resolve) => {
+        resolve({ ok: true, output: data });
+      });
     }
 
-    return new Promise<unknown>((resolve, reject) => {
+    return new InputStreamOncePromise<unknown>((resolve, reject) => {
       const waiter: OnceWaiter = { resolve, reject };
 
       // Handle abort signal
@@ -110,11 +117,14 @@ export class StandardInputStreamManager implements InputStreamManager {
         );
       }
 
-      // Handle timeout
+      // Handle timeout — resolve with error result instead of rejecting
       if (options?.timeoutMs) {
         waiter.timeoutHandle = setTimeout(() => {
           this.#removeOnceWaiter(streamId, waiter);
-          reject(new Error(`Timeout waiting for input stream "${streamId}" after ${options.timeoutMs}ms`));
+          resolve({
+            ok: false,
+            error: new InputStreamTimeoutError(streamId, options.timeoutMs!),
+          });
         }, options.timeoutMs);
       }
 
@@ -256,7 +266,7 @@ export class StandardInputStreamManager implements InputStreamManager {
       if (waiter.timeoutHandle) {
         clearTimeout(waiter.timeoutHandle);
       }
-      waiter.resolve(data);
+      waiter.resolve({ ok: true, output: data });
       // Also invoke persistent handlers
       this.#invokeHandlers(streamId, data);
       return;
diff --git a/packages/core/src/v3/inputStreams/noopManager.ts b/packages/core/src/v3/inputStreams/noopManager.ts
@@ -1,4 +1,4 @@
-import { InputStreamManager } from "./types.js";
+import { InputStreamManager, InputStreamOncePromise } from "./types.js";
 import { InputStreamOnceOptions } from "../realtimeStreams/types.js";
 
 export class NoopInputStreamManager implements InputStreamManager {
@@ -8,8 +8,8 @@ export class NoopInputStreamManager implements InputStreamManager {
     return { off: () => {} };
   }
 
-  once(_streamId: string, _options?: InputStreamOnceOptions): Promise<unknown> {
-    return new Promise(() => {
+  once(_streamId: string, _options?: InputStreamOnceOptions): InputStreamOncePromise<unknown> {
+    return new InputStreamOncePromise(() => {
       // Never resolves in noop mode
     });
   }
diff --git a/packages/core/src/v3/inputStreams/types.ts b/packages/core/src/v3/inputStreams/types.ts
@@ -1,5 +1,42 @@
 import { InputStreamOnceOptions } from "../realtimeStreams/types.js";
 
+export class InputStreamTimeoutError extends Error {
+  constructor(
+    public readonly streamId: string,
+    public readonly timeoutMs: number
+  ) {
+    super(`Timeout waiting for input stream "${streamId}" after ${timeoutMs}ms`);
+    this.name = "InputStreamTimeoutError";
+  }
+}
+
+export type InputStreamOnceResult<TData> =
+  | { ok: true; output: TData }
+  | { ok: false; error: InputStreamTimeoutError };
+
+export class InputStreamOncePromise<TData> extends Promise<InputStreamOnceResult<TData>> {
+  constructor(
+    executor: (
+      resolve: (
+        value: InputStreamOnceResult<TData> | PromiseLike<InputStreamOnceResult<TData>>
+      ) => void,
+      reject: (reason?: any) => void
+    ) => void
+  ) {
+    super(executor);
+  }
+
+  unwrap(): Promise<TData> {
+    return this.then((result) => {
+      if (result.ok) {
+        return result.output;
+      } else {
+        throw result.error;
+      }
+    });
+  }
+}
+
 export interface InputStreamManager {
   /**
    * Set the current run ID and streams version. The tail connection will be
@@ -15,8 +52,10 @@ export interface InputStreamManager {
 
   /**
    * Wait for the next piece of data on the given input stream.
+   * Returns a result object `{ ok, output }` or `{ ok, error }`.
+   * Chain `.unwrap()` to get the data directly or throw on timeout.
    */
-  once(streamId: string, options?: InputStreamOnceOptions): Promise<unknown>;
+  once(streamId: string, options?: InputStreamOnceOptions): InputStreamOncePromise<unknown>;
 
   /**
    * Non-blocking peek at the most recent data on the given input stream.
diff --git a/packages/core/src/v3/realtimeStreams/types.ts b/packages/core/src/v3/realtimeStreams/types.ts
@@ -1,4 +1,7 @@
 import { AnyZodFetchOptions, ApiRequestOptions } from "../apiClient/core.js";
+import type { InputStreamOncePromise } from "../inputStreams/types.js";
+export { InputStreamOncePromise, InputStreamTimeoutError } from "../inputStreams/types.js";
+export type { InputStreamOnceResult } from "../inputStreams/types.js";
 import { AsyncIterableStream } from "../streams/asyncIterableStream.js";
 import { Prettify } from "../types/utils.js";
 import type { ManualWaitpointPromise } from "../waitpoints/index.js";
@@ -162,9 +165,10 @@ export type RealtimeDefinedInputStream<TData> = {
   on: (handler: (data: TData) => void | Promise<void>) => InputStreamSubscription;
   /**
    * Wait for the next piece of data on this input stream.
-   * Resolves with the data when it arrives.
+   * Returns a result object `{ ok, output }` or `{ ok, error }`.
+   * Chain `.unwrap()` to get the data directly or throw on timeout.
    */
-  once: (options?: InputStreamOnceOptions) => Promise<TData>;
+  once: (options?: InputStreamOnceOptions) => InputStreamOncePromise<TData>;
   /**
    * Non-blocking peek at the most recent data received on this input stream.
    * Returns `undefined` if no data has been received yet.
diff --git a/packages/trigger-sdk/src/v3/streams.ts b/packages/trigger-sdk/src/v3/streams.ts
@@ -22,6 +22,8 @@ import {
   type RealtimeDefinedInputStream,
   type InputStreamSubscription,
   type InputStreamOnceOptions,
+  InputStreamOncePromise,
+  type InputStreamOnceResult,
   type InputStreamWaitOptions,
   type SendInputStreamOptions,
   type InferInputStreamType,
@@ -706,26 +708,33 @@ function input<TData>(opts: { id: string }): RealtimeDefinedInputStream<TData> {
       const ctx = taskContext.ctx;
       const runId = ctx?.run.id;
 
-      return tracer.startActiveSpan(
-        `inputStream.once()`,
-        async () => {
-          return inputStreams.once(opts.id, options) as Promise<TData>;
-        },
-        {
-          attributes: {
-            [SemanticInternalAttributes.STYLE_ICON]: "streams",
-            [SemanticInternalAttributes.ENTITY_TYPE]: "input-stream",
-            ...(runId
-              ? { [SemanticInternalAttributes.ENTITY_ID]: `${runId}:${opts.id}` }
-              : {}),
-            streamId: opts.id,
-            ...accessoryAttributes({
-              items: [{ text: opts.id, variant: "normal" }],
-              style: "codepath",
-            }),
-          },
-        }
-      );
+      const innerPromise = inputStreams.once(opts.id, options);
+
+      return new InputStreamOncePromise<TData>((resolve, reject) => {
+        tracer
+          .startActiveSpan(
+            `inputStream.once()`,
+            async () => {
+              const result = await innerPromise;
+              resolve(result as InputStreamOnceResult<TData>);
+            },
+            {
+              attributes: {
+                [SemanticInternalAttributes.STYLE_ICON]: "streams",
+                [SemanticInternalAttributes.ENTITY_TYPE]: "input-stream",
+                ...(runId
+                  ? { [SemanticInternalAttributes.ENTITY_ID]: `${runId}:${opts.id}` }
+                  : {}),
+                streamId: opts.id,
+                ...accessoryAttributes({
+                  items: [{ text: opts.id, variant: "normal" }],
+                  style: "codepath",
+                }),
+              },
+            }
+          )
+          .catch(reject);
+      });
     },
     peek() {
       return inputStreams.peek(opts.id) as TData | undefined;
diff --git a/references/hello-world/src/trigger/inputStreams.ts b/references/hello-world/src/trigger/inputStreams.ts
@@ -10,24 +10,34 @@ const messageStream = streams.input<{ text: string }>({ id: "messages" });
 /**
  * Coordinator task that exercises all input stream patterns end-to-end.
  *
- * 1. .once() — trigger a child, send it data via SSE tail, poll until complete
- * 2. .on()   — trigger a child, send it multiple messages, poll until complete
- * 3. .wait() — trigger a child, send it data (completes its waitpoint), poll until complete
- * 4. .wait() race — send data before child calls .wait(), verify race handling
+ * 1a. .once().unwrap() — trigger a child, send data, verify unwrap returns TData
+ * 1b. .once() result  — trigger a child, send data, verify result object { ok, output }
+ * 2.  .on()            — trigger a child, send it multiple messages, poll until complete
+ * 3.  .wait()          — trigger a child, send it data (completes its waitpoint), poll until complete
+ * 4.  .wait() race     — send data before child calls .wait(), verify race handling
  */
 export const inputStreamCoordinator = task({
   id: "input-stream-coordinator",
   run: async () => {
     const results: Record<string, unknown> = {};
 
-    // --- Test 1: .once() ----
-    logger.info("Test 1: .once()");
-    const onceHandle = await inputStreamOnce.trigger({});
+    // --- Test 1a: .once() with .unwrap() ---
+    logger.info("Test 1a: .once().unwrap()");
+    const onceUnwrapHandle = await inputStreamOnceUnwrap.trigger({});
     await wait.for({ seconds: 5 });
-    await approvalStream.send(onceHandle.id, { approved: true, reviewer: "coordinator-once" });
-    const onceRun = await runs.poll(onceHandle, { pollIntervalMs: 1000 });
-    results.once = onceRun.output;
-    logger.info("Test 1 passed", { output: onceRun.output });
+    await approvalStream.send(onceUnwrapHandle.id, { approved: true, reviewer: "coordinator-unwrap" });
+    const onceUnwrapRun = await runs.poll(onceUnwrapHandle, { pollIntervalMs: 1000 });
+    results.onceUnwrap = onceUnwrapRun.output;
+    logger.info("Test 1a passed", { output: onceUnwrapRun.output });
+
+    // --- Test 1b: .once() with result object ---
+    logger.info("Test 1b: .once() result object");
+    const onceResultHandle = await inputStreamOnceResult.trigger({});
+    await wait.for({ seconds: 5 });
+    await approvalStream.send(onceResultHandle.id, { approved: true, reviewer: "coordinator-result" });
+    const onceResultRun = await runs.poll(onceResultHandle, { pollIntervalMs: 1000 });
+    results.onceResult = onceResultRun.output;
+    logger.info("Test 1b passed", { output: onceResultRun.output });
 
     // --- Test 2: .on() with multiple messages ---
     logger.info("Test 2: .on()");
@@ -64,18 +74,40 @@ export const inputStreamCoordinator = task({
 });
 
 /**
- * Uses .once() to wait for a single input stream message.
+ * Uses .once().unwrap() — returns TData directly, throws InputStreamTimeoutError on timeout.
  */
-export const inputStreamOnce = task({
-  id: "input-stream-once",
+export const inputStreamOnceUnwrap = task({
+  id: "input-stream-once-unwrap",
   run: async (_payload: Record<string, never>) => {
-    logger.info("Waiting for approval via .once()");
-    const approval = await approvalStream.once();
+    logger.info("Waiting for approval via .once().unwrap()");
+
+    const approval = await approvalStream.once({ timeoutMs: 30_000 }).unwrap();
+
     logger.info("Received approval", { approval });
     return { approval };
   },
 });
 
+/**
+ * Uses .once() with result object — check result.ok to handle timeout without try/catch.
+ */
+export const inputStreamOnceResult = task({
+  id: "input-stream-once-result",
+  run: async (_payload: Record<string, never>) => {
+    logger.info("Waiting for approval via .once() result object");
+
+    const result = await approvalStream.once({ timeoutMs: 30_000 });
+
+    if (!result.ok) {
+      logger.error("Timed out waiting for approval", { error: result.error.message });
+      return { approval: null, timedOut: true };
+    }
+
+    logger.info("Received approval", { approval: result.output });
+    return { approval: result.output };
+  },
+});
+
 /**
  * Uses .on() to subscribe and collect multiple messages.
  */

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`import { getGlobal, registerGlobal } from "../utils/globals.js";`
`2`	`2`	`import { NoopInputStreamManager } from "./noopManager.js";`
`3`		`-import { InputStreamManager } from "./types.js";`
	`3`	`+import { InputStreamManager, InputStreamOncePromise } from "./types.js";`
`4`	`4`	`import { InputStreamOnceOptions } from "../realtimeStreams/types.js";`
`5`	`5`
`6`	`6`	`const API_NAME = "input-streams";`
`@@ -39,7 +39,7 @@ export class InputStreamsAPI implements InputStreamManager {`
`39`	`39`	`return this.#getManager().on(streamId, handler);`
`40`	`40`	`}`
`41`	`41`
`42`		`- public once(streamId: string, options?: InputStreamOnceOptions): Promise<unknown> {`
	`42`	`+ public once(streamId: string, options?: InputStreamOnceOptions): InputStreamOncePromise<unknown> {`
`43`	`43`	`return this.#getManager().once(streamId, options);`
`44`	`44`	`}`
`45`	`45`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-import { InputStreamManager } from "./types.js";`
	`1`	`+import { InputStreamManager, InputStreamOncePromise } from "./types.js";`
`2`	`2`	`import { InputStreamOnceOptions } from "../realtimeStreams/types.js";`
`3`	`3`
`4`	`4`	`export class NoopInputStreamManager implements InputStreamManager {`
`@@ -8,8 +8,8 @@ export class NoopInputStreamManager implements InputStreamManager {`
`8`	`8`	`return { off: () => {} };`
`9`	`9`	`}`
`10`	`10`
`11`		`- once(_streamId: string, _options?: InputStreamOnceOptions): Promise<unknown> {`
`12`		`- return new Promise(() => {`
	`11`	`+ once(_streamId: string, _options?: InputStreamOnceOptions): InputStreamOncePromise<unknown> {`
	`12`	`+ return new InputStreamOncePromise(() => {`
`13`	`13`	`// Never resolves in noop mode`
`14`	`14`	`});`
`15`	`15`	`}`