feat(agent): attribute LLM gateway events to the customer team

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);
+    });
+  });
 });

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

@@ -112,11 +112,28 @@ function buildMcpServers(
 }
 
 function buildEnvironment(): Record<string, string> {
-  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<string, string> {
     ...(mcpNonblocking !== undefined && {
       MCP_CONNECTION_NONBLOCKING: mcpNonblocking,
     }),
-    // Route to AWS Bedrock as a fallback when Anthropic returns 5xx
     ANTHROPIC_CUSTOM_HEADERS: customHeaders,
   };
 }

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();
 

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,