From 1f0f5d90cbb88e23a93fd80fe581b4735edfdba2 Mon Sep 17 00:00:00 2001
From: Lifei Zhou <lifei@squareup.com>
Date: Wed, 17 Jun 2026 13:39:25 +1000
Subject: [PATCH] fix: removed window goosed when closing window (#9818)

---
 ui/desktop/src/main.ts | 83 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 71 insertions(+), 12 deletions(-)

diff --git a/ui/desktop/src/main.ts b/ui/desktop/src/main.ts
index 562d6526c..6685f5421 100644
--- a/ui/desktop/src/main.ts
+++ b/ui/desktop/src/main.ts
@@ -802,6 +802,55 @@ const windowMap = new Map<number, BrowserWindow>();
 const goosedClients = new Map<number, Client>();
 const appWindows = new Map<string, BrowserWindow>();
 
+interface GoosedLease {
+  client: Client;
+  cleanup: () => Promise<void>;
+  windowIds: Set<number>;
+  cleanedUp: boolean;
+}
+
+const goosedLeasesByWindowId = new Map<number, GoosedLease>();
+
+const cleanupGoosedLease = async (lease: GoosedLease) => {
+  if (lease.cleanedUp) {
+    return;
+  }
+
+  lease.cleanedUp = true;
+  for (const windowId of lease.windowIds) {
+    goosedLeasesByWindowId.delete(windowId);
+    goosedClients.delete(windowId);
+  }
+  lease.windowIds.clear();
+
+  try {
+    await lease.cleanup();
+  } catch (error) {
+    log.error('Failed to cleanup goosed server:', error);
+  }
+};
+
+const attachWindowToGoosedLease = (windowId: number, lease: GoosedLease) => {
+  lease.windowIds.add(windowId);
+  goosedLeasesByWindowId.set(windowId, lease);
+  goosedClients.set(windowId, lease.client);
+};
+
+const releaseWindowGoosedLease = async (windowId: number) => {
+  const lease = goosedLeasesByWindowId.get(windowId);
+  goosedLeasesByWindowId.delete(windowId);
+  goosedClients.delete(windowId);
+
+  if (!lease) {
+    return;
+  }
+
+  lease.windowIds.delete(windowId);
+  if (lease.windowIds.size === 0) {
+    await cleanupGoosedLease(lease);
+  }
+};
+
 const windowPowerSaveBlockers = new Map<number, number>(); // windowId -> blockerId
 // Track pending initial messages per window
 const pendingInitialMessages = new Map<number, string>(); // windowId -> initialMessage
@@ -871,11 +920,6 @@ const createChat = async (app: App, options: CreateChatOptions = {}) => {
     pinnedCertFingerprint = goosedResult.certFingerprint;
   }
 
-  app.on('will-quit', async () => {
-    log.info('App quitting, terminating goosed server');
-    await goosedResult.cleanup();
-  });
-
   const {
     baseUrl,
     workingDir,
@@ -958,7 +1002,16 @@ const createChat = async (app: App, options: CreateChatOptions = {}) => {
       },
     })
   );
-  goosedClients.set(mainWindow.id, goosedClient);
+  const goosedLease: GoosedLease = {
+    client: goosedClient,
+    cleanup: goosedResult.cleanup,
+    windowIds: new Set<number>(),
+    cleanedUp: false,
+  };
+  attachWindowToGoosedLease(mainWindow.id, goosedLease);
+  mainWindow.once('closed', () => {
+    void releaseWindowGoosedLease(mainWindow.id);
+  });
 
   const serverReady = await checkServerStatus(goosedClient, errorLog, {
     onEvent: recordStartupEvent,
@@ -2710,9 +2763,9 @@ async function appMain() {
       }
 
       const launchingWindowId = launchingWindow.id;
-      const launchingClient = goosedClients.get(launchingWindowId);
-      if (!launchingClient) {
-        throw new Error('No client found for launching window');
+      const launchingLease = goosedLeasesByWindowId.get(launchingWindowId);
+      if (!launchingLease) {
+        throw new Error('No goosed lease found for launching window');
       }
 
       const appWindow = new BrowserWindow({
@@ -2730,11 +2783,11 @@ async function appMain() {
         },
       });
 
-      goosedClients.set(appWindow.id, launchingClient);
+      attachWindowToGoosedLease(appWindow.id, launchingLease);
       appWindows.set(gooseApp.name, appWindow);
 
-      appWindow.on('close', () => {
-        goosedClients.delete(appWindow.id);
+      appWindow.on('closed', () => {
+        void releaseWindowGoosedLease(appWindow.id);
         appWindows.delete(gooseApp.name);
       });
 
@@ -2840,6 +2893,12 @@ app.on('will-quit', async () => {
   // Stop the mesh child process if we spawned one.
   mesh.cleanup();
 
+  const goosedLeases = new Set(goosedLeasesByWindowId.values());
+  if (goosedLeases.size > 0) {
+    log.info(`App quitting, terminating ${goosedLeases.size} goosed server(s)`);
+    await Promise.all([...goosedLeases].map(cleanupGoosedLease));
+  }
+
   for (const [windowId, blockerId] of windowPowerSaveBlockers.entries()) {
     try {
       powerSaveBlocker.stop(blockerId);