From 1ce21c750908f9d0ed3bdde6dbf0f2d1cef1e899 Mon Sep 17 00:00:00 2001 From: Rafa Audibert Date: Tue, 28 Jul 2026 00:36:19 -0300 Subject: [PATCH 1/4] feat(agent): route wizard cloud runs to the onboarding gateway product Setup wizard cloud runs (origin_product "onboarding") fell through to the posthog_code default, so their generations billed against the customer's PostHog Code credits. Route them to the new unbilled `onboarding` product instead. origin_product alone does not authorize the free route: any task:write holder can POST a task claiming it. A run must also carry the wizard_config state key that only the server-side wizard flow stamps, and that the task API refuses to let a caller set. Without the marker the run stays on posthog_code. Co-Authored-By: Claude Opus 5 (1M context) --- ...agent-server.configure-environment.test.ts | 30 +++++++++++++++++++ packages/agent/src/server/agent-server.ts | 14 ++++++++- packages/agent/src/utils/gateway.test.ts | 20 +++++++++++++ packages/agent/src/utils/gateway.ts | 14 ++++++++- 4 files changed, 76 insertions(+), 2 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 d16673f4aa..9366bfc29d 100644 --- a/packages/agent/src/server/agent-server.configure-environment.test.ts +++ b/packages/agent/src/server/agent-server.configure-environment.test.ts @@ -7,6 +7,7 @@ interface TestableServer { configureEnvironment(args?: { isInternal?: boolean; originProduct?: Task["origin_product"] | null; + isWizardCloudRun?: boolean; signalReportId?: string | null; aiStage?: string | null; taskId?: string | null; @@ -183,6 +184,35 @@ describe("AgentServer.configureEnvironment", () => { }, ); + it("tags as onboarding when a wizard cloud run reaches the gateway", () => { + const env = buildServer("background").configureEnvironment({ + isInternal: false, + originProduct: "onboarding", + isWizardCloudRun: true, + }); + + expect(env.anthropicBaseUrl).toBe( + "https://gateway.us.posthog.com/onboarding", + ); + expect(env.openaiBaseUrl).toBe( + "https://gateway.us.posthog.com/onboarding/v1", + ); + }); + + // onboarding is unbilled, so a task that merely claims the origin (anyone with task:write + // can) must not reach it without the server-stamped wizard_config marker. + it("keeps a marker-less onboarding task on posthog_code", () => { + const env = buildServer("background").configureEnvironment({ + isInternal: false, + originProduct: "onboarding", + isWizardCloudRun: false, + }); + + expect(env.anthropicBaseUrl).toBe( + "https://gateway.us.posthog.com/posthog_code", + ); + }); + // The codex/OpenAI path sets provider http_headers rather than // ANTHROPIC_CUSTOM_HEADERS, so the same task metadata must be exposed as a // record — including team_id, which the Claude path adds separately in diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index bb08232f2a..6f9b1d2cf1 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -1545,6 +1545,12 @@ export class AgentServer { const gatewayEnv = this.configureEnvironment({ isInternal: preTask?.internal === true, originProduct: preTask?.origin_product, + // Only the server-side wizard flow can put `wizard_config` on a run: the run PATCH + // allowlist drops it and the run-create body has no state field at all. That makes it + // the trustworthy half of the onboarding check, unlike origin_product. + isWizardCloudRun: + (preTaskRun?.state as Record | undefined) + ?.wizard_config !== undefined, signalReportId: preTask?.signal_report, aiStage: getTaskRunStateString(preTaskRun, "ai_stage"), taskId: payload.task_id, @@ -3839,6 +3845,7 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} private configureEnvironment({ isInternal = false, originProduct, + isWizardCloudRun = false, signalReportId, aiStage, taskId, @@ -3848,6 +3855,7 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} }: { isInternal?: boolean; originProduct?: Task["origin_product"] | null; + isWizardCloudRun?: boolean; signalReportId?: string | null; aiStage?: string | null; taskId?: string | null; @@ -3856,7 +3864,11 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} taskTitle?: string | null; } = {}): GatewayEnv { const { apiKey, apiUrl, projectId } = this.config; - const product = resolveGatewayProduct({ isInternal, originProduct }); + const product = resolveGatewayProduct({ + isInternal, + originProduct, + isWizardCloudRun, + }); const { baseUrl: gatewayUrl, isAiGateway, diff --git a/packages/agent/src/utils/gateway.test.ts b/packages/agent/src/utils/gateway.test.ts index fcedf4cb81..e6af6f0966 100644 --- a/packages/agent/src/utils/gateway.test.ts +++ b/packages/agent/src/utils/gateway.test.ts @@ -79,6 +79,26 @@ describe("resolveGatewayProduct", () => { ); }, ); + + // `onboarding` is unbilled, and origin_product is caller-supplied: any task:write holder can + // POST a task claiming it. Only the wizard_config marker, which the task API will not let a + // caller set, may unlock the free product. Everything else stays on the billed one. + it.each([ + { isInternal: false, isWizardCloudRun: true, expected: "onboarding" }, + { isInternal: false, isWizardCloudRun: false, expected: "posthog_code" }, + { isInternal: true, isWizardCloudRun: false, expected: "posthog_code" }, + ] as const)( + "originProduct=onboarding isWizardCloudRun=$isWizardCloudRun isInternal=$isInternal -> $expected", + ({ isInternal, isWizardCloudRun, expected }) => { + expect( + resolveGatewayProduct({ + isInternal, + originProduct: "onboarding", + isWizardCloudRun, + }), + ).toBe(expected); + }, + ); }); describe("resolveLlmGatewayUrl", () => { diff --git a/packages/agent/src/utils/gateway.ts b/packages/agent/src/utils/gateway.ts index cf9c53c09c..70272328c5 100644 --- a/packages/agent/src/utils/gateway.ts +++ b/packages/agent/src/utils/gateway.ts @@ -4,14 +4,23 @@ export type GatewayProduct = | "signals" | "slack_app" | "posthog_ai" - | "conversations"; + | "conversations" + | "onboarding"; export function resolveGatewayProduct({ isInternal, originProduct, + isWizardCloudRun, }: { isInternal?: boolean; originProduct?: string | null; + /** + * Whether the run carries the `wizard_config` state key that only the server-side + * wizard flow stamps. Required for `onboarding` because that product is unbilled: + * `origin_product` on its own is caller-supplied data we should not spend PostHog's + * money on, so a run that claims the origin without the marker stays on posthog_code. + */ + isWizardCloudRun?: boolean; } = {}): GatewayProduct { if (originProduct === "slack") { return "slack_app"; @@ -28,6 +37,9 @@ export function resolveGatewayProduct({ if (originProduct === "loop") { return "posthog_code"; } + if (originProduct === "onboarding") { + return isWizardCloudRun ? "onboarding" : "posthog_code"; + } if (isInternal) { return "background_agents"; } From 7dd5df371078f2fc0fd01b64ef8c77de46f5a0a1 Mon Sep 17 00:00:00 2001 From: Rafa Audibert Date: Wed, 29 Jul 2026 00:19:54 -0300 Subject: [PATCH 2/4] fix: Simplify removing check for cloud run This is enough, just gotta confirm it's internal --- .../agent-server.configure-environment.test.ts | 16 ---------------- packages/agent/src/server/agent-server.ts | 9 --------- packages/agent/src/utils/gateway.test.ts | 14 +++++--------- packages/agent/src/utils/gateway.ts | 12 ++---------- 4 files changed, 7 insertions(+), 44 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 9366bfc29d..0ae96970f6 100644 --- a/packages/agent/src/server/agent-server.configure-environment.test.ts +++ b/packages/agent/src/server/agent-server.configure-environment.test.ts @@ -7,7 +7,6 @@ interface TestableServer { configureEnvironment(args?: { isInternal?: boolean; originProduct?: Task["origin_product"] | null; - isWizardCloudRun?: boolean; signalReportId?: string | null; aiStage?: string | null; taskId?: string | null; @@ -188,7 +187,6 @@ describe("AgentServer.configureEnvironment", () => { const env = buildServer("background").configureEnvironment({ isInternal: false, originProduct: "onboarding", - isWizardCloudRun: true, }); expect(env.anthropicBaseUrl).toBe( @@ -199,20 +197,6 @@ describe("AgentServer.configureEnvironment", () => { ); }); - // onboarding is unbilled, so a task that merely claims the origin (anyone with task:write - // can) must not reach it without the server-stamped wizard_config marker. - it("keeps a marker-less onboarding task on posthog_code", () => { - const env = buildServer("background").configureEnvironment({ - isInternal: false, - originProduct: "onboarding", - isWizardCloudRun: false, - }); - - expect(env.anthropicBaseUrl).toBe( - "https://gateway.us.posthog.com/posthog_code", - ); - }); - // The codex/OpenAI path sets provider http_headers rather than // ANTHROPIC_CUSTOM_HEADERS, so the same task metadata must be exposed as a // record — including team_id, which the Claude path adds separately in diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 6f9b1d2cf1..281bf9477b 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -1545,12 +1545,6 @@ export class AgentServer { const gatewayEnv = this.configureEnvironment({ isInternal: preTask?.internal === true, originProduct: preTask?.origin_product, - // Only the server-side wizard flow can put `wizard_config` on a run: the run PATCH - // allowlist drops it and the run-create body has no state field at all. That makes it - // the trustworthy half of the onboarding check, unlike origin_product. - isWizardCloudRun: - (preTaskRun?.state as Record | undefined) - ?.wizard_config !== undefined, signalReportId: preTask?.signal_report, aiStage: getTaskRunStateString(preTaskRun, "ai_stage"), taskId: payload.task_id, @@ -3845,7 +3839,6 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} private configureEnvironment({ isInternal = false, originProduct, - isWizardCloudRun = false, signalReportId, aiStage, taskId, @@ -3855,7 +3848,6 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} }: { isInternal?: boolean; originProduct?: Task["origin_product"] | null; - isWizardCloudRun?: boolean; signalReportId?: string | null; aiStage?: string | null; taskId?: string | null; @@ -3867,7 +3859,6 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} const product = resolveGatewayProduct({ isInternal, originProduct, - isWizardCloudRun, }); const { baseUrl: gatewayUrl, diff --git a/packages/agent/src/utils/gateway.test.ts b/packages/agent/src/utils/gateway.test.ts index e6af6f0966..5ec68fb18e 100644 --- a/packages/agent/src/utils/gateway.test.ts +++ b/packages/agent/src/utils/gateway.test.ts @@ -80,21 +80,17 @@ describe("resolveGatewayProduct", () => { }, ); - // `onboarding` is unbilled, and origin_product is caller-supplied: any task:write holder can - // POST a task claiming it. Only the wizard_config marker, which the task API will not let a - // caller set, may unlock the free product. Everything else stays on the billed one. + // `onboarding` is unbilled but make sure it's only used if the task is internal. it.each([ - { isInternal: false, isWizardCloudRun: true, expected: "onboarding" }, - { isInternal: false, isWizardCloudRun: false, expected: "posthog_code" }, - { isInternal: true, isWizardCloudRun: false, expected: "posthog_code" }, + { isInternal: true, expected: "onboarding" }, + { isInternal: false, expected: "posthog_code" }, ] as const)( - "originProduct=onboarding isWizardCloudRun=$isWizardCloudRun isInternal=$isInternal -> $expected", - ({ isInternal, isWizardCloudRun, expected }) => { + "originProduct=onboarding isInternal=$isInternal -> $expected", + ({ isInternal, expected }) => { expect( resolveGatewayProduct({ isInternal, originProduct: "onboarding", - isWizardCloudRun, }), ).toBe(expected); }, diff --git a/packages/agent/src/utils/gateway.ts b/packages/agent/src/utils/gateway.ts index 70272328c5..68314cfb71 100644 --- a/packages/agent/src/utils/gateway.ts +++ b/packages/agent/src/utils/gateway.ts @@ -10,17 +10,9 @@ export type GatewayProduct = export function resolveGatewayProduct({ isInternal, originProduct, - isWizardCloudRun, }: { isInternal?: boolean; originProduct?: string | null; - /** - * Whether the run carries the `wizard_config` state key that only the server-side - * wizard flow stamps. Required for `onboarding` because that product is unbilled: - * `origin_product` on its own is caller-supplied data we should not spend PostHog's - * money on, so a run that claims the origin without the marker stays on posthog_code. - */ - isWizardCloudRun?: boolean; } = {}): GatewayProduct { if (originProduct === "slack") { return "slack_app"; @@ -37,8 +29,8 @@ export function resolveGatewayProduct({ if (originProduct === "loop") { return "posthog_code"; } - if (originProduct === "onboarding") { - return isWizardCloudRun ? "onboarding" : "posthog_code"; + if (originProduct === "onboarding" && isInternal) { + return "onboarding"; } if (isInternal) { return "background_agents"; From 18594cf398ca92e64e6bac9d8cb7de5a83d41e87 Mon Sep 17 00:00:00 2001 From: Rafa Audibert Date: Wed, 29 Jul 2026 00:34:43 -0300 Subject: [PATCH 3/4] refactor(gateway): streamline origin product resolution with a mapping object Replaced multiple conditional checks with a mapping object to simplify the resolution of gateway products based on the origin product. This enhances code readability and maintainability. --- packages/agent/src/utils/gateway.ts | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/packages/agent/src/utils/gateway.ts b/packages/agent/src/utils/gateway.ts index 68314cfb71..c9ef17df70 100644 --- a/packages/agent/src/utils/gateway.ts +++ b/packages/agent/src/utils/gateway.ts @@ -14,23 +14,18 @@ export function resolveGatewayProduct({ isInternal?: boolean; originProduct?: string | null; } = {}): GatewayProduct { - if (originProduct === "slack") { - return "slack_app"; - } - if (originProduct === "posthog_ai") { - return "posthog_ai"; - } - if (originProduct === "signal_report" || originProduct === "signals_scout") { - return "signals"; - } - if (originProduct === "support_reply") { - return "conversations"; - } - if (originProduct === "loop") { - return "posthog_code"; - } - if (originProduct === "onboarding" && isInternal) { - return "onboarding"; + const originProductToGatewayProductMap: Record = { + loop: "posthog_code", + onboarding: "onboarding", + posthog_ai: "posthog_ai", + signal_report: "signals", + signals_scout: "signals", + slack: "slack_app", + support_reply: "conversations", + }; + + if (originProduct && originProduct in originProductToGatewayProductMap) { + return originProductToGatewayProductMap[originProduct]; } if (isInternal) { return "background_agents"; From 06f44ab5a963fd020e7bcc7136d8eccda3ffc88e Mon Sep 17 00:00:00 2001 From: Rafa Audibert Date: Wed, 29 Jul 2026 00:45:10 -0300 Subject: [PATCH 4/4] fix: Simplify code --- packages/agent/src/server/agent-server.ts | 5 +---- packages/agent/src/utils/gateway.test.ts | 16 ---------------- 2 files changed, 1 insertion(+), 20 deletions(-) diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 281bf9477b..bb08232f2a 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -3856,10 +3856,7 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} taskTitle?: string | null; } = {}): GatewayEnv { const { apiKey, apiUrl, projectId } = this.config; - const product = resolveGatewayProduct({ - isInternal, - originProduct, - }); + const product = resolveGatewayProduct({ isInternal, originProduct }); const { baseUrl: gatewayUrl, isAiGateway, diff --git a/packages/agent/src/utils/gateway.test.ts b/packages/agent/src/utils/gateway.test.ts index 5ec68fb18e..fcedf4cb81 100644 --- a/packages/agent/src/utils/gateway.test.ts +++ b/packages/agent/src/utils/gateway.test.ts @@ -79,22 +79,6 @@ describe("resolveGatewayProduct", () => { ); }, ); - - // `onboarding` is unbilled but make sure it's only used if the task is internal. - it.each([ - { isInternal: true, expected: "onboarding" }, - { isInternal: false, expected: "posthog_code" }, - ] as const)( - "originProduct=onboarding isInternal=$isInternal -> $expected", - ({ isInternal, expected }) => { - expect( - resolveGatewayProduct({ - isInternal, - originProduct: "onboarding", - }), - ).toBe(expected); - }, - ); }); describe("resolveLlmGatewayUrl", () => {