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..d7eeb7b07 --- /dev/null +++ b/apps/staged/src-tauri/src/app_lifecycle.rs @@ -0,0 +1,1084 @@ +//! 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`, 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`] +//! 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. 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 +//! 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 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}; +use std::time::{Duration, Instant}; + +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}; + +/// 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 +/// 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 { + /// 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`] 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 + /// 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, +} + +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); + } + + /// Clear any pending prompt, returning whether one was pending. + 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; + } +} + +/// [`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 { + /// 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. + 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.session_labels.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) { + if let WindowEvent::CloseRequested { api, .. } = event { + on_close_requested(window, api); + } +} + +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.is_quitting() { + 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. 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, QuitTrigger::WindowClose); +} + +/// 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 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), 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 Some(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()) + else { + log::warn!("No window left to show"); + return; + }; + + 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); +} + +/// 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 +// ============================================================================= + +/// 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, trigger: QuitTrigger) { + let quit_state = app.state::(); + + // 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; + } + + // 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.is_quitting() { + return; + } + + let blockers = collect_quit_blockers(app); + if !should_prompt(&blockers) { + spawn_quit(app); + 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); +} + +/// 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`. +/// +/// 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. +/// +/// Deliberately absent from the web-mode `dispatch` table — a browser client +/// must not be able to terminate the desktop host. +#[tauri::command] +pub fn quit_app(app_handle: AppHandle) { + request_quit(&app_handle, QuitTrigger::Explicit); +} + +/// Run the quit sequence off the main thread so the bounded waits never freeze +/// 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 || { + shutdown_cleanup(&app); + app.exit(0); + }); +} + +/// Snapshot what a quit would interrupt. +fn collect_quit_blockers(app: &AppHandle) -> QuitBlockers { + let session_labels = match app_store(app) { + Some(store) => owned_active_sessions(&store) + .iter() + .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(), + }; + + 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 { + 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 +// ============================================================================= + +/// 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; + }; + + 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); + }); +} + +/// 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 = 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, 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_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() +} + +/// 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. +/// +/// 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, + 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; + use std::sync::Barrier; + + fn active_session(session_type: Option<&str>, status: SessionStatus) -> ActiveSessionInfo { + ActiveSessionInfo { + session_id: "s1".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(), + } + } + + 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 active_sessions_prompt() { + assert!(should_prompt(&blockers(&["review on fix-login"], 0))); + } + + #[test] + fn running_actions_alone_do_not_prompt() { + assert!(!should_prompt(&blockers(&[], 3))); + } + + #[test] + fn nothing_active_does_not_prompt() { + assert!(!should_prompt(&QuitBlockers::default())); + } + + /// 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(); + + 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 `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, + /// `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.is_quitting()); + + state.run_cleanup_once(|| { + assert!( + state.is_quitting(), + "shutdown was under way but the flag still read false" + ); + }); + + assert!(state.is_quitting()); + } + + /// 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] + 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") + ); + + 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] + 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(); + + assert_eq!(sweep_sessions(&store), 2); + + 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_eq!(sweep_sessions(&store), 0); + + let session = store.get_session(&completed.id).unwrap().unwrap(); + assert_eq!(session.status, SessionStatus::Completed); + assert_eq!( + session.completion_reason, + 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/lib.rs b/apps/staged/src-tauri/src/lib.rs index a88c56dfa..e266d45a4 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,17 @@ pub fn run() { log::warn!("Failed to open window from menu: {e}"); } } + MenuDispatch::RequestQuit => { + app_lifecycle::request_quit(app, app_lifecycle::QuitTrigger::Explicit) + } + MenuDispatch::ShowWindow => app_lifecycle::show_a_window(app), MenuDispatch::Drop => {} } }) .on_window_event(|window, event| { + // Close-to-hide / the quit gate (`CloseRequested`). + 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 +2273,9 @@ 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 + // this so a browser client can't quit the host. + app_lifecycle::quit_app, list_projects, create_project, list_project_repos, @@ -2462,17 +2478,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 +2637,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/pikchr_mcp.rs b/apps/staged/src-tauri/src/pikchr_mcp.rs index 48fbe98fc..99115aa57 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() @@ -1006,14 +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. 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, @@ -1124,29 +1194,111 @@ 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); } -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 +1609,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 +1626,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 +1730,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(); @@ -1592,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"); 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/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 cc534ba44..60be553d6 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() @@ -3116,6 +3116,25 @@ 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. + // + // `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); + } + let queued = store .get_queued_sessions_for_branch(&branch_id) .map_err(|e| e.to_string())?; @@ -3133,6 +3152,17 @@ 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. 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; + } + let started = start_queued_session_for_branch( Arc::clone(&store), Arc::clone(®istry), @@ -3354,6 +3384,13 @@ async fn start_queued_session_for_branch( .get_image_ids_for_session(&session_id) .unwrap_or_default(); + // 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", @@ -3403,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 052b911e1..1948675cf 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() @@ -263,6 +276,75 @@ 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`]. + /// + /// `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. +/// +/// 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 +361,39 @@ 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. + /// + /// 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` + /// 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 +403,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 +425,82 @@ 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. + /// + /// 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 + /// 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, + /// 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_before_run`]. + /// + /// 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, + startup: impl FnOnce() -> Result, + ) -> Result<(CancellationToken, T), StartupFailure> { + let token = self.register(session_id); + match startup() { + Ok(value) => Ok((token, value)), + Err(error) => Err(StartupFailure { + error, + accepted_cancellation: self.deregister_reporting_cancellation(session_id), + }), + } + } + /// 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 { @@ -313,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 @@ -337,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); } } @@ -427,12 +622,63 @@ 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 /// `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), @@ -540,6 +786,167 @@ pub struct SessionConfig { pub background_hold: Option, } +/// 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"); + // 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, + 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, +/// 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. +/// +/// 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, + project_id: Option, + store: &Store, + app_handle: &AppHandle, + completion_reason: CompletionReason, +) -> bool { + let transitioned = store + .transition_from_running( + session_id, + SessionStatus::Cancelled, + None, + Some(&completion_reason), + ) + .unwrap_or(false); + log::info!( + "Session {session_id} was cancelled before its run started (transition won: \ + {transitioned})" + ); + emit_status( + app_handle, + session_id, + SessionStatus::Cancelled.as_str(), + None, + Some(&completion_reason), + branch_id, + project_id, + ); + transitioned +} + /// Start a session: persist the user message, spawn the agent, stream to DB. /// /// Returns immediately — the actual agent work happens on a background task. @@ -553,87 +960,151 @@ pub fn start_session( store: Arc, 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) { +) -> 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 + // 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 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 + // 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 (cancel_token, (driver, resolved_provider_id)) = match started { + Ok(started) => started, + Err(failure) => { + if let Some(completion_reason) = failure.accepted_cancellation { + let transitioned = finish_cancelled_before_run( + &config.session_id, + config.branch_id.clone(), + config.project_id.clone(), + &store, + &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); } }; 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(); @@ -883,6 +1354,32 @@ pub fn start_session( }; include_images = false; + // 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 + // — 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 @@ -1135,13 +1632,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` @@ -1228,7 +1723,7 @@ pub fn start_session( } }); - Ok(()) + Ok(SessionStartOutcome::Started) } #[derive(Debug, Clone, PartialEq, Eq)] @@ -1401,7 +1896,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); @@ -1452,7 +1966,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(), @@ -1472,7 +1986,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) @@ -1513,54 +2031,33 @@ 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. - 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(), @@ -1618,7 +2115,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(), @@ -1652,7 +2149,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(), @@ -1665,24 +2162,96 @@ pub fn start_pipeline_session( } }); - Ok(()) + 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 @@ -1915,7 +2484,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, @@ -1942,7 +2521,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}" ), } } @@ -1960,7 +2539,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}" ), } } @@ -2529,8 +3108,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, @@ -3568,8 +4147,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, @@ -4129,31 +4711,175 @@ 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] + 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] @@ -4296,6 +5022,405 @@ 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, ()), StartupFailure> = + registry.register_for_startup("session-failing", || Err("no agent found".to_string())); + + 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), + "shutdown must not wait on a session whose startup failed" + ); + } + + /// 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" + ); + } + + /// 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. + #[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 + ); + } + + /// 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}-{}", 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); diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index d6b0869c0..ccd48d9c0 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -559,6 +559,15 @@ 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..0aa5144ab 100644 --- a/apps/staged/src/App.svelte +++ b/apps/staged/src/App.svelte @@ -586,8 +586,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)); } diff --git a/apps/staged/src/lib/commands.ts b/apps/staged/src/lib/commands.ts index 2f983052e..3943cf657 100644 --- a/apps/staged/src/lib/commands.ts +++ b/apps/staged/src/lib/commands.ts @@ -70,6 +70,23 @@ export function confirmResetStore(): Promise { return invokeCommand('confirm_reset_store'); } +// ============================================================================= +// App lifecycle +// ============================================================================= + +/** + * 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'); +} + // ============================================================================= // Projects // ============================================================================= 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 {