diff --git a/CHANGELOG.md b/CHANGELOG.md index 41550d5..20f1433 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,6 +60,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- **`muxa upgrade` no longer strands an old daemon when no service manager is + usable.** A capable muxad now drains and re-execs itself on the exact socket + being upgraded, preserving its pid and launch context. The CLI confirms the + replacement by an IPC generation advance instead of a connect-only timing + guess, and source, Homebrew, and release-binary upgrades share the same + restart path. Native service managers remain the compatibility fallback for + older or stopped daemons. SIGTERM/SIGINT wins atomically over any in-flight + restart request, so an explicit stop cannot be accidentally re-armed. - **A stale or non-participant row no longer disables watch's messaging.** `muxa register` task rows and just-stopped agents look like agents in the topology but are not collaboration participants; pointing at one now costs diff --git a/PROTOCOL.md b/PROTOCOL.md index ecaaee6..4037802 100644 --- a/PROTOCOL.md +++ b/PROTOCOL.md @@ -116,6 +116,32 @@ Response: "health": { "version": "0.0.1", "protocol": 1 } } ``` +#### `restart` + +Ask the daemon to drain and re-execute itself onto the binary now resolved by +its original `argv[0]`. + +```json +{ "protocol": 4, "kind": "restart" } +``` + +The daemon commits to the restart before replying. An `ok` response therefore +means accepted, not that the replacement image is already serving: + +```json +{ "ok": true, "protocol": 4 } +``` + +Only a server advertising the `restart` capability accepts this method. An +embedded server without a restart controller refuses instead of draining. A +daemon that has already received SIGTERM/SIGINT refuses too; an in-flight IPC +handler cannot reverse an operator-requested stop. + +Re-exec preserves the pid, argv, environment and working directory, and the +old listener may finish an already accepted request during the drain. Clients +must therefore confirm completion by observing `hello.generation` advance, +not by checking the pid or merely connecting to the socket. + #### `hello` Capability handshake. Optional, but clients SHOULD send it as the first @@ -144,7 +170,8 @@ Response: "protocol": 2, "min_protocol": 1, "max_protocol": 2, - "capabilities": ["waiting_choice", "needs_choice", "rate_limited"] + "capabilities": ["waiting_choice", "needs_choice", "rate_limited", "restart"], + "generation": 0 } ``` @@ -157,6 +184,8 @@ Response: features the server supports. Clients SHOULD feature-gate on these rather than comparing `protocol` integers, so adding a new tag is always non-breaking. +- `generation`: present only with the `restart` capability. It starts at zero + for a fresh daemon and increments across each self-reexec. Capability tags currently advertised: @@ -165,6 +194,7 @@ Capability tags currently advertised: | `waiting_choice` | server emits `AgentState::waiting_choice` (otherwise: `waiting_input`). | | `needs_choice` | server emits `NotificationLevel::needs_choice` (otherwise: `needs_input`). | | `rate_limited` | server emits the `rate_limited` event type and the `rate_limit_*` fields on `Agent`. | +| `restart` | server accepts `restart` and can re-exec itself in place. | #### v1-compat downgrade diff --git a/crates/muxa-cli/src/upgrade.rs b/crates/muxa-cli/src/upgrade.rs index f2601c4..1df707b 100644 --- a/crates/muxa-cli/src/upgrade.rs +++ b/crates/muxa-cli/src/upgrade.rs @@ -32,8 +32,10 @@ //! `muxad` in place (previous binaries kept as `.bak`), restart. //! //! The daemon restart + socket verification tail is shared by all three. +//! The running daemon is asked to drain and re-exec itself first; the native +//! service manager is retained as a compatibility fallback for an older or +//! absent daemon. -use crate::init::util::wait_for_muxad; use anyhow::{anyhow, Context, Result}; use clap::Parser; use std::path::{Path, PathBuf}; @@ -65,11 +67,9 @@ pub struct Args { /// Entry point dispatched from `main.rs`. /// -/// Async to match the rest of the CLI's command-dispatch shape (we -/// `.await` it in `main.rs` next to `init::run`); the body itself -/// is synchronous because cargo / git / launchctl are all blocking -/// child processes. -#[allow(clippy::unused_async)] +/// Async because the shared restart tail performs bounded IPC round trips. +/// Build/install commands remain blocking: this command is interactive and +/// has no other useful work to schedule while they run. pub async fn run(args: Args, socket: PathBuf) -> Result<()> { let _ = cliclack::intro("muxa upgrade"); @@ -81,7 +81,7 @@ pub async fn run(args: Args, socket: PathBuf) -> Result<()> { let Some(repo) = find_repo_root(&cwd) else { // No checkout in sight — most users. Resolve the channel from the // running binary instead of sending them off to clone a repo. - return run_without_repo(&args, &socket); + return run_without_repo(&args, &socket).await; }; let _ = cliclack::log::info(format!("repo: {}", repo.display())); @@ -110,53 +110,21 @@ pub async fn run(args: Args, socket: PathBuf) -> Result<()> { let _ = cliclack::log::step("building muxa-cli"); cargo_install(&repo, "crates/muxa-cli").context("cargo install muxa-cli")?; - let restart_completed = if plan.do_restart { + let restart = if plan.do_restart { let _ = cliclack::log::step("restarting daemon"); - match restart_daemon() { - RestartOutcome::Restarted => true, - RestartOutcome::ManualRequired(reason) => { - let _ = cliclack::log::warning(format!( - "{reason}; restart muxad manually to load the upgraded daemon" - )); - false - } - } + restart_daemon(&socket).await } else { let _ = cliclack::log::info("daemon restart skipped (--no-restart)"); - false + RestartOutcome::Skipped }; - // Verification only makes sense after a service manager reported - // a successful restart. Otherwise we'd be reporting on a stale - // process that the user still needs to restart manually. - if restart_completed { - let _ = cliclack::log::step("verifying"); - if wait_for_muxad(&socket, Duration::from_secs(3)) { - let _ = cliclack::log::success(format!("muxad responsive on {}", socket.display())); - } else { - let _ = cliclack::log::warning(format!( - "muxad did not respond on {} within 3s — check /tmp/muxad.log", - socket.display() - )); - } - } - let head = current_head(&repo).unwrap_or_else(|| "HEAD".into()); - if plan.do_restart && !restart_completed { - let _ = cliclack::outro(format!( - "Upgraded to {head} — manual muxad restart required." - )); - } else { - let _ = cliclack::outro(format!( - "Upgraded to {head} — try `muxa doctor` for a health check." - )); - } - Ok(()) + finish(&restart, &head, &socket) } /// Upgrade without a source checkout: Homebrew delegation or GitHub /// release self-update, chosen from where the running binary lives. -fn run_without_repo(args: &Args, socket: &Path) -> Result<()> { +async fn run_without_repo(args: &Args, socket: &Path) -> Result<()> { let exe = std::env::current_exe().context("resolving current executable")?; let exe = exe.canonicalize().unwrap_or(exe); @@ -166,8 +134,8 @@ fn run_without_repo(args: &Args, socket: &Path) -> Result<()> { let _ = cliclack::note( "Plan", "brew upgrade muxa -restart muxad -verify IPC socket", +ask muxad to re-exec itself (service-manager fallback for an older/stopped daemon) +verify IPC generation/socket", ); let _ = cliclack::outro("Dry run — no changes made."); return Ok(()); @@ -175,8 +143,8 @@ verify IPC socket", let _ = cliclack::log::step("brew upgrade muxa"); run_streaming(Command::new("brew").args(["upgrade", "muxa"])) .context("brew upgrade muxa")?; - finish_with_restart(args, socket, &format!("v{}", latest_installed_version())); - return Ok(()); + return finish_with_restart(args, socket, &format!("v{}", latest_installed_version())) + .await; } let triple = release_target_triple(std::env::consts::OS, std::env::consts::ARCH) @@ -212,8 +180,8 @@ verify IPC socket", "download muxa-{latest}-{triple}.tar.gz + .sha256 verify checksum replace muxa + muxad in {} -restart muxad -verify IPC socket", +ask muxad to re-exec itself (service-manager fallback for an older/stopped daemon) +verify IPC generation/socket", install_dir.display() ), ); @@ -260,45 +228,52 @@ verify IPC socket", } let _ = std::fs::remove_dir_all(&staging); - finish_with_restart(args, socket, &latest); - Ok(()) + finish_with_restart(args, socket, &latest).await } /// Shared tail: restart the daemon (unless opted out), verify the /// socket, and close the flow with the version we ended on. -fn finish_with_restart(args: &Args, socket: &Path, version: &str) { - let restart_completed = if args.no_restart { +async fn finish_with_restart(args: &Args, socket: &Path, version: &str) -> Result<()> { + let restart = if args.no_restart { let _ = cliclack::log::info("daemon restart skipped (--no-restart)"); - false + RestartOutcome::Skipped } else { let _ = cliclack::log::step("restarting daemon"); - match restart_daemon() { - RestartOutcome::Restarted => true, - RestartOutcome::ManualRequired(reason) => { - let _ = cliclack::log::warning(format!( - "{reason}; restart muxad manually to load the upgraded daemon" - )); - false - } - } + restart_daemon(socket).await }; - if restart_completed { - let _ = cliclack::log::step("verifying"); - if wait_for_muxad(socket, Duration::from_secs(3)) { - let _ = cliclack::log::success(format!("muxad responsive on {}", socket.display())); - } else { + finish(&restart, version, socket) +} + +/// Convert the restart result into honest user messaging and an exit status. +fn finish(restart: &RestartOutcome, version: &str, socket: &Path) -> Result<()> { + match restart { + RestartOutcome::ManualRequired(reason) => { let _ = cliclack::log::warning(format!( - "muxad did not respond on {} within 3s — check /tmp/muxad.log", + "{reason}. Start muxad yourself to run the new build:\n muxad --socket {}", socket.display() )); + let _ = cliclack::outro(format!( + "Upgraded to {version} — muxad needs a manual restart." + )); + Ok(()) + } + RestartOutcome::Failed(reason) => { + let _ = cliclack::outro(format!("Upgraded to {version} — muxad is down.")); + Err(anyhow!("{reason}")) + } + RestartOutcome::Restarted => { + let _ = cliclack::log::success(format!("muxad responsive on {}", socket.display())); + let _ = cliclack::outro(format!( + "Upgraded to {version} — try `muxa doctor` for a health check." + )); + Ok(()) + } + RestartOutcome::Skipped => { + let _ = cliclack::outro(format!( + "Upgraded to {version} — try `muxa doctor` for a health check." + )); + Ok(()) } - let _ = cliclack::outro(format!( - "Upgraded to {version} — try `muxa doctor` for a health check." - )); - } else { - let _ = cliclack::outro(format!( - "Upgraded to {version} — manual muxad restart required." - )); } } @@ -455,10 +430,13 @@ fn render_plan_for_os(plan: &Plan, os: &str) -> String { if plan.do_restart { let cmd = restart_command_args(os).join(" "); if cmd.is_empty() { - lines.push("restart: manual restart required (no supported service manager)".into()); + lines.push( + "restart: ask muxad to re-exec itself (older/stopped daemon requires a manual restart)" + .into(), + ); } else { lines.push(format!( - "restart: {cmd} (manual restart required if unavailable or unsuccessful)" + "restart: ask muxad to re-exec itself (fallback: `{cmd}` for an older/stopped daemon)" )); } } else { @@ -574,31 +552,131 @@ fn uid_string() -> String { enum RestartOutcome { Restarted, + /// Nothing was stopped. The new binary is installed, but an older daemon + /// (or no daemon) could not be started automatically. ManualRequired(String), + /// A restart was accepted or a service manager claimed success, but no + /// replacement daemon became responsive. + Failed(String), + /// `--no-restart`: preserve the caller's daemon state exactly. + Skipped, } -/// Restart through the OS-native service manager. If the manager is -/// unavailable or unsuccessful, leave all muxad processes untouched -/// and tell the caller that a manual restart is required. -fn restart_daemon() -> RestartOutcome { +/// Prefer an in-place daemon re-exec and prove it by observing a higher image +/// generation. The service manager remains the compatibility path for a daemon +/// that predates the restart capability, or when no daemon is running. +async fn restart_daemon(socket: &Path) -> RestartOutcome { + let client = muxa::ipc::Client::new(socket.to_path_buf()); + let before = match client.hello(Duration::from_secs(5)).await { + Ok(hello) + if hello.capabilities.iter().any(|cap| cap == "restart") + && hello.generation.is_some() => + { + hello.generation.expect("checked above") + } + Ok(_) => { + let _ = cliclack::log::info( + "the running daemon predates self-restart; trying the service manager", + ); + return via_service_manager(socket).await; + } + Err(muxa::ipc::RuntimeError::NotConnected(_)) => { + let _ = cliclack::log::info("no daemon was running on that socket"); + return via_service_manager(socket).await; + } + Err(error) => { + let _ = cliclack::log::info(format!( + "could not identify the daemon on {} ({error}); trying the service manager", + socket.display() + )); + return via_service_manager(socket).await; + } + }; + + if let Err(error) = client.restart(Duration::from_secs(5)).await { + // The daemon commits to restart before replying, so a lost response is + // ambiguous. Falling back here could race the re-exec for the socket; + // the generation check below is the authoritative outcome. + let _ = cliclack::log::info(format!( + "no answer to the restart request ({error}); waiting for a new generation" + )); + } + + match wait_for_new_generation(socket, before, Duration::from_secs(30)).await { + Some(after) => { + let _ = cliclack::log::info(format!("muxad came back as generation {after}")); + RestartOutcome::Restarted + } + None => RestartOutcome::Failed(format!( + "muxad did not come back on {} — check /tmp/muxad.log and start it manually", + socket.display() + )), + } +} + +/// Wait until a daemon with an image identity newer than `before` answers. +/// Socket connectability and pid are insufficient: a draining listener can +/// still accept one more request, while `exec` deliberately preserves pid. +async fn wait_for_new_generation(socket: &Path, before: u64, deadline: Duration) -> Option { + let client = muxa::ipc::Client::new(socket.to_path_buf()); + let start = std::time::Instant::now(); + while start.elapsed() < deadline { + if let Ok(hello) = client.hello(Duration::from_secs(1)).await { + if let Some(now) = hello.generation.filter(|now| *now > before) { + return Some(now); + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + None +} + +/// Use the OS service manager for an older or absent daemon, then require a +/// real IPC round trip before reporting success. +async fn via_service_manager(socket: &Path) -> RestartOutcome { + match restart_via_service_manager() { + Ok(()) if wait_for_daemon_serving(socket, Duration::from_secs(30)).await => { + RestartOutcome::Restarted + } + Ok(()) => RestartOutcome::Failed(format!( + "the service manager reported success but muxad is not answering on {} — check /tmp/muxad.log", + socket.display() + )), + Err(reason) => RestartOutcome::ManualRequired(reason), + } +} + +fn restart_via_service_manager() -> std::result::Result<(), String> { let args = restart_command_args(std::env::consts::OS); let Some((prog, rest)) = args.split_first() else { - return RestartOutcome::ManualRequired( - "no supported service manager for this operating system".into(), - ); + return Err("no supported service manager for this operating system".into()); }; if which::which(prog).is_err() { - return RestartOutcome::ManualRequired(format!("`{prog}` is not available")); + return Err(format!("`{prog}` is not available")); } match Command::new(prog).args(rest).status() { - Ok(status) if status.success() => RestartOutcome::Restarted, - Ok(status) => RestartOutcome::ManualRequired(format!("`{prog}` exited with {status}")), - Err(error) => RestartOutcome::ManualRequired(format!("could not run `{prog}`: {error}")), + Ok(status) if status.success() => Ok(()), + Ok(status) => Err(format!("`{prog}` exited with {status}")), + Err(error) => Err(format!("could not run `{prog}`: {error}")), } } +/// A serving daemon must complete a `hello` round trip; a connect-only probe +/// is fooled by a bound Unix socket whose process has stopped accepting. +async fn wait_for_daemon_serving(socket: &Path, deadline: Duration) -> bool { + let client = muxa::ipc::Client::new(socket.to_path_buf()); + let start = std::time::Instant::now(); + while start.elapsed() < deadline { + if client.hello(Duration::from_secs(1)).await.is_ok() { + return true; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + false +} + /// Short SHA of HEAD in `repo`. None when `git` is missing or the /// directory isn't a git checkout (e.g. someone exported a tarball). fn current_head(repo: &Path) -> Option { @@ -814,23 +892,24 @@ version = "0.1.0" assert!(s.contains("git pull")); assert!(s.contains("cargo install --path crates/muxad --locked --force")); assert!(s.contains("cargo install --path crates/muxa-cli --locked --force")); - assert!(s.contains("manual restart required")); + assert!(s.contains("re-exec")); + assert!(s.contains("systemctl --user restart muxad")); assert!(!s.contains("pkill")); assert!(!s.contains("SIGUSR1")); } #[test] - fn dry_run_requires_manual_restart_without_service_manager() { + fn dry_run_uses_self_restart_without_service_manager() { let plan = Plan { repo: PathBuf::from("/tmp/fake-muxa"), do_pull: true, do_restart: true, }; let s = render_plan_for_os(&plan, "plan9"); - assert!(s.contains("manual restart required")); - assert!(s.contains("no supported service manager")); + assert!(s.contains("re-exec")); + assert!(s.contains("older/stopped daemon requires a manual restart")); assert!(!s.contains("pkill")); - assert!(!s.contains("spawn")); + assert!(!s.contains("SIGUSR1")); } #[test] @@ -844,4 +923,72 @@ version = "0.1.0" assert!(s.contains("git pull (skipped)")); assert!(s.contains("restart (skipped)")); } + + async fn serve_at_generation( + socket: &Path, + generation: u64, + ) -> ( + tokio::sync::broadcast::Sender<()>, + tokio::task::JoinHandle<()>, + ) { + let (shutdown, receiver) = tokio::sync::broadcast::channel(1); + let restart = std::sync::Arc::new(muxa::ipc::RestartController::new( + generation, + shutdown.clone(), + )); + let server = muxa::ipc::Server::new(socket.to_path_buf(), muxa::Store::shared()) + .with_restart_controller(restart); + let handle = tokio::spawn(async move { + let _ = server.run(receiver).await; + }); + for _ in 0..100 { + if muxa::ipc::Client::new(socket.to_path_buf()) + .hello(Duration::from_millis(100)) + .await + .is_ok() + { + return (shutdown, handle); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("test server never came up on {}", socket.display()); + } + + #[tokio::test] + async fn same_generation_does_not_satisfy_restart_verification() { + let dir = tempdir().unwrap(); + let socket = dir.path().join("same.sock"); + let (shutdown, handle) = serve_at_generation(&socket, 4).await; + + let observed = wait_for_new_generation(&socket, 4, Duration::from_millis(300)).await; + assert_eq!(observed, None); + + let _ = shutdown.send(()); + let _ = handle.await; + } + + #[tokio::test] + async fn newer_generation_satisfies_restart_verification() { + let dir = tempdir().unwrap(); + let socket = dir.path().join("newer.sock"); + let (shutdown, handle) = serve_at_generation(&socket, 5).await; + + let observed = wait_for_new_generation(&socket, 4, Duration::from_secs(2)).await; + assert_eq!(observed, Some(5)); + + let _ = shutdown.send(()); + let _ = handle.await; + } + + #[tokio::test] + async fn serving_probe_requires_an_ipc_answer() { + let dir = tempdir().unwrap(); + let socket = dir.path().join("bound-but-silent.sock"); + let _listener = std::os::unix::net::UnixListener::bind(&socket).unwrap(); + + assert!( + !wait_for_daemon_serving(&socket, Duration::from_millis(300)).await, + "a listener with nobody accepting must not count as muxad", + ); + } } diff --git a/crates/muxa/src/ipc.rs b/crates/muxa/src/ipc.rs index 46f637f..786a83c 100644 --- a/crates/muxa/src/ipc.rs +++ b/crates/muxa/src/ipc.rs @@ -40,6 +40,7 @@ use crate::tmux::PaneInfo; use serde::{Deserialize, Serialize}; use std::os::unix::fs::PermissionsExt; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU8, Ordering as AtomicOrdering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt, BufReader}; @@ -145,6 +146,11 @@ enum RequestBody { limit: Option, }, Health, + /// Ask the daemon to drain and re-exec itself onto the binary currently + /// installed at its argv[0]. Opt-in: only the real daemon installs a + /// restart controller; embedders refuse rather than shutting down with no + /// way to come back. + Restart, /// Capability handshake. Optional first message; opts the connection /// into negotiated-protocol mode. The server replies with its /// `[min, max]` supported range and a list of capability tags, then @@ -358,6 +364,10 @@ const CAPABILITIES: &[&str] = &[ "collaboration_provenance", ]; +/// Advertised only when the server has the controller required to come back +/// after draining. A server without one refuses `restart`. +const RESTART_CAPABILITY: &str = "restart"; + #[derive(Debug, Serialize)] pub struct Response { pub ok: bool, @@ -376,6 +386,11 @@ pub struct Response { pub max_protocol: Option, #[serde(skip_serializing_if = "Option::is_none")] pub capabilities: Option>, + /// Present only when the daemon can restart itself. It increments across + /// each re-exec so a client can distinguish the replacement image from + /// the old daemon still finishing an in-flight response. + #[serde(skip_serializing_if = "Option::is_none")] + pub generation: Option, #[serde(skip_serializing_if = "Option::is_none")] pub sessions: Option>, #[serde(skip_serializing_if = "Option::is_none")] @@ -436,6 +451,7 @@ impl Response { min_protocol: None, max_protocol: None, capabilities: None, + generation: None, sessions: None, session: None, terminal: None, @@ -546,15 +562,84 @@ impl Response { r.ask_entry = Some(entry); r } - fn hello() -> Self { + fn hello(restart: Option<&RestartController>) -> Self { let mut r = Self::ok(); r.min_protocol = Some(MIN_PROTOCOL_VERSION); r.max_protocol = Some(PROTOCOL_VERSION); - r.capabilities = Some(CAPABILITIES.to_vec()); + let mut capabilities = CAPABILITIES.to_vec(); + if restart.is_some() { + capabilities.push(RESTART_CAPABILITY); + } + r.capabilities = Some(capabilities); + r.generation = restart.map(RestartController::generation); r } } +const RESTART_RUNNING: u8 = 0; +const RESTART_REQUESTED: u8 = 1; +const RESTART_STOPPING: u8 = 2; + +/// Coordinates daemon shutdown and self-restart without allowing an already +/// open IPC handler to undo an operator's later SIGTERM/SIGINT. +/// +/// The state transition is monotonic: `running -> restart_requested -> +/// stopping`, while a signal may move `running -> stopping` directly. Once +/// stopping, a restart request is refused permanently. This closes the race +/// in which a signal cleared a boolean and a draining handler set it again. +#[derive(Debug)] +pub struct RestartController { + generation: u64, + state: AtomicU8, + trigger: broadcast::Sender<()>, +} + +impl RestartController { + #[must_use] + pub fn new(generation: u64, trigger: broadcast::Sender<()>) -> Self { + Self { + generation, + state: AtomicU8::new(RESTART_RUNNING), + trigger, + } + } + + #[must_use] + pub fn generation(&self) -> u64 { + self.generation + } + + /// Commit to a normal stop and wake every shutdown subscriber. A later + /// IPC request cannot move the state back to restart-requested. + pub fn stop(&self) { + self.state.store(RESTART_STOPPING, AtomicOrdering::SeqCst); + let _ = self.trigger.send(()); + } + + #[must_use] + pub fn restart_requested(&self) -> bool { + self.state.load(AtomicOrdering::SeqCst) == RESTART_REQUESTED + } + + /// Returns false only after an explicit stop has won. Repeated restart + /// requests are idempotently accepted while the first request drains. + fn request_restart(&self) -> bool { + match self.state.compare_exchange( + RESTART_RUNNING, + RESTART_REQUESTED, + AtomicOrdering::SeqCst, + AtomicOrdering::SeqCst, + ) { + Ok(_) => { + let _ = self.trigger.send(()); + true + } + Err(RESTART_REQUESTED) => true, + Err(_) => false, + } + } +} + /// Daemon-side server. Construct once, call `run` under the tokio runtime. pub struct Server { socket_path: PathBuf, @@ -571,6 +656,7 @@ pub struct Server { collaboration: Arc, collaboration_audit: Arc, ask: Arc, + restart: Option>, handler_limit: usize, } @@ -586,6 +672,7 @@ impl Server { collaboration: CollaborationStore::in_memory(CollaborationOptions::default()), collaboration_audit: CollaborationAuditLog::in_memory(), ask: crate::ask::AskStore::in_memory(crate::ask::AskOptions::default()), + restart: None, handler_limit: MAX_INFLIGHT_HANDLERS, } } @@ -640,6 +727,14 @@ impl Server { self } + /// Allow this server to accept the restart control method. Kept opt-in so + /// embedded servers and tests never drain unless they can re-exec. + #[must_use] + pub fn with_restart_controller(mut self, restart: Arc) -> Self { + self.restart = Some(restart); + self + } + #[cfg(test)] #[must_use] fn with_handler_limit(mut self, handler_limit: usize) -> Self { @@ -732,6 +827,7 @@ impl Server { let collaboration = self.collaboration.clone(); let collaboration_audit = self.collaboration_audit.clone(); let ask = self.ask.clone(); + let restart = self.restart.clone(); handlers.spawn(async move { // Held for the handler's lifetime; released here on exit. let _permit = permit; @@ -745,6 +841,7 @@ impl Server { collaboration, collaboration_audit, ask, + restart, )) .await { @@ -1348,7 +1445,8 @@ async fn record_collaboration_audit( sessions, collaboration, collaboration_audit, - ask + ask, + restart ) )] #[allow(clippy::too_many_arguments, clippy::too_many_lines)] // IPC dispatch table and its shared daemon state @@ -1361,6 +1459,7 @@ async fn handle( collaboration: Arc, collaboration_audit: Arc, ask: Arc, + restart: Option>, ) -> Result<(), RuntimeError> { let mut collaboration_actor = observe_collaboration_actor(&stream); let (reader, mut writer) = stream.into_split(); @@ -1455,7 +1554,7 @@ async fn handle( protocol = requested, "hello" ); - let mut r = Response::hello(); + let mut r = Response::hello(restart.as_deref()); r.protocol = requested; r } else { @@ -1549,6 +1648,24 @@ async fn handle( kind = "health"; Response::health() } + RequestBody::Restart => { + kind = "restart"; + match &restart { + Some(controller) if controller.request_restart() => { + tracing::info!( + generation = controller.generation(), + "restart requested over IPC", + ); + Response::ok() + } + Some(_) => { + Response::err("daemon is already stopping; restart request refused") + } + None => Response::err( + "this server cannot restart itself (no restart controller installed)", + ), + } + } RequestBody::BackendPaneSnapshot { panes } => { kind = "backend_pane_snapshot"; let count = panes.len(); @@ -2119,6 +2236,13 @@ pub struct Client { collaboration_client_kind: CollaborationClientKind, } +/// Identity and feature information returned by the daemon's `hello` method. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Hello { + pub capabilities: Vec, + pub generation: Option, +} + /// The result of a [`Client::send_prompt`]: the two non-atomic keystroke /// injections (the text, then the optional submit CR) reported distinctly. /// @@ -2251,6 +2375,50 @@ impl Client { Ok(decode_agents(&resp)) } + /// Ask the daemon which additive features it supports and, when it can + /// self-restart, which process-image generation is currently serving. + pub async fn hello(&self, deadline: Duration) -> Result { + let req = serde_json::json!({ + "protocol": PROTOCOL_VERSION, + "kind": "hello", + "client": self.collaboration_client_kind.hello_label(), + }); + let resp = self.call_with_timeout(&req, deadline).await?; + if !resp["ok"].as_bool().unwrap_or(false) { + return Err(RuntimeError::Json(serde::de::Error::custom(format!( + "hello rejected: {}", + resp["error"].as_str().unwrap_or("(no error message)") + )))); + } + Ok(Hello { + capabilities: resp["capabilities"] + .as_array() + .map(|capabilities| { + capabilities + .iter() + .filter_map(|capability| capability.as_str().map(str::to_string)) + .collect() + }) + .unwrap_or_default(), + generation: resp["generation"].as_u64(), + }) + } + + /// Ask the daemon on this socket to drain and re-exec itself. Acceptance + /// is not completion; callers confirm completion by waiting for `hello`'s + /// generation to advance. + pub async fn restart(&self, deadline: Duration) -> Result<(), RuntimeError> { + let req = serde_json::json!({ "protocol": PROTOCOL_VERSION, "kind": "restart" }); + let resp = self.call_with_timeout(&req, deadline).await?; + if !resp["ok"].as_bool().unwrap_or(false) { + return Err(RuntimeError::Json(serde::de::Error::custom(format!( + "restart rejected: {}", + resp["error"].as_str().unwrap_or("(no error message)") + )))); + } + Ok(()) + } + /// Ask the daemon to delete fully orphaned rows (no pane, surface, or /// pid) idle longer than `max_age`. `max_age = Duration::ZERO` removes /// every orphan regardless of age. Returns the number removed. Backs @@ -3497,6 +3665,7 @@ mod tests { CollaborationStore::in_memory(CollaborationOptions::default()), CollaborationAuditLog::in_memory(), crate::ask::AskStore::in_memory(crate::ask::AskOptions::default()), + None, )); let req = serde_json::json!({ @@ -3660,6 +3829,114 @@ mod tests { assert!(caps.contains(&"waiting_choice")); assert!(caps.contains(&"needs_choice")); assert!(caps.contains(&"rate_limited")); + assert!(!caps.contains(&RESTART_CAPABILITY)); + assert!(resp["generation"].is_null()); + + tx.send(()).unwrap(); + handle.await.unwrap(); + } + + #[tokio::test] + async fn restart_is_advertised_accepted_and_drained() { + let dir = tempdir().unwrap(); + let sock = dir.path().join("muxa-restart.sock"); + let store = Store::shared(); + let (tx, rx) = broadcast::channel(1); + let restart = Arc::new(RestartController::new(7, tx)); + let server = Server::new(sock.clone(), store).with_restart_controller(Arc::clone(&restart)); + let handle = tokio::spawn(async move { server.run(rx).await.unwrap() }); + wait_for_socket(&sock).await; + + let client = Client::new(sock.clone()); + let hello = client + .hello(Duration::from_secs(2)) + .await + .expect("hello answers"); + assert!(hello + .capabilities + .iter() + .any(|cap| cap == RESTART_CAPABILITY)); + assert_eq!(hello.generation, Some(7)); + + client + .restart(Duration::from_secs(2)) + .await + .expect("daemon accepts restart"); + tokio::time::timeout(Duration::from_secs(2), handle) + .await + .expect("daemon drains after accepting restart") + .unwrap(); + assert!(restart.restart_requested()); + assert!(!sock.exists(), "drained server removes its socket"); + } + + #[tokio::test] + async fn signal_stop_cannot_be_rearmed_by_an_inflight_restart() { + let dir = tempdir().unwrap(); + let sock = dir.path().join("muxa-stopping.sock"); + let store = Store::shared(); + let (tx, rx) = broadcast::channel(1); + let restart = Arc::new(RestartController::new(0, tx)); + let server = Server::new(sock.clone(), store).with_restart_controller(Arc::clone(&restart)); + let handle = tokio::spawn(async move { server.run(rx).await.unwrap() }); + wait_for_socket(&sock).await; + + // Get a handler accepted and parked mid-request before the stop. This + // is the exact ordering that could re-arm the old AtomicBool design. + let mut stream = tokio::net::UnixStream::connect(&sock).await.unwrap(); + let mut request = serde_json::to_vec(&serde_json::json!({ + "protocol": PROTOCOL_VERSION, + "kind": "restart", + })) + .unwrap(); + request.push(b'\n'); + let split = request.len() - 1; + stream.write_all(&request[..split]).await.unwrap(); + stream.flush().await.unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + + restart.stop(); + stream.write_all(&request[split..]).await.unwrap(); + stream.flush().await.unwrap(); + + let mut reader = BufReader::new(stream); + let mut line = String::new(); + reader.read_line(&mut line).await.unwrap(); + let response: serde_json::Value = serde_json::from_str(line.trim()).unwrap(); + assert_eq!(response["ok"], false); + assert!(response["error"] + .as_str() + .unwrap() + .contains("already stopping")); + drop(reader); + + tokio::time::timeout(Duration::from_secs(2), handle) + .await + .expect("normal stop drains the in-flight handler") + .unwrap(); + assert!( + !restart.restart_requested(), + "an in-flight restart must not override SIGTERM/SIGINT", + ); + } + + #[tokio::test] + async fn restart_is_refused_without_a_controller() { + let dir = tempdir().unwrap(); + let sock = dir.path().join("muxa-no-restart.sock"); + let store = Store::shared(); + let server = Server::new(sock.clone(), store); + let (tx, rx) = broadcast::channel(1); + let handle = tokio::spawn(async move { server.run(rx).await.unwrap() }); + wait_for_socket(&sock).await; + + let client = Client::new(sock.clone()); + let error = client + .restart(Duration::from_secs(2)) + .await + .expect_err("embedded server refuses restart"); + assert!(error.to_string().contains("restart")); + assert!(UnixStream::connect(&sock).await.is_ok()); tx.send(()).unwrap(); handle.await.unwrap(); diff --git a/crates/muxad/src/main.rs b/crates/muxad/src/main.rs index b533dd7..e566aa5 100644 --- a/crates/muxad/src/main.rs +++ b/crates/muxad/src/main.rs @@ -21,7 +21,7 @@ use muxa::config::CollaborationWake; use muxa::config::{DashboardAuthMode, NotifierBackend}; use muxa::dashboard::{DashboardConfig, DashboardOverrides}; use muxa::history::{HistoryOptions, PaneSessionCache, PromptHistory}; -use muxa::ipc::{harden_permissions, Client, Server}; +use muxa::ipc::{harden_permissions, Client, RestartController, Server}; use muxa::notify::Notifier; use muxa::reconcile::Reconciler; use muxa::sinks::{webhook as webhook_sink, OhMyPromptSink, WebhookSink}; @@ -45,6 +45,8 @@ const STOPPED_AGENT_TTL_MINUTES: i64 = 60; const GC_SWEEP_INTERVAL_SECONDS: u64 = 60; const PANE_SESSION_CACHE_INTERVAL_SECONDS: u64 = 5; const SHUTDOWN_TASK_TIMEOUT_SECONDS: u64 = 2; +/// Carries the daemon image identity across an in-place re-exec. +const RESTART_GENERATION_ENV: &str = "MUXA_RESTART_GENERATION"; #[derive(Debug, Parser)] #[command(name = "muxad", version, about = "muxa daemon")] @@ -144,7 +146,11 @@ async fn main() -> Result<()> { let (activity_transition_shutdown_tx, _) = broadcast::channel::<()>(1); let (writer_shutdown_tx, _) = broadcast::channel::<()>(1); - install_shutdown_signal_handler(shutdown_tx.clone()); + let restart = Arc::new(RestartController::new( + restart_generation(), + shutdown_tx.clone(), + )); + install_shutdown_signal_handler(Arc::clone(&restart)); // Prompt history must exist before the store: every PromptSubmitted // event fans out into history alongside the live agent record. @@ -346,7 +352,8 @@ async fn main() -> Result<()> { .with_sessions(sessions) .with_collaboration(collaboration) .with_collaboration_audit(collaboration_audit) - .with_ask(ask); + .with_ask(ask) + .with_restart_controller(Arc::clone(&restart)); let handle = tokio::spawn(server.run(shutdown_tx.subscribe())); // Harden socket permissions once the listener exists. We poll briefly @@ -409,9 +416,39 @@ async fn main() -> Result<()> { let _ = snap_shutdown_tx.send(()); await_shutdown_task("state snapshotter", snap_handle).await; server_result??; + if restart.restart_requested() { + return Err(reexec_self().into()); + } Ok(()) } +/// Image generation advertised in IPC `hello`: zero on a fresh process and +/// incremented for each successful self-reexec. +fn restart_generation() -> u64 { + std::env::var(RESTART_GENERATION_ENV) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(0) +} + +/// Replace this process with the binary now resolved by its original argv[0]. +/// `exec` preserves pid, argv, environment, working directory and the service +/// manager's ownership while loading the newly installed inode. +fn reexec_self() -> std::io::Error { + use std::os::unix::process::CommandExt; + + let mut argv = std::env::args_os(); + let Some(program) = argv.next() else { + return std::io::Error::other("cannot restart: argv[0] is missing"); + }; + let next = restart_generation().saturating_add(1); + tracing::info!(?program, generation = next, "restarting: re-executing self"); + std::process::Command::new(&program) + .args(argv) + .env(RESTART_GENERATION_ENV, next.to_string()) + .exec() +} + async fn await_shutdown_task(name: &'static str, handle: Option>) { let Some(mut handle) = handle else { return; @@ -725,7 +762,7 @@ fn spawn_gc_task( }) } -fn install_shutdown_signal_handler(shutdown_tx: broadcast::Sender<()>) { +fn install_shutdown_signal_handler(restart: Arc) { tokio::spawn(async move { let mut term = signal(SignalKind::terminate()).expect("install SIGTERM handler"); let mut int = signal(SignalKind::interrupt()).expect("install SIGINT handler"); @@ -733,7 +770,9 @@ fn install_shutdown_signal_handler(shutdown_tx: broadcast::Sender<()>) { _ = term.recv() => tracing::info!("SIGTERM received"), _ = int.recv() => tracing::info!("SIGINT received"), } - let _ = shutdown_tx.send(()); + // Monotonic stop state: once a signal wins, an already-open IPC + // handler cannot re-arm a restart during the drain. + restart.stop(); }); }