feat(telemetry): use session key as trace ID and improve AI event inputs

Use the session key directly as the PostHog trace ID so feedback and generation events share the same trace. Extract non-assistant messages from conversations for $ai_input.

extensions/posthog-analytics/lib/trace-context.ts

@@ -27,7 +27,7 @@ export class TraceContextManager {
 
   startTrace(sessionKey: string, runId: string): void {
     this.traces.set(sessionKey, {
-      traceId: randomUUID(),
+      traceId: sessionKey,
       sessionId: sessionKey,
       runId,
       startedAt: Date.now(),

src/telemetry/event-mappers.test.ts

@@ -1,6 +1,6 @@
 import { describe, it, expect, beforeEach, vi } from "vitest";
 import { TraceContextManager } from "../../extensions/posthog-analytics/lib/trace-context.js";
-import { emitGeneration, emitToolSpan, emitTrace, emitCustomEvent, extractToolNamesFromMessages, extractUsageFromMessages, normalizeOutputForPostHog, buildTraceState } from "../../extensions/posthog-analytics/lib/event-mappers.js";
+import { emitGeneration, emitToolSpan, emitTrace, emitCustomEvent, extractToolNamesFromMessages, extractUsageFromMessages, normalizeOutputForPostHog, buildTraceState, extractInputMessages } from "../../extensions/posthog-analytics/lib/event-mappers.js";
 
 function createMockPostHog() {
   return {
@@ -131,6 +131,32 @@ describe("normalizeOutputForPostHog", () => {
   });
 });
 
+describe("extractInputMessages", () => {
+  it("extracts non-assistant messages from a conversation", () => {
+    const messages = [
+      { role: "system", content: "You are helpful" },
+      { role: "user", content: "What is 2+2?" },
+      { role: "assistant", content: "4" },
+      { role: "tool", name: "calc", content: "4" },
+    ];
+    const result = extractInputMessages(messages);
+    expect(result).toEqual([
+      { role: "system", content: "You are helpful" },
+      { role: "user", content: "What is 2+2?" },
+      { role: "tool", name: "calc", content: "4" },
+    ]);
+  });
+
+  it("returns undefined for non-array input", () => {
+    expect(extractInputMessages(null)).toBeUndefined();
+    expect(extractInputMessages(undefined)).toBeUndefined();
+  });
+
+  it("returns undefined when only assistant messages exist", () => {
+    expect(extractInputMessages([{ role: "assistant", content: "hi" }])).toBeUndefined();
+  });
+});
+
 describe("emitGeneration", () => {
   let ph: ReturnType<typeof createMockPostHog>;
   let traceCtx: TraceContextManager;
@@ -286,6 +312,45 @@ describe("emitGeneration", () => {
     emitGeneration(ph, traceCtx, "s", { durationMs: 5000 }, true);
     expect(ph.capture.mock.calls[0][0].properties.$ai_latency).toBe(5);
   });
+
+  it("includes user messages from event.messages in $ai_input even when trace.input is empty", () => {
+    traceCtx.startTrace("s", "r");
+
+    const messages = [
+      { role: "user", content: "what is this" },
+      { role: "assistant", content: "It's a config file." },
+    ];
+    emitGeneration(ph, traceCtx, "s", { messages }, false);
+
+    const input = ph.capture.mock.calls[0][0].properties.$ai_input;
+    expect(input).toEqual([{ role: "user", content: "what is this" }]);
+  });
+
+  it("prefers event.messages over trace.input for $ai_input", () => {
+    traceCtx.startTrace("s", "r");
+    traceCtx.setInput("s", [{ role: "system", content: "tool config" }], false);
+
+    const messages = [
+      { role: "user", content: "hello" },
+      { role: "assistant", content: "hi" },
+    ];
+    emitGeneration(ph, traceCtx, "s", { messages }, false);
+
+    const input = ph.capture.mock.calls[0][0].properties.$ai_input;
+    expect(input).toEqual([{ role: "user", content: "hello" }]);
+  });
+
+  it("falls back to trace.input when event.messages has no input messages", () => {
+    traceCtx.startTrace("s", "r");
+    traceCtx.setInput("s", [{ role: "user", content: "from trace" }], false);
+
+    emitGeneration(ph, traceCtx, "s", {
+      output: [{ role: "assistant", content: "reply" }],
+    }, false);
+
+    const input = ph.capture.mock.calls[0][0].properties.$ai_input;
+    expect(input).toEqual([{ role: "user", content: "from trace" }]);
+  });
 });
 
 describe("emitToolSpan", () => {

src/telemetry/trace-context.test.ts

@@ -28,12 +28,13 @@ describe("TraceContextManager", () => {
 
   // ── Trace lifecycle (session-keyed) ──
 
-  it("generates a unique UUID traceId for each trace (ensures PostHog trace grouping)", () => {
+  it("uses sessionKey as traceId so feedback and generation events share the same trace", () => {
     ctx.startTrace("session-1", "run-1");
     ctx.startTrace("session-2", "run-2");
     const t1 = ctx.getTrace("session-1")!;
     const t2 = ctx.getTrace("session-2")!;
-    expect(t1.traceId).toMatch(/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/);
+    expect(t1.traceId).toBe("session-1");
+    expect(t2.traceId).toBe("session-2");
     expect(t1.traceId).not.toBe(t2.traceId);
   });