fix(pi): converge worker failures and thinking state

This commit is contained in:
2026-08-25 11:45:14 +08:00
parent 274187e3cf
commit 61817b161f
28 changed files with 2019 additions and 118 deletions

View File

@@ -19,7 +19,7 @@ import {
PiWorkerPool,
type PiConversationWorker,
} from '../../electron/coding-runtime/pi/worker-pool';
import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
import { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
import type {
PiRpcCommand,
PiRpcEvent,
@@ -40,6 +40,9 @@ class RuntimeFakeWorker implements PiConversationWorker {
tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};
private commandsData: unknown = { commands: [] };
private thinkingLevelsData: unknown = {
levels: ['off', 'minimal', 'low', 'medium', 'high'],
};
private failType: string | null = null;
private readonly responseGates = new Map<string, Promise<void>>();
private readonly events = new Set<(event: PiRpcEvent) => void>();
@@ -66,6 +69,15 @@ class RuntimeFakeWorker implements PiConversationWorker {
this.failType = null;
throw new Error(`fake ${command.type} rejection`);
}
if (command.type === 'set_thinking_level') {
const levels = (this.thinkingLevelsData as { levels?: unknown }).levels;
if (Array.isArray(levels) && levels.includes(command.level)) {
this.stateData = {
...(this.stateData as Record<string, unknown>),
thinkingLevel: command.level,
};
}
}
const data = command.type === 'get_state'
? this.stateData
: command.type === 'get_entries'
@@ -74,7 +86,9 @@ class RuntimeFakeWorker implements PiConversationWorker {
? this.statsData
: command.type === 'get_commands'
? this.commandsData
: undefined;
: command.type === 'get_available_thinking_levels'
? this.thinkingLevelsData
: undefined;
return {
type: 'response',
id: `${this.id}-${this.requests.length}`,
@@ -87,11 +101,20 @@ class RuntimeFakeWorker implements PiConversationWorker {
this.requests.push(structuredClone(command));
}
setSessionData(input: { state?: unknown; entries?: unknown; stats?: unknown; commands?: unknown }): void {
setSessionData(input: {
state?: unknown;
entries?: unknown;
stats?: unknown;
commands?: unknown;
thinkingLevels?: unknown;
}): void {
if (input.state !== undefined) this.stateData = structuredClone(input.state);
if (input.entries !== undefined) this.entriesData = structuredClone(input.entries);
if (input.stats !== undefined) this.statsData = structuredClone(input.stats);
if (input.commands !== undefined) this.commandsData = structuredClone(input.commands);
if (input.thinkingLevels !== undefined) {
this.thinkingLevelsData = structuredClone(input.thinkingLevels);
}
}
failNext(type: string): void { this.failType = type; }
@@ -116,6 +139,10 @@ class RuntimeFakeWorker implements PiConversationWorker {
for (const listener of this.events) listener(event);
}
invalidate(error = new PiProcessError('PI_RPC_EXITED', 'fake worker crashed')): void {
for (const listener of this.invalidations) listener(error);
}
async stop() {
return { mode: 'stdin-close' as const, code: 0, signal: null };
}
@@ -215,7 +242,8 @@ describe('Pi Conversation runtime', () => {
createId: (kind) => `${kind}-fixed`,
extensionHost: trackingHost,
resolveModel: async (candidate) => {
if (candidate.accountId !== 'account-b' || candidate.modelId !== 'model-b') {
if (candidate.accountId !== 'account-b'
|| !['model-b', 'model-c'].includes(candidate.modelId)) {
throw new Error('model unavailable');
}
return {
@@ -239,10 +267,16 @@ describe('Pi Conversation runtime', () => {
]);
expect(await runtime.listCommands('00000000-0000-4000-8000-000000000000')).toEqual([]);
const leftGenerationOneSeqs: number[] = [];
const crashPatchOps: string[] = [];
const unsubscribe = runtime.subscribe((envelope) => {
if (envelope.conversationId === left.id && envelope.workerGeneration === 1) {
leftGenerationOneSeqs.push(envelope.seq);
}
if (envelope.conversationId === left.id
&& envelope.workerGeneration === 3
&& (envelope.patch.op === 'run.state' || envelope.patch.op === 'worker.state')) {
crashPatchOps.push(envelope.patch.op);
}
});
const releasePromptAcceptance = workers.get(left.id)!.holdNext('prompt');
@@ -258,7 +292,7 @@ describe('Pi Conversation runtime', () => {
return value;
});
await expect.poll(() => workers.get(left.id)!.requests.at(-1)?.type).toBe('prompt');
expect(workers.get(left.id)!.requests).toHaveLength(5);
expect(workers.get(left.id)!.requests).toHaveLength(6);
expect(acceptanceResolved).toBe(false);
releasePromptAcceptance();
const accepted = await acceptance;
@@ -354,6 +388,7 @@ describe('Pi Conversation runtime', () => {
type: 'set_model',
}));
expect(workers.get(left.id)!.requests.map(({ type }) => type).sort()).toEqual([
'get_available_thinking_levels',
'get_entries',
'get_session_stats',
'get_state',
@@ -363,12 +398,31 @@ describe('Pi Conversation runtime', () => {
expect((await runtime.getSnapshot(right.id)).conversation.model.model).toEqual(model);
expect((await runtime.getSnapshot(left.id)).nodes).toEqual(settledNodes);
const generationBeforeSameAccountModelChange = pool.getState(left.id)!.generation;
const sameAccountChanged = await runtime.setModel({
conversationId: left.id,
accountId: 'account-b',
modelId: 'model-c',
});
expect(sameAccountChanged.model).toEqual({
accountId: 'account-b',
modelId: 'model-c',
thinkingLevel: 'medium',
});
expect(pool.getState(left.id)!.generation).toBe(generationBeforeSameAccountModelChange);
expect(workerHistory.get(left.id)).toHaveLength(2);
expect(workers.get(left.id)!.requests.slice(-3).map(({ type }) => type)).toEqual([
'set_model',
'get_available_thinking_levels',
'get_state',
]);
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);
expect((await runtime.getSnapshot(left.id)).conversation.model).toEqual(sameAccountChanged);
await runtime.prompt({
clientRequestId: 'request-left-active',
@@ -398,22 +452,26 @@ describe('Pi Conversation runtime', () => {
text: 'Then run the focused test',
attachments: [],
})).resolves.toMatchObject({ accepted: true, mode: 'follow-up', queuePosition: 2 });
const generationBeforeThinking = pool.getState(left.id)!.generation;
await expect(runtime.setThinking({
conversationId: left.id,
thinkingLevel: 'high',
})).resolves.toMatchObject({
model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' },
model: { accountId: 'account-b', modelId: 'model-c', thinkingLevel: 'high' },
});
expect(pool.getState(left.id)!.generation).toBe(generationBeforeThinking);
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([
expect(workers.get(left.id)!.requests.slice(-8).map(({ type }) => type)).toEqual([
'steer',
'steer',
'follow_up',
'get_available_thinking_levels',
'set_thinking_level',
'get_state',
'abort',
'abort',
]);
@@ -421,6 +479,7 @@ describe('Pi Conversation runtime', () => {
'get_state',
'get_entries',
'get_session_stats',
'get_available_thinking_levels',
]);
workers.get(left.id)!.emit({ type: 'agent_settled' });
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('idle');
@@ -487,6 +546,7 @@ describe('Pi Conversation runtime', () => {
seq: 0,
});
expect(workers.get(right.id)!.requests.map(({ type }) => type).sort()).toEqual([
'get_available_thinking_levels',
'get_entries',
'get_session_stats',
'get_state',
@@ -525,6 +585,97 @@ 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' });
workers.get(left.id)!.setSessionData({
thinkingLevels: { levels: ['high'] },
state: {
sessionId: `session-${left.id}`,
thinkingLevel: 'high',
isStreaming: false,
isCompacting: false,
pendingMessageCount: 0,
},
});
await expect(runtime.setThinking({
conversationId: left.id,
thinkingLevel: 'off',
})).rejects.toMatchObject({
publicError: { code: 'CODING_MODEL_UNAVAILABLE', recoverable: true },
});
expect((await runtime.getSnapshot(left.id)).conversation.model).toMatchObject({
model: { thinkingLevel: 'medium' },
availableThinkingLevels: ['high'],
});
const rightBeforeCrash = await runtime.getSnapshot(right.id);
const requestsBeforeCrash = workers.get(left.id)!.requests.length;
await runtime.prompt({
clientRequestId: 'request-worker-exit',
conversationId: left.id,
mode: 'prompt',
text: 'This accepted prompt must not be replayed',
attachments: [],
});
expect(trackingHost.runs.has(left.id)).toBe(true);
crashPatchOps.length = 0;
workers.get(left.id)!.invalidate();
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('error');
expect((await runtime.getSnapshot(left.id))).toMatchObject({
run: {
status: 'error',
terminalReason: 'failed',
error: {
code: 'CODING_RUNTIME_START_FAILED',
message: '本地 Agent 已中断,原请求未自动重发。',
recoverable: true,
},
},
worker: { status: 'error', generation: 3 },
});
expect(trackingHost.runs.has(left.id)).toBe(false);
expect(crashPatchOps).toEqual(['run.state', 'worker.state']);
await expect(runtime.abort(left.id)).resolves.toBeUndefined();
expect((await runtime.getSnapshot(left.id)).run.status).toBe('error');
expect(await runtime.getSnapshot(right.id)).toEqual(rightBeforeCrash);
await runtime.recover(left.id);
expect(workers.get(left.id)!.requests.slice(0, requestsBeforeCrash)).not.toContainEqual({
type: 'prompt',
message: 'This accepted prompt must not be replayed',
});
await runtime.prompt({
clientRequestId: 'request-settled-before-close',
conversationId: left.id,
mode: 'prompt',
text: 'Settle before close',
attachments: [],
});
workers.get(left.id)!.emit({ type: 'agent_settled' });
workers.get(left.id)!.invalidate();
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('idle');
expect((await runtime.getSnapshot(left.id)).run.terminalReason).toBe('completed');
await runtime.recover(left.id);
await runtime.prompt({
clientRequestId: 'request-protocol-invalidation',
conversationId: left.id,
mode: 'prompt',
text: 'Fail this generation without replay',
attachments: [],
});
workers.get(left.id)!.invalidate(new PiProcessError(
'PI_RPC_PROTOCOL_ERROR',
'strict JSONL protocol failure',
));
await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('error');
expect((await runtime.getSnapshot(left.id)).run).toMatchObject({
terminalReason: 'failed',
error: { code: 'CODING_RUNTIME_PROTOCOL_ERROR', recoverable: true },
});
workers.get(left.id)!.emit({ type: 'agent_settled' });
await new Promise((resolve) => setImmediate(resolve));
expect((await runtime.getSnapshot(left.id)).run.status).toBe('error');
unsubscribe();
});
});