fix(pi): retain ownership for uncertain mutations
This commit is contained in:
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user