fix(coding-runtime): close Pi extension lifecycle races

electron/coding-runtime/pi/extension-host.ts

@@ -63,9 +63,11 @@ export class PiManagedExtensionHost {
   private readonly leases: PiProjectWriteLeaseCoordinator;
   private readonly registrations = new Map<string, WorkerRegistrationRecord>();
   private readonly runBindings = new Map<string, string>();
+  private readonly requestFlights = new Set<Promise<void>>();
   private server: Server | null = null;
   private bridgeUrl: string | null = null;
   private startFlight: Promise<string> | null = null;
+  private closing = false;
 
   constructor(leases = new PiProjectWriteLeaseCoordinator()) {
     this.leases = leases;
@@ -131,16 +133,20 @@ export class PiManagedExtensionHost {
   }
 
   async close(): Promise<void> {
-    for (const record of [...this.registrations.values()]) this.disposeRecord(record);
-    this.runBindings.clear();
     const server = this.server;
     this.server = null;
     this.bridgeUrl = null;
     this.startFlight = null;
-    if (!server) return;
-    await new Promise<void>((resolve, reject) => {
+    this.closing = true;
+    const closed = server ? new Promise<void>((resolve, reject) => {
       server.close((error) => error ? reject(error) : resolve());
-    });
+    }) : Promise.resolve();
+    for (const record of [...this.registrations.values()]) this.disposeRecord(record);
+    this.runBindings.clear();
+    await Promise.allSettled([...this.requestFlights]);
+    server?.closeIdleConnections();
+    await closed;
+    this.closing = false;
   }
 
   private start(): Promise<string> {
@@ -148,7 +154,10 @@ export class PiManagedExtensionHost {
     if (this.startFlight) return this.startFlight;
     this.startFlight = new Promise<string>((resolve, reject) => {
       const server = createServer((request, response) => {
-        void this.handle(request, response);
+        const flight = this.handle(request, response).finally(() => {
+          this.requestFlights.delete(flight);
+        });
+        this.requestFlights.add(flight);
       });
       server.once('error', reject);
       server.listen(0, '127.0.0.1', () => {
@@ -159,6 +168,7 @@ export class PiManagedExtensionHost {
           return;
         }
         this.server = server;
+        this.closing = false;
         this.bridgeUrl = `http://127.0.0.1:${address.port}/v1/worker`;
         resolve(this.bridgeUrl);
       });
@@ -226,6 +236,7 @@ export class PiManagedExtensionHost {
         );
         if (record.runId !== value.runId || this.registrations.get(token) !== record) {
           lease.release();
+          this.respond(response, 409, { error: 'Worker run identity is stale' });
           return;
         }
         record.leases.set(value.resourceId, lease);
@@ -266,7 +277,10 @@ export class PiManagedExtensionHost {
 
   private respond(response: ServerResponse, status: number, body: Record<string, unknown>): void {
     if (response.writableEnded) return;
-    response.writeHead(status, { 'content-type': 'application/json; charset=utf-8' });
+    response.writeHead(status, {
+      'content-type': 'application/json; charset=utf-8',
+      ...(this.closing ? { connection: 'close' } : {}),
+    });
     response.end(JSON.stringify(body));
   }