From a85084d0d96f1646181ac36279f436b9a5bafc4e Mon Sep 17 00:00:00 2001
From: Miguel Palhas <mpalhas@gmail.com>
Date: Mon, 24 Aug 2026 11:49:36 +0100
Subject: [PATCH] refactor(maestro): extract shared probe and retain helpers
 from the kill path

probe_worktree and record_retention will be reused by the natural-exit
path (#82); kill's behavior is unchanged.
---
 crates/maestro-daemon/src/session.rs | 86 +++++++++++++++++-----------
 1 file changed, 51 insertions(+), 35 deletions(-)

diff --git a/crates/maestro-daemon/src/session.rs b/crates/maestro-daemon/src/session.rs
index 9edd43c..0f5769d 100644
--- a/crates/maestro-daemon/src/session.rs
+++ b/crates/maestro-daemon/src/session.rs
@@ -441,21 +441,8 @@ impl SessionManager {
                 ),
             }
         };
-        // The probe runs before anything is torn down: both refusing (#79)
-        // and recording the retention (#81) must see the tree as it stands,
-        // and a refused kill leaves the session running as the operator left
-        // it.
         let probe = match &worktree {
-            Some((_, path)) => {
-                let path = path.clone();
-                Some(
-                    tmux_blocking(move || crate::worktree_probe::in_the_way(&path))
-                        .await
-                        .map_err(|e| {
-                            SessionError::Worktree(crate::worktree::WorktreeError::from(e))
-                        })?,
-                )
-            }
+            Some(wt) => Some(probe_worktree(wt).await?),
             None => None,
         };
         let guarded = probe.as_ref().is_some_and(|p| !p.is_clear());
@@ -463,29 +450,15 @@ impl SessionManager {
             return Err(SessionError::KillGuarded(probe.expect("guarded ⇒ probed")));
         }
 
-        // Record the retention BEFORE tearing down: a failed store write then
-        // leaves the session running and every artifact untouched, instead of
-        // ending a session whose work outlives it unrecorded.
         let mut retained = None;
         if guarded {
-            let store =
-                self.inner.retained.clone().ok_or_else(|| {
-                    SessionError::RetainFailed("no retained store is wired".into())
-                })?;
-            let (base_repo, path) = worktree.as_ref().expect("guarded ⇒ worktree");
-            let branch = known_branch
-                .or_else(|| crate::worktree_probe::current_branch(path))
-                .unwrap_or_default();
-            let entry = {
-                let store = store.clone();
-                let base_repo = base_repo.clone();
-                let path = path.clone();
-                tmux_blocking(move || {
-                    store.retain(&path, &base_repo, &branch, RetainedReason::ForceKilled)
-                })
-                .await
-                .map_err(|e| SessionError::RetainFailed(e.to_string()))?
-            };
+            let entry = record_retention(
+                &self.inner,
+                worktree.as_ref().expect("guarded ⇒ worktree"),
+                known_branch,
+                RetainedReason::ForceKilled,
+            )
+            .await?;
             retained = Some((entry, probe.expect("guarded ⇒ probed")));
         }
 
@@ -671,6 +644,49 @@ where
         .expect("tmux blocking task panicked")
 }
 
+/// What would be lost if this session's worktree were removed right now.
+///
+/// The single [`crate::worktree_probe::in_the_way`] call site behind both
+/// endings of a session — the kill guard (#79) and the natural-exit
+/// classifier (#82) — so neither can grow its own answer. Runs on the
+/// blocking pool; callers invoke it before anything is torn down, because
+/// refusing (#79), recording a retention (#81) and classifying an exit (#82)
+/// must all see the tree as it stands.
+async fn probe_worktree(
+    (_, path): &(PathBuf, PathBuf),
+) -> Result<maestro_protocol::InTheWay, SessionError> {
+    let path = path.clone();
+    tmux_blocking(move || crate::worktree_probe::in_the_way(&path))
+        .await
+        .map_err(|e| SessionError::Worktree(crate::worktree::WorktreeError::from(e)))
+}
+
+/// Record a guarded worktree into the retained store with `reason`, BEFORE
+/// the caller tears anything down: a failed write then leaves every artifact
+/// untouched instead of destroying work nothing records anymore. Shared by
+/// the forced kill (`ForceKilled`, #81) and the natural exit
+/// (`ExitedWithUnsavedWork`, #82) so the two paths cannot drift apart.
+async fn record_retention(
+    inner: &Arc<Inner>,
+    (base_repo, path): &(PathBuf, PathBuf),
+    known_branch: Option<String>,
+    reason: RetainedReason,
+) -> Result<crate::retained_store::RetainedWorktreeRecord, SessionError> {
+    let store = inner
+        .retained
+        .clone()
+        .ok_or_else(|| SessionError::RetainFailed("no retained store is wired".into()))?;
+    // Sessions adopted from a sidecar carry no branch; ask the tree itself.
+    let branch = known_branch
+        .or_else(|| crate::worktree_probe::current_branch(path))
+        .unwrap_or_default();
+    let base_repo = base_repo.clone();
+    let path = path.clone();
+    tmux_blocking(move || store.retain(&path, &base_repo, &branch, reason))
+        .await
+        .map_err(|e| SessionError::RetainFailed(e.to_string()))
+}
+
 /// One pass of the liveness watcher: diff the in-memory session set against
 /// the tmux-reported one, and reap anything we still track but tmux has lost.
 async fn run_liveness_pass(inner: &Arc<Inner>) {