diff --git a/packages/agent/src/adapters/claude/session/options.test.ts b/packages/agent/src/adapters/claude/session/options.test.ts index 7c843dc593..e412e7e124 100644 --- a/packages/agent/src/adapters/claude/session/options.test.ts +++ b/packages/agent/src/adapters/claude/session/options.test.ts @@ -1,6 +1,6 @@ import * as os from "node:os"; import * as path from "node:path"; -import { describe, expect, it } from "vitest"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { Logger } from "../../../utils/logger"; import { SUBAGENT_REWRITES } from "../hooks"; import { buildSessionOptions } from "./options"; @@ -70,4 +70,67 @@ describe("buildSessionOptions", () => { expect(options.agents?.["ph-explore"]).toEqual(override); }); + + describe("ANTHROPIC_CUSTOM_HEADERS", () => { + const originalProjectId = process.env.POSTHOG_PROJECT_ID; + const originalCustomHeaders = process.env.ANTHROPIC_CUSTOM_HEADERS; + + beforeEach(() => { + delete process.env.POSTHOG_PROJECT_ID; + delete process.env.ANTHROPIC_CUSTOM_HEADERS; + }); + + afterEach(() => { + for (const [key, value] of [ + ["POSTHOG_PROJECT_ID", originalProjectId], + ["ANTHROPIC_CUSTOM_HEADERS", originalCustomHeaders], + ] as const) { + if (value === undefined) { + delete process.env[key]; + } else { + process.env[key] = value; + } + } + }); + + it.each([ + { + name: "omits the team_id header when POSTHOG_PROJECT_ID is unset", + projectId: undefined, + existingHeaders: undefined, + expected: "x-posthog-use-bedrock-fallback: true", + }, + { + name: "forwards POSTHOG_PROJECT_ID as the team_id attribution header", + projectId: "42", + existingHeaders: undefined, + expected: [ + "x-posthog-property-team_id: 42", + "x-posthog-use-bedrock-fallback: true", + ].join("\n"), + }, + { + name: "preserves pre-existing custom headers ahead of the team_id header", + projectId: "42", + existingHeaders: "x-posthog-property-task_id: task-abc", + expected: [ + "x-posthog-property-task_id: task-abc", + "x-posthog-property-team_id: 42", + "x-posthog-use-bedrock-fallback: true", + ].join("\n"), + }, + ])("$name", ({ projectId, existingHeaders, expected }) => { + if (projectId !== undefined) { + process.env.POSTHOG_PROJECT_ID = projectId; + } + if (existingHeaders !== undefined) { + process.env.ANTHROPIC_CUSTOM_HEADERS = existingHeaders; + } + + const headers = buildSessionOptions(makeParams()).env + ?.ANTHROPIC_CUSTOM_HEADERS; + + expect(headers).toBe(expected); + }); + }); }); diff --git a/packages/agent/src/adapters/claude/session/options.ts b/packages/agent/src/adapters/claude/session/options.ts index b3c2683a7a..c87fb4f096 100644 --- a/packages/agent/src/adapters/claude/session/options.ts +++ b/packages/agent/src/adapters/claude/session/options.ts @@ -112,11 +112,28 @@ function buildMcpServers( } function buildEnvironment(): Record { - const bedrockFallbackHeader = "x-posthog-use-bedrock-fallback: true"; + // 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 + // chokepoint every session (desktop and cloud) funnels through. + const headerLines: string[] = []; const existingCustomHeaders = process.env.ANTHROPIC_CUSTOM_HEADERS; - const customHeaders = existingCustomHeaders - ? `${existingCustomHeaders}\n${bedrockFallbackHeader}` - : bedrockFallbackHeader; + if (existingCustomHeaders) { + headerLines.push(existingCustomHeaders); + } + // Attribute every captured $ai_generation event to the customer's team. The + // gateway authenticates with a shared key, so without this the spend lands on + // the key owner's team. The gateway lifts `x-posthog-property-*` headers onto + // the event; both entrypoints export POSTHOG_PROJECT_ID before this runs + // (apps/code auth-adapter.ts, server/agent-server.ts). Mirrors django's + // get_llm_client(team_id=...). + const projectId = process.env.POSTHOG_PROJECT_ID; + if (projectId) { + headerLines.push(`x-posthog-property-team_id: ${projectId}`); + } + // Route to AWS Bedrock as a fallback when Anthropic returns 5xx + headerLines.push("x-posthog-use-bedrock-fallback: true"); + const customHeaders = headerLines.join("\n"); // SDK 0.3.142 made MCP servers connect in the background by default. That // default is what we want: a slow or unreachable user MCP server (PostHog @@ -136,7 +153,6 @@ function buildEnvironment(): Record { ...(mcpNonblocking !== undefined && { MCP_CONNECTION_NONBLOCKING: mcpNonblocking, }), - // Route to AWS Bedrock as a fallback when Anthropic returns 5xx ANTHROPIC_CUSTOM_HEADERS: customHeaders, }; } 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 fd81726ad8..69871e5cb9 100644 --- a/packages/agent/src/server/agent-server.configure-environment.test.ts +++ b/packages/agent/src/server/agent-server.configure-environment.test.ts @@ -18,6 +18,7 @@ const ENV_KEYS_UNDER_TEST = [ "ANTHROPIC_BASE_URL", "OPENAI_BASE_URL", "ANTHROPIC_CUSTOM_HEADERS", + "POSTHOG_PROJECT_ID", ] as const; describe("AgentServer.configureEnvironment", () => { @@ -75,6 +76,15 @@ describe("AgentServer.configureEnvironment", () => { ); }); + // The Claude session builder reads POSTHOG_PROJECT_ID to emit the + // `x-posthog-property-team_id` attribution header (see + // adapters/claude/session/options.ts), so the cloud path must export it. + it("exports POSTHOG_PROJECT_ID for the team_id attribution header", () => { + buildServer("background").configureEnvironment({ isInternal: false }); + + expect(process.env.POSTHOG_PROJECT_ID).toBe("1"); + }); + it("tags as posthog_code when isInternal is omitted (getTask failure fallback)", () => { buildServer("background").configureEnvironment(); diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index c5307e44ff..8c0f4aa72b 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -1893,7 +1893,9 @@ ${signedCommitInstructions} // Forward task metadata as `x-posthog-property-*` headers so the gateway // lifts them onto the $ai_generation event. Routes through the Anthropic // SDK's ANTHROPIC_CUSTOM_HEADERS env var; the OpenAI/codex path has no - // equivalent today. + // equivalent today. (The `team_id` attribution header is added downstream + // in the Claude session builder from POSTHOG_PROJECT_ID — see + // adapters/claude/session/options.ts.) const customHeaders = buildGatewayPropertyHeaders({ task_origin_product: originProduct, task_internal: isInternal,