fix(chat): Improve agent loop tracing

TELEMETRY.md

@@ -134,7 +134,7 @@ Events: `agent_message_in`, `agent_message_out`, `agent_turn_timeout`,
 `assistant_reply_generation_failed`
 
 Spans: `ai.generate_assistant_reply`, `ai.chat_completion`,
-`chat.route_thinking`, `ai.invoke_advisor`
+`chat.route_thinking`, `ai.invoke_advisor`, `gen_ai.chat`
 
 Attributes: `gen_ai.operation.name`, `gen_ai.request.model`,
 `gen_ai.response.finish_reasons`, `app.ai.outcome`,
@@ -153,7 +153,7 @@ Events: `agent_tool_call_failed`, `mcp_tool_call_failed`,
 Spans: `execute_tool <toolName>`, `sandbox.acquire`, `sandbox.create`,
 `sandbox.snapshot.resolve`, `sandbox.sync_skills`, `bash`
 
-Attributes: `gen_ai.tool.name`, `gen_ai.tool.call.id`, `mcp.method.name`,
+Attributes: `gen_ai.tool.name`, `gen_ai.tool.call.id`, `gen_ai.tool.call.result`, `mcp.method.name`,
 `process.executable.name`, `process.exit.code`, `app.sandbox.source`,
 `app.sandbox.snapshot.resolve_outcome`
 

packages/junior/src/chat/pi/traced-stream.ts

@@ -0,0 +1,128 @@
+import type { StreamFn } from "@mariozechner/pi-agent-core";
+import {
+  type Api,
+  type AssistantMessage,
+  type Context,
+  type Model,
+  streamSimple,
+} from "@mariozechner/pi-ai";
+import * as Sentry from "@/chat/sentry";
+import {
+  extractGenAiUsageAttributes,
+  getLogContextAttributes,
+  serializeGenAiAttribute,
+} from "@/chat/logging";
+import { GEN_AI_PROVIDER_NAME } from "@/chat/pi/client";
+
+// Compose only the OTel GenAI attributes that are knowable at span start
+// (request-shape + system instructions). End-of-call attributes such as
+// usage and finish reasons are set after the stream resolves.
+function buildChatStartAttributes(
+  model: Model<Api>,
+  context: Context,
+): Record<string, string> {
+  const attributes: Record<string, string> = {
+    "gen_ai.operation.name": "chat",
+    "gen_ai.provider.name": GEN_AI_PROVIDER_NAME,
+    "gen_ai.request.model": model.id,
+  };
+
+  const inputMessages = serializeGenAiAttribute(context.messages);
+  if (inputMessages) {
+    attributes["gen_ai.input.messages"] = inputMessages;
+  }
+
+  if (context.systemPrompt) {
+    const systemInstructions = serializeGenAiAttribute([
+      { type: "text", content: context.systemPrompt },
+    ]);
+    if (systemInstructions) {
+      attributes["gen_ai.system_instructions"] = systemInstructions;
+    }
+  }
+
+  return attributes;
+}
+
+// Composes post-stream attributes for the chat span.
+// Known gap: `gen_ai.response.finish_reasons` emits pi-ai's raw StopReason
+// values (e.g. "toolUse", "aborted") instead of the OTel canonical set
+// ("tool_use", "max_tokens"). Tracked separately, out of scope here.
+function buildChatEndAttributes(
+  message: AssistantMessage,
+): Record<string, string | string[] | number> {
+  const attributes: Record<string, string | string[] | number> = {};
+
+  const outputMessages = serializeGenAiAttribute([message]);
+  if (outputMessages) {
+    attributes["gen_ai.output.messages"] = outputMessages;
+  }
+
+  Object.assign(attributes, extractGenAiUsageAttributes(message));
+
+  if (message.stopReason) {
+    attributes["gen_ai.response.finish_reasons"] = [message.stopReason];
+  }
+
+  if (message.model) {
+    attributes["gen_ai.response.model"] = message.model;
+  }
+
+  return attributes;
+}
+
+/**
+ * Wraps pi-ai's `streamSimple` so each LLM call inside a pi-agent-core agent
+ * loop produces its own `gen_ai.chat` Sentry span. The returned function is
+ * passed to `new Agent({ streamFn: ... })` and runs once per loop iteration.
+ *
+ * The base argument exists so tests can inject a stub stream function.
+ */
+export function createTracedStreamFn(base: StreamFn = streamSimple): StreamFn {
+  return async (model, context, options) => {
+    const span = Sentry.startInactiveSpan({
+      name: `chat ${model.id}`,
+      op: "gen_ai.chat",
+      attributes: {
+        ...getLogContextAttributes(),
+        ...buildChatStartAttributes(model, context),
+      },
+    });
+
+    try {
+      const stream = await Sentry.withActiveSpan(span, () =>
+        Promise.resolve(base(model, context, options)),
+      );
+
+      stream
+        .result()
+        .then(
+          (finalMessage) => {
+            try {
+              for (const [key, value] of Object.entries(
+                buildChatEndAttributes(finalMessage),
+              )) {
+                span.setAttribute(key, value);
+              }
+            } finally {
+              span.end();
+            }
+          },
+          () => {
+            span.setStatus({ code: 2, message: "LLM stream failed" });
+            span.end();
+          },
+        )
+        .catch(() => {
+          // setAttribute is best-effort; suppress unexpected attribute-write
+          // errors so they don't surface as unhandled promise rejections.
+        });
+
+      return stream;
+    } catch (error) {
+      span.setStatus({ code: 2, message: "LLM call failed" });
+      span.end();
+      throw error;
+    }
+  };
+}

