@desplega.ai/agent-swarm 1.173.0 → 2.1.0

package/src/tests/pi-mono-adapter.test.ts

@@ -1,13 +1,26 @@
 import { afterAll, afterEach, beforeAll, describe, expect, spyOn, test } from "bun:test";
-import { existsSync, mkdirSync, readFileSync, rmSync, symlinkSync, writeFileSync } from "node:fs";
+import {
+  existsSync,
+  mkdirSync,
+  mkdtempSync,
+  readFileSync,
+  rmSync,
+  symlinkSync,
+  writeFileSync,
+} from "node:fs";
+import { tmpdir } from "node:os";
 import { join } from "node:path";
+import { OPENAI_COMPATIBLE_PROVIDERS } from "@desplega/model-catalog";
 import * as piCodingAgent from "@earendil-works/pi-coding-agent";
 import { configureHttpResolver, resetHttpResolver } from "../prompts/resolver";
 import {
+  checkPiMonoCredentials,
   createPiRuntimeAuth,
   extractPiAssistantText,
   PiMonoAdapter,
   resolveModel,
+  type SessionModelsJson,
+  writeSessionModelsJson,
 } from "../providers/pi-mono-adapter";
 import type { ProviderSessionConfig } from "../providers/types";
 
@@ -151,8 +164,8 @@ describe("PiMonoAdapter.createSession — OPENROUTER_BASE_URL gateway", () => {
     const model = capturedOptions?.model as { baseUrl?: string; id?: string } | undefined;
     expect(model?.id).toBe("google/gemini-3-flash-preview");
     expect(model?.baseUrl).toBe(GATEWAY);
-    // The override was materialized for the runtime to load.
-    expect(existsSync(join(agentDir, "models.json"))).toBe(true);
+    // The override lives in the session's own copy, never the shared file.
+    expect(existsSync(join(agentDir, "models.json"))).toBe(false);
   });
 
   test("session model keeps the default openrouter.ai baseUrl when env is unset", async () => {
@@ -310,7 +323,7 @@ describe("PiMonoAdapter.createSession — MCP setup", () => {
       const adapter = new PiMonoAdapter();
       await adapter.createSession(makeMcpConfig(`http://localhost:${server.port}`));
       const options = createAgentSessionSpy.mock.calls[0]?.[0];
-      expect(options?.customTools?.map((tool) => tool.name)).toEqual(["get-tasks"]);
+      expect(options?.customTools?.map((tool) => tool.name)).toEqual(["get-tasks", "bash"]);
       expect(initializeSessions).toEqual([null, null]);
       expect(new Set(seen)).toEqual(new Set(["pi-retry-runtime"]));
     } finally {
@@ -330,7 +343,9 @@ describe("PiMonoAdapter.createSession — MCP setup", () => {
       await adapter.createSession(makeMcpConfig(`http://localhost:${server.port}`));
       expect(initializeSessions).toEqual([null, null, null]);
       expect(failures).toEqual(["initialize"]);
-      expect(createAgentSessionSpy.mock.calls[0]?.[0]?.customTools).toEqual([]);
+      expect(
+        createAgentSessionSpy.mock.calls[0]?.[0]?.customTools?.map((tool) => tool.name),
+      ).toEqual(["bash"]);
     } finally {
       server.stop(true);
     }
@@ -542,6 +557,371 @@ describe("createPiRuntimeAuth", () => {
   });
 });
 
+// ─── NaN / Helmcode custom providers (models.json, runtime keys, resolve) ───
+
+const NAN_ENV = { NAN_API_KEY: "example-nan-picked" };
+const HELMCODE_ENV = { HELMCODE_API_KEY: "example-helmcode-picked" };
+async function readModelsJson(path: string | undefined): Promise<Record<string, any>> {
+  if (!path) throw new Error("expected a per-session models.json");
+  return JSON.parse(await Bun.file(path).text());
+}
+
+function withTempDir(prefix: string, fn: (dir: string) => Promise<void>): () => Promise<void> {
+  return async () => {
+    const dir = mkdtempSync(join(tmpdir(), prefix));
+    try {
+      await fn(dir);
+    } finally {
+      rmSync(dir, { recursive: true, force: true });
+    }
+  };
+}
+
+/** Compose a session's models.json and always remove it afterwards. */
+async function withSessionModels(
+  dir: string,
+  env: Record<string, string | undefined>,
+  model: string | undefined,
+  fn: (session: SessionModelsJson) => Promise<void>,
+): Promise<void> {
+  const session = await writeSessionModelsJson(dir, env, model);
+  try {
+    await fn(session);
+  } finally {
+    await session.dispose();
+  }
+}
+
+describe("writeSessionModelsJson (pi)", () => {
+  test(
+    "writes both blocks with $VAR apiKeys to a private file, never the key value or the shared file",
+    withTempDir("pi-custom-fresh-", async (dir) => {
+      await withSessionModels(dir, { ...NAN_ENV, ...HELMCODE_ENV }, undefined, async (s) => {
+        expect(s.modelsPath).toBeDefined();
+        expect(s.modelsPath?.startsWith(dir)).toBe(false);
+        expect(existsSync(join(dir, "models.json"))).toBe(false);
+        const raw = await Bun.file(s.modelsPath as string).text();
+        expect(raw).not.toContain("example-nan-picked");
+        expect(raw).not.toContain("example-helmcode-picked");
+        const written = JSON.parse(raw);
+        expect(written.providers.nan).toMatchObject({
+          baseUrl: OPENAI_COMPATIBLE_PROVIDERS.nan.baseUrl,
+          api: "openai-completions",
+          apiKey: "$NAN_API_KEY",
+        });
+        expect(written.providers.helmcode).toMatchObject({
+          baseUrl: OPENAI_COMPATIBLE_PROVIDERS.helmcode.baseUrl,
+          api: "openai-completions",
+          apiKey: "$HELMCODE_API_KEY",
+        });
+        expect(written.providers.helmcode.models.map((m: { id: string }) => m.id)).toEqual(
+          OPENAI_COMPATIBLE_PROVIDERS.helmcode.models.map((m) => m.id),
+        );
+        const glm = written.providers.nan.models.find((m: { id: string }) => m.id === "glm5.3");
+        expect(glm).toMatchObject({
+          contextWindow: 1000000,
+          maxTokens: 131072,
+          reasoning: true,
+          compat: { supportsDeveloperRole: true },
+        });
+      });
+    }),
+  );
+
+  test(
+    "dispose removes the private file",
+    withTempDir("pi-custom-dispose-", async (dir) => {
+      const s = await writeSessionModelsJson(dir, NAN_ENV);
+      expect(existsSync(s.modelsPath as string)).toBe(true);
+      await s.dispose();
+      await s.dispose();
+      expect(existsSync(s.modelsPath as string)).toBe(false);
+    }),
+  );
+
+  test(
+    "only the HELMCODE_API_NAN alias set: nan block still uses $NAN_API_KEY",
+    withTempDir("pi-custom-alias-", async (dir) => {
+      await withSessionModels(dir, { HELMCODE_API_NAN: "example-alias" }, undefined, async (s) => {
+        const written = await readModelsJson(s.modelsPath);
+        expect(written.providers.nan.apiKey).toBe("$NAN_API_KEY");
+        expect(written.providers.helmcode).toBeUndefined();
+        expect(JSON.stringify(written)).not.toContain("example-alias");
+      });
+    }),
+  );
+
+  test(
+    "no key, default gateway, no marker: the session loads the shared file",
+    withTempDir("pi-custom-noop-", async (dir) => {
+      const userConfig = { providers: { nan: { baseUrl: "https://mine.example/v1" } } };
+      await Bun.write(join(dir, "models.json"), JSON.stringify(userConfig));
+      const s = await writeSessionModelsJson(dir, {});
+      expect(s.modelsPath).toBeUndefined();
+      expect(JSON.parse(readFileSync(join(dir, "models.json"), "utf-8"))).toEqual(userConfig);
+    }),
+  );
+
+  test(
+    "a user-authored providers.nan in the shared file wins; helmcode is still added",
+    withTempDir("pi-custom-user-", async (dir) => {
+      const userNan = { baseUrl: "https://mine.example/v1", apiKey: "$MY_VAR", models: [] };
+      const shared = JSON.stringify({ providers: { nan: userNan } });
+      await Bun.write(join(dir, "models.json"), shared);
+      await withSessionModels(dir, { ...NAN_ENV, ...HELMCODE_ENV }, undefined, async (s) => {
+        const written = await readModelsJson(s.modelsPath);
+        expect(written.providers.nan).toEqual(userNan);
+        expect(written.providers.helmcode.apiKey).toBe("$HELMCODE_API_KEY");
+      });
+      expect(readFileSync(join(dir, "models.json"), "utf-8")).toBe(shared);
+    }),
+  );
+
+  test(
+    "keeps other providers and top-level keys; the gateway comes from this session's env",
+    withTempDir("pi-custom-merge-", async (dir) => {
+      const gateway = "https://control-plane.example/proxy/v1";
+      await Bun.write(
+        join(dir, "models.json"),
+        JSON.stringify({
+          $schema: "https://pi.dev/schemas/models.schema.json",
+          providers: {
+            custom: { baseUrl: "https://custom.example/v1", api: "openai-completions" },
+            openrouter: { baseUrl: gateway },
+          },
+        }),
+      );
+      // An earlier version wrote that gateway into the shared file.
+      await Bun.write(
+        join(dir, ".agent-swarm-openrouter-override.json"),
+        JSON.stringify({ baseUrl: gateway }),
+      );
+      const sharedBefore = readFileSync(join(dir, "models.json"), "utf-8");
+      await withSessionModels(dir, NAN_ENV, undefined, async (s) => {
+        // This session uses the default gateway: the leftover override
+        // recorded in the legacy marker is reverted in its copy.
+        const written = await readModelsJson(s.modelsPath);
+        expect(written.$schema).toBe("https://pi.dev/schemas/models.schema.json");
+        expect(written.providers.openrouter).toBeUndefined();
+        expect(written.providers.custom.baseUrl).toBe("https://custom.example/v1");
+        expect(written.providers.nan.apiKey).toBe("$NAN_API_KEY");
+      });
+      const other = "https://other-gateway.example/v1";
+      await withSessionModels(dir, { OPENROUTER_BASE_URL: other }, undefined, async (s) => {
+        expect((await readModelsJson(s.modelsPath)).providers.openrouter).toEqual({
+          baseUrl: other,
+        });
+      });
+      // The shared file is never written by the composer.
+      expect(readFileSync(join(dir, "models.json"), "utf-8")).toBe(sharedBefore);
+    }),
+  );
+
+  test(
+    "sibling sessions get their own endpoints, providers and model ids",
+    withTempDir("pi-custom-siblings-", async (dir) => {
+      const a = await writeSessionModelsJson(
+        dir,
+        { HELMCODE_API_KEY: "a", HELMCODE_BASE_URL: "http://127.0.0.1:1/v1", ...NAN_ENV },
+        "nan/glm5.2",
+      );
+      const b = await writeSessionModelsJson(
+        dir,
+        { HELMCODE_API_KEY: "b", HELMCODE_BASE_URL: "http://127.0.0.1:2/v1" },
+        "helmcode/glm5.3",
+      );
+      try {
+        expect(a.modelsPath).not.toBe(b.modelsPath);
+        const first = await readModelsJson(a.modelsPath);
+        const second = await readModelsJson(b.modelsPath);
+        expect(first.providers.helmcode.baseUrl).toBe("http://127.0.0.1:1/v1");
+        expect(second.providers.helmcode.baseUrl).toBe("http://127.0.0.1:2/v1");
+        // B has no NaN key: no nan block, even while A's session holds one.
+        expect(first.providers.nan.models.map((m: { id: string }) => m.id)).toContain("glm5.2");
+        expect(second.providers.nan).toBeUndefined();
+      } finally {
+        await a.dispose();
+        await b.dispose();
+      }
+    }),
+  );
+
+  test(
+    "corrupt shared models.json: the copy holds our blocks only, the shared file is untouched",
+    withTempDir("pi-custom-corrupt-", async (dir) => {
+      await Bun.write(join(dir, "models.json"), "{not json");
+      const warn = spyOn(console, "warn").mockImplementation(() => {});
+      try {
+        await withSessionModels(dir, NAN_ENV, undefined, async (s) => {
+          expect((await readModelsJson(s.modelsPath)).providers.nan.apiKey).toBe("$NAN_API_KEY");
+        });
+      } finally {
+        warn.mockRestore();
+      }
+      expect(readFileSync(join(dir, "models.json"), "utf-8")).toBe("{not json");
+    }),
+  );
+
+  test(
+    "ModelRuntime resolves the written blocks as nan/<id> and helmcode/<id>",
+    withTempDir("pi-custom-runtime-", async (dir) => {
+      await withSessionModels(dir, { ...NAN_ENV, ...HELMCODE_ENV }, undefined, async (s) => {
+        const runtime = await piCodingAgent.ModelRuntime.create({
+          authPath: join(dir, "auth.json"),
+          modelsPath: s.modelsPath,
+        });
+        expect(runtime.getError()).toBeUndefined();
+        expect(runtime.getModel("nan", "qwen3.8-flash")?.baseUrl).toBe(
+          OPENAI_COMPATIBLE_PROVIDERS.nan.baseUrl,
+        );
+        expect(runtime.getModel("helmcode", "glm5.3")?.contextWindow).toBe(1048576);
+      });
+    }),
+  );
+});
+
+describe("createPiRuntimeAuth — NaN and Helmcode", () => {
+  test(
+    "HELMCODE_API_NAN alone sets the nan runtime key; NAN_API_KEY wins when both are set",
+    withTempDir("pi-custom-auth-", async (dir) => {
+      const models = await writeSessionModelsJson(dir, { NAN_API_KEY: "x", HELMCODE_API_KEY: "y" });
+      // pi's own agent-dir env var points auth.json at the temp dir.
+      const origAgentDir = process.env.PI_CODING_AGENT_DIR;
+      process.env.PI_CODING_AGENT_DIR = dir;
+      try {
+        const path = models.modelsPath;
+        const aliasOnly = await createPiRuntimeAuth({ HELMCODE_API_NAN: "a" }, path);
+        await expect(aliasOnly.getAuth("nan")).resolves.toMatchObject({ auth: { apiKey: "a" } });
+
+        const both = await createPiRuntimeAuth({ HELMCODE_API_NAN: "a", NAN_API_KEY: "b" }, path);
+        await expect(both.getAuth("nan")).resolves.toMatchObject({ auth: { apiKey: "b" } });
+
+        const helm = await createPiRuntimeAuth({ HELMCODE_API_KEY: "h" }, path);
+        await expect(helm.getAuth("helmcode")).resolves.toMatchObject({ auth: { apiKey: "h" } });
+      } finally {
+        await models.dispose();
+        if (origAgentDir === undefined) delete process.env.PI_CODING_AGENT_DIR;
+        else process.env.PI_CODING_AGENT_DIR = origAgentDir;
+      }
+
+      // Runtime keys never land in auth.json.
+      const auth = (await Bun.file(join(dir, "auth.json")).exists())
+        ? await Bun.file(join(dir, "auth.json")).text()
+        : "";
+      expect(auth).not.toContain('"a"');
+      expect(auth).not.toContain('"b"');
+      expect(auth).not.toContain('"h"');
+    }),
+  );
+});
+
+describe("resolveModel — models.json custom providers", () => {
+  test(
+    "falls back to the runtime for nan/qwen3.8-flash and helmcode/glm5.3",
+    withTempDir("pi-custom-resolve-", async (dir) => {
+      const models = await writeSessionModelsJson(dir, { ...NAN_ENV, ...HELMCODE_ENV });
+      const runtime = await piCodingAgent.ModelRuntime.create({
+        authPath: join(dir, "auth.json"),
+        modelsPath: models.modelsPath,
+      });
+      await models.dispose();
+      const nan = resolveModel("nan/qwen3.8-flash", NAN_ENV, runtime);
+      expect(nan?.provider).toBe("nan");
+      expect(nan?.id).toBe("qwen3.8-flash");
+      const helm = resolveModel("helmcode/glm5.3", HELMCODE_ENV, runtime);
+      expect(helm?.provider).toBe("helmcode");
+      expect(helm?.id).toBe("glm5.3");
+      expect(helm?.baseUrl).toBe(OPENAI_COMPATIBLE_PROVIDERS.helmcode.baseUrl);
+    }),
+  );
+
+  test(
+    "a nan/<id> outside the table (nan/glm5.2) gets a bare entry and resolves",
+    withTempDir("pi-custom-bare-", async (dir) => {
+      await withSessionModels(dir, NAN_ENV, "nan/glm5.2", async (s) => {
+        const written = await readModelsJson(s.modelsPath);
+        const ids = written.providers.nan.models.map((m: { id: string }) => m.id);
+        expect(ids).toContain("glm5.2");
+        expect(ids).toContain("glm5.3");
+        const runtime = await piCodingAgent.ModelRuntime.create({
+          authPath: join(dir, "auth.json"),
+          modelsPath: s.modelsPath,
+        });
+        const model = resolveModel("nan/glm5.2", NAN_ENV, runtime);
+        expect(model?.provider).toBe("nan");
+        expect(model?.id).toBe("glm5.2");
+        expect(model?.baseUrl).toBe(OPENAI_COMPATIBLE_PROVIDERS.nan.baseUrl);
+      });
+
+      // The bare entry belongs to that session: the next task's model never sees it.
+      await withSessionModels(dir, NAN_ENV, "nan/glm5.3", async (s) => {
+        const next = await readModelsJson(s.modelsPath);
+        expect(next.providers.nan.models.map((m: { id: string }) => m.id)).not.toContain("glm5.2");
+      });
+    }),
+  );
+
+  test("builtin hits are unchanged when a runtime is passed", async () => {
+    const runtime = { getModel: () => undefined };
+    const model = resolveModel("anthropic/claude-sonnet-5", {}, runtime);
+    expect(model?.provider).toBe("anthropic");
+  });
+
+  test("an unknown provider returns undefined and logs a scrubbed warning", () => {
+    const warn = spyOn(console, "warn").mockImplementation(() => {});
+    try {
+      const runtime = { getModel: () => undefined };
+      expect(resolveModel("nope/model-x", {}, runtime)).toBeUndefined();
+      expect(warn).toHaveBeenCalledTimes(1);
+      expect(String(warn.mock.calls[0]?.[0])).toContain('"nope/model-x"');
+      expect(String(warn.mock.calls[0]?.[0])).toContain("default model");
+    } finally {
+      warn.mockRestore();
+    }
+  });
+
+  test("without a runtime, nan/<id> does not resolve (pi-ai has no builtin nan)", () => {
+    const warn = spyOn(console, "warn").mockImplementation(() => {});
+    try {
+      expect(resolveModel("nan/glm5.3", NAN_ENV)).toBeUndefined();
+    } finally {
+      warn.mockRestore();
+    }
+  });
+});
+
+describe("checkPiMonoCredentials — NaN and Helmcode", () => {
+  const opts = { homeDir: "/home/worker", fs: { existsSync: () => false } };
+
+  test("MODEL_OVERRIDE=nan/* accepts NAN_API_KEY or the HELMCODE_API_NAN alias", async () => {
+    const env = { MODEL_OVERRIDE: "nan/glm5.3" };
+    expect((await checkPiMonoCredentials(env, opts)).ready).toBe(false);
+    expect((await checkPiMonoCredentials({ ...env, NAN_API_KEY: "x" }, opts)).ready).toBe(true);
+    expect((await checkPiMonoCredentials({ ...env, HELMCODE_API_NAN: "x" }, opts)).ready).toBe(
+      true,
+    );
+    // An OpenRouter key does not satisfy a nan/* model.
+    expect((await checkPiMonoCredentials({ ...env, OPENROUTER_API_KEY: "x" }, opts)).ready).toBe(
+      false,
+    );
+  });
+
+  test("MODEL_OVERRIDE=helmcode/* requires HELMCODE_API_KEY", async () => {
+    const env = { MODEL_OVERRIDE: "helmcode/glm5.3" };
+    const missing = await checkPiMonoCredentials({ ...env, NAN_API_KEY: "x" }, opts);
+    expect(missing.ready).toBe(false);
+    expect(missing.missing).toContain("HELMCODE_API_KEY");
+    expect((await checkPiMonoCredentials({ ...env, HELMCODE_API_KEY: "x" }, opts)).ready).toBe(
+      true,
+    );
+  });
+
+  test("permissive case accepts a NaN-only or Helmcode-only worker", async () => {
+    expect((await checkPiMonoCredentials({ HELMCODE_API_NAN: "x" }, opts)).ready).toBe(true);
+    expect((await checkPiMonoCredentials({ HELMCODE_API_KEY: "x" }, opts)).ready).toBe(true);
+  });
+});
+
 describe("Pi-mono event normalization", () => {
   test("extractPiAssistantText ignores user messages", () => {
     const text = extractPiAssistantText({

package/src/tests/pi-shell-env.test.ts

@@ -0,0 +1,170 @@
+import {
+  afterAll,
+  afterEach,
+  beforeAll,
+  beforeEach,
+  describe,
+  expect,
+  spyOn,
+  test,
+} from "bun:test";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import type { AgentSession } from "@earendil-works/pi-coding-agent";
+import { fetchResolvedEnv } from "../commands/runner";
+import { buildPiShellEnv, PiMonoAdapter } from "../providers/pi-mono-adapter";
+
+/**
+ * pi's Bash tool must run with the env the runner resolved for this task,
+ * not the worker's process.env: a row admitted after boot (exportToEnv
+ * flipped on) must reach the shell, and a row revoked after boot (flipped off)
+ * must leave it, even though process.env still holds the boot value.
+ */
+
+type ConfigRow = { key: string; value: string; isSecret: boolean; exportToEnv: boolean };
+
+let server: ReturnType<typeof Bun.serve>;
+let apiUrl = "";
+let rows: ConfigRow[] = [];
+
+beforeAll(() => {
+  server = Bun.serve({
+    port: 0,
+    async fetch(req) {
+      const url = new URL(req.url);
+      if (url.pathname === "/api/config/resolved") return Response.json({ configs: rows });
+      if (url.pathname !== "/mcp") return Response.json({ servers: [] });
+      const body = (await req.json()) as { id?: number; method?: string };
+      const result =
+        body.method === "tools/list"
+          ? { tools: [] }
+          : { protocolVersion: "2025-03-26", capabilities: {} };
+      return Response.json({ jsonrpc: "2.0", id: body.id ?? 1, result });
+    },
+  });
+  apiUrl = `http://localhost:${server.port}`;
+});
+
+afterAll(() => {
+  server.stop(true);
+});
+
+describe("pi Bash tool env", () => {
+  const envKeys = [
+    "PI_CODING_AGENT_DIR",
+    "OPENROUTER_API_KEY",
+    "OPENROUTER_BASE_URL",
+    "SHELL_REVOKED_CANARY",
+    "SHELL_ADMITTED_CANARY",
+  ] as const;
+  const saved: Record<string, string | undefined> = {};
+  let agentDir = "";
+  let log: ReturnType<typeof spyOn>;
+
+  beforeEach(() => {
+    for (const key of envKeys) saved[key] = process.env[key];
+    agentDir = mkdtempSync(join(tmpdir(), "pi-shell-env-"));
+    process.env.PI_CODING_AGENT_DIR = agentDir;
+    process.env.OPENROUTER_API_KEY = "example-test-key";
+    process.env.OPENROUTER_BASE_URL = "http://127.0.0.1:9/api/v1";
+    log = spyOn(console, "log").mockImplementation(() => {});
+  });
+  afterEach(() => {
+    log.mockRestore();
+    for (const key of envKeys) {
+      if (saved[key] === undefined) delete process.env[key];
+      else process.env[key] = saved[key];
+    }
+    rmSync(agentDir, { recursive: true, force: true });
+  });
+
+  async function sessionWithEnv(env: Record<string, string | undefined>): Promise<AgentSession> {
+    const session = await new PiMonoAdapter().createSession({
+      prompt: "hello",
+      systemPrompt: "test",
+      model: "openrouter/google/gemini-3-flash-preview",
+      role: "worker",
+      agentId: "test-agent",
+      taskId: "test-task",
+      apiUrl,
+      apiKey: "example-test-key",
+      cwd: agentDir,
+      logFile: join(agentDir, `session-${crypto.randomUUID()}.log`),
+      env,
+    });
+    return (session as unknown as { agentSession: AgentSession }).agentSession;
+  }
+
+  async function shell(session: AgentSession, command: string): Promise<string> {
+    const bash = session.getToolDefinition("bash");
+    if (!bash) throw new Error("bash tool missing");
+    const result = await bash.execute(
+      "call-1",
+      { command },
+      undefined,
+      undefined,
+      undefined as never,
+    );
+    return result.content
+      .map((part) => ("text" in part ? part.text : ""))
+      .join("")
+      .trim();
+  }
+
+  test("admission and revocation after boot reach the shell", async () => {
+    // Boot exported SHELL_REVOKED_CANARY and withheld SHELL_ADMITTED_CANARY; the
+    // operator then flipped both rows.
+    process.env.SHELL_REVOKED_CANARY = "boot-revoked";
+    delete process.env.SHELL_ADMITTED_CANARY;
+    rows = [
+      { key: "SHELL_REVOKED_CANARY", value: "row-revoked", isSecret: true, exportToEnv: false },
+      { key: "SHELL_ADMITTED_CANARY", value: "row-admitted", isSecret: true, exportToEnv: true },
+    ];
+    const { env } = await fetchResolvedEnv(apiUrl, "key", "agent-1", process.env);
+    const session = await sessionWithEnv(env);
+    try {
+      const out = await shell(
+        session,
+        // The shell's own environment: /proc on Linux, `env` on macOS (no procfs).
+        'echo "admitted=$SHELL_ADMITTED_CANARY revoked=$SHELL_REVOKED_CANARY."; { if [ -r /proc/$$/environ ]; then tr "\\0" "\\n" < /proc/$$/environ; else env; fi; } | grep -c "^SHELL_REVOKED_CANARY=" || true',
+      );
+      expect(out).toContain("admitted=row-admitted revoked=.");
+      expect(out.split("\n").at(-1)).toBe("0");
+    } finally {
+      session.dispose();
+    }
+  });
+
+  test("concurrent sessions keep their own env", async () => {
+    const first = await sessionWithEnv({ ...process.env, SHELL_ADMITTED_CANARY: "first" });
+    const second = await sessionWithEnv({ ...process.env, SHELL_ADMITTED_CANARY: undefined });
+    try {
+      const [a, b] = await Promise.all([
+        shell(first, 'echo "v=$SHELL_ADMITTED_CANARY"'),
+        shell(second, 'echo "v=$SHELL_ADMITTED_CANARY"'),
+      ]);
+      expect(a).toBe("v=first");
+      expect(b).toBe("v=");
+      expect(process.env.SHELL_ADMITTED_CANARY).toBeUndefined();
+    } finally {
+      first.dispose();
+      second.dispose();
+    }
+  });
+});
+
+describe("buildPiShellEnv", () => {
+  test("keeps the session env and takes only PATH and pi's session keys from its spawn env", () => {
+    const env = buildPiShellEnv(
+      { HOME: "/home/worker", PATH: "/usr/bin", TASK_ONLY: "yes" },
+      { PATH: "/pi/bin:/usr/bin", PI_SESSION_ID: "s1", PI_OTHER: "leak", PROCESS_ONLY: "leak" },
+    );
+    expect(env).toEqual({
+      HOME: "/home/worker",
+      PATH: "/pi/bin:/usr/bin",
+      TASK_ONLY: "yes",
+      PI_SESSION_ID: "s1",
+    });
+  });
+});