Add warmTimeoutInSeconds option

Author: ericallam

packages/trigger-sdk/src/v3/ai.ts

@@ -406,6 +406,18 @@ export type ChatTaskOptions<TIdentifier extends string> = Omit<
    * @default "1h"
    */
   turnTimeout?: string;
+
+  /**
+   * How long (in seconds) to keep the run warm after each turn before suspending.
+   * During this window the run stays active and can respond instantly to the
+   * next message. After this timeout, the run suspends (frees compute) and waits
+   * via `inputStream.wait()`.
+   *
+   * Set to `0` to suspend immediately after each turn.
+   *
+   * @default 30
+   */
+  warmTimeoutInSeconds?: number;
 };
 
 /**
@@ -438,7 +450,13 @@ export type ChatTaskOptions<TIdentifier extends string> = Omit<
 function chatTask<TIdentifier extends string>(
   options: ChatTaskOptions<TIdentifier>
 ): Task<TIdentifier, ChatTaskPayload, unknown> {
-  const { run: userRun, maxTurns = 100, turnTimeout = "1h", ...restOptions } = options;
+  const {
+    run: userRun,
+    maxTurns = 100,
+    turnTimeout = "1h",
+    warmTimeoutInSeconds = 30,
+    ...restOptions
+  } = options;
 
   return createTask<TIdentifier, ChatTaskPayload, unknown>({
     ...restOptions,
@@ -512,15 +530,29 @@ function chatTask<TIdentifier extends string>(
             continue;
           }
 
-          // Suspend the task (frees compute) until the next message arrives
+          // Phase 1: Keep the run warm for quick response to the next message.
+          // The run stays active (using compute) during this window.
+          if (warmTimeoutInSeconds > 0) {
+            const warm = await messagesInput.once({
+              timeoutMs: warmTimeoutInSeconds * 1000,
+            });
+
+            if (warm.ok) {
+              // Message arrived while warm — respond instantly
+              currentPayload = warm.output;
+              continue;
+            }
+          }
+
+          // Phase 2: Suspend the task (frees compute) until the next message arrives
           const next = await messagesInput.wait({ timeout: turnTimeout });
 
           if (!next.ok) {
             // Timed out waiting for the next message — end the conversation
             return;
           }
 
-          currentPayload = next.output as ChatTaskPayload;
+          currentPayload = next.output;
         }
       } finally {
         stopSub.off();

packages/trigger-sdk/src/v3/streams.ts

@@ -750,23 +750,20 @@ function input<TData>(opts: { id: string }): RealtimeDefinedInputStream<TData> {
 
           const apiClient = apiClientManager.clientOrThrow();
 
+          // Create the waitpoint before the span so we have the entity ID upfront
+          const response = await apiClient.createInputStreamWaitpoint(ctx.run.id, {
+            streamId: opts.id,
+            timeout: options?.timeout,
+            idempotencyKey: options?.idempotencyKey,
+            idempotencyKeyTTL: options?.idempotencyKeyTTL,
+            tags: options?.tags,
+            lastSeqNum: inputStreams.lastSeqNum(opts.id),
+          });
+
           const result = await tracer.startActiveSpan(
             `inputStream.wait()`,
             async (span) => {
-              // 1. Create a waitpoint linked to this input stream
-              const response = await apiClient.createInputStreamWaitpoint(ctx.run.id, {
-                streamId: opts.id,
-                timeout: options?.timeout,
-                idempotencyKey: options?.idempotencyKey,
-                idempotencyKeyTTL: options?.idempotencyKeyTTL,
-                tags: options?.tags,
-                lastSeqNum: inputStreams.lastSeqNum(opts.id),
-              });
-
-              // Set the entity ID now that we have the waitpoint ID
-              span.setAttribute(SemanticInternalAttributes.ENTITY_ID, response.waitpointId);
-
-              // 2. Block the run on the waitpoint
+              // 1. Block the run on the waitpoint
               const waitResponse = await apiClient.waitForWaitpointToken({
                 runFriendlyId: ctx.run.id,
                 waitpointFriendlyId: response.waitpointId,
@@ -776,10 +773,10 @@ function input<TData>(opts: { id: string }): RealtimeDefinedInputStream<TData> {
                 throw new Error("Failed to block on input stream waitpoint");
               }
 
-              // 3. Suspend the task
+              // 2. Suspend the task
               const waitResult = await runtime.waitUntil(response.waitpointId);
 
-              // 4. Parse the output
+              // 3. Parse the output
               const data =
                 waitResult.output !== undefined
                   ? await conditionallyImportAndParsePacket(
@@ -806,6 +803,7 @@ function input<TData>(opts: { id: string }): RealtimeDefinedInputStream<TData> {
               attributes: {
                 [SemanticInternalAttributes.STYLE_ICON]: "wait",
                 [SemanticInternalAttributes.ENTITY_TYPE]: "waitpoint",
+                [SemanticInternalAttributes.ENTITY_ID]: response.waitpointId,
                 streamId: opts.id,
                 ...accessoryAttributes({
                   items: [

references/ai-chat/src/trigger/chat.ts

@@ -64,6 +64,7 @@ declare const Deno: unknown;
 
 export const aiChat = chat.task({
   id: "ai-chat",
+  warmTimeoutInSeconds: 10,
   run: async ({ messages, stopSignal }) => {
     return streamText({
       model: openai("gpt-4o-mini"),
@@ -72,6 +73,9 @@ export const aiChat = chat.task({
       tools: { inspectEnvironment },
       stopWhen: stepCountIs(10),
       abortSignal: stopSignal,
+      experimental_telemetry: {
+        isEnabled: true,
+      }
     });
   },
 });