fix: settle Pi child capacity reservations

electron/coding-runtime/pi/worker-pool.ts

@@ -230,6 +230,7 @@ export class PiWorkerPool {
   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;
@@ -274,6 +275,7 @@ export class PiWorkerPool {
       .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;
@@ -328,6 +330,7 @@ export class PiWorkerPool {
         lease,
       );
       this.workers.set(conversation.conversationId, record);
+      this.notifyReclaimableWorker();
       await this.trimIdleWorkers();
       return this.publicState(record);
     })().finally(() => {
@@ -344,13 +347,19 @@ export class PiWorkerPool {
     return record ? this.publicState(record) : null;
   }
 
-  async reclaimIdleWorker(): Promise<boolean> {
-    const record = [...this.workers.values()]
-      .filter((candidate) => candidate.state === 'ready' || candidate.state === 'idle')
-      .sort((left, right) => left.lastUsed - right.lastUsed)[0];
-    if (!record) return false;
-    await this.evict(record);
-    return true;
+  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 {
@@ -406,6 +415,7 @@ export class PiWorkerPool {
       const record = this.workers.get(conversationId);
       if (record && record.state !== 'crashed') {
         record.state = 'idle';
+        this.notifyReclaimableWorker();
         try {
           this.revisions.settleRun(record.revisionWorkerId);
         } catch {
@@ -605,6 +615,7 @@ export class PiWorkerPool {
       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();
       }
@@ -619,6 +630,7 @@ export class PiWorkerPool {
     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')) {
@@ -716,6 +728,7 @@ export class PiWorkerPool {
         ? 'running'
         : 'ready';
       this.workers.set(conversationId, replacement);
+      if (replacement.state === 'ready') this.notifyReclaimableWorker();
       this.emit({
         type: 'worker.replaced',
         conversationId,
@@ -796,15 +809,16 @@ export class PiWorkerPool {
     await Promise.all(idle.slice(0, excess).map((record) => this.evict(record)));
   }
 
-  private async evict(record: WorkerRecord): Promise<void> {
+  private async evict(record: WorkerRecord): Promise<boolean> {
     const conversationId = record.conversation.conversationId;
-    if (this.workers.get(conversationId) !== record) return;
+    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 {
@@ -902,9 +916,39 @@ export class PiWorkerPool {
   }
 
   private async stopAndRelease(record: WorkerRecord): Promise<void> {
-    await record.worker.stop();
-    record.processLease?.release();
-    record.processLease = null;
+    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 {