Skip to content

Commit 4d12a67

Browse files
joshsnyjoshsny
authored
feat(agent): attribute LLM gateway events to the customer team
Every LLM call from the agent package flows through the PostHog LLM gateway, which authenticates with a shared key. Without a per-request team override, every captured `$ai_generation` event is attributed to the key owner's team, so per-customer spend can't be rolled up. Mirror django's `get_llm_client(team_id=...)` (PostHog/posthog#60745): forward the team as an `x-posthog-property-team_id` header, which the gateway lifts onto the captured event as a `team_id` property. - Cloud path (`agent-server.ts`): add `team_id` to the task-metadata headers already built in `configureEnvironment`. - Desktop path (`agent.ts`): set the `team_id` property header in `_configureLlmGateway`, deduping by header name so re-configuring across sessions doesn't append it twice and existing headers (e.g. the Bedrock fallback) are preserved. - Expose `PostHogAPIClient.getProjectId()` for the desktop path. Generated-By: PostHog Code Task-Id: a2593f98-9dc7-4fa7-a532-5257b2d5c6b9
1 parent 7d82e7f commit 4d12a67

6 files changed

Lines changed: 148 additions & 2 deletions

File tree

‎packages/agent/src/agent.configure-gateway.test.ts‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
2+
import { Agent } from "./agent";
3+
4+
interface TestableAgent {
5+
_configureLlmGateway(
6+
overrideUrl?: string,
7+
): Promise<{ gatewayUrl: string; apiKey: string } | null>;
8+
}
9+
10+
const ENV_KEYS_UNDER_TEST = [
11+
"ANTHROPIC_BASE_URL",
12+
"ANTHROPIC_AUTH_TOKEN",
13+
"ANTHROPIC_CUSTOM_HEADERS",
14+
"OPENAI_BASE_URL",
15+
"OPENAI_API_KEY",
16+
] as const;
17+
18+
describe("Agent._configureLlmGateway", () => {
19+
const originalEnv: Partial<Record<string, string | undefined>> = {};
20+
21+
beforeEach(() => {
22+
for (const key of ENV_KEYS_UNDER_TEST) {
23+
originalEnv[key] = process.env[key];
24+
delete process.env[key];
25+
}
26+
});
27+
28+
afterEach(() => {
29+
for (const key of ENV_KEYS_UNDER_TEST) {
30+
const value = originalEnv[key];
31+
if (value === undefined) {
32+
delete process.env[key];
33+
} else {
34+
process.env[key] = value;
35+
}
36+
}
37+
});
38+
39+
const buildAgent = (): TestableAgent =>
40+
new Agent({
41+
skipLogPersistence: true,
42+
posthog: {
43+
apiUrl: "https://us.posthog.com",
44+
getApiKey: vi.fn().mockResolvedValue("test-token"),
45+
projectId: 99,
46+
},
47+
}) as unknown as TestableAgent;
48+
49+
it("forwards the team_id as an x-posthog-property header", async () => {
50+
await buildAgent()._configureLlmGateway();
51+
52+
expect(process.env.ANTHROPIC_CUSTOM_HEADERS).toBe(
53+
"x-posthog-property-team_id: 99",
54+
);
55+
});
56+
57+
it("preserves pre-existing custom headers and dedupes the team_id line", async () => {
58+
process.env.ANTHROPIC_CUSTOM_HEADERS = [
59+
"x-posthog-property-team_id: 1",
60+
"x-posthog-use-bedrock-fallback: true",
61+
].join("\n");
62+
63+
await buildAgent()._configureLlmGateway();
64+
65+
expect(process.env.ANTHROPIC_CUSTOM_HEADERS).toBe(
66+
[
67+
"x-posthog-property-team_id: 99",
68+
"x-posthog-use-bedrock-fallback: true",
69+
].join("\n"),
70+
);
71+
});
72+
});

‎packages/agent/src/agent.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
import { PostHogAPIClient, type TaskRunUpdate } from "./posthog-api";
1212
import { SessionLogWriter } from "./session-log-writer";
1313
import type { AgentConfig, TaskExecutionOptions } from "./types";
14+
import { buildGatewayPropertyHeaders } from "./utils/gateway";
1415
import { Logger } from "./utils/logger";
1516

1617
export class Agent {
@@ -67,13 +68,56 @@ export class Agent {
6768
process.env.ANTHROPIC_BASE_URL = gatewayUrl;
6869
process.env.ANTHROPIC_AUTH_TOKEN = apiKey;
6970

71+
// Attribute every captured $ai_generation event to this team. The gateway
72+
// authenticates with a shared key, so without the `team_id` property the
73+
// spend lands on the key owner's team. Forwarded as an
74+
// `x-posthog-property-team_id` header that the gateway lifts onto the
75+
// event (the Claude session builder appends its own headers to this in
76+
// adapters/claude/session/options.ts). Mirrors the cloud path in
77+
// server/agent-server.ts and django's get_llm_client(team_id=...).
78+
this._applyGatewayPropertyHeaders({
79+
team_id: this.posthogAPI.getProjectId(),
80+
});
81+
7082
return { gatewayUrl, apiKey };
7183
} catch (error) {
7284
this.logger.error("Failed to configure LLM gateway", error);
7385
throw error;
7486
}
7587
}
7688

