993 lines
34 KiB
TypeScript
993 lines
34 KiB
TypeScript
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>>;
|
|
send(command: PiRpcCommand): Promise<void>;
|
|
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;
|
|
priority: 'normal' | 'child';
|
|
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,
|
|
priority: 'normal' | 'child' = 'normal',
|
|
): 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,
|
|
priority,
|
|
...(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 childIndex = this.waiters.findIndex(({ priority }) => priority === 'child');
|
|
const [waiter] = this.waiters.splice(childIndex >= 0 ? childIndex : 0, 1);
|
|
if (!waiter) return;
|
|
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 readonly reclaimWaiters = new Set<() => 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);
|
|
this.notifyReclaimableWorker();
|
|
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);
|
|
this.notifyReclaimableWorker();
|
|
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;
|
|
}
|
|
|
|
async reclaimIdleWorker(signal?: AbortSignal): Promise<boolean> {
|
|
while (true) {
|
|
if (signal?.aborted) throw new Error('Pi idle worker reclaim cancelled');
|
|
if (this.shuttingDown) throw new Error('Pi worker pool is shutting down');
|
|
const record = [...this.workers.values()]
|
|
.filter((candidate) => candidate.state === 'ready' || candidate.state === 'idle')
|
|
.sort((left, right) => left.lastUsed - right.lastUsed)[0];
|
|
if (record) {
|
|
if (await this.evict(record)) return true;
|
|
continue;
|
|
}
|
|
await this.waitForReclaimableWorker(signal);
|
|
}
|
|
}
|
|
|
|
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();
|
|
}
|
|
}
|
|
|
|
async send(conversationId: string, command: PiRpcCommand): Promise<void> {
|
|
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;
|
|
await record.worker.send(command);
|
|
}
|
|
|
|
getActiveRun(conversationId: string): { runId: string; generation: number } | null {
|
|
const active = this.activeRuns.get(conversationId);
|
|
return active ? { ...active } : null;
|
|
}
|
|
|
|
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';
|
|
this.notifyReclaimableWorker();
|
|
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;
|
|
let record = this.workers.get(conversationId);
|
|
if (!record) throw new Error('Conversation worker is not prepared');
|
|
this.cancelConversationRuns(conversationId, new Error('Conversation worker is recovering'));
|
|
while (record.rebuildFlight) {
|
|
const existingFlight = record.rebuildFlight;
|
|
try {
|
|
return this.publicState(await existingFlight);
|
|
} catch (error) {
|
|
if (this.shuttingDown) throw error;
|
|
const current = this.workers.get(conversationId);
|
|
if (current && current !== record) {
|
|
record = current;
|
|
continue;
|
|
}
|
|
if (record.rebuildFlight !== existingFlight) continue;
|
|
record.rebuildFlight = undefined;
|
|
break;
|
|
}
|
|
}
|
|
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';
|
|
this.notifyReclaimableWorker();
|
|
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';
|
|
this.notifyReclaimableWorker();
|
|
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);
|
|
if (replacement.state === 'ready') this.notifyReclaimableWorker();
|
|
this.emit({
|
|
type: 'worker.replaced',
|
|
conversationId,
|
|
generation,
|
|
state: this.publicState(replacement),
|
|
});
|
|
await this.trimIdleWorkers();
|
|
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<boolean> {
|
|
const conversationId = record.conversation.conversationId;
|
|
if (this.workers.get(conversationId) !== record) return false;
|
|
record.unsubscribeEvent();
|
|
record.unsubscribeInvalidation();
|
|
this.cancelGenerationResources(record);
|
|
this.revisions.removeWorker(record.revisionWorkerId);
|
|
this.workers.delete(conversationId);
|
|
await this.stopAndRelease(record);
|
|
return true;
|
|
}
|
|
|
|
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> {
|
|
try {
|
|
await record.worker.stop();
|
|
} finally {
|
|
record.processLease?.release();
|
|
record.processLease = null;
|
|
}
|
|
}
|
|
|
|
private waitForReclaimableWorker(signal?: AbortSignal): Promise<void> {
|
|
return new Promise((resolve, reject) => {
|
|
const cleanup = () => {
|
|
this.reclaimWaiters.delete(wake);
|
|
signal?.removeEventListener('abort', cancel);
|
|
this.shutdownController.signal.removeEventListener('abort', cancel);
|
|
};
|
|
const wake = () => {
|
|
cleanup();
|
|
resolve();
|
|
};
|
|
const cancel = () => {
|
|
cleanup();
|
|
reject(new Error(
|
|
this.shuttingDown ? 'Pi worker pool is shutting down' : 'Pi idle worker reclaim cancelled',
|
|
));
|
|
};
|
|
this.reclaimWaiters.add(wake);
|
|
signal?.addEventListener('abort', cancel, { once: true });
|
|
this.shutdownController.signal.addEventListener('abort', cancel, { once: true });
|
|
});
|
|
}
|
|
|
|
private notifyReclaimableWorker(): void {
|
|
for (const wake of [...this.reclaimWaiters]) wake();
|
|
}
|
|
|
|
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 } : {}),
|
|
};
|
|
}
|
|
}
|