feat: add Pi conversation worker pool
This commit is contained in:
321
tests/unit/pi-conversation-runtime.test.ts
Normal file
321
tests/unit/pi-conversation-runtime.test.ts
Normal file
@@ -0,0 +1,321 @@
|
||||
// @vitest-environment node
|
||||
|
||||
import { mkdtemp, rm } from 'node:fs/promises';
|
||||
import { tmpdir } from 'node:os';
|
||||
import path from 'node:path';
|
||||
import { afterEach, describe, expect, it } from 'vitest';
|
||||
import { createCodingConversationStore } from '../../electron/coding-projects/conversation-store';
|
||||
import { createCodingProjectAgent } from '../../electron/coding-projects/project-config';
|
||||
import {
|
||||
createCodingProjectStore,
|
||||
createLocalCodingProject,
|
||||
createMemoryCodingProjectStorage,
|
||||
} from '../../electron/coding-projects/project-store';
|
||||
import { PiConversationRuntime } from '../../electron/coding-runtime/pi/runtime';
|
||||
import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry';
|
||||
import {
|
||||
PiWorkerPool,
|
||||
type PiConversationWorker,
|
||||
} from '../../electron/coding-runtime/pi/worker-pool';
|
||||
import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
|
||||
import type {
|
||||
PiRpcCommand,
|
||||
PiRpcEvent,
|
||||
PiRpcRequestOptions,
|
||||
PiRpcResponse,
|
||||
} from '../../electron/coding-runtime/pi/rpc-client';
|
||||
|
||||
const roots: string[] = [];
|
||||
const NOW = '2026-08-22T15:00:00.000Z';
|
||||
|
||||
class RuntimeFakeWorker implements PiConversationWorker {
|
||||
readonly requests: PiRpcCommand[] = [];
|
||||
private failType: string | null = null;
|
||||
private readonly responseGates = new Map<string, Promise<void>>();
|
||||
private readonly events = new Set<(event: PiRpcEvent) => void>();
|
||||
private readonly invalidations = new Set<(error: PiProcessError) => void>();
|
||||
|
||||
constructor(readonly id: string, readonly generation: number) {}
|
||||
|
||||
async request<T = unknown>(
|
||||
command: PiRpcCommand,
|
||||
_options?: PiRpcRequestOptions,
|
||||
): Promise<PiRpcResponse<T>> {
|
||||
this.requests.push(structuredClone(command));
|
||||
await this.responseGates.get(command.type);
|
||||
this.responseGates.delete(command.type);
|
||||
if (command.type === this.failType) {
|
||||
this.failType = null;
|
||||
throw new Error(`fake ${command.type} rejection`);
|
||||
}
|
||||
return { type: 'response', id: `${this.id}-${this.requests.length}`, success: true };
|
||||
}
|
||||
|
||||
failNext(type: string): void { this.failType = type; }
|
||||
|
||||
holdNext(type: string): () => void {
|
||||
let release!: () => void;
|
||||
this.responseGates.set(type, new Promise<void>((resolve) => { release = resolve; }));
|
||||
return release;
|
||||
}
|
||||
|
||||
subscribe(listener: (event: PiRpcEvent) => void): () => void {
|
||||
this.events.add(listener);
|
||||
return () => this.events.delete(listener);
|
||||
}
|
||||
|
||||
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void {
|
||||
this.invalidations.add(listener);
|
||||
return () => this.invalidations.delete(listener);
|
||||
}
|
||||
|
||||
emit(event: PiRpcEvent): void {
|
||||
for (const listener of this.events) listener(event);
|
||||
}
|
||||
|
||||
async stop() {
|
||||
return { mode: 'stdin-close' as const, code: 0, signal: null };
|
||||
}
|
||||
}
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
|
||||
});
|
||||
|
||||
describe('Pi Conversation runtime', () => {
|
||||
it('separates RPC acceptance from settle and changes only the target Conversation model', async () => {
|
||||
const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-runtime-'));
|
||||
roots.push(projectPath);
|
||||
const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
|
||||
createId: () => 'project-a',
|
||||
now: () => NOW,
|
||||
});
|
||||
await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
|
||||
await createCodingProjectAgent(projectPath, {
|
||||
id: 'agent-a',
|
||||
avatarId: 'avatar-01',
|
||||
roleName: 'Implementer',
|
||||
name: 'Agent A',
|
||||
model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
|
||||
modelResolution: 'resolved',
|
||||
responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
|
||||
}, { now: NOW });
|
||||
const store = createCodingConversationStore(projectPath, {
|
||||
createId: (() => {
|
||||
const ids = [
|
||||
'f47ac10b-58cc-4372-a567-0e02b2c3d479',
|
||||
'8b1a9953-c461-4d88-9c3e-7e1f8f3f2c11',
|
||||
'3d594650-3436-4a8c-8b38-7d1c5e3f9a20',
|
||||
];
|
||||
return () => ids.shift() as string;
|
||||
})(),
|
||||
now: () => NOW,
|
||||
});
|
||||
const model = { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' as const };
|
||||
const left = await store.create({ agentId: 'agent-a', title: 'Left', model, modelResolution: 'resolved' });
|
||||
const right = await store.create({ agentId: 'agent-a', title: 'Right', model, modelResolution: 'resolved' });
|
||||
const forkTarget = await store.create({ agentId: 'agent-a', title: 'Fork', model, modelResolution: 'resolved' });
|
||||
const inputs = [left, right, forkTarget].map((item) => ({
|
||||
conversationId: item.id,
|
||||
projectId: 'project-a',
|
||||
agentId: 'agent-a',
|
||||
title: item.title,
|
||||
model: { model: item.model, modelResolution: item.modelResolution },
|
||||
}));
|
||||
const workers = new Map<string, RuntimeFakeWorker>();
|
||||
const workerHistory = new Map<string, RuntimeFakeWorker[]>();
|
||||
const openInputs: Array<{ conversationId: string; forkSource?: string; sourceEntryId?: string }> = [];
|
||||
const pool = new PiWorkerPool({
|
||||
maxIdle: 4,
|
||||
openWorker: async ({ conversation, fork, generation, existingSession }) => {
|
||||
openInputs.push({
|
||||
conversationId: conversation.conversationId,
|
||||
...(fork ? {
|
||||
forkSource: fork.sourceSession.piSessionId,
|
||||
sourceEntryId: fork.sourceEntryId,
|
||||
} : {}),
|
||||
});
|
||||
const worker = new RuntimeFakeWorker(`worker-${conversation.conversationId}-${generation}`, generation);
|
||||
workers.set(conversation.conversationId, worker);
|
||||
workerHistory.set(conversation.conversationId, [
|
||||
...(workerHistory.get(conversation.conversationId) ?? []),
|
||||
worker,
|
||||
]);
|
||||
return {
|
||||
worker,
|
||||
session: existingSession ?? {
|
||||
piSessionId: `session-${conversation.conversationId}`,
|
||||
sessionKey: `key-${conversation.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
const runtime = new PiConversationRuntime({
|
||||
pool,
|
||||
registry: new PiSessionRegistry({ projectStore }),
|
||||
createId: (kind) => `${kind}-fixed`,
|
||||
resolveModel: async (candidate) => {
|
||||
if (candidate.accountId !== 'account-b' || candidate.modelId !== 'model-b') {
|
||||
throw new Error('model unavailable');
|
||||
}
|
||||
return {
|
||||
accountId: candidate.accountId,
|
||||
runtimeProviderId: 'runtime-account-b',
|
||||
modelId: candidate.modelId,
|
||||
thinkingLevel: candidate.thinkingLevel,
|
||||
input: ['text'],
|
||||
};
|
||||
},
|
||||
});
|
||||
await Promise.all(inputs.slice(0, 2).map((input) => runtime.prepare(input)));
|
||||
|
||||
const releasePromptAcceptance = workers.get(left.id)!.holdNext('prompt');
|
||||
let acceptanceResolved = false;
|
||||
const acceptance = runtime.prompt({
|
||||
clientRequestId: 'request-left',
|
||||
conversationId: left.id,
|
||||
mode: 'prompt',
|
||||
text: 'Implement the change',
|
||||
attachments: [],
|
||||
}).then((value) => {
|
||||
acceptanceResolved = true;
|
||||
return value;
|
||||
});
|
||||
await expect.poll(() => workers.get(left.id)!.requests.length).toBe(1);
|
||||
expect(acceptanceResolved).toBe(false);
|
||||
releasePromptAcceptance();
|
||||
const accepted = await acceptance;
|
||||
expect(accepted).toMatchObject({ accepted: true, runId: 'run-fixed', mode: 'prompt' });
|
||||
expect((await runtime.getSnapshot(left.id)).run.status).toBe('running');
|
||||
workers.get(left.id)!.emit({ type: 'agent_end' });
|
||||
expect((await runtime.getSnapshot(left.id)).run.status).toBe('running');
|
||||
workers.get(left.id)!.emit({ type: 'agent_settled' });
|
||||
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('idle');
|
||||
|
||||
const changed = await runtime.setModel({
|
||||
conversationId: left.id,
|
||||
accountId: 'account-b',
|
||||
modelId: 'model-b',
|
||||
});
|
||||
expect(changed.model).toEqual({ accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'medium' });
|
||||
expect(workerHistory.get(left.id)).toHaveLength(2);
|
||||
expect(workerHistory.get(left.id)![0]!.requests).not.toContainEqual(expect.objectContaining({
|
||||
type: 'set_model',
|
||||
}));
|
||||
expect(workers.get(left.id)!.requests.map(({ type }) => type).sort()).toEqual([
|
||||
'get_entries',
|
||||
'get_state',
|
||||
]);
|
||||
expect(pool.getState(left.id)).toMatchObject({ generation: 2, state: 'ready' });
|
||||
expect(pool.getState(right.id)).toMatchObject({ generation: 1, state: 'ready' });
|
||||
expect((await runtime.getSnapshot(right.id)).conversation.model.model).toEqual(model);
|
||||
|
||||
await expect(runtime.setModel({
|
||||
conversationId: left.id,
|
||||
accountId: 'account-missing',
|
||||
modelId: 'model-missing',
|
||||
})).rejects.toThrow('model unavailable');
|
||||
expect((await runtime.getSnapshot(left.id)).conversation.model).toEqual(changed);
|
||||
|
||||
await runtime.prompt({
|
||||
clientRequestId: 'request-left-active',
|
||||
conversationId: left.id,
|
||||
mode: 'prompt',
|
||||
text: 'Keep working',
|
||||
attachments: [],
|
||||
});
|
||||
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('running');
|
||||
workers.get(left.id)!.failNext('steer');
|
||||
await expect(runtime.steer({
|
||||
clientRequestId: 'request-left-rejected-steer',
|
||||
conversationId: left.id,
|
||||
text: 'This command is rejected',
|
||||
attachments: [],
|
||||
})).rejects.toThrow('fake steer rejection');
|
||||
expect((await runtime.getSnapshot(left.id)).queue.items).toEqual([]);
|
||||
await expect(runtime.steer({
|
||||
clientRequestId: 'request-left-steer',
|
||||
conversationId: left.id,
|
||||
text: 'Use the smaller seam',
|
||||
attachments: [],
|
||||
})).resolves.toMatchObject({ accepted: true, mode: 'steer', queuePosition: 1 });
|
||||
await expect(runtime.followUp({
|
||||
clientRequestId: 'request-left-follow-up',
|
||||
conversationId: left.id,
|
||||
text: 'Then run the focused test',
|
||||
attachments: [],
|
||||
})).resolves.toMatchObject({ accepted: true, mode: 'follow-up', queuePosition: 2 });
|
||||
await expect(runtime.setThinking({
|
||||
conversationId: left.id,
|
||||
thinkingLevel: 'high',
|
||||
})).resolves.toMatchObject({
|
||||
model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' },
|
||||
});
|
||||
workers.get(left.id)!.failNext('abort');
|
||||
await expect(runtime.abort(left.id)).rejects.toThrow('fake abort rejection');
|
||||
expect((await runtime.getSnapshot(left.id)).run.status).toBe('running');
|
||||
await runtime.abort(left.id);
|
||||
expect((await runtime.getSnapshot(left.id)).run.status).toBe('aborting');
|
||||
expect(workers.get(left.id)!.requests.slice(-6).map(({ type }) => type)).toEqual([
|
||||
'steer',
|
||||
'steer',
|
||||
'follow_up',
|
||||
'set_thinking_level',
|
||||
'abort',
|
||||
'abort',
|
||||
]);
|
||||
expect(workers.get(right.id)!.requests).toHaveLength(0);
|
||||
workers.get(left.id)!.emit({ type: 'agent_settled' });
|
||||
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('idle');
|
||||
expect((await runtime.getSnapshot(left.id)).queue.items).toEqual([]);
|
||||
|
||||
workers.get(right.id)!.failNext('compact');
|
||||
await expect(runtime.compact(right.id)).rejects.toThrow('fake compact rejection');
|
||||
expect((await runtime.getSnapshot(right.id)).run.status).toBe('error');
|
||||
await runtime.compact(right.id);
|
||||
expect((await runtime.getSnapshot(right.id)).run.status).toBe('compacting');
|
||||
expect(workers.get(right.id)!.requests.at(-1)).toEqual({ type: 'compact' });
|
||||
expect((await runtime.getSnapshot(left.id)).run.status).toBe('idle');
|
||||
workers.get(right.id)!.emit({ type: 'agent_end' });
|
||||
expect((await runtime.getSnapshot(right.id)).run.status).toBe('compacting');
|
||||
workers.get(right.id)!.emit({ type: 'agent_settled' });
|
||||
await expect.poll(async () => (await runtime.getSnapshot(right.id)).run.status).toBe('idle');
|
||||
|
||||
await expect(runtime.recover(right.id)).resolves.toMatchObject({
|
||||
conversationId: right.id,
|
||||
status: 'ready',
|
||||
workerGeneration: 2,
|
||||
});
|
||||
expect((await runtime.getSnapshot(right.id)).cursor).toMatchObject({
|
||||
workerGeneration: 2,
|
||||
seq: 0,
|
||||
});
|
||||
expect(workers.get(right.id)!.requests.map(({ type }) => type).sort()).toEqual([
|
||||
'get_entries',
|
||||
'get_state',
|
||||
]);
|
||||
expect(pool.getState(left.id)).toMatchObject({ generation: 2, state: 'idle' });
|
||||
|
||||
const sourceRequestCount = workers.get(left.id)!.requests.length;
|
||||
const forked = await runtime.fork({
|
||||
sourceConversationId: left.id,
|
||||
sourceEntryId: 'entry-a',
|
||||
conversation: inputs[2]!,
|
||||
});
|
||||
expect(forked.conversationId).toBe(forkTarget.id);
|
||||
expect(workers.get(forkTarget.id)?.id).toBe(`worker-${forkTarget.id}-1`);
|
||||
expect(workers.get(left.id)!.requests).toHaveLength(sourceRequestCount);
|
||||
expect(openInputs.at(-1)).toEqual({
|
||||
conversationId: forkTarget.id,
|
||||
forkSource: `session-${left.id}`,
|
||||
sourceEntryId: 'entry-a',
|
||||
});
|
||||
|
||||
await runtime.dispose(forkTarget.id);
|
||||
expect(pool.getState(forkTarget.id)).toBeNull();
|
||||
expect(pool.getState(left.id)).toMatchObject({ generation: 2, state: 'idle' });
|
||||
await expect(runtime.getSnapshot(forkTarget.id)).rejects.toMatchObject({
|
||||
publicError: { code: 'CODING_CONVERSATION_NOT_FOUND' },
|
||||
});
|
||||
});
|
||||
});
|
||||
173
tests/unit/pi-managed-worker-opener.test.ts
Normal file
173
tests/unit/pi-managed-worker-opener.test.ts
Normal file
@@ -0,0 +1,173 @@
|
||||
// @vitest-environment node
|
||||
|
||||
import { mkdtemp, readFile, rm } from 'node:fs/promises';
|
||||
import { tmpdir } from 'node:os';
|
||||
import path from 'node:path';
|
||||
import { afterEach, describe, expect, it } from 'vitest';
|
||||
import { createCodingConversationStore } from '../../electron/coding-projects/conversation-store';
|
||||
import { createCodingProjectAgent } from '../../electron/coding-projects/project-config';
|
||||
import {
|
||||
createCodingProjectStore,
|
||||
createLocalCodingProject,
|
||||
createMemoryCodingProjectStorage,
|
||||
} from '../../electron/coding-projects/project-store';
|
||||
import type { ProviderAccount } from '../../electron/shared/providers/types';
|
||||
import {
|
||||
createPiManagedWorkerOpener,
|
||||
type PiWorkerProcessAdapter,
|
||||
} from '../../electron/coding-runtime/pi/runtime';
|
||||
import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry';
|
||||
import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
|
||||
import type {
|
||||
PiRpcCommand,
|
||||
PiRpcEvent,
|
||||
PiRpcRequestOptions,
|
||||
PiRpcResponse,
|
||||
} from '../../electron/coding-runtime/pi/rpc-client';
|
||||
import type { PiWorkerProcessOptions } from '../../electron/coding-runtime/pi/worker-process';
|
||||
import type { PiRuntimeTelemetryEvent } from '../../electron/coding-runtime/pi/telemetry';
|
||||
|
||||
const roots: string[] = [];
|
||||
const NOW = '2026-08-22T16:00:00.000Z';
|
||||
|
||||
class OpenerFakeProcess implements PiWorkerProcessAdapter {
|
||||
readonly generation = 1;
|
||||
constructor(
|
||||
private readonly sessionId: string,
|
||||
private readonly options: PiWorkerProcessOptions,
|
||||
) {}
|
||||
|
||||
async start() { return this; }
|
||||
|
||||
async request<T = unknown>(command: PiRpcCommand, _options?: PiRpcRequestOptions): Promise<PiRpcResponse<T>> {
|
||||
if (command.type !== 'get_state') {
|
||||
return { type: 'response', id: 'fake', success: true };
|
||||
}
|
||||
return {
|
||||
type: 'response',
|
||||
id: 'fake-state',
|
||||
success: true,
|
||||
data: {
|
||||
sessionId: this.sessionId,
|
||||
sessionFile: path.join(this.options.sessionDir, `${this.sessionId}.jsonl`),
|
||||
} as T,
|
||||
};
|
||||
}
|
||||
|
||||
subscribe(_listener: (event: PiRpcEvent) => void): () => void { return () => undefined; }
|
||||
subscribeInvalidation(_listener: (error: PiProcessError) => void): () => void { return () => undefined; }
|
||||
async stop() { return { mode: 'stdin-close' as const, code: 0, signal: null }; }
|
||||
}
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
|
||||
});
|
||||
|
||||
describe('managed Pi worker opener', () => {
|
||||
it('reuses PI-040 credential, managed-resource, and persistent-session projections', async () => {
|
||||
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-opener-'));
|
||||
roots.push(root);
|
||||
const projectPath = path.join(root, 'project');
|
||||
const userDataDir = path.join(root, 'user-data');
|
||||
const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
|
||||
createId: () => 'project-a',
|
||||
now: () => NOW,
|
||||
});
|
||||
await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
|
||||
await createCodingProjectAgent(projectPath, {
|
||||
id: 'agent-a',
|
||||
avatarId: 'avatar-01',
|
||||
roleName: 'Implementer',
|
||||
name: 'Agent A',
|
||||
model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'high' },
|
||||
modelResolution: 'resolved',
|
||||
responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
|
||||
prompt: 'PRIVATE MANAGED PROMPT',
|
||||
skillIds: ['grilling'],
|
||||
}, { now: NOW });
|
||||
const conversationStore = createCodingConversationStore(projectPath, {
|
||||
createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d479',
|
||||
now: () => NOW,
|
||||
});
|
||||
const created = await conversationStore.create({
|
||||
agentId: 'agent-a',
|
||||
title: 'Conversation A',
|
||||
model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'high' },
|
||||
modelResolution: 'resolved',
|
||||
});
|
||||
const input = {
|
||||
conversationId: created.id,
|
||||
projectId: 'project-a',
|
||||
agentId: 'agent-a',
|
||||
title: created.title,
|
||||
model: { model: created.model, modelResolution: created.modelResolution },
|
||||
} as const;
|
||||
const account: ProviderAccount = {
|
||||
id: 'account-a',
|
||||
vendorId: 'custom',
|
||||
label: 'Account A',
|
||||
authMode: 'api_key',
|
||||
apiProtocol: 'openai-completions',
|
||||
baseUrl: 'https://provider.example/v1',
|
||||
model: 'model-a',
|
||||
enabled: true,
|
||||
isDefault: true,
|
||||
createdAt: NOW,
|
||||
updatedAt: NOW,
|
||||
};
|
||||
const processOptions: PiWorkerProcessOptions[] = [];
|
||||
const telemetry: PiRuntimeTelemetryEvent[] = [];
|
||||
const registry = new PiSessionRegistry({ projectStore });
|
||||
const opener = createPiManagedWorkerOpener({
|
||||
registry,
|
||||
executablePath: 'electron.exe',
|
||||
cliPath: 'pi-cli.js',
|
||||
userDataDir,
|
||||
bundledSkillsDir: path.resolve('resources/coding-skills'),
|
||||
loadProviderInput: async () => ({ accounts: [account], modelSummaries: [] }),
|
||||
resolveCredential: async () => 'provider-secret-value',
|
||||
createSessionKey: () => 'session-key-a',
|
||||
onTelemetry: (event) => telemetry.push(event),
|
||||
createProcess: (options) => {
|
||||
processOptions.push(options);
|
||||
const sessionIndex = options.additionalArgs?.indexOf('--session-id') ?? -1;
|
||||
return new OpenerFakeProcess(options.additionalArgs?.[sessionIndex + 1] ?? '', options);
|
||||
},
|
||||
});
|
||||
|
||||
const first = await opener({
|
||||
conversation: input,
|
||||
generation: 1,
|
||||
revision: { provider: 1, resources: 1 },
|
||||
});
|
||||
const reopened = await opener({
|
||||
conversation: input,
|
||||
generation: 2,
|
||||
revision: { provider: 2, resources: 1 },
|
||||
existingSession: first.session,
|
||||
});
|
||||
|
||||
expect(first.session).toEqual({ piSessionId: 'session-key-a', sessionKey: 'session-key-a' });
|
||||
expect(reopened.session).toEqual(first.session);
|
||||
expect(processOptions).toHaveLength(2);
|
||||
for (const options of processOptions) {
|
||||
const argv = JSON.stringify(options.additionalArgs);
|
||||
expect(argv).toContain('--system-prompt');
|
||||
expect(argv).toContain('grilling');
|
||||
expect(argv).toContain('--session-id');
|
||||
expect(argv).not.toContain('PRIVATE MANAGED PROMPT');
|
||||
expect(argv).not.toContain('provider-secret-value');
|
||||
expect(Object.values(options.env ?? {})).toContain('provider-secret-value');
|
||||
expect(options.sensitiveValues).toContain('provider-secret-value');
|
||||
}
|
||||
const modelsFile = path.join(userDataDir, 'coding-runtime', 'pi', 'config', 'models.json');
|
||||
expect(await readFile(modelsFile, 'utf8')).not.toContain('provider-secret-value');
|
||||
expect(telemetry.map(({ milestone }) => milestone)).toEqual([
|
||||
'resources.ready', 'worker.spawn', 'rpc.ready', 'session.open',
|
||||
'resources.ready', 'worker.spawn', 'rpc.ready', 'session.open',
|
||||
]);
|
||||
expect(JSON.stringify(telemetry)).not.toContain(created.id);
|
||||
expect(JSON.stringify(telemetry)).not.toContain('PRIVATE MANAGED PROMPT');
|
||||
expect(JSON.stringify(telemetry)).not.toContain('provider-secret-value');
|
||||
});
|
||||
});
|
||||
@@ -233,12 +233,15 @@ describe('Pi worker process', () => {
|
||||
|
||||
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('keeps only bounded redacted stderr diagnostics', async () => {
|
||||
|
||||
130
tests/unit/pi-runtime-auth-recovery.test.ts
Normal file
130
tests/unit/pi-runtime-auth-recovery.test.ts
Normal file
@@ -0,0 +1,130 @@
|
||||
// @vitest-environment node
|
||||
|
||||
import { mkdtemp, rm } from 'node:fs/promises';
|
||||
import { tmpdir } from 'node:os';
|
||||
import path from 'node:path';
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest';
|
||||
import { createCodingConversationStore } from '../../electron/coding-projects/conversation-store';
|
||||
import { createCodingProjectAgent } from '../../electron/coding-projects/project-config';
|
||||
import {
|
||||
createCodingProjectStore,
|
||||
createLocalCodingProject,
|
||||
createMemoryCodingProjectStorage,
|
||||
} from '../../electron/coding-projects/project-store';
|
||||
import { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
|
||||
import { PiConversationRuntime } from '../../electron/coding-runtime/pi/runtime';
|
||||
import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry';
|
||||
import { PiWorkerPool, type PiConversationWorker } from '../../electron/coding-runtime/pi/worker-pool';
|
||||
import type {
|
||||
PiRpcCommand,
|
||||
PiRpcEvent,
|
||||
PiRpcRequestOptions,
|
||||
PiRpcResponse,
|
||||
} from '../../electron/coding-runtime/pi/rpc-client';
|
||||
|
||||
const roots: string[] = [];
|
||||
const NOW = '2026-08-22T17:00:00.000Z';
|
||||
|
||||
class AuthFailureWorker implements PiConversationWorker {
|
||||
readonly generation = 1;
|
||||
readonly requests: PiRpcCommand[] = [];
|
||||
constructor(readonly id: string) {}
|
||||
|
||||
async request<T = unknown>(command: PiRpcCommand, _options?: PiRpcRequestOptions): Promise<PiRpcResponse<T>> {
|
||||
this.requests.push(structuredClone(command));
|
||||
if (command.type === 'prompt') {
|
||||
throw new PiProcessError('PI_RPC_RESPONSE_ERROR', '401 provider authentication failed');
|
||||
}
|
||||
return { type: 'response', id: 'fake', success: true };
|
||||
}
|
||||
|
||||
subscribe(_listener: (event: PiRpcEvent) => void): () => void { return () => undefined; }
|
||||
subscribeInvalidation(_listener: (error: PiProcessError) => void): () => void { return () => undefined; }
|
||||
async stop() { return { mode: 'stdin-close' as const, code: 0, signal: null }; }
|
||||
}
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
|
||||
});
|
||||
|
||||
describe('Pi runtime Provider authentication recovery', () => {
|
||||
it('refreshes and reopens once, then exposes the second authentication failure without looping', async () => {
|
||||
const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-auth-'));
|
||||
roots.push(projectPath);
|
||||
const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
|
||||
createId: () => 'project-auth',
|
||||
now: () => NOW,
|
||||
});
|
||||
await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
|
||||
await createCodingProjectAgent(projectPath, {
|
||||
id: 'agent-auth',
|
||||
avatarId: 'avatar-01',
|
||||
roleName: 'Implementer',
|
||||
name: 'Auth Agent',
|
||||
model: { accountId: 'account-auth', modelId: 'model-auth', thinkingLevel: 'medium' },
|
||||
modelResolution: 'resolved',
|
||||
responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
|
||||
}, { now: NOW });
|
||||
const store = createCodingConversationStore(projectPath, {
|
||||
createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d479',
|
||||
now: () => NOW,
|
||||
});
|
||||
const conversation = await store.create({
|
||||
agentId: 'agent-auth',
|
||||
title: 'Auth Conversation',
|
||||
model: { accountId: 'account-auth', modelId: 'model-auth', thinkingLevel: 'medium' },
|
||||
modelResolution: 'resolved',
|
||||
});
|
||||
const workers: AuthFailureWorker[] = [];
|
||||
const pool = new PiWorkerPool({
|
||||
maxIdle: 2,
|
||||
openWorker: async ({ conversation: input, existingSession }) => {
|
||||
const worker = new AuthFailureWorker(`worker-${workers.length + 1}`);
|
||||
workers.push(worker);
|
||||
return {
|
||||
worker,
|
||||
session: existingSession ?? {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
const refreshCredential = vi.fn(async () => undefined);
|
||||
const runtime = new PiConversationRuntime({
|
||||
pool,
|
||||
registry: new PiSessionRegistry({ projectStore }),
|
||||
resolveModel: async () => { throw new Error('not used'); },
|
||||
refreshCredential,
|
||||
isAuthenticationError: (error) => (
|
||||
error instanceof PiProcessError && error.message.includes('401')
|
||||
),
|
||||
createId: () => 'run-auth',
|
||||
});
|
||||
await runtime.prepare({
|
||||
conversationId: conversation.id,
|
||||
projectId: 'project-auth',
|
||||
agentId: 'agent-auth',
|
||||
title: conversation.title,
|
||||
model: { model: conversation.model, modelResolution: conversation.modelResolution },
|
||||
});
|
||||
|
||||
await expect(runtime.prompt({
|
||||
clientRequestId: 'request-auth',
|
||||
conversationId: conversation.id,
|
||||
mode: 'prompt',
|
||||
text: 'Do not leak credentials',
|
||||
attachments: [],
|
||||
})).rejects.toMatchObject({
|
||||
code: 'PI_RPC_RESPONSE_ERROR',
|
||||
});
|
||||
|
||||
await expect.poll(() => workers.length).toBe(2);
|
||||
await expect.poll(async () => (await runtime.getSnapshot(conversation.id)).run.status).toBe('error');
|
||||
expect(refreshCredential).toHaveBeenCalledTimes(1);
|
||||
expect(workers).toHaveLength(2);
|
||||
expect(workers.map((worker) => worker.requests.filter(({ type }) => type === 'prompt').length))
|
||||
.toEqual([1, 1]);
|
||||
expect(pool.getState(conversation.id)).toMatchObject({ state: 'idle', generation: 2 });
|
||||
});
|
||||
});
|
||||
95
tests/unit/pi-session-registry.test.ts
Normal file
95
tests/unit/pi-session-registry.test.ts
Normal file
@@ -0,0 +1,95 @@
|
||||
// @vitest-environment node
|
||||
|
||||
import { mkdtemp, rm } from 'node:fs/promises';
|
||||
import { tmpdir } from 'node:os';
|
||||
import path from 'node:path';
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest';
|
||||
import { createCodingConversationStore } from '../../electron/coding-projects/conversation-store';
|
||||
import { createCodingProjectAgent } from '../../electron/coding-projects/project-config';
|
||||
import {
|
||||
createCodingProjectStore,
|
||||
createLocalCodingProject,
|
||||
createMemoryCodingProjectStorage,
|
||||
} from '../../electron/coding-projects/project-store';
|
||||
import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry';
|
||||
|
||||
const roots: string[] = [];
|
||||
const NOW = '2026-08-22T14:00:00.000Z';
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
|
||||
});
|
||||
|
||||
describe('Pi session registry', () => {
|
||||
it('persists one binding and target-only model state across registry reopen', async () => {
|
||||
const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-registry-'));
|
||||
roots.push(projectPath);
|
||||
const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
|
||||
createId: () => 'project-a',
|
||||
now: () => NOW,
|
||||
});
|
||||
const { project } = await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
|
||||
await createCodingProjectAgent(projectPath, {
|
||||
id: 'agent-a',
|
||||
avatarId: 'avatar-01',
|
||||
roleName: 'Implementer',
|
||||
name: 'Agent A',
|
||||
model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
|
||||
modelResolution: 'resolved',
|
||||
responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
|
||||
prompt: 'Managed prompt',
|
||||
skillIds: ['tdd'],
|
||||
}, { now: NOW });
|
||||
const conversations = createCodingConversationStore(projectPath, {
|
||||
createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d479',
|
||||
now: () => NOW,
|
||||
});
|
||||
const created = await conversations.create({
|
||||
agentId: 'agent-a',
|
||||
title: 'Conversation A',
|
||||
model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
|
||||
modelResolution: 'resolved',
|
||||
});
|
||||
const input = {
|
||||
conversationId: created.id,
|
||||
projectId: 'project-a',
|
||||
agentId: 'agent-a',
|
||||
title: created.title,
|
||||
model: { model: created.model, modelResolution: created.modelResolution },
|
||||
} as const;
|
||||
const registry = new PiSessionRegistry({ projectStore });
|
||||
const createBinding = vi.fn(async () => ({
|
||||
piSessionId: 'pi-session-a',
|
||||
sessionKey: 'session-key-a',
|
||||
}));
|
||||
|
||||
const [first, duplicate] = await Promise.all([
|
||||
registry.ensureBinding(input, createBinding),
|
||||
registry.ensureBinding(input, createBinding),
|
||||
]);
|
||||
expect(createBinding).toHaveBeenCalledTimes(1);
|
||||
expect(first.session).toEqual({ piSessionId: 'pi-session-a', sessionKey: 'session-key-a' });
|
||||
expect(duplicate).toEqual(first);
|
||||
|
||||
await registry.setModel(created.id, {
|
||||
model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' },
|
||||
modelResolution: 'resolved',
|
||||
});
|
||||
const reopened = await new PiSessionRegistry({ projectStore }).prepare({
|
||||
...input,
|
||||
model: {
|
||||
model: { thinkingLevel: 'high', modelId: 'model-b', accountId: 'account-b' },
|
||||
modelResolution: 'resolved',
|
||||
},
|
||||
});
|
||||
expect(reopened).toMatchObject({
|
||||
projectPath: project.path,
|
||||
agent: { id: 'agent-a', prompt: 'Managed prompt', skillIds: ['tdd'] },
|
||||
conversation: {
|
||||
id: created.id,
|
||||
model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' },
|
||||
},
|
||||
session: { piSessionId: 'pi-session-a', sessionKey: 'session-key-a' },
|
||||
});
|
||||
});
|
||||
});
|
||||
163
tests/unit/pi-worker-pool-process-integration.test.ts
Normal file
163
tests/unit/pi-worker-pool-process-integration.test.ts
Normal file
@@ -0,0 +1,163 @@
|
||||
// @vitest-environment node
|
||||
|
||||
import { mkdir, mkdtemp, rm } from 'node:fs/promises';
|
||||
import { tmpdir } from 'node:os';
|
||||
import path from 'node:path';
|
||||
import { afterEach, describe, expect, it } from 'vitest';
|
||||
import type { PrepareConversationInput } from '../../electron/coding-runtime/contracts';
|
||||
import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
|
||||
import type {
|
||||
PiRpcCommand,
|
||||
PiRpcEvent,
|
||||
PiRpcRequestOptions,
|
||||
PiRpcResponse,
|
||||
} from '../../electron/coding-runtime/pi/rpc-client';
|
||||
import {
|
||||
PiProcessBudget,
|
||||
PiWorkerPool,
|
||||
type PiConversationWorker,
|
||||
} from '../../electron/coding-runtime/pi/worker-pool';
|
||||
import {
|
||||
PiWorkerProcess,
|
||||
type PiWorkerStopResult,
|
||||
} from '../../electron/coding-runtime/pi/worker-process';
|
||||
|
||||
const scratchRoots: string[] = [];
|
||||
|
||||
function conversation(conversationId: string): PrepareConversationInput {
|
||||
return {
|
||||
conversationId,
|
||||
projectId: 'project-process',
|
||||
agentId: 'agent-process',
|
||||
title: conversationId,
|
||||
model: {
|
||||
model: { accountId: 'account-process', modelId: 'model-process', thinkingLevel: 'medium' },
|
||||
modelResolution: 'resolved',
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
class ProcessBackedWorker implements PiConversationWorker {
|
||||
constructor(
|
||||
readonly id: string,
|
||||
readonly generation: number,
|
||||
private readonly process: PiWorkerProcess,
|
||||
) {}
|
||||
|
||||
request<T = unknown>(command: PiRpcCommand, options?: PiRpcRequestOptions): Promise<PiRpcResponse<T>> {
|
||||
return this.process.request(command, options);
|
||||
}
|
||||
|
||||
subscribe(listener: (event: PiRpcEvent) => void): () => void {
|
||||
return this.process.subscribe(listener);
|
||||
}
|
||||
|
||||
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void {
|
||||
return this.process.subscribeInvalidation(listener);
|
||||
}
|
||||
|
||||
stop(): Promise<PiWorkerStopResult> {
|
||||
return this.process.stop();
|
||||
}
|
||||
}
|
||||
|
||||
afterEach(async () => {
|
||||
await Promise.all(scratchRoots.splice(0).map((root) => rm(root, {
|
||||
recursive: true,
|
||||
force: true,
|
||||
maxRetries: 3,
|
||||
})));
|
||||
});
|
||||
|
||||
describe('Pi worker pool process integration', () => {
|
||||
it('runs two child processes concurrently and aborts only the addressed Conversation', async () => {
|
||||
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-pool-process-'));
|
||||
scratchRoots.push(root);
|
||||
const processBudget = new PiProcessBudget(8);
|
||||
const events: Array<{ conversationId: string; event: PiRpcEvent }> = [];
|
||||
const pool = new PiWorkerPool({
|
||||
maxRunning: 4,
|
||||
maxIdle: 4,
|
||||
processBudget,
|
||||
openWorker: async ({ conversation: input, generation }) => {
|
||||
const workerRoot = path.join(root, input.conversationId);
|
||||
const configDir = path.join(workerRoot, 'config');
|
||||
const sessionDir = path.join(workerRoot, 'sessions');
|
||||
const cwd = path.join(workerRoot, 'project');
|
||||
await Promise.all([
|
||||
mkdir(configDir, { recursive: true }),
|
||||
mkdir(sessionDir, { recursive: true }),
|
||||
mkdir(cwd, { recursive: true }),
|
||||
]);
|
||||
const child = await new PiWorkerProcess({
|
||||
executablePath: process.execPath,
|
||||
cliPath: path.resolve('tests/fixtures/fake-pi-pool-child.mjs'),
|
||||
cwd,
|
||||
configDir,
|
||||
sessionDir,
|
||||
commandTimeoutMs: 2_000,
|
||||
shutdownGraceMs: 500,
|
||||
}).start();
|
||||
return {
|
||||
worker: new ProcessBackedWorker(
|
||||
`worker-${input.conversationId}-${generation}`,
|
||||
generation,
|
||||
child,
|
||||
),
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
const unsubscribe = pool.subscribe((event) => {
|
||||
if (event.type === 'worker.event') {
|
||||
events.push({ conversationId: event.conversationId, event: event.event });
|
||||
}
|
||||
});
|
||||
|
||||
try {
|
||||
await Promise.all([
|
||||
pool.prepare(conversation('conversation-left')),
|
||||
pool.prepare(conversation('conversation-right')),
|
||||
]);
|
||||
const left = pool.startTopLevel({
|
||||
conversationId: 'conversation-left',
|
||||
runId: 'run-left',
|
||||
command: { type: 'prompt', message: 'left', delayMs: 1_000 },
|
||||
});
|
||||
const right = pool.startTopLevel({
|
||||
conversationId: 'conversation-right',
|
||||
runId: 'run-right',
|
||||
command: { type: 'prompt', message: 'right', delayMs: 150 },
|
||||
});
|
||||
await Promise.all([left.accepted, right.accepted]);
|
||||
expect(pool.getState('conversation-left')?.state).toBe('running');
|
||||
expect(pool.getState('conversation-right')?.state).toBe('running');
|
||||
|
||||
await pool.request('conversation-left', { type: 'abort' });
|
||||
await expect.poll(() => pool.getState('conversation-left')?.state).toBe('idle');
|
||||
expect(pool.getState('conversation-right')?.state).toBe('running');
|
||||
await expect.poll(() => pool.getState('conversation-right')?.state).toBe('idle');
|
||||
|
||||
expect(events).toEqual(expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
conversationId: 'conversation-left',
|
||||
event: expect.objectContaining({ type: 'agent_settled', marker: 'left', reason: 'aborted' }),
|
||||
}),
|
||||
expect.objectContaining({
|
||||
conversationId: 'conversation-right',
|
||||
event: expect.objectContaining({ type: 'agent_settled', marker: 'right', reason: 'completed' }),
|
||||
}),
|
||||
]));
|
||||
expect(events.some(({ conversationId, event }) => (
|
||||
conversationId === 'conversation-right' && event.reason === 'aborted'
|
||||
))).toBe(false);
|
||||
} finally {
|
||||
unsubscribe();
|
||||
await pool.shutdown();
|
||||
}
|
||||
expect(processBudget.activeCount).toBe(0);
|
||||
}, 10_000);
|
||||
});
|
||||
531
tests/unit/pi-worker-pool.test.ts
Normal file
531
tests/unit/pi-worker-pool.test.ts
Normal file
@@ -0,0 +1,531 @@
|
||||
// @vitest-environment node
|
||||
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import type { PrepareConversationInput } from '../../electron/coding-runtime/contracts';
|
||||
import {
|
||||
PiProcessBudget,
|
||||
PiWorkerPool,
|
||||
type PiConversationWorker,
|
||||
type PiWorkerOpenResult,
|
||||
} from '../../electron/coding-runtime/pi/worker-pool';
|
||||
import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
|
||||
import { PiProcessError as PiProcessFailure } from '../../electron/coding-runtime/pi/process-errors';
|
||||
import type { PiRpcCommand, PiRpcEvent } from '../../electron/coding-runtime/pi/rpc-client';
|
||||
import type { PiRuntimeTelemetryEvent } from '../../electron/coding-runtime/pi/telemetry';
|
||||
|
||||
const MODEL = {
|
||||
model: {
|
||||
accountId: 'account-a',
|
||||
modelId: 'model-a',
|
||||
thinkingLevel: 'medium' as const,
|
||||
},
|
||||
modelResolution: 'resolved' as const,
|
||||
};
|
||||
|
||||
function conversation(conversationId: string): PrepareConversationInput {
|
||||
return {
|
||||
conversationId,
|
||||
projectId: 'project-a',
|
||||
agentId: 'agent-a',
|
||||
title: conversationId,
|
||||
model: MODEL,
|
||||
};
|
||||
}
|
||||
|
||||
function deferred(): { promise: Promise<void>; resolve(): void } {
|
||||
let resolve!: () => void;
|
||||
const promise = new Promise<void>((done) => { resolve = done; });
|
||||
return { promise, resolve };
|
||||
}
|
||||
|
||||
class FakeWorker implements PiConversationWorker {
|
||||
readonly generation = 1;
|
||||
readonly requests: PiRpcCommand[] = [];
|
||||
stopped = false;
|
||||
private readonly eventListeners = new Set<(event: PiRpcEvent) => void>();
|
||||
private readonly invalidationListeners = new Set<(error: PiProcessError) => void>();
|
||||
|
||||
constructor(readonly id: string) {}
|
||||
|
||||
async request() {
|
||||
this.requests.push(arguments[0] as PiRpcCommand);
|
||||
return { type: 'response' as const, id: 'fake', success: true };
|
||||
}
|
||||
|
||||
subscribe(listener: (event: PiRpcEvent) => void): () => void {
|
||||
this.eventListeners.add(listener);
|
||||
return () => this.eventListeners.delete(listener);
|
||||
}
|
||||
|
||||
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void {
|
||||
this.invalidationListeners.add(listener);
|
||||
return () => this.invalidationListeners.delete(listener);
|
||||
}
|
||||
|
||||
emit(event: PiRpcEvent): void {
|
||||
for (const listener of this.eventListeners) listener(event);
|
||||
}
|
||||
|
||||
invalidate(error = new PiProcessFailure('PI_RPC_EXITED', 'fake worker crashed')): void {
|
||||
for (const listener of this.invalidationListeners) listener(error);
|
||||
}
|
||||
|
||||
async stop() {
|
||||
this.stopped = true;
|
||||
return { mode: 'stdin-close' as const, code: 0, signal: null };
|
||||
}
|
||||
}
|
||||
|
||||
describe('Pi worker pool', () => {
|
||||
it('single-flights prepare per Conversation and never shares its worker with another Conversation', async () => {
|
||||
const gate = deferred();
|
||||
const opened: string[] = [];
|
||||
const pool = new PiWorkerPool({
|
||||
openWorker: async ({ conversation: input }): Promise<PiWorkerOpenResult> => {
|
||||
opened.push(input.conversationId);
|
||||
await gate.promise;
|
||||
return {
|
||||
worker: new FakeWorker(`worker-${input.conversationId}`),
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
maxIdle: 6,
|
||||
});
|
||||
|
||||
const first = pool.prepare(conversation('conversation-a'));
|
||||
const duplicate = pool.prepare(conversation('conversation-a'));
|
||||
const other = pool.prepare(conversation('conversation-b'));
|
||||
await expect.poll(() => opened).toEqual(['conversation-a', 'conversation-b']);
|
||||
gate.resolve();
|
||||
|
||||
const [firstState, duplicateState, otherState] = await Promise.all([first, duplicate, other]);
|
||||
expect(firstState).toEqual(duplicateState);
|
||||
expect(firstState).toMatchObject({
|
||||
conversationId: 'conversation-a',
|
||||
workerId: 'worker-conversation-a',
|
||||
state: 'ready',
|
||||
generation: 1,
|
||||
});
|
||||
expect(otherState).toMatchObject({
|
||||
conversationId: 'conversation-b',
|
||||
workerId: 'worker-conversation-b',
|
||||
state: 'ready',
|
||||
generation: 1,
|
||||
});
|
||||
});
|
||||
|
||||
it('starts only four top-level runs and advances the remaining queue fairly on agent_settled', async () => {
|
||||
const workers = new Map<string, FakeWorker>();
|
||||
const pool = new PiWorkerPool({
|
||||
maxIdle: 6,
|
||||
openWorker: async ({ conversation: input }) => {
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}`);
|
||||
workers.set(input.conversationId, worker);
|
||||
return {
|
||||
worker,
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
const ids = ['a', 'b', 'c', 'd', 'e', 'f'].map((id) => `conversation-${id}`);
|
||||
await Promise.all(ids.map((id) => pool.prepare(conversation(id))));
|
||||
|
||||
const runs = ids.map((conversationId, index) => pool.startTopLevel({
|
||||
conversationId,
|
||||
runId: `run-${index + 1}`,
|
||||
command: { type: 'prompt', message: conversationId },
|
||||
}));
|
||||
|
||||
expect(runs.map((run) => run.queuePosition)).toEqual([undefined, undefined, undefined, undefined, 1, 2]);
|
||||
await expect.poll(() => ids.map((id) => workers.get(id)!.requests.length))
|
||||
.toEqual([1, 1, 1, 1, 0, 0]);
|
||||
workers.get('conversation-a')!.emit({ type: 'agent_end' });
|
||||
expect(workers.get('conversation-e')!.requests).toHaveLength(0);
|
||||
|
||||
workers.get('conversation-a')!.emit({ type: 'agent_settled' });
|
||||
await expect.poll(() => workers.get('conversation-e')!.requests.length).toBe(1);
|
||||
expect(workers.get('conversation-f')!.requests).toHaveLength(0);
|
||||
workers.get('conversation-b')!.emit({ type: 'agent_settled' });
|
||||
await expect.poll(() => workers.get('conversation-f')!.requests.length).toBe(1);
|
||||
});
|
||||
|
||||
it('evicts only the least-recent idle worker and reopens it on demand', async () => {
|
||||
const workers = new Map<string, FakeWorker[]>();
|
||||
const pool = new PiWorkerPool({
|
||||
maxIdle: 2,
|
||||
openWorker: async ({ conversation: input }) => {
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}-${(workers.get(input.conversationId)?.length ?? 0) + 1}`);
|
||||
workers.set(input.conversationId, [...(workers.get(input.conversationId) ?? []), worker]);
|
||||
return {
|
||||
worker,
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
await pool.prepare(conversation('conversation-a'));
|
||||
await pool.prepare(conversation('conversation-b'));
|
||||
await pool.prepare(conversation('conversation-c'));
|
||||
await expect.poll(() => workers.get('conversation-a')![0]!.stopped).toBe(true);
|
||||
expect(pool.getState('conversation-a')).toBeNull();
|
||||
expect(pool.getState('conversation-b')?.state).toBe('ready');
|
||||
expect(pool.getState('conversation-c')?.state).toBe('ready');
|
||||
|
||||
const reopened = await pool.prepare(conversation('conversation-a'));
|
||||
expect(reopened).toMatchObject({
|
||||
workerId: 'worker-conversation-a-2',
|
||||
generation: 2,
|
||||
state: 'ready',
|
||||
});
|
||||
});
|
||||
|
||||
it('shares a fair total-process budget and starts the next worker only after a lease is released', async () => {
|
||||
const processBudget = new PiProcessBudget(2);
|
||||
const opened: string[] = [];
|
||||
const pool = new PiWorkerPool({
|
||||
maxIdle: 3,
|
||||
processBudget,
|
||||
openWorker: async ({ conversation: input }) => {
|
||||
opened.push(input.conversationId);
|
||||
return {
|
||||
worker: new FakeWorker(`worker-${input.conversationId}`),
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
await Promise.all([
|
||||
pool.prepare(conversation('conversation-a')),
|
||||
pool.prepare(conversation('conversation-b')),
|
||||
]);
|
||||
const third = pool.prepare(conversation('conversation-c'));
|
||||
await expect.poll(() => processBudget.waitingCount).toBe(1);
|
||||
expect(opened).toEqual(['conversation-a', 'conversation-b']);
|
||||
|
||||
await pool.dispose('conversation-a');
|
||||
await expect(third).resolves.toMatchObject({ conversationId: 'conversation-c', state: 'ready' });
|
||||
expect(opened).toEqual(['conversation-a', 'conversation-b', 'conversation-c']);
|
||||
expect(processBudget.activeCount).toBe(2);
|
||||
|
||||
await pool.shutdown();
|
||||
expect(processBudget.activeCount).toBe(0);
|
||||
});
|
||||
|
||||
it('never evicts a running worker when the warm-idle LRU exceeds its cap', async () => {
|
||||
const workers = new Map<string, FakeWorker>();
|
||||
const pool = new PiWorkerPool({
|
||||
maxIdle: 1,
|
||||
openWorker: async ({ conversation: input }) => {
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}`);
|
||||
workers.set(input.conversationId, worker);
|
||||
return {
|
||||
worker,
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
await pool.prepare(conversation('conversation-running'));
|
||||
const running = pool.startTopLevel({
|
||||
conversationId: 'conversation-running',
|
||||
runId: 'run-running',
|
||||
command: { type: 'prompt', message: 'keep alive' },
|
||||
});
|
||||
await running.accepted;
|
||||
await pool.prepare(conversation('conversation-old-idle'));
|
||||
await pool.prepare(conversation('conversation-new-idle'));
|
||||
|
||||
expect(workers.get('conversation-running')!.stopped).toBe(false);
|
||||
expect(pool.getState('conversation-running')).toMatchObject({ state: 'running' });
|
||||
expect(workers.get('conversation-old-idle')!.stopped).toBe(true);
|
||||
expect(pool.getState('conversation-new-idle')).toMatchObject({ state: 'ready' });
|
||||
});
|
||||
|
||||
it('cleans only the crashed generation and releases its permit for the next Conversation', async () => {
|
||||
const workers = new Map<string, FakeWorker>();
|
||||
const pool = new PiWorkerPool({
|
||||
maxRunning: 2,
|
||||
maxIdle: 3,
|
||||
openWorker: async ({ conversation: input }) => {
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}`);
|
||||
workers.set(input.conversationId, worker);
|
||||
return {
|
||||
worker,
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
await Promise.all(['a', 'b', 'c'].map((id) => pool.prepare(conversation(`conversation-${id}`))));
|
||||
pool.startTopLevel({ conversationId: 'conversation-a', runId: 'run-a', command: { type: 'prompt', message: 'a' } });
|
||||
pool.startTopLevel({ conversationId: 'conversation-b', runId: 'run-b', command: { type: 'prompt', message: 'b' } });
|
||||
pool.startTopLevel({ conversationId: 'conversation-c', runId: 'run-c', command: { type: 'prompt', message: 'c' } });
|
||||
await expect.poll(() => workers.get('conversation-a')!.requests.length).toBe(1);
|
||||
|
||||
const cancelled: string[] = [];
|
||||
for (const kind of ['command', 'interaction', 'child'] as const) {
|
||||
pool.trackGenerationResource({
|
||||
conversationId: 'conversation-a',
|
||||
kind,
|
||||
id: `${kind}-a`,
|
||||
cancel: () => { cancelled.push(kind); },
|
||||
});
|
||||
}
|
||||
pool.trackGenerationResource({
|
||||
conversationId: 'conversation-b',
|
||||
kind: 'child',
|
||||
id: 'child-b',
|
||||
cancel: () => { cancelled.push('other'); },
|
||||
});
|
||||
|
||||
workers.get('conversation-a')!.invalidate();
|
||||
|
||||
expect(cancelled.sort()).toEqual(['child', 'command', 'interaction']);
|
||||
expect(pool.getState('conversation-a')).toMatchObject({ state: 'crashed', generation: 1 });
|
||||
expect(pool.getState('conversation-b')).toMatchObject({ state: 'running', generation: 1 });
|
||||
await expect.poll(() => workers.get('conversation-c')!.requests.length).toBe(1);
|
||||
expect(cancelled).not.toContain('other');
|
||||
});
|
||||
|
||||
it('rebuilds stale idle workers before prompt and lets running workers settle first', async () => {
|
||||
const workers = new Map<string, FakeWorker[]>();
|
||||
const revisions: Array<{ conversationId: string; provider: number; resources: number }> = [];
|
||||
const pool = new PiWorkerPool({
|
||||
maxIdle: 4,
|
||||
openWorker: async ({ conversation: input, revision }) => {
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}-${(workers.get(input.conversationId)?.length ?? 0) + 1}`);
|
||||
workers.set(input.conversationId, [...(workers.get(input.conversationId) ?? []), worker]);
|
||||
revisions.push({ conversationId: input.conversationId, ...revision });
|
||||
return {
|
||||
worker,
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
await Promise.all([
|
||||
pool.prepare(conversation('conversation-running')),
|
||||
pool.prepare(conversation('conversation-idle')),
|
||||
]);
|
||||
pool.startTopLevel({
|
||||
conversationId: 'conversation-running',
|
||||
runId: 'run-running',
|
||||
command: { type: 'prompt', message: 'running' },
|
||||
});
|
||||
await expect.poll(() => workers.get('conversation-running')![0]!.requests.length).toBe(1);
|
||||
|
||||
pool.markProviderStale();
|
||||
pool.startTopLevel({
|
||||
conversationId: 'conversation-idle',
|
||||
runId: 'run-idle',
|
||||
command: { type: 'prompt', message: 'idle' },
|
||||
});
|
||||
|
||||
await expect.poll(() => workers.get('conversation-idle')?.length).toBe(2);
|
||||
expect(workers.get('conversation-idle')![0]!.stopped).toBe(true);
|
||||
await expect.poll(() => workers.get('conversation-idle')![1]!.requests.length).toBe(1);
|
||||
expect(workers.get('conversation-running')).toHaveLength(1);
|
||||
expect(workers.get('conversation-running')![0]!.stopped).toBe(false);
|
||||
|
||||
workers.get('conversation-running')![0]!.emit({ type: 'agent_settled' });
|
||||
await expect.poll(() => workers.get('conversation-running')?.length).toBe(2);
|
||||
expect(workers.get('conversation-running')![0]!.stopped).toBe(true);
|
||||
expect(revisions).toEqual(expect.arrayContaining([
|
||||
{ conversationId: 'conversation-idle', provider: 2, resources: 1 },
|
||||
{ conversationId: 'conversation-running', provider: 2, resources: 1 },
|
||||
]));
|
||||
});
|
||||
|
||||
it('rejects queued work and stops every parent worker during app shutdown', async () => {
|
||||
const workers: FakeWorker[] = [];
|
||||
const pool = new PiWorkerPool({
|
||||
maxRunning: 1,
|
||||
maxIdle: 3,
|
||||
openWorker: async ({ conversation: input }) => {
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}`);
|
||||
workers.push(worker);
|
||||
return {
|
||||
worker,
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
await Promise.all(['a', 'b', 'c'].map((id) => pool.prepare(conversation(`conversation-${id}`))));
|
||||
pool.startTopLevel({ conversationId: 'conversation-a', runId: 'run-a', command: { type: 'prompt', message: 'a' } });
|
||||
const queued = pool.startTopLevel({
|
||||
conversationId: 'conversation-b',
|
||||
runId: 'run-b',
|
||||
command: { type: 'prompt', message: 'b' },
|
||||
});
|
||||
let childCancelled = false;
|
||||
pool.trackGenerationResource({
|
||||
conversationId: 'conversation-a',
|
||||
kind: 'child',
|
||||
id: 'child-a',
|
||||
cancel: () => { childCancelled = true; },
|
||||
});
|
||||
|
||||
await pool.shutdown();
|
||||
|
||||
await expect(queued.accepted).rejects.toThrow('shutting down');
|
||||
expect(childCancelled).toBe(true);
|
||||
expect(workers.every((worker) => worker.stopped)).toBe(true);
|
||||
expect(pool.getState('conversation-a')).toBeNull();
|
||||
await expect(pool.prepare(conversation('conversation-after-quit')))
|
||||
.rejects.toThrow('shutting down');
|
||||
});
|
||||
|
||||
it('waits for an in-flight fork open and stops that process before shutdown completes', async () => {
|
||||
const forkGate = deferred();
|
||||
const workers: FakeWorker[] = [];
|
||||
let forkOpenStarted = false;
|
||||
const pool = new PiWorkerPool({
|
||||
maxIdle: 3,
|
||||
openWorker: async ({ conversation: input, existingSession }) => {
|
||||
if (input.conversationId === 'conversation-fork') {
|
||||
forkOpenStarted = true;
|
||||
await forkGate.promise;
|
||||
}
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}`);
|
||||
workers.push(worker);
|
||||
return {
|
||||
worker,
|
||||
session: existingSession ?? {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
await pool.prepare(conversation('conversation-source'));
|
||||
const fork = pool.fork('conversation-source', conversation('conversation-fork'));
|
||||
await expect.poll(() => forkOpenStarted).toBe(true);
|
||||
const forkOutcome = fork.then(
|
||||
() => 'resolved',
|
||||
(error: unknown) => error instanceof Error ? error.message : String(error),
|
||||
);
|
||||
let shutdownCompleted = false;
|
||||
const shutdown = pool.shutdown().then(() => { shutdownCompleted = true; });
|
||||
await Promise.resolve();
|
||||
expect(shutdownCompleted).toBe(false);
|
||||
|
||||
forkGate.resolve();
|
||||
await shutdown;
|
||||
expect(await forkOutcome).toContain('shutting down');
|
||||
expect(workers).toHaveLength(2);
|
||||
expect(workers.every((worker) => worker.stopped)).toBe(true);
|
||||
});
|
||||
|
||||
it('recovers the target session with a new generation and disposes no sibling worker', async () => {
|
||||
const workers = new Map<string, FakeWorker[]>();
|
||||
const pool = new PiWorkerPool({
|
||||
maxIdle: 4,
|
||||
openWorker: async ({ conversation: input, existingSession }) => {
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}-${(workers.get(input.conversationId)?.length ?? 0) + 1}`);
|
||||
workers.set(input.conversationId, [...(workers.get(input.conversationId) ?? []), worker]);
|
||||
return {
|
||||
worker,
|
||||
session: existingSession ?? {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
await Promise.all([
|
||||
pool.prepare(conversation('conversation-a')),
|
||||
pool.prepare(conversation('conversation-b')),
|
||||
]);
|
||||
workers.get('conversation-a')![0]!.invalidate();
|
||||
|
||||
const recovered = await pool.recover('conversation-a');
|
||||
expect(recovered).toMatchObject({
|
||||
workerId: 'worker-conversation-a-2',
|
||||
generation: 2,
|
||||
state: 'ready',
|
||||
session: { piSessionId: 'session-conversation-a', sessionKey: 'key-conversation-a' },
|
||||
});
|
||||
expect(workers.get('conversation-a')![0]!.stopped).toBe(true);
|
||||
expect(workers.get('conversation-b')![0]!.stopped).toBe(false);
|
||||
|
||||
await pool.dispose('conversation-a');
|
||||
expect(workers.get('conversation-a')![1]!.stopped).toBe(true);
|
||||
expect(pool.getState('conversation-a')).toBeNull();
|
||||
expect(pool.getState('conversation-b')).toMatchObject({ state: 'ready', generation: 1 });
|
||||
});
|
||||
|
||||
it('records privacy-safe queue wait and RPC prompt acceptance spans', async () => {
|
||||
let now = 0;
|
||||
const telemetry: PiRuntimeTelemetryEvent[] = [];
|
||||
const workers = new Map<string, FakeWorker>();
|
||||
const pool = new PiWorkerPool({
|
||||
maxRunning: 1,
|
||||
maxIdle: 2,
|
||||
now: () => now,
|
||||
onTelemetry: (event) => telemetry.push(event),
|
||||
openWorker: async ({ conversation: input }) => {
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}`);
|
||||
workers.set(input.conversationId, worker);
|
||||
return {
|
||||
worker,
|
||||
session: {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
await Promise.all([
|
||||
pool.prepare(conversation('f47ac10b-58cc-4372-a567-0e02b2c3d479')),
|
||||
pool.prepare(conversation('8b1a9953-c461-4d88-9c3e-7e1f8f3f2c11')),
|
||||
]);
|
||||
const first = pool.startTopLevel({
|
||||
conversationId: 'f47ac10b-58cc-4372-a567-0e02b2c3d479',
|
||||
runId: 'run-first-1234567890',
|
||||
command: { type: 'prompt', message: 'private first prompt' },
|
||||
});
|
||||
now = 5;
|
||||
const second = pool.startTopLevel({
|
||||
conversationId: '8b1a9953-c461-4d88-9c3e-7e1f8f3f2c11',
|
||||
runId: 'run-second-1234567890',
|
||||
command: { type: 'prompt', message: 'private second prompt' },
|
||||
});
|
||||
await first.accepted;
|
||||
now = 25;
|
||||
workers.get('f47ac10b-58cc-4372-a567-0e02b2c3d479')!.emit({ type: 'agent_settled' });
|
||||
await second.accepted;
|
||||
|
||||
expect(telemetry.map(({ milestone }) => milestone)).toEqual([
|
||||
'prompt.accepted',
|
||||
'worker.queue_wait',
|
||||
'prompt.accepted',
|
||||
]);
|
||||
expect(telemetry[1]).toMatchObject({ durationMs: 20, workerGeneration: 1, cold: true });
|
||||
const serialized = JSON.stringify(telemetry);
|
||||
expect(serialized).not.toContain('private first prompt');
|
||||
expect(serialized).not.toContain('private second prompt');
|
||||
expect(serialized).not.toContain('f47ac10b-58cc-4372-a567-0e02b2c3d479');
|
||||
expect(serialized).not.toContain('8b1a9953-c461-4d88-9c3e-7e1f8f3f2c11');
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user