diff --git a/packages/agent/src/server/agent-server.configure-environment.test.ts b/packages/agent/src/server/agent-server.configure-environment.test.ts new file mode 100644 index 000000000..e7d903b86 --- /dev/null +++ b/packages/agent/src/server/agent-server.configure-environment.test.ts @@ -0,0 +1,97 @@ +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { AgentServer } from "./agent-server"; + +interface TestableServer { + configureEnvironment(args?: { isInternal?: boolean }): void; +} + +const ENV_KEYS_UNDER_TEST = [ + "LLM_GATEWAY_URL", + "ANTHROPIC_BASE_URL", + "OPENAI_BASE_URL", +] as const; + +describe("AgentServer.configureEnvironment", () => { + const originalEnv: Partial> = {}; + + beforeEach(() => { + for (const key of ENV_KEYS_UNDER_TEST) { + originalEnv[key] = process.env[key]; + delete process.env[key]; + } + }); + + afterEach(() => { + for (const key of ENV_KEYS_UNDER_TEST) { + const value = originalEnv[key]; + if (value === undefined) { + delete process.env[key]; + } else { + process.env[key] = value; + } + } + }); + + const buildServer = (mode: "background" | "interactive"): TestableServer => + new AgentServer({ + port: 0, + jwtPublicKey: "test-key", + apiUrl: "https://us.posthog.com", + apiKey: "test-api-key", + projectId: 1, + mode, + taskId: "test-task-id", + runId: "test-run-id", + }) as unknown as TestableServer; + + it("tags as background_agents when the task is internal", () => { + buildServer("interactive").configureEnvironment({ isInternal: true }); + + expect(process.env.LLM_GATEWAY_URL).toBe( + "https://gateway.us.posthog.com/background_agents", + ); + expect(process.env.ANTHROPIC_BASE_URL).toBe( + "https://gateway.us.posthog.com/background_agents", + ); + expect(process.env.OPENAI_BASE_URL).toBe( + "https://gateway.us.posthog.com/background_agents/v1", + ); + }); + + it("tags as posthog_code when the task is not internal", () => { + buildServer("background").configureEnvironment({ isInternal: false }); + + expect(process.env.LLM_GATEWAY_URL).toBe( + "https://gateway.us.posthog.com/posthog_code", + ); + }); + + it("tags as posthog_code when isInternal is omitted (getTask failure fallback)", () => { + buildServer("background").configureEnvironment(); + + expect(process.env.LLM_GATEWAY_URL).toBe( + "https://gateway.us.posthog.com/posthog_code", + ); + }); + + it("ignores mode when picking the gateway product", () => { + buildServer("background").configureEnvironment({ isInternal: false }); + const fromBackground = process.env.LLM_GATEWAY_URL; + + buildServer("interactive").configureEnvironment({ isInternal: false }); + const fromInteractive = process.env.LLM_GATEWAY_URL; + + expect(fromBackground).toBe(fromInteractive); + expect(fromBackground).toBe("https://gateway.us.posthog.com/posthog_code"); + }); + + it("respects the LLM_GATEWAY_URL override regardless of internal flag", () => { + process.env.LLM_GATEWAY_URL = "http://ngrok.test/proxy"; + + buildServer("background").configureEnvironment({ isInternal: true }); + + expect(process.env.LLM_GATEWAY_URL).toBe("http://ngrok.test/proxy"); + expect(process.env.ANTHROPIC_BASE_URL).toBe("http://ngrok.test/proxy"); + expect(process.env.OPENAI_BASE_URL).toBe("http://ngrok.test/proxy/v1"); + }); +}); diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 9e5425fc9..b614d5c69 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -46,7 +46,7 @@ import type { } from "../types"; import { resourceLink } from "../utils/acp-content"; import { AsyncMutex } from "../utils/async-mutex"; -import { getLlmGatewayUrl } from "../utils/gateway"; +import { type GatewayProduct, getLlmGatewayUrl } from "../utils/gateway"; import { Logger } from "../utils/logger"; import { logAgentshRuntimeInfo } from "./agentsh-runtime"; import { @@ -778,8 +778,6 @@ export class AgentServer { name: process.env.HOSTNAME || "cloud-sandbox", }; - this.configureEnvironment(); - const [preTaskRun, preTask] = await Promise.all([ this.posthogAPI .getTaskRun(payload.task_id, payload.run_id) @@ -800,6 +798,8 @@ export class AgentServer { }), ]); + this.configureEnvironment({ isInternal: preTask?.internal === true }); + const prUrl = getTaskRunStateString(preTaskRun, "slack_notified_pr_url"); if (prUrl) { @@ -1710,10 +1710,15 @@ ${attributionInstructions} } } - private configureEnvironment(): void { + private configureEnvironment({ + isInternal = false, + }: { + isInternal?: boolean; + } = {}): void { const { apiKey, apiUrl, projectId } = this.config; - const product = - this.config.mode === "background" ? "background_agents" : "posthog_code"; + const product: GatewayProduct = isInternal + ? "background_agents" + : "posthog_code"; const gatewayUrl = process.env.LLM_GATEWAY_URL || getLlmGatewayUrl(apiUrl, product); const openaiBaseUrl = gatewayUrl.endsWith("/v1") diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 00a904f56..18e5572c0 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -41,6 +41,7 @@ export interface Task { github_integration?: number | null; repository: string; // Format: "organization/repository" (e.g., "posthog/posthog-js") json_schema?: Record | null; // JSON schema for task output validation + internal?: boolean; created_at: string; updated_at: string; created_by?: {