fix: propagate kaos env overlays (#654)

.changeset/propagate-exec-env-overrides.md

@@ -0,0 +1,7 @@
+---
+"@moonshot-ai/kaos": patch
+"@moonshot-ai/acp-adapter": patch
+"@moonshot-ai/kimi-code": patch
+---
+
+Propagate configured execution environment overrides across spawned processes.

packages/acp-adapter/src/kaos-acp.ts

@@ -89,6 +89,10 @@ export class AcpKaos implements Kaos {
     return new AcpKaos(this.conn, this.sessionId, this.inner.withCwd(cwd));
   }
 
+  withEnv(env: Record<string, string>): Kaos {
+    return new AcpKaos(this.conn, this.sessionId, this.inner.withEnv(env));
+  }
+
   stat(path: string, options?: { followSymlinks?: boolean }): Promise<StatResult> {
     return this.inner.stat(path, options);
   }

packages/acp-adapter/test/kaos-acp.test.ts

@@ -70,6 +70,7 @@ interface MockInnerKaos extends Kaos {
     getcwdCalls: number;
     chdirCalls: string[];
     withCwdCalls: string[];
+    withEnvCalls: Array<Record<string, string>>;
     statCalls: Array<{ path: string; options?: { followSymlinks?: boolean } }>;
     iterdirCalls: string[];
     globCalls: Array<{ path: string; pattern: string; options?: { caseSensitive?: boolean } }>;
@@ -91,6 +92,7 @@ function makeMockInner(opts?: { pathClass?: 'posix' | 'win32' }): MockInnerKaos
     getcwdCalls: 0,
     chdirCalls: [] as string[],
     withCwdCalls: [] as string[],
+    withEnvCalls: [] as Array<Record<string, string>>,
     statCalls: [] as Array<{ path: string; options?: { followSymlinks?: boolean } }>,
     iterdirCalls: [] as string[],
     globCalls: [] as Array<{ path: string; pattern: string; options?: { caseSensitive?: boolean } }>,
@@ -132,6 +134,11 @@ function makeMockInner(opts?: { pathClass?: 'posix' | 'win32' }): MockInnerKaos
       const child = makeMockInner();
       return child;
     },
+    withEnv: (env: Record<string, string>) => {
+      spy.withEnvCalls.push(env);
+      const child = makeMockInner();
+      return child;
+    },
     stat: async (path: string, options?: { followSymlinks?: boolean }) => {
       spy.statCalls.push({ path, options });
       return {
@@ -450,6 +457,24 @@ describe('AcpKaos', () => {
     });
   });
 
+  describe('withEnv', () => {
+    it('returns an AcpKaos that delegates env to inner and keeps the ACP bridge', async () => {
+      const conn = makeMockConn({
+        readHandler: async () => ({ content: 'BRIDGED' }),
+      });
+      const inner = makeMockInner();
+      const kaos = new AcpKaos(conn.asConn(), 's1', inner);
+      const env = { FOO: 'bar' };
+      const child = kaos.withEnv(env);
+
+      expect(child).toBeInstanceOf(AcpKaos);
+      const text = await child.readText('/foo.ts');
+      expect(text).toBe('BRIDGED');
+      expect(conn.readCalls).toEqual([{ sessionId: 's1', path: '/foo.ts' }]);
+      expect(inner.__spy.withEnvCalls).toEqual([env]);
+    });
+  });
+
   describe('pass-through delegation', () => {
     it('delegates pathClass, normpath, gethome, getcwd to inner', () => {
       const conn = makeMockConn({});

packages/agent-core/test/agent/harness/agent.ts

@@ -979,6 +979,7 @@ function createResumeNoSideEffectKaos(initialCwd: string): Kaos {
     gethome: () => '/home/test',
     getcwd: () => cwd,
     withCwd: (next: string) => createResumeNoSideEffectKaos(next),
+    withEnv: () => createResumeNoSideEffectKaos(cwd),
     chdir: async (next: string) => {
       cwd = next;
     },

packages/agent-core/test/tools/bash-env.test.ts

@@ -64,4 +64,37 @@ describe('BashTool noninteractive env semantics', () => {
       if (previous !== undefined) process.env['GIT_TERMINAL_PROMPT'] = previous;
     }
   });
+
+  it('lets kaos-level env override BashTool env and observes in-place updates', async () => {
+    const execWithEnv = vi.fn().mockResolvedValue(fakeProcess());
+    const sessionEnv = {
+      GIT_TERMINAL_PROMPT: 'configured',
+      KIMI_CODE_ENV: 'initial',
+    };
+    const kaos = createFakeKaos({ execWithEnv, osEnv: posixEnv }).withEnv(sessionEnv);
+    const tool = new BashTool(kaos, '/workspace');
+
+    await executeTool(tool, {
+      turnId: '0',
+      toolCallId: 'tc_kaos_env_1',
+      args: { command: 'true', timeout: 1000 },
+      signal,
+    });
+
+    const firstEnv = execWithEnv.mock.calls[0]?.[1] as Record<string, string>;
+    expect(firstEnv['GIT_TERMINAL_PROMPT']).toBe('configured');
+    expect(firstEnv['KIMI_CODE_ENV']).toBe('initial');
+
+    sessionEnv.KIMI_CODE_ENV = 'updated';
+    await executeTool(tool, {
+      turnId: '0',
+      toolCallId: 'tc_kaos_env_2',
+      args: { command: 'true', timeout: 1000 },
+      signal,
+    });
+
+    const secondEnv = execWithEnv.mock.calls[1]?.[1] as Record<string, string>;
+    expect(secondEnv['GIT_TERMINAL_PROMPT']).toBe('configured');
+    expect(secondEnv['KIMI_CODE_ENV']).toBe('updated');
+  });
 });

packages/agent-core/test/tools/fixtures/fake-kaos.ts

@@ -28,7 +28,10 @@ export const FAKE_OS_ENV: Environment = {
   shellPath: '/bin/bash',
 };
 
-export function createFakeKaos(overrides?: Partial<Kaos>): Kaos {
+export function createFakeKaos(
+  overrides?: Partial<Kaos>,
+  envLayers: readonly Record<string, string>[] = [],
+): Kaos {
   // Hold cwd in a closure so `chdir` (which `config.update({cwd})` now
   // routes through) can mutate it and later `getcwd()` calls see the
   // update — mirroring real-kaos semantics without needing a backing fs.
@@ -40,7 +43,9 @@ export function createFakeKaos(overrides?: Partial<Kaos>): Kaos {
     normpath: (p: string) => p,
     gethome: () => '/home/test',
     getcwd: () => cwd,
-    withCwd: (next: string) => createFakeKaos({ ...overrides, getcwd: () => next }),
+    withCwd: (next: string) => createFakeKaos({ ...overrides, getcwd: () => next }, envLayers),
+    withEnv: (env: Record<string, string>) =>
+      createFakeKaos({ ...overrides, getcwd: () => cwd }, [...envLayers, env]),
     chdir: async (next: string) => {
       cwd = next;
     },
@@ -54,9 +59,31 @@ export function createFakeKaos(overrides?: Partial<Kaos>): Kaos {
     writeText: () => notImplemented('writeText'),
     mkdir: () => notImplemented('mkdir'),
     exec: () => notImplemented('exec'),
-    execWithEnv: () => notImplemented('execWithEnv'),
+    execWithEnv: (args, invocationEnv) => {
+      const mergedEnv = mergeEnvLayers(invocationEnv, envLayers);
+      if (overrides?.execWithEnv) return overrides.execWithEnv(args, mergedEnv);
+      return notImplemented('execWithEnv');
+    },
   };
-  return { ...base, ...overrides } as Kaos;
+  return {
+    ...base,
+    ...overrides,
+    execWithEnv: base.execWithEnv,
+    withCwd: base.withCwd,
+    withEnv: base.withEnv,
+  } as Kaos;
+}
+
+function mergeEnvLayers(
+  invocationEnv: Record<string, string> | undefined,
+  envLayers: readonly Record<string, string>[],
+): Record<string, string> | undefined {
+  if (envLayers.length === 0) return invocationEnv;
+  const merged: Record<string, string> = { ...invocationEnv };
+  for (const layer of envLayers) {
+    Object.assign(merged, layer);
+  }
+  return merged;
 }
 
 export const PERMISSIVE_WORKSPACE: WorkspaceConfig = {

packages/kaos/src/kaos.ts

@@ -37,6 +37,13 @@ export interface Kaos {
   chdir(path: string): Promise<void>;
   /** Return a new Kaos with the given `cwd`. */
   withCwd(cwd: string): Kaos;
+  /**
+   * Return a new Kaos that overlays `env` onto every spawned process.
+   *
+   * The provided record is read when a process is spawned, so callers may
+   * mutate a stable record to update future executions.
+   */
+  withEnv(env: Record<string, string>): Kaos;
   /** Return stat metadata for `path`. */
   stat(path: string, options?: { followSymlinks?: boolean }): Promise<StatResult>;
   /** Yield entry names in the directory at `path`. */

packages/kaos/src/local.ts

@@ -149,14 +149,20 @@ export class LocalKaos implements Kaos {
   readonly name: string = 'local';
   readonly osEnv: Environment;
   private _cwd: string;
+  private readonly _envLayers: readonly Record<string, string>[];
 
-  private constructor(osEnv: Environment, cwd?: string) {
+  private constructor(
+    osEnv: Environment,
+    cwd?: string,
+    envLayers: readonly Record<string, string>[] = [],
+  ) {
     // After construction we never touch `process.cwd()` / `process.chdir()`
     // — all path resolution goes through `this._cwd`. The default seeds
     // from `process.cwd()` but callers can pin to anything via `withCwd`
     // (or supplying `cwd` directly).
     this._cwd = normalize(cwd ?? process.cwd());
     this.osEnv = osEnv;
+    this._envLayers = envLayers;
   }
 
   /**
@@ -172,7 +178,11 @@ export class LocalKaos implements Kaos {
   }
 
   withCwd(cwd: string): LocalKaos {
-    return new LocalKaos(this.osEnv, cwd);
+    return new LocalKaos(this.osEnv, cwd, this._envLayers);
+  }
+
+  withEnv(env: Record<string, string>): LocalKaos {
+    return new LocalKaos(this.osEnv, this._cwd, [...this._envLayers, env]);
   }
 
   private _resolvePath(path: string): string {
@@ -533,6 +543,7 @@ export class LocalKaos implements Kaos {
       // (`taskkill /T` handles the tree there). We do not call `child.unref()`
       // because the parent still waits on the child's exit through `wait()`.
       detached: !isWindows,
+      env: this._buildExecEnv(),
     });
     await waitForSpawn(child);
     return new LocalProcess(child);
@@ -550,11 +561,23 @@ export class LocalKaos implements Kaos {
       cwd: this._cwd,
       stdio: ['pipe', 'pipe', 'pipe'],
       detached: !isWindows,
-      env,
+      env: this._buildExecEnv(env),
     });
     await waitForSpawn(child);
     return new LocalProcess(child);
   }
+
+  private _buildExecEnv(invocationEnv?: Record<string, string>): Record<string, string> | undefined {
+    if (this._envLayers.length === 0) return invocationEnv;
+    const merged: Record<string, string> = {
+      ...(process.env as Record<string, string>),
+      ...invocationEnv,
+    };
+    for (const layer of this._envLayers) {
+      Object.assign(merged, layer);
+    }
+    return merged;
+  }
 }
 
 // Wait for a freshly spawned ChildProcess to either emit 'spawn' (success) or

packages/kaos/src/ssh.ts

@@ -430,6 +430,7 @@ export class SSHKaos implements Kaos {
   private _sftp: SFTPWrapper;
   private _home: string;
   private _cwd: string;
+  private readonly _envLayers: readonly Record<string, string>[];
 
   // Stub: real wiring (probing the remote host via `uname` / `$SHELL` over the
   // SSH transport) is deferred.
@@ -439,15 +440,26 @@ export class SSHKaos implements Kaos {
     );
   }
 
-  private constructor(client: Client, sftp: SFTPWrapper, home: string, cwd: string) {
+  private constructor(
+    client: Client,
+    sftp: SFTPWrapper,
+    home: string,
+    cwd: string,
+    envLayers: readonly Record<string, string>[] = [],
+  ) {
     this._client = client;
     this._sftp = sftp;
     this._home = home;
     this._cwd = cwd;
+    this._envLayers = envLayers;
   }
 
   withCwd(cwd: string): SSHKaos {
-    return new SSHKaos(this._client, this._sftp, this._home, cwd);
+    return new SSHKaos(this._client, this._sftp, this._home, cwd, this._envLayers);
+  }
+
+  withEnv(env: Record<string, string>): SSHKaos {
+    return new SSHKaos(this._client, this._sftp, this._home, this._cwd, [...this._envLayers, env]);
   }
 
   private _resolvePath(path: string): string {
@@ -834,7 +846,7 @@ export class SSHKaos implements Kaos {
         'SSHKaos.exec(): at least one argument (the command to run) is required.',
       );
     }
-    return this._execInternal(args);
+    return this._execInternal(args, this._buildExecEnv());
   }
 
   execWithEnv(args: string[], env?: Record<string, string>): Promise<KaosProcess> {
@@ -843,7 +855,16 @@ export class SSHKaos implements Kaos {
         'SSHKaos.execWithEnv(): at least one argument (the command to run) is required.',
       );
     }
-    return this._execInternal(args, env);
+    return this._execInternal(args, this._buildExecEnv(env));
+  }
+
+  private _buildExecEnv(invocationEnv?: Record<string, string>): Record<string, string> | undefined {
+    if (this._envLayers.length === 0) return invocationEnv;
+    const merged: Record<string, string> = { ...invocationEnv };
+    for (const layer of this._envLayers) {
+      Object.assign(merged, layer);
+    }
+    return merged;
   }
 
   /**

packages/kaos/test/local.test.ts

@@ -640,6 +640,35 @@ describe('LocalKaos', () => {
       expect(exitCode).not.toBe(0);
     });
   });
+
+  describe('withEnv', () => {
+    it('overlays every spawned process and can be updated in place', async () => {
+      const env = {
+        KAOS_BASE_ENV: 'initial',
+        KAOS_COLLISION_ENV: 'configured',
+      };
+      const envKaos = kaos.withEnv(env);
+      const printEnv =
+        'process.stdout.write(`${process.env.KAOS_BASE_ENV}|${process.env.KAOS_COLLISION_ENV}|${process.env.KAOS_CALL_ENV}`)';
+
+      const first = await envKaos.exec('node', '-e', printEnv);
+      expect(await first.wait()).toBe(0);
+      expect((await streamToBuffer(first.stdout)).toString('utf-8')).toBe('initial|configured|undefined');
+
+      const second = await envKaos.execWithEnv(['node', '-e', printEnv], {
+        ...(process.env as Record<string, string>),
+        KAOS_COLLISION_ENV: 'host',
+        KAOS_CALL_ENV: 'call',
+      });
+      expect(await second.wait()).toBe(0);
+      expect((await streamToBuffer(second.stdout)).toString('utf-8')).toBe('initial|configured|call');
+
+      env.KAOS_BASE_ENV = 'updated';
+      const third = await envKaos.exec('node', '-e', printEnv);
+      expect(await third.wait()).toBe(0);
+      expect((await streamToBuffer(third.stdout)).toString('utf-8')).toBe('updated|configured|undefined');
+    });
+  });
 });
 
 describe('LocalKaos instance isolation', () => {

packages/kaos/test/ssh.test.ts

@@ -1280,10 +1280,16 @@ describe('SSHKaos mock success paths', () => {
       },
     };
     const instance = Object.create(SSHKaos.prototype) as SSHKaos;
-    const internal = instance as unknown as { _client: unknown; _cwd: string; _home: string };
+    const internal = instance as unknown as {
+      _client: unknown;
+      _cwd: string;
+      _home: string;
+      _envLayers: readonly Record<string, string>[];
+    };
     internal._client = fakeClient;
     internal._cwd = '/home/tester';
     internal._home = '/home/tester';
+    internal._envLayers = [];
 
     await expect(instance.execWithEnv(['echo', 'hi'], { FOO: 'bar' })).rejects.toThrow('stop');
 
@@ -1442,10 +1448,12 @@ describe('SSHKaos.close lifecycle', () => {
       _cwd: string;
       _home: string;
       _sftp: { end(): void };
+      _envLayers: readonly Record<string, string>[];
     };
     internals._client = new FakeClient();
     internals._cwd = '/tmp';
     internals._home = '/tmp';
+    internals._envLayers = [];
     internals._sftp = {
       end(): void {
         // no-op