From 905a0985f95faff8eeb818887ecb3f17f64b7307 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Wed, 20 May 2026 15:23:05 +0800 Subject: [PATCH] Forward create agent env to provider launches (#1112) --- packages/cli/src/commands/agent/run.ts | 47 ++++++++++-- packages/server/src/client/daemon-client.ts | 3 + .../src/server/agent/agent-manager.test.ts | 75 +++++++++++++++++++ .../server/src/server/agent/agent-manager.ts | 6 +- .../agent/providers/opencode-agent.test.ts | 34 +++++++++ .../server/agent/providers/opencode-agent.ts | 10 ++- .../agent/providers/opencode/runtime.ts | 5 +- .../providers/opencode/server-manager.ts | 35 ++++++++- .../providers/opencode/test-server-manager.ts | 7 +- .../test-utils/test-opencode-runtime.ts | 13 +++- .../server/agent/providers/pi/agent.test.ts | 16 ++++ .../src/server/agent/providers/pi/agent.ts | 3 +- .../src/server/agent/providers/pi/runtime.ts | 9 ++- packages/server/src/server/messages.test.ts | 41 +++++++++- packages/server/src/server/session.ts | 2 + packages/server/src/shared/messages.ts | 1 + 16 files changed, 286 insertions(+), 21 deletions(-) diff --git a/packages/cli/src/commands/agent/run.ts b/packages/cli/src/commands/agent/run.ts index 782eb516f..b934013a0 100644 --- a/packages/cli/src/commands/agent/run.ts +++ b/packages/cli/src/commands/agent/run.ts @@ -46,6 +46,12 @@ export function addRunOptions(cmd: Command): Command { [], ) .option("--cwd ", "Working directory (default: current)") + .option( + "--env ", + "Set environment variable(s) for the agent process (can be used multiple times)", + collectMultiple, + [], + ) .option( "--label ", "Add label(s) to the agent (can be used multiple times)", @@ -95,6 +101,7 @@ export interface AgentRunOptions extends CommandOptions { base?: string; image?: string[]; cwd?: string; + env?: string[]; label?: string[]; waitTimeout?: string; outputSchema?: string; @@ -320,15 +327,41 @@ function loadRunImages( } function parseRunLabels(labelFlags: string[] | undefined): Record { + return parseKeyValueFlags(labelFlags, { + flagName: "--label", + code: "INVALID_LABEL", + noun: "label", + pluralNoun: "Labels", + }); +} + +function parseRunEnv(envFlags: string[] | undefined): Record { + return parseKeyValueFlags(envFlags, { + flagName: "--env", + code: "INVALID_ENV", + noun: "environment variable", + pluralNoun: "Environment variables", + }); +} + +function parseKeyValueFlags( + flags: string[] | undefined, + options: { + flagName: string; + code: CommandError["code"]; + noun: string; + pluralNoun: string; + }, +): Record { const labels: Record = {}; - if (!labelFlags) return labels; - for (const labelStr of labelFlags) { + if (!flags) return labels; + for (const labelStr of flags) { const eqIndex = labelStr.indexOf("="); if (eqIndex === -1) { throw { - code: "INVALID_LABEL", - message: `Invalid label format: ${labelStr}`, - details: "Labels must be in key=value format", + code: options.code, + message: `Invalid ${options.noun} format: ${labelStr}`, + details: `${options.pluralNoun} must be in key=value format`, } satisfies CommandError; } const key = labelStr.slice(0, eqIndex); @@ -394,6 +427,8 @@ export async function runRunCommand( : undefined; const labels = parseRunLabels(options.label); + const env = parseRunEnv(options.env); + const requestEnv = Object.keys(env).length > 0 ? env : undefined; if (outputSchema) { let structuredAgent: AgentSnapshotPayload | null = null; @@ -410,6 +445,7 @@ export async function runRunCommand( initialPrompt: structuredPrompt, outputSchema, images, + env: requestEnv, git, worktreeName: options.worktree, labels: Object.keys(labels).length > 0 ? labels : undefined, @@ -479,6 +515,7 @@ export async function runRunCommand( thinkingOptionId, initialPrompt: prompt, images, + env: requestEnv, git, worktreeName: options.worktree, labels: Object.keys(labels).length > 0 ? labels : undefined, diff --git a/packages/server/src/client/daemon-client.ts b/packages/server/src/client/daemon-client.ts index 9ce4a06ea..382eb3d0f 100644 --- a/packages/server/src/client/daemon-client.ts +++ b/packages/server/src/client/daemon-client.ts @@ -249,6 +249,7 @@ export interface CreateAgentRequestOptions extends AgentConfigOverrides { config?: AgentSessionConfig; provider?: AgentProvider; cwd?: string; + env?: CreateAgentRequestMessage["env"]; workspaceId?: string; initialPrompt?: string; clientMessageId?: string; @@ -1837,6 +1838,7 @@ export class DaemonClient { type: "create_agent_request", requestId, config, + ...(options.env ? { env: options.env } : {}), ...(options.workspaceId !== undefined ? { workspaceId: options.workspaceId } : {}), ...(options.initialPrompt ? { initialPrompt: options.initialPrompt } : {}), ...(options.clientMessageId ? { clientMessageId: options.clientMessageId } : {}), @@ -4667,6 +4669,7 @@ function resolveAgentConfig(options: CreateAgentRequestOptions): AgentSessionCon config, provider, cwd, + env: _env, workspaceId: _workspaceId, initialPrompt: _initialPrompt, images: _images, diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index b22cbffee..be71518f2 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -1,4 +1,5 @@ import { expect, test, vi } from "vitest"; +import { spawn } from "node:child_process"; import { mkdtempSync, rmSync } from "node:fs"; import { join } from "node:path"; import { tmpdir } from "node:os"; @@ -147,6 +148,46 @@ class TestAgentClient implements AgentClient { } } +class EnvProbeAgentClient extends TestAgentClient { + probe: Promise<{ probe: string | null; agentId: string | null }> | null = null; + + override async createSession( + config: AgentSessionConfig, + launchContext?: AgentLaunchContext, + ): Promise { + const script = ` + process.stdout.write(JSON.stringify({ + probe: process.env.CHUNK14_PROBE ?? null, + agentId: process.env.PASEO_AGENT_ID ?? null + })); + `; + const child = spawn(process.execPath, ["-e", script], { + cwd: config.cwd, + env: { ...process.env, ...launchContext?.env }, + stdio: ["ignore", "pipe", "pipe"], + }); + this.probe = new Promise((resolve, reject) => { + let stdout = ""; + let stderr = ""; + child.stdout.on("data", (chunk: Buffer) => { + stdout += chunk.toString(); + }); + child.stderr.on("data", (chunk: Buffer) => { + stderr += chunk.toString(); + }); + child.on("error", reject); + child.on("close", (code) => { + if (code !== 0) { + reject(new Error(`env probe exited ${code}: ${stderr}`)); + return; + } + resolve(JSON.parse(stdout) as { probe: string | null; agentId: string | null }); + }); + }); + return new TestAgentSession(config); + } +} + class TestAgentSession implements AgentSession { readonly provider = "codex" as const; readonly capabilities = TEST_CAPABILITIES; @@ -373,6 +414,40 @@ test("normalizeConfig injects the provider default model when omitted", async () expect(snapshot.config.modeId).toBe("auto"); }); +test("createAgent forwards request env into the spawned provider process", async () => { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-env-test-")); + const client = new EnvProbeAgentClient(); + const manager = new AgentManager({ + clients: { + codex: client, + }, + logger, + idFactory: () => "00000000-0000-4000-8000-00000000e001", + }); + + try { + await manager.createAgent( + { + provider: "codex", + cwd: workdir, + }, + undefined, + { + env: { + CHUNK14_PROBE: "expected", + }, + }, + ); + + await expect(client.probe).resolves.toEqual({ + probe: "expected", + agentId: "00000000-0000-4000-8000-00000000e001", + }); + } finally { + rmSync(workdir, { recursive: true, force: true }); + } +}); + test("normalizeConfig strips legacy 'default' model id", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-test-")); const storagePath = join(workdir, "agents"); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 3c3a95d79..00b3d5fc8 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -805,6 +805,7 @@ export class AgentManager { labels?: Record; workspaceId?: string; initialPrompt?: string; + env?: Record; persistSession?: boolean; }, ): Promise { @@ -826,7 +827,7 @@ export class AgentManager { const normalizedConfig = this.applyDaemonAppendSystemPrompt( await this.normalizeConfig(injectedConfig), ); - const launchContext = this.buildLaunchContext(resolvedAgentId); + const launchContext = this.buildLaunchContext(resolvedAgentId, options?.env); const client = await this.requireAvailableClient({ provider: normalizedConfig.provider, }); @@ -3469,10 +3470,11 @@ export class AgentManager { : next; } - private buildLaunchContext(agentId: string): AgentLaunchContext { + private buildLaunchContext(agentId: string, env?: Record): AgentLaunchContext { return { agentId, env: { + ...env, PASEO_AGENT_ID: agentId, }, }; diff --git a/packages/server/src/server/agent/providers/opencode-agent.test.ts b/packages/server/src/server/agent/providers/opencode-agent.test.ts index 215417ba6..bd2c72afb 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.test.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.test.ts @@ -1138,6 +1138,40 @@ describe("OpenCode adapter startTurn error handling", () => { }); }); +describe("OpenCodeAgentClient env", () => { + test("passes launch-context env to env-specific server acquisition", async () => { + const runtime = new TestOpenCodeRuntime(); + const openCodeClient = new TestOpenCodeClient(); + runtime.enqueueClient(openCodeClient); + const cwd = tmpCwd(); + const client = new OpenCodeAgentClient(createTestLogger(), undefined, { runtime }); + + try { + const session = await client.createSession( + { + provider: "opencode", + cwd, + }, + { + env: { + CHUNK14_PROBE: "expected", + }, + }, + ); + await session.close(); + + expect(runtime.acquisitions[0]).toMatchObject({ + force: false, + env: { + CHUNK14_PROBE: "expected", + }, + }); + } finally { + rmSync(cwd, { recursive: true, force: true }); + } + }); +}); + describe("OpenCode persisted sessions", () => { test("listPersistedAgents returns only sessions whose cwd matches the requested cwd", async () => { const runtime = new TestOpenCodeRuntime(); diff --git a/packages/server/src/server/agent/providers/opencode-agent.ts b/packages/server/src/server/agent/providers/opencode-agent.ts index 9e28b4b73..ec28c29c0 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.ts @@ -961,7 +961,10 @@ interface OpenCodeAgentClientDeps { class ProductionOpenCodeRuntime implements OpenCodeRuntime { constructor(private readonly serverManager: OpenCodeServerManager) {} - async acquireServer(options: { force: boolean }): Promise { + async acquireServer(options: { + force: boolean; + env?: Record; + }): Promise { return this.serverManager.acquire(options); } @@ -1007,7 +1010,10 @@ export class OpenCodeAgentClient implements AgentClient { options?: AgentCreateSessionOptions, ): Promise { const openCodeConfig = this.assertConfig(config); - const acquisition = await this.runtime.acquireServer({ force: false }); + const acquisition = await this.runtime.acquireServer({ + force: false, + env: launchContext?.env, + }); const { url } = acquisition.server; const client = this.runtime.createClient({ baseUrl: url, diff --git a/packages/server/src/server/agent/providers/opencode/runtime.ts b/packages/server/src/server/agent/providers/opencode/runtime.ts index ff30f72f7..4c4cb6440 100644 --- a/packages/server/src/server/agent/providers/opencode/runtime.ts +++ b/packages/server/src/server/agent/providers/opencode/runtime.ts @@ -10,7 +10,10 @@ export interface OpenCodeServerAcquisition { } export interface OpenCodeRuntime { - acquireServer(options: { force: boolean }): Promise; + acquireServer(options: { + force: boolean; + env?: Record; + }): Promise; ensureServerRunning(): Promise<{ port: number; url: string }>; createClient(options: { baseUrl: string; directory: string }): OpencodeClient; shutdown(): Promise; diff --git a/packages/server/src/server/agent/providers/opencode/server-manager.ts b/packages/server/src/server/agent/providers/opencode/server-manager.ts index 1705ca5fa..1a4b4db94 100644 --- a/packages/server/src/server/agent/providers/opencode/server-manager.ts +++ b/packages/server/src/server/agent/providers/opencode/server-manager.ts @@ -21,7 +21,10 @@ export interface OpenCodeServerAcquisition { export interface OpenCodeServerManagerLike { ensureRunning(): Promise<{ port: number; url: string }>; - acquire(options: { force: boolean }): Promise; + acquire(options: { + force: boolean; + env?: Record; + }): Promise; } export interface OpenCodeServerGeneration { @@ -91,10 +94,22 @@ export class OpenCodeServerManager implements OpenCodeServerManagerLike { return acquisition.server; } - async acquire(options: { force: boolean }): Promise { + async acquire(options: { + force: boolean; + env?: Record; + }): Promise { + if (options.env) { + const server = await this.startDedicatedServer(options.env); + return this.acquireServer(server); + } + const server = options.force ? await this.getForcedRefreshServer() : await this.getCurrentServer(); + return this.acquireServer(server); + } + + private acquireServer(server: OpenCodeServerGeneration): OpenCodeServerAcquisition { server.refCount += 1; let released = false; return { @@ -164,7 +179,16 @@ export class OpenCodeServerManager implements OpenCodeServerManagerLike { } } - private async startServer(): Promise { + private async startDedicatedServer( + env: Record, + ): Promise { + const server = await this.startServer(env); + server.retired = true; + this.retiredServers.add(server); + return server; + } + + private async startServer(launchEnv?: Record): Promise { const port = await findAvailablePort(); const url = `http://127.0.0.1:${port}`; const launchPrefix = await resolveProviderCommandPrefix( @@ -179,7 +203,10 @@ export class OpenCodeServerManager implements OpenCodeServerManagerLike { { detached: process.platform !== "win32", stdio: ["ignore", "pipe", "pipe"], - ...createProviderEnvSpec({ runtimeSettings: this.runtimeSettings }), + ...createProviderEnvSpec({ + runtimeSettings: this.runtimeSettings, + overlays: [launchEnv], + }), }, ); diff --git a/packages/server/src/server/agent/providers/opencode/test-server-manager.ts b/packages/server/src/server/agent/providers/opencode/test-server-manager.ts index baab73f9b..cda20879d 100644 --- a/packages/server/src/server/agent/providers/opencode/test-server-manager.ts +++ b/packages/server/src/server/agent/providers/opencode/test-server-manager.ts @@ -2,6 +2,7 @@ import type { OpenCodeServerAcquisition, OpenCodeServerManagerLike } from "./ser export interface TestOpenCodeServerAcquisition { force: boolean; + env?: Record; released: boolean; } @@ -15,9 +16,13 @@ export class TestOpenCodeServerManager implements OpenCodeServerManagerLike { return this.server; } - async acquire(options: { force: boolean }): Promise { + async acquire(options: { + force: boolean; + env?: Record; + }): Promise { const acquisition: TestOpenCodeServerAcquisition = { force: options.force, + env: options.env, released: false, }; this.acquisitions.push(acquisition); diff --git a/packages/server/src/server/agent/providers/opencode/test-utils/test-opencode-runtime.ts b/packages/server/src/server/agent/providers/opencode/test-utils/test-opencode-runtime.ts index 0499dbde4..15e1c8ebd 100644 --- a/packages/server/src/server/agent/providers/opencode/test-utils/test-opencode-runtime.ts +++ b/packages/server/src/server/agent/providers/opencode/test-utils/test-opencode-runtime.ts @@ -8,7 +8,11 @@ interface OpenCodeResponse { } export class TestOpenCodeRuntime implements OpenCodeRuntime { - readonly acquisitions: Array<{ force: boolean; releaseCount: number }> = []; + readonly acquisitions: Array<{ + force: boolean; + env?: Record; + releaseCount: number; + }> = []; readonly clientCreations: Array<{ baseUrl: string; directory: string }> = []; private readonly clients: TestOpenCodeClient[] = []; @@ -18,8 +22,11 @@ export class TestOpenCodeRuntime implements OpenCodeRuntime { this.clients.push(client); } - async acquireServer(options: { force: boolean }): Promise { - const acquisition = { force: options.force, releaseCount: 0 }; + async acquireServer(options: { + force: boolean; + env?: Record; + }): Promise { + const acquisition = { force: options.force, env: options.env, releaseCount: 0 }; this.acquisitions.push(acquisition); return { server: this.server, diff --git a/packages/server/src/server/agent/providers/pi/agent.test.ts b/packages/server/src/server/agent/providers/pi/agent.test.ts index 30a2d6d55..0984a7095 100644 --- a/packages/server/src/server/agent/providers/pi/agent.test.ts +++ b/packages/server/src/server/agent/providers/pi/agent.test.ts @@ -42,6 +42,22 @@ async function createSession(pi = new FakePi()): Promise<{ return { pi, session, events }; } +test("forwards launch-context env to the Pi process launch", async () => { + const pi = new FakePi(); + const client = createClient(pi); + const session = await client.createSession(createConfig(), { + env: { + CHUNK14_PROBE: "expected", + }, + }); + + expect(pi.recordedLaunches[0]?.env).toEqual({ + CHUNK14_PROBE: "expected", + }); + + await session.close(); +}); + class SessionEvents { private readonly events: AgentStreamEvent[] = []; private readonly waiters: Array<{ diff --git a/packages/server/src/server/agent/providers/pi/agent.ts b/packages/server/src/server/agent/providers/pi/agent.ts index 3c4607c99..1f5dc012b 100644 --- a/packages/server/src/server/agent/providers/pi/agent.ts +++ b/packages/server/src/server/agent/providers/pi/agent.ts @@ -843,7 +843,7 @@ export class PiRpcAgentClient implements AgentClient { async createSession( config: AgentSessionConfig, - _launchContext?: AgentLaunchContext, + launchContext?: AgentLaunchContext, ): Promise { const mcpConfig = await this.prepareMcpConfig(config.cwd, config.mcpServers); let runtimeSession: PiRuntimeSession; @@ -857,6 +857,7 @@ export class PiRpcAgentClient implements AgentClient { config.systemPrompt, config.daemonAppendSystemPrompt, ), + env: launchContext?.env, mcpConfigPath: mcpConfig?.path, }); } catch (error) { diff --git a/packages/server/src/server/agent/providers/pi/runtime.ts b/packages/server/src/server/agent/providers/pi/runtime.ts index 0e7a50b5e..8925ed1a0 100644 --- a/packages/server/src/server/agent/providers/pi/runtime.ts +++ b/packages/server/src/server/agent/providers/pi/runtime.ts @@ -21,6 +21,7 @@ export interface PiRuntimeLaunch { export interface PiStartSessionInput { cwd: string; + env?: Record; model?: string; thinkingOptionId?: string; session?: string; @@ -84,7 +85,13 @@ export function buildPiLaunch(input: { return { cwd: input.session.cwd, argv, - env: input.runtimeSettings?.env, + env: + input.runtimeSettings?.env || input.session.env + ? { + ...input.runtimeSettings?.env, + ...input.session.env, + } + : undefined, model: input.session.model, thinkingOptionId: input.session.thinkingOptionId, session: input.session.session, diff --git a/packages/server/src/server/messages.test.ts b/packages/server/src/server/messages.test.ts index cfcfb285c..411d81c78 100644 --- a/packages/server/src/server/messages.test.ts +++ b/packages/server/src/server/messages.test.ts @@ -1,9 +1,48 @@ import { describe, expect, test } from "vitest"; import type { AgentStreamEvent } from "./agent/agent-sdk-types.js"; -import { serializeAgentStreamEvent } from "./messages.js"; +import { SessionInboundMessageSchema, serializeAgentStreamEvent } from "./messages.js"; describe("serializeAgentStreamEvent", () => { + test("accepts create_agent_request env records", () => { + const parsed = SessionInboundMessageSchema.parse({ + type: "create_agent_request", + requestId: "req-env", + config: { + provider: "codex", + cwd: "/tmp", + }, + env: { + CHUNK14_PROBE: "expected", + }, + attachments: [], + }); + + expect(parsed).toMatchObject({ + type: "create_agent_request", + env: { + CHUNK14_PROBE: "expected", + }, + }); + }); + + test("rejects non-string create_agent_request env values", () => { + const parsed = SessionInboundMessageSchema.safeParse({ + type: "create_agent_request", + requestId: "req-env", + config: { + provider: "codex", + cwd: "/tmp", + }, + env: { + CHUNK14_PROBE: 14, + }, + attachments: [], + }); + + expect(parsed.success).toBe(false); + }); + test("preserves user_message text as-is", () => { const event: AgentStreamEvent = { type: "timeline", diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 46a0978cc..edfc7c18c 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -3084,6 +3084,7 @@ export class Session { images, attachments, labels, + env, } = msg; this.sessionLogger.info( { cwd: config.cwd, provider: config.provider, worktreeName }, @@ -3124,6 +3125,7 @@ export class Session { labels, workspaceId: resolvedWorkspace.workspaceId, initialPrompt: trimmedPrompt, + env, }); await this.forwardAgentUpdate(snapshot); diff --git a/packages/server/src/shared/messages.ts b/packages/server/src/shared/messages.ts index 574a1eb05..497b241ad 100644 --- a/packages/server/src/shared/messages.ts +++ b/packages/server/src/shared/messages.ts @@ -1061,6 +1061,7 @@ export type GitSetupOptions = z.infer; export const CreateAgentRequestMessageSchema = z.object({ type: z.literal("create_agent_request"), config: AgentSessionConfigSchema, + env: z.record(z.string()).optional(), workspaceId: z.string().optional(), worktreeName: z.string().optional(), initialPrompt: z.string().optional(),