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 e7d903b86..23e3ed33d 100644 --- a/packages/agent/src/server/agent-server.configure-environment.test.ts +++ b/packages/agent/src/server/agent-server.configure-environment.test.ts @@ -2,13 +2,20 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentServer } from "./agent-server"; interface TestableServer { - configureEnvironment(args?: { isInternal?: boolean }): void; + configureEnvironment(args?: { + isInternal?: boolean; + originProduct?: string | null; + taskId?: string | null; + taskRunId?: string | null; + taskUserId?: number | null; + }): void; } const ENV_KEYS_UNDER_TEST = [ "LLM_GATEWAY_URL", "ANTHROPIC_BASE_URL", "OPENAI_BASE_URL", + "ANTHROPIC_CUSTOM_HEADERS", ] as const; describe("AgentServer.configureEnvironment", () => { @@ -85,6 +92,62 @@ describe("AgentServer.configureEnvironment", () => { expect(fromBackground).toBe("https://gateway.us.posthog.com/posthog_code"); }); + it("tags as signals when an internal task has origin_product 'signal_report'", () => { + buildServer("background").configureEnvironment({ + isInternal: true, + originProduct: "signal_report", + }); + + expect(process.env.LLM_GATEWAY_URL).toBe( + "https://gateway.us.posthog.com/signals", + ); + expect(process.env.ANTHROPIC_BASE_URL).toBe( + "https://gateway.us.posthog.com/signals", + ); + expect(process.env.OPENAI_BASE_URL).toBe( + "https://gateway.us.posthog.com/signals/v1", + ); + }); + + it("does not tag as signals when origin_product is 'signal_report' but the task is not internal", () => { + buildServer("background").configureEnvironment({ + isInternal: false, + originProduct: "signal_report", + }); + + expect(process.env.LLM_GATEWAY_URL).toBe( + "https://gateway.us.posthog.com/posthog_code", + ); + }); + + it("forwards task metadata as ANTHROPIC_CUSTOM_HEADERS", () => { + buildServer("background").configureEnvironment({ + isInternal: true, + originProduct: "signal_report", + taskId: "task-abc", + taskRunId: "run-xyz", + taskUserId: 42, + }); + + expect(process.env.ANTHROPIC_CUSTOM_HEADERS).toBe( + [ + "x-posthog-property-task_origin_product: signal_report", + "x-posthog-property-task_internal: true", + "x-posthog-property-task_id: task-abc", + "x-posthog-property-task_run_id: run-xyz", + "x-posthog-property-task_user_id: 42", + ].join("\n"), + ); + }); + + it("omits optional task metadata from ANTHROPIC_CUSTOM_HEADERS when not provided", () => { + buildServer("background").configureEnvironment({ isInternal: false }); + + expect(process.env.ANTHROPIC_CUSTOM_HEADERS).toBe( + "x-posthog-property-task_internal: false", + ); + }); + it("respects the LLM_GATEWAY_URL override regardless of internal flag", () => { process.env.LLM_GATEWAY_URL = "http://ngrok.test/proxy"; diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index af4998a4a..58e7d9e61 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -47,7 +47,11 @@ import type { } from "../types"; import { resourceLink } from "../utils/acp-content"; import { AsyncMutex } from "../utils/async-mutex"; -import { type GatewayProduct, getLlmGatewayUrl } from "../utils/gateway"; +import { + buildGatewayPropertyHeaders, + getLlmGatewayUrl, + resolveGatewayProduct, +} from "../utils/gateway"; import { Logger } from "../utils/logger"; import { logAgentshRuntimeInfo } from "./agentsh-runtime"; import { @@ -844,7 +848,13 @@ export class AgentServer { }), ]); - this.configureEnvironment({ isInternal: preTask?.internal === true }); + this.configureEnvironment({ + isInternal: preTask?.internal === true, + originProduct: preTask?.origin_product, + taskId: payload.task_id, + taskRunId: payload.run_id, + taskUserId: payload.user_id, + }); const prUrl = getTaskRunStateString(preTaskRun, "slack_notified_pr_url"); @@ -1804,18 +1814,35 @@ ${attributionInstructions} private configureEnvironment({ isInternal = false, + originProduct, + taskId, + taskRunId, + taskUserId, }: { isInternal?: boolean; + originProduct?: string | null; + taskId?: string | null; + taskRunId?: string | null; + taskUserId?: number | null; } = {}): void { const { apiKey, apiUrl, projectId } = this.config; - const product: GatewayProduct = isInternal - ? "background_agents" - : "posthog_code"; + const product = resolveGatewayProduct({ isInternal, originProduct }); const gatewayUrl = process.env.LLM_GATEWAY_URL || getLlmGatewayUrl(apiUrl, product); const openaiBaseUrl = gatewayUrl.endsWith("/v1") ? gatewayUrl : `${gatewayUrl}/v1`; + // 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. + const customHeaders = buildGatewayPropertyHeaders({ + task_origin_product: originProduct, + task_internal: isInternal, + task_id: taskId, + task_run_id: taskRunId, + task_user_id: taskUserId, + }); Object.assign(process.env, { // PostHog @@ -1828,6 +1855,7 @@ ${attributionInstructions} ANTHROPIC_API_KEY: apiKey, ANTHROPIC_AUTH_TOKEN: apiKey, ANTHROPIC_BASE_URL: gatewayUrl, + ANTHROPIC_CUSTOM_HEADERS: customHeaders, // OpenAI (for models like GPT-4, o1, etc.) OPENAI_API_KEY: apiKey, OPENAI_BASE_URL: openaiBaseUrl, diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 18e5572c0..4a5864f6d 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -37,7 +37,8 @@ export interface Task { | "eval_clusters" | "user_created" | "support_queue" - | "session_summaries"; + | "session_summaries" + | "signal_report"; github_integration?: number | null; repository: string; // Format: "organization/repository" (e.g., "posthog/posthog-js") json_schema?: Record | null; // JSON schema for task output validation diff --git a/packages/agent/src/utils/gateway.test.ts b/packages/agent/src/utils/gateway.test.ts new file mode 100644 index 000000000..5329a712d --- /dev/null +++ b/packages/agent/src/utils/gateway.test.ts @@ -0,0 +1,70 @@ +import { describe, expect, it } from "vitest"; +import { buildGatewayPropertyHeaders, resolveGatewayProduct } from "./gateway"; + +describe("resolveGatewayProduct", () => { + it.each([ + { isInternal: false, originProduct: undefined, expected: "posthog_code" }, + { + isInternal: undefined, + originProduct: undefined, + expected: "posthog_code", + }, + { + isInternal: false, + originProduct: "signal_report", + expected: "posthog_code", + }, + { + isInternal: true, + originProduct: undefined, + expected: "background_agents", + }, + { + isInternal: true, + originProduct: "session_summaries", + expected: "background_agents", + }, + { isInternal: true, originProduct: "signal_report", expected: "signals" }, + ] as const)( + "isInternal=$isInternal originProduct=$originProduct -> $expected", + ({ isInternal, originProduct, expected }) => { + expect(resolveGatewayProduct({ isInternal, originProduct })).toBe( + expected, + ); + }, + ); +}); + +describe("buildGatewayPropertyHeaders", () => { + it("renders each property as an x-posthog-property header line", () => { + expect( + buildGatewayPropertyHeaders({ + task_origin_product: "signal_report", + task_internal: true, + }), + ).toBe( + "x-posthog-property-task_origin_product: signal_report\nx-posthog-property-task_internal: true", + ); + }); + + it("drops null and undefined values but keeps falsy primitives", () => { + expect( + buildGatewayPropertyHeaders({ + task_origin_product: null, + task_internal: false, + task_count: 0, + }), + ).toBe( + "x-posthog-property-task_internal: false\nx-posthog-property-task_count: 0", + ); + }); + + it("returns an empty string when no usable properties remain", () => { + expect( + buildGatewayPropertyHeaders({ + task_origin_product: null, + task_internal: undefined, + }), + ).toBe(""); + }); +}); diff --git a/packages/agent/src/utils/gateway.ts b/packages/agent/src/utils/gateway.ts index 161247538..ff27758b2 100644 --- a/packages/agent/src/utils/gateway.ts +++ b/packages/agent/src/utils/gateway.ts @@ -1,4 +1,34 @@ -export type GatewayProduct = "posthog_code" | "background_agents"; +export type GatewayProduct = "posthog_code" | "background_agents" | "signals"; + +export function resolveGatewayProduct({ + isInternal, + originProduct, +}: { + isInternal?: boolean; + originProduct?: string | null; +} = {}): GatewayProduct { + if (isInternal) { + return originProduct === "signal_report" ? "signals" : "background_agents"; + } + return "posthog_code"; +} + +/** + * Build `x-posthog-property-: ` header lines that the LLM + * gateway lifts onto the `$ai_generation` event it captures for each call + * (see `services/llm-gateway/src/llm_gateway/request_context.py`). + * + * Returns a newline-joined string ready for `ANTHROPIC_CUSTOM_HEADERS`. + * `null`/`undefined` property values are dropped. + */ +export function buildGatewayPropertyHeaders( + properties: Record, +): string { + return Object.entries(properties) + .filter(([, value]) => value !== null && value !== undefined) + .map(([key, value]) => `x-posthog-property-${key}: ${value}`) + .join("\n"); +} function getGatewayBaseUrl(posthogHost: string): string { const url = new URL(posthogHost);