fix(pi): keep active runs alive during background sleep

This commit is contained in:
2026-08-25 20:56:17 +08:00
parent 92f4c91088
commit 621ebb1781
22 changed files with 1061 additions and 68 deletions

View File

@@ -19,6 +19,7 @@ import {
PiWorkerPool,
type PiConversationWorker,
} from '../../electron/coding-runtime/pi/worker-pool';
import type { PiWorkerStopReason } from '../../electron/coding-runtime/pi/worker-process';
import { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
import type {
PiRpcCommand,
@@ -33,6 +34,7 @@ const NOW = '2026-08-22T15:00:00.000Z';
class RuntimeFakeWorker implements PiConversationWorker {
readonly requests: PiRpcCommand[] = [];
readonly stopReasons: PiWorkerStopReason[] = [];
private stateData: unknown;
private entriesData: unknown = { entries: [], leafId: null };
private statsData: unknown = {
@@ -143,7 +145,8 @@ class RuntimeFakeWorker implements PiConversationWorker {
for (const listener of this.invalidations) listener(error);
}
async stop() {
async stop(reason: PiWorkerStopReason) {
this.stopReasons.push(reason);
return { mode: 'stdin-close' as const, code: 0, signal: null };
}
}
@@ -236,11 +239,25 @@ describe('Pi Conversation runtime', () => {
},
});
const trackingHost = new TrackingExtensionHost();
let acquiredBackgroundLeases = 0;
let releasedBackgroundLeases = 0;
let activeBackgroundLeases = 0;
const runtime = new PiConversationRuntime({
pool,
registry: new PiSessionRegistry({ projectStore }),
createId: (kind) => `${kind}-fixed`,
extensionHost: trackingHost,
acquireBackgroundLease: () => {
acquiredBackgroundLeases += 1;
activeBackgroundLeases += 1;
let released = false;
return () => {
if (released) throw new Error('background lease released twice');
released = true;
releasedBackgroundLeases += 1;
activeBackgroundLeases -= 1;
};
},
resolveModel: async (candidate) => {
if (candidate.accountId !== 'account-b'
|| !['model-b', 'model-c'].includes(candidate.modelId)) {
@@ -297,6 +314,7 @@ describe('Pi Conversation runtime', () => {
releasePromptAcceptance();
const accepted = await acceptance;
expect(accepted).toMatchObject({ accepted: true, runId: 'run-fixed', mode: 'prompt' });
expect(activeBackgroundLeases).toBe(1);
expect((await runtime.getSnapshot(left.id)).run.status).toBe('running');
for (const event of PI_084_TEXT_TURN.events) {
workers.get(left.id)!.emit(structuredClone(event));
@@ -372,6 +390,7 @@ describe('Pi Conversation runtime', () => {
workers.get(left.id)!.emit({ type: 'agent_settled' });
await expect.poll(async () => (await runtime.getSnapshot(left.id)).queue.items).toEqual([]);
expect((await runtime.getSnapshot(left.id)).run.status).toBe('idle');
expect(activeBackgroundLeases).toBe(0);
const settledNodes = (await runtime.getSnapshot(left.id)).nodes;
expect(settledNodes).toEqual(checkpoint.nodes);
expect(leftGenerationOneSeqs).toEqual(leftGenerationOneSeqs.map((_, index) => index + 1));
@@ -483,13 +502,16 @@ describe('Pi Conversation runtime', () => {
]);
workers.get(left.id)!.emit({ type: 'agent_settled' });
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('idle');
expect(activeBackgroundLeases).toBe(0);
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');
expect(activeBackgroundLeases).toBe(0);
await runtime.compact(right.id);
expect((await runtime.getSnapshot(right.id)).run.status).toBe('compacting');
expect(activeBackgroundLeases).toBe(1);
expect(workers.get(right.id)!.requests.at(-1)).toEqual({ type: 'compact' });
expect((await runtime.getSnapshot(left.id)).run.status).toBe('idle');
const rightDurable = {
@@ -510,6 +532,7 @@ describe('Pi Conversation runtime', () => {
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');
expect(activeBackgroundLeases).toBe(0);
const leftBeforeFailedRecovery = await runtime.getSnapshot(left.id);
const rightBeforeFailedRecovery = await runtime.getSnapshot(right.id);
@@ -568,7 +591,7 @@ describe('Pi Conversation runtime', () => {
sourceEntryId: 'entry-a',
});
await runtime.dispose(forkTarget.id);
await runtime.dispose(forkTarget.id, 'test_injection');
expect(pool.getState(forkTarget.id)).toBeNull();
expect(pool.getState(left.id)).toMatchObject({ generation: 2, state: 'idle' });
await expect(runtime.getSnapshot(forkTarget.id)).rejects.toMatchObject({
@@ -585,6 +608,7 @@ describe('Pi Conversation runtime', () => {
await runtime.recover(left.id);
expect(trackingHost.runs.has(left.id)).toBe(false);
expect((await runtime.getSnapshot(left.id)).run).toEqual({ status: 'idle' });
expect(activeBackgroundLeases).toBe(0);
workers.get(left.id)!.setSessionData({
thinkingLevels: { levels: ['high'] },
@@ -620,6 +644,7 @@ describe('Pi Conversation runtime', () => {
crashPatchOps.length = 0;
workers.get(left.id)!.invalidate();
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('error');
expect(activeBackgroundLeases).toBe(0);
expect((await runtime.getSnapshot(left.id))).toMatchObject({
run: {
status: 'error',
@@ -669,6 +694,7 @@ describe('Pi Conversation runtime', () => {
'strict JSONL protocol failure',
));
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('error');
expect(activeBackgroundLeases).toBe(0);
expect((await runtime.getSnapshot(left.id)).run).toMatchObject({
terminalReason: 'failed',
error: { code: 'CODING_RUNTIME_PROTOCOL_ERROR', recoverable: true },
@@ -676,6 +702,47 @@ describe('Pi Conversation runtime', () => {
workers.get(left.id)!.emit({ type: 'agent_settled' });
await new Promise((resolve) => setImmediate(resolve));
expect((await runtime.getSnapshot(left.id)).run.status).toBe('error');
await runtime.recover(left.id);
const siblingBeforeForcedDispose = await runtime.getSnapshot(right.id);
const bindingBeforeForcedDispose = await store.get(left.id);
const forcedDisposePatches: string[] = [];
const unsubscribeForcedDispose = runtime.subscribe((envelope) => {
if (envelope.conversationId === left.id && envelope.patch.op === 'run.state') {
forcedDisposePatches.push(JSON.stringify(envelope.patch));
}
});
const forcedWorker = workers.get(left.id)!;
await runtime.prompt({
clientRequestId: 'request-project-deactivated',
conversationId: left.id,
mode: 'prompt',
text: 'Do not replay this accepted prompt',
attachments: [],
});
const forcedPromptCount = forcedWorker.requests.filter(({ type }) => type === 'prompt').length;
await runtime.dispose(left.id, 'project_deactivated');
unsubscribeForcedDispose();
expect(forcedDisposePatches.at(-1)).toContain('CODING_RUNTIME_START_FAILED');
expect(forcedDisposePatches.at(-1)).toContain('"recoverable":true');
expect(forcedWorker.stopReasons.at(-1)).toBe('project_deactivated');
expect(activeBackgroundLeases).toBe(0);
expect(await runtime.getSnapshot(right.id)).toEqual(siblingBeforeForcedDispose);
expect((await store.get(left.id))?.sessionKey).toBe(bindingBeforeForcedDispose?.sessionKey);
const persistedAfterForcedDispose = (await store.get(left.id))!;
await runtime.prepare({
...inputs[0]!,
model: {
model: persistedAfterForcedDispose.model,
modelResolution: persistedAfterForcedDispose.modelResolution,
},
});
expect(forcedWorker.requests.filter(({ type }) => type === 'prompt')).toHaveLength(
forcedPromptCount,
);
expect(acquiredBackgroundLeases).toBe(releasedBackgroundLeases);
expect(activeBackgroundLeases).toBe(0);
unsubscribe();
});
});