89+
/**
90+
* Merge `x-posthog-property-*` header lines into `ANTHROPIC_CUSTOM_HEADERS`,
91+
* deduping by header name so re-configuring across sessions doesn't append
92+
* the same property twice. Existing non-property lines are preserved.
93+
*/
94+
private _applyGatewayPropertyHeaders(
95+
properties: Record<string, string | number | boolean | null | undefined>,
96+
): void {
97+
const lines = new Map<string, string>();
98+
const existing = process.env.ANTHROPIC_CUSTOM_HEADERS;
99+
if (existing) {
100+
for (const line of existing.split("\n")) {
101+
const name = line.slice(0, line.indexOf(":")).trim();
102+
if (name) {
103+
lines.set(name, line);
104+
}
105+
}
106+
}
107+
108+
const additions = buildGatewayPropertyHeaders(properties);
109+
if (additions) {
110+
for (const line of additions.split("\n")) {
111+
const name = line.slice(0, line.indexOf(":")).trim();
112+
lines.set(name, line);
113+
}
114+
}
115+
116+
process.env.ANTHROPIC_CUSTOM_HEADERS = Array.from(lines.values()).join(
117+
"\n",
118+
);
119+
}
120+
77121
async run(
78122
taskId: string,
79123
taskRunId: string,

‎packages/agent/src/posthog-api.test.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,16 @@ describe("PostHogAPIClient", () => {
1010
vi.clearAllMocks();
1111
});
1212

13+
it("exposes the configured project id", () => {
14+
const client = new PostHogAPIClient({
15+
apiUrl: "https://app.posthog.com",
16+
getApiKey: vi.fn().mockResolvedValue("token"),
17+
projectId: 42,
18+
});
19+
20+
expect(client.getProjectId()).toBe(42);
21+
});
22+
1323
it("refreshes once when fetching task run logs gets an auth failure", async () => {
1424
const getApiKey = vi.fn().mockResolvedValue("stale-token");
1525
const refreshApiKey = vi.fn().mockResolvedValue("fresh-token");

‎packages/agent/src/posthog-api.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,10 @@ export class PostHogAPIClient {
133133
return this.config.projectId;
134134
}
135135

136+
getProjectId(): number {
137+
return this.config.projectId;
138+
}
139+
136140
async getApiKey(forceRefresh = false): Promise<string> {
137141
return this.resolveApiKey(forceRefresh);
138142
}

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ describe("AgentServer.configureEnvironment", () => {
134134

135135
expect(process.env.ANTHROPIC_CUSTOM_HEADERS).toBe(
136136
[
137+
"x-posthog-property-team_id: 1",
137138
"x-posthog-property-task_origin_product: signal_report",
138139
"x-posthog-property-task_internal: true",
139140
"x-posthog-property-signal_report_id: report-123",
@@ -144,6 +145,14 @@ describe("AgentServer.configureEnvironment", () => {
144145
);
145146
});
146147

148+
it("always forwards the team_id as an x-posthog-property header", () => {
149+
buildServer("background").configureEnvironment({ isInternal: false });
150+
151+
expect(process.env.ANTHROPIC_CUSTOM_HEADERS).toContain(
152+
"x-posthog-property-team_id: 1",
153+
);
154+
});
155+
147156
it("omits signal_report_id from ANTHROPIC_CUSTOM_HEADERS for non-report tasks", () => {
148157
buildServer("background").configureEnvironment({
149158
isInternal: false,
@@ -159,7 +168,10 @@ describe("AgentServer.configureEnvironment", () => {
159168
buildServer("background").configureEnvironment({ isInternal: false });
160169

161170
expect(process.env.ANTHROPIC_CUSTOM_HEADERS).toBe(
162-
"x-posthog-property-task_internal: false",
171+
[
172+
"x-posthog-property-team_id: 1",
173+
"x-posthog-property-task_internal: false",
174+
].join("\n"),
163175
);
164176
});
165177

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1893,8 +1893,12 @@ ${signedCommitInstructions}
18931893
// Forward task metadata as `x-posthog-property-*` headers so the gateway
18941894
// lifts them onto the $ai_generation event. Routes through the Anthropic
18951895
// SDK's ANTHROPIC_CUSTOM_HEADERS env var; the OpenAI/codex path has no
1896-
// equivalent today.
1896+
// equivalent today. `team_id` attributes every captured generation to the
1897+
// customer's PostHog team (the gateway authenticates with a shared key, so
1898+
// without this the spend lands on the key owner's team — see the django
1899+
// `get_llm_client(team_id=...)` equivalent in posthog/llm/gateway_client.py).
18971900
const customHeaders = buildGatewayPropertyHeaders({
1901+
team_id: projectId,
18981902
task_origin_product: originProduct,
18991903
task_internal: isInternal,
19001904
signal_report_id: signalReportId,

0 commit comments

Comments
 (0)