MailUI4Agents/plugins/pi-mail-bridge/test/pool.test.mjs

/**
 * 工作进程池的调度不变量。
 *
 * 这些用例**真的 fork 子进程**，用一个极小的桩 worker（不装 pi SDK）：
 * 要验的是调度（并发上限、同会话串行、排队、硬超时、状态跨 worker 存活），
 * 而不是模型怎么跑。用桩让每个用例在几百毫秒内完成。
 *
 * 桩 worker 的协议与真 worker 一致：`ready` → 收 `job` → 按 `__hold` 停一会儿
 * → 发 `done`。它还把收到的 job 载荷原样回声成一行 `JOB {…}` 日志 ——
 * 判据因此能落在「主进程真的把 sessionFile / grants / 命名指纹传下来了」上，
 * 而不是一个间接的计数（计数在字段被丢掉时依然会给出绿色）。
 */

import { test } from "node:test";
import assert from "node:assert/strict";
import { mkdtempSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";

// ─── 桩 worker ───
//
// 落到临时目录而不是仓库里：它是测试脚手架，不该被 check-shared-libs 之类的
// 一致性脚本看到，也不该让人误以为是第二个真 worker。
const STUB_DIR = mkdtempSync(join(tmpdir(), "pi-pool-test-"));
const STUB = join(STUB_DIR, "stub-worker.mjs");
writeFileSync(
	STUB,
	`
process.on('message', (msg) => {
  if (msg?.type === 'job') {
    const hold = msg.data?.__hold ?? 30;
    process.send({ type: 'log', line: 'JOB ' + JSON.stringify({
      mailID: msg.data.mail_id,
      kind: msg.kind,
      sessionFile: msg.session?.sessionFile || '',
      cwd: msg.session?.cwd || '',
      grants: msg.grants || [],
      lastSyncedName: msg.lastSyncedName || '',
      turnTimeoutMs: msg.config?.turnTimeoutMs ?? null,
    }) });
    process.send({ type: 'session_opened', piSessionId: 'pi-' + msg.data.mail_id,
                   sessionFile: '/tmp/f-' + msg.data.mail_id + '.jsonl', cwd: '/tmp', reused: false });
    if (msg.data?.__grant) {
      process.send({ type: 'permission_grant', toolName: msg.data.__grant });
    }
    if (msg.data?.__name) {
      process.send({ type: 'name_synced', signature: msg.data.__name });
    }
    if (msg.data?.__reopen) {
      // 模型降级换会话：同一个 worker 里第二次 session_opened
      process.send({ type: 'session_opened', piSessionId: msg.data.__reopen,
                     sessionFile: '/tmp/f2.jsonl', cwd: '/tmp', reused: false });
    }
    if (msg.data?.__pending) {
      process.send({ type: 'permission_pending', relayKey: msg.data.__pending });
      return;                                   // 等决策，见下面的分支
    }
    if (msg.data?.__crash) {
      // 未回报 done 就退出：验证主进程会重投（而不是静默丢信）
      process.exit(1);
    }
    // 每 40ms 报一次心跳：并发的判据必须是「两个进程真的同时在干活」，
    // 而不是「running map 里有两个条目」—— fork 返回后立即就有两个条目了。
    const beat = setInterval(() => process.send({ type: 'log', line: 'TICK ' + msg.data.mail_id }), 40);
    if (hold < 0) return;                       // 永不结束，用来验硬超时
    setTimeout(() => {
      clearInterval(beat);
      process.send({ type: 'done', ok: true, error: '' });
      setTimeout(() => process.exit(0), 20);
    }, hold);
    return;
  }
  if (msg?.type === 'permission_decision') {
    process.send({ type: 'log', line: 'DECISION ' + msg.relayKey + '=' + msg.decision });
    process.send({ type: 'done', ok: true, error: '' });
    setTimeout(() => process.exit(0), 20);
    return;
  }
  if (msg?.type === 'shutdown') {
    process.send({ type: 'log', line: 'SHUTDOWN' });
    setTimeout(() => process.exit(0), 10);
  }
});
process.send({ type: 'ready' });
`,
);

const { createWorkerPool } = await import("../src/pool.mjs");

/** 建一个用桩 worker 的池。 */
function makePool(opts = {}) {
	const lines = [];
	const pool = createWorkerPool({
		log: (...a) => lines.push(a.join(" ")),
		config: () => ({ turnTimeoutMs: 1000, ...(opts.config || {}) }),
		onReconfigure: opts.onReconfigure || (() => {}),
		maxWorkers: opts.maxWorkers ?? 2,
		workerMaxMs: opts.workerMaxMs ?? 5000,
		maxAttempts: opts.maxAttempts,
		workerPath: opts.workerPath || STUB,
	});
	return { pool, lines };
}

const sleep = (ms) => new Promise((r) => setTimeout(r, ms));

/** 轮询到条件成立或超时 —— 比固定 sleep 稳。 */
async function until(fn, timeoutMs = 4000) {
	const t0 = Date.now();
	while (Date.now() - t0 < timeoutMs) {
		if (fn()) return true;
		await sleep(20);
	}
	return false;
}

/** 从日志行里取出桩 worker 回声的 job 载荷。 */
function jobs(lines) {
	const out = [];
	for (const l of lines) {
		const i = l.indexOf("JOB ");
		if (i === -1) continue;
		try {
			out.push(JSON.parse(l.slice(i + 4)));
		} catch {
			/* 不是完整一行 */
		}
	}
	return out;
}

/** 某个 mailID 的心跳出现过几次。 */
const ticks = (lines, id) =>
	lines.filter((l) => l.includes(`TICK ${id}`)).length;

test("并发上限被遵守：第三条会话要等前面空出来", async () => {
	const { pool } = makePool({ maxWorkers: 2 });
	for (const id of ["a", "b", "c"]) {
		pool.submit("mail", { mail_id: id, session_id: `S-${id}`, __hold: 250 });
	}

	let peak = 0;
	const t = setInterval(() => {
		peak = Math.max(peak, pool.stats().running);
	}, 15);
	const sawQueue = await until(() => pool.stats().queued > 0, 1000);
	await until(() => pool.stats().running === 0 && pool.stats().queued === 0);
	clearInterval(t);
	pool.stop();

	assert.ok(peak <= 2, `同时跑的 worker 峰值 ${peak}，不该超过 maxWorkers=2`);
	assert.ok(sawQueue, "满载时第三封该进队列而不是被丢掉");
});

test("同一会话串行：两个 worker 的心跳不得重叠", async () => {
	const { pool, lines } = makePool({ maxWorkers: 3 });
	pool.submit("mail", { mail_id: "m1", session_id: "SAME", __hold: 220 });
	pool.submit("mail", { mail_id: "m2", session_id: "SAME", __hold: 60 });

	// 判据一：任一时刻只有一个 worker 在跑。
	let everTwo = false;
	const t = setInterval(() => {
		if (pool.stats().running > 1) everTwo = true;
	}, 10);
	await until(
		() => jobs(lines).length === 2 && pool.stats().running === 0,
		5000,
	);
	clearInterval(t);
	pool.stop();

	assert.equal(everTwo, false, "同一条会话不得有两个 worker 同时装载会话文件");
	// 判据二：m2 一次心跳都没能在 m1 结束前发出 —— m1 的心跳数应当远多于 m2。
	assert.ok(
		ticks(lines, "m1") >= 3,
		`m1 该跑满 220ms，实际心跳 ${ticks(lines, "m1")} 次`,
	);
	assert.equal(jobs(lines).length, 2, "两封都要被处理，不能因为串行而丢掉一封");
});

test("不同会话真并发：两个进程的心跳在同一段时间里交错", async () => {
	const { pool, lines } = makePool({ maxWorkers: 3 });
	pool.submit("mail", { mail_id: "p", session_id: "S-P", __hold: 300 });
	pool.submit("mail", { mail_id: "q", session_id: "S-Q", __hold: 300 });

	const interleaved = await until(
		() => ticks(lines, "p") >= 2 && ticks(lines, "q") >= 2,
		2500,
	);
	pool.stop();
	assert.ok(
		interleaved,
		`两条不同会话应当并发，实际 p=${ticks(lines, "p")} q=${ticks(lines, "q")} 次心跳`,
	);
});

test("sessionFile 与 cwd 跨 worker 传下去：第二封接着第一封的会话谈", async () => {
	const { pool, lines } = makePool({ maxWorkers: 2 });
	pool.submit("mail", { mail_id: "first", session_id: "KEEP", __hold: 30 });
	await until(() => jobs(lines).length === 1 && pool.stats().running === 0);

	pool.submit("mail", { mail_id: "second", session_id: "KEEP", __hold: 30 });
	await until(() => jobs(lines).length === 2 && pool.stats().running === 0);
	pool.stop();

	const [j1, j2] = jobs(lines);
	assert.equal(j1.sessionFile, "", "第一封时还没有会话文件");
	assert.equal(
		j2.sessionFile,
		"/tmp/f-first.jsonl",
		"第二封必须带上第一封开出来的会话文件，否则每封邮件都从零开始、上下文全丢",
	);
	assert.equal(j2.cwd, "/tmp", "cwd 也要传下去");
});

test("「一直同意」与命名指纹跨 worker 存活", async () => {
	const { pool, lines } = makePool({ maxWorkers: 2 });
	pool.submit("mail", {
		mail_id: "g1",
		session_id: "GRANT",
		__hold: 30,
		__grant: "bash",
		__name: "platform:某名字|某名字",
	});
	await until(() => jobs(lines).length === 1 && pool.stats().running === 0);

	pool.submit("mail", { mail_id: "g2", session_id: "GRANT", __hold: 30 });
	await until(() => jobs(lines).length === 2 && pool.stats().running === 0);
	pool.stop();

	const [j1, j2] = jobs(lines);
	assert.deepEqual(j1.grants, [], "第一封时还没人点过「一直同意」");
	assert.deepEqual(
		j2.grants,
		["bash"],
		"「一直同意」不跨 worker 存活的话，下一封邮件又问一遍 —— 那个选项就是在骗人",
	);
	assert.equal(
		j2.lastSyncedName,
		"platform:某名字|某名字",
		"命名指纹要传下去，否则每封邮件都重新 sync 一次",
	);
});

test("config() 每次派活时重取：allowedModels 随心跳变，不能用快照", async () => {
	let turnTimeoutMs = 111;
	const { pool, lines } = makePool({ maxWorkers: 1, config: {} });
	// makePool 的 config 是固定值，这里换成动态的
	pool.stop();

	const lines2 = [];
	const p2 = createWorkerPool({
		log: (...a) => lines2.push(a.join(" ")),
		config: () => ({ turnTimeoutMs }),
		onReconfigure: () => {},
		maxWorkers: 1,
		workerMaxMs: 5000,
		workerPath: STUB,
	});
	p2.submit("mail", { mail_id: "c1", session_id: "C1", __hold: 20 });
	await until(() => jobs(lines2).length === 1 && p2.stats().running === 0);
	turnTimeoutMs = 222;
	p2.submit("mail", { mail_id: "c2", session_id: "C2", __hold: 20 });
	await until(() => jobs(lines2).length === 2 && p2.stats().running === 0);
	p2.stop();

	const [j1, j2] = jobs(lines2);
	assert.equal(j1.turnTimeoutMs, 111);
	assert.equal(
		j2.turnTimeoutMs,
		222,
		"config() 必须每次重取，否则 worker 用的是上一轮的模型范围",
	);
	assert.equal(lines.length >= 0, true);
});

test("硬超时回收卡死的 worker，且不堵住同会话后续邮件", async () => {
	const { pool, lines } = makePool({ maxWorkers: 2, workerMaxMs: 400 });
	pool.submit("mail", { mail_id: "stuck", session_id: "STUCK", __hold: -1 });
	await until(() => pool.stats().running === 1, 1500);

	const freed = await until(() => pool.stats().running === 0, 3000);
	assert.ok(
		freed,
		"卡死的 worker 必须被硬超时回收，否则那条会话的后续邮件永远排队",
	);
	assert.ok(
		lines.some((l) => l.includes("强杀")),
		`应记下强杀日志，实际:\n${lines.join("\n")}`,
	);

	pool.submit("mail", { mail_id: "after", session_id: "STUCK", __hold: 30 });
	const ran = await until(
		() => jobs(lines).some((j) => j.mailID === "after"),
		2000,
	);
	await until(() => pool.stats().running === 0);
	pool.stop();
	assert.ok(ran, "硬超时后同一会话的后续邮件必须能被处理");
});

test("权限决策路由到发起询问的那个 worker", async () => {
	const { pool, lines } = makePool({ maxWorkers: 2 });
	pool.submit("mail", {
		mail_id: "perm",
		session_id: "PERM",
		__pending: "rk-1",
	});
	await until(() => pool.stats().running === 1, 1500);
	await sleep(150); // 等 permission_pending 到主进程

	assert.equal(pool.routePermission("rk-1", "同意"), true, "应当路由成功");
	await until(() => pool.stats().running === 0, 2000);
	pool.stop();

	assert.ok(
		lines.some((l) => l.includes("DECISION rk-1=同意")),
		`worker 应收到决策原文，实际:\n${lines.join("\n")}`,
	);
});

test("决策发的是选项原文而不是归一化的 allow/deny", async () => {
	const { pool, lines } = makePool({ maxWorkers: 2 });
	pool.submit("mail", { mail_id: "p2", session_id: "P2", __pending: "rk-2" });
	await until(() => pool.stats().running === 1, 1500);
	await sleep(150);
	pool.routePermission("rk-2", "一直同意");
	await until(() => pool.stats().running === 0, 2000);
	pool.stop();

	// 「同意」与「一直同意」语义不同，归一化会让后者退化成单次授权
	assert.ok(
		lines.some((l) => l.includes("DECISION rk-2=一直同意")),
		`必须原文透传，实际:\n${lines.join("\n")}`,
	);
});

test("决策找不到 worker 时返回 false（调用方据此走 B-4.2 降级）", async () => {
	const { pool } = makePool();
	assert.equal(pool.routePermission("never-seen", "同意"), false);
	pool.stop();
});

test("worker 退出后它的权限路由被清掉，不会误投给下一个 worker", async () => {
	const { pool } = makePool({ maxWorkers: 2 });
	pool.submit("mail", {
		mail_id: "gone",
		session_id: "GONE",
		__pending: "rk-gone",
	});
	await until(() => pool.stats().running === 1, 1500);
	await sleep(150);
	// 不给决策，直接停掉它
	pool.stop();
	await until(() => pool.stats().running === 0, 4000);

	assert.equal(
		pool.routePermission("rk-gone", "同意"),
		false,
		"worker 已退出，路由必须返回 false 让调用方走降级路径",
	);
});

test("mailDrivenIDs 报出所有跑过的 pi 会话，且不随 worker 退出而清", async () => {
	const { pool, lines } = makePool({ maxWorkers: 2 });
	pool.submit("mail", { mail_id: "d1", session_id: "D1", __hold: 30 });
	pool.submit("mail", { mail_id: "d2", session_id: "D2", __hold: 30 });
	await until(() => jobs(lines).length === 2 && pool.stats().running === 0);
	pool.stop();

	const ids = pool.mailDrivenIDs();
	assert.ok(ids.has("pi-d1"), "D1 的 pi 会话该被标记为邮件驱动");
	assert.ok(ids.has("pi-d2"), "D2 的 pi 会话该被标记为邮件驱动");
});

test("模型降级换掉的旧 pi 会话仍算邮件驱动", async () => {
	const { pool, lines } = makePool({ maxWorkers: 2 });
	pool.submit("mail", {
		mail_id: "r",
		session_id: "RETIRE",
		__hold: 40,
		__reopen: "pi-new",
	});
	await until(() => jobs(lines).length === 1 && pool.stats().running === 0);
	pool.stop();

	const ids = pool.mailDrivenIDs();
	assert.ok(ids.has("pi-new"), "新会话要在");
	assert.ok(
		ids.has("pi-r"),
		"被换掉的旧会话也参与过邮件往来，磁盘上的文件还在，快照该报它",
	);
});

test("hasSession 只对跑过的邮件会话为真", async () => {
	const { pool } = makePool();
	assert.equal(pool.hasSession("NOPE"), false);
	pool.submit("mail", { mail_id: "h1", session_id: "HAS", __hold: 30 });
	await until(() => pool.hasSession("HAS"), 2000);
	await until(() => pool.stats().running === 0);
	pool.stop();
	assert.equal(pool.hasSession("HAS"), true, "worker 退出后仍该记着这条会话");
});

test("kind 透传：权限通知走 permission 而不是 mail", async () => {
	const { pool, lines } = makePool({ maxWorkers: 2 });
	pool.submit("permission", { mail_id: "k1", session_id: "K1", __hold: 20 });
	await until(() => jobs(lines).length === 1 && pool.stats().running === 0);
	pool.stop();
	assert.equal(
		jobs(lines)[0].kind,
		"permission",
		"kind 决定 worker 用哪套提示词，传错会让模型以为收到一封新邮件",
	);
});

test("stop 之后不再派活", async () => {
	const { pool } = makePool();
	pool.stop();
	pool.submit("mail", { mail_id: "late", session_id: "LATE", __hold: 30 });
	await sleep(200);
	assert.equal(pool.stats().running, 0, "关停后不该再起 worker");
	assert.equal(pool.stats().queued, 0, "关停后队列应为空");
});

test("stop 会先给 worker 发 shutdown（让它 fail closed）再杀", async () => {
	const { pool, lines } = makePool({ maxWorkers: 2 });
	pool.submit("mail", { mail_id: "s1", session_id: "S1", __hold: -1 });
	await until(() => pool.stats().running === 1, 1500);
	pool.stop();
	const gotShutdown = await until(
		() => lines.some((l) => l.includes("SHUTDOWN")),
		2000,
	);
	assert.ok(
		gotShutdown,
		"必须先发 shutdown：直接 SIGKILL 会让 pi 侧那些等权限的 await 永不返回",
	);
});

test("没有 session_id 的事件各占一个 key，不会互相串行", async () => {
	const { pool, lines } = makePool({ maxWorkers: 3 });
	pool.submit("mail", { mail_id: "n1", __hold: 300 });
	pool.submit("mail", { mail_id: "n2", __hold: 300 });
	const both = await until(
		() => ticks(lines, "n1") >= 2 && ticks(lines, "n2") >= 2,
		2500,
	);
	pool.stop();
	assert.ok(both, "无 session_id 的两封不属于同一条会话，应能并发");
});

test("reconfigure 上报被转达给主进程", async () => {
	const STUB2 = join(STUB_DIR, "stub-reconf.mjs");
	writeFileSync(
		STUB2,
		`
process.on('message', (msg) => {
  if (msg?.type === 'job') {
    process.send({ type: 'reconfigure', url: 'http://new:9999', agentKey: 'k2' });
    process.send({ type: 'done', ok: true, error: '' });
    setTimeout(() => process.exit(0), 20);
  }
});
process.send({ type: 'ready' });
`,
	);
	let got = null;
	const { pool } = makePool({
		maxWorkers: 1,
		workerPath: STUB2,
		onReconfigure: (url, key) => {
			got = { url, key };
		},
	});
	pool.submit("mail", { mail_id: "r1", session_id: "R1" });
	await until(() => got !== null, 3000);
	pool.stop();
	assert.deepEqual(
		got,
		{ url: "http://new:9999", key: "k2" },
		"worker 里 connect_to_server 换的坐标必须回到主进程 —— worker 马上就退了，改在它自己身上等于没改",
	);
});

test("worker 未回报 done 就退出：有界重投而不是静默丢信", async () => {
	// maxAttempts=2：首次 + 一次重投，然后放弃。
	// 这封邮件必定崩溃，重投就是在验证「有界」——不然它会变成永久活锁。
	const { pool, lines } = makePool({ maxWorkers: 1, maxAttempts: 2 });
	pool.submit("mail", {
		mail_id: "crashy",
		session_id: "CRASH",
		__crash: true,
	});

	const gaveUp = await until(() => lines.some((l) => l.includes("放弃")), 6000);
	pool.stop();

	assert.ok(gaveUp, `重投到上限后应记下「放弃」，实际:\n${lines.join("\n")}`);
	assert.equal(
		jobs(lines).length,
		2,
		`应当尝试 2 次（首次 + 1 次重投），实际 ${jobs(lines).length} 次`,
	);
	assert.ok(
		lines.some((l) => l.includes("未回报 done 就退出")),
		"必须明说是「未回报 done 就退出」——否则看到 exit code 会误以为是普通崩溃",
	);
});

test("重投上限之下不会无限重投（maxAttempts=1 就是不重投）", async () => {
	const { pool, lines } = makePool({ maxWorkers: 1, maxAttempts: 1 });
	pool.submit("mail", { mail_id: "once", session_id: "ONCE", __crash: true });
	await until(() => lines.some((l) => l.includes("放弃")), 4000);
	await sleep(300); // 再等一会儿，确认没有额外重投
	pool.stop();

	assert.equal(jobs(lines).length, 1, "maxAttempts=1 时只跑一次");
});