From 7aa9f79bb2d0acb846547fcc4fd7fc2fda3a90d3 Mon Sep 17 00:00:00 2001
From: brother7 <7brother7@gmail.com>
Date: Tue, 25 Aug 2026 12:08:53 +0800
Subject: [PATCH] fix(pi): single-flight managed extension materialization

---
 electron/coding-runtime/pi/extension-host.ts | 11 ++++++++++-
 tests/unit/pi-extension-host.test.ts         | 20 ++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/electron/coding-runtime/pi/extension-host.ts b/electron/coding-runtime/pi/extension-host.ts
index 8eb1740..36767d8 100644
--- a/electron/coding-runtime/pi/extension-host.ts
+++ b/electron/coding-runtime/pi/extension-host.ts
@@ -139,6 +139,7 @@ export class PiManagedExtensionHost {
   private readonly registrations = new Map<string, WorkerRegistrationRecord>();
   private readonly runBindings = new Map<string, string>();
   private readonly requestFlights = new Set<Promise<void>>();
+  private readonly extensionMaterializations = new Map<string, Promise<string>>();
   private subagentBridge: PiExtensionSubagentBridge | undefined;
   private productTools: PiProductTools | undefined;
   private server: Server | null = null;
@@ -183,7 +184,15 @@ export class PiManagedExtensionHost {
     }
     const bridgeUrl = await this.start();
     await mkdir(input.extensionsDir, { recursive: true });
-    const extensionPath = await materializeMakelorePiExtension(input.extensionsDir);
+    let materialization = this.extensionMaterializations.get(input.extensionsDir);
+    if (!materialization) {
+      materialization = materializeMakelorePiExtension(input.extensionsDir).catch((error) => {
+        this.extensionMaterializations.delete(input.extensionsDir);
+        throw error;
+      });
+      this.extensionMaterializations.set(input.extensionsDir, materialization);
+    }
+    const extensionPath = await materialization;
     const role = input.role ?? 'parent';
     if (role === 'child' && !input.runId?.trim()) {
       throw new Error('Child extension registration requires a parent run id');
diff --git a/tests/unit/pi-extension-host.test.ts b/tests/unit/pi-extension-host.test.ts
index 63b0470..7c3a996 100644
--- a/tests/unit/pi-extension-host.test.ts
+++ b/tests/unit/pi-extension-host.test.ts
@@ -35,6 +35,26 @@ async function post(
 }
 
 describe('managed Pi extension bridge', () => {
+  it('single-flights the shared managed extension for concurrent worker registration', async () => {
+    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-concurrent-'));
+    roots.push(root);
+    const host = new PiManagedExtensionHost();
+    hosts.push(host);
+
+    const registrations = await Promise.all(Array.from({ length: 8 }, async (_, index) => (
+      await host.registerWorker({
+        conversationId: `conversation-${index + 1}`,
+        generation: 1,
+        projectId: `project-${index + 1}`,
+        extensionsDir: root,
+      })
+    )));
+
+    expect(new Set(registrations.map(({ extensionPath }) => extensionPath)).size).toBe(1);
+    expect(host.getDiagnostics().registrations).toEqual({ parent: 8, child: 0 });
+    await Promise.all(registrations.map(({ dispose }) => dispose()));
+  });
+
   it('finishes a queued HTTP request while closing a holder and waiter', async () => {
     const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-close-'));
     roots.push(root);