packages/junior/src/chat/respond.ts

@@ -51,6 +51,7 @@ import {
   resolveGatewayModel,
 } from "@/chat/pi/client";
 import type { PiMessage } from "@/chat/pi/messages";
+import { createTracedStreamFn } from "@/chat/pi/traced-stream";
 import {
   createSandboxExecutor,
   type SandboxAcquiredState,
@@ -819,6 +820,7 @@ export async function generateAssistantReply(
           conversationId: sessionConversationId,
           logContext: spanContext,
           getTools: () => advisorTools,
+          streamFn: createTracedStreamFn(),
         },
       },
     );
@@ -927,6 +929,7 @@ export async function generateAssistantReply(
     // ── Agent execution ──────────────────────────────────────────────
     agent = new Agent({
       getApiKey: () => getPiGatewayApiKeyOverride(),
+      streamFn: createTracedStreamFn(),
       initialState: {
         systemPrompt: baseInstructions,
         model: resolveGatewayModel(botConfig.modelId),

packages/junior/src/chat/tools/slack/channel-list-messages.ts

@@ -94,13 +94,23 @@ export function createSlackChannelListMessagesTool(
         throw error;
       }
 
-      return {
+      const summary = {
         ok: true,
         channel_id: targetChannelId,
         count: result.messages.length,
         next_cursor: result.nextCursor,
         messages: result.messages,
       };
+
+      return {
+        content: [{ type: "text" as const, text: JSON.stringify(summary) }],
+        details: {
+          ok: true,
+          channel_id: targetChannelId,
+          count: result.messages.length,
+          ...(result.nextCursor ? { next_cursor: result.nextCursor } : {}),
+        },
+      };
     },
   });
 }

packages/junior/tests/integration/slack-channel-tools.test.ts

@@ -143,15 +143,16 @@ describe("slack channel tools", () => {
       max_pages: 3,
     });
 
-    expect(result).toMatchObject({
+    expect(result.details).toMatchObject({
       ok: true,
       channel_id: "C123",
       count: 1,
-      next_cursor: undefined,
-    });
-    expect(result).toMatchObject({
-      messages: [{ ts: "1700000000.300", text: "hello", user: "U1" }],
     });
+    expect(result.details).not.toHaveProperty("next_cursor");
+    const body = JSON.parse(result.content[0].text);
+    expect(body.messages).toMatchObject([
+      { ts: "1700000000.300", text: "hello", user: "U1" },
+    ]);
 
     const historyCalls = getCapturedSlackApiCalls("conversations.history");
     expect(historyCalls).toHaveLength(1);
@@ -213,18 +214,17 @@ describe("slack channel tools", () => {
       max_pages: 3,
     });
 
-    expect(result).toMatchObject({
+    expect(result.details).toMatchObject({
       ok: true,
       channel_id: "C123",
       count: 2,
-      next_cursor: undefined,
-    });
-    expect(result).toMatchObject({
-      messages: [
-        { ts: "1700000000.500", text: "page-1", user: "U1" },
-        { ts: "1700000000.501", text: "page-2", user: "U2" },
-      ],
     });
+    expect(result.details).not.toHaveProperty("next_cursor");
+    const body = JSON.parse(result.content[0].text);
+    expect(body.messages).toMatchObject([
+      { ts: "1700000000.500", text: "page-1", user: "U1" },
+      { ts: "1700000000.501", text: "page-2", user: "U2" },
+    ]);
 
     const historyCalls = getCapturedSlackApiCalls("conversations.history");
     expect(historyCalls).toHaveLength(2);