@desplega.ai/agent-swarm 1.164.1 → 1.166.0

package/src/tests/acp-adapter.test.ts

@@ -1,5 +1,5 @@
 import { afterEach, describe, expect, spyOn, test } from "bun:test";
-import { mkdtempSync, rmSync } from "node:fs";
+import { chmodSync, existsSync, mkdirSync, mkdtempSync, rmSync } from "node:fs";
 import { createServer, type IncomingMessage, type Server, type ServerResponse } from "node:http";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
@@ -22,7 +22,12 @@ import {
   sanitizeAcpConfigOptions,
   toAcpMcpServers,
 } from "../providers/acp-adapter";
-import { AcpTargetResolutionError, resolveAcpTarget } from "../providers/acp-targets";
+import {
+  AcpTargetResolutionError,
+  COPILOT_INSTRUCTIONS_FILE,
+  copilotInstructionsDir,
+  resolveAcpTarget,
+} from "../providers/acp-targets";
 import type { ProviderEvent, ProviderSessionConfig } from "../providers/types";
 import { listenOnFreePort } from "./test-net";
 
@@ -536,6 +541,190 @@ new AgentSideConnection((connection) => new FakeAgent(connection), stream);
     });
   });
 
+  test("Gemini preset supplies command, credentials, trust, and a model environment fallback", () => {
+    const config = baseConfig({
+      model: "gemini-3-flash-preview",
+      env: {
+        PATH: "/bin",
+        HOME: "/home/test",
+        ACP_TARGET: "gemini",
+        GEMINI_API_KEY: "example-test-key",
+        GOOGLE_GENAI_USE_VERTEXAI: "true",
+        GOOGLE_CLOUD_PROJECT: "example-project",
+        OPENAI_API_KEY: "not-for-gemini",
+      },
+    });
+    const target = resolveAcpTarget(config);
+
+    expect(target.target).toBe("gemini");
+    expect(target.command(config)).toEqual(["gemini", "--acp"]);
+    const env = target.env(config);
+    expect(env).toMatchObject({
+      GEMINI_API_KEY: "example-test-key",
+      GOOGLE_GENAI_USE_VERTEXAI: "true",
+      GOOGLE_CLOUD_PROJECT: "example-project",
+      GEMINI_CLI_TRUST_WORKSPACE: "true",
+      GEMINI_MODEL: "gemini-3-flash-preview",
+    });
+    expect(env.OPENAI_API_KEY).toBeUndefined();
+    // No artifact written for this config, so the built-in prompt stays.
+    expect(env.GEMINI_SYSTEM_MD).toBeUndefined();
+  });
+
+  test("Gemini preset writes the system prompt outside the cwd and points GEMINI_SYSTEM_MD at it", async () => {
+    const cwd = mkdtempSync(join(tmpdir(), "acp-gemini-cwd-"));
+    try {
+      const config = baseConfig({
+        cwd,
+        systemPrompt: "You are a swarm worker.",
+        env: { PATH: "/bin", ACP_TARGET: "gemini" },
+      });
+      const target = resolveAcpTarget(config);
+      await target.writeSystemPromptArtifact(config);
+      const path = target.env(config).GEMINI_SYSTEM_MD;
+
+      expect(path).toBeDefined();
+      expect(path!.startsWith(cwd)).toBe(false);
+      expect(await Bun.file(path!).text()).toBe("You are a swarm worker.");
+      // A second session gets its own file.
+      const other = baseConfig({ cwd, systemPrompt: "other", env: config.env });
+      await target.writeSystemPromptArtifact(other);
+      expect(target.env(other).GEMINI_SYSTEM_MD).not.toBe(path);
+      rmSync(join(path!, ".."), { recursive: true, force: true });
+      rmSync(join(target.env(other).GEMINI_SYSTEM_MD!, ".."), { recursive: true, force: true });
+    } finally {
+      rmSync(cwd, { recursive: true, force: true });
+    }
+  });
+
+  test("Gemini preset keeps the built-in prompt when the swarm prompt is empty", async () => {
+    const config = baseConfig({ systemPrompt: "  ", env: { PATH: "/bin", ACP_TARGET: "gemini" } });
+    const target = resolveAcpTarget(config);
+    await target.writeSystemPromptArtifact(config);
+
+    expect(target.env(config).GEMINI_SYSTEM_MD).toBeUndefined();
+  });
+
+  test("Copilot preset forwards its own credentials, pins the model, and never forwards GITHUB_TOKEN", () => {
+    const config = baseConfig({
+      model: "openai/gpt-5.2",
+      env: {
+        PATH: "/bin",
+        HOME: "/home/test",
+        ACP_TARGET: "copilot",
+        COPILOT_GITHUB_TOKEN: "example-copilot-token",
+        COPILOT_PROVIDER_BASE_URL: "https://llm.example/v1",
+        COPILOT_PROVIDER_API_KEY: "example-provider-key",
+        COPILOT_CUSTOM_INSTRUCTIONS_DIRS: "/opt/team-instructions",
+        GITHUB_TOKEN: "example-git-token",
+        GH_TOKEN: "example-gh-token",
+      },
+    });
+    const target = resolveAcpTarget(config);
+
+    expect(target.target).toBe("copilot");
+    expect(target.command(config)).toEqual(["copilot", "--acp"]);
+    const env = target.env(config);
+    expect(env).toMatchObject({
+      COPILOT_GITHUB_TOKEN: "example-copilot-token",
+      COPILOT_PROVIDER_BASE_URL: "https://llm.example/v1",
+      COPILOT_PROVIDER_API_KEY: "example-provider-key",
+      COPILOT_MODEL: "openai/gpt-5.2",
+      COPILOT_AUTO_UPDATE: "false",
+      COPILOT_CUSTOM_INSTRUCTIONS_DIRS: `${copilotInstructionsDir(config)},/opt/team-instructions`,
+    });
+    // The worker's git token lacks the Copilot Requests permission and would
+    // shadow a BYOK or COPILOT_GITHUB_TOKEN setup.
+    expect(env.GITHUB_TOKEN).toBeUndefined();
+    expect(env.GH_TOKEN).toBeUndefined();
+  });
+
+  test("Copilot preset delivers the system prompt outside the repo and removes it after the session", async () => {
+    const cwd = makeTempDir();
+    const binDir = makeTempDir();
+    const captureFile = join(binDir, "capture.json");
+    const agentPath = join(binDir, "fake-copilot-acp.ts");
+    const sdkPath = join(process.cwd(), "node_modules/@agentclientprotocol/sdk/dist/acp.js");
+    const repoAgentsMd = "# Repo rules\nKeep this file as is.\n";
+    await Bun.write(join(cwd, "AGENTS.md"), repoAgentsMd);
+    await Bun.write(
+      agentPath,
+      `
+import { readFileSync } from "node:fs";
+import { join } from "node:path";
+import { Readable, Writable } from "node:stream";
+import { AgentSideConnection, PROTOCOL_VERSION, ndJsonStream } from "${sdkPath}";
+class FakeCopilot {
+  constructor(connection) { this.connection = connection; }
+  async initialize() {
+    return { protocolVersion: PROTOCOL_VERSION, agentCapabilities: { loadSession: false } };
+  }
+  async newSession() {
+    const dir = process.env.COPILOT_CUSTOM_INSTRUCTIONS_DIRS.split(",")[0];
+    await Bun.write(${JSON.stringify(captureFile)}, JSON.stringify({
+      argv: process.argv.slice(2),
+      dir,
+      instructions: readFileSync(join(dir, ${JSON.stringify(COPILOT_INSTRUCTIONS_FILE)}), "utf8"),
+      model: process.env.COPILOT_MODEL,
+    }));
+    return { sessionId: "copilot-session-1" };
+  }
+  async prompt(params) {
+    await this.connection.sessionUpdate({
+      sessionId: params.sessionId,
+      update: { sessionUpdate: "agent_message_chunk", content: { type: "text", text: "done" }, messageId: "m1" },
+    });
+    return { stopReason: "end_turn" };
+  }
+  async cancel() {}
+}
+const stream = ndJsonStream(Writable.toWeb(process.stdout), Readable.toWeb(process.stdin));
+new AgentSideConnection((connection) => new FakeCopilot(connection), stream);
+`,
+    );
+    const copilotStub = join(binDir, "copilot");
+    await Bun.write(copilotStub, `#!/bin/sh\nexec "${process.execPath}" "${agentPath}" "$@"\n`);
+    chmodSync(copilotStub, 0o755);
+    mkdirSync(join(binDir, "home"));
+
+    const { server, apiUrl } = await startTokenStubServer("copilot-token-id", "example-aseph_x");
+    try {
+      const config = baseConfig({
+        cwd,
+        apiUrl,
+        model: "openai/gpt-5.2",
+        systemPrompt: "You are a swarm worker.",
+        env: {
+          PATH: `${binDir}:${process.env.PATH ?? ""}`,
+          HOME: join(binDir, "home"),
+          ACP_TARGET: "copilot",
+        },
+      });
+      const session = await new ACPAdapter().createSession(config);
+      const result = await session.waitForCompletion();
+      expect(result.isError).toBe(false);
+
+      const captured = JSON.parse(await Bun.file(captureFile).text()) as {
+        argv: string[];
+        dir: string;
+        instructions: string;
+        model: string;
+      };
+      expect(captured.argv).toEqual(["--acp"]);
+      expect(captured.dir).toBe(copilotInstructionsDir(config));
+      expect(captured.dir.startsWith(cwd)).toBe(false);
+      expect(captured.instructions).toBe('---\napplyTo: "**"\n---\n\nYou are a swarm worker.\n');
+      expect(captured.model).toBe("openai/gpt-5.2");
+      expect(await Bun.file(join(cwd, "AGENTS.md")).text()).toBe(repoAgentsMd);
+
+      const deadline = Date.now() + 2_000;
+      while (existsSync(captured.dir) && Date.now() < deadline) await Bun.sleep(10);
+      expect(existsSync(captured.dir)).toBe(false);
+    } finally {
+      await new Promise<void>((resolve) => server.close(() => resolve()));
+    }
+  });
+
   test("custom target passes through only explicitly named env and supports a model env fallback", () => {
     const config = baseConfig({
       model: "custom-model",

package/src/tests/boot-scrub-sweep.test.ts

@@ -0,0 +1,390 @@
+import { afterAll, beforeAll, beforeEach, describe, expect, test } from "bun:test";
+import { unlink } from "node:fs/promises";
+import {
+  BOOT_SCRUB_TARGETS,
+  bootScrubCursorKey,
+  bootScrubDoneKey,
+  runBootScrubSweep,
+} from "../be/boot-scrub-sweep";
+import { closeDb, getDb, getDbClient, initDb, isSqliteVecAvailable } from "../be/db";
+import { getMemoryStore } from "../be/memory";
+import { contentSha256 } from "../commands/profile-sync";
+import {
+  refreshSecretScrubberCache,
+  registerVolatileSecret,
+  SCRUBBER_RULES_VERSION,
+} from "../utils/secret-scrubber";
+import { randomToken, type SyntheticSecret, syntheticSecret } from "./synthetic-secret-helpers";
+
+const TEST_DB_PATH = "./test-boot-scrub-sweep.sqlite";
+const AGENT_ID = "5c2b0000-0000-4000-8000-000000000001";
+// Distinctive non-secret words: FTS must still find the row by them after redaction.
+const MARKER_WORD = "quokkazebra";
+
+let known: SyntheticSecret;
+let envName: string;
+let envValue: string;
+let password: string;
+
+/** Text carrying all three secret forms: known value, env value and a `*_PASSWORD=` assignment. */
+function secretText(label: string): string {
+  return `${label} ${MARKER_WORD} token=${known.value} env=${envValue} X_DEMO_PASSWORD=${password} end`;
+}
+
+function expectRedacted(value: string | null | undefined, where: string): void {
+  expect(value, where).toBeString();
+  for (const secret of [known.value, envValue, password]) {
+    expect(value!.includes(secret), `${where} still holds a secret`).toBe(false);
+  }
+  expect(value!, where).toContain("[REDACTED:");
+}
+
+async function writeAudit(): Promise<{ tbl: string; id: string }[]> {
+  return getDbClient().query<{ tbl: string; id: string }>("SELECT tbl, id FROM scrub_write_audit");
+}
+
+async function ftsHits(term: string): Promise<string[]> {
+  const rows = await getDbClient().query<{ memory_id: string }>(
+    "SELECT memory_id FROM memory_fts WHERE memory_fts MATCH ?",
+    [`"${term}"`],
+  );
+  return rows.map((r) => r.memory_id);
+}
+
+const ids = {
+  log: crypto.randomUUID(),
+  logControl: crypto.randomUUID(),
+  task: crypto.randomUUID(),
+  taskControl: crypto.randomUUID(),
+  event: crypto.randomUUID(),
+  run: crypto.randomUUID(),
+  step: crypto.randomUUID(),
+  memory: "",
+};
+
+describe("boot retro-sweep", () => {
+  beforeAll(async () => {
+    for (const suffix of ["", "-wal", "-shm"]) {
+      try {
+        await unlink(TEST_DB_PATH + suffix);
+      } catch {}
+    }
+    initDb(TEST_DB_PATH);
+    // workflow_run_steps.runId references workflow_runs; the sweep never
+    // touches FKs, so seed the step without building a whole workflow.
+    getDb().run("PRAGMA foreign_keys = OFF");
+
+    known = syntheticSecret("sweep");
+    envName = "X_DEMO_SWEEP_API_KEY";
+    envValue = randomToken(32);
+    process.env[envName] = envValue;
+    refreshSecretScrubberCache();
+    password = randomToken(24);
+
+    const db = getDbClient();
+    // Every UPDATE on a swept table lands here: proves untouched rows saw no write.
+    await db.run("CREATE TABLE scrub_write_audit (tbl TEXT NOT NULL, id TEXT NOT NULL)");
+    for (const { table } of BOOT_SCRUB_TARGETS) {
+      await db.run(
+        `CREATE TRIGGER audit_${table} AFTER UPDATE ON ${table}
+         BEGIN INSERT INTO scrub_write_audit (tbl, id) VALUES ('${table}', NEW.id); END`,
+      );
+    }
+
+    const now = new Date().toISOString();
+    await db.run(
+      `INSERT INTO session_logs (id, sessionId, iteration, content, lineNumber, createdAt)
+       VALUES (?, 's1', 1, ?, 1, ?), (?, 's1', 1, ?, 2, ?)`,
+      [ids.log, secretText("log"), now, ids.logControl, "plain log line, nothing to hide", now],
+    );
+    await db.run(
+      `INSERT INTO agent_tasks (id, task, output, failureReason, progress, createdAt, lastUpdatedAt)
+       VALUES (?, ?, ?, ?, ?, ?, ?), (?, 'control task', 'control output', NULL, NULL, ?, ?)`,
+      [
+        ids.task,
+        secretText("task"),
+        secretText("output"),
+        secretText("failure"),
+        secretText("progress"),
+        now,
+        now,
+        ids.taskControl,
+        now,
+        now,
+      ],
+    );
+    await db.run(
+      `INSERT INTO events (id, category, event, source, data, createdAt)
+       VALUES (?, 'test', 'test.event', 'test', ?, ?)`,
+      [ids.event, JSON.stringify({ note: secretText("event"), nested: { n: 1 } }), now],
+    );
+    await db.run(
+      `INSERT INTO workflow_run_steps (id, runId, nodeId, nodeType, input, output, error, diagnostics)
+       VALUES (?, ?, 'n1', 'agent-task', ?, ?, ?, ?)`,
+      [
+        ids.step,
+        ids.run,
+        JSON.stringify({ prompt: secretText("input") }),
+        JSON.stringify({ result: secretText("step-output") }),
+        secretText("error"),
+        JSON.stringify({ trace: secretText("diagnostics") }),
+      ],
+    );
+
+    const store = getMemoryStore();
+    const legacy = {
+      name: `note ${known.value}`,
+      content: secretText("memory"),
+      summary: secretText("summary"),
+    };
+    const memory = await store.store({
+      agentId: AGENT_ID,
+      scope: "agent",
+      source: "manual",
+      ...legacy,
+    });
+    ids.memory = memory.id;
+    // store() now scrubs at write; put the raw text back so the row looks like
+    // one written before that, which is what the sweep exists for.
+    await db.run(
+      "UPDATE agent_memory SET name = ?, content = ?, summary = ?, contentHash = ? WHERE id = ?",
+      [legacy.name, legacy.content, legacy.summary, contentSha256(legacy.content), memory.id],
+    );
+    await db.run(
+      "UPDATE agent_memory_version SET content = ?, contentHash = ? WHERE memory_id = ?",
+      [legacy.content, contentSha256(legacy.content), memory.id],
+    );
+    await db.run("DELETE FROM memory_fts WHERE memory_id = ?", [memory.id]);
+    await db.run("INSERT INTO memory_fts(memory_id, name, content) VALUES (?, ?, ?)", [
+      memory.id,
+      legacy.name,
+      legacy.content,
+    ]);
+    const embedding = new Float32Array(512);
+    embedding[0] = 1;
+    await store.updateEmbedding(memory.id, embedding, "test");
+  });
+
+  // The test preload clears volatile secrets after every test.
+  beforeEach(() => registerVolatileSecret(known.value, known.name));
+
+  afterAll(async () => {
+    delete process.env[envName];
+    known.cleanup();
+    refreshSecretScrubberCache();
+    closeDb();
+    for (const suffix of ["", "-wal", "-shm"]) {
+      try {
+        await unlink(TEST_DB_PATH + suffix);
+      } catch {}
+    }
+  });
+
+  test("redacts every target column, keeps context, JSON and FTS, drops the stale vector", async () => {
+    const db = getDbClient();
+    // Positive controls: the secret is really in the index and the vector exists.
+    expect(await ftsHits(password)).toContain(ids.memory);
+    const before = await db.get<{ embedding: Uint8Array | null }>(
+      "SELECT embedding FROM agent_memory WHERE id = ?",
+      [ids.memory],
+    );
+    expect(before?.embedding).not.toBeNull();
+    if (isSqliteVecAvailable()) {
+      expect(
+        await db.get("SELECT memory_id FROM memory_vec WHERE memory_id = ?", [ids.memory]),
+      ).not.toBeNull();
+    }
+
+    const stats = await runBootScrubSweep({ version: SCRUBBER_RULES_VERSION, reembed: false });
+    expect(stats).not.toBeNull();
+    const byTable = Object.fromEntries(stats!.map((s) => [s.table, s]));
+    expect(byTable.session_logs!.changed).toBe(1);
+    expect(byTable.agent_tasks!.changed).toBe(1);
+    expect(byTable.agent_memory!.changed).toBe(1);
+    expect(byTable.agent_memory_version!.changed).toBe(1);
+    expect(byTable.events!.changed).toBe(1);
+    expect(byTable.workflow_run_steps!.changed).toBe(1);
+
+    const log = await db.get<{ content: string }>("SELECT content FROM session_logs WHERE id = ?", [
+      ids.log,
+    ]);
+    expectRedacted(log?.content, "session_logs.content");
+    expect(log!.content).toContain(MARKER_WORD);
+
+    const task = await db.get<Record<string, string>>(
+      "SELECT task, output, failureReason, progress FROM agent_tasks WHERE id = ?",
+      [ids.task],
+    );
+    for (const col of ["task", "output", "failureReason", "progress"]) {
+      expectRedacted(task?.[col], `agent_tasks.${col}`);
+    }
+
+    const mem = await db.get<{
+      name: string;
+      content: string;
+      summary: string;
+      contentHash: string;
+      embedding: Uint8Array | null;
+    }>("SELECT name, content, summary, contentHash, embedding FROM agent_memory WHERE id = ?", [
+      ids.memory,
+    ]);
+    expectRedacted(mem?.name, "agent_memory.name");
+    expectRedacted(mem?.content, "agent_memory.content");
+    expectRedacted(mem?.summary, "agent_memory.summary");
+    expect(mem!.contentHash).toBe(contentSha256(mem!.content));
+    expect(mem!.embedding).toBeNull();
+    if (isSqliteVecAvailable()) {
+      expect(
+        await db.get("SELECT memory_id FROM memory_vec WHERE memory_id = ?", [ids.memory]),
+      ).toBeNull();
+    }
+    // FTS was rewritten from the scrubbed text, not left stale.
+    expect(await ftsHits(MARKER_WORD)).toContain(ids.memory);
+    expect(await ftsHits(password)).not.toContain(ids.memory);
+
+    const version = await db.get<{ content: string; contentHash: string }>(
+      "SELECT content, contentHash FROM agent_memory_version WHERE memory_id = ?",
+      [ids.memory],
+    );
+    expectRedacted(version?.content, "agent_memory_version.content");
+    expect(version!.contentHash).toBe(contentSha256(version!.content));
+
+    const event = await db.get<{ data: string }>("SELECT data FROM events WHERE id = ?", [
+      ids.event,
+    ]);
+    expectRedacted(event?.data, "events.data");
+    const parsed = JSON.parse(event!.data) as { note: string; nested: { n: number } };
+    expect(parsed.nested.n).toBe(1);
+    expect(parsed.note).toContain(MARKER_WORD);
+
+    const step = await db.get<Record<string, string>>(
+      "SELECT input, output, error, diagnostics FROM workflow_run_steps WHERE id = ?",
+      [ids.step],
+    );
+    for (const col of ["input", "output", "error", "diagnostics"]) {
+      expectRedacted(step?.[col], `workflow_run_steps.${col}`);
+    }
+    for (const col of ["input", "output", "diagnostics"]) {
+      expect(() => JSON.parse(step![col]!)).not.toThrow();
+    }
+  });
+
+  test("control rows without a secret see no write", async () => {
+    const audit = await writeAudit();
+    // Positive control: the trigger fired for the rows that held a secret.
+    expect(audit).toContainEqual({ tbl: "session_logs", id: ids.log });
+    expect(audit).toContainEqual({ tbl: "agent_tasks", id: ids.task });
+    // The controls were never written.
+    expect(audit.map((a) => a.id)).not.toContain(ids.logControl);
+    expect(audit.map((a) => a.id)).not.toContain(ids.taskControl);
+    const control = await getDbClient().get<{ content: string }>(
+      "SELECT content FROM session_logs WHERE id = ?",
+      [ids.logControl],
+    );
+    expect(control?.content).toBe("plain log line, nothing to hide");
+  });
+
+  test("marks the version done and is a no-op on rerun; a new version runs again", async () => {
+    const db = getDbClient();
+    expect(
+      await db.get("SELECT key FROM seed_state WHERE kind = 'maintenance' AND key = ?", [
+        bootScrubDoneKey(SCRUBBER_RULES_VERSION),
+      ]),
+    ).not.toBeNull();
+    const cursors = await db.query("SELECT key FROM seed_state WHERE key LIKE 'boot-scrub-v%:%'");
+    expect(cursors).toHaveLength(0);
+
+    const auditBefore = (await writeAudit()).length;
+    expect(await runBootScrubSweep({ version: SCRUBBER_RULES_VERSION, reembed: false })).toBeNull();
+
+    const bumped = SCRUBBER_RULES_VERSION + 100;
+    const stats = await runBootScrubSweep({ version: bumped, reembed: false });
+    expect(stats).not.toBeNull();
+    expect(stats!.every((s) => s.changed === 0)).toBe(true);
+    expect(stats!.find((s) => s.table === "session_logs")!.scanned).toBeGreaterThanOrEqual(2);
+    expect((await writeAudit()).length).toBe(auditBefore);
+    expect(
+      await db.get("SELECT key FROM seed_state WHERE kind = 'maintenance' AND key = ?", [
+        bootScrubDoneKey(bumped),
+      ]),
+    ).not.toBeNull();
+  });
+
+  test("an interrupted run resumes from the cursor without rescanning batch 1", async () => {
+    const db = getDbClient();
+    const now = new Date().toISOString();
+    for (let i = 0; i < 5; i++) {
+      await db.run(
+        `INSERT INTO session_logs (id, sessionId, iteration, content, lineNumber, createdAt)
+         VALUES (?, 's2', 1, ?, ?, ?)`,
+        [crypto.randomUUID(), secretText(`resume-${i}`), 10 + i, now],
+      );
+    }
+    const version = SCRUBBER_RULES_VERSION + 200;
+    let firstBatch: string[] = [];
+    await expect(
+      runBootScrubSweep({
+        version,
+        batchSize: 2,
+        reembed: false,
+        afterBatch: ({ table, ids: batchIds }) => {
+          if (table !== "session_logs") return;
+          firstBatch = batchIds;
+          throw new Error("interrupted after batch 1");
+        },
+      }),
+    ).rejects.toThrow("interrupted after batch 1");
+    expect(firstBatch).toHaveLength(2);
+
+    const cursor = await db.get<{ seededHash: string }>(
+      "SELECT seededHash FROM seed_state WHERE kind = 'maintenance' AND key = ?",
+      [bootScrubCursorKey(version, "session_logs")],
+    );
+    expect(cursor?.seededHash).toBe(firstBatch[1]);
+    expect(
+      await db.get("SELECT key FROM seed_state WHERE key = ?", [bootScrubDoneKey(version)]),
+    ).toBeNull();
+
+    const resumed: string[] = [];
+    const stats = await runBootScrubSweep({
+      version,
+      batchSize: 2,
+      reembed: false,
+      afterBatch: ({ table, ids: batchIds }) => {
+        if (table === "session_logs") resumed.push(...batchIds);
+      },
+    });
+    expect(resumed.length).toBeGreaterThan(0);
+    for (const id of firstBatch) expect(resumed).not.toContain(id);
+    const total = await db.get<{ n: number }>("SELECT COUNT(*) AS n FROM session_logs");
+    expect(stats!.find((s) => s.table === "session_logs")!.scanned).toBe(total!.n - 2);
+
+    // Batch 1 was committed before the interrupt, the rest by the resumed run.
+    const remaining = await db.query<{ content: string }>(
+      "SELECT content FROM session_logs WHERE sessionId = 's2'",
+    );
+    for (const row of remaining) expectRedacted(row.content, "resumed session_logs.content");
+  });
+
+  test("skips a row whose scrub would break valid JSON", async () => {
+    const db = getDbClient();
+    const id = crypto.randomUUID();
+    const data = JSON.stringify({ msg: "BREAKJSON here" });
+    await db.run(
+      `INSERT INTO events (id, category, event, source, data, createdAt)
+       VALUES (?, 'test', 'test.event', 'test', ?, ?)`,
+      [id, data, new Date().toISOString()],
+    );
+    // Force a scrub result that is not valid JSON (an unbalanced quote).
+    const stats = await runBootScrubSweep({
+      version: SCRUBBER_RULES_VERSION + 300,
+      reembed: false,
+      scrub: (text) => text.replace("BREAKJSON", '"'),
+    });
+    const events = stats!.find((s) => s.table === "events")!;
+    expect(events.skippedInvalidJson).toBe(1);
+    expect(events.changed).toBe(0);
+    const row = await db.get<{ data: string }>("SELECT data FROM events WHERE id = ?", [id]);
+    expect(row?.data).toBe(data);
+  });
+});

package/src/tests/budget-user-scope.test.ts

@@ -181,23 +181,15 @@ async function callPoll(agentId: string): Promise<{
 }> {
   let status = 200;
   let bodyStr = "";
-  const headers: Record<string, string> = {};
-
   const req = {
     method: "GET",
     url: "/api/poll",
     headers: { "x-agent-id": agentId },
   } as unknown as Parameters<typeof handlePoll>[0];
-
   const res = {
-    setHeader(name: string, value: string) {
-      headers[name.toLowerCase()] = value;
-    },
-    writeHead(code: number, h?: Record<string, string>) {
+    setHeader() {},
+    writeHead(code: number) {
       status = code;
-      if (h) {
-        for (const [k, v] of Object.entries(h)) headers[k.toLowerCase()] = v;
-      }
     },
     end(body?: string) {
       bodyStr = body ?? "";
@@ -288,17 +280,11 @@ describe("user budget scope", () => {
     expect(result.allowed).toBe(true);
   });
 
-  test("/mcp-user task is refused at worker admission when user budget is spent", async () => {
+  test("/mcp-user task is refused when the Lead polls after the user budget is spent", async () => {
     const server = createMcpUserTestServer();
     const port = await listenOnFreePort(server, "127.0.0.1");
     try {
       const lead = await createAgent({ name: "lead", isLead: true, status: "idle", maxTasks: 1 });
-      const worker = await createAgent({
-        name: "worker",
-        isLead: false,
-        status: "idle",
-        maxTasks: 1,
-      });
       const user = await createUser({ name: "MCP Budget User", dailyBudgetUsd: 0.5 });
       await upsertBudget("user", user.id, 0.5);
       const token = await mintToken(user.id, "qa", ACTOR);
@@ -329,21 +315,21 @@ describe("user budget scope", () => {
       await createSessionCost({
         sessionId: `sess-${crypto.randomUUID()}`,
         taskId,
-        agentId: worker.id,
+        agentId: lead.id,
         totalCostUsd: 0.5,
         durationMs: 1000,
         numTurns: 1,
         model: "test-model",
       });
 
-      const firstPoll = await callPoll(worker.id);
+      const firstPoll = await callPoll(lead.id);
       expect(firstPoll.status).toBe(200);
       if ("error" in firstPoll.body) throw new Error("unexpected poll error");
       expect(firstPoll.body.trigger?.type).toBe("budget_refused");
       expect((firstPoll.body.trigger as { cause: string }).cause).toBe("user");
       expect((firstPoll.body.trigger as { userSpend: number }).userSpend).toBe(0.5);
       expect((firstPoll.body.trigger as { userBudget: number }).userBudget).toBe(0.5);
-      expect((await getTaskById(taskId))?.status).toBe("unassigned");
+      expect((await getTaskById(taskId))?.status).toBe("pending");
 
       const firstDedup = await getDbClient().get<{
         follow_up_task_id: string | null;
@@ -357,7 +343,7 @@ describe("user budget scope", () => {
       const firstFollowUpId = firstDedup?.follow_up_task_id;
       expect(firstFollowUpId ? (await getTaskById(firstFollowUpId))?.agentId : null).toBe(lead.id);
 
-      const secondPoll = await callPoll(worker.id);
+      const secondPoll = await callPoll(lead.id);
       expect(secondPoll.status).toBe(200);
       if ("error" in secondPoll.body) throw new Error("unexpected poll error");
       expect(secondPoll.body.trigger?.type).toBe("budget_refused");