feat(streams): add inputStream.waitWithWarmup(), warm timeout config in sidebar, preload payload option
This commit is contained in:
@@ -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;
|
||||
|
||||
@@ -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";
|
||||
}
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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>
|
||||
);
|
||||
|
||||
Reference in New Issue
Block a user