From 600069ffa17cbc1701df4c2b78360b1f610c6fc8 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Fri, 6 Mar 2026 14:52:10 +0000 Subject: [PATCH] use locals for the chat pipe counter instead of a module global --- packages/trigger-sdk/src/v3/ai.ts | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 16c6667da..3140402d4 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -17,6 +17,7 @@ import type { StreamWriteResult } from "@trigger.dev/core/v3"; import { convertToModelMessages, dynamicTool, generateId as generateMessageId, jsonSchema, JSONSchema7, Schema, Tool, ToolCallOptions, zodSchema } from "ai"; import { type Attributes, trace } from "@opentelemetry/api"; import { auth } from "./auth.js"; +import { locals } from "./locals.js"; import { metadata } from "./metadata.js"; import { streams } from "./streams.js"; import { createTask } from "./shared.js"; @@ -239,12 +240,11 @@ const messagesInput = streams.input({ id: CHAT_MESSAGES_STR const stopInput = streams.input<{ stop: true; message?: string }>({ id: CHAT_STOP_STREAM_ID }); /** - * Tracks how many times `pipeChat` has been called in the current `chatTask` run. - * Used to prevent double-piping when a user both calls `pipeChat()` manually - * and returns a streamable from their `run` function. + * Run-scoped pipe counter. Stored in locals so concurrent runs in the + * same worker don't share state. * @internal */ -let _chatPipeCount = 0; +const chatPipeCountKey = locals.create("chat.pipeCount"); /** * Options for `pipeChat`. @@ -336,7 +336,7 @@ async function pipeChat( source: UIMessageStreamable | AsyncIterable | ReadableStream, options?: PipeChatOptions ): Promise { - _chatPipeCount++; + locals.set(chatPipeCountKey, (locals.get(chatPipeCountKey) ?? 0) + 1); const streamKey = options?.streamKey ?? CHAT_STREAM_KEY; let stream: AsyncIterable | ReadableStream; @@ -662,7 +662,7 @@ function chatTask( const turnResult = await tracer.startActiveSpan( `chat turn ${turn + 1}`, async () => { - _chatPipeCount = 0; + locals.set(chatPipeCountKey, 0); // Per-turn stop controller (reset each turn) const stopController = new AbortController(); @@ -792,7 +792,7 @@ function chatTask( // Auto-pipe if the run function returned a StreamTextResult or similar, // but only if pipeChat() wasn't already called manually during this turn. // We call toUIMessageStream ourselves to attach onFinish for response capture. - if (_chatPipeCount === 0 && isUIMessageStreamable(result)) { + if ((locals.get(chatPipeCountKey) ?? 0) === 0 && isUIMessageStreamable(result)) { const uiStream = result.toUIMessageStream({ onFinish: ({ responseMessage }: { responseMessage: UIMessage }) => { capturedResponseMessage = responseMessage;