From 83f6b69f82450035d7af05b84074bca3b4e28c66 Mon Sep 17 00:00:00 2001
From: kumarabhirup <hey@kumareth.com>
Date: Thu, 5 Mar 2026 19:09:13 -0800
Subject: [PATCH] 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.
---
 .../posthog-analytics/lib/trace-context.ts    |  2 +-
 src/telemetry/event-mappers.test.ts           | 67 ++++++++++++++++++-
 src/telemetry/trace-context.test.ts           |  5 +-
 3 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/extensions/posthog-analytics/lib/trace-context.ts b/extensions/posthog-analytics/lib/trace-context.ts
index f8a3b7e168c..91e948dd98a 100644
--- a/extensions/posthog-analytics/lib/trace-context.ts
+++ b/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(),
diff --git a/src/telemetry/event-mappers.test.ts b/src/telemetry/event-mappers.test.ts
index b0dc8e63123..ad16f4f83a6 100644
--- a/src/telemetry/event-mappers.test.ts
+++ b/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", () => {
diff --git a/src/telemetry/trace-context.test.ts b/src/telemetry/trace-context.test.ts
index 5365d9de240..7c1d6642cfb 100644
--- a/src/telemetry/trace-context.test.ts
+++ b/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);
   });