fix(security): inject provider env into model streams

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<void> = Promise.resolve();
-
-function serializeCredentialTest<T>(task: () => Promise<T>): Promise<T> {
-  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<string, string | undefined>();
-        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();
   });
 

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({

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<string, string> | undefined;
+  streamEnv?: Record<string, string>;
   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",

src/main/worker/runtime/pi-safe-stream.ts

@@ -28,6 +28,13 @@ type SafePiStreamSource = (
   options?: SimpleStreamOptions,
 ) => AssistantMessageEventStream;
 
+type SafePiStreamSourceWithEnv = (
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+  envPatch?: Readonly<Record<string, string>>,
+) => AssistantMessageEventStream;
+
 type ErrorAssistantMessage = AssistantMessage & {
   stopReason: Extract<StopReason, "error" | "aborted">;
 };
@@ -72,6 +79,61 @@ function pushStreamError(
   stream.end(message);
 }
 
+let scopedEnvChain: Promise<void> = Promise.resolve();
+
+function hasEnvPatch(envPatch: Readonly<Record<string, string>> | undefined): envPatch is Readonly<Record<string, string>> {
+  return !!envPatch && Object.keys(envPatch).length > 0;
+}
+
+async function withProcessEnvPatch(
+  envPatch: Readonly<Record<string, string>>,
+  task: () => Promise<void>,
+): Promise<void> {
+  const saved = new Map<string, string | undefined>();
+  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<Record<string, string>> | undefined,
+  task: () => Promise<void>,
+): Promise<void> {
+  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<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+): Promise<void> {
+  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);

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<string | undefined> = [];
+    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<string | undefined> = [];
+    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<typeof safeStream>) => {
+      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;
+    }
+  });
 });