diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index e39a5bb17..d633d59db 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -406,6 +406,18 @@ export type ChatTaskOptions = 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 = Omit< function chatTask( options: ChatTaskOptions ): Task { - const { run: userRun, maxTurns = 100, turnTimeout = "1h", ...restOptions } = options; + const { + run: userRun, + maxTurns = 100, + turnTimeout = "1h", + warmTimeoutInSeconds = 30, + ...restOptions + } = options; return createTask({ ...restOptions, @@ -512,7 +530,21 @@ function chatTask( 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) { @@ -520,7 +552,7 @@ function chatTask( return; } - currentPayload = next.output as ChatTaskPayload; + currentPayload = next.output; } } finally { stopSub.off(); diff --git a/packages/trigger-sdk/src/v3/streams.ts b/packages/trigger-sdk/src/v3/streams.ts index 68edc2a64..fe3af6e61 100644 --- a/packages/trigger-sdk/src/v3/streams.ts +++ b/packages/trigger-sdk/src/v3/streams.ts @@ -750,23 +750,20 @@ function input(opts: { id: string }): RealtimeDefinedInputStream { 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(opts: { id: string }): RealtimeDefinedInputStream { 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(opts: { id: string }): RealtimeDefinedInputStream { attributes: { [SemanticInternalAttributes.STYLE_ICON]: "wait", [SemanticInternalAttributes.ENTITY_TYPE]: "waitpoint", + [SemanticInternalAttributes.ENTITY_ID]: response.waitpointId, streamId: opts.id, ...accessoryAttributes({ items: [ diff --git a/references/ai-chat/src/trigger/chat.ts b/references/ai-chat/src/trigger/chat.ts index 66d7d734f..d4b3bab64 100644 --- a/references/ai-chat/src/trigger/chat.ts +++ b/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, + } }); }, });