refactor: turn cleanup API, TurnState, impl cascade, tools split

crates/hi-agent/src/agent/delegate_turn.rs

@@ -59,6 +59,8 @@ impl crate::Agent {
                 false,
             );
         }
+        // Budget before runner so exhausted sessions get a clear budget message
+        // even when a runner is attached (and tests that only set the counter).
         if self.subagents.delegate_subagents_used >= MAX_DELEGATE_SUBAGENTS_PER_SESSION {
             return delegate_tool_outcome(
                 format!(
@@ -79,15 +81,16 @@ impl crate::Agent {
                 false,
             );
         };
+        let n = self
+            .subagents
+            .try_begin_delegate(MAX_DELEGATE_SUBAGENTS_PER_SESSION)
+            .expect("budget checked above");
 
         let verify = parsed
             .as_ref()
             .and_then(|v| v.get("verify").and_then(Value::as_str))
             .map(str::to_string)
             .filter(|s| !s.trim().is_empty());
-
-        self.subagents.delegate_subagents_used += 1;
-        let n = self.subagents.delegate_subagents_used;
         let summary: String = task.chars().take(72).collect();
         let ellipsis = if task.chars().count() > 72 { "…" } else { "" };
         ui.subagent_note(&format!(

crates/hi-agent/src/agent/explore_turn.rs

@@ -57,17 +57,18 @@ impl crate::Agent {
                 hi_tools::ToolStatus::Failed,
             );
         }
-        if self.subagents.explore_subagents_used >= MAX_EXPLORE_SUBAGENTS_PER_SESSION {
+        let Some(n) = self
+            .subagents
+            .try_begin_explore(MAX_EXPLORE_SUBAGENTS_PER_SESSION)
+        else {
             return explore_tool_outcome(
                 format!(
                     "explore budget exhausted ({MAX_EXPLORE_SUBAGENTS_PER_SESSION} subagents this \
                      session); investigate directly instead."
                 ),
                 hi_tools::ToolStatus::Denied,
             );
-        }
-        self.subagents.explore_subagents_used += 1;
-        let n = self.subagents.explore_subagents_used;
+        };
         // Prominent callout so the user clearly sees a nested agent run start.
         let summary: String = task.chars().take(72).collect();
         let ellipsis = if task.chars().count() > 72 { "…" } else { "" };
@@ -166,18 +167,16 @@ impl crate::Agent {
                     explore_tool_outcome(answer, status)
                 }
                 Err(err) => {
-                    // Mirror parent failed-turn cleanup: kill child backgrounds and
-                    // type an infrastructure outcome so nested escapes don't leak.
-                    let _ = child.finalize_failed_turn();
+                    // Nested escapes: typed fail cleanup (turn-scoped bg kill).
+                    let _ = child.cleanup_turn(crate::TurnCleanupKind::Fail).await;
                     explore_tool_outcome(
                         format!("explore subagent error: {err}"),
                         hi_tools::ToolStatus::Failed,
                     )
                 }
             }
         };
-        // Always tear down child-owned backgrounds (read-only explore should be
-        // quiet, but bash-in-readonly denial paths still share the kill API).
+        // Throwaway child runtime: full kill (local skeptic + any leftover bg).
         child.kill_background_processes();
         // Fold the child's token usage into the parent's session totals.
         self.add_side_usage(*child.totals());

crates/hi-agent/src/agent/lifecycle.rs

@@ -726,26 +726,84 @@ impl crate::Agent {
         self.stop_local_skeptic_server();
     }
 
-    /// Finalize a turn whose future was cancelled by its frontend. Reconcile
-    /// after rollback/cleanup so reports contain the exact surviving effects
-    /// instead of a fabricated empty list.
+    /// Single entry point for abnormal turn teardown (cancel / infrastructure fail).
     ///
-    /// Mirrors frontend cancel cleanup for turn-scoped background processes when
-    /// the caller has not already killed them (safe if already empty).
+    /// Owns turn-scoped background kill via [`WorkspaceTurnState::active_turn_background_baseline`]
+    /// (taken once — second call is a no-op). Frontends should prefer this over
+    /// ad-hoc kill + finalize sequences.
+    ///
+    /// Normal successful turns clear baselines inside `run_turn` and must not call this.
+    pub async fn cleanup_turn(
+        &mut self,
+        kind: crate::TurnCleanupKind,
+    ) -> Result<crate::TurnCleanupResult> {
+        match kind {
+            crate::TurnCleanupKind::Cancel { session } => {
+                let killed = self.take_and_kill_turn_backgrounds();
+                match session {
+                    crate::SessionRollback::AlreadyApplied => {
+                        // Frontend already rewound transcript/goals; don't truncate again.
+                        let _ = self.workspace.active_turn_message_start.take();
+                    }
+                    crate::SessionRollback::AgentOwned {
+                        checkpoint_count_before,
+                    } => {
+                        if self.checkpoint_count() > checkpoint_count_before
+                            && let Err(err) = self.undo().await
+                        {
+                            eprintln!(
+                                "hi-agent: couldn't roll back cancelled workspace edits: {err:#}"
+                            );
+                        }
+                        if let Some(start) = self.workspace.active_turn_message_start.take() {
+                            self.truncate_messages(start);
+                        }
+                    }
+                }
+                let outcome = self.finalize_cancelled_turn_inner()?;
+                Ok(crate::TurnCleanupResult {
+                    outcome,
+                    killed_backgrounds: killed,
+                })
+            }
+            crate::TurnCleanupKind::Fail => {
+                let killed = self.take_and_kill_turn_backgrounds();
+                let outcome = self.finalize_failed_turn_inner();
+                Ok(crate::TurnCleanupResult {
+                    outcome,
+                    killed_backgrounds: killed,
+                })
+            }
+        }
+    }
+
+    /// Finalize a cancelled turn. Prefer [`Self::cleanup_turn`] so background kill
+    /// and session rollback stay consistent across frontends.
     pub fn finalize_cancelled_turn(&mut self) -> Result<crate::TurnOutcome> {
-        self.cleanup_turn_backgrounds();
+        let _ = self.take_and_kill_turn_backgrounds();
+        self.finalize_cancelled_turn_inner()
+    }
+
+    /// Finalize a failed turn. Prefer [`Self::cleanup_turn`]([`TurnCleanupKind::Fail`]).
+    pub fn finalize_failed_turn(&mut self) -> crate::TurnOutcome {
+        let _ = self.take_and_kill_turn_backgrounds();
+        self.finalize_failed_turn_inner()
+    }
+
+    fn finalize_cancelled_turn_inner(&mut self) -> Result<crate::TurnOutcome> {
+        // Message truncate only if still set (AlreadyApplied path takes it first).
         if let Some(start) = self.workspace.active_turn_message_start.take() {
             self.truncate_messages(start);
         }
         self.runtime.ledger().reconcile()?;
         let baseline = self
-            .workspace.active_turn_ledger_revision
+            .workspace
+            .active_turn_ledger_revision
             .take()
             .unwrap_or_else(|| self.runtime.ledger().revision());
         let changes = self.runtime.ledger().changes_since(baseline);
-        self.workspace.last_changed_files = changes.iter().map(|change| change.path.clone()).collect();
-        self.workspace.last_file_changes = changes;
-        self.report.last_verify = None;
+        self.workspace.record_changes(changes, true);
+        self.report.clear_verify();
         self.workspace.clear_active_baselines();
         let outcome = crate::TurnOutcome {
             status: crate::TurnStatus::Cancelled,
@@ -756,46 +814,38 @@ impl crate::Agent {
             verified_workspace_revision: None,
             effective_route: self.report.last_effective_route.clone(),
         };
-        self.report.last_turn_outcome = Some(outcome.clone());
+        self.report.set_outcome(outcome.clone());
         let _ = self.persist();
         Ok(outcome)
     }
 
-    /// Reconcile and type a turn that escaped through an infrastructure or
-    /// provider error before the normal common finalizer ran. Frontends call
-    /// this before writing reports so late UI/session effects are never
-    /// replaced by a fabricated empty change list.
-    ///
-    /// Also tears down turn-scoped background processes so secondary failure
-    /// paths (model retry fatal, verify re-entry escape, nested explore) match
-    /// the primary cancel cleanup sequence.
-    pub fn finalize_failed_turn(&mut self) -> crate::TurnOutcome {
-        self.cleanup_turn_backgrounds();
+    fn finalize_failed_turn_inner(&mut self) -> crate::TurnOutcome {
         let baseline = self
-            .workspace.active_turn_ledger_revision
+            .workspace
+            .active_turn_ledger_revision
             .take()
             .unwrap_or_else(|| self.runtime.ledger().revision());
         let _ = self.runtime.ledger().reconcile();
         let changes = self.runtime.ledger().changes_since(baseline);
-        self.workspace.last_changed_files = changes.iter().map(|change| change.path.clone()).collect();
-        self.workspace.last_file_changes = changes;
-        self.report.last_verify = None;
+        self.workspace.record_changes(changes, true);
+        self.report.clear_verify();
         self.workspace.clear_active_baselines();
         let route = self.report.last_effective_route.clone();
         let outcome = crate::TurnOutcome::infrastructure_failure(
             route.model,
             route.provider,
             self.workspace.last_changed_files.clone(),
         );
-        self.report.last_turn_outcome = Some(outcome.clone());
+        self.report.set_outcome(outcome.clone());
         outcome
     }
 
-    /// Kill background processes started after this turn's baseline, if any.
-    /// Idempotent: missing baseline is a no-op (caller may have already cleaned up).
-    fn cleanup_turn_backgrounds(&self) {
-        if let Some(before) = self.workspace.active_turn_background_baseline.as_deref() {
-            let _ = self.runtime.background().kill_started_after(before);
+    /// Take the turn background baseline and kill anything started after it.
+    /// Second call is a no-op (baseline already taken).
+    fn take_and_kill_turn_backgrounds(&mut self) -> usize {
+        match self.workspace.active_turn_background_baseline.take() {
+            Some(before) => self.runtime.background().kill_started_after(&before),
+            None => 0,
         }
     }
 
@@ -939,7 +989,7 @@ impl crate::Agent {
         let global = crate::memory::read_global_memory();
         let next = crate::memory::memory_section_for_task(&project, &global, task);
         if next != self.task.memory_context {
-            self.task.memory_context = next;
+            self.task.set_memory_context(next);
         }
     }