fix(pi): converge worker failures and thinking state
This commit is contained in:
@@ -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();
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user