Forward create agent env to provider launches (#1112)

This commit is contained in:
Mohamed Boudra
2026-05-20 15:23:05 +08:00
committed by GitHub
parent ca913728d9
commit 905a0985f9
16 changed files with 286 additions and 21 deletions
+42 -5
View File
@@ -46,6 +46,12 @@ export function addRunOptions(cmd: Command): Command {
[], [],
) )
.option("--cwd <path>", "Working directory (default: current)") .option("--cwd <path>", "Working directory (default: current)")
.option(
"--env <key=value>",
"Set environment variable(s) for the agent process (can be used multiple times)",
collectMultiple,
[],
)
.option( .option(
"--label <key=value>", "--label <key=value>",
"Add label(s) to the agent (can be used multiple times)", "Add label(s) to the agent (can be used multiple times)",
@@ -95,6 +101,7 @@ export interface AgentRunOptions extends CommandOptions {
base?: string; base?: string;
image?: string[]; image?: string[];
cwd?: string; cwd?: string;
env?: string[];
label?: string[]; label?: string[];
waitTimeout?: string; waitTimeout?: string;
outputSchema?: string; outputSchema?: string;
@@ -320,15 +327,41 @@ function loadRunImages(
} }
function parseRunLabels(labelFlags: string[] | undefined): Record<string, string> { function parseRunLabels(labelFlags: string[] | undefined): Record<string, string> {
return parseKeyValueFlags(labelFlags, {
flagName: "--label",
code: "INVALID_LABEL",
noun: "label",
pluralNoun: "Labels",
});
}
function parseRunEnv(envFlags: string[] | undefined): Record<string, string> {
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<string, string> {
const labels: Record<string, string> = {}; const labels: Record<string, string> = {};
if (!labelFlags) return labels; if (!flags) return labels;
for (const labelStr of labelFlags) { for (const labelStr of flags) {
const eqIndex = labelStr.indexOf("="); const eqIndex = labelStr.indexOf("=");
if (eqIndex === -1) { if (eqIndex === -1) {
throw { throw {
code: "INVALID_LABEL", code: options.code,
message: `Invalid label format: ${labelStr}`, message: `Invalid ${options.noun} format: ${labelStr}`,
details: "Labels must be in key=value format", details: `${options.pluralNoun} must be in key=value format`,
} satisfies CommandError; } satisfies CommandError;
} }
const key = labelStr.slice(0, eqIndex); const key = labelStr.slice(0, eqIndex);
@@ -394,6 +427,8 @@ export async function runRunCommand(
: undefined; : undefined;
const labels = parseRunLabels(options.label); const labels = parseRunLabels(options.label);
const env = parseRunEnv(options.env);
const requestEnv = Object.keys(env).length > 0 ? env : undefined;
if (outputSchema) { if (outputSchema) {
let structuredAgent: AgentSnapshotPayload | null = null; let structuredAgent: AgentSnapshotPayload | null = null;
@@ -410,6 +445,7 @@ export async function runRunCommand(
initialPrompt: structuredPrompt, initialPrompt: structuredPrompt,
outputSchema, outputSchema,
images, images,
env: requestEnv,
git, git,
worktreeName: options.worktree, worktreeName: options.worktree,
labels: Object.keys(labels).length > 0 ? labels : undefined, labels: Object.keys(labels).length > 0 ? labels : undefined,
@@ -479,6 +515,7 @@ export async function runRunCommand(
thinkingOptionId, thinkingOptionId,
initialPrompt: prompt, initialPrompt: prompt,
images, images,
env: requestEnv,
git, git,
worktreeName: options.worktree, worktreeName: options.worktree,
labels: Object.keys(labels).length > 0 ? labels : undefined, labels: Object.keys(labels).length > 0 ? labels : undefined,
@@ -249,6 +249,7 @@ export interface CreateAgentRequestOptions extends AgentConfigOverrides {
config?: AgentSessionConfig; config?: AgentSessionConfig;
provider?: AgentProvider; provider?: AgentProvider;
cwd?: string; cwd?: string;
env?: CreateAgentRequestMessage["env"];
workspaceId?: string; workspaceId?: string;
initialPrompt?: string; initialPrompt?: string;
clientMessageId?: string; clientMessageId?: string;
@@ -1837,6 +1838,7 @@ export class DaemonClient {
type: "create_agent_request", type: "create_agent_request",
requestId, requestId,
config, config,
...(options.env ? { env: options.env } : {}),
...(options.workspaceId !== undefined ? { workspaceId: options.workspaceId } : {}), ...(options.workspaceId !== undefined ? { workspaceId: options.workspaceId } : {}),
...(options.initialPrompt ? { initialPrompt: options.initialPrompt } : {}), ...(options.initialPrompt ? { initialPrompt: options.initialPrompt } : {}),
...(options.clientMessageId ? { clientMessageId: options.clientMessageId } : {}), ...(options.clientMessageId ? { clientMessageId: options.clientMessageId } : {}),
@@ -4667,6 +4669,7 @@ function resolveAgentConfig(options: CreateAgentRequestOptions): AgentSessionCon
config, config,
provider, provider,
cwd, cwd,
env: _env,
workspaceId: _workspaceId, workspaceId: _workspaceId,
initialPrompt: _initialPrompt, initialPrompt: _initialPrompt,
images: _images, images: _images,
@@ -1,4 +1,5 @@
import { expect, test, vi } from "vitest"; import { expect, test, vi } from "vitest";
import { spawn } from "node:child_process";
import { mkdtempSync, rmSync } from "node:fs"; import { mkdtempSync, rmSync } from "node:fs";
import { join } from "node:path"; import { join } from "node:path";
import { tmpdir } from "node:os"; 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<AgentSession> {
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 { class TestAgentSession implements AgentSession {
readonly provider = "codex" as const; readonly provider = "codex" as const;
readonly capabilities = TEST_CAPABILITIES; readonly capabilities = TEST_CAPABILITIES;
@@ -373,6 +414,40 @@ test("normalizeConfig injects the provider default model when omitted", async ()
expect(snapshot.config.modeId).toBe("auto"); 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 () => { test("normalizeConfig strips legacy 'default' model id", async () => {
const workdir = mkdtempSync(join(tmpdir(), "agent-manager-test-")); const workdir = mkdtempSync(join(tmpdir(), "agent-manager-test-"));
const storagePath = join(workdir, "agents"); const storagePath = join(workdir, "agents");
@@ -805,6 +805,7 @@ export class AgentManager {
labels?: Record<string, string>; labels?: Record<string, string>;
workspaceId?: string; workspaceId?: string;
initialPrompt?: string; initialPrompt?: string;
env?: Record<string, string>;
persistSession?: boolean; persistSession?: boolean;
}, },
): Promise<ManagedAgent> { ): Promise<ManagedAgent> {
@@ -826,7 +827,7 @@ export class AgentManager {
const normalizedConfig = this.applyDaemonAppendSystemPrompt( const normalizedConfig = this.applyDaemonAppendSystemPrompt(
await this.normalizeConfig(injectedConfig), await this.normalizeConfig(injectedConfig),
); );
const launchContext = this.buildLaunchContext(resolvedAgentId); const launchContext = this.buildLaunchContext(resolvedAgentId, options?.env);
const client = await this.requireAvailableClient({ const client = await this.requireAvailableClient({
provider: normalizedConfig.provider, provider: normalizedConfig.provider,
}); });
@@ -3469,10 +3470,11 @@ export class AgentManager {
: next; : next;
} }
private buildLaunchContext(agentId: string): AgentLaunchContext { private buildLaunchContext(agentId: string, env?: Record<string, string>): AgentLaunchContext {
return { return {
agentId, agentId,
env: { env: {
...env,
PASEO_AGENT_ID: agentId, PASEO_AGENT_ID: agentId,
}, },
}; };
@@ -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", () => { describe("OpenCode persisted sessions", () => {
test("listPersistedAgents returns only sessions whose cwd matches the requested cwd", async () => { test("listPersistedAgents returns only sessions whose cwd matches the requested cwd", async () => {
const runtime = new TestOpenCodeRuntime(); const runtime = new TestOpenCodeRuntime();
@@ -961,7 +961,10 @@ interface OpenCodeAgentClientDeps {
class ProductionOpenCodeRuntime implements OpenCodeRuntime { class ProductionOpenCodeRuntime implements OpenCodeRuntime {
constructor(private readonly serverManager: OpenCodeServerManager) {} constructor(private readonly serverManager: OpenCodeServerManager) {}
async acquireServer(options: { force: boolean }): Promise<OpenCodeServerAcquisition> { async acquireServer(options: {
force: boolean;
env?: Record<string, string>;
}): Promise<OpenCodeServerAcquisition> {
return this.serverManager.acquire(options); return this.serverManager.acquire(options);
} }
@@ -1007,7 +1010,10 @@ export class OpenCodeAgentClient implements AgentClient {
options?: AgentCreateSessionOptions, options?: AgentCreateSessionOptions,
): Promise<AgentSession> { ): Promise<AgentSession> {
const openCodeConfig = this.assertConfig(config); 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 { url } = acquisition.server;
const client = this.runtime.createClient({ const client = this.runtime.createClient({
baseUrl: url, baseUrl: url,
@@ -10,7 +10,10 @@ export interface OpenCodeServerAcquisition {
} }
export interface OpenCodeRuntime { export interface OpenCodeRuntime {
acquireServer(options: { force: boolean }): Promise<OpenCodeServerAcquisition>; acquireServer(options: {
force: boolean;
env?: Record<string, string>;
}): Promise<OpenCodeServerAcquisition>;
ensureServerRunning(): Promise<{ port: number; url: string }>; ensureServerRunning(): Promise<{ port: number; url: string }>;
createClient(options: { baseUrl: string; directory: string }): OpencodeClient; createClient(options: { baseUrl: string; directory: string }): OpencodeClient;
shutdown(): Promise<void>; shutdown(): Promise<void>;
@@ -21,7 +21,10 @@ export interface OpenCodeServerAcquisition {
export interface OpenCodeServerManagerLike { export interface OpenCodeServerManagerLike {
ensureRunning(): Promise<{ port: number; url: string }>; ensureRunning(): Promise<{ port: number; url: string }>;
acquire(options: { force: boolean }): Promise<OpenCodeServerAcquisition>; acquire(options: {
force: boolean;
env?: Record<string, string>;
}): Promise<OpenCodeServerAcquisition>;
} }
export interface OpenCodeServerGeneration { export interface OpenCodeServerGeneration {
@@ -91,10 +94,22 @@ export class OpenCodeServerManager implements OpenCodeServerManagerLike {
return acquisition.server; return acquisition.server;
} }
async acquire(options: { force: boolean }): Promise<OpenCodeServerAcquisition> { async acquire(options: {
force: boolean;
env?: Record<string, string>;
}): Promise<OpenCodeServerAcquisition> {
if (options.env) {
const server = await this.startDedicatedServer(options.env);
return this.acquireServer(server);
}
const server = options.force const server = options.force
? await this.getForcedRefreshServer() ? await this.getForcedRefreshServer()
: await this.getCurrentServer(); : await this.getCurrentServer();
return this.acquireServer(server);
}
private acquireServer(server: OpenCodeServerGeneration): OpenCodeServerAcquisition {
server.refCount += 1; server.refCount += 1;
let released = false; let released = false;
return { return {
@@ -164,7 +179,16 @@ export class OpenCodeServerManager implements OpenCodeServerManagerLike {
} }
} }
private async startServer(): Promise<OpenCodeServerGeneration> { private async startDedicatedServer(
env: Record<string, string>,
): Promise<OpenCodeServerGeneration> {
const server = await this.startServer(env);
server.retired = true;
this.retiredServers.add(server);
return server;
}
private async startServer(launchEnv?: Record<string, string>): Promise<OpenCodeServerGeneration> {
const port = await findAvailablePort(); const port = await findAvailablePort();
const url = `http://127.0.0.1:${port}`; const url = `http://127.0.0.1:${port}`;
const launchPrefix = await resolveProviderCommandPrefix( const launchPrefix = await resolveProviderCommandPrefix(
@@ -179,7 +203,10 @@ export class OpenCodeServerManager implements OpenCodeServerManagerLike {
{ {
detached: process.platform !== "win32", detached: process.platform !== "win32",
stdio: ["ignore", "pipe", "pipe"], stdio: ["ignore", "pipe", "pipe"],
...createProviderEnvSpec({ runtimeSettings: this.runtimeSettings }), ...createProviderEnvSpec({
runtimeSettings: this.runtimeSettings,
overlays: [launchEnv],
}),
}, },
); );
@@ -2,6 +2,7 @@ import type { OpenCodeServerAcquisition, OpenCodeServerManagerLike } from "./ser
export interface TestOpenCodeServerAcquisition { export interface TestOpenCodeServerAcquisition {
force: boolean; force: boolean;
env?: Record<string, string>;
released: boolean; released: boolean;
} }
@@ -15,9 +16,13 @@ export class TestOpenCodeServerManager implements OpenCodeServerManagerLike {
return this.server; return this.server;
} }
async acquire(options: { force: boolean }): Promise<OpenCodeServerAcquisition> { async acquire(options: {
force: boolean;
env?: Record<string, string>;
}): Promise<OpenCodeServerAcquisition> {
const acquisition: TestOpenCodeServerAcquisition = { const acquisition: TestOpenCodeServerAcquisition = {
force: options.force, force: options.force,
env: options.env,
released: false, released: false,
}; };
this.acquisitions.push(acquisition); this.acquisitions.push(acquisition);
@@ -8,7 +8,11 @@ interface OpenCodeResponse {
} }
export class TestOpenCodeRuntime implements OpenCodeRuntime { export class TestOpenCodeRuntime implements OpenCodeRuntime {
readonly acquisitions: Array<{ force: boolean; releaseCount: number }> = []; readonly acquisitions: Array<{
force: boolean;
env?: Record<string, string>;
releaseCount: number;
}> = [];
readonly clientCreations: Array<{ baseUrl: string; directory: string }> = []; readonly clientCreations: Array<{ baseUrl: string; directory: string }> = [];
private readonly clients: TestOpenCodeClient[] = []; private readonly clients: TestOpenCodeClient[] = [];
@@ -18,8 +22,11 @@ export class TestOpenCodeRuntime implements OpenCodeRuntime {
this.clients.push(client); this.clients.push(client);
} }
async acquireServer(options: { force: boolean }): Promise<OpenCodeServerAcquisition> { async acquireServer(options: {
const acquisition = { force: options.force, releaseCount: 0 }; force: boolean;
env?: Record<string, string>;
}): Promise<OpenCodeServerAcquisition> {
const acquisition = { force: options.force, env: options.env, releaseCount: 0 };
this.acquisitions.push(acquisition); this.acquisitions.push(acquisition);
return { return {
server: this.server, server: this.server,
@@ -42,6 +42,22 @@ async function createSession(pi = new FakePi()): Promise<{
return { pi, session, events }; 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 { class SessionEvents {
private readonly events: AgentStreamEvent[] = []; private readonly events: AgentStreamEvent[] = [];
private readonly waiters: Array<{ private readonly waiters: Array<{
@@ -843,7 +843,7 @@ export class PiRpcAgentClient implements AgentClient {
async createSession( async createSession(
config: AgentSessionConfig, config: AgentSessionConfig,
_launchContext?: AgentLaunchContext, launchContext?: AgentLaunchContext,
): Promise<AgentSession> { ): Promise<AgentSession> {
const mcpConfig = await this.prepareMcpConfig(config.cwd, config.mcpServers); const mcpConfig = await this.prepareMcpConfig(config.cwd, config.mcpServers);
let runtimeSession: PiRuntimeSession; let runtimeSession: PiRuntimeSession;
@@ -857,6 +857,7 @@ export class PiRpcAgentClient implements AgentClient {
config.systemPrompt, config.systemPrompt,
config.daemonAppendSystemPrompt, config.daemonAppendSystemPrompt,
), ),
env: launchContext?.env,
mcpConfigPath: mcpConfig?.path, mcpConfigPath: mcpConfig?.path,
}); });
} catch (error) { } catch (error) {
@@ -21,6 +21,7 @@ export interface PiRuntimeLaunch {
export interface PiStartSessionInput { export interface PiStartSessionInput {
cwd: string; cwd: string;
env?: Record<string, string>;
model?: string; model?: string;
thinkingOptionId?: string; thinkingOptionId?: string;
session?: string; session?: string;
@@ -84,7 +85,13 @@ export function buildPiLaunch(input: {
return { return {
cwd: input.session.cwd, cwd: input.session.cwd,
argv, argv,
env: input.runtimeSettings?.env, env:
input.runtimeSettings?.env || input.session.env
? {
...input.runtimeSettings?.env,
...input.session.env,
}
: undefined,
model: input.session.model, model: input.session.model,
thinkingOptionId: input.session.thinkingOptionId, thinkingOptionId: input.session.thinkingOptionId,
session: input.session.session, session: input.session.session,
+40 -1
View File
@@ -1,9 +1,48 @@
import { describe, expect, test } from "vitest"; import { describe, expect, test } from "vitest";
import type { AgentStreamEvent } from "./agent/agent-sdk-types.js"; import type { AgentStreamEvent } from "./agent/agent-sdk-types.js";
import { serializeAgentStreamEvent } from "./messages.js"; import { SessionInboundMessageSchema, serializeAgentStreamEvent } from "./messages.js";
describe("serializeAgentStreamEvent", () => { 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", () => { test("preserves user_message text as-is", () => {
const event: AgentStreamEvent = { const event: AgentStreamEvent = {
type: "timeline", type: "timeline",
+2
View File
@@ -3084,6 +3084,7 @@ export class Session {
images, images,
attachments, attachments,
labels, labels,
env,
} = msg; } = msg;
this.sessionLogger.info( this.sessionLogger.info(
{ cwd: config.cwd, provider: config.provider, worktreeName }, { cwd: config.cwd, provider: config.provider, worktreeName },
@@ -3124,6 +3125,7 @@ export class Session {
labels, labels,
workspaceId: resolvedWorkspace.workspaceId, workspaceId: resolvedWorkspace.workspaceId,
initialPrompt: trimmedPrompt, initialPrompt: trimmedPrompt,
env,
}); });
await this.forwardAgentUpdate(snapshot); await this.forwardAgentUpdate(snapshot);
+1
View File
@@ -1061,6 +1061,7 @@ export type GitSetupOptions = z.infer<typeof GitSetupOptionsSchema>;
export const CreateAgentRequestMessageSchema = z.object({ export const CreateAgentRequestMessageSchema = z.object({
type: z.literal("create_agent_request"), type: z.literal("create_agent_request"),
config: AgentSessionConfigSchema, config: AgentSessionConfigSchema,
env: z.record(z.string()).optional(),
workspaceId: z.string().optional(), workspaceId: z.string().optional(),
worktreeName: z.string().optional(), worktreeName: z.string().optional(),
initialPrompt: z.string().optional(), initialPrompt: z.string().optional(),