diff --git a/src/main/ipc-credentials-handlers.ts b/src/main/ipc-credentials-handlers.ts index 051e241e..3981b361 100644 --- a/src/main/ipc-credentials-handlers.ts +++ b/src/main/ipc-credentials-handlers.ts @@ -1,5 +1,5 @@ import { ipcMain, type IpcMainInvokeEvent } from "electron"; -import { streamSimple, type SimpleStreamOptions } from "@mariozechner/pi-ai"; +import type { SimpleStreamOptions } from "@mariozechner/pi-ai"; import type { DesktopResult } from "@shared/contracts"; import { IPC_CHANNELS } from "@shared/contracts"; import { @@ -16,6 +16,7 @@ import type { CredentialStore } from "./credentials/credential-store"; import type { StoredCredentialPayload } from "./credentials/types"; import { ok } from "./ipc-core-common"; import { classifyUnknownError } from "./errors/index.js"; +import { safeStreamSimpleWithEnv } from "./worker/runtime/pi-safe-stream"; /** Strip credential-shaped substrings from an error message before sending it * back to the renderer. Matches are anchored to issuer-specific prefixes so @@ -51,23 +52,6 @@ interface CredentialHandlerDeps { logger: { info: (...args: unknown[]) => void; warn: (...args: unknown[]) => void; error: (...args: unknown[]) => void }; } -/** Serializes `credentialsTest` calls that must mutate `process.env` to inject - * per-provider configuration for pi-ai (Azure / AWS / GCP ADC). pi-ai's - * `SimpleStreamOptions` has no per-request env escape hatch — its providers - * read `process.env` directly — so two concurrent tests would clobber each - * other's env and surface cross-credential failures. A module-level promise - * chain guarantees at most one `extraEnv`-mutating test runs at a time. */ -let credentialTestChain: Promise = Promise.resolve(); - -function serializeCredentialTest(task: () => Promise): Promise { - const run = credentialTestChain.then(task, task); - credentialTestChain = run.then( - () => undefined, - () => undefined, - ); - return run; -} - export function registerCredentialIpcHandlers(deps: CredentialHandlerDeps): () => void { const channels: string[] = [ IPC_CHANNELS.credentialsList, @@ -185,13 +169,14 @@ export function registerCredentialIpcHandlers(deps: CredentialHandlerDeps): () = maxTokens: 1, ...(extraHeaders ? { headers: extraHeaders } : {}), }; - const stream = streamSimple( + const stream = safeStreamSimpleWithEnv( model, { systemPrompt: "Reply with the single character 'k'.", messages: [{ role: "user", content: "ping", timestamp: Date.now() }], }, options, + resolved.extraEnv, ); const final = await stream.result(); if (final.stopReason === "error" || final.stopReason === "aborted") { @@ -207,28 +192,6 @@ export function registerCredentialIpcHandlers(deps: CredentialHandlerDeps): () = } }; - // Providers that require per-request env (Azure/AWS/GCP ADC) must mutate - // `process.env` because pi-ai reads it directly. Serialize those tests so - // concurrent invocations cannot observe each other's env. - if (resolved.extraEnv) { - const envPatch = resolved.extraEnv; - return serializeCredentialTest(async () => { - const saved = new Map(); - for (const [k, v] of Object.entries(envPatch)) { - saved.set(k, process.env[k]); - process.env[k] = v; - } - try { - return await runTest(); - } finally { - for (const [k, prev] of saved) { - if (prev === undefined) delete process.env[k]; - else process.env[k] = prev; - } - } - }); - } - return runTest(); }); diff --git a/src/main/worker/runtime/index.ts b/src/main/worker/runtime/index.ts index e9b6888b..5458b731 100644 --- a/src/main/worker/runtime/index.ts +++ b/src/main/worker/runtime/index.ts @@ -176,6 +176,7 @@ export class WorkerRunner { thinkingLevel: resolveThinkingLevel(job.thinkingEffort), apiKey: prepared.apiKey, extraHeaders, + streamEnv: prepared.extraEnv, rpc, onContextCompaction: (telemetry) => { post({ diff --git a/src/main/worker/runtime/pi-agent-factory.ts b/src/main/worker/runtime/pi-agent-factory.ts index 0c1192df..302edb5a 100644 --- a/src/main/worker/runtime/pi-agent-factory.ts +++ b/src/main/worker/runtime/pi-agent-factory.ts @@ -21,7 +21,7 @@ import type { Api, Model } from "@mariozechner/pi-ai"; import { buildAgentHooks } from "@main/agent-runtime/agent-hooks"; import type { ContextCompactionTelemetry } from "@main/agent-runtime/agent-hooks"; import type { WorkerRpcClient } from "../entry/rpc"; -import { safeStreamSimple } from "./pi-safe-stream"; +import { safeStreamSimpleWithEnv } from "./pi-safe-stream"; // Side-effect import: drags the hook signature contract into the typecheck // graph so any drift in pi-agent-core hook shapes breaks the build here. import "./pi-hook-contract"; @@ -38,6 +38,7 @@ export interface PiAgentFactoryInput { thinkingLevel: ThinkingLevel; apiKey: string; extraHeaders: Record | undefined; + streamEnv?: Record; rpc: WorkerRpcClient; onContextCompaction?: (telemetry: ContextCompactionTelemetry) => void; } @@ -76,11 +77,16 @@ export function createPiAgent(input: PiAgentFactoryInput): Agent { }, sessionId: input.taskSessionId, streamFn: (model, messages, options) => - safeStreamSimple(model, messages, { - ...options, - apiKey: input.apiKey, - headers: { ...(options?.headers ?? {}), ...(input.extraHeaders ?? {}) }, - }), + safeStreamSimpleWithEnv( + model, + messages, + { + ...options, + apiKey: input.apiKey, + headers: { ...(options?.headers ?? {}), ...(input.extraHeaders ?? {}) }, + }, + input.streamEnv, + ), beforeToolCall: hooks.beforeToolCall, transformContext: hooks.transformContext, toolExecution: "sequential", diff --git a/src/main/worker/runtime/pi-safe-stream.ts b/src/main/worker/runtime/pi-safe-stream.ts index 98c1430c..4f88f7ef 100644 --- a/src/main/worker/runtime/pi-safe-stream.ts +++ b/src/main/worker/runtime/pi-safe-stream.ts @@ -28,6 +28,13 @@ type SafePiStreamSource = ( options?: SimpleStreamOptions, ) => AssistantMessageEventStream; +type SafePiStreamSourceWithEnv = ( + model: Model, + context: Context, + options?: SimpleStreamOptions, + envPatch?: Readonly>, +) => AssistantMessageEventStream; + type ErrorAssistantMessage = AssistantMessage & { stopReason: Extract; }; @@ -72,6 +79,61 @@ function pushStreamError( stream.end(message); } +let scopedEnvChain: Promise = Promise.resolve(); + +function hasEnvPatch(envPatch: Readonly> | undefined): envPatch is Readonly> { + return !!envPatch && Object.keys(envPatch).length > 0; +} + +async function withProcessEnvPatch( + envPatch: Readonly>, + task: () => Promise, +): Promise { + const saved = new Map(); + for (const [key, value] of Object.entries(envPatch)) { + saved.set(key, process.env[key]); + process.env[key] = value; + } + try { + await task(); + } finally { + for (const [key, previous] of saved) { + if (previous === undefined) delete process.env[key]; + else process.env[key] = previous; + } + } +} + +function runWithScopedEnv( + envPatch: Readonly> | undefined, + task: () => Promise, +): Promise { + if (!hasEnvPatch(envPatch)) return task(); + const run = scopedEnvChain.then( + () => withProcessEnvPatch(envPatch, task), + () => withProcessEnvPatch(envPatch, task), + ); + scopedEnvChain = run.then( + () => undefined, + () => undefined, + ); + return run; +} + +async function pipeSourceStream( + source: PiStreamSource, + target: AssistantMessageEventStream, + model: Model, + context: Context, + options?: SimpleStreamOptions, +): Promise { + const sourceStream = await source(model, context, options); + for await (const event of sourceStream) { + target.push(event); + } + target.end(); +} + /** * Build a safe stream wrapper. When `gate` is provided, every call * acquires a slot keyed by `(model.provider, options.apiKey)` before @@ -84,7 +146,15 @@ export function createSafePiStream( source: PiStreamSource, gate?: CredentialConcurrencyGate, ): SafePiStreamSource { - return (model, context, options) => { + const scoped = createSafePiStreamWithEnv(source, gate); + return (model, context, options) => scoped(model, context, options); +} + +export function createSafePiStreamWithEnv( + source: PiStreamSource, + gate?: CredentialConcurrencyGate, +): SafePiStreamSourceWithEnv { + return (model, context, options, envPatch) => { const target = createAssistantMessageEventStream(); const apiKey = typeof options?.apiKey === "string" ? options.apiKey : ""; const acquireSlot = gate && apiKey @@ -94,11 +164,7 @@ export function createSafePiStream( .then(acquireSlot) .then(async (release) => { try { - const sourceStream = await source(model, context, options); - for await (const event of sourceStream) { - target.push(event); - } - target.end(); + await runWithScopedEnv(envPatch, () => pipeSourceStream(source, target, model, context, options)); } finally { release(); } @@ -115,3 +181,4 @@ const defaultGate: CredentialConcurrencyGate = new InMemoryCredentialConcurrency ); export const safeStreamSimple = createSafePiStream(streamSimple, defaultGate); +export const safeStreamSimpleWithEnv = createSafePiStreamWithEnv(streamSimple, defaultGate); diff --git a/tests/unit/pi-safe-stream.test.ts b/tests/unit/pi-safe-stream.test.ts index a689cc78..20d9c5c2 100644 --- a/tests/unit/pi-safe-stream.test.ts +++ b/tests/unit/pi-safe-stream.test.ts @@ -6,7 +6,7 @@ import { type Model, } from "@mariozechner/pi-ai"; import { describe, expect, it } from "vitest"; -import { createSafePiStream } from "../../src/main/worker/runtime/pi-safe-stream"; +import { createSafePiStream, createSafePiStreamWithEnv } from "../../src/main/worker/runtime/pi-safe-stream"; import { InMemoryCredentialConcurrencyGate } from "../../src/main/worker/runtime/credential-concurrency-gate"; function createModel(): Model<"openai-responses"> { @@ -188,4 +188,79 @@ describe("pi safe stream", () => { expect(result.stopReason).toBe("error"); expect(result.errorMessage).toBe("stream setup failed"); }); + + it("keeps scoped provider env active until the stream finishes", async () => { + const key = "AGENTCOMPANY_TEST_PROVIDER_ENV"; + const previous = process.env[key]; + process.env[key] = "host"; + const observed: Array = []; + const safeStream = createSafePiStreamWithEnv(() => { + observed.push(process.env[key]); + const stream = createAssistantMessageEventStream(); + setTimeout(() => { + observed.push(process.env[key]); + const message = createAssistantMessage(); + stream.push({ type: "done", reason: "stop", message }); + stream.end(message); + }, 5); + return stream; + }); + + try { + const stream = safeStream(createModel(), emptyContext, undefined, { [key]: "scoped" }); + for await (const _event of stream) { + // drain + } + await stream.result(); + + expect(observed).toEqual(["scoped", "scoped"]); + expect(process.env[key]).toBe("host"); + } finally { + if (previous === undefined) delete process.env[key]; + else process.env[key] = previous; + } + }); + + it("serializes streams that patch process env", async () => { + const key = "AGENTCOMPANY_TEST_SERIALIZED_ENV"; + const previous = process.env[key]; + delete process.env[key]; + const observed: Array = []; + let activeCalls = 0; + let peakActive = 0; + const safeStream = createSafePiStreamWithEnv(() => { + activeCalls += 1; + peakActive = Math.max(peakActive, activeCalls); + observed.push(process.env[key]); + const stream = createAssistantMessageEventStream(); + setTimeout(() => { + observed.push(process.env[key]); + const message = createAssistantMessage(); + stream.push({ type: "done", reason: "stop", message }); + stream.end(message); + activeCalls -= 1; + }, 5); + return stream; + }); + const drain = async (stream: ReturnType) => { + for await (const _event of stream) { + // drain + } + await stream.result(); + }; + + try { + await Promise.all([ + drain(safeStream(createModel(), emptyContext, undefined, { [key]: "one" })), + drain(safeStream(createModel(), emptyContext, undefined, { [key]: "two" })), + ]); + + expect(peakActive).toBe(1); + expect(observed).toEqual(["one", "one", "two", "two"]); + expect(process.env[key]).toBeUndefined(); + } finally { + if (previous === undefined) delete process.env[key]; + else process.env[key] = previous; + } + }); });