diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index ab5efbef8a6..ed832d77ef8 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -678,6 +678,7 @@ pub struct SettingsSnapshot { global_shortcut: String, codex_custom_sessions_dirs: Vec, agent_sessions_enabled: bool, + stay_awake_enabled: bool, agent_session_ssh_hosts: Vec, hooks_enabled: bool, http_proxy_enabled: bool, @@ -800,6 +801,7 @@ impl From for SettingsSnapshot { global_shortcut: settings.global_shortcut, codex_custom_sessions_dirs: settings.codex_custom_sessions_dirs, agent_sessions_enabled: settings.agent_sessions_enabled, + stay_awake_enabled: settings.stay_awake_enabled, agent_session_ssh_hosts: settings.agent_session_ssh_hosts, hooks_enabled: settings.hooks_enabled, http_proxy_enabled: settings.http_proxy_enabled, diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index df0e6621a9e..e5cb97d4b27 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -43,6 +43,7 @@ pub struct SettingsUpdate { pub global_shortcut: Option, pub codex_custom_sessions_dirs: Option>, pub agent_sessions_enabled: Option, + pub stay_awake_enabled: Option, pub agent_session_ssh_hosts: Option>, pub hooks_enabled: Option, pub http_proxy_enabled: Option, @@ -327,6 +328,9 @@ impl SettingsUpdate { if let Some(v) = self.agent_sessions_enabled { settings.agent_sessions_enabled = v; } + if let Some(v) = self.stay_awake_enabled { + settings.stay_awake_enabled = v; + } if let Some(v) = self.agent_session_ssh_hosts.clone() { settings.agent_session_ssh_hosts = codexbar::agent_sessions::RemoteSessionFetcher::sanitized_hosts(&v); @@ -509,6 +513,7 @@ pub async fn update_settings( let clear_local_usage_cache = patch.codex_custom_sessions_dirs.is_some(); let rebuild_tray_menu = patch.rebuilds_tray_menu(); let refresh_tray_presentation = patch.refreshes_tray_presentation(); + let stay_awake_changed = patch.stay_awake_enabled.is_some(); let tray_promotion_changed = patch.changes_tray_promotion(); let tray_panel_always_on_top_changed = patch.tray_panel_always_on_top.is_some(); let previous_promoted = settings.promote_tray_icon; @@ -536,6 +541,9 @@ pub async fn update_settings( } crate::floatbar::after_settings_saved(&app, &float_bar_patch, &settings, notify_float_bar); + if stay_awake_changed { + crate::stay_awake::settings_changed(&app, settings.stay_awake_enabled); + } if rebuild_tray_menu { crate::tray_bridge::rebuild_tray_menu(&app); } diff --git a/apps/desktop-tauri/src-tauri/src/events.rs b/apps/desktop-tauri/src-tauri/src/events.rs index b65e5eb6296..a9f9025371b 100644 --- a/apps/desktop-tauri/src-tauri/src/events.rs +++ b/apps/desktop-tauri/src-tauri/src/events.rs @@ -20,6 +20,7 @@ pub const REFRESH_COMPLETE: &str = "refresh-complete"; pub const UPDATE_STATE_CHANGED: &str = "update-state-changed"; pub const LOCALE_CHANGED: &str = "locale-changed"; pub const SETTINGS_CHANGED: &str = "settings-changed"; +pub const STAY_AWAKE_CHANGED: &str = "stay-awake-changed"; pub const CODEX_ACCOUNTS_UPDATED: &str = "codex-accounts-updated"; pub const LOGIN_PHASE: &str = "login-phase"; @@ -112,6 +113,10 @@ pub fn emit_settings_changed(app: &AppHandle) { let _emit_settings = app.emit(SETTINGS_CHANGED, ()); } +pub fn emit_stay_awake_changed(app: &AppHandle, held: bool) { + let _emit_stay_awake = app.emit(STAY_AWAKE_CHANGED, held); +} + pub fn emit_login_phase(app: &AppHandle, provider_id: &str, phase: &str, auth_link: Option<&str>) { let _emit_login_phase = app.emit( LOGIN_PHASE, diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index 02f48672b6a..39f4a88654a 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -14,6 +14,7 @@ mod proof_harness; mod shell; mod shortcut_bridge; mod state; +mod stay_awake; mod surface; mod surface_target; mod tray_accounts; @@ -159,6 +160,7 @@ fn main() { commands::get_bootstrap_state, commands::get_provider_catalog, commands::get_settings_snapshot, + stay_awake::get_stay_awake_status, commands::list_agent_sessions, commands::focus_agent_session, commands::update_settings, @@ -288,6 +290,7 @@ fn main() { shortcut_bridge::register(app.handle()); floatbar::install(app.handle()); auto_refresh::install(app.handle().clone()); + stay_awake::install(app.handle().clone()); if settings.powertoys_status_pipe_enabled { powertoys::install(app.handle().clone()); } @@ -415,8 +418,13 @@ fn main() { _ => {} } }) - .run(tauri::generate_context!()) - .expect("failed to run CodexBar desktop shell"); + .build(tauri::generate_context!()) + .expect("failed to build CodexBar desktop shell") + .run(|_app, event| { + if matches!(event, tauri::RunEvent::Exit) { + stay_awake::shutdown(); + } + }); } #[cfg(test)] diff --git a/apps/desktop-tauri/src-tauri/src/stay_awake.rs b/apps/desktop-tauri/src-tauri/src/stay_awake.rs new file mode 100644 index 00000000000..898dd7b4c84 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/stay_awake.rs @@ -0,0 +1,101 @@ +//! Stay Awake scheduler: keeps Windows from idle-sleeping while a local agent +//! (Codex, Claude, Pi/OMP) process is running. +//! +//! While `stay_awake_enabled` is on, a local-only process scan runs every 30 +//! seconds, independent of `agent_sessions_enabled`. It never touches SSH or +//! Tailscale hosts and never reads transcript files, so file-only rollouts and +//! remote sessions cannot hold the machine awake. The policy itself lives in +//! `codexbar::power_assertion`. The setting reaches the controller through +//! [`settings_changed`] (every save goes through `update_settings`) and once at +//! startup; the scan loop never re-reads the file, so it cannot re-enable a +//! toggle the user just turned off. + +use std::sync::{Mutex, MutexGuard, OnceLock}; +use std::time::Duration; + +use codexbar::agent_sessions::LocalAgentSessionScanner; +use codexbar::power_assertion::{StayAwake, SystemPowerAssertion}; +use codexbar::settings::Settings; +use tokio::sync::Notify; + +const SCAN_INTERVAL: Duration = Duration::from_secs(30); + +type Controller = StayAwake; + +static CONTROLLER: OnceLock> = OnceLock::new(); +static WAKE: Notify = Notify::const_new(); + +fn controller() -> MutexGuard<'static, Controller> { + CONTROLLER + .get_or_init(|| Mutex::new(StayAwake::new(SystemPowerAssertion::new()))) + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +/// Whether an idle-sleep prevention is currently held (tray status line). +pub fn is_held() -> bool { + controller().is_held() +} + +#[tauri::command] +pub fn get_stay_awake_status() -> bool { + is_held() +} + +fn publish_status_change(app: &tauri::AppHandle) { + crate::tray_bridge::rebuild_tray_menu(app); + crate::events::emit_stay_awake_changed(app, is_held()); +} + +/// Apply a saved settings change: turning Stay Awake off releases right away, +/// turning it on triggers a scan without waiting for the next interval. +pub fn settings_changed(app: &tauri::AppHandle, enabled: bool) { + let changed = controller().set_enabled(enabled); + if changed { + publish_status_change(app); + } + if enabled { + WAKE.notify_one(); + } +} + +/// Release for good on app exit. +pub fn shutdown() { + controller().shutdown(); +} + +pub fn install(app: tauri::AppHandle) { + controller().set_enabled(Settings::load().stay_awake_enabled); + tauri::async_runtime::spawn(async move { + loop { + scan_once(&app).await; + // Sleep one interval, or wake early when the setting is turned on. + tokio::time::timeout(SCAN_INTERVAL, WAKE.notified()) + .await + .ok(); + } + }); +} + +async fn scan_once(app: &tauri::AppHandle) { + if !controller().is_enabled() { + return; + } + let changed = + match tokio::task::spawn_blocking(LocalAgentSessionScanner::has_live_agent_process).await { + // The controller drops this result if the setting was turned off or + // the app shut down while the scan was running. + Ok(Ok(live)) => controller().observe(live), + Ok(Err(error)) => { + tracing::warn!("Stay Awake process scan failed: {error}"); + false + } + Err(error) => { + tracing::warn!("Stay Awake scan task failed: {error}"); + false + } + }; + if changed { + publish_status_change(app); + } +} diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index aebb7c21153..7ad4e455fa9 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -20,7 +20,7 @@ use crate::surface::SurfaceMode; use crate::surface_target::SurfaceTarget; #[cfg(test)] use crate::tray_menu::build_tray_menu; -use crate::tray_menu::{TrayMenuEntry, build_tray_menu_with}; +use crate::tray_menu::{MenuState, TrayMenuEntry, build_tray_menu_with}; #[derive(Debug, Clone, Copy)] struct MonitorScaleInfo { @@ -134,7 +134,10 @@ fn build_native_tray_menu( providers, status_labels, &enabled, - settings.float_bar_enabled, + MenuState { + float_bar_enabled: settings.float_bar_enabled, + stay_awake_held: crate::stay_awake::is_held(), + }, settings.ui_language, ); crate::tray_accounts::prepend_account_menus(&mut spec, &settings); diff --git a/apps/desktop-tauri/src-tauri/src/tray_menu.rs b/apps/desktop-tauri/src-tauri/src/tray_menu.rs index 688e56b42b8..9f60c843b51 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_menu.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_menu.rs @@ -82,6 +82,14 @@ impl TrayMenuEntry { } } +/// Live app state the menu mirrors beyond the provider catalog. +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct MenuState { + pub(crate) float_bar_enabled: bool, + /// Stay Awake currently holds an idle-sleep prevention. + pub(crate) stay_awake_held: bool, +} + #[cfg(test)] pub(crate) fn build_tray_menu( providers: &[ProviderCatalogEntry], @@ -92,7 +100,7 @@ pub(crate) fn build_tray_menu( providers, status_labels, enabled_providers, - false, + MenuState::default(), Language::English, ) } @@ -101,7 +109,7 @@ pub(crate) fn build_tray_menu_with( providers: &[ProviderCatalogEntry], status_labels: &[(String, String)], enabled_providers: &HashSet, - float_bar_enabled: bool, + state: MenuState, lang: Language, ) -> Vec { let mut menu: Vec = Vec::new(); @@ -111,7 +119,13 @@ pub(crate) fn build_tray_menu_with( for (id, label) in status_labels { menu.push(TrayMenuEntry::status_row(format!("status_{id}"), label)); } - if !status_labels.is_empty() { + if state.stay_awake_held { + menu.push(TrayMenuEntry::status_row( + "stay_awake_active", + text(LocaleKey::TrayStayAwakeActive), + )); + } + if !status_labels.is_empty() || state.stay_awake_held { menu.push(TrayMenuEntry::separator()); } @@ -130,7 +144,7 @@ pub(crate) fn build_tray_menu_with( menu.push(TrayMenuEntry::check_item( "toggle_float_bar", text(LocaleKey::TrayShowFloatBar), - float_bar_enabled, + state.float_bar_enabled, )); menu.push(TrayMenuEntry::separator()); @@ -239,7 +253,10 @@ mod tests { &sample_provider_catalog(), &[], &both_enabled(), - /* float_bar_enabled = */ true, + MenuState { + float_bar_enabled: true, + ..MenuState::default() + }, Language::English, ); let toggle = menu_on @@ -253,7 +270,7 @@ mod tests { &sample_provider_catalog(), &[], &both_enabled(), - /* float_bar_enabled = */ false, + MenuState::default(), Language::English, ); let toggle = menu_off @@ -269,7 +286,7 @@ mod tests { &sample_provider_catalog(), &[], &both_enabled(), - false, + MenuState::default(), Language::Japanese, ); fn label_for<'a>(menu: &'a [TrayMenuEntry], id: &'a str) -> &'a str { @@ -311,4 +328,27 @@ mod tests { // Third item should be a separator. assert!(menu[2].is_separator); } + + #[test] + fn stay_awake_row_appears_only_while_held() { + let build = |stay_awake_held| { + build_tray_menu_with( + &sample_provider_catalog(), + &[], + &both_enabled(), + MenuState { + stay_awake_held, + ..MenuState::default() + }, + Language::English, + ) + }; + assert!(!menu_contains(&build(false), "stay_awake_active")); + + let held = build(true); + assert_eq!(held[0].id.as_deref(), Some("stay_awake_active")); + assert_eq!(held[0].label, "Stay Awake: local agent session is live"); + assert!(held[0].disabled); + assert!(held[1].is_separator); + } } diff --git a/apps/desktop-tauri/src/hooks/useStayAwakeStatus.ts b/apps/desktop-tauri/src/hooks/useStayAwakeStatus.ts new file mode 100644 index 00000000000..c7d67e58c29 --- /dev/null +++ b/apps/desktop-tauri/src/hooks/useStayAwakeStatus.ts @@ -0,0 +1,46 @@ +import { useEffect, useState } from "react"; +import { listen } from "@tauri-apps/api/event"; +import { getStayAwakeStatus } from "../lib/tauri"; + +const STAY_AWAKE_CHANGED = "stay-awake-changed"; + +export function useStayAwakeStatus(): boolean { + const [held, setHeld] = useState(false); + + useEffect(() => { + let cancelled = false; + let updates = 0; + let unlisten = () => {}; + + const initialize = async () => { + try { + unlisten = await listen(STAY_AWAKE_CHANGED, (event) => { + updates += 1; + if (!cancelled) setHeld(event.payload); + }); + } catch { + // The status query below still provides a useful initial value. + } + if (cancelled) { + unlisten(); + return; + } + + const observedUpdates = updates; + try { + const current = await getStayAwakeStatus(); + if (!cancelled && updates === observedUpdates) setHeld(current); + } catch { + // Keep the last event value if the status query is unavailable. + } + }; + + void initialize(); + return () => { + cancelled = true; + unlisten(); + }; + }, []); + + return held; +} diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 195620e3bff..6eb0f17f1e7 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -226,6 +226,7 @@ export const ALL_LOCALE_KEYS = [ "TrayPopOutDashboard", "TrayShowWindow", "TrayShowFloatBar", + "TrayStayAwakeActive", "TrayRefreshAll", "TrayProviders", "TraySettings", @@ -443,6 +444,8 @@ export const ALL_LOCALE_KEYS = [ "AgentSessionsTitle", "AgentSessionsEnableLabel", "AgentSessionsEnableHelper", + "AgentSessionsStayAwakeLabel", + "AgentSessionsStayAwakeHelper", "AgentSessionsSshHostsLabel", "AgentSessionsSshHostsHelper", "AgentSessionsLoading", diff --git a/apps/desktop-tauri/src/lib/tauri.ts b/apps/desktop-tauri/src/lib/tauri.ts index 6ccd63c0de6..faf370ea7f7 100644 --- a/apps/desktop-tauri/src/lib/tauri.ts +++ b/apps/desktop-tauri/src/lib/tauri.ts @@ -87,6 +87,10 @@ export function getSettingsSnapshot(): Promise { return invoke("get_settings_snapshot"); } +export function getStayAwakeStatus(): Promise { + return invoke("get_stay_awake_status"); +} + export function updateSettings( patch: SettingsUpdate, ): Promise { diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index 8e156a34bf1..428b0d4c60f 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -6,6 +6,7 @@ const tauriMocks = vi.hoisted(() => ({ refreshProviders: vi.fn(), refreshProvidersIfStale: vi.fn(), getSettingsSnapshot: vi.fn(), + getStayAwakeStatus: vi.fn().mockResolvedValue(false), updateSettings: vi.fn(), getUpdateState: vi.fn(), checkForUpdates: vi.fn(), diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index ccc2e8fe2f4..dc2265cf473 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -15,6 +15,7 @@ import { TRAY_SCALE_STEP, useTrayPanelController, } from "../hooks/useTrayPanelController"; +import { useStayAwakeStatus } from "../hooks/useStayAwakeStatus"; import MenuCard from "../components/MenuCard"; import MenuSurface, { MenuEmpty } from "../components/MenuSurface"; import UpdateBanner from "../components/UpdateBanner"; @@ -52,6 +53,7 @@ const HAS_STATUS_PAGE = new Set([ * 2. Detail: click a provider in grid → show only that provider's card */ export default function TrayPanel({ state }: { state: BootstrapState }) { + const stayAwakeHeld = useStayAwakeStatus(); const { t, settings, @@ -154,6 +156,11 @@ export default function TrayPanel({ state }: { state: BootstrapState }) { : null; const canSwitchClaudeAccount = selectedProvider !== null && hasSuccessfulClaudeCliQuota(selectedProvider); + const stayAwakeStatus = stayAwakeHeld ? ( +

+ {t("TrayStayAwakeActive")} +

+ ) : null; if (sorted.length === 0) { return ( @@ -168,6 +175,7 @@ export default function TrayPanel({ state }: { state: BootstrapState }) { footerRows={footerRows} style={{ zoom: trayScale }} > + {stayAwakeStatus} {settings.agentSessionsEnabled && } + {stayAwakeStatus} {settings.agentSessionsEnabled && } { ).toHaveLength(0); expect(screen.getAllByText("HooksEnableLabel")).toHaveLength(1); }); + + it("toggles Stay Awake independently of agent session discovery", () => { + const set = vi.fn(); + render( + , + ); + + const toggle = screen + .getByText("AgentSessionsStayAwakeLabel") + .closest(".settings-field") + ?.querySelector('input[type="checkbox"]'); + expect(toggle).not.toBeNull(); + expect(toggle).not.toBeDisabled(); + expect(toggle).not.toBeChecked(); + + fireEvent.click(toggle as Element); + expect(set).toHaveBeenCalledWith({ stayAwakeEnabled: true }); + }); }); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx index fcf98d09c30..ee278cc3716 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx @@ -192,6 +192,17 @@ export default function AdvancedTab({ settings, set, saving }: TabProps) { onChange={(v) => set({ agentSessionsEnabled: v })} /> + + set({ stayAwakeEnabled: v })} + /> + Result { + Ok(Self::holds_system_awake(&snapshot_processes()?)) + } + + /// Pure predicate behind [`Self::has_live_agent_process`]. + pub(super) fn holds_system_awake(processes: &[AgentProcessRecord]) -> bool { + processes + .iter() + .any(|process| process.pid > 0 && process.is_agent()) + } +} + +#[cfg(windows)] +fn snapshot_processes() -> Result, String> { + use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle}; + use windows::Win32::Foundation::{ERROR_NO_MORE_FILES, HANDLE}; + use windows::Win32::System::Diagnostics::ToolHelp::{ + CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW, + TH32CS_SNAPPROCESS, + }; + + // SAFETY: a successful call returns a unique snapshot handle owned below. + let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) } + .map_err(|error| format!("process snapshot failed: {error}"))?; + // SAFETY: `snapshot` is a valid, unique handle returned just above; the + // `OwnedHandle` closes it when dropped. + let snapshot = unsafe { OwnedHandle::from_raw_handle(snapshot.0) }; + let handle = HANDLE(snapshot.as_raw_handle()); + + let mut entry = PROCESSENTRY32W { + dwSize: u32::try_from(std::mem::size_of::()) + .map_err(|error| format!("invalid process entry size: {error}"))?, + ..Default::default() + }; + let mut records = Vec::new(); + // SAFETY: the snapshot handle is live and `entry.dwSize` is correct. + match unsafe { Process32FirstW(handle, &mut entry) } { + Ok(()) => {} + Err(error) if error.code() == ERROR_NO_MORE_FILES.to_hresult() => return Ok(records), + Err(error) => return Err(format!("process enumeration failed: {error}")), + } + loop { + let len = entry + .szExeFile + .iter() + .position(|&unit| unit == 0) + .unwrap_or(entry.szExeFile.len()); + let name = String::from_utf16_lossy(&entry.szExeFile[..len]); + records.push(WindowsProcessOutputParser::record( + entry.th32ProcessID, + entry.th32ParentProcessID, + None, + Some(name), + None, + )); + // SAFETY: `entry` is valid for the next snapshot entry. + match unsafe { Process32NextW(handle, &mut entry) } { + Ok(()) => {} + Err(error) if error.code() == ERROR_NO_MORE_FILES.to_hresult() => break, + Err(error) => return Err(format!("process enumeration failed: {error}")), + } + } + Ok(records) +} + +#[cfg(not(windows))] +fn snapshot_processes() -> Result, String> { + Err("native process snapshots are only available on Windows".to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn record(pid: u32, name: &str) -> AgentProcessRecord { + WindowsProcessOutputParser::record(pid, 1, None, Some(name.to_string()), None) + } + + #[test] + fn agent_processes_hold_the_system_awake() { + for name in ["claude.exe", "codex.exe"] { + let processes = [record(10, "explorer.exe"), record(11, name)]; + assert!( + LocalAgentSessionScanner::holds_system_awake(&processes), + "{name} should count" + ); + } + } + + #[test] + fn helpers_and_unrelated_processes_do_not_hold_the_system_awake() { + let processes = [ + record(10, "explorer.exe"), + record(11, "Codex (Renderer)"), + record(12, "codex app-server"), + record(13, "node.exe"), + ]; + assert!(!LocalAgentSessionScanner::holds_system_awake(&processes)); + assert!(!LocalAgentSessionScanner::holds_system_awake(&[])); + } + + #[test] + fn pid_zero_never_counts() { + let processes = [record(0, "claude.exe")]; + assert!(!LocalAgentSessionScanner::holds_system_awake(&processes)); + } + + #[cfg(windows)] + #[test] + fn native_snapshot_lists_the_current_process() { + let processes = snapshot_processes().expect("snapshot"); + assert!(processes.iter().any(|p| p.pid == std::process::id())); + } +} diff --git a/rust/src/agent_sessions/parsers.rs b/rust/src/agent_sessions/parsers.rs index 49e236e8e6e..b29b4ea598f 100644 --- a/rust/src/agent_sessions/parsers.rs +++ b/rust/src/agent_sessions/parsers.rs @@ -139,32 +139,49 @@ impl WindowsProcessOutputParser { .filter_map(|value| serde_json::from_value::(value).ok()) .filter(|process| process.process_id > 0 && seen.insert(process.process_id)) .map(|process| { - let display_name = process - .name - .as_deref() - .filter(|name| !name.trim().is_empty()) - .or(process.executable_path.as_deref()) - .unwrap_or_default(); - let classification = classify_process_command(display_name); - AgentProcessRecord { - pid: process.process_id, - ppid: process.parent_process_id, - started_at: process + Self::record( + process.process_id, + process.parent_process_id, + process .creation_date .as_deref() .and_then(parse_windows_creation_date), - provider: classification.provider, - source: classification.source, - executable: process.name.unwrap_or(classification.executable), - kind: classification.kind, - // Windows deliberately does not harvest other processes' - // command lines (privacy); ps-based remote scans carry - // them, and tests inject them. - command: None, - } + process.name, + process.executable_path.as_deref(), + ) }) .collect()) } + + /// Classify one Windows process. The image name is preferred over the + /// executable path, matching the CIM query and the native snapshot. + pub(super) fn record( + pid: u32, + ppid: u32, + started_at: Option>, + name: Option, + executable_path: Option<&str>, + ) -> AgentProcessRecord { + let display_name = name + .as_deref() + .filter(|name| !name.trim().is_empty()) + .or(executable_path) + .unwrap_or_default(); + let classification = classify_process_command(display_name); + AgentProcessRecord { + pid, + ppid, + started_at, + provider: classification.provider, + source: classification.source, + executable: name.unwrap_or(classification.executable), + kind: classification.kind, + // Windows deliberately does not harvest other processes' + // command lines (privacy); ps-based remote scans carry + // them, and tests inject them. + command: None, + } + } } fn parse_windows_creation_date(value: &str) -> Option> { diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 1dfa17cff5d..b528d1fef72 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -19,6 +19,7 @@ pub mod login; #[cfg(windows)] pub mod managed_process; pub mod notifications; +pub mod power_assertion; pub mod providers; pub mod secure_file; pub mod settings; diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e6ec918ac42..658ecad57ab 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -454,6 +454,7 @@ locale_keys! { TrayPopOutDashboard, TrayShowWindow, TrayShowFloatBar, + TrayStayAwakeActive, TrayRefreshAll, TrayProviders, TraySettings, @@ -687,6 +688,8 @@ locale_keys! { AgentSessionsTitle, AgentSessionsEnableLabel, AgentSessionsEnableHelper, + AgentSessionsStayAwakeLabel, + AgentSessionsStayAwakeHelper, AgentSessionsSshHostsLabel, AgentSessionsSshHostsHelper, AgentSessionsLoading, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8b..977a040d7c9 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -211,6 +211,7 @@ TrayOpenCodexBar = Pop Out Dashboard TrayPopOutDashboard = Pop Out Dashboard TrayShowWindow = Show Window TrayShowFloatBar = Show Float Bar +TrayStayAwakeActive = Stay Awake: local agent session is live TrayRefreshAll = Refresh All TrayProviders = Providers TraySettings = Settings... @@ -380,6 +381,8 @@ AgentSessionsProviderPi = Pi AgentSessionsProviderOmp = OMP AgentSessionsEnableLabel = Show active Codex, Claude, and Pi/OMP sessions AgentSessionsEnableHelper = Discover local sessions and optional SSH hosts. Disabled by default. +AgentSessionsStayAwakeLabel = Stay awake while a local agent session is live +AgentSessionsStayAwakeHelper = Prevents idle sleep while a local Codex, Claude, or Pi/OMP process is running. The display can still turn off, and closing the lid still sleeps the PC. Disabled by default. AgentSessionsSshHostsLabel = SSH hosts AgentSessionsSshHostsHelper = Comma-separated SSH targets. Connections use non-interactive mode and short timeouts. AgentSessionsLoading = Looking for agent sessions... diff --git a/rust/src/power_assertion.rs b/rust/src/power_assertion.rs new file mode 100644 index 00000000000..2c84f008e22 --- /dev/null +++ b/rust/src/power_assertion.rs @@ -0,0 +1,391 @@ +//! Stay Awake: hold one idle-sleep prevention while a local agent session is live. +//! +//! Port of upstream CodexBar 0.67.0 `AgentSessionPowerAssertion` and the +//! `updatePowerAssertion` policy in `AgentSessionsStore`. [`PowerAssertion`] is +//! the platform seam (injectable for tests); [`StayAwake`] is the pure policy +//! that decides when to acquire and release it. +//! +//! Only the system is kept awake, never the display: no display-required flag +//! is set, and lid close or an explicit sleep still sleeps the machine. + +/// Platform seam for one system idle-sleep prevention. +pub trait PowerAssertion: Send { + /// Start preventing idle system sleep. Returns `false` when the OS refused. + fn acquire(&mut self) -> bool; + /// Stop preventing idle system sleep. Safe to call when nothing is held. + fn release(&mut self); +} + +/// Stay Awake policy: at most one assertion, held only while enabled and a +/// live local agent session exists. +/// +/// Every method returns whether the held state changed so callers can refresh +/// UI that mirrors it. Scan results arrive asynchronously, so [`observe`] is +/// ignored once the feature is disabled or shut down: a stale scan completion +/// can never acquire after the toggle was turned off. +/// +/// [`observe`]: StayAwake::observe +pub struct StayAwake { + assertion: A, + enabled: bool, + shut_down: bool, + held: bool, +} + +impl StayAwake { + pub fn new(assertion: A) -> Self { + Self { + assertion, + enabled: false, + shut_down: false, + held: false, + } + } + + pub fn is_held(&self) -> bool { + self.held + } + + pub fn is_enabled(&self) -> bool { + self.enabled + } + + /// Apply the user setting. Turning it off releases immediately; turning it + /// on acquires nothing until a scan reports a live session. + pub fn set_enabled(&mut self, enabled: bool) -> bool { + if self.shut_down { + return false; + } + self.enabled = enabled; + if enabled { false } else { self.release() } + } + + /// Record one scan outcome. A failed acquire leaves nothing held, so the + /// next scan that still sees a live session retries. + pub fn observe(&mut self, live_local_session: bool) -> bool { + if !self.enabled || self.shut_down { + return false; + } + match (live_local_session, self.held) { + (true, false) => { + self.held = self.assertion.acquire(); + self.held + } + (false, true) => self.release(), + _ => false, + } + } + + /// Release for good on app exit; later scans and toggles are ignored. + pub fn shutdown(&mut self) -> bool { + self.shut_down = true; + self.enabled = false; + self.release() + } + + fn release(&mut self) -> bool { + if !self.held { + return false; + } + self.assertion.release(); + self.held = false; + true + } +} + +#[cfg(windows)] +pub use windows_impl::SystemPowerAssertion; + +/// Holds `ES_SYSTEM_REQUIRED` for the process. +/// +/// `SetThreadExecutionState(ES_CONTINUOUS | ...)` is scoped to the calling +/// thread and lapses when that thread ends, so a dedicated long-lived thread +/// owns it. Calling it from a tokio worker would drop the request whenever the +/// runtime retires or reuses that thread. +#[cfg(windows)] +mod windows_impl { + use super::PowerAssertion; + use std::sync::mpsc::{self, Sender}; + use std::thread::JoinHandle; + use std::time::Duration; + use windows::Win32::System::Power::{ + ES_CONTINUOUS, ES_SYSTEM_REQUIRED, EXECUTION_STATE, SetThreadExecutionState, + }; + + const ACQUIRE_TIMEOUT: Duration = Duration::from_secs(2); + + enum Request { + Hold(Sender), + Clear, + } + + struct Worker { + requests: Sender, + thread: JoinHandle<()>, + } + + #[derive(Default)] + pub struct SystemPowerAssertion { + worker: Option, + } + + impl SystemPowerAssertion { + pub fn new() -> Self { + Self::default() + } + + fn worker(&mut self) -> Option<&Worker> { + if self.worker.is_none() { + let (requests, inbox) = mpsc::channel(); + let thread = std::thread::Builder::new() + .name("codexbar-stay-awake".to_string()) + .spawn(move || run(inbox)) + .map_err(|error| tracing::warn!("Stay Awake thread failed to start: {error}")) + .ok()?; + self.worker = Some(Worker { requests, thread }); + } + self.worker.as_ref() + } + } + + impl PowerAssertion for SystemPowerAssertion { + fn acquire(&mut self) -> bool { + let (reply, answer) = mpsc::channel(); + let Some(requests) = self.worker().map(|worker| worker.requests.clone()) else { + return false; + }; + if requests.send(Request::Hold(reply)).is_err() { + // The thread is gone; the next attempt starts a fresh one. + self.worker = None; + return false; + } + match answer.recv_timeout(ACQUIRE_TIMEOUT) { + Ok(acquired) => acquired, + Err(mpsc::RecvTimeoutError::Timeout) => { + // A timed-out reply does not cancel the queued Win32 call. + // Clear after it so a late success cannot strand the hold. + if requests.send(Request::Clear).is_err() { + self.worker = None; + } + false + } + Err(mpsc::RecvTimeoutError::Disconnected) => false, + } + } + + fn release(&mut self) { + if let Some(worker) = &self.worker + && worker.requests.send(Request::Clear).is_err() + { + self.worker = None; + } + } + } + + impl Drop for SystemPowerAssertion { + fn drop(&mut self) { + if let Some(Worker { requests, thread }) = self.worker.take() { + drop(requests); + thread.join().ok(); + } + } + } + + fn run(inbox: mpsc::Receiver) { + while let Ok(request) = inbox.recv() { + match request { + Request::Hold(reply) => { + reply + .send(set_state(ES_CONTINUOUS | ES_SYSTEM_REQUIRED)) + .ok(); + } + Request::Clear => { + set_state(ES_CONTINUOUS); + } + } + } + // Channel closed: clear explicitly instead of relying on thread exit. + set_state(ES_CONTINUOUS); + } + + /// `SetThreadExecutionState` returns the previous state, or 0 on failure. + fn set_state(flags: EXECUTION_STATE) -> bool { + // SAFETY: plain Win32 call with valid flags; it only touches the + // calling thread's execution state. + unsafe { SetThreadExecutionState(flags) }.0 != 0 + } +} + +#[cfg(not(windows))] +pub use fallback::SystemPowerAssertion; + +/// Other platforms have no implementation; nothing is ever held. +#[cfg(not(windows))] +mod fallback { + use super::PowerAssertion; + + #[derive(Default)] + pub struct SystemPowerAssertion; + + impl SystemPowerAssertion { + pub fn new() -> Self { + Self + } + } + + impl PowerAssertion for SystemPowerAssertion { + fn acquire(&mut self) -> bool { + false + } + + fn release(&mut self) {} + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Arc, Mutex}; + + #[derive(Default)] + struct Counts { + acquires: u32, + releases: u32, + fail_next_acquires: u32, + } + + #[derive(Clone, Default)] + struct FakeAssertion(Arc>); + + impl FakeAssertion { + fn counts(&self) -> (u32, u32) { + let counts = self.0.lock().unwrap(); + (counts.acquires, counts.releases) + } + + fn fail_next(&self, times: u32) { + self.0.lock().unwrap().fail_next_acquires = times; + } + } + + impl PowerAssertion for FakeAssertion { + fn acquire(&mut self) -> bool { + let mut counts = self.0.lock().unwrap(); + counts.acquires += 1; + if counts.fail_next_acquires > 0 { + counts.fail_next_acquires -= 1; + return false; + } + true + } + + fn release(&mut self) { + self.0.lock().unwrap().releases += 1; + } + } + + fn enabled() -> (StayAwake, FakeAssertion) { + let fake = FakeAssertion::default(); + let mut stay_awake = StayAwake::new(fake.clone()); + stay_awake.set_enabled(true); + (stay_awake, fake) + } + + #[test] + fn disabled_never_acquires() { + let fake = FakeAssertion::default(); + let mut stay_awake = StayAwake::new(fake.clone()); + assert!(!stay_awake.observe(true)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (0, 0)); + } + + #[test] + fn enabling_alone_acquires_nothing() { + let (stay_awake, fake) = enabled(); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (0, 0)); + } + + #[test] + fn live_session_holds_exactly_one_assertion_until_it_disappears() { + let (mut stay_awake, fake) = enabled(); + assert!(stay_awake.observe(true), "appearing session acquires"); + assert!(!stay_awake.observe(true), "idle session keeps the hold"); + assert!(!stay_awake.observe(true)); + assert!(stay_awake.is_held()); + assert_eq!(fake.counts(), (1, 0)); + + assert!(stay_awake.observe(false), "disappearing session releases"); + assert!(!stay_awake.observe(false)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (1, 1)); + + assert!(stay_awake.observe(true), "a later session holds again"); + assert_eq!(fake.counts(), (2, 1)); + } + + #[test] + fn toggling_off_releases_immediately() { + let (mut stay_awake, fake) = enabled(); + stay_awake.observe(true); + assert!(stay_awake.set_enabled(false)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (1, 1)); + assert!(!stay_awake.set_enabled(false), "already released"); + assert_eq!(fake.counts(), (1, 1)); + } + + #[test] + fn stale_scan_after_disable_does_not_acquire() { + let (mut stay_awake, fake) = enabled(); + stay_awake.set_enabled(false); + assert!(!stay_awake.observe(true)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (0, 0)); + } + + #[test] + fn shutdown_releases_and_blocks_later_acquires() { + let (mut stay_awake, fake) = enabled(); + stay_awake.observe(true); + assert!(stay_awake.shutdown()); + assert_eq!(fake.counts(), (1, 1)); + + assert!(!stay_awake.observe(true)); + stay_awake.set_enabled(true); + assert!(!stay_awake.observe(true)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (1, 1)); + } + + #[test] + fn failed_acquire_retries_on_the_next_scan() { + let (mut stay_awake, fake) = enabled(); + fake.fail_next(1); + assert!(!stay_awake.observe(true), "refused acquire is not held"); + assert!(!stay_awake.is_held()); + assert!(stay_awake.observe(true), "next scan retries"); + assert!(stay_awake.is_held()); + assert_eq!(fake.counts(), (2, 0)); + } + + #[test] + fn failed_acquire_then_no_session_never_releases() { + let (mut stay_awake, fake) = enabled(); + fake.fail_next(1); + stay_awake.observe(true); + assert!(!stay_awake.observe(false)); + assert_eq!(fake.counts(), (1, 0)); + } + + #[cfg(windows)] + #[test] + fn system_assertion_acquires_and_releases_on_its_own_thread() { + let mut assertion = SystemPowerAssertion::new(); + assert!(assertion.acquire()); + assertion.release(); + assert!(assertion.acquire(), "can be re-acquired after release"); + drop(assertion); + } +} diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 9863f194ceb..03764a68a58 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -220,6 +220,10 @@ pub struct Settings { #[serde(default)] pub agent_sessions_enabled: bool, + /// Keep the system awake (no idle sleep) while a local agent session is live. + #[serde(default)] + pub stay_awake_enabled: bool, + /// SSH targets queried for remote agent sessions. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub agent_session_ssh_hosts: Vec, @@ -562,6 +566,7 @@ impl Default for Settings { global_shortcut: default_global_shortcut(), // Ctrl+Shift+U by default codex_custom_sessions_dirs: Vec::new(), agent_sessions_enabled: false, + stay_awake_enabled: false, agent_session_ssh_hosts: Vec::new(), hooks_enabled: false, http_proxy_enabled: false, diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 850c90e9602..b5a52223546 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -121,6 +121,8 @@ pub(super) struct RawSettings { global_shortcut: String, codex_custom_sessions_dirs: Vec, agent_sessions_enabled: bool, + #[serde(default)] + stay_awake_enabled: bool, agent_session_ssh_hosts: Vec, #[serde(default)] hooks_enabled: bool, @@ -255,6 +257,7 @@ impl Default for RawSettings { global_shortcut: s.global_shortcut, codex_custom_sessions_dirs: s.codex_custom_sessions_dirs, agent_sessions_enabled: s.agent_sessions_enabled, + stay_awake_enabled: s.stay_awake_enabled, agent_session_ssh_hosts: s.agent_session_ssh_hosts, hooks_enabled: s.hooks_enabled, http_proxy_enabled: s.http_proxy_enabled, @@ -558,6 +561,7 @@ impl From for Settings { global_shortcut: raw.global_shortcut, codex_custom_sessions_dirs: raw.codex_custom_sessions_dirs, agent_sessions_enabled: raw.agent_sessions_enabled, + stay_awake_enabled: raw.stay_awake_enabled, agent_session_ssh_hosts: raw.agent_session_ssh_hosts, hooks_enabled: raw.hooks_enabled, http_proxy_enabled: raw.http_proxy_enabled,