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

@@ -100,6 +100,14 @@ function handle(command) {
if (command.type === 'crash') {
process.exit(7);
}
if (command.type === 'crash_with_stderr') {
const secret = process.env.FAKE_PI_SECRET ?? 'missing-secret';
process.stderr.write(`Authorization: Bearer ${secret}\n`);
process.stderr.write(`token=${secret}\n`);
process.stderr.write(`cwd=${process.cwd()}\n`);
process.stderr.write('worker-exit-marker\n', () => process.exit(9));
return;
}
writeRecord({
type: 'response',
id: command.id,

View File

@@ -1,4 +1,4 @@
import { fireEvent, render, screen, waitFor } from '@testing-library/react';
import { fireEvent, render, screen, waitFor, within } from '@testing-library/react';
import { describe, expect, it, vi } from 'vitest';
import type { ProviderAccount, ProviderVendorInfo } from '@/lib/providers';
@@ -102,6 +102,41 @@ describe('PI-130 feature-complete Coding UI', () => {
expect(onSubmit).toHaveBeenCalledOnce();
});
it('unlocks editing and offers recovery after the local Agent exits', async () => {
const onRecover = vi.fn();
const { CodingComposer } = await import('@/pages/Chat/CodingComposer');
render(
<CodingComposer
value="Draft remains editable"
editable
canSend={false}
preparing={false}
recovering={false}
runStatus="error"
mode="prompt"
queue={{ items: [] }}
error="本地 Agent 已中断,原请求未自动重发。"
recoverableError
acceptedCount={0}
attachments={[]}
submitting={false}
placeholder="Prompt"
onChange={vi.fn()}
onModeChange={vi.fn()}
onSubmit={vi.fn()}
onRecover={onRecover}
onAddFiles={vi.fn()}
onRemoveAttachment={vi.fn()}
/>,
);
expect(screen.getByRole('textbox')).toBeEnabled();
expect(screen.getByRole('alert')).toHaveTextContent('本地 Agent 已中断,原请求未自动重发。');
expect(screen.queryByText('当前对话正在处理。')).not.toBeInTheDocument();
fireEvent.click(screen.getByRole('button', { name: '重试' }));
expect(onRecover).toHaveBeenCalledOnce();
});
it('answers product interactions and explains stale responses', async () => {
interactionApi.respond.mockResolvedValueOnce(undefined).mockRejectedValueOnce(new Error('409 stale'));
const { CodingInteractionPanel } = await import('@/pages/Chat/CodingInteractionPanel');
@@ -198,8 +233,9 @@ describe('PI-130 feature-complete Coding UI', () => {
agentId: 'agent-1',
title: 'Feature UI',
model: {
model: { accountId: 'account-1', modelId: 'model-a', thinkingLevel: 'medium' },
model: { accountId: 'account-1', modelId: 'model-a', thinkingLevel: 'high' },
modelResolution: 'resolved',
availableThinkingLevels: ['high'],
},
},
nodes: [],
@@ -227,9 +263,13 @@ describe('PI-130 feature-complete Coding UI', () => {
await waitFor(() => expect(interactionApi.model).toHaveBeenCalledWith('conversation-1', {
accountId: 'account-1',
modelId: 'model-b',
thinkingLevel: 'medium',
thinkingLevel: 'high',
}));
await waitFor(() => expect(callbacks.refresh).toHaveBeenCalledOnce());
const thinkingSelect = screen.getByRole('combobox', { name: '当前对话思考级别' });
expect(within(thinkingSelect).getAllByRole('option')).toHaveLength(1);
expect(within(thinkingSelect).getByRole('option', { name: '高思考' })).toBeInTheDocument();
expect(within(thinkingSelect).queryByRole('option', { name: '关闭思考' })).not.toBeInTheDocument();
fireEvent.click(screen.getByRole('button', { name: '中止' }));
await waitFor(() => expect(interactionApi.abort).toHaveBeenCalledWith('conversation-1'));
fireEvent.click(screen.getByRole('button', { name: '创建分支' }));

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();
});
});

View File

@@ -269,6 +269,49 @@ describe('managed Pi extension bridge', () => {
expect(currentWorker.status).toBe(200);
});
it('releases a crashed worker write lease so only its same-project waiter advances', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-crash-lease-'));
roots.push(root);
const leases = new PiProjectWriteLeaseCoordinator();
const host = new PiManagedExtensionHost(leases);
hosts.push(host);
const crashed = await host.registerWorker({
conversationId: 'conversation-crashed', generation: 1, projectId: 'project-a', extensionsDir: root,
});
const waiter = await host.registerWorker({
conversationId: 'conversation-waiter', generation: 1, projectId: 'project-a', extensionsDir: root,
});
const other = await host.registerWorker({
conversationId: 'conversation-other', generation: 1, projectId: 'project-b', extensionsDir: root,
});
await Promise.all([
host.bindRun('conversation-crashed', 1, 'run-crashed'),
host.bindRun('conversation-waiter', 1, 'run-waiter'),
host.bindRun('conversation-other', 1, 'run-other'),
]);
expect((await post(crashed, {
action: 'lease.acquire', conversationId: 'conversation-crashed', workerGeneration: 1,
runId: 'run-crashed', resourceId: 'crashed-write',
})).status).toBe(200);
const waiting = post(waiter, {
action: 'lease.acquire', conversationId: 'conversation-waiter', workerGeneration: 1,
runId: 'run-waiter', resourceId: 'waiting-write',
});
expect((await post(other, {
action: 'lease.acquire', conversationId: 'conversation-other', workerGeneration: 1,
runId: 'run-other', resourceId: 'other-write',
})).status).toBe(200);
await expect.poll(() => leases.waitingCount('project-a')).toBe(1);
await crashed.dispose();
await expect(waiting).resolves.toMatchObject({ status: 200 });
expect(leases.waitingCount()).toBe(0);
expect(leases.activeCount).toBe(2);
await Promise.all([waiter.dispose(), other.dispose()]);
expect(leases.activeCount).toBe(0);
});
it('joins a coding child to the same project write lease as its parent', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-child-lease-'));
roots.push(root);

View File

@@ -200,9 +200,9 @@ describe('managed Pi worker opener', () => {
expect(JSON.stringify(telemetry)).not.toContain(created.id);
expect(JSON.stringify(telemetry)).not.toContain('PRIVATE MANAGED PROMPT');
expect(JSON.stringify(telemetry)).not.toContain('provider-secret-value');
await first.worker.stop();
await reopened.worker.stop();
await restarted.worker.stop();
await first.worker.stop('test_injection');
await reopened.worker.stop('test_injection');
await restarted.worker.stop('test_injection');
await extensionHost.close();
});

View File

@@ -179,6 +179,47 @@ describe('Pi Provider catalog', () => {
})).toThrowError(PiProviderConfigError);
});
it('projects the managed DeepSeek capability contract into Pi models.json', () => {
const provider = account({
id: 'niancode-user-models',
vendorId: 'custom',
apiProtocol: 'openai-completions',
baseUrl: 'http://127.0.0.1:54321/api/ai-proxy/v1',
model: 'deepseek/deepseek-v4-pro',
metadata: {
worksSquareCredentialMode: 'works_square_ai_gateway_proxy',
customModels: ['deepseek/deepseek-v4-pro'],
},
});
const catalog = buildPiProviderCatalog({ accounts: [provider] });
const descriptor = catalog.descriptors[0]!.models[0]!;
const written = catalog.modelsFile.providers[resolvePiRuntimeProviderId(provider.id)]!.models[0]!;
expect(descriptor).toMatchObject({
id: 'deepseek-v4-pro',
reasoning: true,
contextWindow: 1_000_000,
maxOutputTokens: 384_000,
compat: {
thinkingFormat: 'deepseek',
requiresReasoningContentOnAssistantMessages: true,
},
thinkingLevelMap: {
off: null,
minimal: null,
low: null,
medium: null,
high: 'high',
},
});
expect(written).toMatchObject({
reasoning: true,
thinkingLevelMap: descriptor.thinkingLevelMap,
compat: descriptor.compat,
});
});
it('does not replace an existing catalog when model selection is unavailable', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-provider-'));
temporaryRoots.push(root);

View File

@@ -13,6 +13,7 @@ import {
buildPiRpcArgs,
sanitizePiDiagnostic,
buildPiWorkerEnvironment,
type PiWorkerLifecycleEvent,
} from '../../electron/coding-runtime/pi/worker-process';
const fakeChildPath = resolve('tests/fixtures/fake-pi-rpc-child.mjs');
@@ -56,7 +57,7 @@ async function processAlive(pid: number): Promise<boolean> {
}
afterEach(async () => {
await Promise.all(workers.splice(0).map((worker) => worker.stop().catch(() => undefined)));
await Promise.all(workers.splice(0).map((worker) => worker.stop('test_injection').catch(() => undefined)));
await Promise.all(scratchRoots.splice(0).map((root) => rm(root, {
recursive: true,
force: true,
@@ -259,6 +260,87 @@ describe('Pi worker process', () => {
expect(invalidations).toEqual(['PI_RPC_EXITED']);
});
it('classifies unexpected exit with bounded redacted stderr and generation correlation', async () => {
const secret = 'unexpected-exit-secret';
const lifecycle: PiWorkerLifecycleEvent[] = [];
const worker = await makeWorker({
conversationId: 'conversation-exit',
workerGeneration: 7,
env: { FAKE_PI_SECRET: secret },
sensitiveValues: [secret],
diagnosticBytes: 160,
onLifecycleEvent: (event) => lifecycle.push(event),
});
await expect(worker.request({ type: 'crash_with_stderr' })).rejects.toMatchObject({
code: 'PI_RPC_EXITED',
generation: 7,
exitCode: 9,
signal: null,
diagnostic: expect.stringContaining('worker-exit-marker'),
});
const close = lifecycle.find((event) => (
event.classification === 'unexpected_exit' && event.stage === 'close'
));
expect(close).toMatchObject({
conversationId: 'conversation-exit',
generation: 7,
code: 'PI_RPC_EXITED',
exitCode: 9,
signal: null,
});
expect(JSON.stringify(close)).not.toContain(secret);
expect(JSON.stringify(close)).not.toContain(scratchRoots.at(-1));
expect(close?.diagnostic).toContain('cwd=[REDACTED]');
expect(Buffer.byteLength(close?.diagnostic ?? '')).toBeLessThanOrEqual(160);
});
it('keeps protocol invalidation and deliberate stop as distinct reasoned lifecycle events', async () => {
const protocolLifecycle: PiWorkerLifecycleEvent[] = [];
const protocolWorker = await makeWorker({
conversationId: 'conversation-protocol',
workerGeneration: 3,
onLifecycleEvent: (event) => protocolLifecycle.push(event),
});
await expect(protocolWorker.request({ type: 'malformed' })).rejects.toMatchObject({
code: 'PI_RPC_PROTOCOL_ERROR',
generation: 3,
});
expect(protocolLifecycle).toContainEqual(expect.objectContaining({
classification: 'protocol_invalidation',
stage: 'protocol',
conversationId: 'conversation-protocol',
generation: 3,
code: 'PI_RPC_PROTOCOL_ERROR',
}));
expect(protocolLifecycle).not.toContainEqual(expect.objectContaining({
classification: 'unexpected_exit',
}));
const stopLifecycle: PiWorkerLifecycleEvent[] = [];
const stoppedWorker = await makeWorker({
conversationId: 'conversation-stop',
workerGeneration: 5,
onLifecycleEvent: (event) => stopLifecycle.push(event),
});
await expect(stoppedWorker.stop('test_injection')).resolves.toMatchObject({
mode: 'stdin-close',
});
expect(stopLifecycle).toContainEqual(expect.objectContaining({
classification: 'intentional_stop',
stage: 'stop_completed',
reason: 'test_injection',
conversationId: 'conversation-stop',
generation: 5,
code: 'PI_WORKER_STOPPED',
exitCode: 0,
signal: null,
}));
expect(stopLifecycle).not.toContainEqual(expect.objectContaining({
classification: 'unexpected_exit',
}));
});
it('keeps only bounded redacted stderr diagnostics', async () => {
const secret = 'credential-that-must-not-leak';
const worker = await makeWorker({
@@ -314,7 +396,7 @@ describe('Pi worker process', () => {
const descendantPid = response.data?.pid;
expect(descendantPid).toBeTypeOf('number');
await expect(worker.stop()).resolves.toMatchObject({ mode: 'forced-tree-kill' });
await expect(worker.stop('test_injection')).resolves.toMatchObject({ mode: 'forced-tree-kill' });
for (let attempt = 0; attempt < 20 && await processAlive(descendantPid!); attempt += 1) {
await new Promise((resolvePromise) => setTimeout(resolvePromise, 25));
}

View File

@@ -147,8 +147,8 @@ describe('managed Pi subagent child opener', () => {
summary: 'managed child summary',
usage: { inputTokens: 7, outputTokens: 11, cacheReadTokens: 2 },
});
await readOnly.stop();
await coding.stop();
await readOnly.stop('test_injection');
await coding.stop('test_injection');
expect(processes.every(({ stopped }) => stopped)).toBe(true);
await host.close();
});
@@ -210,8 +210,8 @@ describe('managed Pi subagent child opener', () => {
expect(await readFile(path.join(userDataDir, 'coding-runtime', 'pi', 'config', 'models.json'), 'utf8'))
.not.toContain('host-token-');
await first.stop();
await second.stop();
await first.stop('test_injection');
await second.stop('test_injection');
await host.close();
});
});

View File

@@ -160,7 +160,7 @@ describe('Pi subagent scheduler', () => {
it('aborts every unfinished child with the parent and leaves no permit or process lease', async () => {
const processBudget = new PiProcessBudget(8);
const stopped: string[] = [];
const stopped: Array<{ taskId: string; reason: string }> = [];
const scheduler = new PiSubagentScheduler({
processBudget,
openChild: async (input) => ({
@@ -173,7 +173,7 @@ describe('Pi subagent scheduler', () => {
});
return { summary: 'unreachable' };
},
async stop() { stopped.push(input.taskId); },
async stop(reason) { stopped.push({ taskId: input.taskId, reason }); },
}),
});
const controller = new AbortController();
@@ -186,6 +186,7 @@ describe('Pi subagent scheduler', () => {
const result = await flight;
expect(result.details.tasks.map(({ status }) => status)).toEqual(['aborted', 'aborted']);
expect(stopped).toHaveLength(2);
expect(stopped.map(({ reason }) => reason)).toEqual(['subagent_abort', 'subagent_abort']);
expect(processBudget.activeCount).toBe(0);
await scheduler.close();
});

View File

@@ -19,6 +19,7 @@ import {
} from '../../electron/coding-runtime/pi/worker-pool';
import {
PiWorkerProcess,
type PiWorkerStopReason,
type PiWorkerStopResult,
} from '../../electron/coding-runtime/pi/worker-process';
@@ -60,8 +61,8 @@ class ProcessBackedWorker implements PiConversationWorker {
return this.process.subscribeInvalidation(listener);
}
stop(): Promise<PiWorkerStopResult> {
return this.process.stop();
stop(reason: PiWorkerStopReason): Promise<PiWorkerStopResult> {
return this.process.stop(reason);
}
}

View File

@@ -13,6 +13,10 @@ import { PiProcessError as PiProcessFailure } from '../../electron/coding-runtim
import type { PiRpcCommand, PiRpcEvent } from '../../electron/coding-runtime/pi/rpc-client';
import type { PiRuntimeTelemetryEvent } from '../../electron/coding-runtime/pi/telemetry';
import { PiSubagentScheduler } from '../../electron/coding-runtime/pi/subagent';
import type {
PiWorkerProofFailure,
PiWorkerStopReason,
} from '../../electron/coding-runtime/pi/worker-process';
const MODEL = {
model: {
@@ -43,6 +47,7 @@ class FakeWorker implements PiConversationWorker {
readonly generation = 1;
readonly requests: PiRpcCommand[] = [];
stopped = false;
readonly stopReasons: PiWorkerStopReason[] = [];
private readonly eventListeners = new Set<(event: PiRpcEvent) => void>();
private readonly invalidationListeners = new Set<(error: PiProcessError) => void>();
@@ -75,13 +80,50 @@ class FakeWorker implements PiConversationWorker {
for (const listener of this.invalidationListeners) listener(error);
}
async stop() {
async injectFailureForProof(failure: PiWorkerProofFailure): Promise<void> {
this.invalidate(new PiProcessFailure(
failure === 'protocol_invalidation' ? 'PI_RPC_PROTOCOL_ERROR' : 'PI_RPC_EXITED',
'injected proof failure',
));
}
async stop(reason: PiWorkerStopReason) {
this.stopped = true;
this.stopReasons.push(reason);
return { mode: 'stdin-close' as const, code: 0, signal: null };
}
}
describe('Pi worker pool', () => {
it('injects a proof failure only into the requested current generation', async () => {
const workers = new Map<string, FakeWorker>();
const pool = new PiWorkerPool({
maxIdle: 4,
openWorker: async ({ conversation: input }) => {
const worker = new FakeWorker(`worker-${input.conversationId}`);
workers.set(input.conversationId, worker);
return {
worker,
session: {
piSessionId: `session-${input.conversationId}`,
sessionKey: `key-${input.conversationId}`,
},
};
},
});
await Promise.all([
pool.prepare(conversation('conversation-target')),
pool.prepare(conversation('conversation-other')),
]);
await expect(pool.injectFailureForProof('conversation-target', 'unexpected_exit'))
.resolves.toEqual({ generation: 1 });
expect(pool.getState('conversation-target')).toMatchObject({ state: 'crashed', generation: 1 });
expect(pool.getState('conversation-other')).toMatchObject({ state: 'ready', generation: 1 });
expect(workers.get('conversation-other')?.stopped).toBe(false);
});
it('single-flights prepare per Conversation and never shares its worker with another Conversation', async () => {
const gate = deferred();
const opened: string[] = [];
@@ -418,9 +460,11 @@ describe('Pi worker pool', () => {
it('cleans only the crashed generation and releases its permit for the next Conversation', async () => {
const workers = new Map<string, FakeWorker>();
const processBudget = new PiProcessBudget(8);
const pool = new PiWorkerPool({
maxRunning: 2,
maxIdle: 3,
processBudget,
openWorker: async ({ conversation: input }) => {
const worker = new FakeWorker(`worker-${input.conversationId}`);
workers.set(input.conversationId, worker);
@@ -461,6 +505,7 @@ describe('Pi worker pool', () => {
expect(pool.getState('conversation-a')).toMatchObject({ state: 'crashed', generation: 1 });
expect(pool.getState('conversation-b')).toMatchObject({ state: 'running', generation: 1 });
await expect.poll(() => workers.get('conversation-c')!.requests.length).toBe(1);
await expect.poll(() => processBudget.activeCount).toBe(2);
expect(cancelled).not.toContain('other');
});
@@ -468,6 +513,7 @@ describe('Pi worker pool', () => {
const workers = new Map<string, FakeWorker[]>();
const revisions: Array<{ conversationId: string; provider: number; resources: number }> = [];
const telemetry: PiRuntimeTelemetryEvent[] = [];
const replacementReasons: string[] = [];
const pool = new PiWorkerPool({
maxIdle: 4,
onTelemetry: (event) => telemetry.push(event),
@@ -484,6 +530,9 @@ describe('Pi worker pool', () => {
};
},
});
pool.subscribe((event) => {
if (event.type === 'worker.replaced') replacementReasons.push(event.reason);
});
await Promise.all([
pool.prepare(conversation('conversation-running')),
pool.prepare(conversation('conversation-idle')),
@@ -520,6 +569,46 @@ describe('Pi worker pool', () => {
{ conversationId: 'conversation-idle', provider: 2, resources: 1 },
{ conversationId: 'conversation-running', provider: 2, resources: 1 },
]));
expect(replacementReasons).toEqual([
'stale_resource_rebuild',
'stale_resource_rebuild',
]);
expect(workers.get('conversation-idle')![0]!.stopReasons).toEqual([
'stale_resource_rebuild',
]);
expect(workers.get('conversation-running')![0]!.stopReasons).toEqual([
'stale_resource_rebuild',
]);
});
it('records an explicit recover replacement reason without changing the session binding', async () => {
const workers: FakeWorker[] = [];
const replacements: Array<{ reason: string; generation: number }> = [];
const pool = new PiWorkerPool({
openWorker: async ({ conversation: input, generation, existingSession }) => {
const worker = new FakeWorker(`worker-${generation}`);
workers.push(worker);
return {
worker,
session: existingSession ?? {
piSessionId: `session-${input.conversationId}`,
sessionKey: `key-${input.conversationId}`,
},
};
},
});
pool.subscribe((event) => {
if (event.type === 'worker.replaced') {
replacements.push({ reason: event.reason, generation: event.generation });
}
});
const before = await pool.prepare(conversation('conversation-recover'));
const recovered = await pool.recover('conversation-recover');
expect(recovered.session).toEqual(before.session);
expect(replacements).toEqual([{ reason: 'recover', generation: 2 }]);
expect(workers[0]!.stopReasons).toEqual(['recover']);
});
it('re-applies the idle LRU after a running stale worker rebuilds on settle', async () => {

View File

@@ -119,9 +119,9 @@ describe('locked Pi worker process smoke', () => {
'changed_file',
'runtime_context',
]));
await expect(worker.stop()).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
await expect(worker.stop('test_injection')).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
} finally {
await worker.stop().catch(() => undefined);
await worker.stop('test_injection').catch(() => undefined);
await extension.dispose();
await extensionHost.close();
}
@@ -173,9 +173,9 @@ describe('locked Pi worker process smoke', () => {
});
expect(worker.stderrDiagnostic).not.toContain('Failed to load extension');
expect(await readActiveTools(probe.resultPath)).toEqual(['read', 'grep', 'find', 'ls']);
await expect(worker.stop()).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
await expect(worker.stop('test_injection')).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
} finally {
await worker.stop().catch(() => undefined);
await worker.stop('test_injection').catch(() => undefined);
await extension.dispose();
await extensionHost.close();
}
@@ -235,9 +235,9 @@ describe('locked Pi worker process smoke', () => {
});
expect(await readActiveTools(probe.resultPath)).toEqual(['read', 'grep', 'find', 'ls']);
expect(worker.stderrDiagnostic).not.toContain('Failed to load extension');
await expect(worker.stop()).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
await expect(worker.stop('test_injection')).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
} finally {
await worker.stop().catch(() => undefined);
await worker.stop('test_injection').catch(() => undefined);
await extension.dispose();
await extensionHost.close();
}
@@ -294,9 +294,9 @@ describe('locked Pi worker process smoke', () => {
'runtime_context',
]));
expect(worker.stderrDiagnostic).not.toContain('Failed to load extension');
await expect(worker.stop()).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
await expect(worker.stop('test_injection')).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
} finally {
await worker.stop().catch(() => undefined);
await worker.stop('test_injection').catch(() => undefined);
await extension.dispose();
await extensionHost.close();
}