/** * 工作进程池的调度不变量。 * * 这些用例**真的 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; // 等决策,见下面的分支 } // 每 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, 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 马上就退了,改在它自己身上等于没改'); });