feat: keep screen awake during inference, deactivate when idle (#210)

* feat: keep screen awake during inference and deactivate when idle

* fix(test): waiting for checkbox test

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 

ios/Podfile.lock

@@ -2267,7 +2267,7 @@ SPEC CHECKSUMS:
   llama-rn: eb844b9cc4b240409b0411f16836416e567374f8
   nanopb: fad817b59e0457d11a5dfbde799381cd727c1275
   PromisesObjC: f5707f49cb48b9636751c5b2e7d227e43fba9f47
-  RCT-Folly: bf5c0376ffe4dd2cf438dcf86db385df9fdce648
+  RCT-Folly: 84578c8756030547307e4572ab1947de1685c599
   RCTDeprecation: 2c5e1000b04ab70b53956aa498bf7442c3c6e497
   RCTRequired: 5f785a001cf68a551c5f5040fb4c415672dbb481
   RCTTypeSafety: 6b98db8965005d32449605c0d005ecb4fee8a0f7

src/components/Checkbox/__tests__/Checkbox.test.tsx

@@ -1,6 +1,6 @@
 import React from 'react';
 import {Checkbox} from '../Checkbox';
-import {fireEvent, render, waitFor} from '../../../../jest/test-utils';
+import {fireEvent, render} from '../../../../jest/test-utils';
 
 describe('Checkbox', () => {
   it('renders correctly in unchecked state', () => {
@@ -18,10 +18,10 @@ describe('Checkbox', () => {
       <Checkbox checked={true} onPress={onPress} />,
     );
 
-    await waitFor(() => {
-      const checkIcon = findByTestId('check-icon');
-      expect(checkIcon).toBeDefined();
+    const checkIcon = await findByTestId('check-icon', {
+      includeHiddenElements: true,
     });
+    expect(checkIcon).toBeDefined();
   });
 
   it('calls onPress when clicked and not disabled', () => {

src/hooks/__tests__/useTheme.test.ts

@@ -1,4 +1,5 @@
 import {renderHook} from '@testing-library/react-native';
+import {act} from 'react-test-renderer';
 
 jest.unmock('../useTheme');
 jest.unmock('../../store');
@@ -23,10 +24,16 @@ describe('useTheme', () => {
     );
   });
 
-  it('should return dark theme when colorScheme is dark', () => {
+  it('should return dark theme when colorScheme is dark', async () => {
     uiStore.setColorScheme('dark');
 
     const {result} = renderHook(() => useTheme());
+
+    // Wait for the next update to ensure the theme change is applied
+    await act(async () => {
+      await Promise.resolve();
+    });
+
     expect(result.current).toEqual(
       expect.objectContaining({
         ...darkTheme,

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) {
+  console.warn(
+    '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;
+  }
+};