482 lines
18 KiB
TypeScript
482 lines
18 KiB
TypeScript
// @vitest-environment node
|
|
|
|
import { afterEach, describe, expect, it, vi } 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,
|
|
buildPiWorkerEnvironment,
|
|
type PiWorkerLifecycleEvent,
|
|
} 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 () => {
|
|
vi.useRealTimers();
|
|
await Promise.all(workers.splice(0).map((worker) => worker.stop('test_injection').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('can deterministically hold one proof response past the mutation timeout', async () => {
|
|
vi.useFakeTimers();
|
|
let written = '';
|
|
const writable = new Writable({
|
|
write(chunk, _encoding, callback) {
|
|
written += chunk.toString();
|
|
callback();
|
|
},
|
|
});
|
|
const lateResults: Array<{ success: boolean }> = [];
|
|
const client = new PiRpcClient(writable, { generation: 1, defaultTimeoutMs: 10_000 });
|
|
client.delayNextResponseForProof('prompt', 12_000);
|
|
const requested = client.request(
|
|
{ type: 'prompt', message: 'proof response hold' },
|
|
{
|
|
retainAfterTimeout: true,
|
|
onLateResult: (result) => lateResults.push({ success: result.response?.success === true }),
|
|
},
|
|
);
|
|
void requested.catch(() => undefined);
|
|
await Promise.resolve();
|
|
const command = JSON.parse(written) as { id: string };
|
|
client.accept({ type: 'response', id: command.id, success: true });
|
|
|
|
await vi.advanceTimersByTimeAsync(10_000);
|
|
await expect(requested).rejects.toMatchObject({ code: 'PI_RPC_TIMEOUT' });
|
|
expect(client.pendingCount).toBe(1);
|
|
await vi.advanceTimersByTimeAsync(2_000);
|
|
expect(client.pendingCount).toBe(0);
|
|
expect(lateResults).toEqual([{ success: true }]);
|
|
});
|
|
|
|
it('keeps a mutation correlated after the 10 second confirmation timeout', async () => {
|
|
vi.useFakeTimers();
|
|
let written = '';
|
|
const writable = new Writable({
|
|
write(chunk, _encoding, callback) {
|
|
written += chunk.toString();
|
|
callback();
|
|
},
|
|
});
|
|
const lateResults: Array<{ success: boolean; code?: string }> = [];
|
|
const client = new PiRpcClient(writable, { generation: 1, defaultTimeoutMs: 10_000 });
|
|
const requested = client.request(
|
|
{ type: 'prompt', message: 'continue after slow preflight' },
|
|
{
|
|
retainAfterTimeout: true,
|
|
onLateResult: (result: { response?: { success: boolean }; error?: PiProcessError }) => {
|
|
lateResults.push({
|
|
success: result.response?.success === true,
|
|
...(result.error ? { code: result.error.code } : {}),
|
|
});
|
|
},
|
|
} as Parameters<PiRpcClient['request']>[1] & {
|
|
retainAfterTimeout: true;
|
|
onLateResult(result: {
|
|
response?: { success: boolean };
|
|
error?: PiProcessError;
|
|
}): void;
|
|
},
|
|
);
|
|
void requested.catch(() => undefined);
|
|
await Promise.resolve();
|
|
|
|
await vi.advanceTimersByTimeAsync(10_000);
|
|
await expect(requested).rejects.toMatchObject({ code: 'PI_RPC_TIMEOUT' });
|
|
expect(client.pendingCount).toBe(1);
|
|
|
|
const command = JSON.parse(written) as { id: string };
|
|
client.accept({ type: 'response', id: command.id, success: true });
|
|
|
|
expect(client.pendingCount).toBe(0);
|
|
expect(lateResults).toEqual([{ success: true }]);
|
|
});
|
|
|
|
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',
|
|
'--tools',
|
|
'read,bash,edit,write,grep,find,ls,ask_user,subagent,agent_browser,game_asset_browser,game_asset_review,task_state,changed_file,runtime_context',
|
|
'--model', 'model-a',
|
|
]);
|
|
expect(buildPiRpcArgs('sessions', ['--no-session'], ['read', 'grep', 'find', 'ls']))
|
|
.toEqual([
|
|
'--mode', 'rpc',
|
|
'--offline',
|
|
'--session-dir', 'sessions',
|
|
'--no-extensions',
|
|
'--no-skills',
|
|
'--no-prompt-templates',
|
|
'--no-themes',
|
|
'--no-context-files',
|
|
'--no-approve',
|
|
'--tools', 'read,grep,find,ls',
|
|
'--no-session',
|
|
]);
|
|
});
|
|
|
|
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 invalidations: string[] = [];
|
|
worker.subscribeInvalidation((error) => invalidations.push(error.code));
|
|
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);
|
|
expect(invalidations).toEqual(['PI_RPC_EXITED']);
|
|
});
|
|
|
|
it('classifies unexpected exit with bounded redacted stderr and generation correlation', async () => {
|
|
const secret = 'unexpected-exit-secret';
|
|
const lifecycle: PiWorkerLifecycleEvent[] = [];
|
|
const worker = await makeWorker({
|
|
conversationId: 'conversation-exit',
|
|
workerGeneration: 7,
|
|
env: { FAKE_PI_SECRET: secret },
|
|
sensitiveValues: [secret],
|
|
diagnosticBytes: 160,
|
|
onLifecycleEvent: (event) => lifecycle.push(event),
|
|
});
|
|
|
|
await expect(worker.request({ type: 'crash_with_stderr' })).rejects.toMatchObject({
|
|
code: 'PI_RPC_EXITED',
|
|
generation: 7,
|
|
exitCode: 9,
|
|
signal: null,
|
|
diagnostic: expect.stringContaining('worker-exit-marker'),
|
|
});
|
|
const close = lifecycle.find((event) => (
|
|
event.classification === 'unexpected_exit' && event.stage === 'close'
|
|
));
|
|
expect(close).toMatchObject({
|
|
conversationId: 'conversation-exit',
|
|
generation: 7,
|
|
code: 'PI_RPC_EXITED',
|
|
exitCode: 9,
|
|
signal: null,
|
|
});
|
|
expect(JSON.stringify(close)).not.toContain(secret);
|
|
expect(JSON.stringify(close)).not.toContain(scratchRoots.at(-1));
|
|
expect(close?.diagnostic).toContain('cwd=[REDACTED]');
|
|
expect(Buffer.byteLength(close?.diagnostic ?? '')).toBeLessThanOrEqual(160);
|
|
});
|
|
|
|
it('keeps protocol invalidation and deliberate stop as distinct reasoned lifecycle events', async () => {
|
|
const protocolLifecycle: PiWorkerLifecycleEvent[] = [];
|
|
const protocolWorker = await makeWorker({
|
|
conversationId: 'conversation-protocol',
|
|
workerGeneration: 3,
|
|
onLifecycleEvent: (event) => protocolLifecycle.push(event),
|
|
});
|
|
await expect(protocolWorker.request({ type: 'malformed' })).rejects.toMatchObject({
|
|
code: 'PI_RPC_PROTOCOL_ERROR',
|
|
generation: 3,
|
|
});
|
|
expect(protocolLifecycle).toContainEqual(expect.objectContaining({
|
|
classification: 'protocol_invalidation',
|
|
stage: 'protocol',
|
|
conversationId: 'conversation-protocol',
|
|
generation: 3,
|
|
code: 'PI_RPC_PROTOCOL_ERROR',
|
|
}));
|
|
expect(protocolLifecycle).not.toContainEqual(expect.objectContaining({
|
|
classification: 'unexpected_exit',
|
|
}));
|
|
|
|
const stopLifecycle: PiWorkerLifecycleEvent[] = [];
|
|
const stoppedWorker = await makeWorker({
|
|
conversationId: 'conversation-stop',
|
|
workerGeneration: 5,
|
|
onLifecycleEvent: (event) => stopLifecycle.push(event),
|
|
});
|
|
await expect(stoppedWorker.stop('test_injection')).resolves.toMatchObject({
|
|
mode: 'stdin-close',
|
|
});
|
|
expect(stopLifecycle).toContainEqual(expect.objectContaining({
|
|
classification: 'intentional_stop',
|
|
stage: 'stop_completed',
|
|
reason: 'test_injection',
|
|
conversationId: 'conversation-stop',
|
|
generation: 5,
|
|
code: 'PI_WORKER_STOPPED',
|
|
exitCode: 0,
|
|
signal: null,
|
|
}));
|
|
expect(stopLifecycle).not.toContainEqual(expect.objectContaining({
|
|
classification: 'unexpected_exit',
|
|
}));
|
|
});
|
|
|
|
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]');
|
|
expect(sanitizePiDiagnostic('custom-header=q', ['q'])).toBe('custom-header=[REDACTED]');
|
|
});
|
|
|
|
it('inherits only the worker-safe environment allowlist', () => {
|
|
const env = buildPiWorkerEnvironment(
|
|
'D:\\managed-pi',
|
|
{ MAKELore_PI_SELECTED_API_KEY: 'selected-secret' },
|
|
{
|
|
PATH: 'D:\\tools',
|
|
OPENAI_API_KEY: 'unrelated-openai-secret',
|
|
ANTHROPIC_API_KEY: 'unrelated-anthropic-secret',
|
|
CUSTOM_APPLICATION_SECRET: 'unrelated-custom-secret',
|
|
},
|
|
);
|
|
expect(env).toMatchObject({
|
|
PATH: 'D:\\tools',
|
|
MAKELore_PI_SELECTED_API_KEY: 'selected-secret',
|
|
PI_CODING_AGENT_DIR: 'D:\\managed-pi',
|
|
PI_OFFLINE: '1',
|
|
PI_TELEMETRY: '0',
|
|
ELECTRON_RUN_AS_NODE: '1',
|
|
});
|
|
expect(env).not.toHaveProperty('OPENAI_API_KEY');
|
|
expect(env).not.toHaveProperty('ANTHROPIC_API_KEY');
|
|
expect(env).not.toHaveProperty('CUSTOM_APPLICATION_SECRET');
|
|
});
|
|
|
|
it('refuses to put a selected worker credential in argv', async () => {
|
|
const secret = 'argv-secret-value';
|
|
await expect(makeWorker({
|
|
additionalArgs: ['--api-key', secret],
|
|
sensitiveValues: [secret],
|
|
})).rejects.toThrow('arguments contain a sensitive value');
|
|
});
|
|
|
|
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('test_injection')).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);
|
|
});
|
|
});
|