Improve power save blocker mechanism (#3698)

ui/desktop/src/hooks/useChatEngine.ts

@@ -49,6 +49,7 @@ export const useChatEngine = ({
   const [sessionOutputTokens, setSessionOutputTokens] = useState<number>(0);
   const [localInputTokens, setLocalInputTokens] = useState<number>(0);
   const [localOutputTokens, setLocalOutputTokens] = useState<number>(0);
+  const [powerSaveTimeoutId, setPowerSaveTimeoutId] = useState<number | null>(null);
 
   // Store message in global history when it's added (if enabled)
   const storeMessageInHistory = useCallback(
@@ -63,6 +64,20 @@ export const useChatEngine = ({
     [enableLocalStorage]
   );
 
+  const stopPowerSaveBlocker = useCallback(() => {
+    try {
+      window.electron.stopPowerSaveBlocker();
+    } catch (error) {
+      console.error('Failed to stop power save blocker:', error);
+    }
+
+    // Clear timeout if it exists
+    if (powerSaveTimeoutId) {
+      window.clearTimeout(powerSaveTimeoutId);
+      setPowerSaveTimeoutId(null);
+    }
+  }, [powerSaveTimeoutId]);
+
   const {
     messages,
     append: originalAppend,
@@ -84,7 +99,7 @@ export const useChatEngine = ({
     initialMessages: chat.messages,
     body: { session_id: chat.id, session_working_dir: window.appConfig.get('GOOSE_WORKING_DIR') },
     onFinish: async (_message, _reason) => {
-      window.electron.stopPowerSaveBlocker();
+      stopPowerSaveBlocker();
 
       const timeSinceLastInteraction = Date.now() - lastInteractionTime;
       window.electron.logInfo('last interaction:' + lastInteractionTime);
@@ -110,6 +125,8 @@ export const useChatEngine = ({
       onMessageStreamFinish?.();
     },
     onError: (error) => {
+      stopPowerSaveBlocker();
+
       console.log(
         'CHAT ENGINE RECEIVED ERROR FROM MESSAGE STREAM:',
         JSON.stringify(
@@ -206,40 +223,67 @@ export const useChatEngine = ({
     }
   }, [sessionMetadata]);
 
+  useEffect(() => {
+    return () => {
+      if (powerSaveTimeoutId) {
+        window.clearTimeout(powerSaveTimeoutId);
+      }
+      try {
+        window.electron.stopPowerSaveBlocker();
+      } catch (error) {
+        console.error('Failed to stop power save blocker during cleanup:', error);
+      }
+    };
+  }, [powerSaveTimeoutId]);
+
   // Handle submit
   const handleSubmit = useCallback(
     (combinedTextFromInput: string, onSummaryReset?: () => void) => {
       if (combinedTextFromInput.trim()) {
-        window.electron.startPowerSaveBlocker();
+        try {
+          window.electron.startPowerSaveBlocker();
+        } catch (error) {
+          console.error('Failed to start power save blocker:', error);
+        }
+
         setLastInteractionTime(Date.now());
 
+        // Set a timeout to automatically stop the power save blocker after 15 minutes
+        const timeoutId = window.setTimeout(
+          () => {
+            console.warn('Power save blocker timeout - stopping automatically after 15 minutes');
+            stopPowerSaveBlocker();
+          },
+          15 * 60 * 1000
+        );
+
+        setPowerSaveTimeoutId(timeoutId);
+
         const userMessage = createUserMessage(combinedTextFromInput.trim());
 
         if (onSummaryReset) {
           onSummaryReset();
-          setTimeout(() => {
+          window.setTimeout(() => {
             append(userMessage);
-            // Call onMessageSent after the message is sent
             onMessageSent?.();
           }, 150);
         } else {
           append(userMessage);
-          // Call onMessageSent after the message is sent
           onMessageSent?.();
         }
       } else {
         // If nothing was actually submitted (e.g. empty input and no images pasted)
-        window.electron.stopPowerSaveBlocker();
+        stopPowerSaveBlocker();
       }
     },
-    [append, onMessageSent]
+    [append, onMessageSent, stopPowerSaveBlocker]
   );
 
   // Handle stopping the message stream
   const onStopGoose = useCallback(() => {
     stop();
     setLastInteractionTime(Date.now());
-    window.electron.stopPowerSaveBlocker();
+    stopPowerSaveBlocker();
 
     // Handle stopping the message stream
     const lastMessage = messages[messages.length - 1];
@@ -330,7 +374,7 @@ export const useChatEngine = ({
         setMessages([...messages, responseMessage]);
       }
     }
-  }, [stop, messages, _setInput, setMessages, enableLocalStorage]);
+  }, [stop, messages, _setInput, setMessages, stopPowerSaveBlocker, enableLocalStorage]);
 
   const filteredMessages = useMemo(() => {
     return [...ancestorMessages, ...messages].filter((message) => message.display ?? true);

ui/desktop/src/main.ts

@@ -503,8 +503,8 @@ let appConfig = {
 let windowCounter = 0;
 const windowMap = new Map<number, BrowserWindow>();
 
-// Track power save blocker ID globally
-let powerSaveBlockerId: number | null = null;
+// Track power save blockers per window
+const windowPowerSaveBlockers = new Map<number, number>(); // windowId -> blockerId
 
 const createChat = async (
   app: App,
@@ -838,6 +838,23 @@ const createChat = async (
   // Handle window closure
   mainWindow.on('closed', () => {
     windowMap.delete(windowId);
+
+    if (windowPowerSaveBlockers.has(windowId)) {
+      const blockerId = windowPowerSaveBlockers.get(windowId)!;
+      try {
+        powerSaveBlocker.stop(blockerId);
+        console.log(
+          `[Main] Stopped power save blocker ${blockerId} for closing window ${windowId}`
+        );
+      } catch (error) {
+        console.error(
+          `[Main] Failed to stop power save blocker ${blockerId} for window ${windowId}:`,
+          error
+        );
+      }
+      windowPowerSaveBlockers.delete(windowId);
+    }
+
     if (goosedProcess && typeof goosedProcess === 'object' && 'kill' in goosedProcess) {
       goosedProcess.kill();
     }
@@ -1220,10 +1237,22 @@ ipcMain.handle('set-wakelock', async (_event, enable: boolean) => {
     settings.enableWakelock = enable;
     saveSettings(settings);
 
-    // Stop any existing power save blocker when disabling the setting
-    if (!enable && powerSaveBlockerId !== null) {
-      powerSaveBlocker.stop(powerSaveBlockerId);
-      powerSaveBlockerId = null;
+    // Stop all existing power save blockers when disabling the setting
+    if (!enable) {
+      for (const [windowId, blockerId] of windowPowerSaveBlockers.entries()) {
+        try {
+          powerSaveBlocker.stop(blockerId);
+          console.log(
+            `[Main] Stopped power save blocker ${blockerId} for window ${windowId} due to wakelock setting disabled`
+          );
+        } catch (error) {
+          console.error(
+            `[Main] Failed to stop power save blocker ${blockerId} for window ${windowId}:`,
+            error
+          );
+        }
+      }
+      windowPowerSaveBlockers.clear();
     }
 
     return true;
@@ -2026,21 +2055,36 @@ app.whenReady().then(async () => {
     }
   });
 
-  ipcMain.handle('start-power-save-blocker', () => {
-    if (powerSaveBlockerId === null) {
-      powerSaveBlockerId = powerSaveBlocker.start('prevent-app-suspension');
+  ipcMain.handle('start-power-save-blocker', (event) => {
+    const window = BrowserWindow.fromWebContents(event.sender);
+    const windowId = window?.id;
+
+    if (windowId && !windowPowerSaveBlockers.has(windowId)) {
+      const blockerId = powerSaveBlocker.start('prevent-app-suspension');
+      windowPowerSaveBlockers.set(windowId, blockerId);
+      console.log(`[Main] Started power save blocker ${blockerId} for window ${windowId}`);
       return true;
     }
 
+    if (windowId && windowPowerSaveBlockers.has(windowId)) {
+      console.log(`[Main] Power save blocker already active for window ${windowId}`);
+    }
+
     return false;
   });
 
-  ipcMain.handle('stop-power-save-blocker', () => {
-    if (powerSaveBlockerId !== null) {
-      powerSaveBlocker.stop(powerSaveBlockerId);
-      powerSaveBlockerId = null;
+  ipcMain.handle('stop-power-save-blocker', (event) => {
+    const window = BrowserWindow.fromWebContents(event.sender);
+    const windowId = window?.id;
+
+    if (windowId && windowPowerSaveBlockers.has(windowId)) {
+      const blockerId = windowPowerSaveBlockers.get(windowId)!;
+      powerSaveBlocker.stop(blockerId);
+      windowPowerSaveBlockers.delete(windowId);
+      console.log(`[Main] Stopped power save blocker ${blockerId} for window ${windowId}`);
       return true;
     }
+
     return false;
   });
 
@@ -2153,11 +2197,24 @@ async function getAllowList(): Promise<string[]> {
 }
 
 app.on('will-quit', async () => {
+  for (const [windowId, blockerId] of windowPowerSaveBlockers.entries()) {
+    try {
+      powerSaveBlocker.stop(blockerId);
+      console.log(
+        `[Main] Stopped power save blocker ${blockerId} for window ${windowId} during app quit`
+      );
+    } catch (error) {
+      console.error(
+        `[Main] Failed to stop power save blocker ${blockerId} for window ${windowId}:`,
+        error
+      );
+    }
+  }
+  windowPowerSaveBlockers.clear();
+
   // Unregister all shortcuts when quitting
   globalShortcut.unregisterAll();
 
-  // Clean up the temp directory on app quit
-  console.log('[Main] App "will-quit". Cleaning up temporary image directory...');
   try {
     await fs.access(gooseTempDir); // Check if directory exists to avoid error on fs.rm if it doesn't