feat(streams): add inputStream.waitWithWarmup(), warm timeout config in sidebar, preload payload option

ericallam · ericallam · commit 0415ccbc782d · 2026-03-09T17:23:37.000Z
diff --git a/packages/core/src/v3/realtimeStreams/types.ts b/packages/core/src/v3/realtimeStreams/types.ts
@@ -193,6 +193,14 @@ export type RealtimeDefinedInputStream<TData> = {
    * Uses a waitpoint token internally. Can only be called inside a task.run().
    */
   wait: (options?: InputStreamWaitOptions) => ManualWaitpointPromise<TData>;
+  /**
+   * Wait for data with a warm phase before suspending.
+   *
+   * Keeps the task warm (active, using compute) for `warmTimeoutInSeconds`,
+   * then suspends via `.wait()` if no data arrives. If data arrives during
+   * the warm phase the task responds instantly without suspending.
+   */
+  waitWithWarmup: (options: InputStreamWaitWithWarmupOptions) => Promise<{ ok: true; output: TData } | { ok: false; error?: any }>;
   /**
    * Send data to this input stream on a specific run.
    * This is used from outside the task (e.g., from your backend or another task).
@@ -249,6 +257,15 @@ export type InputStreamWaitOptions = {
   spanName?: string;
 };
 
+export type InputStreamWaitWithWarmupOptions = {
+  /** Seconds to keep the task warm before suspending. */
+  warmTimeoutInSeconds: number;
+  /** Maximum time to wait after suspending (duration string, e.g. "1h"). */
+  timeout?: string;
+  /** Override the default span name for the outer operation. */
+  spanName?: string;
+};
+
 export type InferInputStreamType<T> = T extends RealtimeDefinedInputStream<infer TData>
   ? TData
   : unknown;
diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts
@@ -320,6 +320,8 @@ export type ChatTaskWirePayload<TMessage extends UIMessage = UIMessage, TMetadat
   continuation?: boolean;
   /** The run ID of the previous run (only set when `continuation` is true). */
   previousRunId?: string;
+  /** Override warm timeout for this run (seconds). Set by transport.preload(). */
+  warmTimeoutInSeconds?: number;
 };
 
 /**
@@ -941,40 +943,26 @@ function chatTask<
 
           // Wait for the first real message — use preload-specific timeouts if configured
           const effectivePreloadWarmTimeout =
-            (metadata.get(WARM_TIMEOUT_METADATA_KEY) as number | undefined)
+            payload.warmTimeoutInSeconds
             ?? preloadWarmTimeoutInSeconds
             ?? warmTimeoutInSeconds;
 
-          let firstMessage: ChatTaskWirePayload | undefined;
+          const effectivePreloadTimeout =
+            (metadata.get(TURN_TIMEOUT_METADATA_KEY) as string | undefined)
+            ?? preloadTimeout
+            ?? turnTimeout;
 
-          if (effectivePreloadWarmTimeout > 0) {
-            const warm = await messagesInput.once({
-              timeoutMs: effectivePreloadWarmTimeout * 1000,
-              spanName: "preload wait (warm)",
-            });
+          const preloadResult = await messagesInput.waitWithWarmup({
+            warmTimeoutInSeconds: effectivePreloadWarmTimeout,
+            timeout: effectivePreloadTimeout,
+            spanName: "waiting for first message",
+          });
 
-            if (warm.ok) {
-              firstMessage = warm.output;
-            }
+          if (!preloadResult.ok) {
+            return; // Timed out waiting for first message — end run
           }
 
-          if (!firstMessage) {
-            const effectivePreloadTimeout =
-              (metadata.get(TURN_TIMEOUT_METADATA_KEY) as string | undefined)
-              ?? preloadTimeout
-              ?? turnTimeout;
-
-            const suspended = await messagesInput.wait({
-              timeout: effectivePreloadTimeout,
-              spanName: "preload wait (suspended)",
-            });
-
-            if (!suspended.ok) {
-              return; // Timed out waiting for first message — end run
-            }
-
-            firstMessage = suspended.output;
-          }
+          let firstMessage = preloadResult.output;
 
           currentWirePayload = firstMessage;
         }
@@ -1335,35 +1323,19 @@ function chatTask<
                 return "continue";
               }
 
-              // Phase 1: Keep the run warm for quick response to the next message.
-              // The run stays active (using compute) during this window.
+              // Wait for the next message — stay warm briefly, then suspend
               const effectiveWarmTimeout =
                 (metadata.get(WARM_TIMEOUT_METADATA_KEY) as number | undefined) ?? warmTimeoutInSeconds;
-
-              if (effectiveWarmTimeout > 0) {
-                const warm = await messagesInput.once({
-                  timeoutMs: effectiveWarmTimeout * 1000,
-                  spanName: "waiting (warm)",
-                });
-
-                if (warm.ok) {
-                  // Message arrived while warm — respond instantly
-                  currentWirePayload = warm.output;
-                  return "continue";
-                }
-              }
-
-              // Phase 2: Suspend the task (frees compute) until the next message arrives
               const effectiveTurnTimeout =
                 (metadata.get(TURN_TIMEOUT_METADATA_KEY) as string | undefined) ?? turnTimeout;
 
-              const next = await messagesInput.wait({
+              const next = await messagesInput.waitWithWarmup({
+                warmTimeoutInSeconds: effectiveWarmTimeout,
                 timeout: effectiveTurnTimeout,
-                spanName: "waiting (suspended)",
+                spanName: "waiting for next message",
               });
 
               if (!next.ok) {
-                // Timed out waiting for the next message — end the conversation
                 return "exit";
               }
 
diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts
@@ -464,7 +464,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
    *
    * No-op if a session already exists for this chatId.
    */
