fix: confirm bash cancellation cleanup

crates/tools/src/bash.rs

@@ -24,31 +24,62 @@ pub(crate) struct BashTool {
     state: Arc<Mutex<BashExecutionState>>,
 }
 
+#[derive(Clone)]
+struct ActiveCommand {
+    call_id: String,
+    execution_nonce: u64,
+    handle: CommandHandle,
+}
+
 #[derive(Default)]
 struct BashExecutionState {
-    active: HashMap<String, CommandHandle>,
+    active: HashMap<String, ActiveCommand>,
     cancellation_requested: HashSet<String>,
+    legacy_cancellation_requested: HashSet<String>,
+    next_execution_nonce: u64,
 }
 
 struct CommandGuard {
     session: WorkdirSessionHandle,
     state: Arc<Mutex<BashExecutionState>>,
-    call_id: String,
+    execution_id: String,
+    execution_nonce: u64,
     handle: Option<CommandHandle>,
 }
 
 impl Drop for CommandGuard {
     fn drop(&mut self) {
-        let mut state = self.state.lock().unwrap();
-        state.active.remove(&self.call_id);
-        state.cancellation_requested.remove(&self.call_id);
-        drop(state);
-        if let Some(handle) = self.handle.take() {
-            let workdir = self.session.clone();
-            tokio::spawn(async move {
-                let _ = workdir.cancel_command(handle).await;
-            });
-        }
+        let Some(handle) = self.handle.take() else {
+            return;
+        };
+        let workdir = self.session.clone();
+        let state = Arc::clone(&self.state);
+        let execution_id = self.execution_id.clone();
+        let execution_nonce = self.execution_nonce;
+        // A dropped provider future is not terminal confirmation. Keep the live
+        // execution registered until cleanup has both requested cancellation and
+        // observed terminal command output, so cancellation/session teardown
+        // cannot race with an apparently empty registry.
+        tokio::spawn(async move {
+            let _ = workdir.cancel_command(handle.clone()).await;
+            let _ = workdir
+                .command_output(CommandOutputRequest {
+                    handle,
+                    cursor: 0,
+                    limit: INLINE_BYTE_BUDGET,
+                    wait: true,
+                })
+                .await;
+            let mut state = state.lock().unwrap();
+            if state
+                .active
+                .get(&execution_id)
+                .is_some_and(|active| active.execution_nonce == execution_nonce)
+            {
+                state.active.remove(&execution_id);
+                state.cancellation_requested.remove(&execution_id);
+            }
+        });
     }
 }
 
