Added better telemetry support to view turns

ericallam · ericallam · commit c804a98b6d69 · 2026-03-04T16:08:20.000Z
diff --git a/packages/core/src/v3/realtimeStreams/types.ts b/packages/core/src/v3/realtimeStreams/types.ts
@@ -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>
diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts
@@ -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;
+          // Extract turn-level context before entering the span
+          const { metadata: wireMetadata, messages: uiMessages, ...restWire } = currentWirePayload;
+          const lastUserMessage = extractLastUserMessageText(uiMessages);
+
+          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",
+          };
+
+          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",
+              })
+            );
+          }
 
-          // Per-turn stop controller (reset each turn)
-          const stopController = new AbortController();
-          currentStopController = stopController;
+          if (wireMetadata !== undefined) {
+            turnAttributes["chat.client_data"] =
+              typeof wireMetadata === "string" ? wireMetadata : JSON.stringify(wireMetadata);
+          }
 
-          // Three signals for the user's run function
-          const stopSignal = stopController.signal;
-          const cancelSignal = runSignal;
-          const combinedSignal = AbortSignal.any([runSignal, stopController.signal]);
+          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();
+              }
 
-          // Buffer messages that arrive during streaming
-          const pendingMessages: ChatTaskWirePayload[] = [];
-          const msgSub = messagesInput.on((msg) => {
-            pendingMessages.push(msg);
-          });
+              if (runSignal.aborted) return "exit";
 
-          // Convert wire payload to user-facing payload
-          const { metadata: wireMetadata, messages: uiMessages, ...restWire } = currentWirePayload;
-          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 });
-            }
-          } catch (error) {
-            // Handle AbortError from streamText gracefully
-            if (error instanceof Error && error.name === "AbortError") {
-              if (runSignal.aborted) {
-                return; // Full run cancellation — 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";
               }
-              // Stop generation — fall through to continue the loop
-            } else {
-              throw error;
-            }
-          } finally {
-            msgSub.off();
-          }
 
-          if (runSignal.aborted) return;
+              // 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";
+                }
+              }
 
-          // Write turn-complete control chunk so frontend closes its stream
-          await writeTurnCompleteChunk();
+              // Phase 2: Suspend the task (frees compute) until the next message arrives
+              const next = await messagesInput.wait({
+                timeout: turnTimeout,
+                spanName: "waiting (suspended)",
+              });
 
-          // If messages arrived during streaming, use the first one immediately
-          if (pendingMessages.length > 0) {
-            currentWirePayload = pendingMessages[0]!;
-            continue;
-          }
+              if (!next.ok) {
+                // Timed out waiting for the next message — end the conversation
+                return "exit";
+              }
 
-          // 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;
+              currentWirePayload = next.output;
+              return "continue";
+            },
+            {
+              attributes: turnAttributes,
             }
-          }
-
-          // 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;
+}
diff --git a/packages/trigger-sdk/src/v3/streams.ts b/packages/trigger-sdk/src/v3/streams.ts
@@ -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({