-  async preload(chatId: string): Promise<void> {
+  async preload(chatId: string, options?: { warmTimeoutInSeconds?: number }): Promise<void> {
     // Don't preload if session already exists
     if (this.sessions.get(chatId)?.runId) return;
 
@@ -473,6 +473,9 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
       chatId,
       trigger: "preload" as const,
       metadata: this.defaultMetadata,
+      ...(options?.warmTimeoutInSeconds !== undefined
+        ? { warmTimeoutInSeconds: options.warmTimeoutInSeconds }
+        : {}),
     };
 
     const currentToken = await this.resolveAccessToken();
diff --git a/packages/trigger-sdk/src/v3/streams.ts b/packages/trigger-sdk/src/v3/streams.ts
@@ -25,6 +25,7 @@ import {
   InputStreamOncePromise,
   type InputStreamOnceResult,
   type InputStreamWaitOptions,
+  type InputStreamWaitWithWarmupOptions,
   type SendInputStreamOptions,
   type InferInputStreamType,
   type StreamWriteResult,
@@ -767,6 +768,7 @@ function input<TData>(opts: { id: string }): RealtimeDefinedInputStream<TData> {
           const result = await tracer.startActiveSpan(
             options?.spanName ?? `inputStream.wait()`,
             async (span) => {
+
               // 1. Block the run on the waitpoint
               const waitResponse = await apiClient.waitForWaitpointToken({
                 runFriendlyId: ctx.run.id,
@@ -786,7 +788,7 @@ function input<TData>(opts: { id: string }): RealtimeDefinedInputStream<TData> {
               // 3. Suspend the task
               const waitResult = await runtime.waitUntil(response.waitpointId);
 
-              // 3. Parse the output
+              // 4. Parse the output
               const data =
                 waitResult.output !== undefined
                   ? await conditionallyImportAndParsePacket(
@@ -840,6 +842,45 @@ function input<TData>(opts: { id: string }): RealtimeDefinedInputStream<TData> {
         }
       });
     },
+    async waitWithWarmup(options) {
+      const self = this;
+      const spanName = options.spanName ?? `inputStream.waitWithWarmup()`;
+
+      return tracer.startActiveSpan(
+        spanName,
+        async (span) => {
+          // Warm phase: keep compute alive
+          if (options.warmTimeoutInSeconds > 0) {
+            const warm = await inputStreams.once(opts.id, {
+              timeoutMs: options.warmTimeoutInSeconds * 1000,
+            });
+            if (warm.ok) {
+              span.setAttribute("wait.resolved", "warm");
+              return { ok: true as const, output: warm.output as TData };
+            }
+          }
+
+          // Cold phase: suspend via .wait() — creates a child span
+          span.setAttribute("wait.resolved", "suspended");
+          const waitResult = await self.wait({
+            timeout: options.timeout,
+            spanName: "suspended",
+          });
+
+          return waitResult;
+        },
+        {
+          attributes: {
+            [SemanticInternalAttributes.STYLE_ICON]: "streams",
+            streamId: opts.id,
+            ...accessoryAttributes({
+              items: [{ text: opts.id, variant: "normal" }],
+              style: "codepath",
+            }),
+          },
+        }
+      );
+    },
     async send(runId, data, options) {
       return tracer.startActiveSpan(
         `inputStream.send()`,
diff --git a/references/ai-chat/src/components/chat-app.tsx b/references/ai-chat/src/components/chat-app.tsx
@@ -52,6 +52,7 @@ export function ChatApp({
   // Model for new chats (before first message is sent)
   const [newChatModel, setNewChatModel] = useState(DEFAULT_MODEL);
   const [preloadEnabled, setPreloadEnabled] = useState(true);
+  const [warmTimeoutInSeconds, setWarmTimeoutInSeconds] = useState(60);
 
   const handleSessionChange = useCallback(
     (chatId: string, session: SessionInfo | null) => {
@@ -101,7 +102,7 @@ export function ChatApp({
     setNewChatModel(DEFAULT_MODEL);
     if (preloadEnabled) {
       // Eagerly start the run — onPreload fires immediately for initialization
-      transport.preload(id);
+      transport.preload(id, { warmTimeoutInSeconds });
     }
   }
 
@@ -154,6 +155,8 @@ export function ChatApp({
         onDeleteChat={handleDeleteChat}
         preloadEnabled={preloadEnabled}
         onPreloadChange={setPreloadEnabled}
+        warmTimeoutInSeconds={warmTimeoutInSeconds}
+        onWarmTimeoutChange={setWarmTimeoutInSeconds}
       />
       <div className="flex-1">
         {activeChatId ? (
diff --git a/references/ai-chat/src/components/chat-sidebar.tsx b/references/ai-chat/src/components/chat-sidebar.tsx
@@ -26,6 +26,8 @@ type ChatSidebarProps = {
   onDeleteChat: (id: string) => void;
   preloadEnabled: boolean;
   onPreloadChange: (enabled: boolean) => void;
+  warmTimeoutInSeconds: number;
+  onWarmTimeoutChange: (seconds: number) => void;
 };
 
 export function ChatSidebar({
@@ -36,6 +38,8 @@ export function ChatSidebar({
   onDeleteChat,
   preloadEnabled,
   onPreloadChange,
+  warmTimeoutInSeconds,
+  onWarmTimeoutChange,
 }: ChatSidebarProps) {
   const sorted = [...chats].sort((a, b) => b.updatedAt - a.updatedAt);
 
@@ -82,7 +86,7 @@ export function ChatSidebar({
         ))}
       </div>
 
-      <div className="shrink-0 border-t border-gray-200 px-3 py-2.5">
+      <div className="shrink-0 border-t border-gray-200 px-3 py-2.5 space-y-2">
         <label className="flex items-center gap-2 text-xs text-gray-500 cursor-pointer select-none">
           <input
             type="checkbox"
@@ -92,6 +96,18 @@ export function ChatSidebar({
           />
           Preload new chats
         </label>
+        <div className="flex items-center gap-2 text-xs text-gray-500">
+          <span className="shrink-0">Warm timeout</span>
+          <input
+            type="number"
+            min={0}
+            step={5}
+            value={warmTimeoutInSeconds}
+            onChange={(e) => onWarmTimeoutChange(Number(e.target.value))}
+            className="w-16 rounded border border-gray-300 px-1.5 py-0.5 text-xs text-gray-600 outline-none focus:border-blue-500"
+          />
+          <span>s</span>
+        </div>
       </div>
     </div>
   );
