From 2908b3598d280789bf350b192b25642ae0ce5dee Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Mon, 9 Mar 2026 17:23:37 +0000 Subject: [PATCH] feat(streams): add inputStream.waitWithWarmup(), warm timeout config in sidebar, preload payload option --- packages/core/src/v3/realtimeStreams/types.ts | 17 +++++ packages/trigger-sdk/src/v3/ai.ts | 66 ++++++------------- packages/trigger-sdk/src/v3/chat.ts | 5 +- packages/trigger-sdk/src/v3/streams.ts | 43 +++++++++++- .../ai-chat/src/components/chat-app.tsx | 5 +- .../ai-chat/src/components/chat-sidebar.tsx | 18 ++++- 6 files changed, 103 insertions(+), 51 deletions(-) diff --git a/packages/core/src/v3/realtimeStreams/types.ts b/packages/core/src/v3/realtimeStreams/types.ts index 1b7455ebd..b3c8d8270 100644 --- a/packages/core/src/v3/realtimeStreams/types.ts +++ b/packages/core/src/v3/realtimeStreams/types.ts @@ -193,6 +193,14 @@ export type RealtimeDefinedInputStream = { * Uses a waitpoint token internally. Can only be called inside a task.run(). */ wait: (options?: InputStreamWaitOptions) => ManualWaitpointPromise; + /** + * 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 extends RealtimeDefinedInputStream ? TData : unknown; diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 122e1f716..cdfcae980 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -320,6 +320,8 @@ export type ChatTaskWirePayload 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 index bf269c88e..977679430 100644 --- 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 { * * No-op if a session already exists for this chatId. */ - async preload(chatId: string): Promise { + async preload(chatId: string, options?: { warmTimeoutInSeconds?: number }): Promise { // Don't preload if session already exists if (this.sessions.get(chatId)?.runId) return; @@ -473,6 +473,9 @@ export class TriggerChatTransport implements ChatTransport { 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 index 6bdf862eb..13fd9f822 100644 --- 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(opts: { id: string }): RealtimeDefinedInputStream { 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(opts: { id: string }): RealtimeDefinedInputStream { // 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(opts: { id: string }): RealtimeDefinedInputStream { } }); }, + 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 index 6077fbbea..b7a1d25a5 100644 --- 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} />
{activeChatId ? ( diff --git a/references/ai-chat/src/components/chat-sidebar.tsx b/references/ai-chat/src/components/chat-sidebar.tsx index 50861c112..73136f31c 100644 --- 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({ ))}
-
+
+
+ Warm timeout + 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" + /> + s +
);