Files
makelore/tests/unit/pi-background-lifecycle.test.ts

418 lines
15 KiB
TypeScript

// @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 { 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';
import type { PiWorkerStopReason } from '../../electron/coding-runtime/pi/worker-process';
import { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
import { BackgroundLifecycleController } from '../../electron/main/background-lifecycle';
const roots: string[] = [];
const NOW = '2026-08-25T12:00:00.000Z';
class LifecycleFakeWorker implements PiConversationWorker {
readonly requests: PiRpcCommand[] = [];
readonly stopReasons: PiWorkerStopReason[] = [];
private readonly events = new Set<(event: PiRpcEvent) => void>();
private readonly invalidations = new Set<(error: PiProcessError) => void>();
private timeoutType: string | null = null;
private deferredType: string | null = null;
private deferredReject: ((error: unknown) => void) | null = null;
private compacted = false;
private lateResult: PiRpcRequestOptions['onLateResult'];
constructor(readonly id: string, readonly generation: number) {}
async request<T = unknown>(
command: PiRpcCommand,
options?: PiRpcRequestOptions,
): Promise<PiRpcResponse<T>> {
this.requests.push(structuredClone(command));
if (this.timeoutType === command.type) {
this.timeoutType = null;
this.lateResult = options?.onLateResult;
await new Promise<void>((_resolve, reject) => {
setTimeout(() => reject(new PiProcessError(
'PI_RPC_TIMEOUT',
`fake ${command.type} confirmation timeout`,
{ generation: this.generation },
)), 10_000);
});
}
if (this.deferredType === command.type) {
this.deferredType = null;
await new Promise<void>((_resolve, reject) => {
this.deferredReject = reject;
});
}
const data = command.type === 'get_state'
? {
sessionId: `session-${this.id}`,
thinkingLevel: 'medium',
isStreaming: false,
isCompacting: false,
pendingMessageCount: 0,
}
: command.type === 'get_entries'
? this.compacted
? {
entries: [
{
type: 'message',
id: 'entry-kept',
parentId: null,
timestamp: NOW,
message: { role: 'user', content: 'retained', timestamp: 1 },
},
{
type: 'compaction',
id: 'entry-compaction-proof',
parentId: 'entry-kept',
timestamp: NOW,
summary: 'not exposed',
firstKeptEntryId: 'entry-kept',
tokensBefore: 5_000,
},
],
leafId: 'entry-compaction-proof',
}
: { entries: [], leafId: null }
: command.type === 'get_session_stats'
? {
contextUsage: { tokens: 0, contextWindow: 100_000, percent: 0 },
tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
}
: command.type === 'get_commands'
? { commands: [] }
: command.type === 'get_available_thinking_levels'
? { levels: ['off', 'medium', 'high'] }
: undefined;
return {
type: 'response',
id: `${this.id}-${this.requests.length}`,
success: true,
...(data === undefined ? {} : { data: structuredClone(data) as T }),
};
}
async send(command: PiRpcCommand): Promise<void> {
this.requests.push(structuredClone(command));
}
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 {
if (event.type === 'compaction_end' && event.aborted !== true && typeof event.errorMessage !== 'string') {
this.compacted = true;
}
for (const listener of this.events) listener(event);
}
timeoutNext(type: string): void {
this.timeoutType = type;
}
deferNext(type: string): void {
this.deferredType = type;
}
rejectDeferred(error: unknown): void {
this.deferredReject?.(error);
this.deferredReject = null;
}
completeLateSuccess(): void {
this.lateResult?.({
response: { type: 'response', id: 'late-proof', success: true },
});
this.lateResult = undefined;
}
async stop(reason: PiWorkerStopReason) {
this.stopReasons.push(reason);
return { mode: 'stdin-close' as const, code: 0, signal: null };
}
}
afterEach(async () => {
vi.useRealTimers();
await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
});
describe('Pi run background lifecycle lease', () => {
it('keeps confirmation-uncertain and queued work alive while hidden, then evicts idle workers', async () => {
vi.useFakeTimers();
const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-background-'));
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',
];
return () => ids.shift() as string;
})(),
now: () => NOW,
});
const runningConversation = await store.create({
agentId: 'agent-a',
title: 'Running lease proof',
model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
modelResolution: 'resolved',
});
const queuedConversation = await store.create({
agentId: 'agent-a',
title: 'Queued lease proof',
model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
modelResolution: 'resolved',
});
const inputs = [runningConversation, queuedConversation].map((conversation) => ({
conversationId: conversation.id,
projectId: 'project-a',
agentId: 'agent-a',
title: conversation.title,
model: { model: conversation.model, modelResolution: conversation.modelResolution },
}));
const workers: LifecycleFakeWorker[] = [];
const workersByConversation = new Map<string, LifecycleFakeWorker>();
const pool = new PiWorkerPool({
maxRunning: 1,
openWorker: async ({ conversation: workerConversation, generation, existingSession }) => {
const worker = new LifecycleFakeWorker(
`worker-${workerConversation.conversationId}-${generation}`,
generation,
);
workers.push(worker);
workersByConversation.set(workerConversation.conversationId, worker);
return {
worker,
session: existingSession ?? {
piSessionId: `session-${workerConversation.conversationId}`,
sessionKey: `key-${workerConversation.conversationId}`,
},
};
},
});
let runtime!: PiConversationRuntime;
const onSleep = vi.fn().mockResolvedValue(undefined);
const onStopRuntime = vi.fn(async () => {
await Promise.all(runtime.getDiagnostics().workers.map(async ({ conversationId }) => {
await runtime.dispose(conversationId, 'background_sleep');
}));
});
const controller = new BackgroundLifecycleController({
idleStopMs: 100,
onSleep,
onStopRuntime,
});
const acquiredLeaseIds: string[] = [];
const releasedLeaseIds: string[] = [];
runtime = new PiConversationRuntime({
pool,
registry: new PiSessionRegistry({ projectStore }),
createId: (kind) => `${kind}-lease-proof`,
acquireBackgroundLease: (lease) => {
acquiredLeaseIds.push(lease.id);
const release = controller.acquireLease(lease);
return () => {
releasedLeaseIds.push(lease.id);
release();
};
},
});
await Promise.all(inputs.map(async (input) => await runtime.prepare(input)));
workersByConversation.get(runningConversation.id)!.timeoutNext('prompt');
const uncertain = runtime.prompt({
clientRequestId: 'request-1',
conversationId: runningConversation.id,
mode: 'prompt',
text: 'Keep working while hidden',
attachments: [],
});
void uncertain.catch(() => undefined);
await vi.advanceTimersByTimeAsync(10_000);
await expect(uncertain).rejects.toMatchObject({
publicError: { code: 'CODING_REQUEST_UNCERTAIN' },
});
expect(runtime.getResilienceProofDiagnostics()).toMatchObject({
backgroundLeases: { active: 1 },
pool: { runs: { active: 1, waiting: 0 } },
});
expect((await runtime.getSnapshot(runningConversation.id)).run).toMatchObject({
status: 'running',
error: { code: 'CODING_REQUEST_UNCERTAIN' },
});
await expect(runtime.compact(runningConversation.id)).rejects.toMatchObject({
publicError: { code: 'CODING_REQUEST_UNCERTAIN' },
});
expect(workersByConversation.get(runningConversation.id)!.requests.filter(
({ type }) => type === 'compact',
)).toHaveLength(0);
workersByConversation.get(runningConversation.id)!.completeLateSuccess();
await expect.poll(async () => (
await runtime.getSnapshot(runningConversation.id)
).run.error).toBeUndefined();
expect(controller.getLeaseCount()).toBe(1);
const queued = await runtime.prompt({
clientRequestId: 'request-2',
conversationId: queuedConversation.id,
mode: 'prompt',
text: 'Wait safely while hidden',
attachments: [],
});
expect(queued).toMatchObject({ accepted: true, queuePosition: 1 });
expect(controller.getLeaseCount()).toBe(2);
const releaseChild = pool.trackGenerationResource({
conversationId: runningConversation.id,
kind: 'child',
id: 'child-proof',
cancel: () => undefined,
});
releaseChild();
expect(controller.getLeaseCount()).toBe(2);
controller.setActivity({ visible: false, module: 'programming' });
await vi.advanceTimersByTimeAsync(100);
expect(onStopRuntime).not.toHaveBeenCalled();
expect(workers.every(({ stopReasons }) => stopReasons.length === 0)).toBe(true);
expect((await runtime.getSnapshot(runningConversation.id)).run.status).toBe('running');
expect((await runtime.getSnapshot(queuedConversation.id)).run.status).toBe('queued');
workersByConversation.get(runningConversation.id)!.emit({ type: 'agent_settled' });
expect((await runtime.getSnapshot(runningConversation.id)).run.status).toBe('idle');
await expect.poll(async () => (
await runtime.getSnapshot(queuedConversation.id)
).run.status).toBe('running');
expect(controller.getLeaseCount()).toBe(1);
workersByConversation.get(queuedConversation.id)!.emit({ type: 'agent_settled' });
expect((await runtime.getSnapshot(queuedConversation.id)).run.status).toBe('idle');
expect(controller.getLeaseCount()).toBe(0);
expect(releasedLeaseIds.sort()).toEqual(acquiredLeaseIds.sort());
expect(new Set(releasedLeaseIds).size).toBe(releasedLeaseIds.length);
workersByConversation.get(runningConversation.id)!.timeoutNext('compact');
const compact = runtime.compact(runningConversation.id);
void compact.catch(() => undefined);
await vi.advanceTimersByTimeAsync(10_000);
await expect(compact).rejects.toMatchObject({
publicError: { code: 'CODING_REQUEST_UNCERTAIN' },
});
expect(controller.getLeaseCount()).toBe(1);
expect(onStopRuntime).not.toHaveBeenCalled();
expect((await runtime.getSnapshot(runningConversation.id)).run).toMatchObject({
status: 'compacting',
error: { code: 'CODING_REQUEST_UNCERTAIN' },
});
workersByConversation.get(runningConversation.id)!.emit({
type: 'compaction_start',
reason: 'manual',
});
workersByConversation.get(runningConversation.id)!.emit({
type: 'compaction_end',
reason: 'manual',
result: {
summary: 'controlled summary',
firstKeptEntryId: 'entry-kept',
tokensBefore: 5_000,
},
aborted: false,
willRetry: false,
});
expect((await runtime.getSnapshot(runningConversation.id)).context.compaction).toBe('idle');
workersByConversation.get(runningConversation.id)!.completeLateSuccess();
await expect.poll(async () => (
await runtime.getSnapshot(runningConversation.id)
).run.status).toBe('idle');
expect((await runtime.getSnapshot(runningConversation.id)).nodes).toContainEqual(
expect.objectContaining({ kind: 'compaction', status: 'complete' }),
);
expect(controller.getLeaseCount()).toBe(0);
workersByConversation.get(runningConversation.id)!.deferNext('compact');
const rejectedCompact = runtime.compact(runningConversation.id);
void rejectedCompact.catch(() => undefined);
await vi.advanceTimersByTimeAsync(0);
workersByConversation.get(runningConversation.id)!.emit({
type: 'compaction_start',
reason: 'manual',
});
workersByConversation.get(runningConversation.id)!.emit({
type: 'compaction_end',
reason: 'manual',
aborted: false,
willRetry: false,
errorMessage: 'Compaction failed: controlled rejection',
});
await vi.advanceTimersByTimeAsync(0);
expect((await runtime.getSnapshot(runningConversation.id)).run.status).toBe('error');
workersByConversation.get(runningConversation.id)!.rejectDeferred(new PiProcessError(
'PI_RPC_RESPONSE_ERROR',
'controlled compact rejection',
{ generation: 1 },
));
await expect(rejectedCompact).rejects.toBeInstanceOf(PiProcessError);
expect(controller.getLeaseCount()).toBe(0);
await vi.advanceTimersByTimeAsync(100);
expect(onSleep).toHaveBeenCalledTimes(1);
expect(onStopRuntime).toHaveBeenCalledTimes(1);
expect(workers.map(({ stopReasons }) => stopReasons)).toEqual([
['background_sleep'],
['background_sleep'],
]);
controller.dispose();
});
});