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

This commit is contained in:
Eric Allam
2026-03-09 17:23:37 +00:00
parent 7aa5964141
commit 2908b3598d
6 changed files with 103 additions and 51 deletions
@@ -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;
+19 -47
View File
@@ -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";
}
+4 -1
View File
@@ -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();
+42 -1
View File
@@ -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()`,
@@ -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 ? (
@@ -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>
);