返回 CodeWhale
tool_cancellation.rs
根目录 / crates / tui / src / core / engine / tests / tool_cancellation.rs
1 use super::*;
2 use crate::llm_client::mock::{MockLlmClient, canned};
3 use crate::tools::spec::{ToolCapability, ToolContext, ToolSpec};
4
5 #[cfg(unix)]
6 #[tokio::test]
7 #[allow(clippy::await_holding_lock)]
8 async fn engine_cancel_stops_started_foreground_descendants_and_preserves_background() {
9 let _env = lock_test_env();
10 let workspace = tempdir().expect("workspace");
11 let pid_file = workspace.path().join("descendant.pid");
12 let command = format!(
13 "printf 'foreground-started\\n'; CODEWHALE_SHELL_DESCENDANT_HELPER=1 \
14 CODEWHALE_SHELL_DESCENDANT_PID_FILE={} {} --exact \
15 tools::shell::tests::shell_descendant_helper_process --nocapture",
16 shell_words::quote(&pid_file.display().to_string()),
17 shell_words::quote(&std::env::current_exe().unwrap().display().to_string()),
18 );
19 let arguments = json!({"command": command, "timeout": 3600}).to_string();
20 let mock = Arc::new(MockLlmClient::new(vec![
21 tool_batch_turn(&[
22 ("call-foreground", "bash", &arguments),
23 (
24 "call-skipped",
25 "bash",
26 r#"{"command":"touch must-not-start"}"#,
27 ),
28 ]),
29 canned::simple_text_turn("Next user turn completed."),
30 ]));
31 let config = Config::default();
32 let (engine, handle) = Engine::new_with_model_client(
33 deterministic_engine_config(workspace.path()),
34 &config,
35 mock.clone(),
36 );
37 let shell_manager = engine.shell_manager.clone();
38 let session_id = engine.session.id.clone();
39 let background_id = shell_manager
40 .lock()
41 .unwrap()
42 .execute_with_options_env_for_session(
43 "sleep 30",
44 None,
45 30_000,
46 true,
47 None,
48 false,
49 None,
50 HashMap::new(),
51 &session_id,
52 )
53 .expect("start intentionally backgrounded control")
54 .task_id
55 .unwrap();
56 let task = tokio::spawn(engine.run());
57 let mut op = external_user_message_op("Run the foreground fixture.", AppMode::Agent, &config);
58 if let Op::SendMessage(TurnSpec {
59 trust_mode,
60 auto_approve,
61 approval_mode,
62 ..
63 }) = &mut op
64 {
65 *trust_mode = true;
66 *auto_approve = true;
67 *approval_mode = ApprovalMode::Bypass;
68 }
69 handle.send(op).await.expect("dispatch real shell tool");
70 let descendant: libc::pid_t = tokio::time::timeout(Duration::from_secs(10), async {
71 loop {
72 if let Ok(raw) = fs::read_to_string(&pid_file)
73 && let Ok(pid) = raw.trim().parse()
74 {
75 break pid;
76 }
77 tokio::time::sleep(Duration::from_millis(10)).await;
78 }
79 })
80 .await
81 .expect("the actual descendant must start before cancellation");
82 handle.cancel();
83
84 let mut receipts = Vec::new();
85 let mut skipped = false;
86 tokio::time::timeout(Duration::from_secs(10), async {
87 let mut events = handle.rx_event.write().await;
88 loop {
89 match events.recv().await.expect("engine remains alive") {
90 Event::ToolCallComplete { id, result, .. } if id == "call-foreground" => {
91 receipts.push(result.expect("model-visible cancellation receipt"));
92 }
93 Event::ToolCallComplete { id, result, .. } if id == "call-skipped" => {
94 let result = result.unwrap();
95 assert_eq!(result.metadata.unwrap()["executed"], false);
96 assert!(result.content.contains("before this tool ran"));
97 skipped = true;
98 }
99 Event::TurnComplete { status, error, .. } => {
100 assert_eq!(status, TurnOutcomeStatus::Interrupted, "{error:?}");
101 break;
102 }
103 _ => {}
104 }
105 }
106 })
107 .await
108 .expect("the cancelled turn must settle");
109 tokio::time::timeout(Duration::from_secs(5), async {
110 loop {
111 // This is the PID written by the isolated helper, not an ambient process.
112 if unsafe { libc::kill(descendant, 0) } == -1
113 && std::io::Error::last_os_error().raw_os_error() == Some(libc::ESRCH)
114 {
115 break;
116 }
117 tokio::time::sleep(Duration::from_millis(10)).await;
118 }
119 })
120 .await
121 .expect("foreground descendant must be gone before recovery");
122 assert_eq!(receipts.len(), 1);
123 assert!(
124 skipped,
125 "the unstarted call must have its own truthful receipt"
126 );
127 assert!(!workspace.path().join("must-not-start").exists());
128 assert!(!receipts[0].success);
129 assert!(receipts[0].content.contains("after shell work started"));
130 assert!(receipts[0].content.contains("Killed"));
131 assert!(!receipts[0].content.contains("before this tool ran"));
132 assert_eq!(receipts[0].metadata.as_ref().unwrap()["executed"], true);
133 {
134 let mut manager = shell_manager.lock().unwrap();
135 let jobs = manager.list_jobs_for_session(&session_id);
136 let foreground = jobs
137 .iter()
138 .find(|job| job.origin_tool_call_id.as_deref() == Some("call-foreground"))
139 .expect("the exact foreground owner remains inspectable");
140 assert_eq!(foreground.status, crate::tools::shell::ShellStatus::Killed);
141 assert!(foreground.stdout_tail.contains("foreground-started"));
142 assert_eq!(
143 manager.inspect_job(&background_id).unwrap().snapshot.status,
144 crate::tools::shell::ShellStatus::Running
145 );
146 assert!(
147 !manager.has_finished_unreported_jobs_for_session(&session_id),
148 "cancelled foreground work must not wake an unsolicited model turn"
149 );
150 }
151 assert_eq!(mock.call_count(), 1);
152 let snapshot = tokio::time::timeout(Duration::from_secs(10), handle.get_session_snapshot())
153 .await
154 .expect("cancelled session must remain inspectable")
155 .unwrap();
156 let manager =
157 crate::session_manager::SessionManager::new(workspace.path().join("checkpoints")).unwrap();
158 let saved = crate::session_manager::create_saved_session_with_id_and_mode(
159 session_id.clone(),
160 &snapshot.messages,
161 "mock-model",
162 workspace.path(),
163 0,
164 None,
165 Some("agent"),
166 );
167 manager.save_checkpoint(&saved).unwrap();
168 let restored = manager
169 .load_session_checkpoint(&session_id)
170 .unwrap()
171 .unwrap();
172 assert!(restored.messages.iter().flat_map(|message| &message.content).any(|block| matches!(
173 block, ContentBlock::ToolResult { tool_use_id, content, is_error: Some(true), .. }
174 if tool_use_id == "call-foreground" && content.contains("after shell work started")
175 )));
176
177 handle
178 .send(external_user_message_op(
179 "Continue after cancellation.",
180 AppMode::Agent,
181 &config,
182 ))
183 .await
184 .unwrap();
185 tokio::time::timeout(Duration::from_secs(10), handle.get_session_snapshot())
186 .await
187 .unwrap()
188 .unwrap();
189 assert_eq!(
190 mock.call_count(),
191 2,
192 "only the next explicit user turn may resume"
193 );
194 assert_eq!(
195 guardian_tool_results(&mock.captured_requests()[1], "call-foreground")[0].1,
196 Some(true)
197 );
198 shell_manager.lock().unwrap().kill(&background_id).unwrap();
199 handle.send(Op::Shutdown).await.unwrap();
200 tokio::time::timeout(Duration::from_secs(10), task)
201 .await
202 .expect("engine must stop after shutdown")
203 .unwrap();
204 }
205
206 struct ReturnedResultTool;
207
208 #[async_trait::async_trait]
209 impl ToolSpec for ReturnedResultTool {
210 fn name(&self) -> &str {
211 "returned_result"
212 }
213 fn description(&self) -> &str {
214 "Return the selected success or failure fixture."
215 }
216 fn input_schema(&self) -> Value {
217 json!({"type":"object","properties":{"fail":{"type":"boolean"}},"required":["fail"]})
218 }
219 fn capabilities(&self) -> Vec<ToolCapability> {
220 vec![ToolCapability::ReadOnly]
221 }
222 async fn execute(&self, input: Value, _context: &ToolContext) -> Result<ToolResult, ToolError> {
223 Ok(if input["fail"] == true {
224 ToolResult::error("fixture failure")
225 } else {
226 ToolResult::success("fixture success")
227 })
228 }
229 }
230
231 #[tokio::test]
232 async fn returned_tool_failure_reaches_next_model_request_as_error() {
233 let workspace = tempdir().unwrap();
234 let mock = Arc::new(MockLlmClient::new(vec![
235 tool_batch_turn(&[
236 ("call-failure", "returned_result", r#"{"fail":true}"#),
237 ("call-success", "returned_result", r#"{"fail":false}"#),
238 ]),
239 canned::simple_text_turn("Both tool results received."),
240 ]));
241 let (mut engine, handle) = Engine::new_with_model_client(
242 deterministic_engine_config(workspace.path()),
243 &Config::default(),
244 mock.clone(),
245 );
246 let mut registry = crate::tools::ToolRegistry::new(ToolContext::new(workspace.path()));
247 registry.register(Arc::new(ReturnedResultTool));
248 let tools = Some(registry.to_api_tools_with_cache(true));
249 let surface = test_tool_surface(&engine, registry, tools, AppMode::Agent);
250 let mut turn = TurnContext::new(4);
251 let (status, error) = engine.run_turn(&mut turn, surface, None, None).await;
252 assert_eq!(status, TurnOutcomeStatus::Completed, "{error:?}");
253 let requests = mock.captured_requests();
254 assert_eq!(requests.len(), 2);
255 assert_eq!(
256 guardian_tool_results(&requests[1], "call-failure"),
257 vec![("fixture failure", Some(true))]
258 );
259 assert_eq!(
260 guardian_tool_results(&requests[1], "call-success"),
261 vec![("fixture success", None)]
262 );
263 let mut events = handle.rx_event.write().await;
264 let results = std::iter::from_fn(|| events.try_recv().ok())
265 .filter_map(|event| match event {
266 Event::ToolCallComplete { id, result, .. } => Some((id, result)),
267 _ => None,
268 })
269 .collect::<Vec<_>>();
270 assert!(
271 results.iter().any(|(id, result)| id == "call-failure"
272 && result.as_ref().is_ok_and(|output| !output.success)),
273 "this must cover Ok(ToolResult::error), not Err(ToolError)"
274 );
275 }
276
276 lines RUST