Prevent macOS sleep while streaming

Author: 1jehuang

src/platform.rs

@@ -1,5 +1,130 @@
 use std::path::Path;
 
+#[cfg(target_os = "macos")]
+mod macos_power {
+    use std::ffi::CString;
+    use std::os::raw::{c_char, c_void};
+
+    type CFStringRef = *const c_void;
+    type IOPMAssertionID = u32;
+    type IOReturn = i32;
+
+    const K_CF_STRING_ENCODING_UTF8: u32 = 0x0800_0100;
+    const K_IOPM_ASSERTION_LEVEL_ON: u32 = 255;
+    const K_IO_RETURN_SUCCESS: IOReturn = 0;
+
+    #[link(name = "CoreFoundation", kind = "framework")]
+    unsafe extern "C" {
+        fn CFStringCreateWithCString(
+            alloc: *const c_void,
+            c_str: *const c_char,
+            encoding: u32,
+        ) -> CFStringRef;
+        fn CFRelease(cf: *const c_void);
+    }
+
+    #[link(name = "IOKit", kind = "framework")]
+    unsafe extern "C" {
+        fn IOPMAssertionCreateWithName(
+            assertion_type: CFStringRef,
+            assertion_level: u32,
+            assertion_name: CFStringRef,
+            assertion_id: *mut IOPMAssertionID,
+        ) -> IOReturn;
+        fn IOPMAssertionRelease(assertion_id: IOPMAssertionID) -> IOReturn;
+    }
+
+    fn cf_string(value: &str) -> Option<CFStringRef> {
+        let c_string = CString::new(value).ok()?;
+        let cf = unsafe {
+            CFStringCreateWithCString(
+                std::ptr::null(),
+                c_string.as_ptr(),
+                K_CF_STRING_ENCODING_UTF8,
+            )
+        };
+        (!cf.is_null()).then_some(cf)
+    }
+
+    pub struct PowerAssertion {
+        id: Option<IOPMAssertionID>,
+    }
+
+    impl PowerAssertion {
+        pub fn prevent_user_idle_system_sleep(reason: &str) -> Self {
+            let Some(assertion_type) = cf_string("PreventUserIdleSystemSleep") else {
+                return Self { id: None };
+            };
+            let Some(assertion_name) = cf_string(reason) else {
+                unsafe { CFRelease(assertion_type) };
+                return Self { id: None };
+            };
+
+            let mut id = 0;
+            let result = unsafe {
+                IOPMAssertionCreateWithName(
+                    assertion_type,
+                    K_IOPM_ASSERTION_LEVEL_ON,
+                    assertion_name,
+                    &mut id,
+                )
+            };
+            unsafe {
+                CFRelease(assertion_type);
+                CFRelease(assertion_name);
+            }
+
+            if result == K_IO_RETURN_SUCCESS {
+                crate::logging::info(&format!(
+                    "Created macOS sleep-prevention assertion while streaming (id={id})"
+                ));
+                Self { id: Some(id) }
+            } else {
+                crate::logging::warn(&format!(
+                    "Failed to create macOS sleep-prevention assertion while streaming: IOReturn={result}"
+                ));
+                Self { id: None }
+            }
+        }
+
+        #[cfg(test)]
+        pub fn is_active(&self) -> bool {
+            self.id.is_some()
+        }
+    }
+
+    impl Drop for PowerAssertion {
+        fn drop(&mut self) {
+            if let Some(id) = self.id.take() {
+                let result = unsafe { IOPMAssertionRelease(id) };
+                if result != K_IO_RETURN_SUCCESS {
+                    crate::logging::warn(&format!(
+                        "Failed to release macOS sleep-prevention assertion id={id}: IOReturn={result}"
+                    ));
+                }
+            }
+        }
+    }
+}
+
+#[cfg(not(target_os = "macos"))]
+mod macos_power {
+    pub struct PowerAssertion;
+
+    impl PowerAssertion {
+        pub fn prevent_user_idle_system_sleep(_reason: &str) -> Self {
+            Self
+        }
+
+        #[cfg(test)]
+        pub fn is_active(&self) -> bool {
+            false
+        }
+    }
+}
+
+pub use macos_power::PowerAssertion;
+
 fn desired_nofile_soft_limit(current: u64, hard: u64, minimum: u64) -> Option<u64> {
     let desired = current.max(minimum).min(hard);
     (desired > current).then_some(desired)

src/session/active_pids.rs

@@ -48,13 +48,21 @@ pub fn unmark_streaming(session_id: &str) {
 /// never gets stuck showing a phantom streaming session.
 pub struct StreamingGuard {
     session_id: String,
+    #[allow(dead_code)]
+    sleep_assertion: crate::platform::PowerAssertion,
 }
 
 impl StreamingGuard {
     pub fn new(session_id: impl Into<String>) -> Self {
         let session_id = session_id.into();
         mark_streaming(&session_id);
-        Self { session_id }
+        let sleep_assertion = crate::platform::PowerAssertion::prevent_user_idle_system_sleep(
+            "Jcode streaming model response",
+        );
+        Self {
+            session_id,
+            sleep_assertion,
+        }
     }
 }
 
@@ -174,7 +182,10 @@ mod tests {
 
         let counts = session_counts();
         assert_eq!(counts.total, 3, "three live sessions expected");
-        assert_eq!(counts.streaming, 1, "only one live streaming session expected");
+        assert_eq!(
+            counts.streaming, 1,
+            "only one live streaming session expected"
+        );
 
         // Clearing the streaming marker drops the streaming count.
         unmark_streaming("session_alpha");
@@ -206,4 +217,45 @@ mod tests {
 
         crate::env::remove_var("JCODE_HOME");
     }
+
+    #[cfg(target_os = "macos")]
+    #[test]
+    fn streaming_guard_creates_visible_macos_sleep_assertion() {
+        let _guard = crate::storage::lock_test_env();
+        let temp = tempfile::tempdir().expect("tempdir");
+        crate::env::set_var("JCODE_HOME", temp.path());
+
+        let reason = "Jcode streaming model response";
+        register_active_pid("session_power", std::process::id());
+        {
+            let streaming = StreamingGuard::new("session_power");
+            assert!(
+                streaming.sleep_assertion.is_active(),
+                "macOS should create a native power assertion"
+            );
+
+            let output = std::process::Command::new("pmset")
+                .args(["-g", "assertions"])
+                .output()
+                .expect("pmset -g assertions should run on macOS");
+            assert!(output.status.success(), "pmset should succeed");
+            let stdout = String::from_utf8_lossy(&output.stdout);
+            assert!(
+                stdout.contains(reason),
+                "pmset output should show the streaming assertion; output was:\n{stdout}"
+            );
+        }
+
+        let output = std::process::Command::new("pmset")
+            .args(["-g", "assertions"])
+            .output()
+            .expect("pmset -g assertions should run on macOS");
+        let stdout = String::from_utf8_lossy(&output.stdout);
+        assert!(
+            !stdout.contains(reason),
+            "streaming assertion should be released after guard drop; output was:\n{stdout}"
+        );
+
+        crate::env::remove_var("JCODE_HOME");
+    }
 }