fix(pi): retain ownership for uncertain mutations

This commit is contained in:
2026-08-25 23:53:54 +08:00
parent 621ebb1781
commit 019cbb115a
21 changed files with 1113 additions and 67 deletions

View File

@@ -3,10 +3,11 @@ import type {
ConversationModelState,
PrepareConversationInput,
} from '../contracts';
import type { PiProcessError, PiProcessErrorCode } from './process-errors';
import { PiProcessError, type PiProcessErrorCode } from './process-errors';
import type {
PiRpcCommand,
PiRpcEvent,
PiRpcLateResult,
PiRpcRequestOptions,
PiRpcResponse,
} from './rpc-client';
@@ -108,6 +109,19 @@ export type PiWorkerPoolEvent =
generation: number;
reason: PiWorkerReplacementReason;
state: PiWorkerPoolState;
}
| {
type: 'top-level.confirmed';
conversationId: string;
generation: number;
runId: string;
}
| {
type: 'top-level.failed';
conversationId: string;
generation: number;
runId: string;
error: PiProcessError;
};
export type PiWorkerReplacementReason =
@@ -247,9 +261,13 @@ export class PiWorkerPool {
private readonly waitingRuns: PendingTopLevelRun[] = [];
private readonly activeRuns = new Map<string, {
runId: string;
run: PendingTopLevelRun;
generation: number;
cold: boolean;
acceptedAt?: number;
confirmationStartedAt?: number;
confirmation: 'pending' | 'uncertain' | 'confirmed';
cancelConfirmation?: () => void;
}>();
private readonly generations = new Map<string, number>();
private readonly listeners = new Set<(event: PiWorkerPoolEvent) => void>();
@@ -498,13 +516,14 @@ export class PiWorkerPool {
getActiveRun(conversationId: string): { runId: string; generation: number } | null {
const active = this.activeRuns.get(conversationId);
return active ? { ...active } : null;
return active ? { runId: active.runId, generation: active.generation } : null;
}
failTopLevel(conversationId: string, runId: string, error: Error): void {
const active = this.activeRuns.get(conversationId);
if (active?.runId === runId) {
this.activeRuns.delete(conversationId);
active.cancelConfirmation?.();
this.runningCount -= 1;
const record = this.workers.get(conversationId);
if (record && record.state !== 'crashed') {
@@ -667,8 +686,10 @@ export class PiWorkerPool {
await Promise.allSettled([...this.rebuildFlights]);
const records = [...this.workers.values()];
this.workers.clear();
const activeRuns = [...this.activeRuns.values()];
this.activeRuns.clear();
this.runningCount = 0;
for (const active of activeRuns) active.cancelConfirmation?.();
await Promise.all(records.map(async (record) => {
record.unsubscribeEvent();
record.unsubscribeInvalidation();
@@ -684,8 +705,10 @@ export class PiWorkerPool {
this.runningCount += 1;
this.activeRuns.set(run.conversationId, {
runId: run.runId,
run,
generation: record.generation,
cold: record.coldStart && record.acceptedPromptCount === 0,
confirmation: 'pending',
});
void this.acceptTopLevel(record, run);
}
@@ -708,29 +731,34 @@ export class PiWorkerPool {
'worker.queue_wait',
run.queuedAt === undefined ? 0 : this.now() - run.queuedAt,
);
const acceptedAt = this.now();
const response = await current.worker.request(run.command);
if (run.command.type === 'prompt') {
const acceptedFinishedAt = this.now();
active.acceptedAt = acceptedFinishedAt;
this.recordMilestone(
current,
run,
'prompt.accepted',
acceptedFinishedAt - acceptedAt,
active.cold,
);
current.acceptedPromptCount += 1;
}
active.confirmationStartedAt = this.now();
const confirmationController = new AbortController();
active.cancelConfirmation = () => confirmationController.abort();
const response = await current.worker.request(run.command, {
signal: confirmationController.signal,
retainAfterTimeout: true,
onLateResult: (result) => this.handleLateTopLevelResult(current!, run, result),
});
this.confirmTopLevel(current, run);
run.resolve(response);
} catch (error) {
const active = this.activeRuns.get(run.conversationId);
if (error instanceof PiProcessError
&& error.code === 'PI_RPC_TIMEOUT'
&& active?.runId === run.runId
&& current
&& active.generation === current.generation) {
if (active.confirmation !== 'confirmed') active.confirmation = 'uncertain';
run.reject(error);
return;
}
if (active) {
this.activeRuns.delete(run.conversationId);
active.cancelConfirmation?.();
this.runningCount -= 1;
this.launchWaitingRuns();
}
if (current && this.workers.get(run.conversationId) === current
if (active && current && this.workers.get(run.conversationId) === current
&& current.state !== 'crashed') {
current.state = 'idle';
this.notifyReclaimableWorker();
@@ -741,10 +769,84 @@ export class PiWorkerPool {
}
}
private confirmTopLevel(record: WorkerRecord, run: PendingTopLevelRun): boolean {
const active = this.activeRuns.get(run.conversationId);
if (!active
|| active.runId !== run.runId
|| active.generation !== record.generation
|| active.confirmation === 'confirmed') return false;
active.confirmation = 'confirmed';
active.cancelConfirmation = undefined;
if (run.command.type === 'prompt') {
const acceptedFinishedAt = this.now();
active.acceptedAt = acceptedFinishedAt;
this.recordMilestone(
record,
run,
'prompt.accepted',
acceptedFinishedAt - (active.confirmationStartedAt ?? acceptedFinishedAt),
active.cold,
);
record.acceptedPromptCount += 1;
}
return true;
}
private handleLateTopLevelResult(
record: WorkerRecord,
run: PendingTopLevelRun,
result: PiRpcLateResult,
): void {
const active = this.activeRuns.get(run.conversationId);
if (!active || active.runId !== run.runId || active.generation !== record.generation) return;
if (result.response) {
if (!this.confirmTopLevel(record, run)) return;
this.emit({
type: 'top-level.confirmed',
conversationId: run.conversationId,
generation: record.generation,
runId: run.runId,
});
return;
}
if (result.error.code === 'PI_RPC_EXITED'
|| result.error.code === 'PI_RPC_PROTOCOL_ERROR') return;
this.activeRuns.delete(run.conversationId);
active.cancelConfirmation = undefined;
this.runningCount -= 1;
if (this.workers.get(run.conversationId) === record && record.state !== 'crashed') {
record.state = 'idle';
this.notifyReclaimableWorker();
try {
this.revisions.settleRun(record.revisionWorkerId);
} catch {
// Recover/crash may already have removed this generation.
}
}
this.launchWaitingRuns();
void this.trimIdleWorkers();
this.emit({
type: 'top-level.failed',
conversationId: run.conversationId,
generation: record.generation,
runId: run.runId,
error: result.error,
});
}
private settleTopLevel(record: WorkerRecord): void {
const conversationId = record.conversation.conversationId;
const active = this.activeRuns.get(conversationId);
if (!active || active.generation !== record.generation) return;
const cancelConfirmation = active.cancelConfirmation;
if (active.confirmation !== 'confirmed') {
this.confirmTopLevel(record, active.run);
active.run.resolve({
type: 'response',
id: `agent-settled:${active.runId}`,
success: true,
});
}
if (active.acceptedAt !== undefined) {
this.recordMilestone(
record,
@@ -755,6 +857,7 @@ export class PiWorkerPool {
);
}
this.activeRuns.delete(conversationId);
cancelConfirmation?.();
this.runningCount -= 1;
record.state = 'idle';
this.notifyReclaimableWorker();
@@ -1033,6 +1136,7 @@ export class PiWorkerPool {
const active = this.activeRuns.get(conversationId);
if (active?.generation === record.generation) {
this.activeRuns.delete(conversationId);
active.cancelConfirmation?.();
this.runningCount -= 1;
}
for (let index = this.waitingRuns.length - 1; index >= 0; index -= 1) {
@@ -1060,6 +1164,7 @@ export class PiWorkerPool {
const active = this.activeRuns.get(conversationId);
if (active) {
this.activeRuns.delete(conversationId);
active.cancelConfirmation?.();
this.runningCount -= 1;
}
for (let index = this.waitingRuns.length - 1; index >= 0; index -= 1) {