Improve account switching reliability and UX

Author: Fzhiyu1

src-tauri/src/auth/fs_utils.rs

@@ -1,7 +1,9 @@
 use std::ffi::OsString;
 use std::fs::{self, File, OpenOptions};
+use std::io::Read;
 use std::io::Write;
 use std::path::{Path, PathBuf};
+use std::process::{Command, Stdio};
 use std::thread;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
@@ -32,6 +34,7 @@ impl FileLock {
                     return Ok(Self { path: lock_path });
                 }
                 Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
+                    maybe_clear_stale_lock(&lock_path);
                     thread::sleep(Duration::from_millis(50));
                 }
                 Err(err) => {
@@ -42,8 +45,8 @@ impl FileLock {
             }
         }
 
-        anyhow::bail!("Timed out waiting for file lock: {}", lock_path.display());
-    }
+    anyhow::bail!("Timed out waiting for file lock: {}", lock_path.display());
+}
 }
 
 impl Drop for FileLock {
@@ -106,3 +109,91 @@ fn temp_path_for(path: &Path) -> PathBuf {
         .unwrap_or(0);
     sibling_with_suffix(path, &format!(".tmp-{}-{nanos}", std::process::id()))
 }
+
+fn maybe_clear_stale_lock(lock_path: &Path) {
+    let Ok(Some(pid)) = read_lock_pid(lock_path) else {
+        return;
+    };
+
+    if !process_is_alive(pid) {
+        let _ = fs::remove_file(lock_path);
+    }
+}
+
+fn read_lock_pid(lock_path: &Path) -> Result<Option<u32>> {
+    let mut content = String::new();
+    File::open(lock_path)
+        .with_context(|| format!("Failed to open lock file: {}", lock_path.display()))?
+        .read_to_string(&mut content)
+        .with_context(|| format!("Failed to read lock file: {}", lock_path.display()))?;
+
+    let trimmed = content.trim();
+    if trimmed.is_empty() {
+        return Ok(None);
+    }
+
+    let pid = trimmed
+        .parse::<u32>()
+        .with_context(|| format!("Invalid PID in lock file: {}", lock_path.display()))?;
+    Ok(Some(pid))
+}
+
+#[cfg(unix)]
+fn process_is_alive(pid: u32) -> bool {
+    Command::new("kill")
+        .args(["-0", &pid.to_string()])
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .status()
+        .map(|status| status.success())
+        .unwrap_or(false)
+}
+
+#[cfg(windows)]
+fn process_is_alive(pid: u32) -> bool {
+    Command::new("tasklist")
+        .args(["/FI", &format!("PID eq {pid}")])
+        .stdout(Stdio::piped())
+        .stderr(Stdio::null())
+        .output()
+        .map(|output| {
+            let stdout = String::from_utf8_lossy(&output.stdout);
+            stdout.lines().any(|line| line.contains(&pid.to_string()))
+        })
+        .unwrap_or(false)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_dir() -> PathBuf {
+        let path = std::env::temp_dir().join(format!(
+            "codex-switcher-fs-utils-{}-{}",
+            std::process::id(),
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .map(|duration| duration.as_nanos())
+                .unwrap_or(0)
+        ));
+        fs::create_dir_all(&path).unwrap();
+        path
+    }
+
+    #[test]
+    fn acquire_reclaims_stale_lock_file() {
+        let dir = test_dir();
+        let target = dir.join("auth.json");
+        let lock_path = sibling_with_suffix(&target, ".lock");
+        fs::write(&lock_path, format!("{}\n", u32::MAX)).unwrap();
+
+        let lock = FileLock::acquire(&target).expect("stale lock should be reclaimed");
+        let content = fs::read_to_string(&lock_path).unwrap();
+
+        assert_eq!(content.trim(), std::process::id().to_string());
+
+        drop(lock);
+        assert!(!lock_path.exists());
+        fs::remove_dir_all(dir).unwrap();
+    }
+}

src-tauri/src/commands/account.rs

@@ -55,6 +55,15 @@ const MAX_IMPORT_JSON_BYTES: u64 = 2 * 1024 * 1024;
 const MAX_IMPORT_FILE_BYTES: u64 = 8 * 1024 * 1024;
 const SLIM_IMPORT_CONCURRENCY: usize = 6;
 
+#[derive(Debug, Clone, Copy, Default, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub enum SwitchAccountMode {
+    #[default]
+    RequireRestart,
+    KeepRunning,
+    RestartRunning,
+}
+
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
 struct SlimPayload {
     #[serde(rename = "v")]
@@ -124,7 +133,7 @@ pub async fn add_account_from_file(path: String, name: String) -> Result<Account
 #[tauri::command]
 pub async fn switch_account(
     account_id: String,
-    restart_running_codex: Option<bool>,
+    switch_mode: Option<SwitchAccountMode>,
 ) -> Result<(), String> {
     let store = load_accounts().map_err(|e| e.to_string())?;
     let running_processes = collect_running_codex_processes().map_err(|e| e.to_string())?;
@@ -136,12 +145,8 @@ pub async fn switch_account(
         .find(|a| a.id == account_id)
         .ok_or_else(|| format!("Account not found: {account_id}"))?;
 
-    let should_restart = restart_running_codex.unwrap_or(false);
-    if !running_processes.is_empty() && !should_restart {
-        return Err(String::from(
-            "Codex is currently running. Confirm a graceful restart before switching accounts.",
-        ));
-    }
+    let should_restart =
+        resolve_switch_behavior(!running_processes.is_empty(), switch_mode.unwrap_or_default())?;
 
     if should_restart {
         gracefully_stop_codex_processes(&running_processes).map_err(|e| e.to_string())?;
@@ -158,6 +163,50 @@ pub async fn switch_account(
     Ok(())
 }
 
+fn resolve_switch_behavior(
+    has_running_processes: bool,
+    switch_mode: SwitchAccountMode,
+) -> Result<bool, String> {
+    if !has_running_processes {
+        return Ok(false);
+    }
+
+    match switch_mode {
+        SwitchAccountMode::RequireRestart => Err(String::from(
+            "Codex is currently running. Choose whether to keep current sessions running or restart them before switching accounts.",
+        )),
+        SwitchAccountMode::KeepRunning => Ok(false),
+        SwitchAccountMode::RestartRunning => Ok(true),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn keep_running_mode_allows_switch_with_running_processes() {
+        let should_restart = resolve_switch_behavior(true, SwitchAccountMode::KeepRunning).unwrap();
+
+        assert!(!should_restart);
+    }
+
+    #[test]
+    fn restart_running_mode_requests_restart_when_processes_are_running() {
+        let should_restart =
+            resolve_switch_behavior(true, SwitchAccountMode::RestartRunning).unwrap();
+
+        assert!(should_restart);
+    }
+
+    #[test]
+    fn missing_mode_requires_explicit_choice_when_processes_are_running() {
+        let error = resolve_switch_behavior(true, SwitchAccountMode::RequireRestart).unwrap_err();
+
+        assert!(error.contains("Choose whether to keep current sessions running"));
+    }
+}
+
 #[tauri::command]
 pub async fn get_app_settings() -> Result<AppSettings, String> {
     load_settings().map_err(|e| e.to_string())

src-tauri/src/commands/process.rs

@@ -107,7 +107,7 @@ pub fn gracefully_stop_codex_processes(processes: &[RunningCodexProcess]) -> any
         return Ok(());
     }
 
-    anyhow::bail!("Timed out waiting for Codex processes to close gracefully");
+    anyhow::bail!(format_graceful_shutdown_timeout(processes));
 }
 
 pub fn restart_codex_processes(processes: &[RunningCodexProcess]) -> anyhow::Result<()> {
@@ -157,40 +157,67 @@ fn collect_running_codex_processes_unix() -> anyhow::Result<Vec<RunningCodexProc
     let stdout = String::from_utf8_lossy(&output.stdout);
 
     for line in stdout.lines() {
-        let line = line.trim();
-        if line.is_empty() {
-            continue;
-        }
-
-        if let Some((pid_str, command)) = line.split_once(' ') {
-            let command = command.trim();
-            let executable = command.split_whitespace().next().unwrap_or("");
-            let is_codex = executable == "codex" || executable.ends_with("/codex");
-            let is_background = command.contains(".antigravity")
-                || command.contains("openai.chatgpt")
-                || command.contains(".vscode");
-            let is_switcher =
-                command.contains("codex-switcher") || command.contains("Codex Switcher");
-
-            if is_codex && !is_switcher {
-                if let Ok(pid) = pid_str.trim().parse::<u32>() {
-                    if pid != std::process::id()
-                        && !processes.iter().any(|p: &RunningCodexProcess| p.pid == pid)
-                    {
-                        processes.push(RunningCodexProcess {
-                            pid,
-                            command: command.to_string(),
-                            is_background,
-                        });
-                    }
-                }
+        if let Some(process) = parse_unix_process_line(line) {
+            if process.pid != std::process::id()
+                && !processes.iter().any(|p: &RunningCodexProcess| p.pid == process.pid)
+            {
+                processes.push(process);
             }
         }
     }
 
     Ok(processes)
 }
 
+#[cfg(unix)]
+fn parse_unix_process_line(line: &str) -> Option<RunningCodexProcess> {
+    let line = line.trim();
+    if line.is_empty() {
+        return None;
+    }
+
+    let (pid_str, command) = line.split_once(' ')?;
+    let command = command.trim();
+    let executable = command.split_whitespace().next().unwrap_or("");
+    let is_codex = executable == "codex" || executable.ends_with("/codex");
+    let is_background =
+        command.contains(".antigravity") || command.contains("openai.chatgpt") || command.contains(".vscode");
+    let is_switcher = command.contains("codex-switcher") || command.contains("Codex Switcher");
+    let is_codex_app_server =
+        command.contains("/Codex.app/Contents/Resources/codex app-server")
+            || command.contains("/Applications/Codex.app/Contents/Resources/codex app-server");
+
+    if !is_codex || is_switcher || is_codex_app_server {
+        return None;
+    }
+
+    let pid = pid_str.trim().parse::<u32>().ok()?;
+    Some(RunningCodexProcess {
+        pid,
+        command: command.to_string(),
+        is_background,
+    })
+}
+
+fn format_graceful_shutdown_timeout(processes: &[RunningCodexProcess]) -> String {
+    let details = processes
+        .iter()
+        .map(|process| format!("pid {} ({})", process.pid, summarize_command(&process.command)))
+        .collect::<Vec<_>>()
+        .join(", ");
+    format!("Timed out waiting for Codex processes to close gracefully: {details}")
+}
+
+fn summarize_command(command: &str) -> String {
+    const MAX_LEN: usize = 80;
+    if command.chars().count() <= MAX_LEN {
+        return command.to_string();
+    }
+
+    let summary = command.chars().take(MAX_LEN - 3).collect::<String>();
+    format!("{summary}...")
+}
+
 #[cfg(windows)]
 fn collect_running_codex_processes_windows() -> anyhow::Result<Vec<RunningCodexProcess>> {
     let mut processes = Vec::new();
@@ -221,3 +248,42 @@ fn collect_running_codex_processes_windows() -> anyhow::Result<Vec<RunningCodexP
 
     Ok(processes)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn ignores_codex_app_server_processes() {
+        let line = "5989 /Applications/Codex.app/Contents/Resources/codex app-server --analytics-default-enabled";
+
+        let process = parse_unix_process_line(line);
+
+        assert!(process.is_none());
+    }
+
+    #[test]
+    fn timeout_error_lists_remaining_processes() {
+        let processes = vec![
+            RunningCodexProcess {
+                pid: 100,
+                command: String::from(
+                    "/opt/homebrew/lib/node_modules/@openai/codex/vendor/codex/codex resume 123",
+                ),
+                is_background: false,
+            },
+            RunningCodexProcess {
+                pid: 200,
+                command: String::from("/Applications/Codex.app/Contents/Resources/codex app-server"),
+                is_background: true,
+            },
+        ];
+
+        let message = format_graceful_shutdown_timeout(&processes);
+
+        assert!(message.contains("100"));
+        assert!(message.contains("resume 123"));
+        assert!(message.contains("200"));
+        assert!(message.contains("app-server"));
+    }
+}