Added better telemetry support to view turns

This commit is contained in:
Eric Allam
2026-03-04 16:08:20 +00:00
parent 09b7f4d085
commit 6c3c1937cc
3 changed files with 189 additions and 85 deletions
@@ -71,6 +71,10 @@ export type PipeStreamOptions = {
* Additional request options for the API call.
*/
requestOptions?: ApiRequestOptions;
/** Override the default span name for this operation. */
spanName?: string;
/** When true, the span will be collapsed in the dashboard. */
collapsed?: boolean;
};
/**
@@ -199,6 +203,8 @@ export type InputStreamSubscription = {
export type InputStreamOnceOptions = {
signal?: AbortSignal;
timeoutMs?: number;
/** Override the default span name for this operation. */
spanName?: string;
};
export type SendInputStreamOptions = {
@@ -234,6 +240,9 @@ export type InputStreamWaitOptions = {
* and filtering waitpoints via `wait.listTokens()`.
*/
tags?: string[];
/** Override the default span name for this operation. */
spanName?: string;
};
export type InferInputStreamType<T> = T extends RealtimeDefinedInputStream<infer TData>
+175 -81
View File
@@ -1,6 +1,8 @@
import {
accessoryAttributes,
AnyTask,
isSchemaZodEsque,
SemanticInternalAttributes,
Task,
type inferSchemaIn,
type PipeStreamOptions,
@@ -11,10 +13,12 @@ import {
} from "@trigger.dev/core/v3";
import type { ModelMessage, UIMessage } from "ai";
import { convertToModelMessages, dynamicTool, jsonSchema, JSONSchema7, Schema, Tool, ToolCallOptions, zodSchema } from "ai";
import { type Attributes, trace } from "@opentelemetry/api";
import { auth } from "./auth.js";
import { metadata } from "./metadata.js";
import { streams } from "./streams.js";
import { createTask } from "./shared.js";
import { tracer } from "./tracer.js";
import {
CHAT_STREAM_KEY as _CHAT_STREAM_KEY,
CHAT_MESSAGES_STREAM_ID,
@@ -281,6 +285,9 @@ export type PipeChatOptions = {
* @default "self" (current run)
*/
target?: string;
/** Override the default span name for this operation. */
spanName?: string;
};
/**
@@ -373,6 +380,9 @@ async function pipeChat(
if (options?.target) {
pipeOptions.target = options.target;
}
if (options?.spanName) {
pipeOptions.spanName = options.spanName;
}
const { waitUntilComplete } = streams.pipe(streamKey, stream, pipeOptions);
await waitUntilComplete();
@@ -478,6 +488,12 @@ function chatTask<TIdentifier extends string>(
return createTask<TIdentifier, ChatTaskWirePayload, unknown>({
...restOptions,
run: async (payload: ChatTaskWirePayload, { signal: runSignal }) => {
// Set gen_ai.conversation.id on the run-level span for dashboard context
const activeSpan = trace.getActiveSpan();
if (activeSpan) {
activeSpan.setAttribute("gen_ai.conversation.id", payload.chatId);
}
let currentWirePayload = payload;
// Mutable reference to the current turn's stop controller so the
@@ -491,92 +507,142 @@ function chatTask<TIdentifier extends string>(
try {
for (let turn = 0; turn < maxTurns; turn++) {
_chatPipeCount = 0;
// Per-turn stop controller (reset each turn)
const stopController = new AbortController();
currentStopController = stopController;
// Three signals for the user's run function
const stopSignal = stopController.signal;
const cancelSignal = runSignal;
const combinedSignal = AbortSignal.any([runSignal, stopController.signal]);
// Buffer messages that arrive during streaming
const pendingMessages: ChatTaskWirePayload[] = [];
const msgSub = messagesInput.on((msg) => {
pendingMessages.push(msg);
});
// Convert wire payload to user-facing payload
// Extract turn-level context before entering the span
const { metadata: wireMetadata, messages: uiMessages, ...restWire } = currentWirePayload;
const sanitized = sanitizeMessages(uiMessages);
const modelMessages = await convertToModelMessages(sanitized);
const lastUserMessage = extractLastUserMessageText(uiMessages);
try {
const result = await userRun({
...restWire,
messages: modelMessages,
uiMessages: sanitized,
clientData: wireMetadata,
signal: combinedSignal,
cancelSignal,
stopSignal,
});
const turnAttributes: Attributes = {
"turn.number": turn + 1,
"gen_ai.conversation.id": currentWirePayload.chatId,
"gen_ai.operation.name": "chat",
"chat.trigger": currentWirePayload.trigger,
[SemanticInternalAttributes.STYLE_ICON]: "tabler-message-chatbot",
[SemanticInternalAttributes.ENTITY_TYPE]: "chat-turn",
};
// Auto-pipe if the run function returned a StreamTextResult or similar,
// but only if pipeChat() wasn't already called manually during this turn
if (_chatPipeCount === 0 && isUIMessageStreamable(result)) {
await pipeChat(result, { signal: combinedSignal });
}
} catch (error) {
// Handle AbortError from streamText gracefully
if (error instanceof Error && error.name === "AbortError") {
if (runSignal.aborted) {
return; // Full run cancellation — exit
if (lastUserMessage) {
turnAttributes["chat.user_message"] = lastUserMessage;
// Show a truncated preview of the user message as an accessory
const preview =
lastUserMessage.length > 80
? lastUserMessage.slice(0, 80) + "..."
: lastUserMessage;
Object.assign(
turnAttributes,
accessoryAttributes({
items: [{ text: preview, variant: "normal" }],
style: "codepath",
})
);
}
if (wireMetadata !== undefined) {
turnAttributes["chat.client_data"] =
typeof wireMetadata === "string" ? wireMetadata : JSON.stringify(wireMetadata);
}
const turnResult = await tracer.startActiveSpan(
`chat turn ${turn + 1}`,
async () => {
_chatPipeCount = 0;
// Per-turn stop controller (reset each turn)
const stopController = new AbortController();
currentStopController = stopController;
// Three signals for the user's run function
const stopSignal = stopController.signal;
const cancelSignal = runSignal;
const combinedSignal = AbortSignal.any([runSignal, stopController.signal]);
// Buffer messages that arrive during streaming
const pendingMessages: ChatTaskWirePayload[] = [];
const msgSub = messagesInput.on((msg) => {
pendingMessages.push(msg);
});
// Convert wire payload to user-facing payload
const sanitized = sanitizeMessages(uiMessages);
const modelMessages = await convertToModelMessages(sanitized);
try {
const result = await userRun({
...restWire,
messages: modelMessages,
uiMessages: sanitized,
clientData: wireMetadata,
signal: combinedSignal,
cancelSignal,
stopSignal,
});
// Auto-pipe if the run function returned a StreamTextResult or similar,
// but only if pipeChat() wasn't already called manually during this turn
if (_chatPipeCount === 0 && isUIMessageStreamable(result)) {
await pipeChat(result, { signal: combinedSignal, spanName: "stream response" });
}
} catch (error) {
// Handle AbortError from streamText gracefully
if (error instanceof Error && error.name === "AbortError") {
if (runSignal.aborted) {
return "exit"; // Full run cancellation — exit
}
// Stop generation — fall through to continue the loop
} else {
throw error;
}
} finally {
msgSub.off();
}
// Stop generation — fall through to continue the loop
} else {
throw error;
if (runSignal.aborted) return "exit";
// Write turn-complete control chunk so frontend closes its stream
await writeTurnCompleteChunk(currentWirePayload.chatId);
// If messages arrived during streaming, use the first one immediately
if (pendingMessages.length > 0) {
currentWirePayload = pendingMessages[0]!;
return "continue";
}
// 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,
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 next = await messagesInput.wait({
timeout: turnTimeout,
spanName: "waiting (suspended)",
});
if (!next.ok) {
// Timed out waiting for the next message — end the conversation
return "exit";
}
currentWirePayload = next.output;
return "continue";
},
{
attributes: turnAttributes,
}
} finally {
msgSub.off();
}
);
if (runSignal.aborted) return;
// Write turn-complete control chunk so frontend closes its stream
await writeTurnCompleteChunk();
// If messages arrived during streaming, use the first one immediately
if (pendingMessages.length > 0) {
currentWirePayload = pendingMessages[0]!;
continue;
}
// 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
currentWirePayload = 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;
}
currentWirePayload = next.output;
if (turnResult === "exit") return;
// "continue" means proceed to next iteration
}
} finally {
stopSub.off();
@@ -621,11 +687,39 @@ export const chat = {
* The frontend transport intercepts this to close the ReadableStream for the current turn.
* @internal
*/
async function writeTurnCompleteChunk(): Promise<void> {
async function writeTurnCompleteChunk(chatId?: string): Promise<void> {
const { waitUntilComplete } = streams.writer(CHAT_STREAM_KEY, {
spanName: "turn complete",
collapsed: true,
execute: ({ write }) => {
write({ type: "__trigger_turn_complete" });
},
});
await waitUntilComplete();
}
/**
* Extracts the text content of the last user message from a UIMessage array.
* Returns undefined if no user message is found.
* @internal
*/
function extractLastUserMessageText(messages: UIMessage[]): string | undefined {
for (let i = messages.length - 1; i >= 0; i--) {
const msg = messages[i]!;
if (msg.role !== "user") continue;
// UIMessage uses parts array
if (msg.parts) {
const textParts = msg.parts
.filter((p: any) => p.type === "text" && p.text)
.map((p: any) => p.text as string);
if (textParts.length > 0) {
return textParts.join("\n");
}
}
break;
}
return undefined;
}
+5 -4
View File
@@ -139,7 +139,7 @@ function pipe<T>(
opts = valueOrOptions as PipeStreamOptions | undefined;
}
return pipeInternal(key, value, opts, "streams.pipe()");
return pipeInternal(key, value, opts, opts?.spanName ?? "streams.pipe()");
}
/**
@@ -167,6 +167,7 @@ function pipeInternal<T>(
[SemanticInternalAttributes.ENTITY_TYPE]: "realtime-stream",
[SemanticInternalAttributes.ENTITY_ID]: `${runId}:${key}`,
[SemanticInternalAttributes.STYLE_ICON]: "streams",
...(opts?.collapsed ? { [SemanticInternalAttributes.COLLAPSED]: true } : {}),
...accessoryAttributes({
items: [
{
@@ -640,7 +641,7 @@ function writerInternal<TPart>(key: string, options: WriterStreamOptions<TPart>)
}
});
return pipeInternal(key, stream, options, "streams.writer()");
return pipeInternal(key, stream, options, options.spanName ?? "streams.writer()");
}
export type RealtimeDefineStreamOptions = {
@@ -713,7 +714,7 @@ function input<TData>(opts: { id: string }): RealtimeDefinedInputStream<TData> {
return new InputStreamOncePromise<TData>((resolve, reject) => {
tracer
.startActiveSpan(
`inputStream.once()`,
options?.spanName ?? `inputStream.once()`,
async () => {
const result = await innerPromise;
resolve(result as InputStreamOnceResult<TData>);
@@ -761,7 +762,7 @@ function input<TData>(opts: { id: string }): RealtimeDefinedInputStream<TData> {
});
const result = await tracer.startActiveSpan(
`inputStream.wait()`,
options?.spanName ?? `inputStream.wait()`,
async (span) => {
// 1. Block the run on the waitpoint
const waitResponse = await apiClient.waitForWaitpointToken({