@@ -66,11 +97,18 @@ impl Tool for BashTool {
             .unwrap_or(DEFAULT_TIMEOUT_SECS)
             .clamp(1, MAX_TIMEOUT_SECS);
         let cmd_summary = truncate_for_summary(&params.command);
+        let execution_id = ctx.execution_id();
         let call_id = ctx.call_id;
+        let execution_nonce = {
+            let mut state = self.state.lock().unwrap();
+            state.next_execution_nonce = state.next_execution_nonce.wrapping_add(1);
+            state.next_execution_nonce
+        };
         let mut guard = CommandGuard {
             session: self.session.clone(),
             state: self.state.clone(),
-            call_id: call_id.clone(),
+            execution_id: execution_id.clone(),
+            execution_nonce,
             handle: None,
         };
         let handle = self
@@ -85,8 +123,16 @@ impl Tool for BashTool {
             .map_err(crate::ToolsError::from)?;
         let cancel_after_start = {
             let mut state = self.state.lock().unwrap();
-            state.active.insert(call_id.clone(), handle.clone());
-            state.cancellation_requested.contains(&call_id)
+            state.active.insert(
+                execution_id.clone(),
+                ActiveCommand {
+                    call_id: call_id.clone(),
+                    execution_nonce,
+                    handle: handle.clone(),
+                },
+            );
+            state.cancellation_requested.contains(&execution_id)
+                || state.legacy_cancellation_requested.contains(&call_id)
         };
         guard.handle = Some(handle.clone());
         if cancel_after_start {
@@ -107,8 +153,18 @@ impl Tool for BashTool {
             .map_err(crate::ToolsError::from)?;
         let cancellation_requested = {
             let mut state = self.state.lock().unwrap();
-            state.active.remove(&call_id);
-            state.cancellation_requested.remove(&call_id)
+            let owns_registration = state
+                .active
+                .get(&execution_id)
+                .is_some_and(|active| active.execution_nonce == execution_nonce);
+            let exact = if owns_registration {
+                state.active.remove(&execution_id);
+                state.cancellation_requested.remove(&execution_id)
+            } else {
+                false
+            };
+            let legacy = state.legacy_cancellation_requested.remove(&call_id);
+            exact || legacy
         };
         guard.handle = None;
 
@@ -149,10 +205,39 @@ impl Tool for BashTool {
     }
 
     async fn cancel(&self, call_id: &str) -> Result<(), ToolError> {
+        let handles = {
+            let mut state = self.state.lock().unwrap();
+            state
+                .legacy_cancellation_requested
+                .insert(call_id.to_string());
+            state
+                .active
+                .values()
+                .filter(|active| active.call_id == call_id)
+                .map(|active| active.handle.clone())
+                .collect::<Vec<_>>()
+        };
+        for handle in handles {
+            self.session
+                .cancel_command(handle)
+                .await
+                .map_err(crate::ToolsError::from)?;
+        }
+        Ok(())
+    }
+
+    async fn cancel_execution(
+        &self,
+        ctx: &agen::tool::ToolExecutionContext,
+    ) -> Result<(), ToolError> {
+        let execution_id = ctx.execution_id();
         let handle = {
             let mut state = self.state.lock().unwrap();
-            state.cancellation_requested.insert(call_id.to_string());
-            state.active.get(call_id).cloned()
+            state.cancellation_requested.insert(execution_id.clone());
+            state
+                .active
+                .get(&execution_id)
+                .map(|active| active.handle.clone())
         };
         if let Some(handle) = handle {
             self.session

crates/tools/tests/integration.rs

@@ -7,7 +7,10 @@
 use std::path::Path;
 use std::sync::Arc;
 
-use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta};
+use agen::tool::{
+    Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolExecutionHandle,
+    ToolExecutionTerminal, ToolMeta,
+};
 use manifest::{Permission, Scope, ScopeConfig, ScopeRule};
 use serde_json::json;
 use tempfile::TempDir;
@@ -403,20 +406,23 @@ async fn bash_provider_output_does_not_expose_internal_paths() {
 
 #[tokio::test]
 async fn bash_cancellation_returns_bounded_progress_as_terminal_output() {
-    let (_dir, _spill, reg) = setup();
+    let (dir, _spill, reg) = setup();
+    let marker = dir.path().join("must-not-run-after-cancel");
+    let command = format!(
+        "printf 'before\\n'; printf 'err-before\\n' >&2; sleep 1; touch {}; printf 'after\\n'",
+        marker.display()
+    );
+    let input = serde_json::to_string(&json!({ "command": command })).unwrap();
+    let context = ToolExecutionContext::new("call-heavy", "attempt-heavy", 0);
     let bash = reg.get("Bash");
     let executing = bash.clone();
-    let execution = tokio::spawn(async move {
-        executing
-            .execute(
-                r#"{"command":"printf 'before\\n'; printf 'err-before\\n' >&2; sleep 5; printf 'after\\n'"}"#,
-                Default::default(),
-            )
-            .await
-    });
+    let execution_context = context.clone();
+    let execution = tokio::spawn(async move { executing.execute(&input, execution_context).await });
 
     tokio::time::sleep(std::time::Duration::from_millis(100)).await;
-    bash.cancel("direct").await.expect("signal cancellation");
+    bash.cancel_execution(&context)
+        .await
+        .expect("signal exact execution cancellation");
     let error = tokio::time::timeout(std::time::Duration::from_secs(2), execution)
         .await
         .expect("cancelled Bash should terminate inside the Engine grace budget")
@@ -439,6 +445,42 @@ async fn bash_cancellation_returns_bounded_progress_as_terminal_output() {
         "post-cancel output leaked: {content}"
     );
     assert!(content.len() <= 16 * 1024, "output must remain bounded");
+
+    tokio::time::sleep(std::time::Duration::from_millis(1_100)).await;
+    assert!(
+        !marker.exists(),
+        "the cancelled command continued executing after terminal confirmation"
+    );
+}
+
+#[tokio::test]
+async fn bash_force_close_cleanup_stops_command_and_keeps_session_reusable() {
+    let (dir, _spill, reg) = setup();
+    let marker = dir.path().join("must-not-survive-force-close");
+    let command = format!("sleep 1; touch {}", marker.display());
+    let input = serde_json::to_string(&json!({ "command": command })).unwrap();
+    let bash = reg.get("Bash");
+    let context = ToolExecutionContext::new("call-force", "attempt-force", 0);
+    let (handle, terminal) = ToolExecutionHandle::start(bash.clone(), input, context);
+
+    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+    handle.force_close();
+    assert!(matches!(
+        terminal.await,
+        ToolExecutionTerminal::OutcomeUnknown
+    ));
+
+    tokio::time::sleep(std::time::Duration::from_millis(1_100)).await;
+    assert!(
+        !marker.exists(),
+        "CommandGuard cleanup allowed a force-closed command to continue"
+    );
+
+    let output = bash
+        .execute(r#"{"command":"printf 'reused'"}"#, Default::default())
+        .await
+        .expect("workdir session remains reusable after cleanup");
+    assert_eq!(output.content.as_deref(), Some("reused"));
 }
 
 // Sanity: unused Path import guard