feat(aio): declare session id and per-turn traceparent to the llm gateway (#3927)

Author: carlos-marchal-ph

packages/agent/src/adapters/claude/session/options.test.ts

@@ -376,14 +376,18 @@ describe("buildSessionOptions", () => {
         name: "omits the team_id header when POSTHOG_PROJECT_ID is unset",
         projectId: undefined,
         existingHeaders: undefined,
-        expected: "x-posthog-use-bedrock-fallback: true",
+        expected: [
+          "x-posthog-property-$ai_session_id: test-session",
+          "x-posthog-use-bedrock-fallback: true",
+        ].join("\n"),
       },
       {
         name: "forwards POSTHOG_PROJECT_ID as the team_id attribution header",
         projectId: "42",
         existingHeaders: undefined,
         expected: [
           "x-posthog-property-team_id: 42",
+          "x-posthog-property-$ai_session_id: test-session",
           "x-posthog-use-bedrock-fallback: true",
         ].join("\n"),
       },
@@ -394,6 +398,7 @@ describe("buildSessionOptions", () => {
         expected: [
           "x-posthog-property-task_id: task-abc",
           "x-posthog-property-team_id: 42",
+          "x-posthog-property-$ai_session_id: test-session",
           "x-posthog-use-bedrock-fallback: true",
         ].join("\n"),
       },
@@ -411,6 +416,101 @@ describe("buildSessionOptions", () => {
       expect(headers).toBe(expected);
     });
   });
+
+  describe("gateway turn tracing env", () => {
+    const KEYS = [
+      "CLAUDE_CODE_ENABLE_TELEMETRY",
+      "CLAUDE_CODE_ENHANCED_TELEMETRY_BETA",
+      "CLAUDE_CODE_PROPAGATE_TRACEPARENT",
+      "OTEL_TRACES_EXPORTER",
+      "OTEL_EXPORTER_OTLP_PROTOCOL",
+      "OTEL_EXPORTER_OTLP_ENDPOINT",
+      "TRACEPARENT",
+      "TRACESTATE",
+    ] as const;
+    const original: Partial<Record<string, string | undefined>> = {};
+
+    beforeEach(() => {
+      for (const key of KEYS) {
+        original[key] = process.env[key];
+        delete process.env[key];
+      }
+    });
+
+    afterEach(() => {
+      for (const key of KEYS) {
+        const value = original[key];
+        if (value === undefined) {
+          delete process.env[key];
+        } else {
+          process.env[key] = value;
+        }
+      }
+    });
+
+    const gatewayEnv = {
+      anthropicBaseUrl: "https://gateway.example",
+      anthropicAuthToken: "tok",
+      openaiBaseUrl: "https://gateway.example/v1",
+      openaiApiKey: "tok",
+    };
+
+    it("enables per-turn traceparent when routed through the gateway", () => {
+      const env = buildSessionOptions({ ...makeParams(), gatewayEnv }).env;
+
+      expect(env?.CLAUDE_CODE_ENABLE_TELEMETRY).toBe("1");
+      expect(env?.CLAUDE_CODE_ENHANCED_TELEMETRY_BETA).toBe("1");
+      expect(env?.CLAUDE_CODE_PROPAGATE_TRACEPARENT).toBe("1");
+      expect(env?.OTEL_TRACES_EXPORTER).toBe("otlp");
+      expect(env?.OTEL_EXPORTER_OTLP_PROTOCOL).toBe("http/json");
+      expect(env?.OTEL_EXPORTER_OTLP_ENDPOINT).toBe("http://127.0.0.1:9");
+    });
+
+    it("honors a caller-supplied OTLP endpoint", () => {
+      process.env.OTEL_EXPORTER_OTLP_ENDPOINT =
+        "http://collector.internal:4318";
+
+      const env = buildSessionOptions({ ...makeParams(), gatewayEnv }).env;
+
+      expect(env?.OTEL_EXPORTER_OTLP_ENDPOINT).toBe(
+        "http://collector.internal:4318",
+      );
+    });
+
+    it("pins exporter and protocol so an inherited none can't disable tracing", () => {
+      process.env.OTEL_TRACES_EXPORTER = "none";
+      process.env.OTEL_EXPORTER_OTLP_PROTOCOL = "grpc";
+
+      const env = buildSessionOptions({ ...makeParams(), gatewayEnv }).env;
+
+      expect(env?.OTEL_TRACES_EXPORTER).toBe("otlp");
+      expect(env?.OTEL_EXPORTER_OTLP_PROTOCOL).toBe("http/json");
+    });
+
+    it("strips inherited TRACEPARENT so turns keep distinct trace ids", () => {
+      process.env.TRACEPARENT =
+        "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";
+      process.env.TRACESTATE = "vendor=x";
+
+      const env = buildSessionOptions({ ...makeParams(), gatewayEnv }).env;
+
+      expect(env?.TRACEPARENT).toBeUndefined();
+      expect(env?.TRACESTATE).toBeUndefined();
+    });
+
+    it("leaves BYOK sessions untouched", () => {
+      process.env.TRACEPARENT =
+        "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";
+
+      const env = buildSessionOptions(makeParams()).env;
+
+      expect(env?.CLAUDE_CODE_ENABLE_TELEMETRY).toBeUndefined();
+      expect(env?.CLAUDE_CODE_PROPAGATE_TRACEPARENT).toBeUndefined();
+      expect(env?.TRACEPARENT).toBe(
+        "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+      );
+    });
+  });
 });
 
 describe("buildSystemPrompt", () => {

packages/agent/src/adapters/claude/session/options.ts

@@ -151,7 +151,10 @@ function buildMcpServers(
   };
 }
 
-function buildEnvironment(gateway?: GatewayEnv): Record<string, string> {
+function buildEnvironment(
+  gateway?: GatewayEnv,
+  sessionId?: string,
+): Record<string, string> {
   // Custom HTTP headers reach the model only through the Claude CLI subprocess,
   // which reads them from this env var (newline-delimited `name: value` lines)
   // — the SDK has no direct header option. We finalize them here, the single
@@ -174,6 +177,11 @@ function buildEnvironment(gateway?: GatewayEnv): Record<string, string> {
   if (projectId) {
     headerLines.push(buildGatewayPropertyHeaders({ team_id: projectId }));
   }
+  if (sessionId) {
+    headerLines.push(
+      buildGatewayPropertyHeaders({ $ai_session_id: sessionId }),
+    );
+  }
   // Route to AWS Bedrock as a fallback when Anthropic returns 5xx
   headerLines.push("x-posthog-use-bedrock-fallback: true");
   const customHeaders = headerLines.join("\n");
@@ -185,8 +193,31 @@ function buildEnvironment(gateway?: GatewayEnv): Record<string, string> {
   // sessions that genuinely need MCP tools available on turn 1.
   const mcpNonblocking = process.env.MCP_CONNECTION_NONBLOCKING;
 
-  return {
+  // Every var is load-bearing (ablation-tested): the CLI stamps the per-turn
+  // traceparent only once its OTel tracer initializes, and the dead endpoint
+  // keeps the throwaway spans off any local collector. Exporter and protocol
+  // are pinned rather than inherited — an ambient OTEL_TRACES_EXPORTER=none or
+  // unknown protocol registers no tracer and silently drops the traceparent;
+  // the endpoint stays overridable for a real collector.
+  // Residual risk: a repo's .claude/settings.json `env` is applied over these
+  // inside the CLI and can redirect the endpoint or turn on content capture
+  // (OTEL_LOG_TOOL_CONTENT, …) — pre-existing settingSources exposure, not
+  // closable from here; hardening tracked separately.
+  const gatewayTracing: Record<string, string> = gateway?.anthropicBaseUrl
+    ? {
+        CLAUDE_CODE_ENABLE_TELEMETRY: "1",
+        CLAUDE_CODE_ENHANCED_TELEMETRY_BETA: "1",
+        CLAUDE_CODE_PROPAGATE_TRACEPARENT: "1",
+        OTEL_TRACES_EXPORTER: "otlp",
+        OTEL_EXPORTER_OTLP_PROTOCOL: "http/json",
+        OTEL_EXPORTER_OTLP_ENDPOINT:
+          process.env.OTEL_EXPORTER_OTLP_ENDPOINT ?? "http://127.0.0.1:9",
+      }
+    : {};
+
+  const env: Record<string, string> = {
     ...process.env,
+    ...gatewayTracing,
     // Explicit gateway values win over whatever happens to be in process.env.
     // This prevents concurrent Agent instances from clobbering each other's
     // gateway config when process.env was mutated globally.
@@ -212,6 +243,13 @@ function buildEnvironment(gateway?: GatewayEnv): Record<string, string> {
     }),
     ANTHROPIC_CUSTOM_HEADERS: customHeaders,
   };
+  if (gateway?.anthropicBaseUrl) {
+    // The CLI parents every turn under an inherited ambient TRACEPARENT,
+    // collapsing the per-turn trace ids this block exists to produce.
+    delete env.TRACEPARENT;
+    delete env.TRACESTATE;
+  }
+  return env;
 }
 
 function buildHooks(
@@ -473,7 +511,7 @@ export function buildSessionOptions(params: BuildOptionsParams): Options {
       params.mcpServers,
       loadUserClaudeJsonMcpServers(params.cwd, params.logger),
     ),
-    env: buildEnvironment(params.gatewayEnv),
+    env: buildEnvironment(params.gatewayEnv, params.sessionId),
     hooks: buildHooks(
       params.userProvidedOptions?.hooks,
       params.onModeChange,

packages/agent/src/adapters/codex-app-server/spawn.test.ts

@@ -57,6 +57,18 @@ describe("buildAppServerArgs", () => {
     );
   });
 
+  it("quotes $-prefixed posthog property header keys in the TOML table", () => {
+    const args = buildAppServerArgs({
+      binaryPath: "/bundle/codex",
+      apiBaseUrl: "https://gateway.example/v1",
+      httpHeaders: { "x-posthog-property-$ai_session_id": "task-123" },
+    });
+
+    expect(args).toContain(
+      'model_providers.posthog.http_headers={ "x-posthog-property-$ai_session_id" = "task-123" }',
+    );
+  });
+
   it("omits http_headers when none are provided or the provider is unset", () => {
     const withoutHeaders = buildAppServerArgs({
       binaryPath: "/bundle/codex",

packages/agent/src/agent.ts

@@ -14,6 +14,7 @@ import {
 import { PostHogAPIClient, type TaskRunUpdate } from "./posthog-api";
 import { SessionLogWriter } from "./session-log-writer";
 import type { AgentConfig, TaskExecutionOptions } from "./types";
+import { buildGatewayPropertyHeaderRecord } from "./utils/gateway";
 import { Logger } from "./utils/logger";
 
 export class Agent {
@@ -148,6 +149,9 @@ export class Agent {
               model: sanitizedModel,
               reasoningEffort: options.reasoningEffort,
               developerInstructions: options.developerInstructions,
+              httpHeaders: taskId
+                ? buildGatewayPropertyHeaderRecord({ $ai_session_id: taskId })
+                : undefined,
               additionalDirectories: options.additionalDirectories,
             }
           : undefined,

packages/agent/src/server/agent-server.configure-environment.test.ts

@@ -223,6 +223,7 @@ describe("AgentServer.configureEnvironment", () => {
       "x-posthog-property-task_user_id": "42",
       "x-posthog-property-task_title": "Fix the bug",
       "x-posthog-property-team_id": "1",
+      "x-posthog-property-$ai_session_id": "task-abc",
     });
   });
 
@@ -340,6 +341,25 @@ describe("AgentServer.configureEnvironment", () => {
     );
   });
 
+  it("folds the task id into the codex session header only", () => {
+    const env = buildServer("interactive").configureEnvironment({
+      taskId: "task-123",
+    });
+
+    expect(env.openaiCustomHeaders?.["x-posthog-property-$ai_session_id"]).toBe(
+      "task-123",
+    );
+    expect(env.anthropicCustomHeaders ?? "").not.toContain("$ai_session_id");
+  });
+
+  it("omits the codex session header without a task id", () => {
+    const env = buildServer("interactive").configureEnvironment({});
+
+    expect(
+      env.openaiCustomHeaders?.["x-posthog-property-$ai_session_id"],
+    ).toBeUndefined();
+  });
+
   it("appends the resolved product to a LLM_GATEWAY_URL override base", () => {
     // The override is treated as a base URL. The product slug is always
     // appended so the gateway routes to the correct product config — a bare

packages/agent/src/server/agent-server.ts

@@ -3975,9 +3975,12 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions}${
       openaiCustomHeaders = buildGatewayPropertiesHeaderRecord(properties);
     } else {
       customHeaders = buildGatewayPropertyHeaders(gatewayProperties);
+      // No $ai_session_id on the Go-gateway path above: it strips $-prefixed
+      // blob keys, so the session id would be silently dropped there.
       openaiCustomHeaders = buildGatewayPropertyHeaderRecord({
         ...gatewayProperties,
         team_id: projectId,
+        $ai_session_id: taskId,
       });
     }