diff --git a/packages/agent/src/adapters/claude/session/options.test.ts b/packages/agent/src/adapters/claude/session/options.test.ts index f2a0b66c99..418b2c8d9e 100644 --- a/packages/agent/src/adapters/claude/session/options.test.ts +++ b/packages/agent/src/adapters/claude/session/options.test.ts @@ -376,7 +376,10 @@ 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", @@ -384,6 +387,7 @@ describe("buildSessionOptions", () => { 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> = {}; + + 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", () => { diff --git a/packages/agent/src/adapters/claude/session/options.ts b/packages/agent/src/adapters/claude/session/options.ts index a5e08ccc79..25c6743d2f 100644 --- a/packages/agent/src/adapters/claude/session/options.ts +++ b/packages/agent/src/adapters/claude/session/options.ts @@ -151,7 +151,10 @@ function buildMcpServers( }; } -function buildEnvironment(gateway?: GatewayEnv): Record { +function buildEnvironment( + gateway?: GatewayEnv, + sessionId?: string, +): Record { // 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 { 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 { // 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 = 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 = { ...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 { }), 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, diff --git a/packages/agent/src/adapters/codex-app-server/spawn.test.ts b/packages/agent/src/adapters/codex-app-server/spawn.test.ts index 6794bef2dd..a5784252d4 100644 --- a/packages/agent/src/adapters/codex-app-server/spawn.test.ts +++ b/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", diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 63b3830849..b7ea903d8d 100644 --- a/packages/agent/src/agent.ts +++ b/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, diff --git a/packages/agent/src/server/agent-server.configure-environment.test.ts b/packages/agent/src/server/agent-server.configure-environment.test.ts index 0ae96970f6..6b18e64f22 100644 --- a/packages/agent/src/server/agent-server.configure-environment.test.ts +++ b/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 diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 990fd388db..e1f075ebcf 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/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, }); }