feat: add Pi conversation worker pool
This commit is contained in:
898
electron/coding-runtime/pi/worker-pool.ts
Normal file
898
electron/coding-runtime/pi/worker-pool.ts
Normal file
@@ -0,0 +1,898 @@
|
||||
import type { PrepareConversationInput } from '../contracts';
|
||||
import type { ConversationModelState } from '../contracts';
|
||||
import type { PiProcessError, PiProcessErrorCode } from './process-errors';
|
||||
import type {
|
||||
PiRpcCommand,
|
||||
PiRpcEvent,
|
||||
PiRpcRequestOptions,
|
||||
PiRpcResponse,
|
||||
} from './rpc-client';
|
||||
import type { PiWorkerStopResult } from './worker-process';
|
||||
import {
|
||||
PiManagedInputRevisionCoordinator,
|
||||
type PiManagedInputRevision,
|
||||
} from './managed-input-revision';
|
||||
import {
|
||||
createPiRuntimeTelemetryEvent,
|
||||
type PiRuntimeTelemetryEvent,
|
||||
} from './telemetry';
|
||||
|
||||
export interface PiConversationWorker {
|
||||
readonly id: string;
|
||||
readonly generation: number;
|
||||
request<T = unknown>(
|
||||
command: PiRpcCommand,
|
||||
options?: PiRpcRequestOptions,
|
||||
): Promise<PiRpcResponse<T>>;
|
||||
subscribe(listener: (event: PiRpcEvent) => void): () => void;
|
||||
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void;
|
||||
stop(): Promise<PiWorkerStopResult>;
|
||||
}
|
||||
|
||||
export interface PiWorkerSessionBinding {
|
||||
piSessionId: string;
|
||||
sessionKey: string;
|
||||
}
|
||||
|
||||
export interface PiWorkerOpenInput {
|
||||
conversation: PrepareConversationInput;
|
||||
generation: number;
|
||||
revision: PiManagedInputRevision;
|
||||
existingSession?: PiWorkerSessionBinding;
|
||||
fork?: {
|
||||
sourceSession: PiWorkerSessionBinding;
|
||||
sourceEntryId?: string;
|
||||
};
|
||||
}
|
||||
|
||||
export interface PiWorkerOpenResult {
|
||||
worker: PiConversationWorker;
|
||||
session: PiWorkerSessionBinding;
|
||||
}
|
||||
|
||||
export interface PiWorkerPoolState {
|
||||
conversationId: string;
|
||||
workerId: string;
|
||||
state: 'spawning' | 'ready' | 'queued' | 'running' | 'idle' | 'crashed';
|
||||
generation: number;
|
||||
session: PiWorkerSessionBinding;
|
||||
failureCode?: PiProcessErrorCode;
|
||||
}
|
||||
|
||||
export interface PiTopLevelRunInput {
|
||||
conversationId: string;
|
||||
runId: string;
|
||||
command: PiRpcCommand;
|
||||
}
|
||||
|
||||
export interface PiTopLevelRunTicket {
|
||||
queuePosition?: number;
|
||||
accepted: Promise<PiRpcResponse>;
|
||||
}
|
||||
|
||||
export type PiGenerationResourceKind = 'command' | 'interaction' | 'child';
|
||||
|
||||
export interface PiGenerationResourceInput {
|
||||
conversationId: string;
|
||||
kind: PiGenerationResourceKind;
|
||||
id: string;
|
||||
cancel(): void;
|
||||
}
|
||||
|
||||
export type PiWorkerPoolEvent =
|
||||
| {
|
||||
type: 'worker.event';
|
||||
conversationId: string;
|
||||
generation: number;
|
||||
event: PiRpcEvent;
|
||||
}
|
||||
| {
|
||||
type: 'worker.crashed';
|
||||
conversationId: string;
|
||||
generation: number;
|
||||
error: PiProcessError;
|
||||
}
|
||||
| {
|
||||
type: 'worker.replaced';
|
||||
conversationId: string;
|
||||
generation: number;
|
||||
state: PiWorkerPoolState;
|
||||
};
|
||||
|
||||
export interface PiWorkerPoolOptions {
|
||||
openWorker(input: PiWorkerOpenInput): Promise<PiWorkerOpenResult>;
|
||||
maxRunning?: number;
|
||||
maxIdle?: number;
|
||||
processBudget?: PiProcessBudget;
|
||||
revisionCoordinator?: PiManagedInputRevisionCoordinator;
|
||||
now?: () => number;
|
||||
onTelemetry?: (event: PiRuntimeTelemetryEvent) => void;
|
||||
}
|
||||
|
||||
export interface PiProcessLease {
|
||||
release(): void;
|
||||
}
|
||||
|
||||
interface PiProcessBudgetWaiter {
|
||||
signal?: AbortSignal;
|
||||
resolve(lease: PiProcessLease): void;
|
||||
reject(error: Error): void;
|
||||
onAbort?: () => void;
|
||||
}
|
||||
|
||||
export class PiProcessBudget {
|
||||
private readonly waiters: PiProcessBudgetWaiter[] = [];
|
||||
private active = 0;
|
||||
|
||||
constructor(readonly maxProcesses = 8) {
|
||||
if (!Number.isSafeInteger(maxProcesses) || maxProcesses <= 0) {
|
||||
throw new Error('maxProcesses must be a positive safe integer');
|
||||
}
|
||||
}
|
||||
|
||||
get activeCount(): number { return this.active; }
|
||||
get waitingCount(): number { return this.waiters.length; }
|
||||
|
||||
acquire(signal?: AbortSignal): Promise<PiProcessLease> {
|
||||
if (signal?.aborted) return Promise.reject(new Error('Pi process budget acquisition cancelled'));
|
||||
if (this.active < this.maxProcesses) return Promise.resolve(this.issueLease());
|
||||
return new Promise<PiProcessLease>((resolve, reject) => {
|
||||
const waiter: PiProcessBudgetWaiter = { resolve, reject, ...(signal ? { signal } : {}) };
|
||||
if (signal) {
|
||||
waiter.onAbort = () => {
|
||||
const index = this.waiters.indexOf(waiter);
|
||||
if (index >= 0) this.waiters.splice(index, 1);
|
||||
reject(new Error('Pi process budget acquisition cancelled'));
|
||||
};
|
||||
signal.addEventListener('abort', waiter.onAbort, { once: true });
|
||||
}
|
||||
this.waiters.push(waiter);
|
||||
});
|
||||
}
|
||||
|
||||
private issueLease(): PiProcessLease {
|
||||
this.active += 1;
|
||||
let released = false;
|
||||
return {
|
||||
release: () => {
|
||||
if (released) return;
|
||||
released = true;
|
||||
this.active -= 1;
|
||||
this.advance();
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
private advance(): void {
|
||||
while (this.active < this.maxProcesses && this.waiters.length > 0) {
|
||||
const waiter = this.waiters.shift() as PiProcessBudgetWaiter;
|
||||
if (waiter.onAbort && waiter.signal) {
|
||||
waiter.signal.removeEventListener('abort', waiter.onAbort);
|
||||
}
|
||||
if (waiter.signal?.aborted) {
|
||||
waiter.reject(new Error('Pi process budget acquisition cancelled'));
|
||||
continue;
|
||||
}
|
||||
waiter.resolve(this.issueLease());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
interface WorkerRecord {
|
||||
conversation: PrepareConversationInput;
|
||||
worker: PiConversationWorker;
|
||||
session: PiWorkerSessionBinding;
|
||||
generation: number;
|
||||
revisionWorkerId: string;
|
||||
lastUsed: number;
|
||||
state: PiWorkerPoolState['state'];
|
||||
unsubscribeEvent: () => void;
|
||||
unsubscribeInvalidation: () => void;
|
||||
failureCode?: PiProcessErrorCode;
|
||||
generationResources: Record<PiGenerationResourceKind, Map<string, () => void>>;
|
||||
rebuildFlight?: Promise<WorkerRecord>;
|
||||
acceptedPromptCount: number;
|
||||
processLease: PiProcessLease | null;
|
||||
processStopFlight?: Promise<void>;
|
||||
reconfigureAfterSettled: boolean;
|
||||
}
|
||||
|
||||
interface PendingTopLevelRun extends PiTopLevelRunInput {
|
||||
resolve(response: PiRpcResponse): void;
|
||||
reject(error: unknown): void;
|
||||
queuedAt?: number;
|
||||
}
|
||||
|
||||
export class PiWorkerPool {
|
||||
private readonly openWorker: PiWorkerPoolOptions['openWorker'];
|
||||
private readonly maxRunning: number;
|
||||
private readonly maxIdle: number;
|
||||
private readonly processBudget: PiProcessBudget;
|
||||
private readonly revisions: PiManagedInputRevisionCoordinator;
|
||||
private readonly now: () => number;
|
||||
private readonly onTelemetry: ((event: PiRuntimeTelemetryEvent) => void) | undefined;
|
||||
private readonly workers = new Map<string, WorkerRecord>();
|
||||
private readonly prepareFlights = new Map<string, Promise<PiWorkerPoolState>>();
|
||||
private readonly rebuildFlights = new Set<Promise<WorkerRecord>>();
|
||||
private readonly waitingRuns: PendingTopLevelRun[] = [];
|
||||
private readonly activeRuns = new Map<string, { runId: string; generation: number }>();
|
||||
private readonly generations = new Map<string, number>();
|
||||
private readonly listeners = new Set<(event: PiWorkerPoolEvent) => void>();
|
||||
private commandSequence = 0;
|
||||
private runningCount = 0;
|
||||
private useSequence = 0;
|
||||
private shuttingDown = false;
|
||||
private shutdownFlight: Promise<void> | null = null;
|
||||
private readonly shutdownController = new AbortController();
|
||||
|
||||
constructor(options: PiWorkerPoolOptions) {
|
||||
this.openWorker = options.openWorker;
|
||||
this.maxRunning = options.maxRunning ?? 4;
|
||||
this.maxIdle = options.maxIdle ?? 4;
|
||||
this.processBudget = options.processBudget ?? new PiProcessBudget();
|
||||
this.revisions = options.revisionCoordinator ?? new PiManagedInputRevisionCoordinator();
|
||||
this.now = options.now ?? Date.now;
|
||||
this.onTelemetry = options.onTelemetry;
|
||||
if (!Number.isSafeInteger(this.maxRunning) || this.maxRunning <= 0) {
|
||||
throw new Error('maxRunning must be a positive safe integer');
|
||||
}
|
||||
if (!Number.isSafeInteger(this.maxIdle) || this.maxIdle < 0) {
|
||||
throw new Error('maxIdle must be a non-negative safe integer');
|
||||
}
|
||||
}
|
||||
|
||||
prepare(conversation: PrepareConversationInput): Promise<PiWorkerPoolState> {
|
||||
if (this.shuttingDown) return Promise.reject(new Error('Pi worker pool is shutting down'));
|
||||
const existing = this.workers.get(conversation.conversationId);
|
||||
if (existing) {
|
||||
this.touch(existing);
|
||||
void this.trimIdleWorkers();
|
||||
return Promise.resolve(this.publicState(existing));
|
||||
}
|
||||
const pending = this.prepareFlights.get(conversation.conversationId);
|
||||
if (pending) return pending;
|
||||
|
||||
const generation = this.nextGeneration(conversation.conversationId);
|
||||
const revision = this.revisions.current;
|
||||
const flight = this.openWithLease({
|
||||
conversation: structuredClone(conversation),
|
||||
generation,
|
||||
revision,
|
||||
})
|
||||
.then(async ({ opened: { worker, session }, lease }) => {
|
||||
const record = this.createRecord(conversation, worker, session, generation, revision, lease);
|
||||
this.workers.set(conversation.conversationId, record);
|
||||
const state = this.publicState(record);
|
||||
await this.trimIdleWorkers();
|
||||
return state;
|
||||
})
|
||||
.finally(() => {
|
||||
if (this.prepareFlights.get(conversation.conversationId) === flight) {
|
||||
this.prepareFlights.delete(conversation.conversationId);
|
||||
}
|
||||
});
|
||||
this.prepareFlights.set(conversation.conversationId, flight);
|
||||
return flight;
|
||||
}
|
||||
|
||||
markProviderStale(): PiManagedInputRevision {
|
||||
return this.revisions.markProviderStale();
|
||||
}
|
||||
|
||||
markResourcesStale(): PiManagedInputRevision {
|
||||
return this.revisions.markResourcesStale();
|
||||
}
|
||||
|
||||
async fork(
|
||||
sourceConversationId: string,
|
||||
conversation: PrepareConversationInput,
|
||||
sourceEntryId?: string,
|
||||
): Promise<PiWorkerPoolState> {
|
||||
if (this.shuttingDown) throw new Error('Pi worker pool is shutting down');
|
||||
const source = this.workers.get(sourceConversationId);
|
||||
if (!source || source.state === 'crashed') throw new Error('Source Conversation worker is not available');
|
||||
if (this.workers.has(conversation.conversationId)
|
||||
|| this.prepareFlights.has(conversation.conversationId)) {
|
||||
throw new Error('Fork target Conversation worker already exists');
|
||||
}
|
||||
const generation = this.nextGeneration(conversation.conversationId);
|
||||
const revision = this.revisions.current;
|
||||
const flight = (async () => {
|
||||
const { opened, lease } = await this.openWithLease({
|
||||
conversation: structuredClone(conversation),
|
||||
generation,
|
||||
revision,
|
||||
fork: {
|
||||
sourceSession: structuredClone(source.session),
|
||||
...(sourceEntryId ? { sourceEntryId } : {}),
|
||||
},
|
||||
});
|
||||
const record = this.createRecord(
|
||||
conversation,
|
||||
opened.worker,
|
||||
opened.session,
|
||||
generation,
|
||||
revision,
|
||||
lease,
|
||||
);
|
||||
this.workers.set(conversation.conversationId, record);
|
||||
await this.trimIdleWorkers();
|
||||
return this.publicState(record);
|
||||
})().finally(() => {
|
||||
if (this.prepareFlights.get(conversation.conversationId) === flight) {
|
||||
this.prepareFlights.delete(conversation.conversationId);
|
||||
}
|
||||
});
|
||||
this.prepareFlights.set(conversation.conversationId, flight);
|
||||
return await flight;
|
||||
}
|
||||
|
||||
getState(conversationId: string): PiWorkerPoolState | null {
|
||||
const record = this.workers.get(conversationId);
|
||||
return record ? this.publicState(record) : null;
|
||||
}
|
||||
|
||||
subscribe(listener: (event: PiWorkerPoolEvent) => void): () => void {
|
||||
this.listeners.add(listener);
|
||||
return () => this.listeners.delete(listener);
|
||||
}
|
||||
|
||||
async request<T = unknown>(
|
||||
conversationId: string,
|
||||
command: PiRpcCommand,
|
||||
options: PiRpcRequestOptions = {},
|
||||
): Promise<PiRpcResponse<T>> {
|
||||
let record = this.workers.get(conversationId);
|
||||
if (!record || record.state === 'crashed') throw new Error('Conversation worker is not available');
|
||||
if (record.rebuildFlight) record = await record.rebuildFlight;
|
||||
const controller = new AbortController();
|
||||
const abort = () => controller.abort();
|
||||
options.signal?.addEventListener('abort', abort, { once: true });
|
||||
const untrack = this.trackGenerationResource({
|
||||
conversationId,
|
||||
kind: 'command',
|
||||
id: `command-${++this.commandSequence}`,
|
||||
cancel: abort,
|
||||
});
|
||||
try {
|
||||
return await record.worker.request<T>(command, {
|
||||
...options,
|
||||
signal: controller.signal,
|
||||
});
|
||||
} finally {
|
||||
options.signal?.removeEventListener('abort', abort);
|
||||
untrack();
|
||||
}
|
||||
}
|
||||
|
||||
failTopLevel(conversationId: string, runId: string, error: Error): void {
|
||||
const active = this.activeRuns.get(conversationId);
|
||||
if (active?.runId === runId) {
|
||||
this.activeRuns.delete(conversationId);
|
||||
this.runningCount -= 1;
|
||||
const record = this.workers.get(conversationId);
|
||||
if (record && record.state !== 'crashed') {
|
||||
record.state = 'idle';
|
||||
try {
|
||||
this.revisions.settleRun(record.revisionWorkerId);
|
||||
} catch {
|
||||
// Recover/crash may already have removed this generation.
|
||||
}
|
||||
}
|
||||
}
|
||||
for (let index = this.waitingRuns.length - 1; index >= 0; index -= 1) {
|
||||
const pending = this.waitingRuns[index];
|
||||
if (pending?.conversationId !== conversationId || pending.runId !== runId) continue;
|
||||
this.waitingRuns.splice(index, 1);
|
||||
pending.reject(error);
|
||||
}
|
||||
this.launchWaitingRuns();
|
||||
}
|
||||
|
||||
updateConversationModel(conversationId: string, model: ConversationModelState): void {
|
||||
const record = this.workers.get(conversationId);
|
||||
if (!record) throw new Error('Conversation worker is not prepared');
|
||||
record.conversation = {
|
||||
...record.conversation,
|
||||
model: structuredClone(model),
|
||||
};
|
||||
}
|
||||
|
||||
async reconfigureConversationModel(
|
||||
conversationId: string,
|
||||
model: ConversationModelState,
|
||||
): Promise<PiWorkerPoolState | null> {
|
||||
if (this.shuttingDown) throw new Error('Pi worker pool is shutting down');
|
||||
let record = this.workers.get(conversationId);
|
||||
if (!record) throw new Error('Conversation worker is not prepared');
|
||||
if (record.rebuildFlight) record = await record.rebuildFlight;
|
||||
record.conversation = {
|
||||
...record.conversation,
|
||||
model: structuredClone(model),
|
||||
};
|
||||
if (this.activeRuns.has(conversationId)) {
|
||||
record.reconfigureAfterSettled = true;
|
||||
return null;
|
||||
}
|
||||
record.rebuildFlight = this.beginRebuild(record, this.revisions.current);
|
||||
return this.publicState(await record.rebuildFlight);
|
||||
}
|
||||
|
||||
trackGenerationResource(input: PiGenerationResourceInput): () => void {
|
||||
const record = this.workers.get(input.conversationId);
|
||||
if (!record || record.state === 'crashed') {
|
||||
throw new Error('Conversation worker generation is not available');
|
||||
}
|
||||
const resources = record.generationResources[input.kind];
|
||||
if (resources.has(input.id)) throw new Error(`Duplicate ${input.kind} resource: ${input.id}`);
|
||||
resources.set(input.id, input.cancel);
|
||||
return () => {
|
||||
if (this.workers.get(input.conversationId) === record) resources.delete(input.id);
|
||||
};
|
||||
}
|
||||
|
||||
startTopLevel(input: PiTopLevelRunInput): PiTopLevelRunTicket {
|
||||
if (this.shuttingDown) throw new Error('Pi worker pool is shutting down');
|
||||
const record = this.workers.get(input.conversationId);
|
||||
if (!record) throw new Error('Conversation worker is not prepared');
|
||||
if (this.activeRuns.has(input.conversationId)
|
||||
|| this.waitingRuns.some((run) => run.conversationId === input.conversationId)) {
|
||||
throw new Error('Conversation already has a top-level run');
|
||||
}
|
||||
|
||||
let resolve!: (response: PiRpcResponse) => void;
|
||||
let reject!: (error: unknown) => void;
|
||||
const accepted = new Promise<PiRpcResponse>((done, fail) => {
|
||||
resolve = done;
|
||||
reject = fail;
|
||||
});
|
||||
const pending: PendingTopLevelRun = {
|
||||
...structuredClone(input),
|
||||
resolve,
|
||||
reject,
|
||||
};
|
||||
if (this.runningCount < this.maxRunning) {
|
||||
this.launchTopLevel(record, pending);
|
||||
return { accepted };
|
||||
}
|
||||
this.waitingRuns.push(pending);
|
||||
pending.queuedAt = this.now();
|
||||
record.state = 'queued';
|
||||
return { queuePosition: this.waitingRuns.length, accepted };
|
||||
}
|
||||
|
||||
shutdown(): Promise<void> {
|
||||
if (!this.shutdownFlight) this.shutdownFlight = this.performShutdown();
|
||||
return this.shutdownFlight;
|
||||
}
|
||||
|
||||
async recover(conversationId: string): Promise<PiWorkerPoolState> {
|
||||
if (this.shuttingDown) throw new Error('Pi worker pool is shutting down');
|
||||
const pendingPrepare = this.prepareFlights.get(conversationId);
|
||||
if (pendingPrepare) await pendingPrepare;
|
||||
const record = this.workers.get(conversationId);
|
||||
if (!record) throw new Error('Conversation worker is not prepared');
|
||||
this.cancelConversationRuns(conversationId, new Error('Conversation worker is recovering'));
|
||||
record.rebuildFlight = this.beginRebuild(record, this.revisions.current);
|
||||
return this.publicState(await record.rebuildFlight);
|
||||
}
|
||||
|
||||
async dispose(conversationId: string): Promise<void> {
|
||||
const pendingPrepare = this.prepareFlights.get(conversationId);
|
||||
if (pendingPrepare) await pendingPrepare.catch(() => undefined);
|
||||
let record = this.workers.get(conversationId);
|
||||
if (!record) return;
|
||||
if (record.rebuildFlight) {
|
||||
record = await record.rebuildFlight.catch(() => record as WorkerRecord);
|
||||
}
|
||||
this.cancelConversationRuns(conversationId, new Error('Conversation worker was disposed'));
|
||||
record.unsubscribeEvent();
|
||||
record.unsubscribeInvalidation();
|
||||
this.cancelGenerationResources(record);
|
||||
this.revisions.removeWorker(record.revisionWorkerId);
|
||||
if (this.workers.get(conversationId) === record) this.workers.delete(conversationId);
|
||||
await this.stopAndRelease(record);
|
||||
}
|
||||
|
||||
private async performShutdown(): Promise<void> {
|
||||
this.shuttingDown = true;
|
||||
this.shutdownController.abort();
|
||||
const error = new Error('Pi worker pool is shutting down');
|
||||
for (const pending of this.waitingRuns.splice(0)) pending.reject(error);
|
||||
await Promise.allSettled([...this.prepareFlights.values()]);
|
||||
await Promise.allSettled([...this.rebuildFlights]);
|
||||
const records = [...this.workers.values()];
|
||||
this.workers.clear();
|
||||
this.activeRuns.clear();
|
||||
this.runningCount = 0;
|
||||
await Promise.all(records.map(async (record) => {
|
||||
record.unsubscribeEvent();
|
||||
record.unsubscribeInvalidation();
|
||||
this.cancelGenerationResources(record);
|
||||
this.revisions.removeWorker(record.revisionWorkerId);
|
||||
await this.stopAndRelease(record);
|
||||
}));
|
||||
}
|
||||
|
||||
private launchTopLevel(record: WorkerRecord, run: PendingTopLevelRun): void {
|
||||
record.state = 'running';
|
||||
this.touch(record);
|
||||
this.runningCount += 1;
|
||||
this.activeRuns.set(run.conversationId, {
|
||||
runId: run.runId,
|
||||
generation: record.generation,
|
||||
});
|
||||
void this.acceptTopLevel(record, run);
|
||||
}
|
||||
|
||||
private async acceptTopLevel(record: WorkerRecord, run: PendingTopLevelRun): Promise<void> {
|
||||
let current: WorkerRecord | undefined;
|
||||
let beganRun = false;
|
||||
try {
|
||||
current = await this.ensureFresh(record);
|
||||
const active = this.activeRuns.get(run.conversationId);
|
||||
if (!active) throw new Error('Top-level run was cancelled before acceptance');
|
||||
active.generation = current.generation;
|
||||
current.state = 'running';
|
||||
this.revisions.beginRun(current.revisionWorkerId);
|
||||
beganRun = true;
|
||||
if (run.queuedAt !== undefined) {
|
||||
this.recordMilestone(current, run, 'worker.queue_wait', this.now() - run.queuedAt);
|
||||
}
|
||||
const acceptedAt = this.now();
|
||||
const response = await current.worker.request(run.command);
|
||||
if (run.command.type === 'prompt') {
|
||||
this.recordMilestone(current, run, 'prompt.accepted', this.now() - acceptedAt);
|
||||
current.acceptedPromptCount += 1;
|
||||
}
|
||||
run.resolve(response);
|
||||
} catch (error) {
|
||||
const active = this.activeRuns.get(run.conversationId);
|
||||
if (active) {
|
||||
this.activeRuns.delete(run.conversationId);
|
||||
this.runningCount -= 1;
|
||||
this.launchWaitingRuns();
|
||||
}
|
||||
if (current && this.workers.get(run.conversationId) === current
|
||||
&& current.state !== 'crashed') {
|
||||
current.state = 'idle';
|
||||
if (beganRun) this.revisions.settleRun(current.revisionWorkerId);
|
||||
void this.trimIdleWorkers();
|
||||
}
|
||||
run.reject(error);
|
||||
}
|
||||
}
|
||||
|
||||
private settleTopLevel(record: WorkerRecord): void {
|
||||
const conversationId = record.conversation.conversationId;
|
||||
const active = this.activeRuns.get(conversationId);
|
||||
if (!active || active.generation !== record.generation) return;
|
||||
this.activeRuns.delete(conversationId);
|
||||
this.runningCount -= 1;
|
||||
record.state = 'idle';
|
||||
const revisionAction = this.revisions.settleRun(record.revisionWorkerId);
|
||||
if (!this.shuttingDown
|
||||
&& (record.reconfigureAfterSettled || revisionAction.action === 'rebuild-after-settled')) {
|
||||
record.rebuildFlight = this.beginRebuild(
|
||||
record,
|
||||
revisionAction.action === 'rebuild-after-settled'
|
||||
? revisionAction.revision
|
||||
: this.revisions.current,
|
||||
);
|
||||
void record.rebuildFlight.catch(() => undefined);
|
||||
}
|
||||
this.launchWaitingRuns();
|
||||
void this.trimIdleWorkers();
|
||||
}
|
||||
|
||||
private async ensureFresh(record: WorkerRecord): Promise<WorkerRecord> {
|
||||
const current = this.workers.get(record.conversation.conversationId);
|
||||
if (!current) throw new Error('Conversation worker is no longer available');
|
||||
if (current !== record) return await this.ensureFresh(current);
|
||||
if (record.rebuildFlight) return await record.rebuildFlight;
|
||||
const action = this.revisions.beforePrompt(record.revisionWorkerId);
|
||||
if (action.action !== 'rebuild-before-prompt') return record;
|
||||
record.rebuildFlight = this.beginRebuild(record, action.revision);
|
||||
return await record.rebuildFlight;
|
||||
}
|
||||
|
||||
private beginRebuild(
|
||||
record: WorkerRecord,
|
||||
revision: PiManagedInputRevision,
|
||||
): Promise<WorkerRecord> {
|
||||
const flight = this.rebuild(record, revision).finally(() => {
|
||||
this.rebuildFlights.delete(flight);
|
||||
});
|
||||
this.rebuildFlights.add(flight);
|
||||
return flight;
|
||||
}
|
||||
|
||||
private async rebuild(
|
||||
record: WorkerRecord,
|
||||
revision: PiManagedInputRevision,
|
||||
): Promise<WorkerRecord> {
|
||||
const conversationId = record.conversation.conversationId;
|
||||
record.state = 'spawning';
|
||||
record.unsubscribeEvent();
|
||||
record.unsubscribeInvalidation();
|
||||
this.cancelGenerationResources(record);
|
||||
this.revisions.removeWorker(record.revisionWorkerId);
|
||||
try {
|
||||
if (record.processStopFlight) await record.processStopFlight;
|
||||
else await record.worker.stop();
|
||||
let lease = record.processLease;
|
||||
if (!lease) lease = await this.processBudget.acquire(this.shutdownController.signal);
|
||||
if (this.shuttingDown) {
|
||||
lease.release();
|
||||
if (record.processLease === lease) record.processLease = null;
|
||||
throw new Error('Pi worker pool is shutting down');
|
||||
}
|
||||
const generation = this.nextGeneration(conversationId);
|
||||
let opened: PiWorkerOpenResult;
|
||||
try {
|
||||
opened = await this.openWorker({
|
||||
conversation: structuredClone(record.conversation),
|
||||
generation,
|
||||
revision: structuredClone(revision),
|
||||
existingSession: structuredClone(record.session),
|
||||
});
|
||||
if (this.shuttingDown) {
|
||||
await opened.worker.stop().catch(() => undefined);
|
||||
lease.release();
|
||||
if (record.processLease === lease) record.processLease = null;
|
||||
throw new Error('Pi worker pool is shutting down');
|
||||
}
|
||||
} catch (error) {
|
||||
lease.release();
|
||||
if (record.processLease === lease) record.processLease = null;
|
||||
throw error;
|
||||
}
|
||||
if (opened.session.piSessionId !== record.session.piSessionId
|
||||
|| opened.session.sessionKey !== record.session.sessionKey) {
|
||||
await opened.worker.stop().catch(() => undefined);
|
||||
lease.release();
|
||||
if (record.processLease === lease) record.processLease = null;
|
||||
throw new Error('Reopened Pi worker returned a different session binding');
|
||||
}
|
||||
record.processLease = null;
|
||||
const replacement = this.createRecord(
|
||||
record.conversation,
|
||||
opened.worker,
|
||||
opened.session,
|
||||
generation,
|
||||
revision,
|
||||
lease,
|
||||
);
|
||||
replacement.state = record.state === 'spawning' && this.activeRuns.has(conversationId)
|
||||
? 'running'
|
||||
: 'ready';
|
||||
this.workers.set(conversationId, replacement);
|
||||
this.emit({
|
||||
type: 'worker.replaced',
|
||||
conversationId,
|
||||
generation,
|
||||
state: this.publicState(replacement),
|
||||
});
|
||||
return replacement;
|
||||
} catch (error) {
|
||||
record.state = 'crashed';
|
||||
record.failureCode = error && typeof error === 'object' && 'code' in error
|
||||
? (error as { code: PiProcessErrorCode }).code
|
||||
: 'PI_WORKER_START_FAILED';
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
private createRecord(
|
||||
conversation: PrepareConversationInput,
|
||||
worker: PiConversationWorker,
|
||||
session: PiWorkerSessionBinding,
|
||||
generation: number,
|
||||
revision: PiManagedInputRevision,
|
||||
processLease: PiProcessLease,
|
||||
): WorkerRecord {
|
||||
const revisionWorkerId = `${conversation.conversationId}:${generation}`;
|
||||
this.revisions.registerWorker(revisionWorkerId, revision);
|
||||
const record: WorkerRecord = {
|
||||
conversation: structuredClone(conversation),
|
||||
worker,
|
||||
session: structuredClone(session),
|
||||
generation,
|
||||
revisionWorkerId,
|
||||
lastUsed: ++this.useSequence,
|
||||
state: 'ready',
|
||||
unsubscribeEvent: () => undefined,
|
||||
unsubscribeInvalidation: () => undefined,
|
||||
generationResources: {
|
||||
command: new Map(),
|
||||
interaction: new Map(),
|
||||
child: new Map(),
|
||||
},
|
||||
acceptedPromptCount: 0,
|
||||
processLease,
|
||||
reconfigureAfterSettled: false,
|
||||
};
|
||||
record.unsubscribeEvent = worker.subscribe((event) => {
|
||||
if (event.type === 'agent_settled') this.settleTopLevel(record);
|
||||
this.emit({
|
||||
type: 'worker.event',
|
||||
conversationId: conversation.conversationId,
|
||||
generation,
|
||||
event: structuredClone(event),
|
||||
});
|
||||
});
|
||||
record.unsubscribeInvalidation = worker.subscribeInvalidation((error) => {
|
||||
this.handleInvalidation(record, error);
|
||||
});
|
||||
return record;
|
||||
}
|
||||
|
||||
private nextGeneration(conversationId: string): number {
|
||||
const generation = (this.generations.get(conversationId) ?? 0) + 1;
|
||||
this.generations.set(conversationId, generation);
|
||||
return generation;
|
||||
}
|
||||
|
||||
private touch(record: WorkerRecord): void {
|
||||
record.lastUsed = ++this.useSequence;
|
||||
}
|
||||
|
||||
private async trimIdleWorkers(): Promise<void> {
|
||||
const idle = [...this.workers.values()]
|
||||
.filter((record) => record.state === 'ready' || record.state === 'idle')
|
||||
.sort((left, right) => left.lastUsed - right.lastUsed);
|
||||
const excess = idle.length - this.maxIdle;
|
||||
if (excess <= 0) return;
|
||||
await Promise.all(idle.slice(0, excess).map((record) => this.evict(record)));
|
||||
}
|
||||
|
||||
private async evict(record: WorkerRecord): Promise<void> {
|
||||
const conversationId = record.conversation.conversationId;
|
||||
if (this.workers.get(conversationId) !== record) return;
|
||||
record.unsubscribeEvent();
|
||||
record.unsubscribeInvalidation();
|
||||
this.cancelGenerationResources(record);
|
||||
this.revisions.removeWorker(record.revisionWorkerId);
|
||||
this.workers.delete(conversationId);
|
||||
await this.stopAndRelease(record);
|
||||
}
|
||||
|
||||
private launchWaitingRuns(): void {
|
||||
while (this.runningCount < this.maxRunning && this.waitingRuns.length > 0) {
|
||||
const pending = this.waitingRuns.shift() as PendingTopLevelRun;
|
||||
const record = this.workers.get(pending.conversationId);
|
||||
if (!record) {
|
||||
pending.reject(new Error('Conversation worker is no longer available'));
|
||||
continue;
|
||||
}
|
||||
this.launchTopLevel(record, pending);
|
||||
}
|
||||
}
|
||||
|
||||
private handleInvalidation(record: WorkerRecord, error: PiProcessError): void {
|
||||
const conversationId = record.conversation.conversationId;
|
||||
if (this.workers.get(conversationId) !== record || record.state === 'crashed') return;
|
||||
record.state = 'crashed';
|
||||
record.failureCode = error.code;
|
||||
record.unsubscribeEvent();
|
||||
this.revisions.removeWorker(record.revisionWorkerId);
|
||||
this.cancelGenerationResources(record);
|
||||
if (!record.processStopFlight) {
|
||||
record.processStopFlight = this.stopAndRelease(record).catch(() => undefined);
|
||||
}
|
||||
|
||||
const active = this.activeRuns.get(conversationId);
|
||||
if (active?.generation === record.generation) {
|
||||
this.activeRuns.delete(conversationId);
|
||||
this.runningCount -= 1;
|
||||
}
|
||||
for (let index = this.waitingRuns.length - 1; index >= 0; index -= 1) {
|
||||
const pending = this.waitingRuns[index];
|
||||
if (pending?.conversationId !== conversationId) continue;
|
||||
this.waitingRuns.splice(index, 1);
|
||||
pending.reject(error);
|
||||
}
|
||||
this.launchWaitingRuns();
|
||||
this.emit({
|
||||
type: 'worker.crashed',
|
||||
conversationId,
|
||||
generation: record.generation,
|
||||
error,
|
||||
});
|
||||
}
|
||||
|
||||
private cancelConversationRuns(conversationId: string, error: Error): void {
|
||||
const active = this.activeRuns.get(conversationId);
|
||||
if (active) {
|
||||
this.activeRuns.delete(conversationId);
|
||||
this.runningCount -= 1;
|
||||
}
|
||||
for (let index = this.waitingRuns.length - 1; index >= 0; index -= 1) {
|
||||
const pending = this.waitingRuns[index];
|
||||
if (pending?.conversationId !== conversationId) continue;
|
||||
this.waitingRuns.splice(index, 1);
|
||||
pending.reject(error);
|
||||
}
|
||||
this.launchWaitingRuns();
|
||||
}
|
||||
|
||||
private cancelGenerationResources(record: WorkerRecord): void {
|
||||
for (const resources of Object.values(record.generationResources)) {
|
||||
for (const cancel of resources.values()) {
|
||||
try {
|
||||
cancel();
|
||||
} catch {
|
||||
// A failed cleanup must not prevent sibling resources from being cancelled.
|
||||
}
|
||||
}
|
||||
resources.clear();
|
||||
}
|
||||
}
|
||||
|
||||
private async openWithLease(
|
||||
input: PiWorkerOpenInput,
|
||||
): Promise<{ opened: PiWorkerOpenResult; lease: PiProcessLease }> {
|
||||
const lease = await this.processBudget.acquire(this.shutdownController.signal);
|
||||
if (this.shuttingDown) {
|
||||
lease.release();
|
||||
throw new Error('Pi worker pool is shutting down');
|
||||
}
|
||||
try {
|
||||
const opened = await this.openWorker(input);
|
||||
if (this.shuttingDown) {
|
||||
await opened.worker.stop().catch(() => undefined);
|
||||
lease.release();
|
||||
throw new Error('Pi worker pool is shutting down');
|
||||
}
|
||||
return { opened, lease };
|
||||
} catch (error) {
|
||||
lease.release();
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
private async stopAndRelease(record: WorkerRecord): Promise<void> {
|
||||
await record.worker.stop();
|
||||
record.processLease?.release();
|
||||
record.processLease = null;
|
||||
}
|
||||
|
||||
private emit(event: PiWorkerPoolEvent): void {
|
||||
for (const listener of this.listeners) {
|
||||
try {
|
||||
listener(event);
|
||||
} catch {
|
||||
// Runtime observers must not affect worker lifecycle.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private recordMilestone(
|
||||
record: WorkerRecord,
|
||||
run: PendingTopLevelRun,
|
||||
milestone: 'worker.queue_wait' | 'prompt.accepted',
|
||||
durationMs: number,
|
||||
): void {
|
||||
if (!this.onTelemetry) return;
|
||||
this.onTelemetry(createPiRuntimeTelemetryEvent({
|
||||
milestone,
|
||||
conversationId: record.conversation.conversationId,
|
||||
workerGeneration: record.generation,
|
||||
runId: run.runId,
|
||||
cold: record.acceptedPromptCount === 0,
|
||||
durationMs,
|
||||
at: this.now(),
|
||||
}));
|
||||
}
|
||||
|
||||
private publicState(record: WorkerRecord): PiWorkerPoolState {
|
||||
return {
|
||||
conversationId: record.conversation.conversationId,
|
||||
workerId: record.worker.id,
|
||||
state: record.state,
|
||||
generation: record.generation,
|
||||
session: structuredClone(record.session),
|
||||
...(record.failureCode ? { failureCode: record.failureCode } : {}),
|
||||
};
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user