From 699d09ae8fc62f560c3ebd7cc823436cfee4279b Mon Sep 17 00:00:00 2001 From: Brandon Leung Date: Tue, 21 Jul 2026 19:08:53 -0400 Subject: [PATCH 1/2] feat(agent): route selected products to the ai-gateway Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_015gyKxP2MuVuNRCK9uZt1Lr --- ...agent-server.configure-environment.test.ts | 154 +++++++++++++- packages/agent/src/server/agent-server.ts | 41 +++- packages/agent/src/utils/gateway.test.ts | 188 ++++++++++++++++++ packages/agent/src/utils/gateway.ts | 109 ++++++++++ .../src/posthog-property-headers.test.ts | 131 ++++++++++++ .../shared/src/posthog-property-headers.ts | 84 ++++++++ 6 files changed, 695 insertions(+), 12 deletions(-) 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 f5b96ccea9..66a67db0a4 100644 --- a/packages/agent/src/server/agent-server.configure-environment.test.ts +++ b/packages/agent/src/server/agent-server.configure-environment.test.ts @@ -16,7 +16,12 @@ interface TestableServer { }): GatewayEnv; } -const ENV_KEYS_UNDER_TEST = ["LLM_GATEWAY_URL", "POSTHOG_PROJECT_ID"] as const; +const ENV_KEYS_UNDER_TEST = [ + "LLM_GATEWAY_URL", + "POSTHOG_PROJECT_ID", + "AI_GATEWAY_URL", + "AI_GATEWAY_PRODUCTS", +] as const; describe("AgentServer.configureEnvironment", () => { const originalEnv: Partial> = {}; @@ -323,3 +328,150 @@ describe("AgentServer.configureEnvironment", () => { ); }); }); + +describe("AgentServer.configureEnvironment on the Go ai-gateway", () => { + const originalEnv: Partial> = {}; + const ENV_KEYS = [ + "LLM_GATEWAY_URL", + "POSTHOG_PROJECT_ID", + "AI_GATEWAY_URL", + "AI_GATEWAY_PRODUCTS", + ]; + const GO_GATEWAY = "https://ai-gateway.us.posthog.com"; + + beforeEach(() => { + for (const key of ENV_KEYS) { + originalEnv[key] = process.env[key]; + delete process.env[key]; + } + process.env.AI_GATEWAY_URL = GO_GATEWAY; + process.env.AI_GATEWAY_PRODUCTS = [ + "signals_scout", + "signals_research", + "signals_implementation", + "signals_repo_selection", + "background_agents", + ].join(","); + }); + + afterEach(() => { + for (const key of ENV_KEYS) { + const value = originalEnv[key]; + if (value === undefined) { + delete process.env[key]; + } else { + process.env[key] = value; + } + } + }); + + const buildServer = (): TestableServer => + new AgentServer({ + port: 0, + jwtPublicKey: "test-key", + apiUrl: "https://us.posthog.com", + apiKey: "test-api-key", + projectId: 42, + mode: "background", + taskId: "test-task-id", + runId: "test-run-id", + }) as unknown as TestableServer; + + const parseBlob = (headerLines: string): Record => { + const prefix = "X-PostHog-Properties: "; + expect(headerLines.startsWith(prefix)).toBe(true); + return JSON.parse(headerLines.slice(prefix.length)); + }; + + it("drops the product slug from the base URLs", () => { + const env = buildServer().configureEnvironment({ + originProduct: "signals_scout", + aiStage: "scout", + }); + + expect(env.anthropicBaseUrl).toBe(GO_GATEWAY); + expect(env.openaiBaseUrl).toBe(`${GO_GATEWAY}/v1`); + }); + + it("honours an explicit AI_GATEWAY_URL with a trailing /v1", () => { + process.env.AI_GATEWAY_URL = "https://ai-gateway.dev.posthog.dev/v1"; + const env = buildServer().configureEnvironment({ + originProduct: "signal_report", + aiStage: "research", + }); + + expect(env.anthropicBaseUrl).toBe("https://ai-gateway.dev.posthog.dev"); + expect(env.openaiBaseUrl).toBe("https://ai-gateway.dev.posthog.dev/v1"); + }); + + it("leaves an unlisted product on the Python gateway, slug and all", () => { + process.env.AI_GATEWAY_PRODUCTS = "signals_scout"; + const env = buildServer().configureEnvironment({ isInternal: false }); + + expect(env.anthropicBaseUrl).toBe( + "https://gateway.us.posthog.com/posthog_code", + ); + expect(env.anthropicCustomHeaders).toContain( + "x-posthog-property-task_internal", + ); + }); + + it.each([ + ["scout", "signals_scout"], + ["research", "signals_research"], + ["implementation", "signals_implementation"], + ["repo_selection", "signals_repo_selection"], + ])("sends ai_product %s as %s", (aiStage, expected) => { + const env = buildServer().configureEnvironment({ + originProduct: "signal_report", + aiStage, + }); + + expect(parseBlob(env.anthropicCustomHeaders ?? "").ai_product).toBe( + expected, + ); + }); + + it("carries stage and team attribution in the blob for both adapters", () => { + const env = buildServer().configureEnvironment({ + originProduct: "signal_report", + aiStage: "scout", + taskId: "task-1", + taskRunId: "run-1", + }); + + const expected = { + task_origin_product: "signal_report", + task_internal: false, + ai_stage: "scout", + task_id: "task-1", + task_run_id: "run-1", + ai_product: "signals_scout", + team_id: 42, + }; + expect(parseBlob(env.anthropicCustomHeaders ?? "")).toMatchObject(expected); + expect( + JSON.parse(env.openaiCustomHeaders?.["X-PostHog-Properties"] ?? "{}"), + ).toMatchObject(expected); + }); + + it("sends no per-property headers, which the Go gateway ignores", () => { + const env = buildServer().configureEnvironment({ + originProduct: "signal_report", + aiStage: "scout", + }); + + expect(env.anthropicCustomHeaders).not.toContain("x-posthog-property-"); + expect(Object.keys(env.openaiCustomHeaders ?? {})).toEqual([ + "X-PostHog-Properties", + ]); + }); + + it("keeps non-signals products on their existing ai_product name", () => { + const env = buildServer().configureEnvironment({ isInternal: true }); + + expect(parseBlob(env.anthropicCustomHeaders ?? "").ai_product).toBe( + "background_agents", + ); + }); +}); diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 684ef146b7..54db43ff32 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -88,10 +88,12 @@ import type { import { resourceLink } from "../utils/acp-content"; import { AsyncMutex } from "../utils/async-mutex"; import { + buildGatewayPropertiesHeader, + buildGatewayPropertiesHeaderRecord, buildGatewayPropertyHeaderRecord, buildGatewayPropertyHeaders, resolveGatewayProduct, - resolveLlmGatewayUrl, + resolveGatewayTarget, } from "../utils/gateway"; import { Logger } from "../utils/logger"; import { logAgentshRuntimeInfo } from "./agentsh-runtime"; @@ -3661,11 +3663,11 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} } = {}): GatewayEnv { const { apiKey, apiUrl, projectId } = this.config; const product = resolveGatewayProduct({ isInternal, originProduct }); - const gatewayUrl = resolveLlmGatewayUrl( - process.env.LLM_GATEWAY_URL, - apiUrl, - product, - ); + const { + baseUrl: gatewayUrl, + slugless, + aiProduct, + } = resolveGatewayTarget({ product, aiStage, posthogHost: apiUrl }); const openaiBaseUrl = gatewayUrl.endsWith("/v1") ? gatewayUrl : `${gatewayUrl}/v1`; @@ -3684,14 +3686,31 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} task_user_id: taskUserId, task_title: taskTitle, }; - const customHeaders = buildGatewayPropertyHeaders(gatewayProperties); // The Claude path appends `team_id` in buildEnvironment from // POSTHOG_PROJECT_ID; the codex path has no such hook, so fold it into the // record here to keep team attribution working for both adapters. - const openaiCustomHeaders = buildGatewayPropertyHeaderRecord({ - ...gatewayProperties, - team_id: projectId, - }); + let customHeaders: string; + let openaiCustomHeaders: Record; + if (slugless) { + // The Go gateway reads one X-PostHog-Properties JSON blob and ignores + // per-property headers, and it has no product route, so `ai_product` + // has to travel in the blob or the spend lands unattributed. `team_id` + // is included for both adapters since the Claude hook sets it as a + // per-property header the Go gateway does not read. + const properties = { + ...gatewayProperties, + ai_product: aiProduct, + team_id: projectId, + }; + customHeaders = buildGatewayPropertiesHeader(properties); + openaiCustomHeaders = buildGatewayPropertiesHeaderRecord(properties); + } else { + customHeaders = buildGatewayPropertyHeaders(gatewayProperties); + openaiCustomHeaders = buildGatewayPropertyHeaderRecord({ + ...gatewayProperties, + team_id: projectId, + }); + } // Server-level constants that don't vary per task — safe to keep in // process.env so spawned tools (PostHog MCP, workspace-server, etc.) can diff --git a/packages/agent/src/utils/gateway.test.ts b/packages/agent/src/utils/gateway.test.ts index 9a2bc0a2aa..f2b184e9c5 100644 --- a/packages/agent/src/utils/gateway.test.ts +++ b/packages/agent/src/utils/gateway.test.ts @@ -1,7 +1,9 @@ import { describe, expect, it } from "vitest"; import { getLlmGatewayUrl, + resolveAiProduct, resolveGatewayProduct, + resolveGatewayTarget, resolveLlmGatewayUrl, } from "./gateway"; @@ -135,3 +137,189 @@ describe("getLlmGatewayUrl", () => { ); }); }); + +describe("resolveLlmGatewayUrl (slugless)", () => { + it("does not append the product slug", () => { + expect( + resolveLlmGatewayUrl( + "https://gateway.us.posthog.com", + "https://us.posthog.com", + "signals", + { slugless: true }, + ), + ).toBe("https://gateway.us.posthog.com"); + }); + + it("strips a trailing /v1 so callers can re-append their own path", () => { + expect( + resolveLlmGatewayUrl( + "https://gateway.us.posthog.com/v1", + "https://us.posthog.com", + "signals", + { slugless: true }, + ), + ).toBe("https://gateway.us.posthog.com"); + }); + + it("strips trailing slashes", () => { + expect( + resolveLlmGatewayUrl( + "https://gateway.us.posthog.com/v1//", + "https://us.posthog.com", + "signals", + { slugless: true }, + ), + ).toBe("https://gateway.us.posthog.com"); + }); + + it("falls back to the region-aware host when no override is set", () => { + expect( + resolveLlmGatewayUrl(undefined, "https://eu.posthog.com", "signals", { + slugless: true, + }), + ).toBe("https://gateway.eu.posthog.com"); + }); + + it("keeps appending the slug when not slugless", () => { + expect( + resolveLlmGatewayUrl( + "https://gateway.us.posthog.com", + "https://us.posthog.com", + "signals", + ), + ).toBe("https://gateway.us.posthog.com/signals"); + }); +}); + +describe("resolveGatewayTarget", () => { + const GO = "https://ai-gateway.us.posthog.com"; + const PY_HOST = "https://us.posthog.com"; + const SIGNALS_ENV = { + AI_GATEWAY_URL: GO, + AI_GATEWAY_PRODUCTS: "signals_scout,signals_research", + }; + + it("routes a listed product to the Go gateway with no slug", () => { + expect( + resolveGatewayTarget({ + product: "signals", + aiStage: "scout", + posthogHost: PY_HOST, + env: SIGNALS_ENV, + }), + ).toEqual({ baseUrl: GO, slugless: true, aiProduct: "signals_scout" }); + }); + + it("leaves an unlisted signals stage on the Python gateway", () => { + expect( + resolveGatewayTarget({ + product: "signals", + aiStage: "implementation", + posthogHost: PY_HOST, + env: SIGNALS_ENV, + }), + ).toEqual({ + baseUrl: "https://gateway.us.posthog.com/signals", + slugless: false, + aiProduct: "signals_implementation", + }); + }); + + it.each(["posthog_code", "background_agents", "slack_app"] as const)( + "leaves %s on the Python gateway while signals migrates", + (product) => { + const target = resolveGatewayTarget({ + product, + posthogHost: PY_HOST, + env: SIGNALS_ENV, + }); + expect(target.slugless).toBe(false); + expect(target.baseUrl).toBe(`https://gateway.us.posthog.com/${product}`); + }, + ); + + it("stays on the Python gateway when no Go URL is set", () => { + expect( + resolveGatewayTarget({ + product: "signals", + aiStage: "scout", + posthogHost: PY_HOST, + env: { AI_GATEWAY_PRODUCTS: "signals_scout" }, + }).slugless, + ).toBe(false); + }); + + it("stays on the Python gateway when the allowlist is empty", () => { + expect( + resolveGatewayTarget({ + product: "signals", + aiStage: "scout", + posthogHost: PY_HOST, + env: { AI_GATEWAY_URL: GO }, + }).slugless, + ).toBe(false); + }); + + it("tolerates whitespace and blanks in the allowlist", () => { + expect( + resolveGatewayTarget({ + product: "signals", + aiStage: "scout", + posthogHost: PY_HOST, + env: { AI_GATEWAY_URL: GO, AI_GATEWAY_PRODUCTS: " , signals_scout , " }, + }).slugless, + ).toBe(true); + }); + + it("honours an LLM_GATEWAY_URL override on the unrouted path", () => { + expect( + resolveGatewayTarget({ + product: "posthog_code", + posthogHost: PY_HOST, + env: { + ...SIGNALS_ENV, + LLM_GATEWAY_URL: "https://gateway.dev.posthog.dev", + }, + }).baseUrl, + ).toBe("https://gateway.dev.posthog.dev/posthog_code"); + }); +}); + +describe("resolveAiProduct", () => { + it.each([ + ["scout", "signals_scout"], + ["research", "signals_research"], + ["implementation", "signals_implementation"], + ["repo_selection", "signals_repo_selection"], + ])("maps the signals %s stage to %s", (aiStage, expected) => { + expect(resolveAiProduct({ product: "signals", aiStage })).toBe(expected); + }); + + it("leaves signals unsplit for an unrecognized stage", () => { + expect(resolveAiProduct({ product: "signals", aiStage: "match" })).toBe( + "signals", + ); + }); + + it("leaves signals unsplit when no stage is set", () => { + expect(resolveAiProduct({ product: "signals", aiStage: null })).toBe( + "signals", + ); + }); + + it("does not split a non-signals product that shares a stage name", () => { + expect( + resolveAiProduct({ product: "posthog_code", aiStage: "implementation" }), + ).toBe("posthog_code"); + }); + + it.each([ + "posthog_code", + "background_agents", + "slack_app", + "posthog_ai", + "conversations", + ] as const)("keeps %s unchanged", (product) => { + expect(resolveAiProduct({ product })).toBe(product); + }); +}); diff --git a/packages/agent/src/utils/gateway.ts b/packages/agent/src/utils/gateway.ts index b258086070..b55a4bf3b2 100644 --- a/packages/agent/src/utils/gateway.ts +++ b/packages/agent/src/utils/gateway.ts @@ -32,6 +32,8 @@ export function resolveGatewayProduct({ } export { + buildPosthogPropertiesHeaderLines as buildGatewayPropertiesHeader, + buildPosthogPropertiesHeaderRecord as buildGatewayPropertiesHeaderRecord, buildPosthogPropertyHeaderLines as buildGatewayPropertyHeaders, buildPosthogPropertyHeaderRecord as buildGatewayPropertyHeaderRecord, } from "@posthog/shared/posthog-property-headers"; @@ -79,16 +81,123 @@ export function resolveLlmGatewayUrl( envUrl: string | undefined, posthogHost: string, product: GatewayProduct = "posthog_code", + { slugless = false }: { slugless?: boolean } = {}, ): string { + if (slugless) { + // The Go gateway routes by credential and header, not path, so the product + // slug would 404. A trailing `/v1` is stripped because callers re-append + // it: the Anthropic SDK adds `/v1/messages`, and the OpenAI base URL is + // derived by appending `/v1`. + const base = (envUrl ?? getGatewayBaseUrl(posthogHost)).replace(/\/+$/, ""); + return base.replace(/\/v1$/, ""); + } if (envUrl) { return `${envUrl.replace(/\/$/, "")}/${product}`; } return getLlmGatewayUrl(posthogHost, product); } +/** Products routed to the Go gateway, from a comma-separated env value. */ +function parseAiGatewayProducts(raw: string | undefined): Set { + return new Set( + (raw ?? "") + .split(",") + .map((entry) => entry.trim()) + .filter(Boolean), + ); +} + +/** + * Sandbox-run signals stages, which are billed per stage rather than under one + * `signals` product. The stage names come from `TaskRun.state.ai_stage`, set in + * `Task._build_task` in posthog/posthog. + */ +const SIGNALS_STAGE_PRODUCTS = new Set([ + "scout", + "research", + "implementation", + "repo_selection", +]); + +/** + * The `ai_product` value sent to the Go gateway, which has no product route and + * so cannot infer it from the path. + * + * Signals splits per stage (`signals_scout`), matching the `signals_grouping` / + * `signals_emission` tags its non-sandbox stages already emit. Every other + * product keeps its existing name so its dashboards stay continuous across the + * cutover. + */ +export function resolveAiProduct({ + product, + aiStage, +}: { + product: GatewayProduct; + aiStage?: string | null; +}): string { + if (product === "signals" && aiStage && SIGNALS_STAGE_PRODUCTS.has(aiStage)) { + return `signals_${aiStage}`; + } + return product; +} + export function getGatewayUsageUrl( posthogHost: string, product: GatewayProduct = "posthog_code", ): string { return `${getGatewayBaseUrl(posthogHost)}/v1/usage/${product}`; } + +export interface GatewayTarget { + /** Base URL: slugless for the Go gateway, product-slugged for the Python one. */ + baseUrl: string; + /** True when the Go gateway is selected, which changes the header format. */ + slugless: boolean; + /** Product tag; only the Go gateway needs it sent explicitly. */ + aiProduct: string; +} + +/** + * Choose which gateway a request targets. + * + * The two gateways run on separate hosts (`ai-gateway.*` vs `gateway.*`), so + * the base URL is what distinguishes them. `AI_GATEWAY_URL` selects the Go + * gateway, and `AI_GATEWAY_PRODUCTS` limits that to named products so one + * product can migrate without moving every other caller sharing this server. + * + * Both must be set: an unlisted product, or a listed one with no URL, stays on + * the Python gateway. Migration is therefore opt-in per product, and clearing + * either value rolls everything back. + */ +export function resolveGatewayTarget({ + product, + aiStage, + posthogHost, + env = process.env, +}: { + product: GatewayProduct; + aiStage?: string | null; + posthogHost: string; + env?: Record; +}): GatewayTarget { + const aiProduct = resolveAiProduct({ product, aiStage }); + const aiGatewayUrl = (env.AI_GATEWAY_URL ?? "").trim(); + const routed = + aiGatewayUrl !== "" && + parseAiGatewayProducts(env.AI_GATEWAY_PRODUCTS).has(aiProduct); + + if (routed) { + return { + baseUrl: resolveLlmGatewayUrl(aiGatewayUrl, posthogHost, product, { + slugless: true, + }), + slugless: true, + aiProduct, + }; + } + return { + baseUrl: resolveLlmGatewayUrl(env.LLM_GATEWAY_URL, posthogHost, product), + slugless: false, + aiProduct, + }; +} diff --git a/packages/shared/src/posthog-property-headers.test.ts b/packages/shared/src/posthog-property-headers.test.ts index 61fec187e2..8f72c26edc 100644 --- a/packages/shared/src/posthog-property-headers.test.ts +++ b/packages/shared/src/posthog-property-headers.test.ts @@ -1,5 +1,8 @@ import { describe, expect, it } from "vitest"; import { + buildPosthogPropertiesBlob, + buildPosthogPropertiesHeaderLines, + buildPosthogPropertiesHeaderRecord, buildPosthogPropertyHeaderLines, buildPosthogPropertyHeaderRecord, } from "./posthog-property-headers"; @@ -174,3 +177,131 @@ describe("buildPosthogPropertyHeaderLines", () => { }, ); }); + +describe("buildPosthogPropertiesBlob", () => { + it("serializes properties as one JSON object", () => { + expect( + buildPosthogPropertiesBlob({ + ai_product: "signals_scout", + ai_stage: "scout", + task_internal: true, + }), + ).toBe( + '{"ai_product":"signals_scout","ai_stage":"scout","task_internal":true}', + ); + }); + + it("drops null and undefined values but keeps falsy primitives", () => { + expect( + buildPosthogPropertiesBlob({ + ai_stage: null, + task_internal: false, + task_count: 0, + skipped: undefined, + }), + ).toBe('{"task_internal":false,"task_count":0}'); + }); + + it("drops reserved $-prefixed keys the gateway strips anyway", () => { + expect( + buildPosthogPropertiesBlob({ + $ai_trace_id: "trace-1", + ai_stage: "research", + }), + ).toBe('{"ai_stage":"research"}'); + }); + + it("returns an empty string when no usable properties remain", () => { + expect( + buildPosthogPropertiesBlob({ ai_stage: null, task_id: undefined }), + ).toBe(""); + }); + + it.each([ + { description: "newlines", title: "Fix the bug\nmalicious: true" }, + { description: "carriage returns", title: "Fix the bug\rmalicious: true" }, + ])("collapses $description in values", ({ title }) => { + expect(buildPosthogPropertiesBlob({ task_title: title })).toBe( + '{"task_title":"Fix the bug malicious: true"}', + ); + }); + + it("drops the longest string value until the blob fits the gateway cap", () => { + const blob = buildPosthogPropertiesBlob({ + ai_product: "signals_scout", + ai_stage: "scout", + task_title: "t".repeat(9000), + }); + expect(blob).toBe('{"ai_product":"signals_scout","ai_stage":"scout"}'); + }); + + it("drops only as many values as the cap requires", () => { + const blob = buildPosthogPropertiesBlob({ + ai_product: "signals_research", + task_title: "t".repeat(5000), + task_description: "d".repeat(5000), + }); + // One 5000-char value fits, so only the other is dropped. + expect(blob).toBe( + `{"ai_product":"signals_research","task_description":"${"d".repeat(5000)}"}`, + ); + }); + + it("keeps attribution when every free-text value is oversized", () => { + const blob = buildPosthogPropertiesBlob({ + ai_product: "signals_research", + task_title: "t".repeat(9000), + task_description: "d".repeat(9000), + }); + expect(blob).toBe('{"ai_product":"signals_research"}'); + }); + + it("stays under the cap it enforces", () => { + const blob = buildPosthogPropertiesBlob({ + ai_product: "signals_scout", + task_title: "t".repeat(20000), + }); + expect(new TextEncoder().encode(blob).length).toBeLessThanOrEqual(8192); + }); + + it("sanitizes keys, which are serialized into the header value", () => { + expect( + buildPosthogPropertiesBlob({ "ai_stage\ud83d\ude80": "scout" }), + ).toBe('{"ai_stage":"scout"}'); + }); + + it("emits only printable ASCII a strict HTTP client accepts", () => { + const blob = buildPosthogPropertiesBlob({ + ai_product: "signals_scout", + task_title: "s\u00ec, perch\u00e9 non funziona? \ud83d\ude80", + }); + expect(blob).toMatch(/^[\x20-\x7e]*$/); + expect(blob).toBe( + '{"ai_product":"signals_scout","task_title":"si, perche non funziona? "}', + ); + }); +}); + +describe("buildPosthogPropertiesHeaderRecord", () => { + it("wraps the blob in the X-PostHog-Properties header", () => { + expect( + buildPosthogPropertiesHeaderRecord({ ai_product: "signals_scout" }), + ).toEqual({ "X-PostHog-Properties": '{"ai_product":"signals_scout"}' }); + }); + + it("returns an empty record when there is nothing to send", () => { + expect(buildPosthogPropertiesHeaderRecord({ ai_stage: null })).toEqual({}); + }); +}); + +describe("buildPosthogPropertiesHeaderLines", () => { + it("emits a single header line", () => { + expect( + buildPosthogPropertiesHeaderLines({ ai_product: "signals_scout" }), + ).toBe('X-PostHog-Properties: {"ai_product":"signals_scout"}'); + }); + + it("returns an empty string when there is nothing to send", () => { + expect(buildPosthogPropertiesHeaderLines({ ai_stage: null })).toBe(""); + }); +}); diff --git a/packages/shared/src/posthog-property-headers.ts b/packages/shared/src/posthog-property-headers.ts index bf644194de..475852cf3b 100644 --- a/packages/shared/src/posthog-property-headers.ts +++ b/packages/shared/src/posthog-property-headers.ts @@ -58,3 +58,87 @@ export function buildPosthogPropertyHeaderLines( .map(([key, value]) => `${key}: ${value}`) .join("\n"); } + +/** Header carrying the whole property set as one JSON object. */ +export const POSTHOG_PROPERTIES_HEADER = "X-PostHog-Properties"; + +/** + * Byte cap the Go gateway enforces on {@link POSTHOG_PROPERTIES_HEADER}; a + * larger blob is rejected outright, losing every property including + * `ai_product`. Keep in sync with `maxPropertiesLen` in the gateway's + * `internal/httpapi/dispatch.go`. + */ +const MAX_PROPERTIES_BYTES = 8192; + +const encoder = new TextEncoder(); + +function byteLength(value: string): number { + return encoder.encode(value).length; +} + +/** + * Serialize properties for {@link POSTHOG_PROPERTIES_HEADER}, or `""` when + * there is nothing to send. + * + * Keys beginning with `$` are dropped: the gateway strips reserved `$ai_*` + * keys at the header boundary, so sending them silently loses them. Values + * are sanitized like the per-property headers so a stray control character + * cannot break the header block. + * + * Callers supply free-text values (task titles), so the blob can exceed + * {@link MAX_PROPERTIES_BYTES}. Rather than let the gateway reject the whole + * header, the longest string values are dropped one at a time until it fits, + * preferring to lose descriptive text over attribution keys, which are short. + */ +export function buildPosthogPropertiesBlob( + properties: PosthogProperties, +): string { + const clean: PosthogProperties = {}; + for (const [key, value] of Object.entries(properties)) { + if (value === null || value === undefined) continue; + if (key.startsWith("$")) continue; + // Keys are sanitized too: unlike the per-property headers, where the key + // becomes the header name, here it is serialized into the header value, + // so a non-ASCII key would make the whole request unsendable. + const safeKey = sanitizeHeaderValue(key); + if (!safeKey) continue; + clean[safeKey] = + typeof value === "string" ? sanitizeHeaderValue(value) : value; + } + if (Object.keys(clean).length === 0) return ""; + + let blob = JSON.stringify(clean); + while (byteLength(blob) > MAX_PROPERTIES_BYTES) { + const longest = Object.entries(clean) + .filter(([, value]) => typeof value === "string") + .sort((a, b) => (b[1] as string).length - (a[1] as string).length)[0]; + // Only non-string values remain and they still overflow: send nothing + // rather than a blob the gateway will reject. + if (!longest) return ""; + delete clean[longest[0]]; + blob = JSON.stringify(clean); + } + return blob; +} + +/** + * {@link buildPosthogPropertiesBlob} as a `fetch()`-ready header record, empty + * when there is nothing to send. + */ +export function buildPosthogPropertiesHeaderRecord( + properties: PosthogProperties, +): Record { + const blob = buildPosthogPropertiesBlob(properties); + return blob ? { [POSTHOG_PROPERTIES_HEADER]: blob } : {}; +} + +/** + * {@link buildPosthogPropertiesBlob} as a single `key: value` line for + * `ANTHROPIC_CUSTOM_HEADERS`, empty when there is nothing to send. + */ +export function buildPosthogPropertiesHeaderLines( + properties: PosthogProperties, +): string { + const blob = buildPosthogPropertiesBlob(properties); + return blob ? `${POSTHOG_PROPERTIES_HEADER}: ${blob}` : ""; +} From b255d8d7597638930e8c01546d00226ab184ee43 Mon Sep 17 00:00:00 2001 From: Brandon Leung Date: Tue, 21 Jul 2026 23:17:02 -0400 Subject: [PATCH 2/2] test(agent): cover blob overflow boundary and correct header-layer assertion Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_015gyKxP2MuVuNRCK9uZt1Lr --- .../server/agent-server.configure-environment.test.ts | 8 +++++++- packages/shared/src/posthog-property-headers.test.ts | 10 ++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) 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 66a67db0a4..c894353f7b 100644 --- a/packages/agent/src/server/agent-server.configure-environment.test.ts +++ b/packages/agent/src/server/agent-server.configure-environment.test.ts @@ -455,13 +455,19 @@ describe("AgentServer.configureEnvironment on the Go ai-gateway", () => { ).toMatchObject(expected); }); - it("sends no per-property headers, which the Go gateway ignores", () => { + it("emits attribution as one X-PostHog-Properties blob, not per-property headers", () => { + // Asserts the gateway env this function produces. The Claude adapter's + // buildEnvironment later appends `x-posthog-property-team_id` and + // `x-posthog-use-bedrock-fallback` as separate header lines; the Go gateway + // reads only the blob (team_id is already in it, and it does Bedrock + // failover itself), so those extra lines are inert on this path. const env = buildServer().configureEnvironment({ originProduct: "signal_report", aiStage: "scout", }); expect(env.anthropicCustomHeaders).not.toContain("x-posthog-property-"); + expect(env.anthropicCustomHeaders?.split("\n")).toHaveLength(1); expect(Object.keys(env.openaiCustomHeaders ?? {})).toEqual([ "X-PostHog-Properties", ]); diff --git a/packages/shared/src/posthog-property-headers.test.ts b/packages/shared/src/posthog-property-headers.test.ts index 8f72c26edc..099151a553 100644 --- a/packages/shared/src/posthog-property-headers.test.ts +++ b/packages/shared/src/posthog-property-headers.test.ts @@ -264,6 +264,16 @@ describe("buildPosthogPropertiesBlob", () => { expect(new TextEncoder().encode(blob).length).toBeLessThanOrEqual(8192); }); + it("returns empty when only non-string values remain and still overflow", () => { + // No string value to drop, so trimming can't shrink the blob: send nothing + // rather than a blob the gateway rejects wholesale. + const props: Record = {}; + for (let i = 0; i < 600; i++) { + props[`numeric_property_number_${i}`] = i; + } + expect(buildPosthogPropertiesBlob(props)).toBe(""); + }); + it("sanitizes keys, which are serialized into the header value", () => { expect( buildPosthogPropertiesBlob({ "ai_stage\ud83d\ude80": "scout" }),