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

@@ -1,5 +1,6 @@
import { randomUUID } from 'node:crypto';
import path from 'node:path';
import { logger } from '../../utils/logger';
import type { ModelSummary, ProviderAccount } from '../../shared/providers/types';
import { validateSessionKey } from '../../coding-projects/conversation-store';
import {
@@ -59,6 +60,8 @@ import type {
import {
PiWorkerProcess,
type PiWorkerProcessOptions,
type PiWorkerProofFailure,
type PiWorkerStopReason,
type PiWorkerStopResult,
} from './worker-process';
import {
@@ -118,7 +121,8 @@ export interface PiWorkerProcessAdapter {
send(command: PiRpcCommand): Promise<void>;
subscribe(listener: (event: PiRpcEvent) => void): () => void;
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void;
stop(): Promise<PiWorkerStopResult>;
stop(reason: PiWorkerStopReason): Promise<PiWorkerStopResult>;
injectFailureForProof?(failure: PiWorkerProofFailure): Promise<void>;
}
export interface PiManagedProviderInput {
@@ -172,14 +176,21 @@ class ManagedPiConversationWorker implements PiConversationWorker {
return this.process.subscribeInvalidation(listener);
}
async stop(): Promise<PiWorkerStopResult> {
async stop(reason: PiWorkerStopReason): Promise<PiWorkerStopResult> {
this.unsubscribeExtensionInvalidation();
try {
return await this.process.stop();
return await this.process.stop(reason);
} finally {
await this.disposeExtension();
}
}
async injectFailureForProof(failure: PiWorkerProofFailure): Promise<void> {
if (!this.process.injectFailureForProof) {
throw new Error('Managed Pi process does not support proof failure injection');
}
await this.process.injectFailureForProof(failure);
}
}
export function createPiManagedWorkerOpener(
@@ -261,6 +272,8 @@ export function createPiManagedWorkerOpener(
],
env: { ...credential.env, ...extension.env },
sensitiveValues: [...credential.sensitiveValues, ...extension.sensitiveValues],
conversationId: input.conversation.conversationId,
workerGeneration: input.generation,
});
let unsubscribeExtensionInvalidation = process.subscribeInvalidation(() => {
void extension.dispose();
@@ -348,7 +361,7 @@ export function createPiManagedWorkerOpener(
};
} catch (error) {
unsubscribeExtensionInvalidation();
await process.stop().catch(() => undefined);
await process.stop('open_failure').catch(() => undefined);
await extension.dispose();
throw error;
}
@@ -377,6 +390,40 @@ function clone<T>(value: T): T {
return structuredClone(value);
}
const PRODUCT_THINKING_LEVELS = new Set<ProductModelRef['thinkingLevel']>([
'off',
'minimal',
'low',
'medium',
'high',
]);
function productThinkingLevel(value: unknown): ProductModelRef['thinkingLevel'] | null {
return typeof value === 'string'
&& PRODUCT_THINKING_LEVELS.has(value as ProductModelRef['thinkingLevel'])
? value as ProductModelRef['thinkingLevel']
: null;
}
function availableThinkingLevels(value: unknown): ProductModelRef['thinkingLevel'][] {
if (!value || typeof value !== 'object' || Array.isArray(value)) return [];
const levels = (value as { levels?: unknown }).levels;
if (!Array.isArray(levels)) return [];
return [...new Set(levels.flatMap((level) => {
const normalized = productThinkingLevel(level);
return normalized ? [normalized] : [];
}))];
}
function effectiveThinkingLevel(value: unknown): ProductModelRef['thinkingLevel'] | null {
if (!value || typeof value !== 'object' || Array.isArray(value)) return null;
return productThinkingLevel((value as { thinkingLevel?: unknown }).thinkingLevel);
}
function runIsTerminal(status: ConversationSnapshot['run']['status']): boolean {
return status === 'idle' || status === 'error';
}
function publicWorkerState(state: PiWorkerPoolState): ConversationSnapshot['worker'] {
if (state.state === 'spawning') return { status: 'starting', generation: state.generation };
if (state.state === 'crashed') {
@@ -415,6 +462,13 @@ function runtimeFailure(error: unknown): CodingRuntimePublicError {
recoverable: true,
};
}
if (error.code === 'PI_RPC_EXITED') {
return {
code: 'CODING_RUNTIME_START_FAILED',
message: '本地 Agent 已中断,原请求未自动重发。',
recoverable: true,
};
}
}
return {
code: 'CODING_RUNTIME_START_FAILED',
@@ -458,6 +512,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
private readonly isAuthenticationError: ((error: unknown) => boolean) | undefined;
private readonly refreshCredential: ((accountId: string) => Promise<void>) | undefined;
private readonly extensionHost: PiManagedExtensionHost | undefined;
private readonly subagentScheduler: PiSubagentScheduler | undefined;
private readonly interactions: PiInteractionStore;
private readonly extensionUi: PiExtensionUiProjector;
private readonly onExtensionUiProjection: ((projection: PiExtensionUiProjection) => void) | undefined;
@@ -493,6 +548,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
this.isAuthenticationError = options.isAuthenticationError;
this.refreshCredential = options.refreshCredential;
this.extensionHost = options.extensionHost;
this.subagentScheduler = options.subagentScheduler;
if (options.subagentScheduler && !this.extensionHost) {
throw new Error('Subagent scheduler requires the managed extension host');
}
@@ -657,6 +713,8 @@ export class PiConversationRuntime implements CodingConversationRuntime {
async abort(conversationId: string): Promise<void> {
await this.waitForProjection(conversationId);
const current = this.snapshot(conversationId).run;
if (runIsTerminal(current.status)) return;
const generation = this.snapshot(conversationId).cursor.workerGeneration;
this.emit(conversationId, {
op: 'run.state',
run: { ...current, status: 'aborting' },
@@ -670,8 +728,14 @@ export class PiConversationRuntime implements CodingConversationRuntime {
this.extensionUi.endRun(conversationId, current.runId);
}
} catch (error) {
await this.waitForProjection(conversationId);
const latest = this.snapshot(conversationId).run;
if (latest.runId === current.runId && latest.status === 'aborting') {
const worker = this.pool.getState(conversationId);
if (latest.runId === current.runId
&& latest.status === 'aborting'
&& (!worker || worker.state === 'crashed' || worker.generation !== generation)) {
this.failRun(conversationId, current.runId!, error, generation);
} else if (latest.runId === current.runId && latest.status === 'aborting') {
this.emit(conversationId, { op: 'run.state', run: current }, current.runId);
}
throw error;
@@ -720,17 +784,31 @@ export class PiConversationRuntime implements CodingConversationRuntime {
}
throw error;
}
return clone(persisted);
return clone(this.snapshot(input.conversationId).conversation.model);
}
await this.pool.request(input.conversationId, {
type: 'set_model',
provider: selection.runtimeProviderId,
modelId: selection.modelId,
});
const persisted = await this.registry.setModel(input.conversationId, model);
this.pool.updateConversationModel(input.conversationId, persisted);
this.replaceModel(input.conversationId, persisted);
return clone(persisted);
const capabilities = await this.pool.request<{ levels?: unknown }>(
input.conversationId,
{ type: 'get_available_thinking_levels' },
{ retry: 'read-only-once' },
);
const state = await this.pool.request(
input.conversationId,
{ type: 'get_state' },
{ retry: 'read-only-once' },
);
return await this.persistEffectiveThinking(
input.conversationId,
model,
state.data,
capabilities.data,
true,
true,
);
}
async setThinking(input: SetThinkingLevelInput): Promise<ConversationModelState> {
@@ -743,17 +821,44 @@ export class PiConversationRuntime implements CodingConversationRuntime {
true,
);
}
const capabilities = await this.pool.request<{ levels?: unknown }>(
input.conversationId,
{ type: 'get_available_thinking_levels' },
{ retry: 'read-only-once' },
);
const available = availableThinkingLevels(capabilities.data);
const withCapabilities: ConversationModelState = {
...clone(current),
...(available.length > 0 ? { availableThinkingLevels: available } : {}),
};
this.replaceModel(input.conversationId, withCapabilities);
if (!available.includes(input.thinkingLevel)) {
throw new CodingRuntimeContractError(
'CODING_MODEL_UNAVAILABLE',
'The selected thinking level is not supported by this model',
true,
);
}
await this.pool.request(input.conversationId, {
type: 'set_thinking_level',
level: input.thinkingLevel,
});
const model: ConversationModelState = {
const state = await this.pool.request(
input.conversationId,
{ type: 'get_state' },
{ retry: 'read-only-once' },
);
const persisted = await this.persistEffectiveThinking(input.conversationId, {
model: { ...current.model, thinkingLevel: input.thinkingLevel },
modelResolution: 'resolved',
};
const persisted = await this.registry.setModel(input.conversationId, model);
this.pool.updateConversationModel(input.conversationId, persisted);
this.replaceModel(input.conversationId, persisted);
}, state.data, capabilities.data, true, true);
if (persisted.model?.thinkingLevel !== input.thinkingLevel) {
throw new CodingRuntimeContractError(
'CODING_MODEL_UNAVAILABLE',
'Pi did not accept the selected thinking level',
true,
);
}
return clone(persisted);
}
@@ -892,6 +997,25 @@ export class PiConversationRuntime implements CodingConversationRuntime {
return this.pool.getDiagnostics();
}
async injectWorkerFailureForProof(
conversationId: string,
failure: PiWorkerProofFailure,
): Promise<{ generation: number }> {
return await this.pool.injectFailureForProof(conversationId, failure);
}
getResilienceProofDiagnostics(): {
pool: ReturnType<PiWorkerPool['getResilienceProofDiagnostics']>;
subagents: ReturnType<PiSubagentScheduler['getDiagnostics']> | null;
extension: ReturnType<PiManagedExtensionHost['getDiagnostics']> | null;
} {
return {
pool: this.pool.getResilienceProofDiagnostics(),
subagents: this.subagentScheduler?.getDiagnostics() ?? null,
extension: this.extensionHost?.getDiagnostics() ?? null,
};
}
markProviderStale(): void {
this.pool.markProviderStale();
}
@@ -1121,14 +1245,29 @@ export class PiConversationRuntime implements CodingConversationRuntime {
return;
}
if (event.type === 'worker.crashed') {
void this.interactions.cancelGeneration(event.conversationId, event.generation);
const runId = this.states.get(event.conversationId)?.snapshot.run.runId;
if (runId) {
this.extensionUi.endRun(event.conversationId, runId);
void this.extensionHost?.clearRun(event.conversationId, event.generation, runId);
}
const state = this.pool.getState(event.conversationId);
if (state) this.emit(event.conversationId, { op: 'worker.state', state: publicWorkerState(state) });
void this.enqueueProjection(event.conversationId, async () => {
const snapshot = this.states.get(event.conversationId)?.snapshot;
if (!snapshot || snapshot.cursor.workerGeneration !== event.generation) return;
await this.interactions.cancelGeneration(event.conversationId, event.generation);
const current = this.snapshot(event.conversationId).run;
if (current.runId && !runIsTerminal(current.status)) {
this.failRun(
event.conversationId,
current.runId,
event.error,
event.generation,
);
}
const state = this.pool.getState(event.conversationId);
if (state?.generation === event.generation) {
this.emit(event.conversationId, {
op: 'worker.state',
state: publicWorkerState(state),
});
}
}).catch((error) => {
this.recordProjectionFailure(event.conversationId, event.generation, error);
});
return;
}
void this.enqueueProjection(event.conversationId, async () => {
@@ -1161,7 +1300,7 @@ export class PiConversationRuntime implements CodingConversationRuntime {
}
if (event.event.type === 'agent_end' || event.event.type === 'agent_settled') {
const worker = this.pool.getState(event.conversationId);
if (worker?.generation === event.generation) {
if (worker?.generation === event.generation && worker.state !== 'crashed') {
await this.hydrateGenerationNow(
event.conversationId,
worker,
@@ -1183,9 +1322,18 @@ export class PiConversationRuntime implements CodingConversationRuntime {
});
}
private failRun(conversationId: string, runId: string, error: unknown): void {
private failRun(
conversationId: string,
runId: string,
error: unknown,
generation?: number,
): void {
const current = this.states.get(conversationId)?.snapshot?.run;
if (current?.runId !== runId) return;
const snapshotGeneration = this.states.get(conversationId)?.snapshot?.cursor.workerGeneration;
if (current?.runId !== runId
|| runIsTerminal(current.status)
|| (generation !== undefined && snapshotGeneration !== generation)) return;
const publicError = runtimeFailure(error);
this.emit(conversationId, {
op: 'run.state',
run: {
@@ -1193,13 +1341,20 @@ export class PiConversationRuntime implements CodingConversationRuntime {
runId,
settledAt: this.now(),
terminalReason: 'failed',
error: runtimeFailure(error),
error: publicError,
},
}, runId);
logger.warn('[PiWorkerLifecycle]', {
event: 'run.failed',
conversationId,
runId,
generation: generation ?? snapshotGeneration,
code: publicError.code,
});
void this.interactions.cancelRun(conversationId, runId, true);
this.extensionUi.endRun(conversationId, runId);
const generation = this.pool.getState(conversationId)?.generation;
if (generation) void this.extensionHost?.clearRun(conversationId, generation, runId);
const runGeneration = generation ?? this.pool.getState(conversationId)?.generation;
if (runGeneration) void this.extensionHost?.clearRun(conversationId, runGeneration, runId);
}
private requestHydration(
@@ -1232,10 +1387,13 @@ export class PiConversationRuntime implements CodingConversationRuntime {
this.replaceWorkerGeneration(conversationId, workerState, preserveRun);
this.resetProjector(conversationId);
}
const [stateResponse, entriesResponse, statsResponse] = await Promise.all([
const [stateResponse, entriesResponse, statsResponse, capabilitiesResponse] = await Promise.all([
this.pool.request(conversationId, { type: 'get_state' }, { retry: 'read-only-once' }),
this.pool.request(conversationId, { type: 'get_entries' }, { retry: 'read-only-once' }),
this.pool.request(conversationId, { type: 'get_session_stats' }, { retry: 'read-only-once' }),
this.pool.request(conversationId, { type: 'get_available_thinking_levels' }, {
retry: 'read-only-once',
}),
]);
const before = this.snapshot(conversationId);
let projected = await projectPiSessionSnapshot({
@@ -1253,9 +1411,54 @@ export class PiConversationRuntime implements CodingConversationRuntime {
queue: clone(before.queue),
};
}
projected = {
...projected,
conversation: {
...projected.conversation,
model: await this.persistEffectiveThinking(
conversationId,
projected.conversation.model,
stateResponse.data,
capabilitiesResponse.data,
false,
),
},
};
this.states.set(conversationId, createConversationReducerState(projected));
}
private async persistEffectiveThinking(
conversationId: string,
requested: ConversationModelState,
stateValue: unknown,
capabilitiesValue: unknown,
replaceSnapshot = true,
forcePersist = false,
): Promise<ConversationModelState> {
if (!requested.model) return clone(requested);
const effective = effectiveThinkingLevel(stateValue) ?? requested.model.thinkingLevel;
const available = availableThinkingLevels(capabilitiesValue);
if (!available.includes(effective)) available.push(effective);
const durable: ConversationModelState = {
model: { ...requested.model, thinkingLevel: effective },
modelResolution: 'resolved',
};
const persisted = !forcePersist && requested.model.thinkingLevel === effective
? durable
: await this.registry.setModel(conversationId, durable);
this.pool.updateConversationModel(conversationId, persisted);
const result: ConversationModelState = {
...persisted,
...(available.length > 0 ? { availableThinkingLevels: available } : {}),
};
if (replaceSnapshot) this.replaceModel(conversationId, result);
else {
const input = this.inputs.get(conversationId);
if (input) input.model = clone(persisted);
}
return clone(result);
}
private enqueueProjection(conversationId: string, action: () => Promise<void>): Promise<void> {
const previous = this.projectionChains.get(conversationId) ?? Promise.resolve();
const flight = previous.then(action);