feat: add Pi process and RPC foundation
This commit is contained in:
brother7 committed
1 parent
b6f693048d
commit
9a31dacb2a
14 files changed
+1955
-26
No files matched your search
@@ -0,0 +1,271 @@
|
||||
// @vitest-environment node
|
||||
|
||||
import { afterEach, describe, expect, it } from 'vitest';
|
||||
import { mkdir, mkdtemp, rm } from 'node:fs/promises';
|
||||
import { tmpdir } from 'node:os';
|
||||
import { join, resolve } from 'node:path';
|
||||
import { Writable } from 'node:stream';
|
||||
import { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
|
||||
import { PiRpcClient } from '../../electron/coding-runtime/pi/rpc-client';
|
||||
import { StrictLfJsonlFramer } from '../../electron/coding-runtime/pi/rpc-framer';
|
||||
import {
|
||||
PiWorkerProcess,
|
||||
buildPiRpcArgs,
|
||||
sanitizePiDiagnostic,
|
||||
} from '../../electron/coding-runtime/pi/worker-process';
|
||||
|
||||
const fakeChildPath = resolve('tests/fixtures/fake-pi-rpc-child.mjs');
|
||||
const scratchRoots: string[] = [];
|
||||
const workers: PiWorkerProcess[] = [];
|
||||
|
||||
async function makeWorker(
|
||||
options: Partial<ConstructorParameters<typeof PiWorkerProcess>[0]> = {},
|
||||
): Promise<PiWorkerProcess> {
|
||||
const root = await mkdtemp(join(tmpdir(), 'makelore-pi-rpc-test-'));
|
||||
scratchRoots.push(root);
|
||||
const configDir = join(root, 'config');
|
||||
const sessionDir = join(root, 'sessions');
|
||||
const cwd = join(root, 'project');
|
||||
await Promise.all([
|
||||
mkdir(configDir),
|
||||
mkdir(sessionDir),
|
||||
mkdir(cwd),
|
||||
]);
|
||||
const worker = new PiWorkerProcess({
|
||||
executablePath: process.execPath,
|
||||
cliPath: fakeChildPath,
|
||||
cwd,
|
||||
configDir,
|
||||
sessionDir,
|
||||
commandTimeoutMs: 1_000,
|
||||
shutdownGraceMs: 500,
|
||||
...options,
|
||||
});
|
||||
workers.push(worker);
|
||||
return await worker.start();
|
||||
}
|
||||
|
||||
async function processAlive(pid: number): Promise<boolean> {
|
||||
try {
|
||||
process.kill(pid, 0);
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(workers.splice(0).map((worker) => worker.stop().catch(() => undefined)));
|
||||
await Promise.all(scratchRoots.splice(0).map((root) => rm(root, {
|
||||
recursive: true,
|
||||
force: true,
|
||||
maxRetries: 3,
|
||||
})));
|
||||
});
|
||||
|
||||
describe('strict Pi LF JSONL framing', () => {
|
||||
it('handles chunk boundaries, multiple records, CRLF, and Unicode separators', () => {
|
||||
const records: unknown[] = [];
|
||||
const framer = new StrictLfJsonlFramer({ onRecord: (record) => records.push(record) });
|
||||
const source = Buffer.from(
|
||||
`${JSON.stringify({ text: 'left\u2028middle\u2029right' })}\n${JSON.stringify({ ok: true })}\r\n`,
|
||||
);
|
||||
|
||||
framer.push(source.subarray(0, 8));
|
||||
framer.push(source.subarray(8, 23));
|
||||
framer.push(source.subarray(23));
|
||||
framer.finish();
|
||||
|
||||
expect(records).toEqual([
|
||||
{ text: 'left\u2028middle\u2029right' },
|
||||
{ ok: true },
|
||||
]);
|
||||
});
|
||||
|
||||
it('fails closed for malformed, blank, invalid UTF-8, oversized, and partial records', () => {
|
||||
const make = (maxLineBytes = 64) => new StrictLfJsonlFramer({
|
||||
maxLineBytes,
|
||||
onRecord: () => undefined,
|
||||
});
|
||||
|
||||
expect(() => make().push('not-json\n')).toThrow(/malformed JSON/);
|
||||
expect(() => make().push('\n')).toThrow(/blank line/);
|
||||
expect(() => make().push(Buffer.from([0xff, 0x0a]))).toThrow(/valid UTF-8/);
|
||||
expect(() => make(4).push('12345')).toThrow(/exceeded 4 bytes/);
|
||||
const partial = make();
|
||||
partial.push('{"ok":true}');
|
||||
expect(() => partial.finish()).toThrow(/partial line/);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Pi RPC client', () => {
|
||||
it('waits for writable completion when the stream applies backpressure', async () => {
|
||||
let written = '';
|
||||
let flush: (() => void) | undefined;
|
||||
const writable = new Writable({
|
||||
highWaterMark: 1,
|
||||
write(chunk, _encoding, callback) {
|
||||
written += chunk.toString();
|
||||
flush = callback;
|
||||
},
|
||||
});
|
||||
const client = new PiRpcClient(writable, { generation: 3, defaultTimeoutMs: 500 });
|
||||
let settled = false;
|
||||
const requested = client.request({ type: 'get_state' }).then((response) => {
|
||||
settled = true;
|
||||
return response;
|
||||
});
|
||||
await new Promise((resolvePromise) => setImmediate(resolvePromise));
|
||||
const command = JSON.parse(written) as { id: string };
|
||||
client.accept({ type: 'response', id: command.id, success: true, data: { ready: true } });
|
||||
await new Promise((resolvePromise) => setImmediate(resolvePromise));
|
||||
expect(settled).toBe(false);
|
||||
flush?.();
|
||||
|
||||
await expect(requested).resolves.toMatchObject({ data: { ready: true } });
|
||||
});
|
||||
|
||||
it('rejects retry policy for prompt-like mutation commands', async () => {
|
||||
const writable = new Writable({ write(_chunk, _encoding, callback) { callback(); } });
|
||||
const client = new PiRpcClient(writable, { generation: 1 });
|
||||
|
||||
await expect(client.request(
|
||||
{ type: 'prompt', message: 'do not replay' },
|
||||
{ retry: 'read-only-once' },
|
||||
)).rejects.toThrow(/not a retryable read-only command/);
|
||||
await expect(client.request(
|
||||
{ type: 'set_model', provider: 'provider-a', modelId: 'model-a' },
|
||||
{ retry: 'read-only-once' },
|
||||
)).rejects.toThrow(/not a retryable read-only command/);
|
||||
});
|
||||
|
||||
it('settles and removes the pending command when the writable fails', async () => {
|
||||
const writable = new Writable({
|
||||
write(_chunk, _encoding, callback) {
|
||||
callback(new Error('closed pipe'));
|
||||
},
|
||||
});
|
||||
const client = new PiRpcClient(writable, { generation: 1 });
|
||||
|
||||
await expect(client.request({ type: 'get_state' }))
|
||||
.rejects.toMatchObject({ code: 'PI_RPC_WRITE_FAILED' });
|
||||
expect(client.pendingCount).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Pi worker process', () => {
|
||||
it('uses the locked offline and no-discovery RPC arguments', () => {
|
||||
expect(buildPiRpcArgs('sessions', ['--model', 'model-a'])).toEqual([
|
||||
'--mode', 'rpc',
|
||||
'--offline',
|
||||
'--session-dir', 'sessions',
|
||||
'--no-extensions',
|
||||
'--no-skills',
|
||||
'--no-prompt-templates',
|
||||
'--no-themes',
|
||||
'--no-context-files',
|
||||
'--no-approve',
|
||||
'--no-tools',
|
||||
'--model', 'model-a',
|
||||
]);
|
||||
});
|
||||
|
||||
it('correlates out-of-order responses, dispatches events, and reassembles partial lines', async () => {
|
||||
const worker = await makeWorker();
|
||||
const events: unknown[] = [];
|
||||
worker.subscribe(() => { throw new Error('consumer failed'); });
|
||||
worker.subscribe((event) => events.push(event));
|
||||
const held = worker.request<{ order: string }>({ type: 'hold' });
|
||||
const released = worker.request<{ order: string }>({ type: 'release' });
|
||||
|
||||
await expect(released).resolves.toMatchObject({ data: { order: 'first' } });
|
||||
await expect(held).resolves.toMatchObject({ data: { order: 'second' } });
|
||||
await expect(worker.request({ type: 'emit_event', marker: 'event-a' }))
|
||||
.resolves.toMatchObject({ data: { emitted: true } });
|
||||
await expect(worker.request({ type: 'partial' }))
|
||||
.resolves.toMatchObject({ data: { partial: true } });
|
||||
expect(events).toContainEqual({ type: 'agent_start', marker: 'event-a' });
|
||||
expect(worker.stderrDiagnostic).toContain('[event-listener] consumer failed');
|
||||
await expect(worker.request({ type: 'echo', value: 'listener-isolated' }))
|
||||
.resolves.toMatchObject({ data: { value: 'listener-isolated' } });
|
||||
});
|
||||
|
||||
it('settles timeout and abort without replaying prompt, while read-only retry runs once', async () => {
|
||||
const worker = await makeWorker({ commandTimeoutMs: 500 });
|
||||
await expect(worker.request(
|
||||
{ type: 'get_state', fakeRetry: true },
|
||||
{ retry: 'read-only-once', timeoutMs: 500 },
|
||||
)).resolves.toMatchObject({ data: { attempts: 2 } });
|
||||
|
||||
await expect(worker.request(
|
||||
{ type: 'prompt', message: 'one attempt' },
|
||||
{ timeoutMs: 100 },
|
||||
)).rejects.toMatchObject({ code: 'PI_RPC_TIMEOUT' });
|
||||
const stats = await worker.request<{ counts: Record<string, number> }>({ type: 'stats' });
|
||||
expect(stats.data?.counts.prompt).toBe(1);
|
||||
|
||||
const controller = new AbortController();
|
||||
const pending = worker.request({ type: 'no_response' }, { signal: controller.signal });
|
||||
const aborted = expect(pending).rejects.toMatchObject({ code: 'PI_RPC_ABORTED' });
|
||||
controller.abort();
|
||||
await aborted;
|
||||
expect(worker.pendingCommandCount).toBe(0);
|
||||
});
|
||||
|
||||
it.each(['malformed', 'blank', 'invalid_utf8', 'trailing_partial', 'large'])(
|
||||
'invalidates only the target worker for %s stdout',
|
||||
async (failureType) => {
|
||||
const left = await makeWorker({ maxLineBytes: 128 });
|
||||
const right = await makeWorker({ maxLineBytes: 128 });
|
||||
const command = failureType === 'large'
|
||||
? { type: failureType, bytes: 512 }
|
||||
: { type: failureType };
|
||||
|
||||
await expect(left.request(command)).rejects.toMatchObject({ code: 'PI_RPC_PROTOCOL_ERROR' });
|
||||
expect(left.generation).toBe(2);
|
||||
expect(left.protocolError).toBeInstanceOf(PiProcessError);
|
||||
expect(left.protocolError?.diagnostic).toContain('[stdout-protocol]');
|
||||
await expect(right.request({ type: 'echo', value: 'still-alive' }))
|
||||
.resolves.toMatchObject({ data: { value: 'still-alive' } });
|
||||
},
|
||||
);
|
||||
|
||||
it('settles every pending command after an unexpected exit', async () => {
|
||||
const worker = await makeWorker();
|
||||
const pending = worker.request({ type: 'no_response' });
|
||||
const crash = worker.request({ type: 'crash' });
|
||||
|
||||
await expect(Promise.all([pending, crash])).rejects.toMatchObject({ code: 'PI_RPC_EXITED' });
|
||||
expect(worker.pendingCommandCount).toBe(0);
|
||||
expect(worker.generation).toBe(2);
|
||||
});
|
||||
|
||||
it('keeps only bounded redacted stderr diagnostics', async () => {
|
||||
const secret = 'credential-that-must-not-leak';
|
||||
const worker = await makeWorker({
|
||||
env: { FAKE_PI_SECRET: secret },
|
||||
sensitiveValues: [secret],
|
||||
diagnosticBytes: 160,
|
||||
});
|
||||
|
||||
await worker.request({ type: 'stderr_secret' });
|
||||
await new Promise((resolvePromise) => setTimeout(resolvePromise, 10));
|
||||
expect(worker.stderrDiagnostic).not.toContain(secret);
|
||||
expect(worker.stderrDiagnostic).toContain('[REDACTED]');
|
||||
expect(Buffer.byteLength(worker.stderrDiagnostic)).toBeLessThanOrEqual(160);
|
||||
expect(sanitizePiDiagnostic(`token=${secret}`, [secret])).toBe('token=[REDACTED]');
|
||||
});
|
||||
|
||||
it('forces the complete child tree down after the graceful deadline', async () => {
|
||||
const worker = await makeWorker({ shutdownGraceMs: 100 });
|
||||
const response = await worker.request<{ pid: number }>({ type: 'spawn_descendant' });
|
||||
const descendantPid = response.data?.pid;
|
||||
expect(descendantPid).toBeTypeOf('number');
|
||||
|
||||
await expect(worker.stop()).resolves.toMatchObject({ mode: 'forced-tree-kill' });
|
||||
for (let attempt = 0; attempt < 20 && await processAlive(descendantPid!); attempt += 1) {
|
||||
await new Promise((resolvePromise) => setTimeout(resolvePromise, 25));
|
||||
}
|
||||
expect(await processAlive(descendantPid!)).toBe(false);
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user