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' },
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user