use locals for the chat pipe counter instead of a module global

This commit is contained in:
Eric Allam
2026-03-06 14:52:10 +00:00
parent 3ae26683a6
commit 600069ffa1
+7 -7
View File
@@ -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<ChatTaskWirePayload>({ 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<number>("chat.pipeCount");
/**
* Options for `pipeChat`.
@@ -336,7 +336,7 @@ async function pipeChat(
source: UIMessageStreamable | AsyncIterable<unknown> | ReadableStream<unknown>,
options?: PipeChatOptions
): Promise<void> {
_chatPipeCount++;
locals.set(chatPipeCountKey, (locals.get(chatPipeCountKey) ?? 0) + 1);
const streamKey = options?.streamKey ?? CHAT_STREAM_KEY;
let stream: AsyncIterable<unknown> | ReadableStream<unknown>;
@@ -662,7 +662,7 @@ function chatTask<TIdentifier extends string>(
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<TIdentifier extends string>(
// 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;