From 0ebfa19cd875ca6d9dd3799f17c87b72bfdc0b79 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Mon, 17 Aug 2026 22:07:36 +1000 Subject: [PATCH 01/13] feat(lifecycle): keep Staged running when its last window closes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closing the window (red button / Cmd+W) terminated the app, taking every running agent session with it — and Cmd+Q was worse: `PredefinedMenuItem::quit` maps to `NSApp terminate:`, which reaches no Tauri hook, so it skipped the action-shutdown handler entirely and left agent CLIs (spawned with their own process group and only `kill_on_drop`) orphaned. New `app_lifecycle` module owns both halves of the fix, adapted to the multi-window model (#928) this lands on top of: - Closing a window with peers still live (visible or hidden) is just a close: sessions belong to the process, so the window is destroyed normally and the existing `Destroyed` hook does the per-window cleanup. Closing the *last* window is where the interception bites: on macOS `CloseRequested` is prevented and the window hidden, so sessions keep streaming; the Dock icon (`RunEvent::Reopen`), `Window ▸ Staged`, or a quit arriving while hidden brings it back — `show_a_window` prefers `main` for its restored geometry but recovers any surviving `win-N` peer. Hiding also drops that window's `tauri-{label}` PR-poll client to its unfocused tier, which a hidden window's missing webview blur would not. Other platforms have no Dock or tray to recover a hidden window, so closing the last window still quits there — now through the confirmation gate. - A custom Quit menu item makes Cmd+Q routable, so `request_quit` can gate it on active sessions and raise a confirmation dialog; confirming cancels each session with `CompletionReason::AppQuit` (the cancel is what runs the ACP child's graceful stop), waits for sessions and actions inside one shared 2s budget, sweeps any still-active rows to cancelled/app_quit, then exits. Queued sessions count as active; running actions are reported but don't gate. Quit and `Window ▸ Staged` route through the shared `dispatch_menu_event` router as focus-independent backend actions — every window being hidden is exactly when they matter. The dialog is addressed to exactly one window (`emit_to` plus a window-scoped frontend listener, the same pattern as menu routing): where the user is, or a window revealed for the purpose. A broadcast would raise one dialog per window, each unaware of the others' answers. The pending-prompt flag remembers its host window and is cleared when that window is destroyed, so the force-on-second- request escape hatch can't fire with no dialog on screen. `RunEvent::Exit` now runs the same idempotent cleanup, which is the only hook on the terminate: path — Dock ▸ Quit and logout stop sessions and actions instead of orphaning them. Ownership is checked against `owner_pid`, so a quit never prompts about or cancels another Staged instance's work. The quit commands are refused in the web-mode dispatch table: a browser client must not be able to terminate the desktop host. The store-incompatibility screens' Close buttons now quit rather than closing a window that would only hide. Phase 4 of the plan (routing Dock ▸ Quit through the prompt via a runtime `applicationShouldTerminate:`) is deliberately left out: the shared cleanup already prevents the process and data damage there, only the prompt is missing. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/app_lifecycle.rs | 646 ++++++++++++++++++ apps/staged/src-tauri/src/lib.rs | 133 ++-- .../staged/src-tauri/src/pr_poll_scheduler.rs | 15 + apps/staged/src-tauri/src/session_commands.rs | 2 +- apps/staged/src-tauri/src/session_runner.rs | 67 ++ apps/staged/src-tauri/src/web_server.rs | 10 + apps/staged/src/App.svelte | 14 +- apps/staged/src/lib/commands.ts | 27 + .../lifecycle/QuitConfirmDialog.svelte | 62 ++ .../features/lifecycle/quitPromptCopy.test.ts | 63 ++ .../lib/features/lifecycle/quitPromptCopy.ts | 60 ++ .../lib/features/projects/ProjectHome.svelte | 5 +- .../src/lib/listeners/quitListener.test.ts | 137 ++++ apps/staged/src/lib/listeners/quitListener.ts | 24 + .../src/lib/stores/quitPrompt.svelte.ts | 61 ++ apps/staged/src/lib/types.ts | 12 + 16 files changed, 1291 insertions(+), 47 deletions(-) create mode 100644 apps/staged/src-tauri/src/app_lifecycle.rs create mode 100644 apps/staged/src/lib/features/lifecycle/QuitConfirmDialog.svelte create mode 100644 apps/staged/src/lib/features/lifecycle/quitPromptCopy.test.ts create mode 100644 apps/staged/src/lib/features/lifecycle/quitPromptCopy.ts create mode 100644 apps/staged/src/lib/listeners/quitListener.test.ts create mode 100644 apps/staged/src/lib/listeners/quitListener.ts create mode 100644 apps/staged/src/lib/stores/quitPrompt.svelte.ts diff --git a/apps/staged/src-tauri/src/app_lifecycle.rs b/apps/staged/src-tauri/src/app_lifecycle.rs new file mode 100644 index 000000000..aae9e7916 --- /dev/null +++ b/apps/staged/src-tauri/src/app_lifecycle.rs @@ -0,0 +1,646 @@ +//! Window close, the quit gate, and shutdown cleanup. +//! +//! Staged's work outlives its windows: agent sessions and long-running actions +//! are child processes this process owns. Two rules follow from that, and this +//! module owns both. +//! +//! **Closing a window is not quitting.** With peer windows still open, a close +//! is just a close — the process lives on in the others, so the window is +//! destroyed normally (`window_commands` owns that cleanup). Closing the *last* +//! window is where the rules bite: on macOS `CloseRequested` is prevented and +//! the window hidden, so sessions keep streaming; the Dock icon +//! (`RunEvent::Reopen`) or `Window ▸ Staged` brings it back. Other platforms +//! have no Dock/tray to recover a hidden window, so closing the last window +//! still quits there — but through the same confirmation gate as `Cmd+Q`. +//! +//! **Quitting with sessions running asks first, then stops them cleanly.** +//! [`request_quit`] gates on active sessions and hands the decision to the +//! frontend dialog, addressed to a single live window (revealed first if every +//! window is hidden); [`shutdown_cleanup`] cancels sessions with +//! [`CompletionReason::AppQuit`] and stops actions. That cancel is the only +//! thing that shuts an agent down: ACP children are spawned with +//! `process_group(0)` and `kill_on_drop`, and `process::exit` runs no +//! destructors, so a bare exit leaves the agent CLIs running. +//! +//! Every exit path funnels into [`shutdown_cleanup`], which runs its work at +//! most once — a confirmed quit calls it directly, `RunEvent::ExitRequested` +//! covers programmatic exits, and `RunEvent::Exit` is the only hook on the +//! `NSApp terminate:` path (Dock ▸ Quit, logout), which never emits +//! `ExitRequested`. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use serde::Serialize; +use tauri::{AppHandle, Emitter, Manager, WebviewWindow, Window, WindowEvent}; + +use crate::actions; +use crate::session_commands::{self, ActiveSessionInfo}; +use crate::session_runner::SessionRegistry; +use crate::store::{CompletionReason, Session, SessionStatus, Store}; + +/// Event that raises the frontend's quit confirmation dialog. +const QUIT_REQUESTED_EVENT: &str = "app:quit-requested"; + +/// Menu id of the app-menu Quit item. Custom rather than +/// `PredefinedMenuItem::quit` so `Cmd+Q` is routable at all: the predefined item +/// maps to `NSApp terminate:`, which reaches no Tauri hook that can gate it. +pub(crate) const QUIT_MENU_ID: &str = "quit"; + +/// Menu id of `Window ▸ Staged`. The recovery path for `Cmd+Tab`-ing to an app +/// whose windows are all hidden — macOS sends no reopen event for that. +pub(crate) const SHOW_WINDOW_MENU_ID: &str = "show_window"; + +/// Label of the cold-start window (the `tauri.conf.json` entry). Secondary +/// windows are `win-N` peers — see `window_commands` — with nothing privileged +/// about `main` beyond being the one whose geometry is restored, which makes it +/// the nicest default to reveal. +const MAIN_WINDOW_LABEL: &str = "main"; + +/// Total budget for stopping sessions and actions. Sessions and actions are +/// signalled first and waited on against this one deadline, because the +/// `RunEvent::Exit` path runs inside `applicationWillTerminate:`, where the OS +/// gives us limited time before killing the process outright. +const SHUTDOWN_BUDGET: Duration = Duration::from_secs(2); + +/// Grace period before an action's process group is escalated to `SIGKILL`. +const ACTION_FORCE_KILL_AFTER: Duration = Duration::from_secs(1); + +/// Quit bookkeeping, managed as Tauri state. +#[derive(Default)] +pub struct QuitState { + /// Set by the first caller into [`shutdown_cleanup`], so the cleanup runs + /// exactly once however many exit events follow it. + quit_in_progress: AtomicBool, + /// Label of the window showing an unanswered confirmation dialog. A quit + /// request arriving while it is set forces the quit — a wedged webview must + /// never be able to trap the app, so a second `Cmd+Q` always gets out. The + /// label is what lets a destroyed host window clear the flag instead of + /// leaving that force path armed with no dialog on screen. + prompt_host: Mutex>, +} + +impl QuitState { + fn set_prompt_host(&self, label: &str) { + *self.prompt_host.lock().unwrap() = Some(label.to_string()); + } + + /// Clear any pending prompt, returning whether one was pending. + fn take_prompt(&self) -> bool { + self.prompt_host.lock().unwrap().take().is_some() + } + + /// Clear the pending prompt if `label` was hosting it. + fn clear_prompt_if_host(&self, label: &str) { + let mut host = self.prompt_host.lock().unwrap(); + if host.as_deref() == Some(label) { + *host = None; + } + } +} + +/// What a quit would interrupt, as sent to the confirmation dialog. +#[derive(Debug, Clone, Default, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct QuitBlockers { + /// Running and queued sessions owned by this process — the only thing that + /// gates a quit. + pub sessions: Vec, + /// Running actions. Reported so the dialog can say they stop too, but they + /// don't gate the quit on their own: a dev server left running is the normal + /// state of a workspace, and blocking `Cmd+Q` on it would be noise. + pub running_action_count: usize, +} + +/// Whether a quit should stop and ask first. +/// +/// Queued sessions count: they're work the user asked for that a quit silently +/// drops, so they belong in the prompt. +fn should_prompt(blockers: &QuitBlockers) -> bool { + !blockers.sessions.is_empty() +} + +// ============================================================================= +// Window events +// ============================================================================= + +/// `Builder::on_window_event` hook — see the module docs for why closing the +/// last window doesn't end the process. +pub fn on_window_event(window: &Window, event: &WindowEvent) { + match event { + WindowEvent::CloseRequested { api, .. } => on_close_requested(window, api), + // A destroyed window takes its webview — and any dialog in it — with + // it. Left set, the pending flag would turn the next quit request into + // a silent force-quit; cleared, that quit just asks again. + WindowEvent::Destroyed => { + if let Some(quit_state) = window.app_handle().try_state::() { + quit_state.clear_prompt_if_host(window.label()); + } + } + _ => {} + } +} + +fn on_close_requested(window: &Window, api: &tauri::CloseRequestApi) { + let app = window.app_handle(); + + // Mid-shutdown, closes are the exit tearing windows down — stay out of the + // way. + if let Some(quit_state) = app.try_state::() { + if quit_state.quit_in_progress.load(Ordering::SeqCst) { + return; + } + } + + // With peer windows still live (visible or hidden), a close is just a + // close: sessions belong to the process, not this window. The `Destroyed` + // hook in lib.rs does the per-window cleanup. + if app.webview_windows().len() > 1 { + return; + } + + // Last window: the window-state plugin has its own `CloseRequested` handler + // and saves geometry there, so preventing the close still persists the + // window's position and size. + api.prevent_close(); + + #[cfg(target_os = "macos")] + hide_window(window); + + // No Dock or tray icon elsewhere, so a hidden window would be unreachable — + // closing the last window still quits, with the confirmation gate in front + // of it. + #[cfg(not(target_os = "macos"))] + request_quit(app, false); +} + +/// Hide the window and drop its PR-poll client to the unfocused tier. +/// +/// `prPollingService` derives focus from `document.hasFocus()` and the webview's +/// focus events, and hiding the native window does not reliably deliver a blur +/// to the webview — so tell the scheduler directly instead of leaving it polling +/// on behalf of a window nobody can see. +#[cfg(target_os = "macos")] +fn hide_window(window: &Window) { + if let Err(e) = window.hide() { + log::warn!("Failed to hide window on close: {e}"); + return; + } + set_native_focus(window.app_handle(), window.label(), false); +} + +/// Bring a window back on screen: the Dock-icon click, `Window ▸ Staged`, and a +/// quit arriving with no visible window all funnel here. +pub fn show_a_window(app: &AppHandle) { + if reveal_a_window(app).is_none() { + log::warn!("No window left to show"); + } +} + +/// Pick a window and make sure it is on screen and focused, returning it. +/// +/// Prefers where the user already is (focused, then visible — reachable when a +/// quit request arrives from the store-incompatibility screen or the web +/// dispatch refusal path while windows are up), then falls back to unhiding one: +/// `main` for its restored geometry, else any. `None` only if every window has +/// been destroyed, which no close path produces — closing the last window hides +/// it instead. +fn reveal_a_window(app: &AppHandle) -> Option { + let windows = app.webview_windows(); + let window = windows + .values() + .find(|window| window.is_focused().unwrap_or(false)) + .or_else(|| { + windows + .values() + .find(|window| window.is_visible().unwrap_or(false)) + }) + .or_else(|| windows.get(MAIN_WINDOW_LABEL)) + .or_else(|| windows.values().next())?; + + if let Err(e) = window.show() { + log::warn!("Failed to show window: {e}"); + } + if let Err(e) = window.unminimize() { + log::warn!("Failed to unminimize window: {e}"); + } + if let Err(e) = window.set_focus() { + log::warn!("Failed to focus window: {e}"); + } + set_native_focus(app, window.label(), true); + Some(window.clone()) +} + +/// Mirror a native window's visibility onto its PR-poll client's focus hint. +/// Paired with the webview's own focus events, which report the same value once +/// the window is back on screen. +fn set_native_focus(app: &AppHandle, window_label: &str, focused: bool) { + if let Some(scheduler) = app.try_state::>() { + crate::pr_poll_scheduler::set_tauri_client_focus(&scheduler, window_label, focused); + } +} + +// ============================================================================= +// Quit gate +// ============================================================================= + +/// Handle a quit request from the app menu, `Cmd+Q`, or (off macOS) the last +/// window's close. Cheap enough for the main thread: it snapshots blockers and +/// either hands off to a background quit or raises the dialog. +pub fn request_quit(app: &AppHandle, force: bool) { + let quit_state = app.state::(); + + // A quit arriving while the dialog is unanswered (a second `Cmd+Q`) is the + // escape hatch from a webview that never rendered or answered it. + if force || quit_state.take_prompt() { + spawn_quit(app); + return; + } + + let blockers = collect_quit_blockers(app); + if !should_prompt(&blockers) { + spawn_quit(app); + return; + } + + // The dialog goes to exactly one window — where the user is, or a window + // revealed for the purpose if the quit arrived with everything hidden. A + // broadcast would raise one dialog per window, each unaware of the others' + // answers. No window at all means nobody to ask, so the quit proceeds. + let Some(host) = reveal_a_window(app) else { + spawn_quit(app); + return; + }; + quit_state.set_prompt_host(host.label()); + + if let Err(e) = app.emit_to(host.label(), QUIT_REQUESTED_EVENT, &blockers) { + log::warn!("Failed to ask for quit confirmation, quitting anyway: {e}"); + quit_state.take_prompt(); + spawn_quit(app); + } +} + +/// Quit from the UI, through the same gate as `Cmd+Q`. +/// +/// Used by the store-incompatibility screens' "Close" button, which has to end +/// the app: closing the last window only hides it, and those screens have no +/// working database behind them to come back to. +#[tauri::command] +pub fn quit_app(app_handle: AppHandle) { + request_quit(&app_handle, false); +} + +/// Quit confirmed in the dialog: stop sessions and actions, then exit. +/// +/// Deliberately absent from the web-mode `dispatch` table — a browser client +/// must not be able to terminate the desktop host. +#[tauri::command] +pub fn confirm_quit(app_handle: AppHandle) { + app_handle.state::().take_prompt(); + spawn_quit(&app_handle); +} + +/// Quit declined in the dialog: sessions keep running. +#[tauri::command] +pub fn cancel_quit(app_handle: AppHandle) { + app_handle.state::().take_prompt(); +} + +/// Run the quit sequence off the main thread so the bounded waits never freeze +/// the event loop — the dialog stays interactive and can render its +/// "Stopping sessions…" state while agents shut down. +fn spawn_quit(app: &AppHandle) { + let app = app.clone(); + std::thread::spawn(move || { + shutdown_cleanup(&app); + app.exit(0); + }); +} + +/// Snapshot what a quit would interrupt. +fn collect_quit_blockers(app: &AppHandle) -> QuitBlockers { + let sessions = match app_store(app) { + Some(store) => owned_active_sessions(&store) + .iter() + .map(|session| session_commands::project_active_session(&store, session)) + .collect(), + None => Vec::new(), + }; + + let running_action_count = match ( + app.try_state::>(), + app.try_state::>(), + ) { + (Some(executor), Some(registry)) => { + actions::commands::get_all_running_actions_impl(&executor, ®istry) + .map(|running| running.len()) + .unwrap_or(0) + } + _ => 0, + }; + + QuitBlockers { + sessions, + running_action_count, + } +} + +// ============================================================================= +// Shutdown cleanup +// ============================================================================= + +/// Stop everything this process owns. Idempotent — the first caller does the +/// work, later ones return immediately. +pub fn shutdown_cleanup(app: &AppHandle) { + let Some(quit_state) = app.try_state::() else { + return; + }; + if quit_state.quit_in_progress.swap(true, Ordering::SeqCst) { + return; + } + + // Signal both kinds of work before waiting on either, so they shut down in + // parallel inside one shared budget instead of one after the other. + let session_ids = cancel_owned_sessions(app); + let execution_ids = stop_running_actions(app); + + let deadline = Instant::now() + SHUTDOWN_BUDGET; + if !session_ids.is_empty() && !wait_for_sessions(app, &session_ids, deadline) { + log::warn!( + "Timed out waiting for {} session(s) to stop during app shutdown", + session_ids.len() + ); + } + if !execution_ids.is_empty() && !wait_for_actions(app, &execution_ids, deadline) { + log::warn!( + "Timed out waiting for {} action(s) to stop during app shutdown", + execution_ids.len() + ); + } + + // Last, so the rows reflect whatever the session threads managed to write + // for themselves first. + sweep_active_sessions(app); +} + +/// Cancel every session this process is running, recording `AppQuit` as the +/// reason. Returns the ids that were signalled. +fn cancel_owned_sessions(app: &AppHandle) -> Vec { + let Some(registry) = app.try_state::>() else { + return Vec::new(); + }; + + let session_ids = registry.running_session_ids(); + for session_id in &session_ids { + registry.cancel_with_completion_reason(session_id, CompletionReason::AppQuit); + } + session_ids +} + +/// Send every running action's process group a hangup, escalating to `SIGKILL` +/// after a grace period. Returns the execution ids that were signalled. +fn stop_running_actions(app: &AppHandle) -> Vec { + let (Some(executor), Some(registry)) = ( + app.try_state::>(), + app.try_state::>(), + ) else { + return Vec::new(); + }; + + actions::commands::stop_all_actions( + &executor, + ®istry, + actions::StopOptions { + force_kill_after: Some(ACTION_FORCE_KILL_AFTER), + }, + ) +} + +fn wait_for_sessions(app: &AppHandle, session_ids: &[String], deadline: Instant) -> bool { + let Some(registry) = app.try_state::>() else { + return true; + }; + registry.wait_for_sessions(session_ids, remaining_until(deadline)) +} + +fn wait_for_actions(app: &AppHandle, execution_ids: &[String], deadline: Instant) -> bool { + let Some(executor) = app.try_state::>() else { + return true; + }; + executor.wait_for_executions(execution_ids, remaining_until(deadline)) +} + +fn remaining_until(deadline: Instant) -> Duration { + deadline.saturating_duration_since(Instant::now()) +} + +/// Mark whatever is still active in the DB as cancelled by the quit. +/// +/// Covers sessions whose thread didn't finish its own terminal write inside the +/// budget, plus queued sessions that never started. Without this the next launch +/// finds them owned by a dead process and reports them as errored sessions. +fn sweep_active_sessions(app: &AppHandle) { + let Some(store) = app_store(app) else { + return; + }; + + let swept = owned_active_sessions(&store) + .iter() + .filter(|session| { + // Guarded CAS per row: a session thread that wrote its own terminal + // status while we were waiting keeps that status. + store + .transition_from_active( + &session.id, + SessionStatus::Cancelled, + None, + Some(&CompletionReason::AppQuit), + ) + .unwrap_or_else(|e| { + log::warn!("Failed to cancel session {} on quit: {e}", session.id); + false + }) + }) + .count(); + + if swept > 0 { + log::info!("Marked {swept} session(s) cancelled (app_quit) during shutdown"); + } +} + +/// Running and queued sessions **this process owns**. +/// +/// The store is shared with any other Staged instance pointed at the same data +/// dir — that's what `owner_pid` is for — so a quit must neither prompt about +/// nor cancel another instance's work. Queued rows carry no owner yet, so they +/// count as ours: claiming one (`transition_queued_to_running`) stamps a pid +/// atomically, which is what takes another instance's claim out of this set. +fn owned_active_sessions(store: &Store) -> Vec { + let sessions = match store.get_active_sessions() { + Ok(sessions) => sessions, + Err(e) => { + log::warn!("Failed to query active sessions during quit: {e}"); + return Vec::new(); + } + }; + + sessions + .into_iter() + .filter(|session| { + session.status == SessionStatus::Queued || session.owner_pid == Some(std::process::id()) + }) + .collect() +} + +fn app_store(app: &AppHandle) -> Option> { + app.try_state::>>>() + .and_then(|slot| slot.lock().unwrap().clone()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::Path; + + fn active_session(status: &str) -> ActiveSessionInfo { + ActiveSessionInfo { + session_id: "s1".to_string(), + project_id: None, + branch_id: None, + session_type: None, + status: status.to_string(), + } + } + + #[test] + fn running_sessions_prompt() { + let blockers = QuitBlockers { + sessions: vec![active_session("running")], + running_action_count: 0, + }; + assert!(should_prompt(&blockers)); + } + + #[test] + fn queued_sessions_prompt() { + let blockers = QuitBlockers { + sessions: vec![active_session("queued")], + running_action_count: 0, + }; + assert!(should_prompt(&blockers)); + } + + #[test] + fn running_actions_alone_do_not_prompt() { + let blockers = QuitBlockers { + sessions: Vec::new(), + running_action_count: 3, + }; + assert!(!should_prompt(&blockers)); + } + + #[test] + fn nothing_active_does_not_prompt() { + assert!(!should_prompt(&QuitBlockers::default())); + } + + /// The pending flag turns the next quit into a force-quit, so it must not + /// outlive the window whose dialog it stands for — but a *peer* window + /// closing must not answer a dialog it isn't showing. + #[test] + fn prompt_clears_only_when_its_host_window_is_destroyed() { + let state = QuitState::default(); + + state.set_prompt_host("win-2"); + state.clear_prompt_if_host("main"); + assert!(state.take_prompt(), "peer destruction dropped the prompt"); + + state.set_prompt_host("win-2"); + state.clear_prompt_if_host("win-2"); + assert!( + !state.take_prompt(), + "host destruction left the prompt armed" + ); + } + + #[test] + fn owned_active_sessions_skips_other_instances_running_sessions() { + let store = Store::in_memory().unwrap(); + + let ours = Session::new_running("ours", Path::new("/tmp")); + store.create_session(&ours).unwrap(); + let queued = Session::new_queued("queued"); + store.create_session(&queued).unwrap(); + let mut theirs = Session::new_running("theirs", Path::new("/tmp")); + theirs.owner_pid = Some(std::process::id().wrapping_add(1)); + store.create_session(&theirs).unwrap(); + + let owned = owned_active_sessions(&store); + assert_eq!(owned.len(), 2); + assert!(owned.iter().any(|session| session.id == ours.id)); + assert!(owned.iter().any(|session| session.id == queued.id)); + } + + /// The DB sweep is what keeps the next launch from reporting these sessions + /// as errors recovered from a dead process. + #[test] + fn sweep_cancels_running_and_queued_sessions() { + let store = Store::in_memory().unwrap(); + + let running = Session::new_running("running", Path::new("/tmp")); + store.create_session(&running).unwrap(); + let queued = Session::new_queued("queued"); + store.create_session(&queued).unwrap(); + + for session in owned_active_sessions(&store) { + assert!(store + .transition_from_active( + &session.id, + SessionStatus::Cancelled, + None, + Some(&CompletionReason::AppQuit), + ) + .unwrap()); + } + + for id in [&running.id, &queued.id] { + let session = store.get_session(id).unwrap().unwrap(); + assert_eq!(session.status, SessionStatus::Cancelled); + assert_eq!(session.completion_reason, Some(CompletionReason::AppQuit)); + } + } + + #[test] + fn sweep_leaves_terminal_sessions_alone() { + let store = Store::in_memory().unwrap(); + + let completed = Session::new_running("completed", Path::new("/tmp")); + store.create_session(&completed).unwrap(); + store + .update_session_status( + &completed.id, + SessionStatus::Completed, + None, + Some(&CompletionReason::TurnComplete), + ) + .unwrap(); + + assert!(owned_active_sessions(&store).is_empty()); + assert!(!store + .transition_from_active( + &completed.id, + SessionStatus::Cancelled, + None, + Some(&CompletionReason::AppQuit), + ) + .unwrap()); + + let session = store.get_session(&completed.id).unwrap().unwrap(); + assert_eq!(session.status, SessionStatus::Completed); + assert_eq!( + session.completion_reason, + Some(CompletionReason::TurnComplete) + ); + } +} diff --git a/apps/staged/src-tauri/src/lib.rs b/apps/staged/src-tauri/src/lib.rs index a88c56dfa..5cc4db542 100644 --- a/apps/staged/src-tauri/src/lib.rs +++ b/apps/staged/src-tauri/src/lib.rs @@ -8,6 +8,7 @@ pub mod acp_tools; pub mod acp_tools_reconciler; pub mod actions; pub mod agent; +pub mod app_lifecycle; pub mod background_sync; pub mod blox; pub mod branches; @@ -48,9 +49,7 @@ pub mod test_utils; use serde::Serialize; use std::path::PathBuf; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; -use std::time::Duration; use store::Store; use tauri::{Emitter, Manager}; @@ -67,11 +66,6 @@ struct DbState { needs_reset: Mutex>, } -#[derive(Default)] -struct ShutdownState { - quit_in_progress: AtomicBool, -} - pub(crate) fn preferences_store_path_buf() -> Option { crate::paths::data_dir().map(|d| d.join("preferences.json")) } @@ -266,29 +260,6 @@ pub(crate) fn get_store( .ok_or_else(|| "Database not initialized — please reset from the startup prompt".into()) } -fn stop_actions_for_app_shutdown(app_handle: &tauri::AppHandle) { - let executor = app_handle.state::>(); - let registry = app_handle.state::>(); - let stopped_execution_ids = actions::commands::stop_all_actions( - &executor, - ®istry, - actions::StopOptions { - force_kill_after: Some(Duration::from_secs(1)), - }, - ); - - if stopped_execution_ids.is_empty() { - return; - } - - if !executor.wait_for_executions(&stopped_execution_ids, Duration::from_secs(2)) { - log::warn!( - "Timed out waiting for {} action(s) to stop during app shutdown", - stopped_execution_ids.len() - ); - } -} - fn start_store_services( store: Arc, pr_scheduler: Arc, @@ -1786,6 +1757,10 @@ enum MenuDispatch { EmitToFocused(&'static str), /// Create a window here in the backend, with no project seed. OpenWindowUnseeded, + /// Run the quit gate in the backend (`app_lifecycle::request_quit`). + RequestQuit, + /// Reveal a window in the backend (`app_lifecycle::show_a_window`). + ShowWindow, /// Nothing to do — unknown item, or a window-scoped item with no target. Drop, } @@ -1804,6 +1779,15 @@ enum MenuDispatch { /// can just create it. That also un-strands the other items: the new window is /// focused, so Settings/Find/zoom route normally again. fn dispatch_menu_event(id: &str, has_focused_window: bool) -> MenuDispatch { + // Lifecycle items are app-scoped and handled in the backend, focus or no + // focus — every window being hidden is exactly when `Window ▸ Staged` and a + // gateable `Cmd+Q` matter most. + match id { + app_lifecycle::QUIT_MENU_ID => return MenuDispatch::RequestQuit, + app_lifecycle::SHOW_WINDOW_MENU_ID => return MenuDispatch::ShowWindow, + _ => {} + } + let event_name = match id { "new_window" => "menu:new-window", "settings" => "menu:settings", @@ -1947,6 +1931,26 @@ pub fn run() { true, Some("CmdOrCtrl+0"), )?; + // Custom rather than `PredefinedMenuItem::quit`: that one maps + // straight to `NSApp terminate:`, which reaches no Tauri hook, + // so Cmd+Q could never be gated on running sessions. + let quit_item = MenuItem::with_id( + handle, + app_lifecycle::QUIT_MENU_ID, + "Quit Staged", + true, + Some("CmdOrCtrl+Q"), + )?; + // Recovery path for an app whose windows are all hidden: Cmd+Tab + // sends no reopen event, so without this the app looks dead (the + // same reason Slack exposes `Window ▸ Slack`). + let show_window_item = MenuItem::with_id( + handle, + app_lifecycle::SHOW_WINDOW_MENU_ID, + "Staged", + true, + None::<&str>, + )?; let app_menu = Submenu::with_items( handle, @@ -1966,7 +1970,7 @@ pub fn run() { &PredefinedMenuItem::hide(handle, None)?, &PredefinedMenuItem::hide_others(handle, None)?, &PredefinedMenuItem::separator(handle)?, - &PredefinedMenuItem::quit(handle, Some("Quit Staged"))?, + &quit_item, ], )?; @@ -2024,6 +2028,8 @@ pub fn run() { &PredefinedMenuItem::maximize(handle, None)?, &PredefinedMenuItem::separator(handle)?, &PredefinedMenuItem::close_window(handle, None)?, + &PredefinedMenuItem::separator(handle)?, + &show_window_item, ], )?; @@ -2161,7 +2167,7 @@ pub fn run() { app.manage(window_commands::UpdaterWindowState::default()); app.manage(Arc::new(actions::ActionExecutor::new())); app.manage(Arc::new(actions::ActionRegistry::new())); - app.manage(ShutdownState::default()); + app.manage(app_lifecycle::QuitState::default()); app.manage(DbState { db_path, needs_reset: Mutex::new(reset_info), @@ -2221,10 +2227,16 @@ pub fn run() { log::warn!("Failed to open window from menu: {e}"); } } + MenuDispatch::RequestQuit => app_lifecycle::request_quit(app, false), + MenuDispatch::ShowWindow => app_lifecycle::show_a_window(app), MenuDispatch::Drop => {} } }) .on_window_event(|window, event| { + // Close-to-hide / the quit gate (`CloseRequested`), and dropping a + // pending quit prompt whose host window went away (`Destroyed`). + app_lifecycle::on_window_event(window, event); + if let tauri::WindowEvent::Destroyed = event { // Native windows have no WS heartbeat and their PR-poll client // ids are exempt from TTL eviction, so a closed window must @@ -2260,6 +2272,11 @@ pub fn run() { window_commands::new_window, window_commands::take_window_seed, window_commands::claim_updater_ownership, + // Lifecycle — desktop only; the web-mode `dispatch` table refuses + // these so a browser client can't quit the host. + app_lifecycle::quit_app, + app_lifecycle::confirm_quit, + app_lifecycle::cancel_quit, list_projects, create_project, list_project_repos, @@ -2462,17 +2479,30 @@ pub fn run() { ]) .build(tauri::generate_context!()) .expect("error while building tauri application") - .run(|app_handle, event| { - if let tauri::RunEvent::ExitRequested { api, .. } = event { - let shutdown = app_handle.state::(); - if shutdown.quit_in_progress.swap(true, Ordering::SeqCst) { - return; - } - - api.prevent_exit(); - stop_actions_for_app_shutdown(app_handle); - app_handle.exit(0); + .run(|app_handle, event| match event { + // Now that window close is intercepted, the only producers are our + // own confirmed quit (which has already cleaned up) and the updater's + // relaunch — which ignores `prevent_exit` anyway, so nothing here + // tries to hold the exit back. + tauri::RunEvent::ExitRequested { .. } => { + app_lifecycle::shutdown_cleanup(app_handle); + } + // The only hook on the `NSApp terminate:` path (Dock ▸ Quit, logout), + // which never emits `ExitRequested`. Without it those quits orphan + // the agent and action child processes. + tauri::RunEvent::Exit => { + app_lifecycle::shutdown_cleanup(app_handle); + } + // Dock-icon click or `open -a Staged` on an app whose windows are + // all hidden. + #[cfg(target_os = "macos")] + tauri::RunEvent::Reopen { + has_visible_windows: false, + .. + } => { + app_lifecycle::show_a_window(app_handle); } + _ => {} }); } @@ -2608,12 +2638,29 @@ mod tests { #[test] fn unknown_menu_events_drop_regardless_of_focus() { - for id in ["", "quit", "menu:new-window", "New Window"] { + for id in ["", "menu:new-window", "New Window"] { assert_eq!(dispatch_menu_event(id, true), MenuDispatch::Drop); assert_eq!(dispatch_menu_event(id, false), MenuDispatch::Drop); } } + /// The lifecycle items must route with no window focused: every window + /// being hidden is exactly when `Window ▸ Staged` and a gateable `Cmd+Q` + /// matter most. + #[test] + fn lifecycle_menu_events_route_to_the_backend_regardless_of_focus() { + for focused in [true, false] { + assert_eq!( + dispatch_menu_event(crate::app_lifecycle::QUIT_MENU_ID, focused), + MenuDispatch::RequestQuit + ); + assert_eq!( + dispatch_menu_event(crate::app_lifecycle::SHOW_WINDOW_MENU_ID, focused), + MenuDispatch::ShowWindow + ); + } + } + fn remote_branch( project_id: &str, id: &str, diff --git a/apps/staged/src-tauri/src/pr_poll_scheduler.rs b/apps/staged/src-tauri/src/pr_poll_scheduler.rs index 893e8247f..41c295e97 100644 --- a/apps/staged/src-tauri/src/pr_poll_scheduler.rs +++ b/apps/staged/src-tauri/src/pr_poll_scheduler.rs @@ -635,6 +635,21 @@ pub fn set_foreground_project( scheduler.set_foreground(client_id, project_id); } +/// Report a native window's focus from the backend, bypassing the frontend. +/// +/// `app_lifecycle` hides and shows windows itself, and a hidden native window +/// does not reliably deliver a blur to its webview — so without this the +/// scheduler would keep polling on the focused tier for a window nobody can +/// see. The id mirrors the frontend's own `tauri-{label}` scheme, so both sides +/// address the same per-window client. +pub(crate) fn set_tauri_client_focus( + scheduler: &PrPollScheduler, + window_label: &str, + focused: bool, +) { + scheduler.set_focus(format!("{TAURI_CLIENT_PREFIX}{window_label}"), focused); +} + /// Report a client's window focus. With no client focused, periodic polling /// pauses (an explicit `refresh_now` still fetches). #[tauri::command(rename_all = "camelCase")] diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index cc534ba44..5002e064b 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -1002,7 +1002,7 @@ pub struct ActiveSessionInfo { /// sessions (pr/push) link no artifact, so their branch comes from the /// session row's own `branch_id` and their type falls back to prompt /// inference. -fn project_active_session(store: &Store, session: &store::Session) -> ActiveSessionInfo { +pub(crate) fn project_active_session(store: &Store, session: &store::Session) -> ActiveSessionInfo { let project_note = store .get_project_note_by_session(&session.id) .ok() diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 052b911e1..18751cc28 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -427,6 +427,44 @@ impl SessionRegistry { .unwrap_or_default() } + /// Ids of every session this process is currently running. + /// + /// The shutdown path uses this to cancel them all: the registry, not the DB, + /// is what says which running rows belong to *this* process's threads. + pub fn running_session_ids(&self) -> Vec { + self.inner.lock().unwrap().running.keys().cloned().collect() + } + + /// Wait until none of `session_ids` are registered as running, or until + /// `timeout` elapses. Returns `true` if they all deregistered in time. + /// + /// Modelled on `ActionExecutor::wait_for_executions`: session threads + /// deregister themselves as they exit, so polling the registry is how the + /// shutdown path learns a cancelled session's agent is actually gone rather + /// than exiting out from under it. + pub fn wait_for_sessions(&self, session_ids: &[String], timeout: Duration) -> bool { + let deadline = std::time::Instant::now() + timeout; + + loop { + let all_stopped = { + let inner = self.inner.lock().unwrap(); + session_ids + .iter() + .all(|session_id| !inner.running.contains_key(session_id)) + }; + + if all_stopped { + return true; + } + + if std::time::Instant::now() >= deadline { + return false; + } + + std::thread::sleep(Duration::from_millis(25)); + } + } + /// Register a session whose work is driven outside `start_session` (e.g. a /// pikchr diagram child session run by a `generate_pikchr` worker thread), /// so a user cancel reaches the actual work instead of taking @@ -4156,6 +4194,35 @@ mod tests { assert_eq!(failed.completion_reason, Some(CompletionReason::Crashed)); } + #[test] + fn wait_for_sessions_returns_once_every_session_deregisters() { + let registry = Arc::new(SessionRegistry::new()); + registry.register("session-1"); + registry.register("session-2"); + let session_ids = registry.running_session_ids(); + assert_eq!(session_ids.len(), 2); + + let deregistering = Arc::clone(®istry); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(50)); + deregistering.deregister("session-1"); + deregistering.deregister("session-2"); + }); + + assert!(registry.wait_for_sessions(&session_ids, Duration::from_secs(2))); + assert!(registry.running_session_ids().is_empty()); + } + + #[test] + fn wait_for_sessions_times_out_while_a_session_is_still_running() { + let registry = SessionRegistry::new(); + registry.register("session-1"); + + assert!(!registry.wait_for_sessions(&["session-1".to_string()], Duration::from_millis(50))); + // Unknown ids count as stopped, so a stale snapshot can't block a quit. + assert!(registry.wait_for_sessions(&["gone".to_string()], Duration::from_millis(50))); + } + #[test] fn running_project_session_cancellation_records_completion_reason_override() { let registry = SessionRegistry::new(); diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index d6b0869c0..2daac5533 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -559,6 +559,16 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result { + Err(format!("{command} is not available in web mode")) + } + // ===================================================================== // Projects // ===================================================================== diff --git a/apps/staged/src/App.svelte b/apps/staged/src/App.svelte index 4d680c14d..3dd41cbe0 100644 --- a/apps/staged/src/App.svelte +++ b/apps/staged/src/App.svelte @@ -19,6 +19,7 @@ import ProjectsList from './lib/features/projects/ProjectsList.svelte'; import ProjectsSidebar from './lib/features/projects/ProjectsSidebar.svelte'; import ProjectDeleteDialog from './lib/features/projects/ProjectDeleteDialog.svelte'; + import QuitConfirmDialog from './lib/features/lifecycle/QuitConfirmDialog.svelte'; import ReposListView from './lib/features/projects/ReposListView.svelte'; import SessionLauncher from './lib/features/sessions/SessionLauncher.svelte'; import SettingsPage from './lib/features/settings/SettingsPage.svelte'; @@ -62,6 +63,7 @@ import { listenForPageLifecycle } from './lib/listeners/pageLifecycleListener'; import { listenForAcpToolsReconciled } from './lib/listeners/acpToolsListener'; import { listenForMenuEvents } from './lib/listeners/menuListener'; + import { listenForQuitRequests } from './lib/listeners/quitListener'; import { darkMode } from './lib/stores/isDark.svelte'; import * as prPollingService from './lib/services/prPollingService'; import type { StoreIncompatibility } from './lib/types'; @@ -77,6 +79,7 @@ let unlistenAcpToolsReconciled: UnlistenFn | undefined; let unlistenStoreReset: UnlistenFn | undefined; let unlistenUpdaterOwnerAvailable: UnlistenFn | undefined; + let unlistenQuitRequests: UnlistenFn | undefined; let unregisterShortcuts: (() => void) | null = null; let stopUpdaterLoop: (() => void) | null = null; let updaterStartPending = false; @@ -356,6 +359,9 @@ // Refresh provider discovery (and any loaded doctor report) once the // backend finishes installing/upgrading the managed ACP bridges. unlistenAcpToolsReconciled = listenForAcpToolsReconciled(); + // Raise the quit confirmation when the backend gates a quit on running + // sessions (Tauri only — see quitListener.ts). + unlistenQuitRequests = listenForQuitRequests(); // Keep the shared project-list cache fresh for the app's lifetime — the // store dedupes, so starting before any view consumes it is safe. projectsDataStore.startListeners(); @@ -568,6 +574,7 @@ unlistenAcpToolsReconciled?.(); unlistenStoreReset?.(); unlistenUpdaterOwnerAvailable?.(); + unlistenQuitRequests?.(); projectsDataStore.stopListeners(); projectRunActionsStore.stopListening(); stopUpdaterLoop?.(); @@ -586,8 +593,10 @@ } } + // Quits rather than closing the window: closing the last window only hides + // it, and there is no usable app behind this screen to come back to. function handleClose() { - getWindowSync().close(); + void commands.quitApp().catch((e) => console.error('Failed to quit:', e)); } @@ -697,6 +706,9 @@ point (sidebar, landing grid, ProjectHome top bar/shortcut). --> + + + {/if} diff --git a/apps/staged/src/lib/commands.ts b/apps/staged/src/lib/commands.ts index 2f983052e..f94fc86a2 100644 --- a/apps/staged/src/lib/commands.ts +++ b/apps/staged/src/lib/commands.ts @@ -70,6 +70,33 @@ export function confirmResetStore(): Promise { return invokeCommand('confirm_reset_store'); } +// ============================================================================= +// App lifecycle +// ============================================================================= + +/** + * Quit Staged, through the same gate as `Cmd+Q` — it raises the confirmation + * dialog when sessions are still running. Closing the window only hides it, so + * UI that means "end the app" (the store-incompatibility screens) needs this. + */ +export function quitApp(): Promise { + return invokeCommand('quit_app'); +} + +/** + * Confirm the quit raised by `app:quit-requested`: the backend stops the active + * sessions and running actions, then exits. Desktop only — the command is not in + * the web-mode dispatch table, so a browser client cannot quit the host. + */ +export function confirmQuit(): Promise { + return invokeCommand('confirm_quit'); +} + +/** Decline the quit; sessions keep running. */ +export function cancelQuit(): Promise { + return invokeCommand('cancel_quit'); +} + // ============================================================================= // Projects // ============================================================================= diff --git a/apps/staged/src/lib/features/lifecycle/QuitConfirmDialog.svelte b/apps/staged/src/lib/features/lifecycle/QuitConfirmDialog.svelte new file mode 100644 index 000000000..db45d8fd9 --- /dev/null +++ b/apps/staged/src/lib/features/lifecycle/QuitConfirmDialog.svelte @@ -0,0 +1,62 @@ + + + + !v && quitPrompt.cancel()}> + + + Quit Staged? + {description} + + + Cancel + quitPrompt.confirm()} + > + {quitPrompt.stopping ? 'Stopping sessions…' : 'Quit & Stop Sessions'} + + + + diff --git a/apps/staged/src/lib/features/lifecycle/quitPromptCopy.test.ts b/apps/staged/src/lib/features/lifecycle/quitPromptCopy.test.ts new file mode 100644 index 000000000..0f1310cee --- /dev/null +++ b/apps/staged/src/lib/features/lifecycle/quitPromptCopy.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from 'vitest'; +import type { ActiveSessionInfo } from '../../types'; +import { quitPromptDescription, quitSessionLabel } from './quitPromptCopy'; + +function session(overrides: Partial = {}): ActiveSessionInfo { + return { + sessionId: 's1', + projectId: 'p1', + branchId: 'b1', + sessionType: 'review', + status: 'running', + ...overrides, + }; +} + +describe('quitSessionLabel', () => { + it('names the session type and where it runs', () => { + expect(quitSessionLabel(session(), 'fix-login')).toBe('review on fix-login'); + }); + + it('marks queued sessions', () => { + expect(quitSessionLabel(session({ status: 'queued' }), 'fix-login')).toBe( + 'review on fix-login (queued)' + ); + }); + + it('falls back to "session" for an unknown or missing type', () => { + expect(quitSessionLabel(session({ sessionType: null }), 'fix-login')).toBe( + 'session on fix-login' + ); + expect(quitSessionLabel(session({ sessionType: 'mystery' }), 'fix-login')).toBe( + 'session on fix-login' + ); + }); + + it('drops the location when there is none to show', () => { + expect(quitSessionLabel(session({ sessionType: 'note' }), null)).toBe('note'); + }); +}); + +describe('quitPromptDescription', () => { + it('reads singular for one session', () => { + expect(quitPromptDescription(['commit on fix-login'], 0)).toBe( + '1 session is still running. Quitting will stop it. commit on fix-login.' + ); + }); + + it('lists every session for a plural count', () => { + expect(quitPromptDescription(['commit on fix-login', 'note on docs'], 0)).toBe( + '2 sessions are still running. Quitting will stop them. commit on fix-login, note on docs.' + ); + }); + + it('mentions running actions only when there are some', () => { + expect(quitPromptDescription(['commit on fix-login'], 1)).toContain( + '1 running action will also stop.' + ); + expect(quitPromptDescription(['commit on fix-login'], 3)).toContain( + '3 running actions will also stop.' + ); + expect(quitPromptDescription(['commit on fix-login'], 0)).not.toContain('action'); + }); +}); diff --git a/apps/staged/src/lib/features/lifecycle/quitPromptCopy.ts b/apps/staged/src/lib/features/lifecycle/quitPromptCopy.ts new file mode 100644 index 000000000..508f3dbca --- /dev/null +++ b/apps/staged/src/lib/features/lifecycle/quitPromptCopy.ts @@ -0,0 +1,60 @@ +/** + * Copy for the quit confirmation dialog. + * + * Kept out of the component so the wording is unit-testable: the dialog resolves + * each session's branch/project names from the stores and passes labels in. + */ + +import type { ActiveSessionInfo } from '../../types'; + +/** How each session type reads in the dialog body. */ +const SESSION_TYPE_LABELS: Record = { + note: 'note', + commit: 'commit', + review: 'review', + pr: 'PR', + push: 'push', + pull: 'pull', +}; + +/** + * Label for one session that a quit would stop, e.g. `review on fix-login` or + * `commit on fix-login (queued)`. + * + * `where` is the branch name when the session belongs to one, otherwise the + * project name — project-level sessions (notes on a project) have no branch. + */ +export function quitSessionLabel(session: ActiveSessionInfo, where: string | null): string { + const kind = session.sessionType ? SESSION_TYPE_LABELS[session.sessionType] : null; + const base = where ? `${kind ?? 'session'} on ${where}` : (kind ?? 'session'); + return session.status === 'queued' ? `${base} (queued)` : base; +} + +/** + * Dialog body: how much stops, what it is, and whether actions go with it. + * + * Actions never gate the quit (see `should_prompt` in `app_lifecycle.rs`), so + * they are mentioned only as a consequence of one. + */ +export function quitPromptDescription(sessionLabels: string[], runningActionCount: number): string { + const count = sessionLabels.length; + const sentences = [ + count === 1 + ? '1 session is still running. Quitting will stop it.' + : `${count} sessions are still running. Quitting will stop them.`, + ]; + + if (sessionLabels.length > 0) { + sentences.push(`${sessionLabels.join(', ')}.`); + } + + if (runningActionCount > 0) { + sentences.push( + runningActionCount === 1 + ? '1 running action will also stop.' + : `${runningActionCount} running actions will also stop.` + ); + } + + return sentences.join(' '); +} diff --git a/apps/staged/src/lib/features/projects/ProjectHome.svelte b/apps/staged/src/lib/features/projects/ProjectHome.svelte index c6beb8ab0..26fda21e2 100644 --- a/apps/staged/src/lib/features/projects/ProjectHome.svelte +++ b/apps/staged/src/lib/features/projects/ProjectHome.svelte @@ -12,7 +12,6 @@ import Pause from '@lucide/svelte/icons/pause'; import Plus from '@lucide/svelte/icons/plus'; import Trash2 from '@lucide/svelte/icons/trash-2'; - import { getWindowSync } from '../../transport'; import type { Project, ProjectRepo, @@ -208,8 +207,10 @@ } } + // Quits rather than closing the window: closing the last window only hides + // it, and there is no usable app behind this screen to come back to. function handleClose() { - getWindowSync().close(); + void commands.quitApp().catch((e) => console.error('Failed to quit:', e)); } function scheduleDeferredTask(callback: () => void, timeout = 1500): () => void { diff --git a/apps/staged/src/lib/listeners/quitListener.test.ts b/apps/staged/src/lib/listeners/quitListener.test.ts new file mode 100644 index 000000000..70c885fb8 --- /dev/null +++ b/apps/staged/src/lib/listeners/quitListener.test.ts @@ -0,0 +1,137 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { ActiveSessionInfo, QuitRequestedPayload } from '../types'; + +const confirmQuit = vi.fn<() => Promise>(); +const cancelQuit = vi.fn<() => Promise>(); +const unlisten = vi.fn(); +// Window-scoped on purpose: the backend addresses the event to one window with +// `emit_to`, and an any-target listener would raise the dialog in all of them. +const listenToWindowEvent = vi.fn(); + +let handlers: Array<(payload: QuitRequestedPayload) => void>; + +/** + * Load the listener and store fresh, with transport in the requested mode. The + * store is a singleton, so each test needs its own module registry. + */ +async function load({ isTauri = true } = {}) { + vi.resetModules(); + vi.doMock('../transport', () => ({ isTauri, listenToWindowEvent })); + vi.doMock('../api/commands', () => ({ confirmQuit, cancelQuit })); + + const { listenForQuitRequests } = await import('./quitListener'); + const { quitPrompt } = await import('../stores/quitPrompt.svelte'); + return { listenForQuitRequests, quitPrompt }; +} + +function session(overrides: Partial = {}): ActiveSessionInfo { + return { + sessionId: 's1', + projectId: 'p1', + branchId: 'b1', + sessionType: 'commit', + status: 'running', + ...overrides, + }; +} + +describe('quitListener', () => { + beforeEach(() => { + // The store's runes compile away in the app build; under vitest they stay + // plain global calls, so stub $state as identity (projectsData.test.ts + // precedent). + vi.stubGlobal('$state', (initial: unknown) => initial); + handlers = []; + confirmQuit.mockReset().mockResolvedValue(undefined); + cancelQuit.mockReset().mockResolvedValue(undefined); + unlisten.mockReset(); + listenToWindowEvent.mockReset().mockImplementation((_event, handler) => { + handlers.push(handler as (payload: QuitRequestedPayload) => void); + return unlisten; + }); + }); + + afterEach(() => { + vi.doUnmock('../transport'); + vi.doUnmock('../api/commands'); + vi.unstubAllGlobals(); + vi.restoreAllMocks(); + }); + + it('opens the prompt with the payload from app:quit-requested', async () => { + const { listenForQuitRequests, quitPrompt } = await load(); + listenForQuitRequests(); + + expect(listenToWindowEvent).toHaveBeenCalledWith('app:quit-requested', expect.any(Function)); + expect(quitPrompt.open).toBe(false); + + handlers[0]({ sessions: [session()], runningActionCount: 2 }); + + expect(quitPrompt.open).toBe(true); + expect(quitPrompt.payload?.sessions).toHaveLength(1); + expect(quitPrompt.payload?.runningActionCount).toBe(2); + expect(quitPrompt.stopping).toBe(false); + }); + + it('registers no listener in web mode', async () => { + const { listenForQuitRequests } = await load({ isTauri: false }); + + // Callable no-op, so App.svelte's teardown needs no extra guard. + listenForQuitRequests()(); + + expect(listenToWindowEvent).not.toHaveBeenCalled(); + }); + + it('confirming invokes confirm_quit and leaves the dialog stopping', async () => { + const { listenForQuitRequests, quitPrompt } = await load(); + listenForQuitRequests(); + handlers[0]({ sessions: [session()], runningActionCount: 0 }); + + await quitPrompt.confirm(); + + expect(confirmQuit).toHaveBeenCalledTimes(1); + // The backend exits the process; until it does, the dialog reports progress + // instead of pretending the app is still usable. + expect(quitPrompt.open).toBe(true); + expect(quitPrompt.stopping).toBe(true); + + await quitPrompt.confirm(); + expect(confirmQuit).toHaveBeenCalledTimes(1); + }); + + it('closes the dialog when confirm_quit fails', async () => { + const { listenForQuitRequests, quitPrompt } = await load(); + listenForQuitRequests(); + handlers[0]({ sessions: [session()], runningActionCount: 0 }); + confirmQuit.mockRejectedValueOnce(new Error('nope')); + vi.spyOn(console, 'error').mockImplementation(() => {}); + + await quitPrompt.confirm(); + + expect(quitPrompt.open).toBe(false); + expect(quitPrompt.stopping).toBe(false); + }); + + it('cancelling invokes cancel_quit and closes the dialog', async () => { + const { listenForQuitRequests, quitPrompt } = await load(); + listenForQuitRequests(); + handlers[0]({ sessions: [session()], runningActionCount: 0 }); + + quitPrompt.cancel(); + + expect(cancelQuit).toHaveBeenCalledTimes(1); + expect(quitPrompt.open).toBe(false); + }); + + it('ignores a cancel once the quit is under way', async () => { + const { listenForQuitRequests, quitPrompt } = await load(); + listenForQuitRequests(); + handlers[0]({ sessions: [session()], runningActionCount: 0 }); + + await quitPrompt.confirm(); + quitPrompt.cancel(); + + expect(cancelQuit).not.toHaveBeenCalled(); + expect(quitPrompt.open).toBe(true); + }); +}); diff --git a/apps/staged/src/lib/listeners/quitListener.ts b/apps/staged/src/lib/listeners/quitListener.ts new file mode 100644 index 000000000..f88e6ac90 --- /dev/null +++ b/apps/staged/src/lib/listeners/quitListener.ts @@ -0,0 +1,24 @@ +/** + * Listener for the backend's `app:quit-requested` event. + * + * `Cmd+Q` / the app-menu Quit item reach `app_lifecycle::request_quit`, which + * emits this event instead of exiting when sessions are still active. The + * backend addresses it to exactly one window with `emit_to`, so this must be a + * *window-scoped* listener — the any-target `listenToEvent` also matches emits + * addressed to other windows, and every window would raise its own dialog. + * Wired at App level so it works on any route, and Tauri-only: quitting is a + * desktop-host action, and the `confirm_quit` command a browser client would + * need is deliberately absent from the web-mode dispatch table. + */ + +import { isTauri, listenToWindowEvent, type UnlistenFn } from '../transport'; +import { quitPrompt } from '../stores/quitPrompt.svelte'; +import type { QuitRequestedPayload } from '../types'; + +export function listenForQuitRequests(): UnlistenFn { + if (!isTauri) return () => {}; + + return listenToWindowEvent('app:quit-requested', (payload) => { + quitPrompt.requested(payload); + }); +} diff --git a/apps/staged/src/lib/stores/quitPrompt.svelte.ts b/apps/staged/src/lib/stores/quitPrompt.svelte.ts new file mode 100644 index 000000000..0ee0dbec7 --- /dev/null +++ b/apps/staged/src/lib/stores/quitPrompt.svelte.ts @@ -0,0 +1,61 @@ +/** + * State behind the quit confirmation dialog. + * + * The backend raises `app:quit-requested` when the user quits with sessions + * still active (see `app_lifecycle.rs`); quitListener.ts feeds that payload in + * here and QuitConfirmDialog renders it. Answering is a round trip back to the + * backend: confirming hands off to the shutdown sequence, which stops the + * sessions and then exits the process — so the dialog stays up, in its + * `stopping` state, until the app goes away underneath it. + */ + +import * as commands from '../api/commands'; +import type { QuitRequestedPayload } from '../types'; + +class QuitPromptStore { + private _payload = $state(null); + /** The quit was confirmed and the backend is stopping sessions. */ + private _stopping = $state(false); + + get payload(): QuitRequestedPayload | null { + return this._payload; + } + + get open(): boolean { + return this._payload !== null; + } + + get stopping(): boolean { + return this._stopping; + } + + /** A quit is waiting on the user's answer. */ + requested(payload: QuitRequestedPayload): void { + this._payload = payload; + this._stopping = false; + } + + /** Quit and stop the listed sessions. */ + async confirm(): Promise { + if (this._stopping) return; + this._stopping = true; + try { + await commands.confirmQuit(); + } catch (e) { + // The quit never started, so drop the dialog rather than leaving it stuck + // on "Stopping sessions…" for an app that isn't going anywhere. + console.error('Failed to confirm quit:', e); + this._payload = null; + this._stopping = false; + } + } + + /** Keep running. Also the Esc / click-outside path. */ + cancel(): void { + if (this._stopping) return; + this._payload = null; + void commands.cancelQuit().catch((e) => console.error('Failed to cancel quit:', e)); + } +} + +export const quitPrompt = new QuitPromptStore(); diff --git a/apps/staged/src/lib/types.ts b/apps/staged/src/lib/types.ts index 88ef5a3b5..cad11908c 100644 --- a/apps/staged/src/lib/types.ts +++ b/apps/staged/src/lib/types.ts @@ -664,6 +664,18 @@ export interface ActiveSessionInfo { status: SessionStatus; } +/** + * Payload of the `app:quit-requested` event: what a quit would interrupt. + * + * Emitted by `app_lifecycle::request_quit` when the user quits with sessions + * still active. Sessions are what gate the quit; running actions are reported + * so the dialog can say they stop too. + */ +export interface QuitRequestedPayload { + sessions: ActiveSessionInfo[]; + runningActionCount: number; +} + /** * Payload emitted by the `pr-created` domain event when a completed PR * session produced a pull request. The backend has already persisted the PR From 313925cef5153700dd1bf73bd39066bcb542a071 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Tue, 1 Sep 2026 16:11:21 +1000 Subject: [PATCH 02/13] refactor(lifecycle): ask before quitting with an unparented native alert MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The quit confirmation lived in the webview: `request_quit` emitted `app:quit-requested` to one chosen window and a Svelte `AlertDialog` rendered it. That forced a window into the quit path, and the case it hurt is the central scenario of this branch — the whole point of closing-window-is-not- quitting is that sessions keep streaming with everything hidden, so "all windows hidden, sessions running, user hits Cmd+Q" isn't a corner to tolerate, it's the primary path the confirmation exists to serve. Reaching it materialised a full application window (restored geometry, hydrating project tree) to host a two-button question, and `cancel_quit` did nothing but clear a flag — so Cmd+Q then Cancel left the user with a visible window they had to close a second time, having asked for neither. Parenting a native alert would not have fixed that: `.parent()` is exactly what makes `tauri-plugin-dialog` render an `NSAlert` as a window-modal *sheet*, so the reveal would have stayed. Unparented is a different widget, not a different modality of the same one — rfd 0.16 routes a parentless dialog to `CFUserNotificationDisplayAlert`, displayed by the system rather than by AppKit. That's what buys window-independence, so the reveal drops out of the quit path entirely: quitting from a hidden state stays hidden, cancelling returns the app to exactly the state the user left it in, and the branch where no window could be revealed and the app quit *without asking* disappears rather than being preserved. Structurally this resolves the review finding about a pending prompt outliving its host window by removing the concept of a host window. `QuitState.prompt_host` collapses to `prompt_pending: AtomicBool`; `clear_prompt_if_host`, the `Destroyed` arm of `on_window_event`, and `reveal_a_window` (now inlined into its one caller, `show_a_window`) all go away. The frontend half goes with them: `QuitConfirmDialog`, the `quitPrompt` store, `quitListener`, `quitPromptCopy`, the `QuitRequestedPayload` type, and the `confirm_quit` / `cancel_quit` commands. `quit_app` stays — the store-incompatibility screens still need it — and stays refused in the web-mode dispatch table. The prompt copy ports to Rust, where `get_branch` / `get_project` resolve the names the Svelte dialog used to read from its stores; the review's suggestion to fold the session list into the preceding sentence with a colon is taken while the wording moves. `OkCancelCustom`'s ok slot is the default (Return) button, so it holds "Keep Running" and the *cancel* slot holds "Quit & Stop Sessions" — a stray Return must not be what kills running agents. Accepted costs, all inherent to the widget: the alert carries generic system chrome rather than Staged's icon; the "Stopping sessions…" progress state is gone, since a native alert dismisses on click while `shutdown_cleanup` runs out its 2s budget; and it is not modal to the app, so work can start behind it. The last resolves correctly — `shutdown_cleanup` re-queries active sessions instead of trusting the prompt's snapshot — and it keeps the force-quit escape hatch dispatchable, which a second Cmd+Q needs. Because that dismissal leaves nothing on screen, `request_quit` now returns early when a shutdown is already under way instead of raising a second alert about sessions the first one is stopping. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/app_lifecycle.rs | 477 +++++++++++++----- apps/staged/src-tauri/src/lib.rs | 7 +- apps/staged/src-tauri/src/web_server.rs | 7 +- apps/staged/src/App.svelte | 10 - apps/staged/src/lib/commands.ts | 24 +- .../lifecycle/QuitConfirmDialog.svelte | 62 --- .../features/lifecycle/quitPromptCopy.test.ts | 63 --- .../lib/features/lifecycle/quitPromptCopy.ts | 60 --- .../src/lib/listeners/quitListener.test.ts | 137 ----- apps/staged/src/lib/listeners/quitListener.ts | 24 - .../src/lib/stores/quitPrompt.svelte.ts | 61 --- apps/staged/src/lib/types.ts | 12 - 12 files changed, 350 insertions(+), 594 deletions(-) delete mode 100644 apps/staged/src/lib/features/lifecycle/QuitConfirmDialog.svelte delete mode 100644 apps/staged/src/lib/features/lifecycle/quitPromptCopy.test.ts delete mode 100644 apps/staged/src/lib/features/lifecycle/quitPromptCopy.ts delete mode 100644 apps/staged/src/lib/listeners/quitListener.test.ts delete mode 100644 apps/staged/src/lib/listeners/quitListener.ts delete mode 100644 apps/staged/src/lib/stores/quitPrompt.svelte.ts diff --git a/apps/staged/src-tauri/src/app_lifecycle.rs b/apps/staged/src-tauri/src/app_lifecycle.rs index aae9e7916..43f715e61 100644 --- a/apps/staged/src-tauri/src/app_lifecycle.rs +++ b/apps/staged/src-tauri/src/app_lifecycle.rs @@ -14,14 +14,25 @@ //! still quits there — but through the same confirmation gate as `Cmd+Q`. //! //! **Quitting with sessions running asks first, then stops them cleanly.** -//! [`request_quit`] gates on active sessions and hands the decision to the -//! frontend dialog, addressed to a single live window (revealed first if every -//! window is hidden); [`shutdown_cleanup`] cancels sessions with -//! [`CompletionReason::AppQuit`] and stops actions. That cancel is the only -//! thing that shuts an agent down: ACP children are spawned with -//! `process_group(0)` and `kill_on_drop`, and `process::exit` runs no +//! [`request_quit`] gates on active sessions and asks; [`shutdown_cleanup`] +//! cancels sessions with [`CompletionReason::AppQuit`] and stops actions. That +//! cancel is the only thing that shuts an agent down: ACP children are spawned +//! with `process_group(0)` and `kill_on_drop`, and `process::exit` runs no //! destructors, so a bare exit leaves the agent CLIs running. //! +//! The question is asked by a native alert with **no parent window**, not by a +//! dialog in a webview. Quitting is scoped to the application, and the case the +//! confirmation exists for is precisely the one where every window is hidden: +//! parenting the alert (which `tauri-plugin-dialog` renders as a window-modal +//! sheet) would drag a full window back on screen — restored geometry, +//! hydrating project tree and all — to host a two-button question, and +//! cancelling would leave it there. Unparented, rfd reaches for +//! `CFUserNotificationDisplayAlert` on macOS instead of `NSAlert`: system +//! chrome rather than the app's, and not modal to the app, in exchange for +//! needing no window at all. So quitting from a hidden state stays hidden, +//! cancelling returns the app to exactly the state the user left it in, and +//! there is no longer any state where a quit can't ask. +//! //! Every exit path funnels into [`shutdown_cleanup`], which runs its work at //! most once — a confirmed quit calls it directly, `RunEvent::ExitRequested` //! covers programmatic exits, and `RunEvent::Exit` is the only hook on the @@ -32,16 +43,28 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -use serde::Serialize; -use tauri::{AppHandle, Emitter, Manager, WebviewWindow, Window, WindowEvent}; +use tauri::{AppHandle, Manager, Window, WindowEvent}; +use tauri_plugin_dialog::{ + DialogExt, MessageDialogButtons, MessageDialogKind, MessageDialogResult, +}; use crate::actions; use crate::session_commands::{self, ActiveSessionInfo}; use crate::session_runner::SessionRegistry; use crate::store::{CompletionReason, Session, SessionStatus, Store}; -/// Event that raises the frontend's quit confirmation dialog. -const QUIT_REQUESTED_EVENT: &str = "app:quit-requested"; +/// Title of the quit confirmation alert. +const QUIT_PROMPT_TITLE: &str = "Quit Staged?"; + +/// Alert button that goes through with the quit. +/// +/// It sits in `OkCancelCustom`'s *cancel* slot, and [`KEEP_RUNNING_BUTTON`] in +/// the ok slot, because the ok slot is the default (`Return`) button — a stray +/// Return must not be what kills a room full of running agents. +const QUIT_BUTTON: &str = "Quit & Stop Sessions"; + +/// Alert button that dismisses the prompt and leaves the sessions alone. +const KEEP_RUNNING_BUTTON: &str = "Keep Running"; /// Menu id of the app-menu Quit item. Custom rather than /// `PredefinedMenuItem::quit` so `Cmd+Q` is routable at all: the predefined item @@ -73,44 +96,33 @@ pub struct QuitState { /// Set by the first caller into [`shutdown_cleanup`], so the cleanup runs /// exactly once however many exit events follow it. quit_in_progress: AtomicBool, - /// Label of the window showing an unanswered confirmation dialog. A quit - /// request arriving while it is set forces the quit — a wedged webview must - /// never be able to trap the app, so a second `Cmd+Q` always gets out. The - /// label is what lets a destroyed host window clear the flag instead of - /// leaving that force path armed with no dialog on screen. - prompt_host: Mutex>, + /// Set while a confirmation alert is unanswered. A quit request arriving + /// while it is set forces the quit, so an alert that never appeared or never + /// came back can't trap the app: a second `Cmd+Q` always gets out. + prompt_pending: AtomicBool, } impl QuitState { - fn set_prompt_host(&self, label: &str) { - *self.prompt_host.lock().unwrap() = Some(label.to_string()); + fn set_prompt_pending(&self) { + self.prompt_pending.store(true, Ordering::SeqCst); } /// Clear any pending prompt, returning whether one was pending. fn take_prompt(&self) -> bool { - self.prompt_host.lock().unwrap().take().is_some() - } - - /// Clear the pending prompt if `label` was hosting it. - fn clear_prompt_if_host(&self, label: &str) { - let mut host = self.prompt_host.lock().unwrap(); - if host.as_deref() == Some(label) { - *host = None; - } + self.prompt_pending.swap(false, Ordering::SeqCst) } } -/// What a quit would interrupt, as sent to the confirmation dialog. -#[derive(Debug, Clone, Default, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct QuitBlockers { - /// Running and queued sessions owned by this process — the only thing that - /// gates a quit. - pub sessions: Vec, - /// Running actions. Reported so the dialog can say they stop too, but they +/// What a quit would interrupt, as the alert describes it. +#[derive(Debug, Default)] +struct QuitBlockers { + /// One label per running or queued session owned by this process, e.g. + /// `review on fix-login` — the only thing that gates a quit. + session_labels: Vec, + /// Running actions. Reported so the alert can say they stop too, but they /// don't gate the quit on their own: a dev server left running is the normal /// state of a workspace, and blocking `Cmd+Q` on it would be noise. - pub running_action_count: usize, + running_action_count: usize, } /// Whether a quit should stop and ask first. @@ -118,7 +130,7 @@ pub struct QuitBlockers { /// Queued sessions count: they're work the user asked for that a quit silently /// drops, so they belong in the prompt. fn should_prompt(blockers: &QuitBlockers) -> bool { - !blockers.sessions.is_empty() + !blockers.session_labels.is_empty() } // ============================================================================= @@ -128,17 +140,8 @@ fn should_prompt(blockers: &QuitBlockers) -> bool { /// `Builder::on_window_event` hook — see the module docs for why closing the /// last window doesn't end the process. pub fn on_window_event(window: &Window, event: &WindowEvent) { - match event { - WindowEvent::CloseRequested { api, .. } => on_close_requested(window, api), - // A destroyed window takes its webview — and any dialog in it — with - // it. Left set, the pending flag would turn the next quit request into - // a silent force-quit; cleared, that quit just asks again. - WindowEvent::Destroyed => { - if let Some(quit_state) = window.app_handle().try_state::() { - quit_state.clear_prompt_if_host(window.label()); - } - } - _ => {} + if let WindowEvent::CloseRequested { api, .. } = event { + on_close_requested(window, api); } } @@ -190,25 +193,16 @@ fn hide_window(window: &Window) { set_native_focus(window.app_handle(), window.label(), false); } -/// Bring a window back on screen: the Dock-icon click, `Window ▸ Staged`, and a -/// quit arriving with no visible window all funnel here. -pub fn show_a_window(app: &AppHandle) { - if reveal_a_window(app).is_none() { - log::warn!("No window left to show"); - } -} - -/// Pick a window and make sure it is on screen and focused, returning it. +/// Bring a window back on screen: the Dock-icon click and `Window ▸ Staged` +/// both funnel here. Deliberately *not* on the quit path — see the module docs. /// -/// Prefers where the user already is (focused, then visible — reachable when a -/// quit request arrives from the store-incompatibility screen or the web -/// dispatch refusal path while windows are up), then falls back to unhiding one: -/// `main` for its restored geometry, else any. `None` only if every window has -/// been destroyed, which no close path produces — closing the last window hides -/// it instead. -fn reveal_a_window(app: &AppHandle) -> Option { +/// Prefers where the user already is (focused, then visible), then falls back to +/// unhiding one: `main` for its restored geometry, else any surviving `win-N` +/// peer. Finds nothing only if every window has been destroyed, which no close +/// path produces — closing the last window hides it instead. +pub fn show_a_window(app: &AppHandle) { let windows = app.webview_windows(); - let window = windows + let Some(window) = windows .values() .find(|window| window.is_focused().unwrap_or(false)) .or_else(|| { @@ -217,7 +211,11 @@ fn reveal_a_window(app: &AppHandle) -> Option { .find(|window| window.is_visible().unwrap_or(false)) }) .or_else(|| windows.get(MAIN_WINDOW_LABEL)) - .or_else(|| windows.values().next())?; + .or_else(|| windows.values().next()) + else { + log::warn!("No window left to show"); + return; + }; if let Err(e) = window.show() { log::warn!("Failed to show window: {e}"); @@ -229,7 +227,6 @@ fn reveal_a_window(app: &AppHandle) -> Option { log::warn!("Failed to focus window: {e}"); } set_native_focus(app, window.label(), true); - Some(window.clone()) } /// Mirror a native window's visibility onto its PR-poll client's focus hint. @@ -247,38 +244,71 @@ fn set_native_focus(app: &AppHandle, window_label: &str, focused: bool) { /// Handle a quit request from the app menu, `Cmd+Q`, or (off macOS) the last /// window's close. Cheap enough for the main thread: it snapshots blockers and -/// either hands off to a background quit or raises the dialog. +/// either hands off to a background quit or raises the alert. pub fn request_quit(app: &AppHandle, force: bool) { let quit_state = app.state::(); - // A quit arriving while the dialog is unanswered (a second `Cmd+Q`) is the - // escape hatch from a webview that never rendered or answered it. + // A quit arriving while the alert is unanswered (a second `Cmd+Q`) is the + // escape hatch from a prompt that never appeared or never came back. The + // system alert isn't app-modal, so that second `Cmd+Q` is still dispatchable + // with the alert on screen. if force || quit_state.take_prompt() { spawn_quit(app); return; } + // Already shutting down, and the alert dismissed on click while cleanup runs + // out its budget — so there is nothing on screen saying so, and a `Cmd+Q` + // here means "I already answered", not "ask me again". + if quit_state.quit_in_progress.load(Ordering::SeqCst) { + return; + } + let blockers = collect_quit_blockers(app); if !should_prompt(&blockers) { spawn_quit(app); return; } - // The dialog goes to exactly one window — where the user is, or a window - // revealed for the purpose if the quit arrived with everything hidden. A - // broadcast would raise one dialog per window, each unaware of the others' - // answers. No window at all means nobody to ask, so the quit proceeds. - let Some(host) = reveal_a_window(app) else { - spawn_quit(app); - return; - }; - quit_state.set_prompt_host(host.label()); + quit_state.set_prompt_pending(); + ask_before_quitting(app, &blockers); +} - if let Err(e) = app.emit_to(host.label(), QUIT_REQUESTED_EVENT, &blockers) { - log::warn!("Failed to ask for quit confirmation, quitting anyway: {e}"); - quit_state.take_prompt(); - spawn_quit(app); - } +/// Raise the confirmation alert and act on the answer. +/// +/// No `.parent()`, which is what keeps this window-independent — see the module +/// docs. `tauri-plugin-dialog` hops to the main thread to start the alert and +/// then runs it on its own thread, so this returns immediately and the event +/// loop keeps turning underneath it. +fn ask_before_quitting(app: &AppHandle, blockers: &QuitBlockers) { + let app = app.clone(); + app.dialog() + .message(quit_prompt_message(blockers)) + .title(QUIT_PROMPT_TITLE) + .kind(MessageDialogKind::Warning) + .buttons(MessageDialogButtons::OkCancelCustom( + KEEP_RUNNING_BUTTON.to_string(), + QUIT_BUTTON.to_string(), + )) + .show_with_result(move |result| { + app.state::().take_prompt(); + // Anything that isn't the quit button — "Keep Running", or the + // system dismissing the alert itself — leaves the sessions alone. + // Nothing to undo on that path: no window was revealed to host the + // question, so the app is already in the state the user left it in. + if quit_confirmed(&result) { + // The snapshot the message was built from may be stale by now: + // the alert is not modal to the app, so a session could have + // started or finished behind it. `shutdown_cleanup` re-queries, + // so it stops what is actually running. + spawn_quit(&app); + } + }); +} + +/// Whether the alert was answered with [`QUIT_BUTTON`]. +fn quit_confirmed(result: &MessageDialogResult) -> bool { + matches!(result, MessageDialogResult::Custom(label) if label == QUIT_BUTTON) } /// Quit from the UI, through the same gate as `Cmd+Q`. @@ -286,30 +316,17 @@ pub fn request_quit(app: &AppHandle, force: bool) { /// Used by the store-incompatibility screens' "Close" button, which has to end /// the app: closing the last window only hides it, and those screens have no /// working database behind them to come back to. -#[tauri::command] -pub fn quit_app(app_handle: AppHandle) { - request_quit(&app_handle, false); -} - -/// Quit confirmed in the dialog: stop sessions and actions, then exit. /// /// Deliberately absent from the web-mode `dispatch` table — a browser client /// must not be able to terminate the desktop host. #[tauri::command] -pub fn confirm_quit(app_handle: AppHandle) { - app_handle.state::().take_prompt(); - spawn_quit(&app_handle); -} - -/// Quit declined in the dialog: sessions keep running. -#[tauri::command] -pub fn cancel_quit(app_handle: AppHandle) { - app_handle.state::().take_prompt(); +pub fn quit_app(app_handle: AppHandle) { + request_quit(&app_handle, false); } /// Run the quit sequence off the main thread so the bounded waits never freeze -/// the event loop — the dialog stays interactive and can render its -/// "Stopping sessions…" state while agents shut down. +/// the event loop — windows keep repainting while agents shut down, and the +/// close events the exit generates are still delivered. fn spawn_quit(app: &AppHandle) { let app = app.clone(); std::thread::spawn(move || { @@ -320,10 +337,14 @@ fn spawn_quit(app: &AppHandle) { /// Snapshot what a quit would interrupt. fn collect_quit_blockers(app: &AppHandle) -> QuitBlockers { - let sessions = match app_store(app) { + let session_labels = match app_store(app) { Some(store) => owned_active_sessions(&store) .iter() - .map(|session| session_commands::project_active_session(&store, session)) + .map(|session| { + let session = session_commands::project_active_session(&store, session); + let location = session_location(&store, &session); + quit_session_label(&session, location.as_deref()) + }) .collect(), None => Vec::new(), }; @@ -341,11 +362,94 @@ fn collect_quit_blockers(app: &AppHandle) -> QuitBlockers { }; QuitBlockers { - sessions, + session_labels, running_action_count, } } +// ============================================================================= +// Prompt copy +// ============================================================================= + +/// How each session type reads in the alert. +fn session_type_label(session_type: &str) -> Option<&'static str> { + match session_type { + "note" => Some("note"), + "commit" => Some("commit"), + "review" => Some("review"), + "pr" => Some("PR"), + "push" => Some("push"), + "pull" => Some("pull"), + _ => None, + } +} + +/// Where a session is running, as the user knows it: its branch name, or its +/// project name for project-level sessions (a note on a project has no branch). +fn session_location(store: &Store, session: &ActiveSessionInfo) -> Option { + let branch_name = session + .branch_id + .as_deref() + .and_then(|id| store.get_branch(id).ok().flatten()) + .map(|branch| branch.branch_name); + + branch_name.or_else(|| { + session + .project_id + .as_deref() + .and_then(|id| store.get_project(id).ok().flatten()) + .map(|project| project.name) + }) +} + +/// Label for one session a quit would stop, e.g. `review on fix-login` or +/// `commit on fix-login (queued)`. +/// +/// Both halves can be missing — an unrecognised session type, or a row whose +/// branch and project have already been deleted — so each falls back rather than +/// dropping the session from the list. +fn quit_session_label(session: &ActiveSessionInfo, location: Option<&str>) -> String { + let kind = session + .session_type + .as_deref() + .and_then(session_type_label) + .unwrap_or("session"); + let base = match location { + Some(location) => format!("{kind} on {location}"), + None => kind.to_string(), + }; + + if session.status == SessionStatus::Queued.as_str() { + format!("{base} (queued)") + } else { + base + } +} + +/// Alert body: how much stops, what it is, and whether actions go with it. +/// +/// Actions never gate the quit (see [`should_prompt`]), so they are mentioned +/// only as a consequence of one. +fn quit_prompt_message(blockers: &QuitBlockers) -> String { + let labels = blockers.session_labels.join(", "); + let mut message = if blockers.session_labels.len() == 1 { + format!("1 session is still running: {labels}. Quitting will stop it.") + } else { + format!( + "{} sessions are still running: {labels}. Quitting will stop them.", + blockers.session_labels.len() + ) + }; + + match blockers.running_action_count { + 0 => {} + 1 => message.push_str(" 1 running action will also stop."), + count => message.push_str(&format!(" {count} running actions will also stop.")), + } + + message +} + // ============================================================================= // Shutdown cleanup // ============================================================================= @@ -503,41 +607,31 @@ mod tests { use super::*; use std::path::Path; - fn active_session(status: &str) -> ActiveSessionInfo { + fn active_session(session_type: Option<&str>, status: SessionStatus) -> ActiveSessionInfo { ActiveSessionInfo { session_id: "s1".to_string(), - project_id: None, - branch_id: None, - session_type: None, - status: status.to_string(), + project_id: Some("p1".to_string()), + branch_id: Some("b1".to_string()), + session_type: session_type.map(str::to_string), + status: status.as_str().to_string(), } } - #[test] - fn running_sessions_prompt() { - let blockers = QuitBlockers { - sessions: vec![active_session("running")], - running_action_count: 0, - }; - assert!(should_prompt(&blockers)); + fn blockers(session_labels: &[&str], running_action_count: usize) -> QuitBlockers { + QuitBlockers { + session_labels: session_labels.iter().map(|s| s.to_string()).collect(), + running_action_count, + } } #[test] - fn queued_sessions_prompt() { - let blockers = QuitBlockers { - sessions: vec![active_session("queued")], - running_action_count: 0, - }; - assert!(should_prompt(&blockers)); + fn active_sessions_prompt() { + assert!(should_prompt(&blockers(&["review on fix-login"], 0))); } #[test] fn running_actions_alone_do_not_prompt() { - let blockers = QuitBlockers { - sessions: Vec::new(), - running_action_count: 3, - }; - assert!(!should_prompt(&blockers)); + assert!(!should_prompt(&blockers(&[], 3))); } #[test] @@ -545,23 +639,128 @@ mod tests { assert!(!should_prompt(&QuitBlockers::default())); } - /// The pending flag turns the next quit into a force-quit, so it must not - /// outlive the window whose dialog it stands for — but a *peer* window - /// closing must not answer a dialog it isn't showing. + /// The pending flag turns the next quit into a force-quit, so answering the + /// alert has to disarm it — otherwise the next `Cmd+Q` quits without asking. #[test] - fn prompt_clears_only_when_its_host_window_is_destroyed() { + fn answering_the_prompt_disarms_the_force_path() { let state = QuitState::default(); - state.set_prompt_host("win-2"); - state.clear_prompt_if_host("main"); - assert!(state.take_prompt(), "peer destruction dropped the prompt"); + assert!(!state.take_prompt(), "nothing pending, nothing to force"); + + state.set_prompt_pending(); + assert!(state.take_prompt(), "pending prompt did not arm the force"); + assert!(!state.take_prompt(), "prompt stayed armed after answering"); + } + + /// The ok slot is the default (`Return`) button, so it holds "Keep Running" + /// and the cancel slot holds the destructive answer. + #[test] + fn only_the_quit_button_confirms() { + assert!(quit_confirmed(&MessageDialogResult::Custom( + QUIT_BUTTON.to_string() + ))); + assert!(!quit_confirmed(&MessageDialogResult::Custom( + KEEP_RUNNING_BUTTON.to_string() + ))); + // What a system-dismissed alert reports. + assert!(!quit_confirmed(&MessageDialogResult::Cancel)); + assert!(!quit_confirmed(&MessageDialogResult::Ok)); + } + + #[test] + fn session_label_names_the_type_and_where_it_runs() { + assert_eq!( + quit_session_label( + &active_session(Some("review"), SessionStatus::Running), + Some("fix-login") + ), + "review on fix-login" + ); + } + + #[test] + fn session_label_marks_queued_sessions() { + assert_eq!( + quit_session_label( + &active_session(Some("review"), SessionStatus::Queued), + Some("fix-login") + ), + "review on fix-login (queued)" + ); + } + + #[test] + fn session_label_falls_back_for_unknown_type_or_missing_location() { + assert_eq!( + quit_session_label(&active_session(None, SessionStatus::Running), Some("docs")), + "session on docs" + ); + assert_eq!( + quit_session_label( + &active_session(Some("mystery"), SessionStatus::Running), + Some("docs") + ), + "session on docs" + ); + assert_eq!( + quit_session_label(&active_session(Some("note"), SessionStatus::Running), None), + "note" + ); + } + + /// A branch session reads as its branch; a project-level session (a note on + /// a project) has no branch, so it reads as its project. + #[test] + fn session_location_prefers_the_branch_then_the_project() { + let store = Store::in_memory().unwrap(); + let mut project = crate::store::Project::new("owner/repo"); + project.name = "Widgets".to_string(); + store.create_project(&project).unwrap(); + let branch = crate::store::Branch::new(&project.id, "fix-login", "main"); + store.create_branch(&branch).unwrap(); + + let mut session = active_session(Some("note"), SessionStatus::Running); + session.project_id = Some(project.id.clone()); + session.branch_id = Some(branch.id.clone()); + assert_eq!( + session_location(&store, &session).as_deref(), + Some("fix-login") + ); - state.set_prompt_host("win-2"); - state.clear_prompt_if_host("win-2"); - assert!( - !state.take_prompt(), - "host destruction left the prompt armed" + session.branch_id = None; + assert_eq!( + session_location(&store, &session).as_deref(), + Some("Widgets") ); + + session.project_id = None; + assert_eq!(session_location(&store, &session), None); + } + + #[test] + fn prompt_message_reads_singular_for_one_session() { + assert_eq!( + quit_prompt_message(&blockers(&["commit on fix-login"], 0)), + "1 session is still running: commit on fix-login. Quitting will stop it." + ); + } + + #[test] + fn prompt_message_lists_every_session_for_a_plural_count() { + assert_eq!( + quit_prompt_message(&blockers(&["commit on fix-login", "note on docs"], 0)), + "2 sessions are still running: commit on fix-login, note on docs. \ + Quitting will stop them." + ); + } + + #[test] + fn prompt_message_mentions_actions_only_when_there_are_some() { + assert!(quit_prompt_message(&blockers(&["commit on fix-login"], 1)) + .ends_with(" 1 running action will also stop.")); + assert!(quit_prompt_message(&blockers(&["commit on fix-login"], 3)) + .ends_with(" 3 running actions will also stop.")); + assert!(!quit_prompt_message(&blockers(&["commit on fix-login"], 0)).contains("action")); } #[test] diff --git a/apps/staged/src-tauri/src/lib.rs b/apps/staged/src-tauri/src/lib.rs index 5cc4db542..66dee3395 100644 --- a/apps/staged/src-tauri/src/lib.rs +++ b/apps/staged/src-tauri/src/lib.rs @@ -2233,8 +2233,7 @@ pub fn run() { } }) .on_window_event(|window, event| { - // Close-to-hide / the quit gate (`CloseRequested`), and dropping a - // pending quit prompt whose host window went away (`Destroyed`). + // Close-to-hide / the quit gate (`CloseRequested`). app_lifecycle::on_window_event(window, event); if let tauri::WindowEvent::Destroyed = event { @@ -2273,10 +2272,8 @@ pub fn run() { window_commands::take_window_seed, window_commands::claim_updater_ownership, // Lifecycle — desktop only; the web-mode `dispatch` table refuses - // these so a browser client can't quit the host. + // this so a browser client can't quit the host. app_lifecycle::quit_app, - app_lifecycle::confirm_quit, - app_lifecycle::cancel_quit, list_projects, create_project, list_project_repos, diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index 2daac5533..ccd48d9c0 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -564,10 +564,9 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result { - Err(format!("{command} is not available in web mode")) - } + // confirmation this gates on is a native alert on the host anyway, with + // no browser client to answer it. + "quit_app" => Err(format!("{command} is not available in web mode")), // ===================================================================== // Projects diff --git a/apps/staged/src/App.svelte b/apps/staged/src/App.svelte index 3dd41cbe0..0aa5144ab 100644 --- a/apps/staged/src/App.svelte +++ b/apps/staged/src/App.svelte @@ -19,7 +19,6 @@ import ProjectsList from './lib/features/projects/ProjectsList.svelte'; import ProjectsSidebar from './lib/features/projects/ProjectsSidebar.svelte'; import ProjectDeleteDialog from './lib/features/projects/ProjectDeleteDialog.svelte'; - import QuitConfirmDialog from './lib/features/lifecycle/QuitConfirmDialog.svelte'; import ReposListView from './lib/features/projects/ReposListView.svelte'; import SessionLauncher from './lib/features/sessions/SessionLauncher.svelte'; import SettingsPage from './lib/features/settings/SettingsPage.svelte'; @@ -63,7 +62,6 @@ import { listenForPageLifecycle } from './lib/listeners/pageLifecycleListener'; import { listenForAcpToolsReconciled } from './lib/listeners/acpToolsListener'; import { listenForMenuEvents } from './lib/listeners/menuListener'; - import { listenForQuitRequests } from './lib/listeners/quitListener'; import { darkMode } from './lib/stores/isDark.svelte'; import * as prPollingService from './lib/services/prPollingService'; import type { StoreIncompatibility } from './lib/types'; @@ -79,7 +77,6 @@ let unlistenAcpToolsReconciled: UnlistenFn | undefined; let unlistenStoreReset: UnlistenFn | undefined; let unlistenUpdaterOwnerAvailable: UnlistenFn | undefined; - let unlistenQuitRequests: UnlistenFn | undefined; let unregisterShortcuts: (() => void) | null = null; let stopUpdaterLoop: (() => void) | null = null; let updaterStartPending = false; @@ -359,9 +356,6 @@ // Refresh provider discovery (and any loaded doctor report) once the // backend finishes installing/upgrading the managed ACP bridges. unlistenAcpToolsReconciled = listenForAcpToolsReconciled(); - // Raise the quit confirmation when the backend gates a quit on running - // sessions (Tauri only — see quitListener.ts). - unlistenQuitRequests = listenForQuitRequests(); // Keep the shared project-list cache fresh for the app's lifetime — the // store dedupes, so starting before any view consumes it is safe. projectsDataStore.startListeners(); @@ -574,7 +568,6 @@ unlistenAcpToolsReconciled?.(); unlistenStoreReset?.(); unlistenUpdaterOwnerAvailable?.(); - unlistenQuitRequests?.(); projectsDataStore.stopListeners(); projectRunActionsStore.stopListening(); stopUpdaterLoop?.(); @@ -706,9 +699,6 @@ point (sidebar, landing grid, ProjectHome top bar/shortcut). --> - - - {/if} diff --git a/apps/staged/src/lib/commands.ts b/apps/staged/src/lib/commands.ts index f94fc86a2..3943cf657 100644 --- a/apps/staged/src/lib/commands.ts +++ b/apps/staged/src/lib/commands.ts @@ -75,28 +75,18 @@ export function confirmResetStore(): Promise { // ============================================================================= /** - * Quit Staged, through the same gate as `Cmd+Q` — it raises the confirmation - * dialog when sessions are still running. Closing the window only hides it, so - * UI that means "end the app" (the store-incompatibility screens) needs this. + * Quit Staged, through the same gate as `Cmd+Q` — the backend raises a native + * confirmation alert when sessions are still running, and owns the answer. + * Closing the window only hides it, so UI that means "end the app" (the + * store-incompatibility screens) needs this. + * + * Desktop only: the command is not in the web-mode dispatch table, so a browser + * client cannot quit the host. */ export function quitApp(): Promise { return invokeCommand('quit_app'); } -/** - * Confirm the quit raised by `app:quit-requested`: the backend stops the active - * sessions and running actions, then exits. Desktop only — the command is not in - * the web-mode dispatch table, so a browser client cannot quit the host. - */ -export function confirmQuit(): Promise { - return invokeCommand('confirm_quit'); -} - -/** Decline the quit; sessions keep running. */ -export function cancelQuit(): Promise { - return invokeCommand('cancel_quit'); -} - // ============================================================================= // Projects // ============================================================================= diff --git a/apps/staged/src/lib/features/lifecycle/QuitConfirmDialog.svelte b/apps/staged/src/lib/features/lifecycle/QuitConfirmDialog.svelte deleted file mode 100644 index db45d8fd9..000000000 --- a/apps/staged/src/lib/features/lifecycle/QuitConfirmDialog.svelte +++ /dev/null @@ -1,62 +0,0 @@ - - - - !v && quitPrompt.cancel()}> - - - Quit Staged? - {description} - - - Cancel - quitPrompt.confirm()} - > - {quitPrompt.stopping ? 'Stopping sessions…' : 'Quit & Stop Sessions'} - - - - diff --git a/apps/staged/src/lib/features/lifecycle/quitPromptCopy.test.ts b/apps/staged/src/lib/features/lifecycle/quitPromptCopy.test.ts deleted file mode 100644 index 0f1310cee..000000000 --- a/apps/staged/src/lib/features/lifecycle/quitPromptCopy.test.ts +++ /dev/null @@ -1,63 +0,0 @@ -import { describe, expect, it } from 'vitest'; -import type { ActiveSessionInfo } from '../../types'; -import { quitPromptDescription, quitSessionLabel } from './quitPromptCopy'; - -function session(overrides: Partial = {}): ActiveSessionInfo { - return { - sessionId: 's1', - projectId: 'p1', - branchId: 'b1', - sessionType: 'review', - status: 'running', - ...overrides, - }; -} - -describe('quitSessionLabel', () => { - it('names the session type and where it runs', () => { - expect(quitSessionLabel(session(), 'fix-login')).toBe('review on fix-login'); - }); - - it('marks queued sessions', () => { - expect(quitSessionLabel(session({ status: 'queued' }), 'fix-login')).toBe( - 'review on fix-login (queued)' - ); - }); - - it('falls back to "session" for an unknown or missing type', () => { - expect(quitSessionLabel(session({ sessionType: null }), 'fix-login')).toBe( - 'session on fix-login' - ); - expect(quitSessionLabel(session({ sessionType: 'mystery' }), 'fix-login')).toBe( - 'session on fix-login' - ); - }); - - it('drops the location when there is none to show', () => { - expect(quitSessionLabel(session({ sessionType: 'note' }), null)).toBe('note'); - }); -}); - -describe('quitPromptDescription', () => { - it('reads singular for one session', () => { - expect(quitPromptDescription(['commit on fix-login'], 0)).toBe( - '1 session is still running. Quitting will stop it. commit on fix-login.' - ); - }); - - it('lists every session for a plural count', () => { - expect(quitPromptDescription(['commit on fix-login', 'note on docs'], 0)).toBe( - '2 sessions are still running. Quitting will stop them. commit on fix-login, note on docs.' - ); - }); - - it('mentions running actions only when there are some', () => { - expect(quitPromptDescription(['commit on fix-login'], 1)).toContain( - '1 running action will also stop.' - ); - expect(quitPromptDescription(['commit on fix-login'], 3)).toContain( - '3 running actions will also stop.' - ); - expect(quitPromptDescription(['commit on fix-login'], 0)).not.toContain('action'); - }); -}); diff --git a/apps/staged/src/lib/features/lifecycle/quitPromptCopy.ts b/apps/staged/src/lib/features/lifecycle/quitPromptCopy.ts deleted file mode 100644 index 508f3dbca..000000000 --- a/apps/staged/src/lib/features/lifecycle/quitPromptCopy.ts +++ /dev/null @@ -1,60 +0,0 @@ -/** - * Copy for the quit confirmation dialog. - * - * Kept out of the component so the wording is unit-testable: the dialog resolves - * each session's branch/project names from the stores and passes labels in. - */ - -import type { ActiveSessionInfo } from '../../types'; - -/** How each session type reads in the dialog body. */ -const SESSION_TYPE_LABELS: Record = { - note: 'note', - commit: 'commit', - review: 'review', - pr: 'PR', - push: 'push', - pull: 'pull', -}; - -/** - * Label for one session that a quit would stop, e.g. `review on fix-login` or - * `commit on fix-login (queued)`. - * - * `where` is the branch name when the session belongs to one, otherwise the - * project name — project-level sessions (notes on a project) have no branch. - */ -export function quitSessionLabel(session: ActiveSessionInfo, where: string | null): string { - const kind = session.sessionType ? SESSION_TYPE_LABELS[session.sessionType] : null; - const base = where ? `${kind ?? 'session'} on ${where}` : (kind ?? 'session'); - return session.status === 'queued' ? `${base} (queued)` : base; -} - -/** - * Dialog body: how much stops, what it is, and whether actions go with it. - * - * Actions never gate the quit (see `should_prompt` in `app_lifecycle.rs`), so - * they are mentioned only as a consequence of one. - */ -export function quitPromptDescription(sessionLabels: string[], runningActionCount: number): string { - const count = sessionLabels.length; - const sentences = [ - count === 1 - ? '1 session is still running. Quitting will stop it.' - : `${count} sessions are still running. Quitting will stop them.`, - ]; - - if (sessionLabels.length > 0) { - sentences.push(`${sessionLabels.join(', ')}.`); - } - - if (runningActionCount > 0) { - sentences.push( - runningActionCount === 1 - ? '1 running action will also stop.' - : `${runningActionCount} running actions will also stop.` - ); - } - - return sentences.join(' '); -} diff --git a/apps/staged/src/lib/listeners/quitListener.test.ts b/apps/staged/src/lib/listeners/quitListener.test.ts deleted file mode 100644 index 70c885fb8..000000000 --- a/apps/staged/src/lib/listeners/quitListener.test.ts +++ /dev/null @@ -1,137 +0,0 @@ -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; -import type { ActiveSessionInfo, QuitRequestedPayload } from '../types'; - -const confirmQuit = vi.fn<() => Promise>(); -const cancelQuit = vi.fn<() => Promise>(); -const unlisten = vi.fn(); -// Window-scoped on purpose: the backend addresses the event to one window with -// `emit_to`, and an any-target listener would raise the dialog in all of them. -const listenToWindowEvent = vi.fn(); - -let handlers: Array<(payload: QuitRequestedPayload) => void>; - -/** - * Load the listener and store fresh, with transport in the requested mode. The - * store is a singleton, so each test needs its own module registry. - */ -async function load({ isTauri = true } = {}) { - vi.resetModules(); - vi.doMock('../transport', () => ({ isTauri, listenToWindowEvent })); - vi.doMock('../api/commands', () => ({ confirmQuit, cancelQuit })); - - const { listenForQuitRequests } = await import('./quitListener'); - const { quitPrompt } = await import('../stores/quitPrompt.svelte'); - return { listenForQuitRequests, quitPrompt }; -} - -function session(overrides: Partial = {}): ActiveSessionInfo { - return { - sessionId: 's1', - projectId: 'p1', - branchId: 'b1', - sessionType: 'commit', - status: 'running', - ...overrides, - }; -} - -describe('quitListener', () => { - beforeEach(() => { - // The store's runes compile away in the app build; under vitest they stay - // plain global calls, so stub $state as identity (projectsData.test.ts - // precedent). - vi.stubGlobal('$state', (initial: unknown) => initial); - handlers = []; - confirmQuit.mockReset().mockResolvedValue(undefined); - cancelQuit.mockReset().mockResolvedValue(undefined); - unlisten.mockReset(); - listenToWindowEvent.mockReset().mockImplementation((_event, handler) => { - handlers.push(handler as (payload: QuitRequestedPayload) => void); - return unlisten; - }); - }); - - afterEach(() => { - vi.doUnmock('../transport'); - vi.doUnmock('../api/commands'); - vi.unstubAllGlobals(); - vi.restoreAllMocks(); - }); - - it('opens the prompt with the payload from app:quit-requested', async () => { - const { listenForQuitRequests, quitPrompt } = await load(); - listenForQuitRequests(); - - expect(listenToWindowEvent).toHaveBeenCalledWith('app:quit-requested', expect.any(Function)); - expect(quitPrompt.open).toBe(false); - - handlers[0]({ sessions: [session()], runningActionCount: 2 }); - - expect(quitPrompt.open).toBe(true); - expect(quitPrompt.payload?.sessions).toHaveLength(1); - expect(quitPrompt.payload?.runningActionCount).toBe(2); - expect(quitPrompt.stopping).toBe(false); - }); - - it('registers no listener in web mode', async () => { - const { listenForQuitRequests } = await load({ isTauri: false }); - - // Callable no-op, so App.svelte's teardown needs no extra guard. - listenForQuitRequests()(); - - expect(listenToWindowEvent).not.toHaveBeenCalled(); - }); - - it('confirming invokes confirm_quit and leaves the dialog stopping', async () => { - const { listenForQuitRequests, quitPrompt } = await load(); - listenForQuitRequests(); - handlers[0]({ sessions: [session()], runningActionCount: 0 }); - - await quitPrompt.confirm(); - - expect(confirmQuit).toHaveBeenCalledTimes(1); - // The backend exits the process; until it does, the dialog reports progress - // instead of pretending the app is still usable. - expect(quitPrompt.open).toBe(true); - expect(quitPrompt.stopping).toBe(true); - - await quitPrompt.confirm(); - expect(confirmQuit).toHaveBeenCalledTimes(1); - }); - - it('closes the dialog when confirm_quit fails', async () => { - const { listenForQuitRequests, quitPrompt } = await load(); - listenForQuitRequests(); - handlers[0]({ sessions: [session()], runningActionCount: 0 }); - confirmQuit.mockRejectedValueOnce(new Error('nope')); - vi.spyOn(console, 'error').mockImplementation(() => {}); - - await quitPrompt.confirm(); - - expect(quitPrompt.open).toBe(false); - expect(quitPrompt.stopping).toBe(false); - }); - - it('cancelling invokes cancel_quit and closes the dialog', async () => { - const { listenForQuitRequests, quitPrompt } = await load(); - listenForQuitRequests(); - handlers[0]({ sessions: [session()], runningActionCount: 0 }); - - quitPrompt.cancel(); - - expect(cancelQuit).toHaveBeenCalledTimes(1); - expect(quitPrompt.open).toBe(false); - }); - - it('ignores a cancel once the quit is under way', async () => { - const { listenForQuitRequests, quitPrompt } = await load(); - listenForQuitRequests(); - handlers[0]({ sessions: [session()], runningActionCount: 0 }); - - await quitPrompt.confirm(); - quitPrompt.cancel(); - - expect(cancelQuit).not.toHaveBeenCalled(); - expect(quitPrompt.open).toBe(true); - }); -}); diff --git a/apps/staged/src/lib/listeners/quitListener.ts b/apps/staged/src/lib/listeners/quitListener.ts deleted file mode 100644 index f88e6ac90..000000000 --- a/apps/staged/src/lib/listeners/quitListener.ts +++ /dev/null @@ -1,24 +0,0 @@ -/** - * Listener for the backend's `app:quit-requested` event. - * - * `Cmd+Q` / the app-menu Quit item reach `app_lifecycle::request_quit`, which - * emits this event instead of exiting when sessions are still active. The - * backend addresses it to exactly one window with `emit_to`, so this must be a - * *window-scoped* listener — the any-target `listenToEvent` also matches emits - * addressed to other windows, and every window would raise its own dialog. - * Wired at App level so it works on any route, and Tauri-only: quitting is a - * desktop-host action, and the `confirm_quit` command a browser client would - * need is deliberately absent from the web-mode dispatch table. - */ - -import { isTauri, listenToWindowEvent, type UnlistenFn } from '../transport'; -import { quitPrompt } from '../stores/quitPrompt.svelte'; -import type { QuitRequestedPayload } from '../types'; - -export function listenForQuitRequests(): UnlistenFn { - if (!isTauri) return () => {}; - - return listenToWindowEvent('app:quit-requested', (payload) => { - quitPrompt.requested(payload); - }); -} diff --git a/apps/staged/src/lib/stores/quitPrompt.svelte.ts b/apps/staged/src/lib/stores/quitPrompt.svelte.ts deleted file mode 100644 index 0ee0dbec7..000000000 --- a/apps/staged/src/lib/stores/quitPrompt.svelte.ts +++ /dev/null @@ -1,61 +0,0 @@ -/** - * State behind the quit confirmation dialog. - * - * The backend raises `app:quit-requested` when the user quits with sessions - * still active (see `app_lifecycle.rs`); quitListener.ts feeds that payload in - * here and QuitConfirmDialog renders it. Answering is a round trip back to the - * backend: confirming hands off to the shutdown sequence, which stops the - * sessions and then exits the process — so the dialog stays up, in its - * `stopping` state, until the app goes away underneath it. - */ - -import * as commands from '../api/commands'; -import type { QuitRequestedPayload } from '../types'; - -class QuitPromptStore { - private _payload = $state(null); - /** The quit was confirmed and the backend is stopping sessions. */ - private _stopping = $state(false); - - get payload(): QuitRequestedPayload | null { - return this._payload; - } - - get open(): boolean { - return this._payload !== null; - } - - get stopping(): boolean { - return this._stopping; - } - - /** A quit is waiting on the user's answer. */ - requested(payload: QuitRequestedPayload): void { - this._payload = payload; - this._stopping = false; - } - - /** Quit and stop the listed sessions. */ - async confirm(): Promise { - if (this._stopping) return; - this._stopping = true; - try { - await commands.confirmQuit(); - } catch (e) { - // The quit never started, so drop the dialog rather than leaving it stuck - // on "Stopping sessions…" for an app that isn't going anywhere. - console.error('Failed to confirm quit:', e); - this._payload = null; - this._stopping = false; - } - } - - /** Keep running. Also the Esc / click-outside path. */ - cancel(): void { - if (this._stopping) return; - this._payload = null; - void commands.cancelQuit().catch((e) => console.error('Failed to cancel quit:', e)); - } -} - -export const quitPrompt = new QuitPromptStore(); diff --git a/apps/staged/src/lib/types.ts b/apps/staged/src/lib/types.ts index cad11908c..88ef5a3b5 100644 --- a/apps/staged/src/lib/types.ts +++ b/apps/staged/src/lib/types.ts @@ -664,18 +664,6 @@ export interface ActiveSessionInfo { status: SessionStatus; } -/** - * Payload of the `app:quit-requested` event: what a quit would interrupt. - * - * Emitted by `app_lifecycle::request_quit` when the user quits with sessions - * still active. Sessions are what gate the quit; running actions are reported - * so the dialog can say they stop too. - */ -export interface QuitRequestedPayload { - sessions: ActiveSessionInfo[]; - runningActionCount: number; -} - /** * Payload emitted by the `pr-created` domain event when a completed PR * session produced a pull request. The backend has already persisted the PR From 1d0688e98feebc214614f647fd0ccd519ac786b5 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Thu, 3 Sep 2026 15:51:25 +1000 Subject: [PATCH 03/13] fix(lifecycle): re-ask instead of force-quitting on a repeated window close MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The force path out of an unanswered quit prompt keyed off nothing but "a quit request arrived while `prompt_pending` is set", which off macOS put it on the last window's close button — the only interactive quit trigger there is without an app menu. Two quick clicks on X therefore killed running agents with the question never answered, and X is exactly the control people click again when a window doesn't shut: reflex aimed at the window, not an answer to a system alert they may not have noticed (unparented, so it carries no app chrome tying it to the click). `request_quit` now takes a `QuitTrigger` instead of the vestigial `force: bool` that no caller ever set, and only `QuitTrigger::Explicit` — the app menu's Quit item / `Cmd+Q`, and the store-incompatibility screens' Close button — may take the force path. A `QuitTrigger::WindowClose` request falls through to the normal gate, so a repeat close re-raises the alert (a second copy of it if the first is still up, which is a cheap thing to dismiss next to a forced quit) built from a fresh blocker snapshot — and if the sessions have finished in the meantime it just quits, with nothing left to warn about. That keeps the escape hatch where it was justified: a deliberate second `Cmd+Q` still gets out of a prompt that never appeared or never came back, which it must, since the alert is not app-modal. Off macOS the close button loses the hatch and needs no replacement — asking again is itself the way out of an invisible prompt, because clicking X puts an answerable question back on screen. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/app_lifecycle.rs | 87 ++++++++++++++++++---- apps/staged/src-tauri/src/lib.rs | 4 +- 2 files changed, 75 insertions(+), 16 deletions(-) diff --git a/apps/staged/src-tauri/src/app_lifecycle.rs b/apps/staged/src-tauri/src/app_lifecycle.rs index 43f715e61..49b58056b 100644 --- a/apps/staged/src-tauri/src/app_lifecycle.rs +++ b/apps/staged/src-tauri/src/app_lifecycle.rs @@ -11,7 +11,9 @@ //! the window hidden, so sessions keep streaming; the Dock icon //! (`RunEvent::Reopen`) or `Window ▸ Staged` brings it back. Other platforms //! have no Dock/tray to recover a hidden window, so closing the last window -//! still quits there — but through the same confirmation gate as `Cmd+Q`. +//! still quits there — but through the same confirmation gate as `Cmd+Q`, and +//! clicking the close button again asks again instead of forcing the quit (see +//! [`QuitTrigger`]). //! //! **Quitting with sessions running asks first, then stops them cleanly.** //! [`request_quit`] gates on active sessions and asks; [`shutdown_cleanup`] @@ -96,9 +98,11 @@ pub struct QuitState { /// Set by the first caller into [`shutdown_cleanup`], so the cleanup runs /// exactly once however many exit events follow it. quit_in_progress: AtomicBool, - /// Set while a confirmation alert is unanswered. A quit request arriving - /// while it is set forces the quit, so an alert that never appeared or never - /// came back can't trap the app: a second `Cmd+Q` always gets out. + /// Set while a confirmation alert is unanswered. An + /// [`Explicit`](QuitTrigger::Explicit) quit arriving while it is set forces + /// the quit, so an alert that never appeared or never came back can't trap + /// the app: a second `Cmd+Q` always gets out. A repeated window close does + /// not force — it asks again. See [`QuitTrigger`]. prompt_pending: AtomicBool, } @@ -173,9 +177,10 @@ fn on_close_requested(window: &Window, api: &tauri::CloseRequestApi) { // No Dock or tray icon elsewhere, so a hidden window would be unreachable — // closing the last window still quits, with the confirmation gate in front - // of it. + // of it. Clicking the X again re-raises that question rather than forcing + // the quit — see `QuitTrigger::may_force`. #[cfg(not(target_os = "macos"))] - request_quit(app, false); + request_quit(app, QuitTrigger::WindowClose); } /// Hide the window and drop its PR-poll client to the unfocused tier. @@ -242,17 +247,56 @@ fn set_native_focus(app: &AppHandle, window_label: &str, focused: bool) { // Quit gate // ============================================================================= +/// What asked the app to quit. +/// +/// A first request is a first request whatever raised it; the trigger decides +/// what a *repeat* means while the confirmation alert is still unanswered. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum QuitTrigger { + /// A quit aimed at the application: the app menu's Quit item (`Cmd+Q`), or + /// the store-incompatibility screens' Close button. + Explicit, + /// The last window's close button — or `Alt+F4`, which takes the same + /// `CloseRequested` path — off macOS, where closing the last window quits. + WindowClose, +} + +impl QuitTrigger { + /// Whether repeating this trigger while the alert is unanswered quits + /// without an answer. + /// + /// Only an [`Explicit`](Self::Explicit) quit may: a second `Cmd+Q` is a + /// deliberate second attempt at quitting, and it has to get out in case the + /// alert never appeared or never came back. A close button carries no such + /// intent. Off macOS it is the only interactive quit trigger there is, and + /// it is the thing users click twice in a second when a window won't shut — + /// reflex, aimed at the window, not an answer to a question they may not + /// have noticed. Killing running agents on that reflex is the opposite of + /// what the gate exists for. + /// + /// Asking again covers the trapped case just as well, without the damage: + /// clicking the X puts an answerable question back on screen — a second + /// copy of it if the first is still up, which is a cheap thing to dismiss + /// next to a forced quit — and if the sessions have finished in the + /// meantime, the repeat quits with nothing left to warn about. + fn may_force(self) -> bool { + matches!(self, Self::Explicit) + } +} + /// Handle a quit request from the app menu, `Cmd+Q`, or (off macOS) the last /// window's close. Cheap enough for the main thread: it snapshots blockers and /// either hands off to a background quit or raises the alert. -pub fn request_quit(app: &AppHandle, force: bool) { +pub fn request_quit(app: &AppHandle, trigger: QuitTrigger) { let quit_state = app.state::(); - // A quit arriving while the alert is unanswered (a second `Cmd+Q`) is the - // escape hatch from a prompt that never appeared or never came back. The - // system alert isn't app-modal, so that second `Cmd+Q` is still dispatchable - // with the alert on screen. - if force || quit_state.take_prompt() { + // An explicit quit arriving while the alert is unanswered (a second + // `Cmd+Q`) is the escape hatch from a prompt that never appeared or never + // came back. The system alert isn't app-modal, so that second `Cmd+Q` is + // still dispatchable with the alert on screen. A repeated window close + // deliberately doesn't take this path and falls through to the gate below, + // which asks again. + if trigger.may_force() && quit_state.take_prompt() { spawn_quit(app); return; } @@ -270,6 +314,9 @@ pub fn request_quit(app: &AppHandle, force: bool) { return; } + // A prompt already pending here means a repeated window close: the flag is + // already set, so setting it again is a no-op, and the alert goes back up + // describing whatever is running *now*. quit_state.set_prompt_pending(); ask_before_quitting(app, &blockers); } @@ -321,7 +368,7 @@ fn quit_confirmed(result: &MessageDialogResult) -> bool { /// must not be able to terminate the desktop host. #[tauri::command] pub fn quit_app(app_handle: AppHandle) { - request_quit(&app_handle, false); + request_quit(&app_handle, QuitTrigger::Explicit); } /// Run the quit sequence off the main thread so the bounded waits never freeze @@ -639,8 +686,18 @@ mod tests { assert!(!should_prompt(&QuitBlockers::default())); } - /// The pending flag turns the next quit into a force-quit, so answering the - /// alert has to disarm it — otherwise the next `Cmd+Q` quits without asking. + /// A second `Cmd+Q` forces the quit; a second click on a close button asks + /// again, because that click is aimed at the window and is exactly the one + /// users repeat by reflex. + #[test] + fn only_an_explicit_quit_can_force_past_the_prompt() { + assert!(QuitTrigger::Explicit.may_force()); + assert!(!QuitTrigger::WindowClose.may_force()); + } + + /// The pending flag turns the next explicit quit into a force-quit, so + /// answering the alert has to disarm it — otherwise the next `Cmd+Q` quits + /// without asking. #[test] fn answering_the_prompt_disarms_the_force_path() { let state = QuitState::default(); diff --git a/apps/staged/src-tauri/src/lib.rs b/apps/staged/src-tauri/src/lib.rs index 66dee3395..e266d45a4 100644 --- a/apps/staged/src-tauri/src/lib.rs +++ b/apps/staged/src-tauri/src/lib.rs @@ -2227,7 +2227,9 @@ pub fn run() { log::warn!("Failed to open window from menu: {e}"); } } - MenuDispatch::RequestQuit => app_lifecycle::request_quit(app, false), + MenuDispatch::RequestQuit => { + app_lifecycle::request_quit(app, app_lifecycle::QuitTrigger::Explicit) + } MenuDispatch::ShowWindow => app_lifecycle::show_a_window(app), MenuDispatch::Drop => {} } From fdb382d99dd1985510026f2362e892c68c7e0431 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Thu, 3 Sep 2026 15:56:39 +1000 Subject: [PATCH 04/13] fix(lifecycle): make a late shutdown_cleanup caller wait for the one in flight MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `shutdown_cleanup` guarded itself with `quit_in_progress.swap(true)`: the first caller did the work and every later one returned immediately. That's the wrong contract for one of its three callers. Confirming "Quit & Stop Sessions" runs the cleanup on a background thread — up to 2s of bounded waits with *nothing on screen*, since the native alert dismisses on click and the webview's "Stopping sessions…" progress state went with it. An impatient Dock ▸ Quit in that window sends `terminate:`; Tauri fires `RunEvent::Exit` on the main thread inside `applicationWillTerminate:`; the guard saw the flag set and returned instantly; `applicationWillTerminate:` returned and the OS killed the process mid-cancel. The agent CLIs were orphaned — own process groups, and neither an OS kill nor `process::exit` runs the `kill_on_drop` destructors — and the DB sweep was skipped, so the next launch found active rows owned by a dead pid and reported them as errored sessions. Exactly the two failure modes this branch exists to prevent, reachable by a second quit gesture the missing feedback invites. The swap-guard becomes a completion flag behind a mutex, extracted as `QuitState::run_cleanup_once` — holding the lock for the duration of the work is the mechanism, so a late caller parks on `lock()`, then sees `done == true` and returns having waited. On the `Exit` path that holds `applicationWillTerminate:` open until sessions are actually stopped and swept, bounded by the same `SHUTDOWN_BUDGET` that was sized for that context in the first place. `quit_in_progress` stays, unchanged in meaning, because its two readers (`request_quit`'s already-shutting-down early return, `on_close_requested`'s stay-out-of-the-way check) run on the main thread while a cleanup may be in flight and need a non-blocking load. It just stops being the once-guard: the mutex claims, the atomic publishes. Blocking the main thread here is safe. It happens only on the `terminate:`/`ExitRequested` paths, where the app is exiting and repainting no longer matters; the cleanup's waits are satisfied by session threads on the tokio runtime and action process-group signals, none of which need the event loop to turn — established behavior, since a Dock ▸ Quit with no confirmed quit in flight already runs the whole cleanup synchronously inside `applicationWillTerminate:`. `spawn_quit`'s reason for existing (keep the event loop turning during a user-initiated quit) is untouched. Poisoning is taken over with `PoisonError::into_inner` deliberately: if the first caller panicked mid-cleanup, `done` is still `false` and the late caller re-runs the work — the right recovery, since every step is idempotent (cancelling a cancelled session is a no-op, the sweep is a guarded CAS per row). That's why `std::sync::Once` is the wrong tool despite matching blocking semantics: a panicked `call_once` poisons the `Once` and makes every later caller panic, and a panic inside `applicationWillTerminate:` aborts with no cleanup at all. Extracting the latch also gives it test coverage the old guard couldn't have, since `shutdown_cleanup` needs an `AppHandle`: a late caller returning only after the running cleanup finished and without re-running it, a sequential second call being a no-op, `quit_in_progress` reading true from inside the work, and a panicked cleanup being retried. This makes the impatient Dock ▸ Quit safe, not visible — there is still no feedback for up to 2s after confirming a quit, which is what invites the second gesture. Restoring a progress indication without a webview (Dock badge, `NSApp` activity) is a separate decision, deliberately not bundled here. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/app_lifecycle.rs | 200 +++++++++++++++++---- 1 file changed, 168 insertions(+), 32 deletions(-) diff --git a/apps/staged/src-tauri/src/app_lifecycle.rs b/apps/staged/src-tauri/src/app_lifecycle.rs index 49b58056b..074fb37a9 100644 --- a/apps/staged/src-tauri/src/app_lifecycle.rs +++ b/apps/staged/src-tauri/src/app_lifecycle.rs @@ -36,10 +36,11 @@ //! there is no longer any state where a quit can't ask. //! //! Every exit path funnels into [`shutdown_cleanup`], which runs its work at -//! most once — a confirmed quit calls it directly, `RunEvent::ExitRequested` -//! covers programmatic exits, and `RunEvent::Exit` is the only hook on the -//! `NSApp terminate:` path (Dock ▸ Quit, logout), which never emits -//! `ExitRequested`. +//! most once and makes later callers wait for it to finish — a confirmed quit +//! calls it directly, `RunEvent::ExitRequested` covers programmatic exits, and +//! `RunEvent::Exit` is the only hook on the `NSApp terminate:` path (Dock ▸ +//! Quit, logout), which never emits `ExitRequested`. That wait is what keeps an +//! impatient second quit gesture from cutting a cleanup already in flight short. use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; @@ -95,8 +96,15 @@ const ACTION_FORCE_KILL_AFTER: Duration = Duration::from_secs(1); /// Quit bookkeeping, managed as Tauri state. #[derive(Default)] pub struct QuitState { - /// Set by the first caller into [`shutdown_cleanup`], so the cleanup runs - /// exactly once however many exit events follow it. + /// Held for the duration of [`shutdown_cleanup`]'s work; `true` once it has + /// completed. The lock is what makes a late caller *wait* rather than skip: + /// see [`run_cleanup_once`](Self::run_cleanup_once). + cleanup_done: Mutex, + /// Cheap "a shutdown is under way" signal, published by the caller that + /// claims the cleanup. Separate from `cleanup_done` because its readers + /// ([`request_quit`], [`on_close_requested`]) run on the main thread while a + /// cleanup may be in flight and must not block on the lock. The mutex + /// claims; this atomic publishes. quit_in_progress: AtomicBool, /// Set while a confirmation alert is unanswered. An /// [`Explicit`](QuitTrigger::Explicit) quit arriving while it is set forces @@ -115,6 +123,41 @@ impl QuitState { fn take_prompt(&self) -> bool { self.prompt_pending.swap(false, Ordering::SeqCst) } + + /// Run `cleanup` at most once. A caller that arrives while it is already + /// running **blocks until it has finished**, then returns without re-running + /// it. + /// + /// Waiting, rather than returning early, is the point. A confirmed quit runs + /// the cleanup on a background thread with nothing on screen for up to + /// [`SHUTDOWN_BUDGET`], and a `terminate:` arriving in that window (an + /// impatient Dock ▸ Quit) reaches [`shutdown_cleanup`] via `RunEvent::Exit`, + /// on the main thread, inside `applicationWillTerminate:`. Returning there + /// would let that method return and the OS kill the process mid-cancel — + /// orphaning the agent CLIs (own process groups, `kill_on_drop` destructors + /// an OS kill never runs) and skipping the DB sweep. Holding it open until + /// the work is done is what those two guarantees need. + /// + /// A poisoned lock is taken over rather than propagated: if the first caller + /// panicked mid-cleanup, `done` is still `false` and the late caller re-runs + /// the work, which is the right recovery given every step is idempotent + /// (cancelling a cancelled session is a no-op; the sweep is a guarded CAS + /// per row). That's also why this isn't a [`std::sync::Once`] despite the + /// matching blocking semantics — a panicked `call_once` poisons the `Once` + /// and makes every later caller panic, and a panic inside + /// `applicationWillTerminate:` aborts with no cleanup at all. + fn run_cleanup_once(&self, cleanup: impl FnOnce()) { + let mut done = self + .cleanup_done + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if *done { + return; + } + self.quit_in_progress.store(true, Ordering::SeqCst); + cleanup(); + *done = true; + } } /// What a quit would interrupt, as the alert describes it. @@ -501,38 +544,39 @@ fn quit_prompt_message(blockers: &QuitBlockers) -> String { // Shutdown cleanup // ============================================================================= -/// Stop everything this process owns. Idempotent — the first caller does the -/// work, later ones return immediately. +/// Stop everything this process owns. Runs its work at most once — the first +/// caller does it, and a caller arriving while it runs waits for it to finish +/// (see [`QuitState::run_cleanup_once`]) rather than returning to an exit that +/// would kill the process mid-cleanup. pub fn shutdown_cleanup(app: &AppHandle) { let Some(quit_state) = app.try_state::() else { return; }; - if quit_state.quit_in_progress.swap(true, Ordering::SeqCst) { - return; - } - // Signal both kinds of work before waiting on either, so they shut down in - // parallel inside one shared budget instead of one after the other. - let session_ids = cancel_owned_sessions(app); - let execution_ids = stop_running_actions(app); - - let deadline = Instant::now() + SHUTDOWN_BUDGET; - if !session_ids.is_empty() && !wait_for_sessions(app, &session_ids, deadline) { - log::warn!( - "Timed out waiting for {} session(s) to stop during app shutdown", - session_ids.len() - ); - } - if !execution_ids.is_empty() && !wait_for_actions(app, &execution_ids, deadline) { - log::warn!( - "Timed out waiting for {} action(s) to stop during app shutdown", - execution_ids.len() - ); - } + quit_state.run_cleanup_once(|| { + // Signal both kinds of work before waiting on either, so they shut down + // in parallel inside one shared budget instead of one after the other. + let session_ids = cancel_owned_sessions(app); + let execution_ids = stop_running_actions(app); + + let deadline = Instant::now() + SHUTDOWN_BUDGET; + if !session_ids.is_empty() && !wait_for_sessions(app, &session_ids, deadline) { + log::warn!( + "Timed out waiting for {} session(s) to stop during app shutdown", + session_ids.len() + ); + } + if !execution_ids.is_empty() && !wait_for_actions(app, &execution_ids, deadline) { + log::warn!( + "Timed out waiting for {} action(s) to stop during app shutdown", + execution_ids.len() + ); + } - // Last, so the rows reflect whatever the session threads managed to write - // for themselves first. - sweep_active_sessions(app); + // Last, so the rows reflect whatever the session threads managed to + // write for themselves first. + sweep_active_sessions(app); + }); } /// Cancel every session this process is running, recording `AppQuit` as the @@ -653,6 +697,7 @@ fn app_store(app: &AppHandle) -> Option> { mod tests { use super::*; use std::path::Path; + use std::sync::Barrier; fn active_session(session_type: Option<&str>, status: SessionStatus) -> ActiveSessionInfo { ActiveSessionInfo { @@ -709,6 +754,97 @@ mod tests { assert!(!state.take_prompt(), "prompt stayed armed after answering"); } + /// The `RunEvent::Exit` case: a `terminate:` arriving while a confirmed + /// quit's cleanup is still running must hold `applicationWillTerminate:` + /// open until that cleanup finishes, not return to an exit that kills the + /// process mid-cancel. + #[test] + fn a_late_caller_waits_for_the_running_cleanup() { + let state = Arc::new(QuitState::default()); + // Trips inside the first closure, so passing it proves the first caller + // holds the lock before the test thread tries to take it. + let entered = Arc::new(Barrier::new(2)); + let finished = Arc::new(AtomicBool::new(false)); + + let first = { + let state = Arc::clone(&state); + let entered = Arc::clone(&entered); + let finished = Arc::clone(&finished); + std::thread::spawn(move || { + state.run_cleanup_once(|| { + entered.wait(); + // Stands in for the bounded waits, so the second caller has + // to block rather than happening to arrive after the fact. + std::thread::sleep(Duration::from_millis(50)); + finished.store(true, Ordering::SeqCst); + }); + }) + }; + + entered.wait(); + let second_ran = AtomicBool::new(false); + state.run_cleanup_once(|| second_ran.store(true, Ordering::SeqCst)); + + assert!( + finished.load(Ordering::SeqCst), + "second call returned before the running cleanup had finished" + ); + assert!( + !second_ran.load(Ordering::SeqCst), + "second call re-ran the cleanup instead of waiting for it" + ); + first.join().unwrap(); + } + + #[test] + fn a_sequential_second_call_does_not_run_the_cleanup_again() { + let state = QuitState::default(); + let runs = AtomicBool::new(false); + + state.run_cleanup_once(|| runs.store(true, Ordering::SeqCst)); + runs.store(false, Ordering::SeqCst); + state.run_cleanup_once(|| runs.store(true, Ordering::SeqCst)); + + assert!(!runs.load(Ordering::SeqCst)); + } + + /// `request_quit`'s "already shutting down" early return and + /// `on_close_requested`'s stay-out-of-the-way check both read this from the + /// main thread while a cleanup is in flight, so it has to be published + /// before the work starts rather than after it. + #[test] + fn quit_in_progress_is_published_while_the_cleanup_runs() { + let state = QuitState::default(); + assert!(!state.quit_in_progress.load(Ordering::SeqCst)); + + state.run_cleanup_once(|| { + assert!( + state.quit_in_progress.load(Ordering::SeqCst), + "shutdown was under way but the flag still read false" + ); + }); + + assert!(state.quit_in_progress.load(Ordering::SeqCst)); + } + + /// Poisoning is taken over rather than propagated, so a cleanup that + /// panicked half-done gets re-run by the next exit event — every step of it + /// is idempotent. + #[test] + fn a_panicked_cleanup_is_retried() { + let state = Arc::new(QuitState::default()); + + let panicked = { + let state = Arc::clone(&state); + std::thread::spawn(move || state.run_cleanup_once(|| panic!("cleanup blew up"))) + }; + assert!(panicked.join().is_err(), "expected the panic to propagate"); + + let retried = AtomicBool::new(false); + state.run_cleanup_once(|| retried.store(true, Ordering::SeqCst)); + assert!(retried.load(Ordering::SeqCst)); + } + /// The ok slot is the default (`Return`) button, so it holds "Keep Running" /// and the cancel slot holds the destructive answer. #[test] From 199ad9208668ea7a1503f73ec48d2e04cbc9c05f Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Thu, 3 Sep 2026 16:02:24 +1000 Subject: [PATCH 05/13] fix(lifecycle): make the quit sweep's cancel CAS ownership-aware MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `shutdown_cleanup` ends by sweeping whatever is still active in the DB to cancelled/`app_quit`, so the next launch doesn't find rows owned by a dead process and report them as errored sessions. The snapshot it works from, `owned_active_sessions`, filters on ownership — running rows stamped with our `owner_pid`, plus all queued rows, which carry no owner and count as ours by convention. The per-row CAS behind it did not: `transition_from_active` guarded only on `status IN ('queued', 'running')`, and that ownership filter doesn't survive the gap between the two. A review of `941614d3` flagged the interleaving. Another instance pointed at the same data dir drains its queue mid-sweep and claims a row via `transition_queued_to_running` — one statement that sets `status='running'` and stamps *its* pid. The row is now that instance's live session, but its status still matches, so our CAS succeeds and cancels another process's in-flight work. The review suggested splitting by snapshotted status — a `queued`-only CAS for rows snapshotted as queued. That makes the *other instance's* claim mutually exclusive with the sweep, but it isn't the only claimant inside the window: the session completion handler drains the branch queue on any terminal transition, including the cancels `shutdown_cleanup` just issued, so this process's own drain can claim a queued row too. Under the status split that claim also fails the narrow CAS, leaving the row `running` under our pid as we exit — precisely the artifact the sweep exists to prevent. The distinction that matters isn't "queued or running at snapshot time", it's "ours or not ours at write time", so the ownership check moves into the statement: WHERE id = ?5 AND (status = 'queued' OR (status = 'running' AND owner_pid = ?6)) `transition_from_active` becomes `transition_from_owned_active` — the sweep is its only caller, so it could change shape freely — and both interleavings now resolve in one predicate. Another instance's claim stamps a different pid and fails it; our own drain's claim stamps ours and still matches, which is right, because this process is exiting and taking that session with it. That argument surfaces the same-process sibling, fixed here too: `drain_queued_sessions_for_branch` now returns early once a shutdown is claimed (`QuitState::is_quitting`, extracted from the `quit_in_progress` load `on_close_requested` already did). Without it, shutdown feeds itself — its cancels trigger drains that claim queued rows and spawn fresh agent children which `app.exit(0)` then orphans, since they run in their own process groups and an exit runs no `kill_on_drop` destructors. The sweep would put those DB rows right; nothing would put the processes right. Unchanged: queued rows still count as every instance's own, so a quit still prompts about and cancels unclaimed queued work. This narrows exactly one thing — a row that stops being queued because someone *else* claimed it mid-sweep. The sweep's snapshot-and-CAS body is extracted as `sweep_sessions(&Store)`, so the tests that used to inline a copy of the loop exercise the real one, and a new case asserts a running row owned by another pid survives it untouched. At the store level, where the CAS fires exactly as it does inside the race window: that row returns `false` and keeps its status, alongside positives for a queued row and for a running row carrying our pid. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/app_lifecycle.rs | 119 ++++++++++++------ apps/staged/src-tauri/src/session_commands.rs | 13 ++ apps/staged/src-tauri/src/store/sessions.rs | 26 +++- apps/staged/src-tauri/src/store/tests.rs | 76 ++++++++++- 4 files changed, 184 insertions(+), 50 deletions(-) diff --git a/apps/staged/src-tauri/src/app_lifecycle.rs b/apps/staged/src-tauri/src/app_lifecycle.rs index 074fb37a9..27be1faec 100644 --- a/apps/staged/src-tauri/src/app_lifecycle.rs +++ b/apps/staged/src-tauri/src/app_lifecycle.rs @@ -20,7 +20,10 @@ //! cancels sessions with [`CompletionReason::AppQuit`] and stops actions. That //! cancel is the only thing that shuts an agent down: ACP children are spawned //! with `process_group(0)` and `kill_on_drop`, and `process::exit` runs no -//! destructors, so a bare exit leaves the agent CLIs running. +//! destructors, so a bare exit leaves the agent CLIs running. Once a shutdown is +//! claimed the branch queue stops draining ([`QuitState::is_quitting`]), so +//! those cancels — terminal transitions like any other — can't feed fresh agent +//! children into an exit that would orphan them. //! //! The question is asked by a native alert with **no parent window**, not by a //! dialog in a webview. Quitting is scoped to the application, and the case the @@ -102,9 +105,9 @@ pub struct QuitState { cleanup_done: Mutex, /// Cheap "a shutdown is under way" signal, published by the caller that /// claims the cleanup. Separate from `cleanup_done` because its readers - /// ([`request_quit`], [`on_close_requested`]) run on the main thread while a - /// cleanup may be in flight and must not block on the lock. The mutex - /// claims; this atomic publishes. + /// ([`request_quit`] and [`on_close_requested`] on the main thread, the + /// queue drain on the tokio runtime) run while a cleanup may be in flight + /// and must not block on the lock. The mutex claims; this atomic publishes. quit_in_progress: AtomicBool, /// Set while a confirmation alert is unanswered. An /// [`Explicit`](QuitTrigger::Explicit) quit arriving while it is set forces @@ -115,6 +118,15 @@ pub struct QuitState { } impl QuitState { + /// Whether a shutdown has been claimed. Non-blocking by construction — see + /// [`quit_in_progress`](Self::quit_in_progress). + /// + /// Read by the queue drain to stop starting new work mid-shutdown, and by + /// [`on_close_requested`] to stay out of the way of the exit's own closes. + pub(crate) fn is_quitting(&self) -> bool { + self.quit_in_progress.load(Ordering::SeqCst) + } + fn set_prompt_pending(&self) { self.prompt_pending.store(true, Ordering::SeqCst); } @@ -198,7 +210,7 @@ fn on_close_requested(window: &Window, api: &tauri::CloseRequestApi) { // Mid-shutdown, closes are the exit tearing windows down — stay out of the // way. if let Some(quit_state) = app.try_state::() { - if quit_state.quit_in_progress.load(Ordering::SeqCst) { + if quit_state.is_quitting() { return; } } @@ -347,7 +359,7 @@ pub fn request_quit(app: &AppHandle, trigger: QuitTrigger) { // Already shutting down, and the alert dismissed on click while cleanup runs // out its budget — so there is nothing on screen saying so, and a `Cmd+Q` // here means "I already answered", not "ask me again". - if quit_state.quit_in_progress.load(Ordering::SeqCst) { + if quit_state.is_quitting() { return; } @@ -640,28 +652,39 @@ fn sweep_active_sessions(app: &AppHandle) { return; }; - let swept = owned_active_sessions(&store) + let swept = sweep_sessions(&store); + if swept > 0 { + log::info!("Marked {swept} session(s) cancelled (app_quit) during shutdown"); + } +} + +/// The sweep itself: snapshot what we own, then CAS each row to cancelled. +/// Returns how many rows it actually moved. +fn sweep_sessions(store: &Store) -> usize { + let owner_pid = std::process::id(); + + owned_active_sessions(store) .iter() .filter(|session| { - // Guarded CAS per row: a session thread that wrote its own terminal - // status while we were waiting keeps that status. + // Guarded CAS per row, on liveness *and* ownership. A session thread + // that wrote its own terminal status while we were waiting keeps + // that status; so does a row another instance claimed since the + // snapshot, which stamped its pid in the same statement that took + // the row off `queued`. store - .transition_from_active( + .transition_from_owned_active( &session.id, SessionStatus::Cancelled, None, Some(&CompletionReason::AppQuit), + owner_pid, ) .unwrap_or_else(|e| { log::warn!("Failed to cancel session {} on quit: {e}", session.id); false }) }) - .count(); - - if swept > 0 { - log::info!("Marked {swept} session(s) cancelled (app_quit) during shutdown"); - } + .count() } /// Running and queued sessions **this process owns**. @@ -671,6 +694,10 @@ fn sweep_active_sessions(app: &AppHandle) { /// nor cancel another instance's work. Queued rows carry no owner yet, so they /// count as ours: claiming one (`transition_queued_to_running`) stamps a pid /// atomically, which is what takes another instance's claim out of this set. +/// +/// A claim can also land *after* this snapshot, which is why the sweep's CAS +/// (`transition_from_owned_active`) re-checks the same ownership rule at write +/// time rather than trusting the list this returns. fn owned_active_sessions(store: &Store) -> Vec { let sessions = match store.get_active_sessions() { Ok(sessions) => sessions, @@ -808,23 +835,23 @@ mod tests { assert!(!runs.load(Ordering::SeqCst)); } - /// `request_quit`'s "already shutting down" early return and - /// `on_close_requested`'s stay-out-of-the-way check both read this from the - /// main thread while a cleanup is in flight, so it has to be published - /// before the work starts rather than after it. + /// `request_quit`'s "already shutting down" early return, + /// `on_close_requested`'s stay-out-of-the-way check, and the queue drain's + /// gate all read this while a cleanup is in flight, so it has to be + /// published before the work starts rather than after it. #[test] fn quit_in_progress_is_published_while_the_cleanup_runs() { let state = QuitState::default(); - assert!(!state.quit_in_progress.load(Ordering::SeqCst)); + assert!(!state.is_quitting()); state.run_cleanup_once(|| { assert!( - state.quit_in_progress.load(Ordering::SeqCst), + state.is_quitting(), "shutdown was under way but the flag still read false" ); }); - assert!(state.quit_in_progress.load(Ordering::SeqCst)); + assert!(state.is_quitting()); } /// Poisoning is taken over rather than propagated, so a cleanup that @@ -985,16 +1012,7 @@ mod tests { let queued = Session::new_queued("queued"); store.create_session(&queued).unwrap(); - for session in owned_active_sessions(&store) { - assert!(store - .transition_from_active( - &session.id, - SessionStatus::Cancelled, - None, - Some(&CompletionReason::AppQuit), - ) - .unwrap()); - } + assert_eq!(sweep_sessions(&store), 2); for id in [&running.id, &queued.id] { let session = store.get_session(id).unwrap().unwrap(); @@ -1019,14 +1037,7 @@ mod tests { .unwrap(); assert!(owned_active_sessions(&store).is_empty()); - assert!(!store - .transition_from_active( - &completed.id, - SessionStatus::Cancelled, - None, - Some(&CompletionReason::AppQuit), - ) - .unwrap()); + assert_eq!(sweep_sessions(&store), 0); let session = store.get_session(&completed.id).unwrap().unwrap(); assert_eq!(session.status, SessionStatus::Completed); @@ -1035,4 +1046,32 @@ mod tests { Some(CompletionReason::TurnComplete) ); } + + /// Another instance's live session survives our quit. The snapshot already + /// filters it out; this asserts the CAS does too, which is what covers a + /// claim landing *after* the snapshot — the interleaving the sweep can't + /// otherwise see. + #[test] + fn sweep_leaves_another_instances_running_session_alone() { + let store = Store::in_memory().unwrap(); + + let mut theirs = Session::new_running("theirs", Path::new("/tmp")); + theirs.owner_pid = Some(std::process::id().wrapping_add(1)); + store.create_session(&theirs).unwrap(); + + assert_eq!(sweep_sessions(&store), 0); + assert!(!store + .transition_from_owned_active( + &theirs.id, + SessionStatus::Cancelled, + None, + Some(&CompletionReason::AppQuit), + std::process::id(), + ) + .unwrap()); + + let session = store.get_session(&theirs.id).unwrap().unwrap(); + assert_eq!(session.status, SessionStatus::Running); + assert_eq!(session.completion_reason, None); + } } diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index 5002e064b..4cdf5fd3d 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -3116,6 +3116,19 @@ pub async fn drain_queued_sessions_for_branch( branch_id: String, provider: Option, ) -> Result { + // A quit cancels every running session, and every terminal transition drains + // the branch queue — including those cancels. Left ungated, shutdown feeds + // itself: it claims queued rows and spawns fresh agent children that + // `app.exit(0)` then orphans, since they run in their own process groups and + // an exit runs no `kill_on_drop` destructors. The sweep would put the DB rows + // right; nothing would put the processes right. + if app_handle + .try_state::() + .is_some_and(|quit_state| quit_state.is_quitting()) + { + return Ok(false); + } + let queued = store .get_queued_sessions_for_branch(&branch_id) .map_err(|e| e.to_string())?; diff --git a/apps/staged/src-tauri/src/store/sessions.rs b/apps/staged/src-tauri/src/store/sessions.rs index ba0c84eb0..12e2e9639 100644 --- a/apps/staged/src-tauri/src/store/sessions.rs +++ b/apps/staged/src-tauri/src/store/sessions.rs @@ -110,17 +110,31 @@ impl Store { Ok(rows > 0) } - /// Transition session status only if it is currently `queued` or `running`. + /// Transition session status only if the row is still active *and ours*: + /// `queued`, or `running` under `owner_pid`. /// /// Returns `true` if the row was updated, `false` if the session already - /// moved to another state or didn't exist. This is the safe path for - /// cancelling work that may still be in the queue. - pub fn transition_from_active( + /// moved to another state, is running under a different pid, or didn't + /// exist. This is the app-quit sweep's path, and the guard mirrors the + /// filter that snapshot uses: the store is shared with any other Staged + /// instance pointed at the same data dir, queued rows carry no owner and + /// count as the caller's by convention, and running rows must carry the + /// caller's pid. + /// + /// Checking ownership *here* rather than only in the snapshot is what closes + /// the gap between the two. A claim landing in that gap + /// (`transition_queued_to_running`) flips a queued row to `running` and + /// stamps a pid atomically: another instance's claim now fails this + /// predicate and its live work is left alone, while our own drain's claim + /// still matches and is swept — right, because this process is exiting and + /// taking the session with it. + pub fn transition_from_owned_active( &self, id: &str, new_status: SessionStatus, error_message: Option<&str>, completion_reason: Option<&CompletionReason>, + owner_pid: u32, ) -> Result { let conn = self.conn.lock().unwrap(); let error_msg = if new_status == SessionStatus::Error { @@ -130,8 +144,8 @@ impl Store { }; let rows = conn.execute( "UPDATE sessions SET status = ?1, error_message = ?2, completion_reason = ?3, updated_at = ?4 - WHERE id = ?5 AND status IN ('queued', 'running')", - params![new_status.as_str(), error_msg, completion_reason.map(|r| r.as_str()), now_timestamp(), id], + WHERE id = ?5 AND (status = 'queued' OR (status = 'running' AND owner_pid = ?6))", + params![new_status.as_str(), error_msg, completion_reason.map(|r| r.as_str()), now_timestamp(), id, owner_pid], )?; Ok(rows > 0) } diff --git a/apps/staged/src-tauri/src/store/tests.rs b/apps/staged/src-tauri/src/store/tests.rs index 40fe8d344..21f7debbe 100644 --- a/apps/staged/src-tauri/src/store/tests.rs +++ b/apps/staged/src-tauri/src/store/tests.rs @@ -495,23 +495,85 @@ fn test_transition_from_running_keeps_message_for_cancelled() { } #[test] -fn test_transition_from_active_succeeds_when_queued() { +fn test_transition_from_owned_active_succeeds_when_queued() { let store = Store::in_memory().unwrap(); + // Queued rows carry no owner, and count as the caller's by convention. let session = Session::new_queued("queued"); store.create_session(&session).unwrap(); let transitioned = store - .transition_from_active(&session.id, SessionStatus::Cancelled, None, None) + .transition_from_owned_active( + &session.id, + SessionStatus::Cancelled, + None, + None, + std::process::id(), + ) + .unwrap(); + assert!(transitioned); + + let final_state = store.get_session(&session.id).unwrap().unwrap(); + assert_eq!(final_state.status, SessionStatus::Cancelled); +} + +#[test] +fn test_transition_from_owned_active_succeeds_for_our_running_session() { + let store = Store::in_memory().unwrap(); + + let session = Session::new_running("ours", Path::new("/tmp")); + store.create_session(&session).unwrap(); + + let transitioned = store + .transition_from_owned_active( + &session.id, + SessionStatus::Cancelled, + None, + Some(&CompletionReason::AppQuit), + std::process::id(), + ) .unwrap(); assert!(transitioned); let final_state = store.get_session(&session.id).unwrap().unwrap(); assert_eq!(final_state.status, SessionStatus::Cancelled); + assert_eq!( + final_state.completion_reason, + Some(CompletionReason::AppQuit) + ); } +/// The race this guard exists for: between the quit sweep's snapshot and its +/// write, another instance can claim a queued row — one statement that flips it +/// to `running` and stamps *its* pid. Liveness alone still matches, so an +/// ownership-blind CAS would cancel that instance's live session. #[test] -fn test_transition_from_active_does_not_overwrite_completed_session() { +fn test_transition_from_owned_active_leaves_another_instances_session_alone() { + let store = Store::in_memory().unwrap(); + + // What that claim leaves behind: running, with someone else's pid on it. + let mut theirs = Session::new_running("theirs", Path::new("/tmp")); + theirs.owner_pid = Some(std::process::id().wrapping_add(1)); + store.create_session(&theirs).unwrap(); + + let transitioned = store + .transition_from_owned_active( + &theirs.id, + SessionStatus::Cancelled, + None, + Some(&CompletionReason::AppQuit), + std::process::id(), + ) + .unwrap(); + assert!(!transitioned); + + let final_state = store.get_session(&theirs.id).unwrap().unwrap(); + assert_eq!(final_state.status, SessionStatus::Running); + assert_eq!(final_state.completion_reason, None); +} + +#[test] +fn test_transition_from_owned_active_does_not_overwrite_completed_session() { let store = Store::in_memory().unwrap(); let session = Session::new_running("completed first", Path::new("/tmp")); @@ -521,7 +583,13 @@ fn test_transition_from_active_does_not_overwrite_completed_session() { .unwrap(); let transitioned = store - .transition_from_active(&session.id, SessionStatus::Cancelled, None, None) + .transition_from_owned_active( + &session.id, + SessionStatus::Cancelled, + None, + None, + std::process::id(), + ) .unwrap(); assert!(!transitioned); From a34a9e75f65ae92d49f6f6703207cbc4dd5777e2 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Thu, 3 Sep 2026 17:40:31 +1000 Subject: [PATCH 06/13] fix(lifecycle): re-check the quit gate at the drain's claim and spawn points MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shutdown-feeds-itself gate on `drain_queued_sessions_for_branch` was one `is_quitting` load at the top of the function, and a review of `abc299a9` flagged what that leaves open: the drain awaits between claims, so a drain already past the gate when a shutdown is claimed — one triggered by a session finishing naturally just as the user confirms the quit — can still claim a queued row and spawn a fresh agent child after `cancel_owned_sessions` has snapshotted the registry. The ownership-aware sweep CAS puts the DB row right (our pid matches), but `app.exit(0)` orphans the process: agent CLIs run in their own process groups, and an exit runs no `kill_on_drop` destructors. The gate is re-checked at the two points that matter, through a shared `app_lifecycle::is_quitting(app)` extracted from the entry check's inline try_state dance: - In the drain loop, immediately before each start — the review's suggestion. For the commit- and git-pipeline paths this is as tight as the check can get: their bodies run synchronously from dispatch to `start_pipeline_session` registering the session's cancellation token, so nothing can interleave past it. - In `start_queued_session_for_branch`, after the claim and immediately before `start_session` — because for agent sessions the loop check is not the last chance. Real awaits sit between the claim and the spawn (`review_tip_sha`, `commit_pre_head_sha`, the remote workdir resolve, and context building before the claim), room for a whole shutdown to start. Past this check the path is synchronous and `start_session` registers the cancellation token before any child spawns. Bailing after the claim is deliberate, and is what `abc299a9` bought: a row left `running` under our pid is exactly the shape the sweep's ownership-aware CAS moves to cancelled/app_quit, and the claim succeeding proves the sweep hasn't processed that row yet — the CAS and the claim guard on the same `status = 'queued'`. The bail also lands before the `session-status-changed: running` emit, so no client hears about a session that will never run. This shrinks the window, not closes it: `run_cleanup_once` publishes `quit_in_progress` before the shutdown snapshots the registry, so the losing interleaving narrows to a shutdown doing both inside the few synchronous statements between the final load and the token registration. Closing it outright would mean shutdown-side changes (a cancel that re-snapshots, or deferring cancels by DB snapshot via `cancel_or_defer`) — a bigger contract change than this finding calls for. No new test: the gates read managed app state, and the drain chain takes the concrete Wry `AppHandle` (through `start_session`'s `SessionConfig`), so a mock-runtime harness can't reach them without genericizing the whole spawn path — the same reason the entry gate landed with store-level tests only. Verified with `just check-all`. Co-Authored-By: Claude Fable 5 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/app_lifecycle.rs | 7 +++++ apps/staged/src-tauri/src/session_commands.rs | 27 ++++++++++++++++--- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/apps/staged/src-tauri/src/app_lifecycle.rs b/apps/staged/src-tauri/src/app_lifecycle.rs index 27be1faec..d7eeb7b07 100644 --- a/apps/staged/src-tauri/src/app_lifecycle.rs +++ b/apps/staged/src-tauri/src/app_lifecycle.rs @@ -172,6 +172,13 @@ impl QuitState { } } +/// [`QuitState::is_quitting`] from anywhere that holds an `AppHandle`. `false` +/// when the state isn't managed (mock apps in tests). +pub(crate) fn is_quitting(app: &AppHandle) -> bool { + app.try_state::() + .is_some_and(|quit_state| quit_state.is_quitting()) +} + /// What a quit would interrupt, as the alert describes it. #[derive(Debug, Default)] struct QuitBlockers { diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index 4cdf5fd3d..fe19ef864 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -3122,10 +3122,7 @@ pub async fn drain_queued_sessions_for_branch( // `app.exit(0)` then orphans, since they run in their own process groups and // an exit runs no `kill_on_drop` destructors. The sweep would put the DB rows // right; nothing would put the processes right. - if app_handle - .try_state::() - .is_some_and(|quit_state| quit_state.is_quitting()) - { + if crate::app_lifecycle::is_quitting(&app_handle) { return Ok(false); } @@ -3146,6 +3143,13 @@ pub async fn drain_queued_sessions_for_branch( break; } + // The entry gate, re-checked before each claim: every start below + // awaits, so a shutdown claimed mid-drain would otherwise keep being + // fed the rows the remaining iterations were about to take. + if crate::app_lifecycle::is_quitting(&app_handle) { + break; + } + let started = start_queued_session_for_branch( Arc::clone(&store), Arc::clone(®istry), @@ -3367,6 +3371,21 @@ async fn start_queued_session_for_branch( .get_image_ids_for_session(&session_id) .unwrap_or_default(); + // Last look before the spawn. The gates above leave real awaits between + // themselves and `start_session` (context building, `review_tip_sha`, + // `commit_pre_head_sha`, the remote workdir resolve) — room for a whole + // shutdown to start. From here the path is synchronous, and + // `start_session` registers the session's cancellation token before any + // child spawns, so the losing interleaving narrows to a shutdown + // publishing `quit_in_progress` and snapshotting the registry inside the + // few statements between this load and that registration. Bailing here + // deliberately strands the claim: a row left `running` under our pid is + // exactly what the sweep's ownership-aware CAS puts right, while a + // spawned child is what nothing puts right. + if crate::app_lifecycle::is_quitting(&app_handle) { + return Ok(false); + } + let session_type_str = match session_type { BranchSessionType::Commit => "commit", BranchSessionType::Note => "note", From f731f971e4e9cc92932029913c91f4759929ca6d Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 4 Sep 2026 08:53:57 +1000 Subject: [PATCH 07/13] fix(lifecycle): register a session's token before its slow driver construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The post-claim quit gate in `start_queued_session_for_branch` claimed the path from its final `is_quitting` load to the registration of the session's cancellation token was synchronous and a few statements wide. A review of `5d6d7441` showed how far that undersold the agent-session case: `start_session` constructed its `AcpDriver` before calling `registry.register`, and construction resolves the agent binary through `doctor::resolve`'s login-shell probes — a 10s timeout per shell, and with no provider pinned, `discover_providers()` probes every known agent. Seconds of wall clock sat inside the claimed "few statements": room for a shutdown on another thread to publish `quit_in_progress`, snapshot a registry that doesn't contain the session, and reach `app.exit(0)` — after which the child spawns unregistered and nothing stops it. The pipeline paths never had the problem: `start_pipeline_session` registers at entry, before any slow work. Of the two shapes the review offered — register at entry, or move the last-look `is_quitting` load next to the registration — registering at entry is the one that covers the slow phase rather than stepping over it. A last look after construction would leave the constructing session invisible to the shutdown snapshot for those seconds, and its bail would sit inside a function whose direct callers write an `Err` up as an *errored* session — the next-launch artifact this branch exists to prevent — where the quit sweep would have written cancelled/app_quit. Registered at entry, the slow phase is visible: shutdown's cancel fires the session's token, the connect path checks it before protocol setup (`run_acp_session` selects on the token ahead of `initialize`), and the child it spawned is torn down by the connection task's `graceful_stop` — a stop `wait_for_sessions` holds the exit open for, since the session now deregisters through its normal terminal path with the recorded AppQuit reason. The mechanism is `SessionRegistry::register_for_startup`: register, run the fallible startup (driver construction and the user-message persist), and deregister on the way out of a failure — the session thread that normally deregisters never spawns on that path, and a stale entry would hold `wait_for_sessions` open for its full budget and misreport `is_running`. This closes the same gap for ordinary cancels, not just quits: a user cancel landing during construction used to take `cancel_session_impl`'s `!was_running` branch — a DB write the thread about to spawn would then contradict by running the agent anyway. Now it reaches the token. The pipeline handoff's reliance on `register()` replacing the pipeline's token without a tokenless gap (`PipelineOutcome::HandedOffToAi`) is preserved; the replacement just happens before the slow work instead of after it. The oversold comment at the drain's post-claim gate is rewritten to match: the losing interleaving is now genuinely a shutdown publishing and snapshotting inside the few statements between that load and the registration — the status emit and the call itself — a window this change shrinks, not closes, and the gate's bail-strands-the-claim reasoning stands unchanged. New tests cover the latch at the registry level, where no AppHandle is needed: a cancel arriving mid-startup fires the token and records its reason while the session stays registered, and a failed startup deregisters so a shutdown doesn't wait on a session that can never stop. Verified with `just check-all`. Co-Authored-By: Claude Fable 5 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/session_commands.rs | 20 +- apps/staged/src-tauri/src/session_runner.rs | 237 ++++++++++++------ 2 files changed, 178 insertions(+), 79 deletions(-) diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index fe19ef864..4dcc22624 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -3374,14 +3374,18 @@ async fn start_queued_session_for_branch( // Last look before the spawn. The gates above leave real awaits between // themselves and `start_session` (context building, `review_tip_sha`, // `commit_pre_head_sha`, the remote workdir resolve) — room for a whole - // shutdown to start. From here the path is synchronous, and - // `start_session` registers the session's cancellation token before any - // child spawns, so the losing interleaving narrows to a shutdown - // publishing `quit_in_progress` and snapshotting the registry inside the - // few statements between this load and that registration. Bailing here - // deliberately strands the claim: a row left `running` under our pid is - // exactly what the sweep's ownership-aware CAS puts right, while a - // spawned child is what nothing puts right. + // shutdown to start. `start_session` registers the session's cancellation + // token at entry, before the slow driver construction (login-shell binary + // probes — seconds, not statements), so past that registration a shutdown + // finds the session in its snapshot: the cancel fires the token, the + // connect path checks it before protocol setup, and the child it spawned + // is stopped by the connection teardown the exit waits on. The losing + // interleaving narrows to a shutdown publishing `quit_in_progress` and + // snapshotting the registry inside the few statements between this load + // and that registration — the status emit below and the call itself. + // Bailing here deliberately strands the claim: a row left `running` under + // our pid is exactly what the sweep's ownership-aware CAS puts right, + // while a spawned child is what nothing puts right. if crate::app_lifecycle::is_quitting(&app_handle) { return Ok(false); } diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 18751cc28..135e45615 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -306,6 +306,41 @@ impl SessionRegistry { self.inner.lock().unwrap().running.remove(session_id); } + /// Register `session_id` around its fallible startup work: the token + /// exists before `startup` runs, and a failed startup deregisters on the + /// way out. + /// + /// Registering *before* the slow half of session startup — driver + /// construction resolves the agent binary through login-shell probes, and + /// probes every known agent when no provider is pinned, so it's seconds of + /// wall clock, not statements — is what makes a session that is still + /// starting up visible to cancellation. The shutdown path cancels a + /// snapshot of registered ids and holds the exit open until they + /// deregister, and a user cancel only reaches a token the registry can + /// find. Unregistered, both land nowhere while the startup goes on to + /// spawn an agent child; registered, they fire this token, which the + /// connect path checks before protocol setup, so the child is stopped by + /// the connection teardown right after it spawns. + /// + /// Deregistering on failure is the other half of the contract: the + /// session thread that normally deregisters never spawns on that path, + /// and a stale entry would hold shutdown's `wait_for_sessions` open for + /// its full budget and misreport `is_running`. + fn register_for_startup( + &self, + session_id: &str, + startup: impl FnOnce() -> Result, + ) -> Result<(CancellationToken, T), String> { + let token = self.register(session_id); + match startup() { + Ok(value) => Ok((token, value)), + Err(e) => { + self.deregister(session_id); + Err(e) + } + } + } + /// Cancel a running session. Returns true if the session was found and /// signalled, false if it wasn't running (already finished or unknown). pub fn cancel(&self, session_id: &str) -> bool { @@ -592,86 +627,99 @@ pub fn start_session( app_handle: AppHandle, registry: Arc, ) -> Result<(), String> { - // Create the driver eagerly so we fail fast if the agent isn't found. - // Local sessions without an explicit provider resolve the first available - // provider and persist it on the session. Review-producing callers resolve - // a concrete provider before creating their session/review rows. - // Also track the provider id the driver actually resolved to. The pikchr - // sub-session (`generate_pikchr`) reuses it so its sub-agent matches the - // agent the user chose, without re-running the (login-shell) discovery. - let (driver, resolved_provider_id): (AcpDriver, Option) = if let Some(ref ws_name) = - config.workspace_name - { - let mut d = AcpDriver::for_workspace(ws_name, config.provider.as_deref())?; - if let Some(ref remote_dir) = config.remote_working_dir { - d = d.with_remote_working_dir(remote_dir.clone()); - } - (d, config.provider.clone()) - } else { - match &config.provider { - Some(id) => (AcpDriver::new(id)?, Some(id.clone())), - None => { - // Resolve the first available provider and backfill it on - // the local session record so consumers see the provider - // that actually ran the agent. - let providers = crate::agent::discover_providers(); - let first = providers.first().ok_or_else(|| { - "No ACP agent found. Install Goose, Claude Code, Codex, Pi, or Amp and ensure it's on your PATH.".to_string() - })?; - if let Err(e) = store.set_session_provider(&config.session_id, &first.id) { + // Registered before the driver is constructed, not after — see + // `register_for_startup` for why the slow construction must run with the + // token already in the registry (a shutdown or user cancel landing during + // it would otherwise miss a session about to spawn an agent child). + // `start_pipeline_session` registers at entry for the same reason, and the + // pipeline handoff's token-replacement contract (see + // `PipelineOutcome::HandedOffToAi`) is preserved: the replacement just + // happens before the slow work instead of after it. + let (cancel_token, (driver, resolved_provider_id)) = registry + .register_for_startup(&config.session_id, || { + // Create the driver eagerly so we fail fast if the agent isn't found. + // Local sessions without an explicit provider resolve the first available + // provider and persist it on the session. Review-producing callers resolve + // a concrete provider before creating their session/review rows. + // Also track the provider id the driver actually resolved to. The pikchr + // sub-session (`generate_pikchr`) reuses it so its sub-agent matches the + // agent the user chose, without re-running the (login-shell) discovery. + let (driver, resolved_provider_id): (AcpDriver, Option) = + if let Some(ref ws_name) = config.workspace_name { + let mut d = AcpDriver::for_workspace(ws_name, config.provider.as_deref())?; + if let Some(ref remote_dir) = config.remote_working_dir { + d = d.with_remote_working_dir(remote_dir.clone()); + } + (d, config.provider.clone()) + } else { + match &config.provider { + Some(id) => (AcpDriver::new(id)?, Some(id.clone())), + None => { + // Resolve the first available provider and backfill it on + // the local session record so consumers see the provider + // that actually ran the agent. + let providers = crate::agent::discover_providers(); + let first = providers.first().ok_or_else(|| { + "No ACP agent found. Install Goose, Claude Code, Codex, Pi, or Amp and ensure it's on your PATH.".to_string() + })?; + if let Err(e) = + store.set_session_provider(&config.session_id, &first.id) + { + log::warn!( + "Failed to backfill provider on session {}: {e}", + config.session_id + ); + } + (AcpDriver::new(&first.id)?, Some(first.id.clone())) + } + } + }; + + // Persist the user message right away so it's visible immediately. + // Include image IDs so the frontend can display them alongside the text. + // We also mark attached images as session-scoped immediately after so they + // don't appear in the branch timeline. Both operations are kept together; + // if set_images_session_id fails we log a warning rather than aborting the + // session, since the message was already persisted. + if let Some(ref queued_message_id) = config.queued_message_id { + store + .add_session_message_with_images_from_queue( + &config.session_id, + MessageRole::User, + &config.prompt, + &config.image_ids, + queued_message_id, + ) + .map_err(|e| format!("Failed to persist queued user message: {e}"))? + } else { + store + .add_session_message_with_images( + &config.session_id, + MessageRole::User, + &config.prompt, + &config.image_ids, + ) + .map_err(|e| format!("Failed to persist user message: {e}"))? + }; + + if !config.image_ids.is_empty() { + if let Err(e) = store.set_images_session_id(&config.image_ids, &config.session_id) + { log::warn!( - "Failed to backfill provider on session {}: {e}", + "Failed to associate images {:?} with session {}: {e}. \ + Images may appear orphaned in the branch timeline.", + config.image_ids, config.session_id ); } - (AcpDriver::new(&first.id)?, Some(first.id.clone())) } - } - }; + + Ok((driver, resolved_provider_id)) + })?; let selected_acp_config_options = crate::acp_config::selected_acp_config_options(config.acp_config_selection.as_ref()); - // Persist the user message right away so it's visible immediately. - // Include image IDs so the frontend can display them alongside the text. - // We also mark attached images as session-scoped immediately after so they - // don't appear in the branch timeline. Both operations are kept together; - // if set_images_session_id fails we log a warning rather than aborting the - // session, since the message was already persisted. - if let Some(ref queued_message_id) = config.queued_message_id { - store - .add_session_message_with_images_from_queue( - &config.session_id, - MessageRole::User, - &config.prompt, - &config.image_ids, - queued_message_id, - ) - .map_err(|e| format!("Failed to persist queued user message: {e}"))? - } else { - store - .add_session_message_with_images( - &config.session_id, - MessageRole::User, - &config.prompt, - &config.image_ids, - ) - .map_err(|e| format!("Failed to persist user message: {e}"))? - }; - - if !config.image_ids.is_empty() { - if let Err(e) = store.set_images_session_id(&config.image_ids, &config.session_id) { - log::warn!( - "Failed to associate images {:?} with session {}: {e}. \ - Images may appear orphaned in the branch timeline.", - config.image_ids, - config.session_id - ); - } - } - - let cancel_token = registry.register(&config.session_id); - // The agent protocol may use !Send futures, so we spin up a dedicated // thread with its own single-threaded Tokio runtime + LocalSet. let session_id_for_status = config.session_id.clone(); @@ -4363,6 +4411,53 @@ mod tests { ); } + /// The shutdown path cancelling its registry snapshot while a session's + /// driver is still constructing: the session is registered for the whole + /// startup, so the cancel fires the token startup hands to the session + /// thread, and the reason survives for the terminal write. + #[test] + fn register_for_startup_makes_the_session_cancellable_during_startup() { + let registry = SessionRegistry::new(); + + let (token, ()) = registry + .register_for_startup("session-starting", || { + assert!(registry.is_running("session-starting")); + assert!(registry + .cancel_with_completion_reason("session-starting", CompletionReason::AppQuit)); + Ok(()) + }) + .unwrap(); + + assert!(token.is_cancelled()); + assert!( + registry.is_running("session-starting"), + "a successful startup must stay registered for its session thread to deregister" + ); + assert_eq!( + registry.cancellation_completion_reason("session-starting"), + Some(CompletionReason::AppQuit) + ); + } + + /// A failed startup never spawns the session thread that normally + /// deregisters, so the failure path has to deregister itself — a stale + /// entry would hold shutdown's `wait_for_sessions` open for its full + /// budget on a session that can never stop. + #[test] + fn register_for_startup_deregisters_when_startup_fails() { + let registry = SessionRegistry::new(); + + let result: Result<(CancellationToken, ()), String> = + registry.register_for_startup("session-failing", || Err("no agent found".to_string())); + + assert_eq!(result.unwrap_err(), "no agent found"); + assert!(!registry.is_running("session-failing")); + assert!( + registry.wait_for_sessions(&["session-failing".to_string()], Duration::ZERO), + "shutdown must not wait on a session whose startup failed" + ); + } + fn make_git_repo(test_name: &str) -> PathBuf { let dir = std::env::temp_dir().join(format!( "staged-{test_name}-{}", From a56e90ceb7619c1644f58a675b0f57ccf028f095 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 4 Sep 2026 09:16:12 +1000 Subject: [PATCH 08/13] fix(lifecycle): hand on a cancellation the registry has already accepted MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A review of `a0c52320` found two places where a Stop the registry had already answered `true` to was then thrown away. They are the same mechanism. When `registry.cancel(id)` finds an entry, `cancel_session_impl` takes the fast path and writes no status, because that entry's observer — the session thread — is expected to record the terminal state itself. The `true` is a commitment, and two paths drop the entry without an observer. `register_for_startup`'s failure arm is the first, and new in `a0c52320`. A Stop landing while the driver resolves fires the token, records its reason and returns `true`; the resolve then fails (pinned provider binary gone, `discover_providers()` empty), the entry is deregistered, and the session thread that would have observed the token never spawns. On the pipeline-handoff path `finish_failed_pipeline_handoff_start` then wins `transition_from_running` and the row lands on `error`/`Crashed` — where before the token was registered at all, `cancel_session_impl`'s `!was_running` fallback wrote `Cancelled` at cancel time and made that transition lose. On the queued-branch path the `Err` reaches nothing but a log line in `drain_queued_sessions_for_branch`'s callers, so the row stays `running` under our pid and the next launch reports it as an errored session — the artifact this branch exists to prevent, produced here by a Stop. `register` replacing a live entry is the second, and pre-existing. `PipelineOutcome::HandedOffToAi` deliberately skips `deregister` so the insert swaps the token with no gap, but the swap carried nothing across: a cancel landing after `run_pipeline` returns and before `start_session` reaches `register` — a window that holds `git_identity_env_from_global_config()`, which shells out to `git config` — fires a token the pipeline thread is already past observing, answers `true`, and is then erased by the replacement. For an app quit that is this branch's core failure mode intact: the id is in `cancel_owned_sessions`' snapshot, so `wait_for_sessions` blocks on a session that will never honour the cancel, times out at `SHUTDOWN_BUDGET`, and `app.exit(0)` orphans the agent child. One rule covers both, named as `RunningSession::accepted_cancellation`: an entry leaving the registry without an observer hands its cancellation to whoever takes over that job. - `register` carries it onto the replacement, alongside the `pending_cancellations` intent it already applied — the same treatment, because both are cancellations recorded against a session id whose token has since been replaced. Only one can be set at a time (`cancel_or_defer` parks an intent only when nothing is registered), so the pending one is simply preferred. Scoping the carry to a *cancelled* predecessor is what keeps an ordinary handoff from starting its AI turn pre-cancelled. - `register_for_startup` reports it rather than dropping it, through a `StartupFailure` pairing the startup's error with it, and `start_session` records `Cancelled` before returning that error. This restores the ordering the pre-registration code had by accident — the cancelled write lands first, the failure's own `error` transition loses — and gets the queued-branch row off `running`. The error is still returned and still emitted by the caller; only the persisted status differs. Finding a live entry to replace means the handoff, not a session winding down: an ordinary session's thread deregisters before its terminal DB write, so while its entry is present the row is still `running` and `transition_to_running` refuses to start another turn on it. Carrying from a live predecessor therefore can't cancel a turn the user just asked for. Not closed: `cancel_with_completion_reason` clones the entry out from under the registry lock and applies the cancellation after releasing it, so a cancel already past that clone when the startup-failure removal runs lands on an `Arc` no longer in the map and goes unreported. That is the same window `start_session`'s post-`deregister` token re-read already documents, and closing it means applying cancellations under the registry lock — a bigger contract change than these findings call for. Tested at the registry level, where no `AppHandle` is needed: a cancel landing mid-startup coming back on the failure, a startup nobody cancelled reporting none, a cancelled predecessor's token and reason arriving on the replacement, and an uncancelled predecessor leaving the replacement clean. Verified with `just check-all`. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/session_runner.rs | 254 ++++++++++++++++++-- 1 file changed, 240 insertions(+), 14 deletions(-) diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 135e45615..99e40721e 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -263,6 +263,32 @@ impl RunningSession { } self.token.cancel(); } + + /// The cancellation this entry has already accepted, if any. + /// + /// `Some` is a commitment, not a status read: the registry answered `true` + /// to the `cancel` that recorded it, so + /// [`cancel_session_impl`](crate::session_commands::cancel_session_impl) + /// took the fast path and wrote no status, leaving this entry's observer + /// to record the terminal state. Any path that drops the entry without an + /// observer has to hand that commitment on — see + /// [`SessionRegistry::register`] and + /// [`SessionRegistry::register_for_startup`]. + fn accepted_cancellation(&self) -> Option { + self.cancellation_completion_reason.lock().unwrap().clone() + } +} + +/// A session startup that failed, and the cancellation (if any) the registry +/// accepted for it while it was running. +/// +/// The cancellation has nowhere else to go: the session thread that would +/// normally observe the fired token and write the terminal state never spawns +/// on this path. See [`SessionRegistry::register_for_startup`]. +#[derive(Debug)] +struct StartupFailure { + error: String, + accepted_cancellation: Option, } impl Default for SessionRegistry { @@ -279,6 +305,30 @@ impl SessionRegistry { } /// Register a new session and return a `CancellationToken` for it. + /// + /// A cancellation already recorded for this session id carries onto the new + /// entry, from either of the two places one can be waiting: + /// + /// - `pending_cancellations`, where [`cancel_or_defer`](Self::cancel_or_defer) + /// parks an intent for a session that hasn't registered yet (DB already + /// `running`, token not yet registered). + /// - the entry this call *replaces*. `start_pipeline_session` deliberately + /// skips `deregister` on an AI handoff so this insert swaps the token with + /// no gap (see [`PipelineOutcome::HandedOffToAi`]), and a cancel landing in + /// that swap window fires the pipeline's token — which the pipeline, past + /// its own cancellation checkpoints, will never observe — and answers + /// `true`, so no status is written. Dropping it here would leave the id in + /// shutdown's cancel snapshot with nothing that can honour the cancel: + /// `wait_for_sessions` would block on it for the whole `SHUTDOWN_BUDGET` + /// and `app.exit(0)` would orphan the agent child. + /// + /// Only one can be set at a time — `cancel_or_defer` parks an intent only + /// when nothing is registered — so the pending intent is simply preferred. + /// + /// Finding a *live* entry to replace means the handoff: the thread of an + /// ordinary session deregisters before its terminal DB write, so while its + /// entry is present the row is still `running` and `transition_to_running` + /// refuses to start another turn on it. fn register(&self, session_id: &str) -> CancellationToken { let token = CancellationToken::new(); let running_session = Arc::new(RunningSession { @@ -288,10 +338,14 @@ impl SessionRegistry { background_hold: std::sync::Mutex::new(acp_client::BackgroundHoldStatus::default()), }); let mut inner = self.inner.lock().unwrap(); - // If a cancellation arrived while this session was still starting up - // (DB already `running` but the token not yet registered), apply it now - // so the startup race can't drop it. - if let Some(completion_reason) = inner.pending_cancellations.remove(session_id) { + let pending = inner.pending_cancellations.remove(session_id); + let carried = pending.or_else(|| { + inner + .running + .get(session_id) + .and_then(|previous| previous.accepted_cancellation()) + }); + if let Some(completion_reason) = carried { running_session.apply_cancellation(completion_reason); } inner @@ -306,6 +360,22 @@ impl SessionRegistry { self.inner.lock().unwrap().running.remove(session_id); } + /// Remove a session from the registry and report the cancellation it had + /// accepted, for a caller that is taking over the entry's job of recording + /// the terminal state. + /// + /// Not what an ordinary [`deregister`](Self::deregister) wants: the session + /// thread deregisters *because* it observed the cancel and is about to write + /// the terminal state itself. + fn deregister_reporting_cancellation(&self, session_id: &str) -> Option { + self.inner + .lock() + .unwrap() + .running + .remove(session_id) + .and_then(|running| running.accepted_cancellation()) + } + /// Register `session_id` around its fallible startup work: the token /// exists before `startup` runs, and a failed startup deregisters on the /// way out. @@ -326,18 +396,29 @@ impl SessionRegistry { /// session thread that normally deregisters never spawns on that path, /// and a stale entry would hold shutdown's `wait_for_sessions` open for /// its full budget and misreport `is_running`. + /// + /// That deregister is also where a cancellation would go missing, so the + /// failure reports it instead of dropping it: the cancel that fired this + /// token was answered `true` and wrote no status, and the thread that would + /// have recorded the terminal state never spawns. The caller records it — + /// see [`finish_cancelled_startup`]. + /// + /// The report is as tight as the removal: a `cancel` that cloned the entry + /// under the lock but hasn't applied yet lands on an `Arc` already out of + /// the map and goes unreported, the same window + /// `start_session`'s post-`deregister` token re-read documents. fn register_for_startup( &self, session_id: &str, startup: impl FnOnce() -> Result, - ) -> Result<(CancellationToken, T), String> { + ) -> Result<(CancellationToken, T), StartupFailure> { let token = self.register(session_id); match startup() { Ok(value) => Ok((token, value)), - Err(e) => { - self.deregister(session_id); - Err(e) - } + Err(error) => Err(StartupFailure { + error, + accepted_cancellation: self.deregister_reporting_cancellation(session_id), + }), } } @@ -613,6 +694,58 @@ pub struct SessionConfig { pub background_hold: Option, } +/// Record the terminal state for a session that was cancelled while its +/// startup was still running and whose startup then failed. +/// +/// Nothing else can. `cancel_session_impl` took the registry's fast path — the +/// token was registered, so `cancel` answered `true` — and wrote no status, +/// leaving the session to record its own terminal state; on a failed startup +/// the thread that would have done that never spawns. Left unrecorded, the +/// Stop produces nothing the user can see: on the pipeline-handoff path +/// `finish_failed_pipeline_handoff_start` wins `transition_from_running` and +/// the row reads `error`/`Crashed`, and on the queued-branch path the `Err` +/// reaches nothing but a log line — or a `let _` — in +/// `drain_queued_sessions_for_branch`'s callers, leaving the row `running` +/// under our pid for the next launch to report as an errored session. +/// +/// Writing `Cancelled` first is also what restores the ordering the +/// pre-registration code had by accident: the `!was_running` fallback wrote it +/// at cancel time, so the startup failure's own `error` transition lost. The +/// error is still returned to the caller and still emitted by it — only the +/// persisted status differs. +/// +/// The event is emitted whether or not the transition won, matching the +/// session thread's terminal emit: a row already moved on (deleted, say) still +/// needs the client's `running` state cleaned up. +fn finish_cancelled_startup( + config: &SessionConfig, + store: &Store, + app_handle: &AppHandle, + completion_reason: CompletionReason, +) { + let transitioned = store + .transition_from_running( + &config.session_id, + SessionStatus::Cancelled, + None, + Some(&completion_reason), + ) + .unwrap_or(false); + log::info!( + "Session {} was cancelled during startup (transition won: {transitioned})", + config.session_id + ); + emit_status( + app_handle, + &config.session_id, + SessionStatus::Cancelled.as_str(), + None, + Some(&completion_reason), + config.branch_id.clone(), + config.project_id.clone(), + ); +} + /// Start a session: persist the user message, spawn the agent, stream to DB. /// /// Returns immediately — the actual agent work happens on a background task. @@ -635,7 +768,7 @@ pub fn start_session( // pipeline handoff's token-replacement contract (see // `PipelineOutcome::HandedOffToAi`) is preserved: the replacement just // happens before the slow work instead of after it. - let (cancel_token, (driver, resolved_provider_id)) = registry + let started = registry .register_for_startup(&config.session_id, || { // Create the driver eagerly so we fail fast if the agent isn't found. // Local sessions without an explicit provider resolve the first available @@ -715,7 +848,16 @@ pub fn start_session( } Ok((driver, resolved_provider_id)) - })?; + }); + let (cancel_token, (driver, resolved_provider_id)) = match started { + Ok(started) => started, + Err(failure) => { + if let Some(completion_reason) = failure.accepted_cancellation { + finish_cancelled_startup(&config, &store, &app_handle, completion_reason); + } + return Err(failure.error); + } + }; let selected_acp_config_options = crate::acp_config::selected_acp_config_options(config.acp_config_selection.as_ref()); @@ -1558,7 +1700,11 @@ pub fn start_pipeline_session( // We intentionally skip deregister here: start_session's register() // call will atomically replace the old cancel token. This avoids a // window where the session has no token registered (during which a - // cancel request would be silently lost). + // cancel request would be silently lost). A cancel that lands on + // the pipeline's entry *before* that replacement isn't lost either + // — `register` carries a cancelled predecessor's reason onto the + // new entry — which matters because this thread is already past + // every point that would have observed the pipeline's token. let pre_head_sha = pre_head_for_pipeline_handoff(&config); let extra_env = if store_for_status .get_commit_by_session(&session_id) @@ -4447,10 +4593,15 @@ mod tests { fn register_for_startup_deregisters_when_startup_fails() { let registry = SessionRegistry::new(); - let result: Result<(CancellationToken, ()), String> = + let result: Result<(CancellationToken, ()), StartupFailure> = registry.register_for_startup("session-failing", || Err("no agent found".to_string())); - assert_eq!(result.unwrap_err(), "no agent found"); + let failure = result.unwrap_err(); + assert_eq!(failure.error, "no agent found"); + assert_eq!( + failure.accepted_cancellation, None, + "a startup nobody cancelled has no cancellation to hand back" + ); assert!(!registry.is_running("session-failing")); assert!( registry.wait_for_sessions(&["session-failing".to_string()], Duration::ZERO), @@ -4458,6 +4609,81 @@ mod tests { ); } + /// A Stop landing while the driver resolves, on a resolve that then fails: + /// `cancel` answered `true` and so wrote no status, and the session thread + /// that would record the terminal state never spawns. The failure hands the + /// cancellation back rather than dropping it, so the caller can write + /// `cancelled` instead of letting the startup's own error be the only + /// outcome the user's Stop produced. + #[test] + fn register_for_startup_reports_a_cancellation_that_landed_during_startup() { + let registry = SessionRegistry::new(); + + let result: Result<(CancellationToken, ()), StartupFailure> = registry + .register_for_startup("session-cancelled-mid-startup", || { + assert!( + registry.cancel_with_completion_reason( + "session-cancelled-mid-startup", + CompletionReason::AppQuit + ), + "the cancel must take the registry's fast path, which writes no status" + ); + Err("No ACP agent found.".to_string()) + }); + + let failure = result.unwrap_err(); + assert_eq!(failure.error, "No ACP agent found."); + assert_eq!( + failure.accepted_cancellation, + Some(CompletionReason::AppQuit) + ); + assert!(!registry.is_running("session-cancelled-mid-startup")); + } + + /// The pipeline handoff replaces a live entry rather than deregistering it, + /// so the replacement has to inherit a cancel the predecessor accepted: the + /// pipeline thread is past every point that would observe its own token, and + /// `cancel` already answered `true`, so nothing else will honour the Stop. + #[test] + fn register_carries_a_cancelled_predecessors_state_onto_its_replacement() { + let registry = SessionRegistry::new(); + + let pipeline_token = registry.register("session-handoff"); + assert!( + registry.cancel_with_completion_reason("session-handoff", CompletionReason::AppQuit) + ); + assert!(pipeline_token.is_cancelled()); + + let ai_token = registry.register("session-handoff"); + + assert!( + ai_token.is_cancelled(), + "the handed-off session must start already cancelled" + ); + assert_eq!( + registry.cancellation_completion_reason("session-handoff"), + Some(CompletionReason::AppQuit), + "and must keep the reason its terminal write has to persist" + ); + } + + /// The carry-forward is scoped to a cancelled predecessor: an ordinary + /// handoff hands over a live session, and starting it pre-cancelled would + /// kill the AI turn the pipeline just asked for. + #[test] + fn register_starts_clean_when_the_entry_it_replaces_was_not_cancelled() { + let registry = SessionRegistry::new(); + + registry.register("session-handoff"); + let ai_token = registry.register("session-handoff"); + + assert!(!ai_token.is_cancelled()); + assert_eq!( + registry.cancellation_completion_reason("session-handoff"), + None + ); + } + fn make_git_repo(test_name: &str) -> PathBuf { let dir = std::env::temp_dir().join(format!( "staged-{test_name}-{}", From 7c13106d683d8c6bf6b456672f257f4d5199c7a7 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 4 Sep 2026 09:46:26 +1000 Subject: [PATCH 09/13] fix(lifecycle): gate the agent spawn on the cancellation the startup can't see MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `a0c52320` registered a session's token before its slow driver construction and described what that buys: a cancel fires the token, `run_acp_session` selects on it ahead of `initialize`, the child is torn down by the connection task's `graceful_stop`, and `wait_for_sessions` holds the exit open for it. A review of that commit showed the chain only completes inside `SHUTDOWN_BUDGET` if the startup it now covers also fits there — and routinely it won't: `doctor` gives each login-shell probe a 10s timeout, and the unpinned-provider branch pays that per provider through `discover_providers()`. That work is blocking and not token-aware, so the cancel can't shorten it. Shutdown therefore waits on a session it previously couldn't see, times out, warns, sweeps and exits — and if the resolve finishes just inside the deadline, `connect` spawns the child moments before `app.exit(0)`: the orphan case again. The weak link is that `connect` spawns unconditionally. Its check is a `select!` *after* the spawn, so the design's answer to "cancelled before we connected" is to start the agent and immediately kill it — correct only while the kill outruns the exit. The session loop now takes a last look at the token before it connects and returns `Ok(AgentRunOutcome::Cancelled)` when it has already fired, so nothing is started that has to be raced back down. The `generate_pikchr` worker has always gated its `driver.run` this way; the main session loop was the path missing the check. Terminal handling is untouched, because that outcome is exactly what the spawn-then-teardown path produced: the row still lands `cancelled` with the registry's recorded reason (`app_quit` for a quit), and shutdown's wait now ends on the session's own deregister instead of on the budget. An ordinary Stop during startup gets the same benefit — the agent CLI never launches at all rather than launching to be killed. The doc comment on `register_for_startup` now states the guarantee it can actually keep. Registration buys visibility, not a bounded stop: a cancel is *observed* only once the blocking startup ends, routinely past the 2s budget, so shutdown still times out on such a session, warns, sweeps its row to cancelled/`app_quit` and exits, killing the still-probing thread with the process. That is a clean end rather than a leak precisely because of the gate — the thread has spawned no agent child, and now never will. Considered and rejected: having shutdown skip the wait for sessions that haven't connected, on the theory that exiting sooner kills the resolving thread before it can spawn anything. It doesn't dominate. The resolve runs as blocking `Command` probes on the thread that called `start_session` (a tokio worker for every production caller), so while it runs it neither observes the token nor spawns anything agent-shaped — doctor's probes are short-lived login shells in their own process groups, which `run_command_with_timeout` reaps. So nothing is leaked by waiting, and nothing is saved by not waiting: the orphan window in both shapes is "the resolve finishes within the teardown-sized slice just before the process dies", which skipping the wait moves earlier rather than closes. The gate closes it in both. What skipping would add is a `connected` bit on `RunningSession`, a shutdown-side filter, and the loss of the session's own terminal write in exchange for a sweep row saying the same thing — latency bought with contract. No new test: the gate sits between a constructed `AcpDriver` and a spawned agent child inside `start_session`'s session thread, which takes a concrete `AppHandle`, so reaching it needs the whole spawn path — the same limit `5d6d7441`'s quit gates landed under. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/session_runner.rs | 37 +++++++++++++++++++-- 1 file changed, 35 insertions(+), 2 deletions(-) diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 99e40721e..39bf11f23 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -389,8 +389,21 @@ impl SessionRegistry { /// deregister, and a user cancel only reaches a token the registry can /// find. Unregistered, both land nowhere while the startup goes on to /// spawn an agent child; registered, they fire this token, which the - /// connect path checks before protocol setup, so the child is stopped by - /// the connection teardown right after it spawns. + /// session thread takes a last look at before it connects — so a cancel + /// that lands during the startup spawns no agent child at all, and one + /// that lands after that look meets the post-spawn check and + /// `graceful_stop`. + /// + /// What registration does *not* buy is a stop bounded by + /// `SHUTDOWN_BUDGET`. The startup it now covers is blocking and not + /// token-aware — `doctor` gives each login-shell probe a 10s timeout, and + /// the unpinned branch pays that per provider — so a cancel is only + /// *observed* once that work ends, routinely past the 2s budget. Shutdown + /// then times out on this session, warns, sweeps its row to + /// cancelled/`app_quit`, and exits, killing the still-probing thread with + /// the process. That is a clean end rather than a leak precisely because + /// of the gate before `connect`: the thread has spawned no agent child, + /// and now never will. /// /// Deregistering on failure is the other half of the contract: the /// session thread that normally deregisters never spawns on that path, @@ -1111,6 +1124,26 @@ pub fn start_session( }; include_images = false; + // Last look before an agent process exists. Everything from the + // token's registration to here is blocking and not token-aware + // (driver construction's login-shell probes, then the env + // snapshot capture), so a cancel that landed during it — a + // Stop, or a quit whose `SHUTDOWN_BUDGET` has since run out — + // is first observable at this point. `connect` spawns the + // child unconditionally: its own check sits *after* the spawn, + // ahead of `initialize`, and leaves `graceful_stop` to take + // the child back down, which only beats a quit's `app.exit(0)` + // if the startup fit in the budget too. Bailing here leaves + // nothing to take down. The `generate_pikchr` worker gates its + // `driver.run` the same way. + // + // Past this check the path is synchronous into `cmd.spawn()`, + // so what remains is a cancel landing inside those statements + // — and that one the post-spawn check still answers. + if cancel_token.is_cancelled() { + return Ok(AgentRunOutcome::Cancelled); + } + // Open a session-scoped connection, then send this turn's // prompt over it. Without a background hold the connection // still tears the bridge process down as soon as the prompt From f4cbc8c73ea923141f38eb34eb1ed92b83005d0a Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 4 Sep 2026 10:15:43 +1000 Subject: [PATCH 10/13] fix(lifecycle): refuse a start the shutdown has already ruled out MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `is_quitting` gated the queue drain and nothing else. A review of `89c821c3` named what that leaves open: `start_session` and `start_pipeline_session` had no gate of their own, so every other way into a session — the `start_session` and `resume_session` commands, `start_branch_session`, `start_project_session` (all four also dispatchable in web mode), and the four `start_pipeline_session` sites in `prs` — could still enter inside `SHUTDOWN_BUDGET`, after `cancel_owned_sessions` had taken its registry snapshot. Such a session does register, so `sweep_active_sessions` puts its DB row right, but it is neither cancelled nor waited on, and `app.exit(0)` orphans its agent child — own process group, and an exit runs no `kill_on_drop` destructors. That is the one outcome the sweep cannot fix. The review classed it as a new decision rather than a defect, and it is the last coverage gap in the invariant this branch exists for, so it is closed the way the review framed it: one gate in each of the two funnels rather than a gate per caller. Every agent start now reaches `register_for_startup` by way of `start_session`, and every pipeline reaches `start_pipeline_session`; the project-MCP `start_repo_session` only ever enqueues a row and drains, so it inherits the drain's gate. A refused start must not be an `Err`. An error from here becomes an errored session one way or another — the pipeline-handoff path writes `error`/`Crashed` outright via `finish_failed_pipeline_handoff_start`, and everywhere else the row is simply left `running` under our pid for the next launch to report as one, which is the artifact this branch exists to prevent. So the refusal does what the sweep would have done and answers with a non-error outcome: - The row is written `cancelled`/`app_quit` by `finish_cancelled_before_run` (`finish_cancelled_startup` generalized to scalars, since `PipelineConfig` needs it too). Leaving that to the sweep would be right only until the sweep runs: it is `shutdown_cleanup`'s last step, so a refusal in the window between it and `app.exit(0)` would strand a `running` row. - The registry entry is dropped, and a cancellation it had already accepted is preferred over `AppQuit`. That matters on one path: the pipeline handoff skips `deregister` so `start_session`'s `register` can swap the token with no gap. A refusal never reaches that `register`, so without this the pipeline's entry would outlive the thread that owned it — `is_running` true forever, and `wait_for_sessions` burning the whole budget on a session nothing can stop. Taking the entry out means taking over its job of recording the terminal state, which is the rule `15735fb1` named. - Both functions return `SessionStartOutcome` instead of `()`. Most callers keep `?;` and are right to — the row and the status event are already correct, so returning their session id says what actually exists. The callers that care are the queued-branch drain and the two queued pipeline starts in `prs`, which now report `Ok(false)` (a refusal leaves the branch as idle as a claim that lost its race), and the pipeline handoff, which resolves its own artifact the way its `Cancelled` arm does. Of the drain's three gates, two stay and one goes. The entry gate and the per-iteration gate bail *before* `transition_queued_to_running`, and an unclaimed row is a different DB state, not just an earlier one: still `queued`, carrying no owner, and so still available to another instance pointed at the same data dir right up until our sweep reaches it — the case `abc299a9` made the sweep's CAS ownership-aware for. Claiming and then refusing would stamp our pid on work we are about to throw away. The third gate, added by `5d6d7441` after the claim, has no such argument: it fires at the same point as the runner's gate and deliberately stranded the claim for the sweep, which the runner's gate has no need to do. It is removed rather than kept as a strictly worse duplicate. Also from that review: the pre-connect gate's inventory of what precedes it named only driver construction and the env snapshot capture. Between those and the check, the async block also starts the project MCP server, starts the pikchr MCP server, and reads plus base64-encodes the attached images. They are awaits rather than blocking work and none of them spawns an agent child, so the gate's reasoning is unaffected — but a session cancelled as early as its registration does stand two localhost HTTP servers up before bailing (tasks on the session thread's runtime, so they go down with it). The comment now says what is there. Deliberately not covered: `generate_pikchr`'s diagram sub-session reaches its agent through `register_external` and `driver.run`, not `start_session`, so a tool call landing after `cancel_owned_sessions`' snapshot registers a token shutdown never fires. It is narrower than anything closed here — the parent session has to be mid-turn and already being cancelled — and it is a different cancellation chain, so it belongs to its own change rather than folded into this one. Tested at the level the gate allows. The registry half needs no app at all: which reason a refusal persists, and that it frees the handoff predecessor's entry so `wait_for_sessions` stops waiting on it. The whole refusal is then driven against a real store with a mock app, which making `emit_status` runtime-generic is what enables. The gate itself is still out of reach — `start_session` takes the concrete `AppHandle` the spawn path needs. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/prs.rs | 12 +- apps/staged/src-tauri/src/session_commands.rs | 50 +-- apps/staged/src-tauri/src/session_runner.rs | 313 ++++++++++++++++-- 3 files changed, 314 insertions(+), 61 deletions(-) diff --git a/apps/staged/src-tauri/src/prs.rs b/apps/staged/src-tauri/src/prs.rs index 5976f9bb8..e540dc733 100644 --- a/apps/staged/src-tauri/src/prs.rs +++ b/apps/staged/src-tauri/src/prs.rs @@ -809,9 +809,9 @@ pub(crate) async fn start_queued_commit_pipeline_for_branch( store, app_handle, Arc::clone(®istry), - )?; - - Ok(true) + ) + // A shutdown refusal is not a start — see `start_queued_session_for_branch`. + .map(session_runner::SessionStartOutcome::started) } /// Insert the session row for a push that runs right now. @@ -1002,9 +1002,9 @@ pub(crate) async fn start_queued_git_pipeline_for_branch( store, app_handle, Arc::clone(®istry), - )?; - - Ok(true) + ) + // A shutdown refusal is not a start — see `start_queued_session_for_branch`. + .map(session_runner::SessionStartOutcome::started) } /// What the branch queue decided to do with a pull request. diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index 4dcc22624..60be553d6 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -3122,6 +3122,15 @@ pub async fn drain_queued_sessions_for_branch( // `app.exit(0)` then orphans, since they run in their own process groups and // an exit runs no `kill_on_drop` destructors. The sweep would put the DB rows // right; nothing would put the processes right. + // + // `start_session` and `start_pipeline_session` now refuse a start of their + // own accord, so this gate is no longer what stops the children — but it is + // still what stops the *claims*. A row this returns without touching is + // still `queued`, carrying no owner, which is what leaves it available to + // another instance pointed at the same data dir right up until our sweep + // reaches it (the sweep's CAS is ownership-aware for exactly that reason). + // Claiming it and then refusing would stamp our pid on work we are about to + // throw away. if crate::app_lifecycle::is_quitting(&app_handle) { return Ok(false); } @@ -3145,7 +3154,11 @@ pub async fn drain_queued_sessions_for_branch( // The entry gate, re-checked before each claim: every start below // awaits, so a shutdown claimed mid-drain would otherwise keep being - // fed the rows the remaining iterations were about to take. + // fed the rows the remaining iterations were about to take. Kept for + // the same reason as the entry gate now that the runner refuses starts + // itself — this is the last point at which a queued row can be left + // unclaimed, and the rows after it are the ones a drain would otherwise + // claim one by one on its way out. if crate::app_lifecycle::is_quitting(&app_handle) { break; } @@ -3371,25 +3384,13 @@ async fn start_queued_session_for_branch( .get_image_ids_for_session(&session_id) .unwrap_or_default(); - // Last look before the spawn. The gates above leave real awaits between - // themselves and `start_session` (context building, `review_tip_sha`, - // `commit_pre_head_sha`, the remote workdir resolve) — room for a whole - // shutdown to start. `start_session` registers the session's cancellation - // token at entry, before the slow driver construction (login-shell binary - // probes — seconds, not statements), so past that registration a shutdown - // finds the session in its snapshot: the cancel fires the token, the - // connect path checks it before protocol setup, and the child it spawned - // is stopped by the connection teardown the exit waits on. The losing - // interleaving narrows to a shutdown publishing `quit_in_progress` and - // snapshotting the registry inside the few statements between this load - // and that registration — the status emit below and the call itself. - // Bailing here deliberately strands the claim: a row left `running` under - // our pid is exactly what the sweep's ownership-aware CAS puts right, - // while a spawned child is what nothing puts right. - if crate::app_lifecycle::is_quitting(&app_handle) { - return Ok(false); - } - + // No last look before the spawn here any more: `start_session` takes one + // itself, a few statements further on (the status emit and the config it is + // handed) and for every caller rather than this one. Both fire after the + // claim, so the only difference between them is what they leave behind — + // this gate stranded a `running` row for the sweep to put right, which is + // fine while the sweep is still to come and wrong in the window between the + // sweep and `app.exit(0)`. The runner's gate records the row itself. let session_type_str = match session_type { BranchSessionType::Commit => "commit", BranchSessionType::Note => "note", @@ -3439,9 +3440,12 @@ async fn start_queued_session_for_branch( store, app_handle, Arc::clone(®istry), - )?; - - Ok(true) + ) + // A start the runner refused because a shutdown is under way is not a + // start: report it like a claim that lost its race, so the drain loop + // re-reads the branch's active kinds instead of counting this session as + // running. Its next iteration stops on the pre-claim gate anyway. + .map(session_runner::SessionStartOutcome::started) } // ============================================================================= diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 39bf11f23..823d2eb2e 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -707,8 +707,105 @@ pub struct SessionConfig { pub background_hold: Option, } -/// Record the terminal state for a session that was cancelled while its -/// startup was still running and whose startup then failed. +/// What a start request did with the session it was handed. +/// +/// A start is not always a start: [`start_session`] and +/// [`start_pipeline_session`] refuse one outright when a shutdown has already +/// been claimed. That refusal is not an error — see +/// [`refuse_start_during_shutdown`] — so it needs a way to say "nothing is +/// running" that no caller can turn into a failed session. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SessionStartOutcome { + /// A runner thread has the session. + Started, + /// A shutdown was already under way, so nothing was started and the row was + /// recorded `cancelled`/`app_quit`. + RefusedShuttingDown, +} + +impl SessionStartOutcome { + /// Whether a runner thread actually took the session. What the queue drain + /// reports as "started" — a refusal leaves the branch exactly as idle as a + /// claim that lost its race. + pub fn started(self) -> bool { + matches!(self, Self::Started) + } +} + +/// Refuse to start `session_id` because a shutdown has been claimed, leaving +/// the row in the state the quit sweep would have given it. +/// +/// The gate this serves is the coverage boundary the drain's own gates leave +/// open. `is_quitting` used to be read only on the queue-drain path, so every +/// other way into a session — the `start_session` and `resume_session` +/// commands, `start_branch_session`, `start_project_session` (all four also +/// dispatchable in web mode), and the four `start_pipeline_session` sites in +/// `prs` — could still enter inside `SHUTDOWN_BUDGET`, *after* +/// `cancel_owned_sessions` took its registry snapshot. Such a session +/// registers, so `sweep_active_sessions` puts its row right, but it is neither +/// cancelled nor waited on: `app.exit(0)` orphans its agent child, and a child +/// in its own process group is the one thing the sweep cannot fix. Every agent +/// start funnels through [`start_session`] and every pipeline through +/// [`start_pipeline_session`], so those two gates cover all of them — including +/// the project-MCP `start_repo_session`, which only ever enqueues a row and +/// drains. +/// +/// Writing the row here rather than leaving it to the sweep is what makes the +/// refusal self-contained. The sweep is `shutdown_cleanup`'s *last* step, so a +/// refusal arriving after it — the window between the sweep and `app.exit(0)` +/// — would strand a `running` row under our pid for the next launch to report +/// as an errored session, which is the artifact this branch exists to prevent. +/// +/// Dropping the registry entry matters on one path: the pipeline handoff skips +/// `deregister` so [`start_session`]'s `register` can swap the token with no +/// gap (see [`PipelineOutcome::HandedOffToAi`]). A refusal never reaches that +/// `register`, so without this the pipeline's entry would outlive the thread +/// that owned it — `is_running` true forever, and shutdown's +/// `wait_for_sessions` burning its whole budget on a session nothing can stop. +/// Taking the entry out also means taking over its job, so a cancellation it +/// had already accepted becomes the reason we persist, in preference to +/// `AppQuit`: the `cancel` that recorded it was answered `true` and wrote no +/// status. +fn refuse_start_during_shutdown( + session_id: &str, + branch_id: Option, + project_id: Option, + store: &Store, + app_handle: &AppHandle, + registry: &SessionRegistry, +) -> SessionStartOutcome { + log::info!("Refusing to start session {session_id}: a shutdown is under way"); + finish_cancelled_before_run( + session_id, + branch_id, + project_id, + store, + app_handle, + refused_start_completion_reason(registry, session_id), + ); + SessionStartOutcome::RefusedShuttingDown +} + +/// The registry half of a refused start: drop the entry the session already +/// had, if any, and answer with the reason its terminal write should carry. +/// +/// A cancellation the dropped entry had accepted wins over `AppQuit` — see +/// [`RunningSession::accepted_cancellation`] for why an entry leaving the +/// registry without an observer has to hand its cancellation on. Only the +/// pipeline handoff has an entry to find here; every other caller reaches +/// [`start_session`] with nothing registered, and gets `AppQuit`. +fn refused_start_completion_reason( + registry: &SessionRegistry, + session_id: &str, +) -> CompletionReason { + registry + .deregister_reporting_cancellation(session_id) + .unwrap_or(CompletionReason::AppQuit) +} + +/// Record the terminal state for a session that never reached its run: one +/// cancelled while its startup was still going and whose startup then failed, +/// or one [`refuse_start_during_shutdown`] turned away. /// /// Nothing else can. `cancel_session_impl` took the registry's fast path — the /// token was registered, so `cancel` answered `true` — and wrote no status, @@ -730,32 +827,34 @@ pub struct SessionConfig { /// The event is emitted whether or not the transition won, matching the /// session thread's terminal emit: a row already moved on (deleted, say) still /// needs the client's `running` state cleaned up. -fn finish_cancelled_startup( - config: &SessionConfig, +fn finish_cancelled_before_run( + session_id: &str, + branch_id: Option, + project_id: Option, store: &Store, - app_handle: &AppHandle, + app_handle: &AppHandle, completion_reason: CompletionReason, ) { let transitioned = store .transition_from_running( - &config.session_id, + session_id, SessionStatus::Cancelled, None, Some(&completion_reason), ) .unwrap_or(false); log::info!( - "Session {} was cancelled during startup (transition won: {transitioned})", - config.session_id + "Session {session_id} was cancelled before its run started (transition won: \ + {transitioned})" ); emit_status( app_handle, - &config.session_id, + session_id, SessionStatus::Cancelled.as_str(), None, Some(&completion_reason), - config.branch_id.clone(), - config.project_id.clone(), + branch_id, + project_id, ); } @@ -772,7 +871,23 @@ pub fn start_session( store: Arc, app_handle: AppHandle, registry: Arc, -) -> Result<(), String> { +) -> Result { + // The gate every agent start passes, whatever raised it — see + // `refuse_start_during_shutdown`. It sits ahead of the registration + // deliberately: a session registered here would have to be deregistered + // again on the way out, and this is the one point where a start can still + // be declined without anything having been spawned. + if crate::app_lifecycle::is_quitting(&app_handle) { + return Ok(refuse_start_during_shutdown( + &config.session_id, + config.branch_id.clone(), + config.project_id.clone(), + &store, + &app_handle, + ®istry, + )); + } + // Registered before the driver is constructed, not after — see // `register_for_startup` for why the slow construction must run with the // token already in the registry (a shutdown or user cancel landing during @@ -866,7 +981,14 @@ pub fn start_session( Ok(started) => started, Err(failure) => { if let Some(completion_reason) = failure.accepted_cancellation { - finish_cancelled_startup(&config, &store, &app_handle, completion_reason); + finish_cancelled_before_run( + &config.session_id, + config.branch_id.clone(), + config.project_id.clone(), + &store, + &app_handle, + completion_reason, + ); } return Err(failure.error); } @@ -1124,18 +1246,24 @@ pub fn start_session( }; include_images = false; - // Last look before an agent process exists. Everything from the - // token's registration to here is blocking and not token-aware - // (driver construction's login-shell probes, then the env - // snapshot capture), so a cancel that landed during it — a - // Stop, or a quit whose `SHUTDOWN_BUDGET` has since run out — - // is first observable at this point. `connect` spawns the - // child unconditionally: its own check sits *after* the spawn, - // ahead of `initialize`, and leaves `graceful_stop` to take - // the child back down, which only beats a quit's `app.exit(0)` - // if the startup fit in the budget too. Bailing here leaves - // nothing to take down. The `generate_pikchr` worker gates its - // `driver.run` the same way. + // Last look before an agent process exists. Nothing between the + // token's registration and here observes it: first the blocking, + // non-token-aware half of startup (driver construction's + // login-shell probes, then the env snapshot capture), then a run + // of awaits — the project MCP server, the pikchr MCP server, and + // reading plus base64-encoding the attached images. So a cancel + // that landed anywhere in there — a Stop, or a quit whose + // `SHUTDOWN_BUDGET` has since run out — is first observable at + // this point. A session cancelled as early as its registration + // therefore still stands both localhost MCP servers up on its way + // here; they are tasks on this thread's runtime, so they go down + // with it once the terminal handling below finishes. `connect` + // spawns the child unconditionally: its own check sits *after* + // the spawn, ahead of `initialize`, and leaves `graceful_stop` + // to take the child back down, which only beats a quit's + // `app.exit(0)` if the startup fit in the budget too. Bailing + // here leaves nothing to take down. The `generate_pikchr` worker + // gates its `driver.run` the same way. // // Past this check the path is synchronous into `cmd.spawn()`, // so what remains is a cancel landing inside those statements @@ -1489,7 +1617,7 @@ pub fn start_session( } }); - Ok(()) + Ok(SessionStartOutcome::Started) } #[derive(Debug, Clone, PartialEq, Eq)] @@ -1662,7 +1790,26 @@ pub fn start_pipeline_session( store: Arc, app_handle: AppHandle, registry: Arc, -) -> Result<(), String> { +) -> Result { + // The pipeline half of the shutdown gate — see + // `refuse_start_during_shutdown`. A pipeline spawns no agent child of its + // own, but it runs command steps in their own process groups and can hand + // off to an AI session, so an exit through the middle of one leaves the + // same mess by a longer route. The artifact resolution matches the + // `PipelineOutcome::Cancelled` arm below, because that is what this is: a + // pipeline cancelled before it ran a step. + if crate::app_lifecycle::is_quitting(&app_handle) { + resolve_pipeline_artifacts_without_ai(&config, &store, false); + return Ok(refuse_start_during_shutdown( + &config.session_id, + config.branch_id.clone(), + config.project_id.clone(), + &store, + &app_handle, + ®istry, + )); + } + let cancel_token = registry.register(&config.session_id); let session_id = config.session_id.clone(); let store_for_status = Arc::clone(&store); @@ -1778,12 +1925,22 @@ pub fn start_pipeline_session( parent_project_note_id: None, background_hold: crate::session_commands::default_background_hold(), }; - if let Err(e) = start_session( + let handoff = start_session( ai_config, store_for_status.clone(), app_handle.clone(), Arc::clone(®istry), - ) { + ); + // A refusal has already taken over this session's registry + // entry — the one this handoff deliberately left in place for + // `register` to swap — and written the row + // cancelled/`app_quit`. What it can't know about is the + // pipeline's own artifact, which still needs the resolution a + // cancelled pipeline gives it. + if matches!(handoff, Ok(SessionStartOutcome::RefusedShuttingDown)) { + resolve_pipeline_artifacts_without_ai(&config, &store_for_status, false); + } + if let Err(e) = handoff { log::error!("Failed to start AI session after pipeline handoff: {e}"); // If the handoff came from an explicit AiHandoff step, mark // it as failed so the UI doesn't show a perpetual spinner. @@ -1930,7 +2087,7 @@ pub fn start_pipeline_session( } }); - Ok(()) + Ok(SessionStartOutcome::Started) } fn finish_failed_pipeline_handoff_start( @@ -3833,8 +3990,11 @@ fn find_closing_fence(text: &str) -> Option { None } -fn emit_status( - app_handle: &AppHandle, +/// Generic over the runtime purely so the paths that end a session *before* it +/// runs (see [`finish_cancelled_before_run`]) can be driven by a mock app in +/// tests. Every production caller passes the concrete handle. +fn emit_status( + app_handle: &AppHandle, session_id: &str, status: &str, error: Option, @@ -4717,6 +4877,95 @@ mod tests { ); } + /// Every way into a session but the pipeline handoff arrives at + /// `start_session` with nothing registered, so an ordinary refusal has only + /// the quit's own reason to record. + #[test] + fn a_refused_start_with_nothing_registered_records_the_quit() { + assert_eq!( + refused_start_completion_reason(&SessionRegistry::new(), "session-never-registered"), + CompletionReason::AppQuit + ); + } + + /// The handoff is the one path that arrives with a live entry: it skips + /// `deregister` so `register` can swap the token with no gap. A refusal + /// never reaches that `register`, so it has to take the entry out itself — + /// left behind, it would hold shutdown's `wait_for_sessions` open for the + /// whole budget on a session no thread is driving. + #[test] + fn a_refused_start_frees_the_handoff_predecessors_entry() { + let registry = SessionRegistry::new(); + registry.register("session-handoff"); + + assert_eq!( + refused_start_completion_reason(®istry, "session-handoff"), + CompletionReason::AppQuit + ); + assert!(!registry.is_running("session-handoff")); + assert!( + registry.wait_for_sessions(&["session-handoff".to_string()], Duration::ZERO), + "shutdown must not wait on a session the refusal took over" + ); + } + + /// And when a Stop had already landed on that entry, the refusal is what + /// takes over its job of recording the terminal state: `cancel` answered + /// `true` and so wrote no status, which makes its reason the one the row + /// has to carry rather than the quit's. + #[test] + fn a_refused_start_persists_a_cancellation_the_predecessor_accepted() { + let registry = SessionRegistry::new(); + registry.register("session-handoff"); + assert!(registry.cancel("session-handoff")); + + assert_eq!( + refused_start_completion_reason(®istry, "session-handoff"), + CompletionReason::Interrupted + ); + } + + /// A mock app so the refusal's status event has somewhere to go. Nothing + /// listens, so the emit is a no-op — the point is to drive the real + /// refusal rather than a test-only copy of it. + fn mock_app() -> tauri::App { + tauri::test::mock_builder() + .build(tauri::test::mock_context(tauri::test::noop_assets())) + .expect("failed to build mock app") + } + + /// The whole refusal against a real store. Every caller has already moved + /// its row to `running` by the time it gets here, and the refusal writes + /// the terminal state itself rather than leaving it to the quit sweep — + /// the sweep is `shutdown_cleanup`'s last step, so a refusal landing after + /// it would strand a `running` row under our pid for the next launch to + /// report as an errored session. + #[test] + fn a_refused_start_records_the_row_the_sweep_would_have() { + let store = Store::in_memory().unwrap(); + let session = crate::store::Session::new_running("prompt", &PathBuf::from("/tmp")); + store.create_session(&session).unwrap(); + let app = mock_app(); + + let outcome = refuse_start_during_shutdown( + &session.id, + None, + None, + &store, + app.handle(), + &SessionRegistry::new(), + ); + + assert_eq!(outcome, SessionStartOutcome::RefusedShuttingDown); + assert!( + !outcome.started(), + "a refusal must not read as a start to the queue drain" + ); + let row = store.get_session(&session.id).unwrap().unwrap(); + assert_eq!(row.status, SessionStatus::Cancelled); + assert_eq!(row.completion_reason, Some(CompletionReason::AppQuit)); + } + fn make_git_repo(test_name: &str) -> PathBuf { let dir = std::env::temp_dir().join(format!( "staged-{test_name}-{}", From 9b3db0162a01f35a7ec24d994c54de79eb62b1e0 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 4 Sep 2026 10:42:57 +1000 Subject: [PATCH 11/13] fix(lifecycle): apply a cancellation under the lock its takeover reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A review of `89c821c3` left three findings on the startup path, all about a cancellation surviving the trip from the registry to the row. **The carry-forward race.** `register` reads `previous.accepted_cancellation()` under the registry lock, but `cancel_with_completion_reason` cloned the entry's `Arc` under the lock and applied the cancellation after releasing it. A cancel caught between those two points reads as `None`: the replacement starts clean, the cancel then lands on an `Arc` the map no longer holds, and `cancel` still answered `true` — so `cancel_session_impl` wrote no status. On the handoff path that is exactly what `15735fb1` set out to close: the id is in shutdown's cancel snapshot, the AI session runs uncancelled, `wait_for_sessions` burns the whole `SHUTDOWN_BUDGET`, and `app.exit(0)` orphans the agent child. The review offered liveness-reporting — have `apply_cancellation` say whether it landed on an entry still in the map, so `cancel` can answer `false` and let `cancel_session_impl` take its store-write fallback. Applying under the lock is the better of the two, and the earlier session was wrong to have called it a bigger contract change: `register` has always applied a carried or pending cancellation with the guard held, and the pre-`86e4efbe` registry cancelled straight out of the map the same way. Nothing in `apply_cancellation` reaches back into the registry — the reason mutex is only ever taken *under* the registry lock, never the reverse, and `CancellationToken::cancel` notifies its waiters, which schedules the waiting tasks rather than running them inline — so the lock scope costs nothing. Liveness-reporting, by contrast, answers `false` to cancels that *were* honoured (the entry's successor or its remover took them over), buying a duplicate store write and a reason downgraded to `Interrupted` for a quit's `AppQuit`. The mechanism is `RegistryInner::cancel_registered`, which takes `&RegistryInner` rather than `&SessionRegistry` so the lock guard is the only way to reach it; `cancel_with_completion_reason` and `cancel_or_defer` both route through it. That closes both sites at once — the `register` carry, and the `deregister_reporting_cancellation` removal that `register_for_startup`'s doc had been documenting as open. The window was narrower than either the review or that doc said, which is worth recording rather than losing. `apply_cancellation` holds the reason mutex across `token.cancel()`, and both takeover readers go through `accepted_cancellation`, which wants that same mutex — so a cancel already *inside* `apply_cancellation` was serialized against them anyway, and what was exposed is the instant between the map lookup and that mutex. That is an accident of two unrelated lock scopes, invisible at both sites and undone by any tidying that releases the reason guard before firing the token. The two new tests pin the invariant to the registry lock rather than to the accident: they park a cancel inside `apply_cancellation` — with a waker that blocks, since `cancel` notifies its waiters synchronously, the only interposition point this race has — and require the swap and the removal to wait for it. Against the old split they pass, for the incidental reason; against the old split with the reason guard released early they both fail; under the registry lock they hold either way. **The failed handoff talking over a Stop.** When the handoff's `start_session` failed with a carried cancellation, `finish_failed_pipeline_handoff_start` correctly lost `transition_from_running` — `finish_cancelled_before_run` had already written `cancelled` and emitted it — but the `emit_status` behind it fired regardless. The last event a client saw was `error`/`Crashed` on a row that says cancelled, and the AiHandoff step was stamped "Failed to start AI session: …" for what was the user's Stop. Both now hang off `transitioned`, the way the other three `PipelineOutcome` arms gate their side effects. The emit keeps one escape the review didn't call for: a row that is *gone* rather than moved on, when the user deletes the pending commit mid-pipeline. That is the case the unconditional emit existed for — no other writer emitted anything, and the event is all that clears the client's `running` state — so gating on `transitioned` alone would have quietly dropped it. All of it moves into `finish_failed_pipeline_handoff_start`, now generic over the runtime and taking the `PipelineConfig`, which is what puts it in reach of a test: a mock app and a real store drive both the win (row `error`/`Crashed`, step `Failed`) and the loss (row still cancelled/`Interrupted`, step still pending). **The startup failure not draining its branch.** The session thread's terminal path drains the branch queue for any `branch_id.is_some()` once its transition wins, cancelled sessions included. The startup-failure path didn't, so a Stop landing during a startup that then fails cleared the row and left the branch's remaining queued sessions parked until some unrelated terminal transition happened to trigger a drain. `finish_cancelled_before_run` now reports whether it won the transition, and `start_session` kicks the drain on that answer, through the pipeline arms' own helper — renamed `drain_queued_after_pipeline_terminal` to `drain_queued_after_terminal_state`, since pipelines are no longer its only callers. Safe on the quit path because the drain is `is_quitting`-gated. The refusal deliberately kicks none: it only ever runs while that gate is closed, and a queued row left unclaimed is what a quit wants to leave behind. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/session_runner.rs | 634 +++++++++++++++++--- 1 file changed, 538 insertions(+), 96 deletions(-) diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 823d2eb2e..ac9a8062f 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -251,6 +251,19 @@ impl RunningSession { /// Record the completion reason to persist and signal cancellation. A /// `ProjectSessionInterrupted` reason overrides a previously stored one so /// an explicit project cancel wins over an in-flight interrupt. + /// + /// Always called with the registry lock held — by + /// [`RegistryInner::cancel_registered`] for an entry already in the map, + /// and by [`SessionRegistry::register`] for one about to be inserted. That + /// is what makes [`accepted_cancellation`](Self::accepted_cancellation) a + /// complete answer: no cancel can be part-way through applying to an entry + /// a lock-holder is reading, replacing or removing. + /// + /// Safe to hold the lock across, because nothing here reaches back into the + /// registry: the reason mutex is only ever taken *under* the registry lock + /// (never the reverse), and `CancellationToken::cancel` notifies its + /// waiters, which schedules the waiting tasks rather than running them + /// inline. fn apply_cancellation(&self, completion_reason: CompletionReason) { let mut stored_reason = self.cancellation_completion_reason.lock().unwrap(); if stored_reason.is_none() @@ -274,11 +287,54 @@ impl RunningSession { /// observer has to hand that commitment on — see /// [`SessionRegistry::register`] and /// [`SessionRegistry::register_for_startup`]. + /// + /// `None` is the matching commitment in the other direction, and only + /// because cancels apply under the registry lock: read by a lock-holder + /// that is about to replace or remove this entry, it means no `cancel` has + /// answered `true` for it, so every later one will find the entry gone, + /// answer `false`, and write its own status. fn accepted_cancellation(&self) -> Option { self.cancellation_completion_reason.lock().unwrap().clone() } } +impl RegistryInner { + /// Apply `completion_reason` to `session_id`'s entry, reporting whether + /// there was one. + /// + /// Taking `&RegistryInner` rather than `&SessionRegistry` is the point: the + /// only way to call it is through the lock guard, so the lookup and the + /// [`apply_cancellation`](RunningSession::apply_cancellation) cannot be + /// split by a `register` or a `deregister` on another thread. Held apart + /// — the entry cloned out under the lock and cancelled after releasing it + /// — a cancel could answer `true` and then land on an `Arc` already out of + /// the map, which is a lost cancellation rather than a late one: the `true` + /// tells `cancel_session_impl` to write no status, and the entry's + /// successor (a handoff's replacement) or its remover + /// ([`SessionRegistry::deregister_reporting_cancellation`]) read `None` and + /// take over nothing. + /// + /// The split was never as wide as it looked, which is why it is worth + /// naming what now holds it shut. `apply_cancellation` keeps the reason + /// mutex across `token.cancel()`, and both takeover readers go through + /// [`accepted_cancellation`](RunningSession::accepted_cancellation), which + /// needs that same mutex — so a cancel already inside `apply_cancellation` + /// was serialized against them anyway, leaving only the instant between the + /// map lookup and that mutex. That is an accident of two unrelated lock + /// scopes, invisible at both sites and undone by any tidying that releases + /// the reason guard before firing the token. This lock is the one that is + /// about the invariant. + fn cancel_registered(&self, session_id: &str, completion_reason: CompletionReason) -> bool { + match self.running.get(session_id) { + Some(running_session) => { + running_session.apply_cancellation(completion_reason); + true + } + None => false, + } + } +} + /// A session startup that failed, and the cancellation (if any) the registry /// accepted for it while it was running. /// @@ -325,6 +381,15 @@ impl SessionRegistry { /// Only one can be set at a time — `cancel_or_defer` parks an intent only /// when nothing is registered — so the pending intent is simply preferred. /// + /// Both reads are complete because this whole body runs under the one lock + /// that cancels are applied under (see + /// [`RegistryInner::cancel_registered`]). A cancel racing the swap is + /// therefore on one side of it or the other: applied before, and carried; + /// or after, when it finds the replacement and fires *its* token. Neither + /// is the lost cancellation a split lookup-then-apply would allow, where + /// the cancel answers `true` to a caller that writes no status and then + /// lands on the predecessor this call has already dropped. + /// /// Finding a *live* entry to replace means the handoff: the thread of an /// ordinary session deregisters before its terminal DB write, so while its /// entry is present the row is still `running` and `transition_to_running` @@ -414,12 +479,13 @@ impl SessionRegistry { /// failure reports it instead of dropping it: the cancel that fired this /// token was answered `true` and wrote no status, and the thread that would /// have recorded the terminal state never spawns. The caller records it — - /// see [`finish_cancelled_startup`]. + /// see [`finish_cancelled_before_run`]. /// - /// The report is as tight as the removal: a `cancel` that cloned the entry - /// under the lock but hasn't applied yet lands on an `Arc` already out of - /// the map and goes unreported, the same window - /// `start_session`'s post-`deregister` token re-read documents. + /// The report misses nothing, because the removal and every cancel take the + /// same lock (see [`RegistryInner::cancel_registered`]): a cancel is either + /// applied before the entry leaves the map and reported here, or it arrives + /// after, finds nothing, answers `false`, and is written straight to the + /// store by `cancel_session_impl`'s fallback. fn register_for_startup( &self, session_id: &str, @@ -442,18 +508,21 @@ impl SessionRegistry { } /// Cancel a running session and remember the completion reason it should persist. + /// + /// The `true` is a commitment `cancel_session_impl` relies on — it writes + /// no status of its own — and [`RegistryInner::cancel_registered`] is what + /// lets the registry keep it: by the time this returns, the cancellation is + /// recorded on an entry that was still in the map when it was applied, so + /// whoever takes that entry out sees it and takes over. pub fn cancel_with_completion_reason( &self, session_id: &str, completion_reason: CompletionReason, ) -> bool { - let running_session = self.inner.lock().unwrap().running.get(session_id).cloned(); - if let Some(running_session) = running_session { - running_session.apply_cancellation(completion_reason); - true - } else { - false - } + self.inner + .lock() + .unwrap() + .cancel_registered(session_id, completion_reason) } /// Cancel `session_id` if it's running, otherwise record the cancellation so @@ -466,23 +535,20 @@ impl SessionRegistry { /// would find nothing and silently drop the cancellation. Deferring the /// intent guarantees it lands however long startup takes (e.g. a remote /// review awaiting a network-bound `git rev-parse`), which a single - /// fixed-delay retry could outlast. The check-or-record happens under one - /// lock so it can't interleave with a concurrent `register`. + /// fixed-delay retry could outlast. The check-or-record *and the cancel it + /// may choose* happen under one lock, so neither can interleave with a + /// concurrent `register`: the intent is parked before a registration can + /// claim to have found none, and a direct cancel is applied before a + /// replacement can read past it. pub fn cancel_or_defer(&self, session_id: &str, completion_reason: CompletionReason) { let mut inner = self.inner.lock().unwrap(); - match inner.running.get(session_id).cloned() { - // Registered already (possibly between an earlier cancel attempt and - // this call) — cancel it directly. - Some(running_session) => { - drop(inner); - running_session.apply_cancellation(completion_reason); - } - // Not registered yet — record the intent for `register` to apply. - None => { - inner - .pending_cancellations - .insert(session_id.to_string(), completion_reason); - } + // Registered already (possibly between an earlier cancel attempt and + // this call) — cancel it directly. Otherwise record the intent for + // `register` to apply. + if !inner.cancel_registered(session_id, completion_reason.clone()) { + inner + .pending_cancellations + .insert(session_id.to_string(), completion_reason); } } @@ -775,7 +841,12 @@ fn refuse_start_during_shutdown( registry: &SessionRegistry, ) -> SessionStartOutcome { log::info!("Refusing to start session {session_id}: a shutdown is under way"); - finish_cancelled_before_run( + // No drain follows this terminal state, unlike every other one: the drain + // is `is_quitting`-gated and a refusal only happens while that is true, so + // it would answer `Ok(false)` without touching a row. The branch's queue + // stays queued and unclaimed, which is what a quit wants — those rows are + // still available to another instance, and to the next launch. + let _ = finish_cancelled_before_run( session_id, branch_id, project_id, @@ -827,6 +898,10 @@ fn refused_start_completion_reason( /// The event is emitted whether or not the transition won, matching the /// session thread's terminal emit: a row already moved on (deleted, say) still /// needs the client's `running` state cleaned up. +/// +/// Returns whether the transition won, which is what says this path owns the +/// terminal state — and so owes the branch queue the drain that state unblocks. +#[must_use] fn finish_cancelled_before_run( session_id: &str, branch_id: Option, @@ -834,7 +909,7 @@ fn finish_cancelled_before_run( store: &Store, app_handle: &AppHandle, completion_reason: CompletionReason, -) { +) -> bool { let transitioned = store .transition_from_running( session_id, @@ -856,6 +931,7 @@ fn finish_cancelled_before_run( branch_id, project_id, ); + transitioned } /// Start a session: persist the user message, spawn the agent, stream to DB. @@ -981,7 +1057,7 @@ pub fn start_session( Ok(started) => started, Err(failure) => { if let Some(completion_reason) = failure.accepted_cancellation { - finish_cancelled_before_run( + let transitioned = finish_cancelled_before_run( &config.session_id, config.branch_id.clone(), config.project_id.clone(), @@ -989,6 +1065,25 @@ pub fn start_session( &app_handle, completion_reason, ); + // This is a terminal state like any other, so it owes the + // branch its drain: the session thread's terminal path kicks + // one for every branch session it ends, cancelled ones + // included, and nothing else will do it for a session whose + // thread never spawned — the `Err` below reaches only a log + // line in `drain_queued_sessions_for_branch`'s callers, and the + // drain that produced this session has already aborted on it. + // Without this the branch's remaining queued rows sit parked + // until some unrelated session happens to finish. + if transitioned { + drain_queued_after_terminal_state( + Arc::clone(&store), + Arc::clone(®istry), + app_handle.clone(), + config.session_id.clone(), + config.branch_id.clone(), + false, + ); + } } return Err(failure.error); } @@ -1524,13 +1619,11 @@ pub fn start_session( if transitioned { let branch_id = config.branch_id.clone(); // Read the token again rather than reusing what the terminal-state - // match saw, to catch a Stop that reached the registry before the - // `deregister` above but only fires the token after that match read - // it: `cancel_with_completion_reason` clones the - // `Arc` out from under the registry lock and calls - // `apply_cancellation` after releasing it, so the flip can land any - // time after the clone — including once the session has left the - // map. + // match saw, to catch a Stop that landed between that read and the + // `deregister` above. Cancels apply under the registry lock, so + // that is now the whole of the window — a Stop can no longer be + // mid-flight past the deregister, holding a clone of an entry the + // map has already dropped. // // A Stop landing *later* — during the seconds-long post-completion // hooks, say — never reaches this token at all: `apply_cancellation` @@ -1860,7 +1953,7 @@ pub fn start_pipeline_session( config.project_id.clone(), ); if transitioned { - drain_queued_after_pipeline_terminal( + drain_queued_after_terminal_state( Arc::clone(&store_for_status), Arc::clone(®istry), app_handle.clone(), @@ -1942,47 +2035,16 @@ pub fn start_pipeline_session( } if let Err(e) = handoff { log::error!("Failed to start AI session after pipeline handoff: {e}"); - // If the handoff came from an explicit AiHandoff step, mark - // it as failed so the UI doesn't show a perpetual spinner. - if let Some(step_idx) = ai_step_index { - if let Ok(Some(session)) = store_for_status.get_session(&session_id) { - if let Some(mut pipeline) = session.pipeline { - if step_idx < pipeline.steps.len() { - pipeline.steps[step_idx].status = StepStatus::Failed; - pipeline.steps[step_idx].error = - Some(format!("Failed to start AI session: {e}")); - pipeline.steps[step_idx].completed_at = - Some(crate::store::now_timestamp()); - let _ = store_for_status - .update_session_pipeline(&session_id, &pipeline); - emit_pipeline_step( - &app_handle, - &session_id, - step_idx, - &pipeline.steps[step_idx], - ); - } - } - } - } resolve_pipeline_artifacts_without_ai(&config, &store_for_status, false); - let transitioned = finish_failed_pipeline_handoff_start( + if finish_failed_pipeline_handoff_start( + &config, &store_for_status, ®istry, - &session_id, - &e, - ); - emit_status( &app_handle, - &session_id, - "error", - Some(e), - Some(&CompletionReason::Crashed), - config.branch_id.clone(), - config.project_id.clone(), - ); - if transitioned { - drain_queued_after_pipeline_terminal( + &e, + ai_step_index, + ) { + drain_queued_after_terminal_state( Arc::clone(&store_for_status), Arc::clone(®istry), app_handle.clone(), @@ -2040,7 +2102,7 @@ pub fn start_pipeline_session( config.project_id.clone(), ); if transitioned { - drain_queued_after_pipeline_terminal( + drain_queued_after_terminal_state( Arc::clone(&store_for_status), Arc::clone(®istry), app_handle.clone(), @@ -2074,7 +2136,7 @@ pub fn start_pipeline_session( config.project_id.clone(), ); if transitioned { - drain_queued_after_pipeline_terminal( + drain_queued_after_terminal_state( Arc::clone(&store_for_status), Arc::clone(®istry), app_handle.clone(), @@ -2090,21 +2152,93 @@ pub fn start_pipeline_session( Ok(SessionStartOutcome::Started) } -fn finish_failed_pipeline_handoff_start( +/// Record and announce the terminal state of a pipeline whose AI handoff +/// couldn't be started, and report whether this path is the one that owns it. +/// +/// Every side effect hangs off winning `transition_from_running`, matching the +/// other three [`PipelineOutcome`] arms, because losing it here means another +/// writer has already recorded a terminal state this one must not talk over. +/// The writer it loses to is usually the startup itself: a Stop that lands +/// while the driver resolves comes back out of [`start_session`] as an `Err` +/// only *after* [`finish_cancelled_before_run`] has written `cancelled` and +/// emitted it. An ungated `error`/`Crashed` behind that leaves the row saying +/// cancelled and the client saying errored until its next refetch, and stamps +/// the AiHandoff step "Failed to start AI session: …" when what happened was +/// the user's Stop. +/// +/// The exception is a row that is *gone* rather than moved on — the user +/// deleted the pending commit mid-pipeline. Nobody else emitted anything for +/// it, and the event is all that clears the client's `running` state, which is +/// what the emit was unconditional for in the first place. +fn finish_failed_pipeline_handoff_start( + config: &PipelineConfig, store: &Store, registry: &SessionRegistry, - session_id: &str, + app_handle: &AppHandle, error: &str, + ai_step_index: Option, ) -> bool { + let session_id = &config.session_id; registry.deregister(session_id); - store + let transitioned = store .transition_from_running( session_id, SessionStatus::Error, Some(error), Some(&CompletionReason::Crashed), ) - .unwrap_or(false) + .unwrap_or(false); + + // If the handoff came from an explicit AiHandoff step, mark it as failed so + // the UI doesn't show a perpetual spinner. + if transitioned { + if let Some(step_index) = ai_step_index { + mark_ai_handoff_step_failed(store, app_handle, session_id, step_index, error); + } + } + + if transitioned || matches!(store.get_session(session_id), Ok(None)) { + emit_status( + app_handle, + session_id, + SessionStatus::Error.as_str(), + Some(error.to_string()), + Some(&CompletionReason::Crashed), + config.branch_id.clone(), + config.project_id.clone(), + ); + } + + transitioned +} + +/// Stamp the AiHandoff step that couldn't start as failed, and publish it. +fn mark_ai_handoff_step_failed( + store: &Store, + app_handle: &AppHandle, + session_id: &str, + step_index: usize, + error: &str, +) { + let Ok(Some(session)) = store.get_session(session_id) else { + return; + }; + let Some(mut pipeline) = session.pipeline else { + return; + }; + if step_index >= pipeline.steps.len() { + return; + } + pipeline.steps[step_index].status = StepStatus::Failed; + pipeline.steps[step_index].error = Some(format!("Failed to start AI session: {error}")); + pipeline.steps[step_index].completed_at = Some(crate::store::now_timestamp()); + let _ = store.update_session_pipeline(session_id, &pipeline); + emit_pipeline_step( + app_handle, + session_id, + step_index, + &pipeline.steps[step_index], + ); } /// Error message for an aborted pipeline, or `None` when the abort is an expected @@ -2337,7 +2471,17 @@ fn finalize_rebase_pipeline_without_ai(config: &PipelineConfig, store: &Store) { } } -fn drain_queued_after_pipeline_terminal( +/// Kick the queue progression a session's terminal state unblocks: its own +/// queued follow-up message when the turn earned one, and the next queued +/// session on its branch. +/// +/// Every caller gates this on winning `transition_from_running` — losing means +/// another writer owns the terminal state, and the drain with it. Shared by the +/// pipeline's four terminal arms and by the startup-failure path, whose row is +/// just as terminal (`cancelled`, written by [`finish_cancelled_before_run`]) +/// and whose branch queue would otherwise sit parked until some unrelated +/// session happened to finish. +fn drain_queued_after_terminal_state( store: Arc, registry: Arc, app_handle: AppHandle, @@ -2364,7 +2508,7 @@ fn drain_queued_after_pipeline_terminal( Ok(true) => log::info!("Drained queued follow-up message for session {session_id}"), Ok(false) => {} Err(e) => log::error!( - "Failed to drain queued follow-up message after pipeline terminal state for session {session_id}: {e}" + "Failed to drain queued follow-up message after a terminal state for session {session_id}: {e}" ), } } @@ -2382,7 +2526,7 @@ fn drain_queued_after_pipeline_terminal( Ok(true) => log::info!("Drained next queued session for branch {branch_id}"), Ok(false) => {} Err(e) => log::error!( - "Failed to drain queued sessions after pipeline terminal state for branch {branch_id}: {e}" + "Failed to drain queued sessions after a terminal state for branch {branch_id}: {e}" ), } } @@ -2951,8 +3095,8 @@ fn send_signal_to_pipeline_process_group(pid: u32, signal: libc::c_int) -> io::R } } -fn emit_pipeline_step( - app_handle: &AppHandle, +fn emit_pipeline_step( + app_handle: &AppHandle, session_id: &str, step_index: usize, step: &crate::store::PipelineStepStatus, @@ -4554,31 +4698,146 @@ mod tests { assert!(!prompt_output.contains("20%")); } - #[test] - fn failed_pipeline_handoff_start_cleans_running_state() { + /// A session at its AI handoff: the row is `running` and the handoff step + /// is still pending, which is the state `start_session` is called in. + fn store_with_pending_ai_handoff() -> (Store, PipelineConfig) { let store = Store::in_memory().unwrap(); - let session = crate::store::Session::new_running("handoff", std::path::Path::new("/tmp")); + let steps = vec![PipelineStep::AiHandoff { + label: "Write PR title and body".to_string(), + prompt_template: "{step_outputs}".to_string(), + }]; + let pipeline = PipelineExecution::from_steps(&steps); + let mut session = + crate::store::Session::new_running("handoff", std::path::Path::new("/tmp")); + session.pipeline = Some(pipeline.clone()); store.create_session(&session).unwrap(); + let config = PipelineConfig { + session_id: session.id, + prompt: "handoff".to_string(), + steps, + pipeline, + working_dir: PathBuf::from("/tmp"), + pre_head_sha: None, + provider: None, + workspace_name: None, + remote_working_dir: None, + branch_id: None, + project_id: None, + }; + (store, config) + } + + #[test] + fn failed_pipeline_handoff_start_cleans_running_state() { + let (store, config) = store_with_pending_ai_handoff(); + let app = mock_app(); + let registry = SessionRegistry::new(); - registry.register(&session.id); - assert!(registry.is_running(&session.id)); + registry.register(&config.session_id); + assert!(registry.is_running(&config.session_id)); - finish_failed_pipeline_handoff_start( + assert!(finish_failed_pipeline_handoff_start( + &config, &store, ®istry, - &session.id, + app.handle(), "provider unavailable", - ); + Some(0), + )); - assert!(!registry.is_running(&session.id)); - let failed = store.get_session(&session.id).unwrap().unwrap(); + assert!(!registry.is_running(&config.session_id)); + let failed = store.get_session(&config.session_id).unwrap().unwrap(); assert_eq!(failed.status, SessionStatus::Error); assert_eq!( failed.error_message.as_deref(), Some("provider unavailable") ); assert_eq!(failed.completion_reason, Some(CompletionReason::Crashed)); + let step = &failed.pipeline.unwrap().steps[0]; + assert_eq!(step.status, StepStatus::Failed); + assert_eq!( + step.error.as_deref(), + Some("Failed to start AI session: provider unavailable") + ); + } + + /// The handoff failure that *is* a Stop: `start_session` returns the + /// startup error only after `finish_cancelled_before_run` has written + /// `cancelled` and emitted it, so this path loses the transition — and + /// everything it would otherwise have said loses with it. Ungated, the row + /// reads cancelled while the last event the client saw says errored, and + /// the AiHandoff step is blamed for the user's Stop. + #[test] + fn a_failed_handoff_leaves_a_cancelled_row_and_its_step_alone() { + let (store, config) = store_with_pending_ai_handoff(); + let app = mock_app(); + + assert!( + finish_cancelled_before_run( + &config.session_id, + None, + None, + &store, + app.handle(), + CompletionReason::Interrupted, + ), + "the Stop's write is the one that takes the row" + ); + + assert!(!finish_failed_pipeline_handoff_start( + &config, + &store, + &SessionRegistry::new(), + app.handle(), + "No ACP agent found.", + Some(0), + )); + + let row = store.get_session(&config.session_id).unwrap().unwrap(); + assert_eq!(row.status, SessionStatus::Cancelled); + assert_eq!(row.completion_reason, Some(CompletionReason::Interrupted)); + let step = &row.pipeline.unwrap().steps[0]; + assert_eq!( + step.status, + StepStatus::Pending, + "a Stop is not the handoff step failing" + ); + assert_eq!(step.error, None); + } + + /// The drain the startup-failure path kicks hangs off this answer: a + /// terminal state is drained by whoever recorded it, and a write that lost + /// the row recorded nothing to drain on. + #[test] + fn a_cancel_before_the_run_reports_whether_it_recorded_the_terminal_state() { + let store = Store::in_memory().unwrap(); + let session = crate::store::Session::new_running("prompt", &PathBuf::from("/tmp")); + store.create_session(&session).unwrap(); + let app = mock_app(); + + assert!(finish_cancelled_before_run( + &session.id, + Some("branch-1".to_string()), + None, + &store, + app.handle(), + CompletionReason::Interrupted, + )); + assert!( + !finish_cancelled_before_run( + &session.id, + Some("branch-1".to_string()), + None, + &store, + app.handle(), + CompletionReason::AppQuit, + ), + "a row that is no longer running was recorded by someone else" + ); + + let row = store.get_session(&session.id).unwrap().unwrap(); + assert_eq!(row.completion_reason, Some(CompletionReason::Interrupted)); } #[test] @@ -4860,6 +5119,189 @@ mod tests { ); } + /// A waker that parks the thread waking it until it is released. + /// + /// `CancellationToken::cancel` notifies its waiters synchronously, so a + /// waker enrolled on a session's token stops a cancelling thread *inside* + /// [`RunningSession::apply_cancellation`], which is the only interposition + /// point this race has. Nothing in production wakes like this. + /// + /// What the two tests below pin is therefore which lock keeps a cancel and + /// a takeover apart, not that they are kept apart at all: they would also + /// pass against a lookup-then-apply split, because `apply_cancellation` + /// holds the reason mutex across `token.cancel()` and every takeover reader + /// wants that mutex too. Release the reason guard before firing the token + /// and the split loses both tests while the registry lock keeps them. + struct ParkingWaker { + entered: std::sync::Mutex>, + release: std::sync::Mutex>, + } + + impl std::task::Wake for ParkingWaker { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + let _ = self.entered.lock().unwrap().send(()); + let _ = self.release.lock().unwrap().recv(); + } + } + + /// A cancel stopped part-way through applying itself, holding whatever the + /// registry handed it. + struct ParkedCancel { + canceller: Option>, + release: std::sync::mpsc::Sender<()>, + /// Kept alive for the whole park: dropping the enrolment early would + /// take the notify lock the parked thread is inside. + _enrolled: std::pin::Pin>, + _waker: std::task::Waker, + } + + impl ParkedCancel { + /// Let the cancel finish, and answer what the registry told it. + fn finish(mut self) -> bool { + self.release + .send(()) + .expect("the parked cancel must still be waiting"); + self.canceller + .take() + .expect("a parked cancel is only finished once") + .join() + .expect("the cancelling thread must not panic") + } + } + + /// Cancel `session_id` on another thread and stop it mid-apply: the + /// completion reason is recorded and the token fired, but the thread has + /// not yet returned the `true` that tells `cancel_session_impl` to write no + /// status of its own. + fn park_a_cancel_mid_apply( + registry: &Arc, + session_id: &str, + token: &CancellationToken, + completion_reason: CompletionReason, + ) -> ParkedCancel { + let (entered_tx, entered_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let waker: std::task::Waker = Arc::new(ParkingWaker { + entered: std::sync::Mutex::new(entered_tx), + release: std::sync::Mutex::new(release_rx), + }) + .into(); + let mut enrolled = Box::pin(token.clone().cancelled_owned()); + assert!( + std::future::Future::poll( + enrolled.as_mut(), + &mut std::task::Context::from_waker(&waker), + ) + .is_pending(), + "the token must still be live for the waker to enrol on" + ); + + let canceller = { + let registry = Arc::clone(registry); + let session_id = session_id.to_string(); + std::thread::spawn(move || { + registry.cancel_with_completion_reason(&session_id, completion_reason) + }) + }; + entered_rx + .recv() + .expect("the cancel must reach the token's waiters"); + + ParkedCancel { + canceller: Some(canceller), + release: release_tx, + _enrolled: enrolled, + _waker: waker, + } + } + + /// How long to leave a thread that must be blocked a chance to prove it + /// isn't. Only ever reported as "still running", so a slow machine can make + /// this test weaker, never wrong. + const BLOCKED_ENOUGH: Duration = Duration::from_millis(50); + + /// The carry has to see a cancel that is still being applied, not just one + /// that finished. Reading past it would give the AI session a clean token + /// while `cancel` answered `true` — so `cancel_session_impl` writes no + /// status, the id stays in shutdown's cancel snapshot with nothing that can + /// honour it, `wait_for_sessions` burns the whole budget, and `app.exit(0)` + /// orphans the agent child. + #[test] + fn register_carries_a_cancellation_that_is_still_being_applied() { + let registry = Arc::new(SessionRegistry::new()); + let pipeline_token = registry.register("session-handoff"); + let parked = park_a_cancel_mid_apply( + ®istry, + "session-handoff", + &pipeline_token, + CompletionReason::AppQuit, + ); + + let swapping = { + let registry = Arc::clone(®istry); + std::thread::spawn(move || registry.register("session-handoff")) + }; + std::thread::sleep(BLOCKED_ENOUGH); + assert!( + !swapping.is_finished(), + "the swap must wait on the in-flight cancel rather than read past it" + ); + + assert!( + parked.finish(), + "the cancel found an entry, so nothing else wrote a status for it" + ); + let ai_token = swapping.join().expect("the swapping thread must not panic"); + assert!( + ai_token.is_cancelled(), + "the handed-off session must start already cancelled" + ); + assert_eq!( + registry.cancellation_completion_reason("session-handoff"), + Some(CompletionReason::AppQuit), + "and must keep the reason its terminal write has to persist" + ); + } + + /// The same for the other way an entry leaves without an observer. A + /// startup-failure removal that read past an in-flight cancel would report + /// `None`, `start_session` would write no `cancelled` row, and the Stop + /// would produce nothing but an errored session. + #[test] + fn a_removal_reports_a_cancellation_that_is_still_being_applied() { + let registry = Arc::new(SessionRegistry::new()); + let token = registry.register("session-failing-startup"); + let parked = park_a_cancel_mid_apply( + ®istry, + "session-failing-startup", + &token, + CompletionReason::Interrupted, + ); + + let removing = { + let registry = Arc::clone(®istry); + std::thread::spawn(move || { + registry.deregister_reporting_cancellation("session-failing-startup") + }) + }; + std::thread::sleep(BLOCKED_ENOUGH); + assert!( + !removing.is_finished(), + "the removal must wait on the in-flight cancel rather than read past it" + ); + + assert!(parked.finish()); + assert_eq!( + removing.join().expect("the removing thread must not panic"), + Some(CompletionReason::Interrupted), + "the entry left without an observer, so its cancellation has to come back" + ); + } + /// The carry-forward is scoped to a cancelled predecessor: an ordinary /// handoff hands over a live session, and starting it pre-cancelled would /// kill the AI turn the pipeline just asked for. From c6f8509ce649051de1d5aaa16f7652fb90fccf11 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 4 Sep 2026 10:59:37 +1000 Subject: [PATCH 12/13] fix(lifecycle): claim a diagram sub-session's slot before asking about the quit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `1775b4e8` gated every session start on the shutdown, one gate in each of the two funnels, and named what it deliberately left out: `generate_pikchr`'s diagram sub-session reaches its agent through `register_external` and `driver.run` rather than `start_session`, so a tool call landing after `cancel_owned_sessions` has snapshotted the registry registers a token nothing will ever fire, and spawns an agent child `app.exit(0)` then orphans — own process group, and an exit runs no `kill_on_drop` destructors. Narrower than anything closed there (the parent session has to be mid-turn and already being cancelled), but the same invariant, and the last way into an agent process that the shutdown could not see. `register_external` is not the chokepoint, despite being the shape the gap is described by. Its only production caller is that worker — the other two are tests, which have no business consulting a quit flag — the registry has no route to the Tauri-managed `QuitState`, and, decisively, the correct order is claim *then* ask, which a gate inside a function whose whole job is to hand out a token cannot express. So the gate goes in the worker, as `pikchr_mcp::reserve_child_session`, and the obligations it discharges are written on `register_external` where the next external registrant will meet them. **Claiming first is what closes the window rather than narrowing it.** The funnel gates in `start_session` / `start_pipeline_session` read `quit_in_progress` with nothing yet registered for a snapshot to find, so they shrink their race and say so. Here both sides are ordered: `shutdown_cleanup` publishes the flag and *then* snapshots the registry; the worker registers and *then* reads the flag. - A claim landing before the snapshot is in it. Shutdown fires the token, `forward_user_cancel` hands it to the worker, the `is_cancelled` check `generate_pikchr_source` takes before each `driver.run` refuses to start an agent at all, and `wait_for_sessions` holds the exit open until the worker deregisters. If the worker's own startup outruns `SHUTDOWN_BUDGET` — the preview server plus `AcpDriver::new`'s login-shell probes — shutdown times out, warns, sweeps the row and exits, killing the thread with the process: the clean end `89c821c3` described for `register_for_startup`, and for the same reason, that the gate means no agent child was ever started. - A claim landing after the snapshot reads a flag already published, and refuses here. Both fail together only if the claim landed after the snapshot *and* the read landed before the publish — impossible given the two program orders (publish → snapshot, register → read) and the total order of the registry mutex. **A refusal leaves nothing behind**, which is what lets it be a plain early return instead of `1775b4e8`'s `finish_cancelled_before_run`. The claim now happens before anything is persisted, so `create_pikchr_child_session` splits into `new_pikchr_child_session` (build the row) and `persist_pikchr_child_session` (write it), and a refused call writes no session row for the sweep to chase and announces no diagram session into the parent's transcript that never drew anything. Holding a registry entry for an id whose row doesn't exist yet is safe by construction: the row and the announcement are what publish the id, both come after the check, and the only thing that can reach an unpublished id is the shutdown walking every registered one, which fires the token and waits for the entry to go. **What the refusal returns** is an error, like every other failure in this tool. There is no MCP way to say "stop your turn", so a clear terminal refusal and a retryable failure are the same wire shape; what makes the retry harmless is that the refusal is free — no row, no registry entry, and above all no agent child — so a loop spins on an atomic load while the parent session, cancelled by the same shutdown, is torn down underneath it. The message says not to bother. **The wait was also measuring the wrong thing.** Shutdown does already cover `register_external` entries — `cancel_owned_sessions` and `wait_for_sessions` walk the whole registry, not just what the runner started — but the guard was held by the parent MCP request future, which only *awaits* the worker and is dropped as soon as the parent session's runtime goes down. On the quit path the parent is cancelled in the same loop as the child, so the two teardowns race, and a parent finishing first retired the child's entry while the specialist's agent CLI was still being stopped: `wait_for_sessions` returns, the sweep runs, and the exit proceeds over the child it was supposed to be waiting for. The guard moves into the worker thread, declared first so it drops last — after that thread's runtime and every task on it, which is where the agent child actually lives. The same move fixes the normal path's inconsistency in the safe direction: the entry now retires a hair after the tool result rather than a hair before the worker is done. Tested where the ordering is reachable without an `AppHandle`, which is why `reserve_child_session` takes the quit check as a parameter: a reservation is visible to a registry snapshot taken at the last instant the gate could still say "carry on", and a refusal releases the slot it claimed so a shutdown that did catch it stops waiting. The gate's own call site stays out of reach — `is_quitting` wants the concrete `AppHandle`, and driving `generate_pikchr` needs an MCP `RequestContext`. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/pikchr_mcp.rs | 193 +++++++++++++++++--- apps/staged/src-tauri/src/session_runner.rs | 13 ++ 2 files changed, 180 insertions(+), 26 deletions(-) diff --git a/apps/staged/src-tauri/src/pikchr_mcp.rs b/apps/staged/src-tauri/src/pikchr_mcp.rs index 48fbe98fc..73faf1e4a 100644 --- a/apps/staged/src-tauri/src/pikchr_mcp.rs +++ b/apps/staged/src-tauri/src/pikchr_mcp.rs @@ -61,7 +61,7 @@ use rmcp::{schemars, tool, tool_handler, tool_router, ErrorData, Peer, RoleServe use crate::agent::AcpDriver; use crate::pikchr_subsession::{CancelReason, GenOutcome, LastRenderSlot, ACCEPT_SENTINEL}; -use crate::session_runner::SessionRegistry; +use crate::session_runner::{ExternalSessionRegistration, SessionRegistry}; use crate::store::{AcpMessageMetadata, CompletionReason, Session, SessionStatus, Store}; /// Wall-clock cap for one `generate_pikchr` call. Each call spins a provider @@ -102,6 +102,19 @@ const PIKCHR_CHILD_SESSION_PROMPT: &str = "Generate Pikchr diagram"; /// only names the child session once the specialist finishes, so this early /// announcement is what lets the UI offer "open diagram session" mid-run. const PIKCHR_SESSION_STARTED_EVENT: &str = "pikchr_session_started"; +/// Handed back when a `generate_pikchr` call arrives after a shutdown has been +/// claimed (see [`reserve_child_session`]). +/// +/// There is no MCP way to say "stop your turn", so this is a failed tool call +/// like any other and the calling agent may well retry it. That's tolerable +/// because the refusal is free — no store row, no registry entry left behind, +/// and above all no agent child — so a retry loop spins on an atomic load and +/// spawns nothing, while the parent session, cancelled by the same shutdown, is +/// being torn down underneath it. The wording still says not to bother, since +/// the only thing an agent can usefully do here is stop asking. +const SHUTDOWN_REFUSAL_MESSAGE: &str = + "Staged is shutting down, so no diagram session was started. Retrying will not help; \ +write the Pikchr by hand or generate the diagram after restarting."; #[derive(serde::Deserialize, schemars::JsonSchema)] struct GeneratePikchrParams { @@ -838,9 +851,36 @@ rendered PNG preview you may open as an optional final check." ), None => (self.provider_id.clone(), Vec::new(), None), }; - let session = create_pikchr_child_session(&self.store, &provider_id) - .map_err(|e| ErrorData::internal_error(e, None))?; + let session = new_pikchr_child_session(&provider_id); + + // Claim the child session's slot in the SessionRegistry, or refuse the + // call because a shutdown has already been claimed. + // + // The slot is what makes the child session cancellable at all: the Stop + // control in the opened diagram session — and, identically, a quit's + // `cancel_owned_sessions` — fires the registered token, which the worker + // forwards onto its own (recording the reason first), instead of taking + // `cancel_session`'s fallback of writing Cancelled to a store row this + // worker never re-reads. + // + // Refusing is the other half. Past the shutdown's registry snapshot + // nothing would ever fire this token, so the worker below would spawn an + // agent child that `app.exit(0)` orphans — own process group, and an + // exit runs no `kill_on_drop` destructors. See `reserve_child_session` + // for why the claim has to come before the question. + let Some(registration) = reserve_child_session(&self.registry, &session.id, || { + crate::app_lifecycle::is_quitting(&self.app_handle) + }) else { + return Err(ErrorData::internal_error( + SHUTDOWN_REFUSAL_MESSAGE.to_string(), + None, + )); + }; + let user_cancel = registration.token().clone(); + let inner_session_id = session.id.clone(); + persist_pikchr_child_session(&self.store, &session) + .map_err(|e| ErrorData::internal_error(e, None))?; announce_pikchr_child_session(&self.store, &self.parent_session_id, &inner_session_id); let store = Arc::clone(&self.store); // The full grammar text is inlined into the sub-agent's prompt rather @@ -865,16 +905,6 @@ rendered PNG preview you may open as an optional final check." let worker_cancel_reason = Arc::new(CancelReason::new()); let _cancel_on_drop = cancel.drop_guard(); - // Register the child session in the SessionRegistry under its own - // token so the Stop control in the opened diagram session terminates - // the actual work: `cancel_session` fires the registered token, which - // the worker forwards onto its own token (recording the reason first) - // — instead of taking the fallback path that just writes Cancelled to - // a store row this worker never re-reads. The registration guard - // deregisters when this call ends, however it ends. - let registration = self.registry.register_external(&inner_session_id); - let user_cancel = registration.token().clone(); - // The ACP driver spawns tasks via `spawn_local`, which requires a // `LocalSet`; the MCP server's request tasks don't run inside one. So // drive the whole generation loop on a dedicated thread with its own @@ -883,6 +913,19 @@ rendered PNG preview you may open as an optional final check." let worker_store = Arc::clone(&store); let worker_session_id = inner_session_id.clone(); std::thread::spawn(move || { + // Declared first so it drops *last* — after this thread's runtime + // and every task on it, which is where the specialist's agent child + // actually lives. + // + // The registry entry is what `wait_for_sessions` polls, so it has to + // span this worker rather than the parent MCP request future that + // opened it. That future only *awaits* the work, and it is dropped + // the moment the parent session's runtime goes down — which on the + // quit path runs in parallel with this teardown, not after it. Held + // over there, a parent that finished first would retire the entry + // while the agent CLI here was still being stopped, and the exit + // would proceed straight over the top of it. + let _registration = registration; let rt = match tokio::runtime::Builder::new_current_thread() .enable_all() .build() @@ -1007,8 +1050,13 @@ accepted a render, so the diagram run was cancelled.", }); // A Stop pressed in the child diagram session fires the token // registered in the SessionRegistry; forward it to the - // worker's token so the run actually terminates. Both watcher - // tasks are dropped with this LocalSet. + // worker's token so the run actually terminates. A quit fires + // the same token — `cancel_owned_sessions` walks every + // registered id — so this is also how a shutdown reaches the + // specialist, and the `is_cancelled` check `generate_pikchr_source` + // takes before each `driver.run` is what keeps it from starting + // an agent the exit would then orphan. Both watcher tasks are + // dropped with this LocalSet. tokio::task::spawn_local(forward_user_cancel( user_cancel, Arc::clone(&worker_cancel_reason), @@ -1138,15 +1186,69 @@ async fn forward_user_cancel( worker_cancel.cancel(); } -fn create_pikchr_child_session(store: &Store, provider_id: &str) -> Result { - let mut session = Session::new_running(PIKCHR_CHILD_SESSION_PROMPT, &std::env::temp_dir()); - if !provider_id.is_empty() { - session = session.with_provider(provider_id); +/// Build the child diagram session row, unpersisted. +/// +/// Split from persisting it so its id can be reserved in the registry first — +/// see [`reserve_child_session`], which is what a refusal that has written +/// nothing depends on. +fn new_pikchr_child_session(provider_id: &str) -> Session { + let session = Session::new_running(PIKCHR_CHILD_SESSION_PROMPT, &std::env::temp_dir()); + if provider_id.is_empty() { + session + } else { + session.with_provider(provider_id) } +} + +fn persist_pikchr_child_session(store: &Store, session: &Session) -> Result<(), String> { store - .create_session(&session) - .map_err(|e| format!("Failed to create Pikchr child session: {e}"))?; - Ok(session) + .create_session(session) + .map_err(|e| format!("Failed to create Pikchr child session: {e}")) +} + +/// Reserve `session_id`'s slot in the [`SessionRegistry`], or report that a +/// shutdown has been claimed and this diagram session must not start. +/// +/// Registering *before* asking is the whole mechanism, and it is why the gate +/// can't just sit at the top of the tool call. +/// [`crate::app_lifecycle::shutdown_cleanup`] publishes `quit_in_progress` and +/// *then* snapshots the registry, so: +/// +/// - a reservation that lands before that snapshot is in it — shutdown fires +/// this token, the worker's pre-`driver.run` check refuses to start an agent, +/// and `wait_for_sessions` holds the exit open until the worker deregisters; +/// - a reservation that lands after it reads a flag already published, and +/// refuses here. +/// +/// Both can only fail together if this register landed after the snapshot *and* +/// this read landed before the publish, which the two program orders (publish +/// → snapshot, register → read) and the registry mutex's total order rule out. +/// So this closes the window rather than narrowing it, unlike the funnel gates +/// in `start_session` / `start_pipeline_session`, which read the flag with +/// nothing yet registered for a snapshot to find. +/// +/// Holding a registry entry for an id whose row doesn't exist yet is +/// deliberate. Nothing can ask about the id in that gap — the store row and the +/// parent-transcript announcement are what publish it, and both come after — +/// and the shutdown that can reach it (which walks every registered id) only +/// fires the token and waits for the entry to go. A refusal therefore leaves +/// nothing behind at all: no row for the sweep to chase, and no diagram session +/// in the parent's transcript that never drew anything. +/// +/// `quitting` is a parameter rather than an inline `app_lifecycle::is_quitting` +/// so the ordering is testable without an `AppHandle`. +fn reserve_child_session( + registry: &Arc, + session_id: &str, + quitting: impl FnOnce() -> bool, +) -> Option { + let registration = registry.register_external(session_id); + if quitting() { + // Dropping the guard deregisters, so a shutdown that did catch this + // entry in its snapshot stops waiting on it immediately. + return None; + } + Some(registration) } /// Write a hidden metadata row into the parent session's transcript naming the @@ -1457,8 +1559,8 @@ arrow from COLL.e to SNOW.w"#; fn create_pikchr_child_session_persists_running_provider_session() { let store = Store::in_memory().expect("in-memory store"); - let session = - create_pikchr_child_session(&store, "fake-agent").expect("create child session"); + let session = new_pikchr_child_session("fake-agent"); + persist_pikchr_child_session(&store, &session).expect("persist child session"); assert_eq!(session.prompt, PIKCHR_CHILD_SESSION_PROMPT); assert_eq!(session.status, SessionStatus::Running); @@ -1474,6 +1576,45 @@ arrow from COLL.e to SNOW.w"#; assert_eq!(persisted.provider.as_deref(), Some("fake-agent")); } + /// The reservation is in the registry *before* the quit gate reads, which + /// is what makes the refusal airtight rather than narrow: a shutdown that + /// snapshots the registry at any point up to this read finds the entry and + /// cancels it, and one that snapshots later has already published the flag + /// this read sees. + #[test] + fn a_reserved_child_session_is_registered_before_the_quit_gate_reads() { + let registry = Arc::new(SessionRegistry::new()); + let snapshot = std::cell::RefCell::new(Vec::new()); + + let registration = reserve_child_session(®istry, "diagram-child", || { + // Stands in for `cancel_owned_sessions` snapshotting the registry + // at the last instant this gate could still say "carry on". + *snapshot.borrow_mut() = registry.running_session_ids(); + false + }) + .expect("no shutdown claimed, so the reservation stands"); + + assert_eq!(snapshot.into_inner(), vec!["diagram-child".to_string()]); + assert!(registry.is_running("diagram-child")); + + // And the guard is what holds it: the worker owning it is what makes + // shutdown's wait span the specialist rather than the MCP request. + drop(registration); + assert!(!registry.is_running("diagram-child")); + } + + /// A refusal leaves nothing behind — in particular no registry entry, so a + /// shutdown that did catch it in its snapshot stops waiting on it. + #[test] + fn refusing_a_child_session_releases_the_slot_it_claimed() { + let registry = Arc::new(SessionRegistry::new()); + + let refused = reserve_child_session(®istry, "diagram-child", || true); + + assert!(refused.is_none()); + assert!(registry.running_session_ids().is_empty()); + } + #[test] fn progress_keepalive_reports_elapsed_seconds_with_no_total() { let token = ProgressToken(rmcp::model::NumberOrString::Number(7)); @@ -1539,8 +1680,8 @@ arrow from COLL.e to SNOW.w"#; async fn registry_stop_terminates_the_run_and_reads_as_a_user_stop() { let registry = Arc::new(SessionRegistry::new()); let store = Arc::new(Store::in_memory().expect("in-memory store")); - let session = - create_pikchr_child_session(&store, "fake-agent").expect("create child session"); + let session = new_pikchr_child_session("fake-agent"); + persist_pikchr_child_session(&store, &session).expect("persist child session"); let registration = registry.register_external(&session.id); let worker_cancel = CancellationToken::new(); diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index ac9a8062f..1948675cf 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -666,6 +666,19 @@ impl SessionRegistry { /// `cancel_session`'s store-write fallback, which the worker never /// observes. Returns a guard exposing the session's cancellation token; /// dropping the guard deregisters the session. + /// + /// The entry carries the same shutdown contract as a `start_session` one, + /// because `cancel_owned_sessions` and `wait_for_sessions` walk the whole + /// registry rather than the sessions the runner started. Two obligations + /// follow, and a caller that spawns an agent process owes both: + /// + /// - Hold the guard on whatever owns that process, for as long as it lives. + /// Released early — by, say, a request future that merely *awaits* the + /// work — the exit is free to proceed over a child still being stopped. + /// - Claim the slot *before* consulting `app_lifecycle::is_quitting`, never + /// after, so the claim is either in the shutdown's snapshot or made + /// against a flag it has already published. See + /// `pikchr_mcp::reserve_child_session`. pub fn register_external(self: &Arc, session_id: &str) -> ExternalSessionRegistration { ExternalSessionRegistration { token: self.register(session_id), From a7ac34f684207a44a3e56afe777d883ff46d6fed Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 4 Sep 2026 11:17:53 +1000 Subject: [PATCH 13/13] fix(lifecycle): look at the registered token before the specialist spawns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `a7f6de3e` gated the diagram sub-session on the shutdown and described the rest of the chain: shutdown fires the reserved token, `forward_user_cancel` hands it to the worker, and the `is_cancelled` check `generate_pikchr_source` takes before each `driver.run` is what keeps an agent from starting. A review of that commit showed the handoff cannot happen in time for the first attempt, which is the only one a shutdown races. `LocalSet` polls the main future first and ticks spawned tasks only once it returns Pending. Between the `spawn_local` and that check nothing yields; the check reads the worker's token rather than the registered one; and the next yield is inside `AcpDriver::run` → `connect`, which has no await and no token check before `cmd.spawn()`. So the forwarder is first polled with the specialist's agent child already running — the spawn-then-teardown shape `89c821c3` deliberately removed from the main session loop, whose gate reads its registered token directly and so never had the problem. Usually that resolves: the forwarder fires at connect's first await, `run_acp_session`'s select aborts, `graceful_stop` kills the child, and `wait_for_sessions` is still holding the exit open. It bites in exactly the case this gate exists for — that teardown outrunning `SHUTDOWN_BUDGET`, after which `app.exit(0)` orphans the child: own process group, and an exit runs no `kill_on_drop` destructors. The worker now takes a synchronous look at the registered token immediately after spawning the forwarder, through `arm_worker_if_user_cancelled` — the record-then-cancel body lifted out of `forward_user_cancel`, so the pre-check and the forwarder can't tell different stories about the same cancellation. A quit and a Stop are indistinguishable here (`cancel_owned_sessions` fires the registered token exactly as `cancel_session` does), so the reason recorded is the one the forwarder would have recorded; both steps are idempotent — first reason wins, a second `cancel` is a no-op — so which of them arrives first doesn't matter. That covers every cancellation up to that point, including the seconds `AcpDriver::new` spends probing login shells, which is the bulk of the startup a quit can land in: the worker reaches the pre-`driver.run` check with its own token already armed and bails having started nothing. What remains is a cancel arriving in the few statements between the look and the spawn, which is the forwarder-and-teardown race again — now a slice of the startup rather than all of it. The comment at the site says that, instead of claiming the coverage the review disproved. Tested where the ordering is reachable without an `AppHandle`: a token fired before the worker body, then the forwarder spawned and the look taken in the production order, against a driver whose `run` records that it was reached. Without the look that driver runs — the assertion fails, since `LocalSet` polls the main future first — and with it the row still lands cancelled/`Interrupted` carrying the Stop message, which is what the spawn-then-teardown path produced. Verified with `just check-all`. Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/pikchr_mcp.rs | 179 ++++++++++++++++++++++-- 1 file changed, 166 insertions(+), 13 deletions(-) diff --git a/apps/staged/src-tauri/src/pikchr_mcp.rs b/apps/staged/src-tauri/src/pikchr_mcp.rs index 73faf1e4a..99115aa57 100644 --- a/apps/staged/src-tauri/src/pikchr_mcp.rs +++ b/apps/staged/src-tauri/src/pikchr_mcp.rs @@ -1049,19 +1049,41 @@ accepted a render, so the diagram run was cancelled.", timeout_cancel.cancel(); }); // A Stop pressed in the child diagram session fires the token - // registered in the SessionRegistry; forward it to the - // worker's token so the run actually terminates. A quit fires - // the same token — `cancel_owned_sessions` walks every - // registered id — so this is also how a shutdown reaches the - // specialist, and the `is_cancelled` check `generate_pikchr_source` - // takes before each `driver.run` is what keeps it from starting - // an agent the exit would then orphan. Both watcher tasks are - // dropped with this LocalSet. + // registered in the SessionRegistry; forward it to the worker's + // token so the run actually terminates. A quit fires the same + // token — `cancel_owned_sessions` walks every registered id — + // so this is also how a shutdown reaches the specialist. Both + // watcher tasks are dropped with this LocalSet. tokio::task::spawn_local(forward_user_cancel( - user_cancel, + user_cancel.clone(), Arc::clone(&worker_cancel_reason), worker_cancel.clone(), )); + // The forwarder cannot be what covers the *first* attempt, + // which is the one a shutdown races. `LocalSet` polls the main + // future first and ticks spawned tasks only once it returns + // Pending, and nothing between here and the `is_cancelled` + // check `generate_pikchr_source` takes before each `driver.run` + // yields — nor does `AcpDriver::connect`, which spawns the + // agent child before its first await. So the forwarder is first + // polled with the specialist's process already running: the + // spawn-then-teardown shape `89c821c3` removed from the main + // session loop, whose gate reads its registered token directly. + // Usually harmless (the forwarder fires at connect's first + // await, the run aborts, `graceful_stop` kills the child, and + // `wait_for_sessions` holds the exit open for it) — but the + // case this gate exists for is exactly the one where that + // teardown outruns `SHUTDOWN_BUDGET` and the exit orphans the + // child: own process group, and an exit runs no `kill_on_drop` + // destructors. + // + // So take a synchronous look at the registered token too. That + // covers every cancel up to this point, including the seconds + // `AcpDriver::new` spends probing login shells above — the bulk + // of the startup a quit can land in. A cancel arriving in the + // few statements after it is back to the forwarder and that + // teardown race, which is the residual, not the common case. + arm_worker_if_user_cancelled(&user_cancel, &worker_cancel_reason, &worker_cancel); crate::pikchr_subsession::generate_pikchr_source( &driver, worker_store, @@ -1172,18 +1194,46 @@ const USER_STOP_CANCEL_MESSAGE: &str = "The diagram session was stopped before the specialist accepted a render, so the \ generate_pikchr call was cancelled."; +/// Arm the worker's token for a cancellation the SessionRegistry has already +/// delivered, recording the reason first so the cancelled session row and the +/// parent tool error read as a deliberate stop rather than caller abandonment. +/// +/// The worker's synchronous pre-check and [`forward_user_cancel`] both go +/// through here, so a Stop caught before the first `driver.run` and one caught +/// during it tell the same story. Which of them gets here first doesn't matter: +/// the first reason wins and a second `cancel` is a no-op. +/// +/// A quit takes this same path — `cancel_owned_sessions` fires the registered +/// token exactly as `cancel_session` does — so the message is the one a Stop +/// gets. Nothing here can tell them apart, and the async forwarder couldn't +/// either. +fn arm_worker_if_user_cancelled( + user_cancel: &CancellationToken, + reason: &CancelReason, + worker_cancel: &CancellationToken, +) { + if !user_cancel.is_cancelled() { + return; + } + reason.record(USER_STOP_CANCEL_MESSAGE.to_string()); + worker_cancel.cancel(); +} + /// Wait for a user Stop on the child diagram session — `cancel_session` fires /// `user_cancel`, the token registered in the SessionRegistry — and forward it -/// to the worker's own token, recording the reason first so the cancelled -/// session and the parent tool error read as a deliberate stop. +/// to the worker's own token. +/// +/// This covers a cancellation arriving once the run is under way. One that +/// arrived before it is the worker's synchronous pre-check to catch, because +/// this task is not polled until the main future yields, which it first does +/// with the specialist's process already spawned. async fn forward_user_cancel( user_cancel: CancellationToken, reason: Arc, worker_cancel: CancellationToken, ) { user_cancel.cancelled().await; - reason.record(USER_STOP_CANCEL_MESSAGE.to_string()); - worker_cancel.cancel(); + arm_worker_if_user_cancelled(&user_cancel, &reason, &worker_cancel); } /// Build the child diagram session row, unpersisted. @@ -1733,6 +1783,109 @@ arrow from COLL.e to SNOW.w"#; ); } + /// Stands in for a specialist that must never be launched. A real + /// `driver.run` spawns the agent child before its first await, so "was this + /// called" is the closest a test gets to "was a process started". + #[derive(Default)] + struct NeverRunDriver { + ran: std::cell::Cell, + } + + #[async_trait::async_trait(?Send)] + impl crate::agent::AgentDriver for NeverRunDriver { + async fn run( + &self, + _session_id: &str, + _prompt: &str, + _images: &[(String, String)], + _working_dir: &std::path::Path, + _store: &Arc, + _writer: &Arc, + _cancel_token: &CancellationToken, + _agent_session_id: Option<&str>, + _config_options: &[acp_client::AcpSessionConfigOptionSelection], + ) -> Result { + self.ran.set(true); + Ok(acp_client::AgentRunOutcome::Completed) + } + } + + /// A cancellation that lands before the worker starts — a quit's + /// `cancel_owned_sessions` reaching the reservation while `AcpDriver::new` + /// is still probing login shells — must stop the specialist from launching + /// at all, not launch it and race the teardown against `SHUTDOWN_BUDGET`. + /// + /// The forwarder can't be what does that: spawned tasks are ticked only + /// after the main future returns Pending, and the first yield on the way to + /// `driver.run` is inside the driver, past the spawn. Hence the synchronous + /// look, which this drives in the same order the worker does. + #[tokio::test] + async fn a_cancel_landing_before_the_forwarder_runs_never_starts_the_specialist() { + let registry = Arc::new(SessionRegistry::new()); + let store = Arc::new(Store::in_memory().expect("in-memory store")); + let session = new_pikchr_child_session("fake-agent"); + persist_pikchr_child_session(&store, &session).expect("persist child session"); + + let registration = registry.register_external(&session.id); + let user_cancel = registration.token().clone(); + let worker_cancel = CancellationToken::new(); + let reason = Arc::new(CancelReason::new()); + let driver = NeverRunDriver::default(); + let slot = LastRenderSlot::new(); + + // The shutdown fires the reserved token while the worker is still + // getting to the lines below. + assert!(registry.cancel(&session.id)); + + let local = tokio::task::LocalSet::new(); + let result = local + .run_until(async { + tokio::task::spawn_local(forward_user_cancel( + user_cancel.clone(), + Arc::clone(&reason), + worker_cancel.clone(), + )); + arm_worker_if_user_cancelled(&user_cancel, &reason, &worker_cancel); + crate::pikchr_subsession::generate_pikchr_source( + &driver, + Arc::clone(&store), + &session.id, + Some("test grammar body"), + "a friendly box", + None, + &[], + None, + &slot, + &worker_cancel, + &reason, + ) + .await + }) + .await; + + assert!( + !driver.ran.get(), + "the specialist's agent must never be started once the token has fired" + ); + // And the refusal still reads as the stop it was, on both the row and + // the tool error — the same story the forwarder would have told. + assert_eq!(result.err().as_deref(), Some(USER_STOP_CANCEL_MESSAGE)); + + let persisted = store + .get_session(&session.id) + .expect("load session") + .expect("session exists"); + assert_eq!(persisted.status, SessionStatus::Cancelled); + assert_eq!( + persisted.error_message.as_deref(), + Some(USER_STOP_CANCEL_MESSAGE) + ); + assert_eq!( + persisted.completion_reason.as_ref(), + Some(&CompletionReason::Interrupted) + ); + } + #[test] fn announce_pikchr_child_session_writes_hidden_parent_metadata_row() { let store = Store::in_memory().expect("in-memory store");