feat: keep screen awake during inference and deactivate when idle

android/app/src/main/java/com/pocketpalai/KeepAwakeModule.kt

@@ -0,0 +1,28 @@
+package com.pocketpal
+
+import android.app.Activity
+import android.view.WindowManager
+import com.facebook.react.bridge.ReactApplicationContext
+import com.facebook.react.bridge.ReactContextBaseJavaModule
+import com.facebook.react.bridge.ReactMethod
+
+class KeepAwakeModule(private val reactContext: ReactApplicationContext) : ReactContextBaseJavaModule(reactContext) {
+
+    override fun getName(): String = "KeepAwakeModule"
+
+    @ReactMethod
+    fun activate() {
+        val activity = reactContext.currentActivity
+        activity?.runOnUiThread {
+            activity.window.addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON)
+        }
+    }
+
+    @ReactMethod
+    fun deactivate() {
+        val activity = reactContext.currentActivity
+        activity?.runOnUiThread {
+            activity.window.clearFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON)
+        }
+    }
+} 

android/app/src/main/java/com/pocketpalai/KeepAwakePackage.kt

@@ -0,0 +1,16 @@
+package com.pocketpal
+
+import com.facebook.react.ReactPackage
+import com.facebook.react.bridge.NativeModule
+import com.facebook.react.bridge.ReactApplicationContext
+import com.facebook.react.uimanager.ViewManager
+
+class KeepAwakePackage : ReactPackage {
+    override fun createNativeModules(reactContext: ReactApplicationContext): List<NativeModule> {
+        return listOf(KeepAwakeModule(reactContext))
+    }
+
+    override fun createViewManagers(reactContext: ReactApplicationContext): List<ViewManager<*, *>> {
+        return emptyList()
+    }
+} 

android/app/src/main/java/com/pocketpalai/MainApplication.kt

@@ -21,6 +21,7 @@ class MainApplication : Application(), ReactApplication {
               // Packages that cannot be autolinked yet can be added manually here, for example:
               // add(MyReactNativePackage())
               add(DeviceInfoPackage())
+              add(KeepAwakePackage())
             }
 
         override fun getJSMainModuleName(): String = "index"

ios/PocketPal/KeepAwakeModule.h

@@ -0,0 +1,4 @@
+#import <React/RCTBridgeModule.h>
+
+@interface KeepAwakeModule : NSObject <RCTBridgeModule>
+@end 

ios/PocketPal/KeepAwakeModule.m

@@ -0,0 +1,22 @@
+#import "KeepAwakeModule.h"
+#import <UIKit/UIKit.h>
+
+@implementation KeepAwakeModule
+
+RCT_EXPORT_MODULE(KeepAwakeModule);
+
+RCT_EXPORT_METHOD(activate)
+{
+    dispatch_async(dispatch_get_main_queue(), ^{
+        [[UIApplication sharedApplication] setIdleTimerDisabled:YES];
+    });
+}
+
+RCT_EXPORT_METHOD(deactivate)
+{
+    dispatch_async(dispatch_get_main_queue(), ^{
+        [[UIApplication sharedApplication] setIdleTimerDisabled:NO];
+    });
+}
+
+@end 

src/hooks/useChatSession.ts

@@ -1,4 +1,5 @@
 import React, {useRef, useCallback} from 'react';
+
 import {toJS} from 'mobx';
 import throttle from 'lodash.throttle';
 
@@ -8,6 +9,7 @@ import {chatSessionStore, modelStore} from '../store';
 
 import {MessageType, User} from '../utils/types';
 import {applyChatTemplate, convertToChatMessages} from '../utils/chat';
+import {activateKeepAwake, deactivateKeepAwake} from '../utils/keepAwake';
 
 export const useChatSession = (
   currentMessageInfo: React.MutableRefObject<{
@@ -86,6 +88,14 @@ export const useChatSession = (
     modelStore.setIsStreaming(false);
     chatSessionStore.setIsGenerating(true);
 
+    // Keep screen awake during completion
+    try {
+      activateKeepAwake();
+    } catch (error) {
+      console.error('Failed to activate keep awake during chat:', error);
+      // Continue with chat even if keep awake fails
+    }
+
     const id = randId();
     const createdAt = Date.now();
     currentMessageInfo.current = {createdAt, id};
@@ -166,6 +176,13 @@ export const useChatSession = (
       } else {
         addSystemMessage(`Completion failed: ${errorMessage}`);
       }
+    } finally {
+      // Always try to deactivate keep awake in finally block
+      try {
+        deactivateKeepAwake();
+      } catch (error) {
+        console.error('Failed to deactivate keep awake after chat:', error);
+      }
     }
   };
 
@@ -190,6 +207,15 @@ export const useChatSession = (
     }
     modelStore.setInferencing(false);
     modelStore.setIsStreaming(false);
+    // Deactivate keep awake when stopping completion
+    try {
+      deactivateKeepAwake();
+    } catch (error) {
+      console.error(
+        'Failed to deactivate keep awake after stopping chat:',
+        error,
+      );
+    }
   };
 
   return {

src/hooks/useKeepAwake.ts

@@ -0,0 +1,32 @@
+import {useEffect} from 'react';
+import {activateKeepAwake, deactivateKeepAwake} from '../utils/keepAwake';
+
+/**
+ * React hook that prevents the screen from going to sleep while the component is mounted.
+ *
+ * @example
+ * ```tsx
+ * function VideoPlayer() {
+ *   useKeepAwake(); // Screen will stay awake while VideoPlayer is mounted
+ *   return <Video source={source} />;
+ * }
+ * ```
+ */
+export function useKeepAwake(): void {
+  useEffect(() => {
+    try {
+      activateKeepAwake();
+      return () => {
+        try {
+          deactivateKeepAwake();
+        } catch (error) {
+          console.error('Failed to deactivate keep awake in cleanup:', error);
+        }
+      };
+    } catch (error) {
+      console.error('Failed to activate keep awake:', error);
+      // We don't rethrow here as it would crash the component
+      // Instead, we log the error and let the screen timeout normally
+    }
+  }, []);
+}

src/utils/keepAwake.ts

@@ -0,0 +1,38 @@
+import {NativeModules} from 'react-native';
+
+const {KeepAwakeModule} = NativeModules;
+
+if (!KeepAwakeModule) {
+  throw new Error(
+    'KeepAwakeModule is not available. Make sure:\n' +
+      '- You rebuilt the app after adding the native modules\n' +
+      '- The native module is properly linked\n' +
+      '- You are not using Expo managed workflow',
+  );
+}
+
+/**
+ * Activates keep awake functionality to prevent the screen from going to sleep
+ * @throws {Error} If the native module fails to activate
+ */
+export const activateKeepAwake = (): void => {
+  try {
+    KeepAwakeModule.activate();
+  } catch (error) {
+    console.error('Failed to activate keep awake:', error);
+    throw error;
+  }
+};
+
+/**
+ * Deactivates keep awake functionality allowing the screen to go to sleep
+ * @throws {Error} If the native module fails to deactivate
+ */
+export const deactivateKeepAwake = (): void => {
+  try {
+    KeepAwakeModule.deactivate();
+  } catch (error) {
+    console.error('Failed to deactivate keep awake:', error);
+    throw error;
+  }
+};