From 79db9d1bfd04caf8dcc52825ba177ebb0dbd1b7b Mon Sep 17 00:00:00 2001 From: Michael Yong Date: Thu, 3 Sep 2026 23:48:48 -0700 Subject: [PATCH] Apply later-turn environment to Pi and ACP --- .../src/bridge/bridge.test.ts | 17 ++++++++++ .../provider-bridge-acp/src/bridge/bridge.ts | 26 +++++++++++++++- plugins/provider-pi/src/bridge/bridge.ts | 18 +++++++++-- .../src/bridge/bridge.turn-options.test.ts | 31 +++++++++++++++++++ .../provider-pi/src/bridge/fake-pi-rpc.mjs | 6 ++++ 5 files changed, 95 insertions(+), 3 deletions(-) diff --git a/packages/provider-bridge-acp/src/bridge/bridge.test.ts b/packages/provider-bridge-acp/src/bridge/bridge.test.ts index 0ea2072880..8923c2afa2 100644 --- a/packages/provider-bridge-acp/src/bridge/bridge.test.ts +++ b/packages/provider-bridge-acp/src/bridge/bridge.test.ts @@ -1561,6 +1561,23 @@ describe("acp bridge", () => { expect(agentMessageTexts()).toContain("echo:hello there"); }); + it("rebuilds the agent with environment from a later turn", async () => { + const envVars = { FAKE_ACP_LOAD_SESSION: "1", FAKE_ACP_PROMPT_ERROR: "1" }; + const { providerThreadId } = await startThread({ envVars }); + const turnId = sendTurnRequest("turn/start", providerThreadId, { + input: [{ type: "text", text: "fresh environment", mentions: [] }], + options: executionOptions({ + envVars: { FAKE_ACP_PROMPT_ERROR: "0" }, + providerOptions: { acpLaunchSpec: acpLaunchSpec({ envVars }) }, + }), + }); + + expect((await waitForResponse(turnId)).error).toBeUndefined(); + expect(await waitForTurnCompleted()).toMatchObject({ status: "completed" }); + expect(agentMessageTexts()).toContain("echo:fresh environment"); + expect(notifications("session/replaced")).toHaveLength(1); + }); + it("authenticates ACP sessions with cached tokens when advertised", async () => { const { providerThreadId } = await startThread({ envVars: { FAKE_ACP_AUTH_METHODS: "cached_token" }, diff --git a/packages/provider-bridge-acp/src/bridge/bridge.ts b/packages/provider-bridge-acp/src/bridge/bridge.ts index 66884bbed1..1a596c5309 100644 --- a/packages/provider-bridge-acp/src/bridge/bridge.ts +++ b/packages/provider-bridge-acp/src/bridge/bridge.ts @@ -36,6 +36,7 @@ import { promises as fs, readFileSync } from "node:fs"; import { createServer, type Server, type Socket } from "node:net"; import { dirname, isAbsolute, basename, relative, resolve } from "node:path"; import { fileURLToPath } from "node:url"; +import { isDeepStrictEqual } from "node:util"; import { z } from "zod"; type DecodedToolCallResponse = ReturnType; @@ -159,6 +160,7 @@ interface AcpPendingTurnInput { interface AcpThreadSession { bbThreadId: string; + construction: AcpSessionParams; providerThreadId: string; cwd: string; dialect: AcpDialect; @@ -1699,6 +1701,7 @@ async function startAgentSession( }); session = { bbThreadId, + construction: params, providerThreadId: "", cwd: params.cwd, dialect, @@ -2590,7 +2593,7 @@ async function handleRequest( case "turn/start": { const params = request.params; - const session = liveSessionForThread(params.threadId); + let session = liveSessionForThread(params.threadId); if (session === undefined) { sendError(request.id, -32000, "No active ACP session"); return; @@ -2599,6 +2602,27 @@ async function handleRequest( sendError(request.id, -32000, "A turn is already active"); return; } + if (Object.keys(params.options.envVars ?? {}).length > 0) { + const envVars = { + ...(decodeLaunchSpec(params.options.providerOptions)?.env ?? {}), + ...params.options.envVars, + }; + if (!isDeepStrictEqual(envVars, session.construction.envVars ?? {})) { + const previousProviderThreadId = session.providerThreadId; + session = await startAgentSession({ + kind: "resume", + params: { ...session.construction, envVars }, + resumeProviderThreadId: previousProviderThreadId, + }); + sendNotification(BRIDGE_NOTIFICATION_METHODS.sessionReplaced, { + threadId: params.threadId, + providerThreadId: session.providerThreadId, + reason: + "Execution settings changed; the ACP session was rebuilt to apply them.", + contextLost: session.providerThreadId !== previousProviderThreadId, + }); + } + } const pending: AcpPendingTurnInput = { clientRequestId: params.clientRequestId, input: params.input, diff --git a/plugins/provider-pi/src/bridge/bridge.ts b/plugins/provider-pi/src/bridge/bridge.ts index cb187f5b53..8c0abf0473 100644 --- a/plugins/provider-pi/src/bridge/bridge.ts +++ b/plugins/provider-pi/src/bridge/bridge.ts @@ -9,6 +9,7 @@ import { } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { isDeepStrictEqual } from "node:util"; import { z } from "zod"; import { BRIDGE_JSON_RPC_ERRORS, @@ -16,6 +17,7 @@ import { PROVIDER_BRIDGE_PROTOCOL_VERSION, THREAD_DELTA_GRAMMAR_V3, THREAD_DELTA_NOTIFICATION_METHOD, + buildShellEnvOverrides, bridgeRequestEnvelopeSchema, createBridgeIo, createBridgeLineHandler, @@ -1031,6 +1033,13 @@ async function reconcileTurnOptions( ): Promise { const turnOptions = buildPiTurnOptions(options); const construction = threadSession.construction; + const shellEnvOverrides = + options.envVars && Object.keys(options.envVars).length > 0 + ? { BB_THREAD_ID: threadId, ...buildShellEnvOverrides(options.envVars) } + : undefined; + const environmentChanged = + shellEnvOverrides !== undefined && + !isDeepStrictEqual(shellEnvOverrides, construction.shellEnvOverrides); const changedModelRequest = turnOptions.model !== undefined && turnOptions.model !== construction.model ? turnOptions.model @@ -1038,7 +1047,11 @@ async function reconcileTurnOptions( const thinkingLevelChanged = turnOptions.thinkingLevel !== undefined && turnOptions.thinkingLevel !== construction.thinkingLevel; - if (changedModelRequest === undefined && !thinkingLevelChanged) { + if ( + !environmentChanged && + changedModelRequest === undefined && + !thinkingLevelChanged + ) { return threadSession; } const nextModel = @@ -1050,11 +1063,12 @@ async function reconcileTurnOptions( (threadSession.constructionModel === undefined || threadSession.constructionModel.provider !== nextModel.provider || threadSession.constructionModel.id !== nextModel.id); - if (!modelChanged && !thinkingLevelChanged) { + if (!environmentChanged && !modelChanged && !thinkingLevelChanged) { return threadSession; } const replacement = await rebuildThreadSession(threadId, threadSession, { ...construction, + ...(shellEnvOverrides === undefined ? {} : { shellEnvOverrides }), ...(turnOptions.model === undefined ? {} : { model: turnOptions.model }), ...(turnOptions.thinkingLevel === undefined ? {} diff --git a/plugins/provider-pi/src/bridge/bridge.turn-options.test.ts b/plugins/provider-pi/src/bridge/bridge.turn-options.test.ts index 3b32510fa1..bf8338bef2 100644 --- a/plugins/provider-pi/src/bridge/bridge.turn-options.test.ts +++ b/plugins/provider-pi/src/bridge/bridge.turn-options.test.ts @@ -1,3 +1,5 @@ +import { readFileSync } from "node:fs"; +import { join } from "node:path"; import { afterEach, beforeEach, expect, it, vi } from "vitest"; import { z } from "zod"; import type { @@ -71,6 +73,35 @@ function turnStart( }); } +it( + "rebuilds the session with environment from a later turn", + async () => { + const threadId = "thr_turn_options_env"; + const envLog = join(harness.workspaceDir, "env.log"); + const options = (marker: string) => ({ + ...MINI, + envVars: { FAKE_PI_ENV_LOG: envLog, FAKE_PI_ENV_MARKER: marker }, + }); + await harness.startThread(threadId, { options: options("first") }); + + expect( + (await turnStart(1, threadId, "first", options("first"))).error, + ).toBeUndefined(); + const seen = await harness.waitForTurnBoundary(threadId, 0); + expect( + (await turnStart(2, threadId, "second", options("second"))).error, + ).toBeUndefined(); + await harness.waitForTurnBoundary(threadId, seen); + + expect(readFileSync(envLog, "utf8").trim().split("\n")).toEqual([ + "first", + "second", + ]); + expect(sessionReplacements(threadId)).toHaveLength(1); + }, + TURN_OPTIONS_TEST_TIMEOUT_MS, +); + it( "rebuilds the session on the model a later turn carries", async () => { diff --git a/plugins/provider-pi/src/bridge/fake-pi-rpc.mjs b/plugins/provider-pi/src/bridge/fake-pi-rpc.mjs index 49ddf70f14..b04aada8e6 100644 --- a/plugins/provider-pi/src/bridge/fake-pi-rpc.mjs +++ b/plugins/provider-pi/src/bridge/fake-pi-rpc.mjs @@ -101,6 +101,12 @@ const extensionPath = flag("--extension"); const processLogPath = process.env.FAKE_PI_PROCESS_LOG; const commandLogPath = process.env.FAKE_PI_COMMAND_LOG; const promptDumpPath = process.env.FAKE_PI_PROMPT_DUMP; +if (process.env.FAKE_PI_ENV_LOG) { + appendFileSync( + process.env.FAKE_PI_ENV_LOG, + `${process.env.FAKE_PI_ENV_MARKER ?? ""}\n`, + ); +} if (sessionFile !== undefined) { mkdirSync(dirname(sessionFile), { recursive: true }); if (!existsSync(sessionFile)) {