fix termination bug

src/daemon/mod.rs

@@ -1,6 +1,7 @@
 mod handler;
 mod state;
 
+use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
 use anyhow::{Context, Result};
@@ -12,6 +13,29 @@ use crate::core::paths::{pid_path_from_env, socket_path_from_env};
 use crate::core::session;
 use state::DaemonState;
 
+struct RuntimePathsGuard {
+    socket_path: PathBuf,
+    pid_path: Option<PathBuf>,
+}
+
+impl RuntimePathsGuard {
+    fn new(socket_path: PathBuf, pid_path: Option<PathBuf>) -> Self {
+        Self {
+            socket_path,
+            pid_path,
+        }
+    }
+}
+
+impl Drop for RuntimePathsGuard {
+    fn drop(&mut self) {
+        remove_runtime_path(&self.socket_path);
+        if let Some(ref pid_path) = self.pid_path {
+            remove_runtime_path(pid_path);
+        }
+    }
+}
+
 pub fn run() -> Result<()> {
     // Validate session before starting
     session::detect_session()?;
@@ -25,7 +49,6 @@ pub fn run() -> Result<()> {
 
 async fn async_run() -> Result<()> {
     let socket_path = socket_path_from_env().context("DESKCTL_SOCKET_PATH not set")?;
-
     let pid_path = pid_path_from_env();
 
     // Clean up stale socket
@@ -33,20 +56,21 @@ async fn async_run() -> Result<()> {
         std::fs::remove_file(&socket_path)?;
     }
 
-    // Write PID file
-    if let Some(ref pid_path) = pid_path {
-        std::fs::write(pid_path, std::process::id().to_string())?;
-    }
-
-    let listener = UnixListener::bind(&socket_path)
-        .context(format!("Failed to bind socket: {}", socket_path.display()))?;
-
     let session = std::env::var("DESKCTL_SESSION").unwrap_or_else(|_| "default".to_string());
     let state = Arc::new(Mutex::new(
         DaemonState::new(session, socket_path.clone())
             .context("Failed to initialize daemon state")?,
     ));
 
+    let listener = UnixListener::bind(&socket_path)
+        .context(format!("Failed to bind socket: {}", socket_path.display()))?;
+    let _runtime_paths = RuntimePathsGuard::new(socket_path.clone(), pid_path.clone());
+
+    // Write PID file only after the daemon is ready to serve requests.
+    if let Some(ref pid_path) = pid_path {
+        std::fs::write(pid_path, std::process::id().to_string())?;
+    }
+
     let shutdown = Arc::new(tokio::sync::Notify::new());
     let shutdown_clone = shutdown.clone();
 
@@ -75,14 +99,6 @@ async fn async_run() -> Result<()> {
         }
     }
 
-    // Cleanup
-    if socket_path.exists() {
-        let _ = std::fs::remove_file(&socket_path);
-    }
-    if let Some(ref pid_path) = pid_path {
-        let _ = std::fs::remove_file(pid_path);
-    }
-
     Ok(())
 }
 
@@ -123,3 +139,11 @@ async fn handle_connection(
 
     Ok(())
 }
+
+fn remove_runtime_path(path: &Path) {
+    if let Err(error) = std::fs::remove_file(path) {
+        if error.kind() != std::io::ErrorKind::NotFound {
+            eprintln!("Failed to remove runtime path {}: {error}", path.display());
+        }
+    }
+}

tests/support/mod.rs

@@ -142,6 +142,10 @@ impl TestSession {
             .expect("TestSession always has an explicit socket path")
     }
 
+    pub fn pid_path(&self) -> PathBuf {
+        self.root.join("deskctl.pid")
+    }
+
     pub fn create_stale_socket(&self) -> Result<()> {
         let listener = UnixListener::bind(self.socket_path())
             .with_context(|| format!("Failed to bind {}", self.socket_path().display()))?;
@@ -187,6 +191,29 @@ impl TestSession {
             )
         })
     }
+
+    pub fn run_daemon<I, K, V>(&self, env: I) -> Result<Output>
+    where
+        I: IntoIterator<Item = (K, V)>,
+        K: AsRef<std::ffi::OsStr>,
+        V: AsRef<std::ffi::OsStr>,
+    {
+        let mut command = Command::new(env!("CARGO_BIN_EXE_deskctl"));
+        command
+            .env("DESKCTL_DAEMON", "1")
+            .env("DESKCTL_SOCKET_PATH", self.socket_path())
+            .env("DESKCTL_PID_PATH", self.pid_path())
+            .env("DESKCTL_SESSION", &self.opts.session)
+            .envs(env);
+
+        command.output().with_context(|| {
+            format!(
+                "Failed to run daemon {} against {}",
+                env!("CARGO_BIN_EXE_deskctl"),
+                self.socket_path().display()
+            )
+        })
+    }
 }
 
 impl Drop for TestSession {
@@ -195,6 +222,9 @@ impl Drop for TestSession {
         if self.socket_path().exists() {
             let _ = std::fs::remove_file(self.socket_path());
         }
+        if self.pid_path().exists() {
+            let _ = std::fs::remove_file(self.pid_path());
+        }
         let _ = std::fs::remove_dir_all(&self.root);
     }
 }

tests/x11_runtime.rs

@@ -114,6 +114,31 @@ fn daemon_start_recovers_from_stale_socket() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn daemon_init_failure_cleans_runtime_state() -> Result<()> {
+    let _guard = env_lock_guard();
+    let session = TestSession::new("daemon-init-failure")?;
+
+    let output = session.run_daemon([("XDG_SESSION_TYPE", "x11"), ("DISPLAY", ":99999")])?;
+    assert!(!output.status.success(), "daemon startup should fail");
+
+    let stderr = String::from_utf8_lossy(&output.stderr);
+    assert!(
+        stderr.contains("Failed to initialize daemon state"),
+        "unexpected stderr: {stderr}"
+    );
+    assert!(
+        !session.socket_path().exists(),
+        "failed startup should remove the socket path"
+    );
+    assert!(
+        !session.pid_path().exists(),
+        "failed startup should remove the pid path"
+    );
+
+    Ok(())
+}
+
 #[test]
 fn wait_window_returns_matched_window_payload() -> Result<()> {
     let _guard = env_lock_guard();