Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: confirm bash cancellation cleanup

Browse Source
This commit is contained in:
Hare
2026-08-27 23:01:06 +09:00
parent 3337cafcdf
commit f5ff0b7c13
2 changed files with 157 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/tools/src/bash.rs
+104 -19
View File
@@ -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
+53 -11
View File
@@ -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