feat(agent): attribute LLM gateway events to the customer team#2454
Open
joshsny wants to merge 4 commits into
Open
feat(agent): attribute LLM gateway events to the customer team#2454joshsny wants to merge 4 commits into
joshsny wants to merge 4 commits into
Conversation
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
…hokepoint Replace the two-site implementation (desktop `agent.ts` + cloud `agent-server.ts`) with a single definition in the Claude session builder's `buildEnvironment`, where `ANTHROPIC_CUSTOM_HEADERS` is already finalized for every session. Both entrypoints already export `POSTHOG_PROJECT_ID` (apps/code auth-adapter.ts, server/agent-server.ts), so reading it there covers desktop and cloud uniformly. `buildEnvironment` returns a fresh env object rather than mutating `process.env`, so the cross-session dedup helper is no longer needed — and the `PostHogAPIClient.getProjectId()` accessor and the `agent.ts` changes are dropped entirely. Net: `team_id` attribution is defined once, the diff is smaller, and the cloud header block keeps a comment pointing to where it now lives. Generated-By: PostHog Code Task-Id: a2593f98-9dc7-4fa7-a532-5257b2d5c6b9
The team_id attribution header is emitted from POSTHOG_PROJECT_ID in the Claude session builder, so the cloud configureEnvironment must export it. The desktop path is already guarded (auth-adapter.test.ts); this closes the matching gap for the cloud path. Generated-By: PostHog Code Task-Id: a2593f98-9dc7-4fa7-a532-5257b2d5c6b9
Contributor
Prompt To Fix All With AIFix the following 1 code review issue. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 1
packages/agent/src/adapters/claude/session/options.test.ts:96-133
The three new tests share the same structure (set env, call `buildSessionOptions`, assert on `ANTHROPIC_CUSTOM_HEADERS`) and differ only in their inputs and expected output, which is an ideal fit for `it.each`. Flattening them into a parameterised table removes the repeated `beforeEach`/`afterEach` scaffolding from the outer describe block and aligns with the team's stated preference for parameterised tests.
```suggestion
it.each([
[
"omits team_id when POSTHOG_PROJECT_ID is unset",
undefined,
undefined,
"x-posthog-use-bedrock-fallback: true",
],
[
"forwards POSTHOG_PROJECT_ID as the team_id attribution header",
"42",
undefined,
["x-posthog-property-team_id: 42", "x-posthog-use-bedrock-fallback: true"].join("\n"),
],
[
"preserves pre-existing custom headers ahead of the team_id header",
"42",
"x-posthog-property-task_id: task-abc",
[
"x-posthog-property-task_id: task-abc",
"x-posthog-property-team_id: 42",
"x-posthog-use-bedrock-fallback: true",
].join("\n"),
],
] as const)(
"%s",
(_, 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);
},
);
```
Reviews (1): Last reviewed commit: "test(agent): guard cloud POSTHOG_PROJECT..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+96
to
+133
| it("forwards POSTHOG_PROJECT_ID as the team_id attribution header", () => { | ||
| process.env.POSTHOG_PROJECT_ID = "42"; | ||
|
|
||
| const headers = buildSessionOptions(makeParams()).env | ||
| ?.ANTHROPIC_CUSTOM_HEADERS; | ||
|
|
||
| expect(headers).toBe( | ||
| [ | ||
| "x-posthog-property-team_id: 42", | ||
| "x-posthog-use-bedrock-fallback: true", | ||
| ].join("\n"), | ||
| ); | ||
| }); | ||
|
|
||
| it("preserves pre-existing custom headers ahead of the team_id header", () => { | ||
| process.env.POSTHOG_PROJECT_ID = "42"; | ||
| process.env.ANTHROPIC_CUSTOM_HEADERS = | ||
| "x-posthog-property-task_id: task-abc"; | ||
|
|
||
| const headers = buildSessionOptions(makeParams()).env | ||
| ?.ANTHROPIC_CUSTOM_HEADERS; | ||
|
|
||
| expect(headers).toBe( | ||
| [ | ||
| "x-posthog-property-task_id: task-abc", | ||
| "x-posthog-property-team_id: 42", | ||
| "x-posthog-use-bedrock-fallback: true", | ||
| ].join("\n"), | ||
| ); | ||
| }); | ||
|
|
||
| it("omits the team_id header when POSTHOG_PROJECT_ID is unset", () => { | ||
| const headers = buildSessionOptions(makeParams()).env | ||
| ?.ANTHROPIC_CUSTOM_HEADERS; | ||
|
|
||
| expect(headers).toBe("x-posthog-use-bedrock-fallback: true"); | ||
| expect(headers).not.toContain("team_id"); | ||
| }); |
Contributor
There was a problem hiding this comment.
buildSessionOptions, assert on ANTHROPIC_CUSTOM_HEADERS) and differ only in their inputs and expected output, which is an ideal fit for it.each. Flattening them into a parameterised table removes the repeated beforeEach/afterEach scaffolding from the outer describe block and aligns with the team's stated preference for parameterised tests.
Suggested change
| it("forwards POSTHOG_PROJECT_ID as the team_id attribution header", () => { | |
| process.env.POSTHOG_PROJECT_ID = "42"; | |
| const headers = buildSessionOptions(makeParams()).env | |
| ?.ANTHROPIC_CUSTOM_HEADERS; | |
| expect(headers).toBe( | |
| [ | |
| "x-posthog-property-team_id: 42", | |
| "x-posthog-use-bedrock-fallback: true", | |
| ].join("\n"), | |
| ); | |
| }); | |
| it("preserves pre-existing custom headers ahead of the team_id header", () => { | |
| process.env.POSTHOG_PROJECT_ID = "42"; | |
| process.env.ANTHROPIC_CUSTOM_HEADERS = | |
| "x-posthog-property-task_id: task-abc"; | |
| const headers = buildSessionOptions(makeParams()).env | |
| ?.ANTHROPIC_CUSTOM_HEADERS; | |
| expect(headers).toBe( | |
| [ | |
| "x-posthog-property-task_id: task-abc", | |
| "x-posthog-property-team_id: 42", | |
| "x-posthog-use-bedrock-fallback: true", | |
| ].join("\n"), | |
| ); | |
| }); | |
| it("omits the team_id header when POSTHOG_PROJECT_ID is unset", () => { | |
| const headers = buildSessionOptions(makeParams()).env | |
| ?.ANTHROPIC_CUSTOM_HEADERS; | |
| expect(headers).toBe("x-posthog-use-bedrock-fallback: true"); | |
| expect(headers).not.toContain("team_id"); | |
| }); | |
| it.each([ | |
| [ | |
| "omits team_id when POSTHOG_PROJECT_ID is unset", | |
| undefined, | |
| undefined, | |
| "x-posthog-use-bedrock-fallback: true", | |
| ], | |
| [ | |
| "forwards POSTHOG_PROJECT_ID as the team_id attribution header", | |
| "42", | |
| undefined, | |
| ["x-posthog-property-team_id: 42", "x-posthog-use-bedrock-fallback: true"].join("\n"), | |
| ], | |
| [ | |
| "preserves pre-existing custom headers ahead of the team_id header", | |
| "42", | |
| "x-posthog-property-task_id: task-abc", | |
| [ | |
| "x-posthog-property-task_id: task-abc", | |
| "x-posthog-property-team_id: 42", | |
| "x-posthog-use-bedrock-fallback: true", | |
| ].join("\n"), | |
| ], | |
| ] as const)( | |
| "%s", | |
| (_, 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); | |
| }, | |
| ); |
Prompt To Fix With AI
This is a comment left during a code review.
Path: packages/agent/src/adapters/claude/session/options.test.ts
Line: 96-133
Comment:
The three new tests share the same structure (set env, call `buildSessionOptions`, assert on `ANTHROPIC_CUSTOM_HEADERS`) and differ only in their inputs and expected output, which is an ideal fit for `it.each`. Flattening them into a parameterised table removes the repeated `beforeEach`/`afterEach` scaffolding from the outer describe block and aligns with the team's stated preference for parameterised tests.
```suggestion
it.each([
[
"omits team_id when POSTHOG_PROJECT_ID is unset",
undefined,
undefined,
"x-posthog-use-bedrock-fallback: true",
],
[
"forwards POSTHOG_PROJECT_ID as the team_id attribution header",
"42",
undefined,
["x-posthog-property-team_id: 42", "x-posthog-use-bedrock-fallback: true"].join("\n"),
],
[
"preserves pre-existing custom headers ahead of the team_id header",
"42",
"x-posthog-property-task_id: task-abc",
[
"x-posthog-property-task_id: task-abc",
"x-posthog-property-team_id: 42",
"x-posthog-use-bedrock-fallback: true",
].join("\n"),
],
] as const)(
"%s",
(_, 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);
},
);
```
How can I resolve this? If you propose a fix, please make it concise.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Flatten the three ANTHROPIC_CUSTOM_HEADERS cases into a single it.each table, per Greptile review feedback and the team's preference for parameterized tests. Keeps the beforeEach/afterEach env isolation so the cases stay order-independent and process.env is restored afterward. Generated-By: PostHog Code Task-Id: a2593f98-9dc7-4fa7-a532-5257b2d5c6b9
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Every LLM call from the
packages/agentagent flows through the PostHog LLM gateway, which authenticates with a single shared personal API key. The gateway attributes each captured$ai_generationevent to the key owner's team unless the request carries a per-team override. That means agent LLM spend can't be rolled up per customer team.This mirrors the issue solved on the django side for API-key events in PostHog/posthog#60745, which added
get_llm_client(team_id=...). That PR forwards the team as anx-posthog-property-team_idheader, and the gateway liftsx-posthog-property-*headers onto the captured event as properties. That gateway mechanism is already live (the agent already uses it to forward task metadata), so this change is standalone and does not depend on the django PR landing.Changes
Forward the PostHog team as an
x-posthog-property-team_idheader on every agent → gateway request, so all$ai_generationevents get ateam_idproperty attributing them to the correct customer team.The header is added in one place:
buildEnvironment()in the Claude session builder (adapters/claude/session/options.ts), which is the single chokepoint every session (desktop and cloud) funnels through and whereANTHROPIC_CUSTOM_HEADERSis already finalized (alongside the Bedrock-fallback header). It readsPOSTHOG_PROJECT_ID, which both entrypoints already export (apps/codeauth-adapter.tsfor desktop,server/agent-server.tsfor cloud).Why this lives in an env var rather than a direct option: the Claude Agent SDK spawns the Claude Code CLI as a child process and that subprocess makes the HTTP calls — the SDK's
Optionstype has no custom-header field, soANTHROPIC_CUSTOM_HEADERS(which the CLI reads and forwards) is the only mechanism. The OpenAI/codex path has no header equivalent today — same known limitation as the existing task-metadata headers.How did you test this?
adapters/claude/session/options.test.ts(run, passing):POSTHOG_PROJECT_IDasx-posthog-property-team_idPOSTHOG_PROJECT_IDis unsetserver/agent-server.configure-environment.test.tsandposthog-api.test.tsstill pass.pnpm --filter agent typecheck— clean.biome checkon the changed files — clean.The remaining failures in the full
pnpm --filter agent testrun are pre-existing and unrelated: test helpers invokegit commit, which is blocked in this signed-commit environment.Automatic notifications
Created with PostHog Code