Skip to content

Commit d394d5e

Browse files
VojtechBartosVojtechBartos
authored
feat(llm-gateway): tag helper calls with $ai_span_name for analytics (#2981)
1 parent 4ee525d commit d394d5e

10 files changed

Lines changed: 260 additions & 114 deletions

File tree

‎packages/agent/src/utils/gateway.test.ts‎

Lines changed: 0 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import { describe, expect, it } from "vitest";
22
import {
3-
buildGatewayPropertyHeaders,
43
getLlmGatewayUrl,
54
resolveGatewayProduct,
65
resolveLlmGatewayUrl,
@@ -60,83 +59,6 @@ describe("resolveGatewayProduct", () => {
6059
);
6160
});
6261

63-
describe("buildGatewayPropertyHeaders", () => {
64-
it("renders each property as an x-posthog-property header line", () => {
65-
expect(
66-
buildGatewayPropertyHeaders({
67-
task_origin_product: "signal_report",
68-
task_internal: true,
69-
}),
70-
).toBe(
71-
"x-posthog-property-task_origin_product: signal_report\nx-posthog-property-task_internal: true",
72-
);
73-
});
74-
75-
it("drops null and undefined values but keeps falsy primitives", () => {
76-
expect(
77-
buildGatewayPropertyHeaders({
78-
task_origin_product: null,
79-
task_internal: false,
80-
task_count: 0,
81-
}),
82-
).toBe(
83-
"x-posthog-property-task_internal: false\nx-posthog-property-task_count: 0",
84-
);
85-
});
86-
87-
it("returns an empty string when no usable properties remain", () => {
88-
expect(
89-
buildGatewayPropertyHeaders({
90-
task_origin_product: null,
91-
task_internal: undefined,
92-
}),
93-
).toBe("");
94-
});
95-
96-
it.each([
97-
{
98-
description: "LF",
99-
title: "Fix the bug\nx-posthog-property-task_internal: true",
100-
},
101-
{
102-
description: "CRLF",
103-
title: "Fix the bug\r\nx-posthog-property-task_internal: true",
104-
},
105-
{
106-
description: "CR",
107-
title: "Fix the bug\rx-posthog-property-task_internal: true",
108-
},
109-
{
110-
description: "consecutive newlines",
111-
title: "Fix the bug\n\nx-posthog-property-task_internal: true",
112-
},
113-
])(
114-
"collapses $description in values so they cannot inject extra headers",
115-
({ title }) => {
116-
expect(
117-
buildGatewayPropertyHeaders({
118-
task_title: title,
119-
task_id: "task-abc",
120-
}),
121-
).toBe(
122-
"x-posthog-property-task_title: Fix the bug x-posthog-property-task_internal: true\nx-posthog-property-task_id: task-abc",
123-
);
124-
},
125-
);
126-
127-
it("strips characters an HTTP header value cannot carry", () => {
128-
expect(buildGatewayPropertyHeaders({ task_title: "don’t🚀ship" })).toBe(
129-
"x-posthog-property-task_title: dontship",
130-
);
131-
});
132-
133-
it("keeps latin1 characters such as accents", () => {
134-
expect(buildGatewayPropertyHeaders({ task_title: "café" })).toBe(
135-
"x-posthog-property-task_title: café",
136-
);
137-
});
138-
});
139-
14062
describe("resolveLlmGatewayUrl", () => {
14163
it("appends the product slug to an env-provided base URL", () => {
14264
expect(

‎packages/agent/src/utils/gateway.ts‎

Lines changed: 1 addition & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -27,37 +27,7 @@ export function resolveGatewayProduct({
2727
return "posthog_code";
2828
}
2929

30-
/**
31-
* Make a value safe to embed in an HTTP header value. Collapses newlines to
32-
* spaces (the header block is newline-delimited) and drops characters outside
33-
* the valid header-byte range — control chars and code points above latin1
34-
* (emoji, smart quotes) — which an HTTP client (e.g. undici) would otherwise
35-
* reject before sending. ASCII is preserved.
36-
*/
37-
function sanitizeHeaderValue(value: string): string {
38-
return value.replace(/[\r\n]+/g, " ").replace(/[^\x20-\x7e\x80-\xff]/g, "");
39-
}
40-
41-
/**
42-
* Build `x-posthog-property-<name>: <value>` header lines that the LLM
43-
* gateway lifts onto the `$ai_generation` event it captures for each call
44-
* (see `services/llm-gateway/src/llm_gateway/request_context.py`).
45-
*
46-
* Returns a newline-joined string ready for `ANTHROPIC_CUSTOM_HEADERS`.
47-
* `null`/`undefined` values are dropped; values are sanitized to be HTTP-header
48-
* safe (see {@link sanitizeHeaderValue}).
49-
*/
50-
export function buildGatewayPropertyHeaders(
51-
properties: Record<string, string | number | boolean | null | undefined>,
52-
): string {
53-
return Object.entries(properties)
54-
.filter(([, value]) => value !== null && value !== undefined)
55-
.map(
56-
([key, value]) =>
57-
`x-posthog-property-${key}: ${sanitizeHeaderValue(String(value))}`,
58-
)
59-
.join("\n");
60-
}
30+
export { buildPosthogPropertyHeaderLines as buildGatewayPropertyHeaders } from "@posthog/shared/posthog-property-headers";
6131

6232
function getGatewayBaseUrl(posthogHost: string): string {
6333
const url = new URL(posthogHost);

‎packages/core/src/git-pr/git-pr.test.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@ describe("GitPrService.generateCommitMessage", () => {
6767
expect(messages[0].content).toContain("modified: x.ts");
6868
expect(messages[0].content).toContain("why context");
6969
expect(options.system).toContain("commit message generator");
70+
expect(options.posthogProperties).toEqual({
71+
$ai_span_name: "commit_message",
72+
});
7073
});
7174
});
7275

@@ -98,6 +101,10 @@ describe("GitPrService.generatePrTitleAndBody", () => {
98101
expect(result.title).toBe("feat: add widget");
99102
expect(result.body).toBe("TL;DR: adds a widget.");
100103
expect(diffSource.fetchFromRemote).toHaveBeenCalledWith("/repo");
104+
const [, options] = (llm.prompt as ReturnType<typeof vi.fn>).mock.calls[0];
105+
expect(options.posthogProperties).toEqual({
106+
$ai_span_name: "pr_description",
107+
});
101108
});
102109
});
103110

‎packages/core/src/git-pr/git-pr.ts‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,11 @@ ${truncatedDiff}${contextSection}`;
100100

101101
const response = await this.llm.prompt(
102102
[{ role: "user", content: userMessage }],
103-
{ system, model: HELPER_GATEWAY_MODEL },
103+
{
104+
system,
105+
model: HELPER_GATEWAY_MODEL,
106+
posthogProperties: { $ai_span_name: "commit_message" },
107+
},
104108
);
105109

106110
return { message: response.content.trim() };
@@ -211,7 +215,12 @@ ${truncatedDiff || "(no diff available)"}${contextSection}`;
211215

212216
const response = await this.llm.prompt(
213217
[{ role: "user", content: userMessage }],
214-
{ system, maxTokens: 2000, model: HELPER_GATEWAY_MODEL },
218+
{
219+
system,
220+
maxTokens: 2000,
221+
model: HELPER_GATEWAY_MODEL,
222+
posthogProperties: { $ai_span_name: "pr_description" },
223+
},
215224
);
216225

217226
const content = response.content.trim();

‎packages/core/src/llm-gateway/llm-gateway.test.ts‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,38 @@ describe("LlmGatewayService.prompt", () => {
9999
expect(body.stream).toBe(false);
100100
});
101101

102+
it("forwards posthogProperties as x-posthog-property-* request headers and skips nulls", async () => {
103+
const fetchMock = vi
104+
.fn()
105+
.mockResolvedValue(createJsonResponse(SUCCESS_BODY));
106+
const { service } = createService(fetchMock);
107+
108+
await service.prompt([{ role: "user", content: "hi" }], {
109+
posthogProperties: {
110+
$ai_span_name: "pr_description",
111+
task_id: 42,
112+
is_dry_run: false,
113+
// Null/undefined values are dropped so the gateway doesn't see
114+
// literal "null" strings on the captured event.
115+
unused: null,
116+
skipped: undefined,
117+
// Newlines and non-latin1 bytes are sanitized so an undici-backed
118+
// fetch doesn't reject the request before it's sent.
119+
rich: "line one\nline two — done 🎉",
120+
},
121+
});
122+
123+
const [, init] = fetchMock.mock.calls[0];
124+
expect(init.headers).toMatchObject({
125+
"x-posthog-property-$ai_span_name": "pr_description",
126+
"x-posthog-property-task_id": "42",
127+
"x-posthog-property-is_dry_run": "false",
128+
"x-posthog-property-rich": "line one line two done ",
129+
});
130+
expect(init.headers).not.toHaveProperty("x-posthog-property-unused");
131+
expect(init.headers).not.toHaveProperty("x-posthog-property-skipped");
132+
});
133+
102134
it("throws a typed LlmGatewayError with parsed error fields on non-ok response", async () => {
103135
const fetchMock = vi.fn().mockResolvedValue(
104136
createJsonResponse(

‎packages/core/src/llm-gateway/llm-gateway.ts‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
11
import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger";
2+
import {
3+
buildPosthogPropertyHeaderRecord,
4+
type PosthogProperties,
5+
} from "@posthog/shared/posthog-property-headers";
26
import { inject, injectable } from "inversify";
37
import {
48
LLM_GATEWAY_HOST,
@@ -58,6 +62,13 @@ export class LlmGatewayService {
5862
model?: string;
5963
signal?: AbortSignal;
6064
timeoutMs?: number;
65+
/**
66+
* Free-form metadata forwarded as `x-posthog-property-<key>` headers.
67+
* The gateway lifts each one onto the `$ai_generation` event it
68+
* captures, so helper callers (commit messages, PR descriptions, etc.)
69+
* can be told apart from the agent's main generations.
70+
*/
71+
posthogProperties?: PosthogProperties;
6172
} = {},
6273
): Promise<PromptOutput> {
6374
const {
@@ -66,6 +77,7 @@ export class LlmGatewayService {
6677
model = this.endpoints.defaultModel,
6778
signal,
6879
timeoutMs = 60_000,
80+
posthogProperties,
6981
} = options;
7082

7183
const auth = await this.auth.getValidAccessToken();
@@ -101,13 +113,18 @@ export class LlmGatewayService {
101113
else signal.addEventListener("abort", onCallerAbort, { once: true });
102114
}
103115

116+
const headers: Record<string, string> = {
117+
"Content-Type": "application/json",
118+
...(posthogProperties
119+
? buildPosthogPropertyHeaderRecord(posthogProperties)
120+
: {}),
121+
};
122+
104123
let response: Response;
105124
try {
106125
response = await this.auth.authenticatedFetch(messagesUrl, {
107126
method: "POST",
108-
headers: {
109-
"Content-Type": "application/json",
110-
},
127+
headers,
111128
body: JSON.stringify(requestBody),
112129
signal: timeoutController.signal,
113130
});

‎packages/shared/package.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@
3939
"./mcp-sandbox-proxy": {
4040
"types": "./dist/mcp-sandbox-proxy.d.ts",
4141
"import": "./dist/mcp-sandbox-proxy.js"
42+
},
43+
"./posthog-property-headers": {
44+
"types": "./dist/posthog-property-headers.d.ts",
45+
"import": "./dist/posthog-property-headers.js"
4246
}
4347
},
4448
"scripts": {

0 commit comments

Comments
 (0)