From f5ff0b7c136a71fb72108d317d438d2e53592610 Mon Sep 17 00:00:00 2001 From: Hare Date: Thu, 27 Aug 2026 23:01:06 +0900 Subject: [PATCH] fix: confirm bash cancellation cleanup --- crates/tools/src/bash.rs | 123 +++++++++++++++++++++++++----- crates/tools/tests/integration.rs | 64 +++++++++++++--- 2 files changed, 157 insertions(+), 30 deletions(-) diff --git a/crates/tools/src/bash.rs b/crates/tools/src/bash.rs index 09cceda7..bc573849 100644 --- a/crates/tools/src/bash.rs +++ b/crates/tools/src/bash.rs @@ -24,31 +24,62 @@ pub(crate) struct BashTool { state: Arc>, } +#[derive(Clone)] +struct ActiveCommand { + call_id: String, + execution_nonce: u64, + handle: CommandHandle, +} + #[derive(Default)] struct BashExecutionState { - active: HashMap, + active: HashMap, cancellation_requested: HashSet, + legacy_cancellation_requested: HashSet, + next_execution_nonce: u64, } struct CommandGuard { session: WorkdirSessionHandle, state: Arc>, - call_id: String, + execution_id: String, + execution_nonce: u64, handle: Option, } 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(¶ms.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::>() + }; + 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 diff --git a/crates/tools/tests/integration.rs b/crates/tools/tests/integration.rs index 55291d52..98e2dab6 100644 --- a/crates/tools/tests/integration.rs +++ b/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