| 1 | //! Low-level tool execution helpers for the engine turn loop. |
| 2 | //! |
| 3 | //! This module keeps the mechanics of MCP dispatch, execution locking, and |
| 4 | //! parallel-tool fanout out of `engine.rs`; the turn loop still owns planning, |
| 5 | //! approval, and how tool results are written back into session state. |
| 6 | |
| 7 | use std::{fs::OpenOptions, io::Write}; |
| 8 | |
| 9 | use super::*; |
| 10 | |
| 11 | /// RAII guard that pauses the TUI's terminal-state ownership for the duration |
| 12 | /// of an interactive tool, then restores it on drop. |
| 13 | /// |
| 14 | /// Background: interactive tools (anything that needs the raw TTY — external |
| 15 | /// editor, `exec_shell` with stdin, etc.) need the TUI to leave alt-screen, |
| 16 | /// disable raw mode, and release mouse capture so the child sees a normal |
| 17 | /// terminal. The TUI listens for `Event::PauseEvents` / `Event::ResumeEvents` |
| 18 | /// and runs `pause_terminal` / `resume_terminal` in response. |
| 19 | /// |
| 20 | /// Earlier code sent `PauseEvents` before tool execution and `ResumeEvents` |
| 21 | /// after. That worked on the happy path, but if the tool's future was dropped |
| 22 | /// — Ctrl+C cancellation, sub-agent abort, parent task cancelled while the |
| 23 | /// tool was awaiting — the second `await` never reached and `ResumeEvents` |
| 24 | /// was never sent. The terminal stayed paused: parent shell scrollbar took |
| 25 | /// over, mouse wheel scrolled the host terminal instead of the transcript, |
| 26 | /// and the TUI rendered as if into a regular cooked-mode buffer. |
| 27 | /// |
| 28 | /// `Drop` runs synchronously and can't await, so we use `try_send` on a |
| 29 | /// **clone of the event channel** to push `ResumeEvents` non-blockingly. The |
| 30 | /// engine event channel is the same one we sent `PauseEvents` on, so by the |
| 31 | /// time we drop there is by construction at least one consumed slot, which |
| 32 | /// keeps `try_send` reliable in practice. |
| 33 | pub(super) struct InteractiveTerminalGuard { |
| 34 | tx: Option<mpsc::Sender<Event>>, |
| 35 | } |
| 36 | |
| 37 | impl InteractiveTerminalGuard { |
| 38 | /// Send `PauseEvents` and arm the guard. If `interactive` is false the |
| 39 | /// guard is a no-op — `Drop` will skip the resume. |
| 40 | pub(super) async fn engage(tx: mpsc::Sender<Event>, interactive: bool) -> Self { |
| 41 | if !interactive { |
| 42 | return Self { tx: None }; |
| 43 | } |
| 44 | // Best-effort: if the receiver is gone the TUI has already shut down |
| 45 | // and there's nothing to restore. Either way we still arm the guard |
| 46 | // so `Drop` symmetrically tries the resume. |
| 47 | let _ = tx.send(Event::PauseEvents).await; |
| 48 | Self { tx: Some(tx) } |
| 49 | } |
| 50 | } |
| 51 | |
| 52 | impl Drop for InteractiveTerminalGuard { |
| 53 | fn drop(&mut self) { |
| 54 | if let Some(tx) = self.tx.take() { |
| 55 | // Synchronous, non-blocking. If the channel is full we still want |
| 56 | // the resume to land — log so a cancellation that loses the |
| 57 | // resume is visible in traces, but don't panic. The TUI also |
| 58 | // re-sends a resume on its own teardown path as a backstop. |
| 59 | if let Err(err) = tx.try_send(Event::ResumeEvents) { |
| 60 | tracing::warn!( |
| 61 | target: "engine.tool_execution", |
| 62 | ?err, |
| 63 | "InteractiveTerminalGuard: try_send(ResumeEvents) failed; \ |
| 64 | terminal may stay in paused state until the next \ |
| 65 | pause/resume cycle" |
| 66 | ); |
| 67 | } |
| 68 | } |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | pub(super) fn emit_tool_audit(event: serde_json::Value) { |
| 73 | let Some(path) = std::env::var_os("DEEPSEEK_TOOL_AUDIT_LOG") else { |
| 74 | return; |
| 75 | }; |
| 76 | let line = match serde_json::to_string(&event) { |
| 77 | Ok(line) => line, |
| 78 | Err(_) => return, |
| 79 | }; |
| 80 | let path = PathBuf::from(path); |
| 81 | if let Some(parent) = path.parent() { |
| 82 | let _ = std::fs::create_dir_all(parent); |
| 83 | } |
| 84 | if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(path) { |
| 85 | let _ = writeln!(file, "{line}"); |
| 86 | } |
| 87 | } |
| 88 | |
| 89 | impl Engine { |
| 90 | pub(super) async fn execute_mcp_tool_with_pool( |
| 91 | pool: Arc<AsyncMutex<McpPool>>, |
| 92 | name: &str, |
| 93 | input: serde_json::Value, |
| 94 | ) -> Result<ToolResult, ToolError> { |
| 95 | let mut pool = pool.lock().await; |
| 96 | let result = pool |
| 97 | .call_tool(name, input) |
| 98 | .await |
| 99 | .map_err(|e| ToolError::execution_failed(format!("MCP tool failed: {e}")))?; |
| 100 | let content = serde_json::to_string_pretty(&result).unwrap_or_else(|_| result.to_string()); |
| 101 | Ok(ToolResult::success(content)) |
| 102 | } |
| 103 | |
| 104 | pub(super) async fn execute_parallel_tool( |
| 105 | &mut self, |
| 106 | input: serde_json::Value, |
| 107 | tool_registry: Option<&crate::tools::ToolRegistry>, |
| 108 | tool_exec_lock: Arc<RwLock<()>>, |
| 109 | ) -> Result<ToolResult, ToolError> { |
| 110 | let calls = parse_parallel_tool_calls(&input)?; |
| 111 | let mcp_pool = if calls.iter().any(|(tool, _)| McpPool::is_mcp_tool(tool)) { |
| 112 | Some(self.ensure_mcp_pool().await?) |
| 113 | } else { |
| 114 | None |
| 115 | }; |
| 116 | let Some(registry) = tool_registry else { |
| 117 | return Err(ToolError::not_available( |
| 118 | "tool registry unavailable for multi_tool_use.parallel", |
| 119 | )); |
| 120 | }; |
| 121 | |
| 122 | let mut tasks = FuturesUnordered::new(); |
| 123 | for (tool_name, tool_input) in calls { |
| 124 | if tool_name == MULTI_TOOL_PARALLEL_NAME { |
| 125 | return Err(ToolError::invalid_input( |
| 126 | "multi_tool_use.parallel cannot call itself", |
| 127 | )); |
| 128 | } |
| 129 | if McpPool::is_mcp_tool(&tool_name) { |
| 130 | if !mcp_tool_is_parallel_safe(&tool_name) { |
| 131 | return Err(ToolError::invalid_input(format!( |
| 132 | "Tool '{tool_name}' is an MCP tool and cannot run in parallel. \ |
| 133 | Allowed MCP tools: list_mcp_resources, list_mcp_resource_templates, \ |
| 134 | mcp_read_resource, read_mcp_resource, mcp_get_prompt." |
| 135 | ))); |
| 136 | } |
| 137 | } else { |
| 138 | let Some(spec) = registry.get(&tool_name) else { |
| 139 | return Err(ToolError::not_available(format!( |
| 140 | "tool '{tool_name}' is not registered" |
| 141 | ))); |
| 142 | }; |
| 143 | if !spec.is_read_only() { |
| 144 | return Err(ToolError::invalid_input(format!( |
| 145 | "Tool '{tool_name}' is not read-only and cannot run in parallel" |
| 146 | ))); |
| 147 | } |
| 148 | if spec.approval_requirement() != ApprovalRequirement::Auto { |
| 149 | return Err(ToolError::invalid_input(format!( |
| 150 | "Tool '{tool_name}' requires approval and cannot run in parallel" |
| 151 | ))); |
| 152 | } |
| 153 | if !spec.supports_parallel() { |
| 154 | return Err(ToolError::invalid_input(format!( |
| 155 | "Tool '{tool_name}' does not support parallel execution" |
| 156 | ))); |
| 157 | } |
| 158 | } |
| 159 | |
| 160 | let registry_ref = registry; |
| 161 | let lock = tool_exec_lock.clone(); |
| 162 | let tx_event = self.tx_event.clone(); |
| 163 | let mcp_pool = mcp_pool.clone(); |
| 164 | tasks.push(async move { |
| 165 | let result = Engine::execute_tool_with_lock( |
| 166 | lock, |
| 167 | true, |
| 168 | false, |
| 169 | tx_event, |
| 170 | tool_name.clone(), |
| 171 | tool_input.clone(), |
| 172 | Some(registry_ref), |
| 173 | mcp_pool, |
| 174 | None, |
| 175 | ) |
| 176 | .await; |
| 177 | (tool_name, result) |
| 178 | }); |
| 179 | } |
| 180 | |
| 181 | let mut results = Vec::new(); |
| 182 | while let Some((tool_name, result)) = tasks.next().await { |
| 183 | match result { |
| 184 | Ok(output) => { |
| 185 | let mut error = None; |
| 186 | if !output.success { |
| 187 | error = Some(output.content.clone()); |
| 188 | } |
| 189 | results.push(ParallelToolResultEntry { |
| 190 | tool_name, |
| 191 | success: output.success, |
| 192 | content: output.content, |
| 193 | error, |
| 194 | }); |
| 195 | } |
| 196 | Err(err) => { |
| 197 | let message = format!("{err}"); |
| 198 | results.push(ParallelToolResultEntry { |
| 199 | tool_name, |
| 200 | success: false, |
| 201 | content: format!("Error: {message}"), |
| 202 | error: Some(message), |
| 203 | }); |
| 204 | } |
| 205 | } |
| 206 | } |
| 207 | |
| 208 | ToolResult::json(&ParallelToolResult { results }) |
| 209 | .map_err(|e| ToolError::execution_failed(e.to_string())) |
| 210 | } |
| 211 | |
| 212 | #[allow(clippy::too_many_arguments)] |
| 213 | pub(super) async fn execute_tool_with_lock( |
| 214 | lock: Arc<RwLock<()>>, |
| 215 | supports_parallel: bool, |
| 216 | interactive: bool, |
| 217 | tx_event: mpsc::Sender<Event>, |
| 218 | tool_name: String, |
| 219 | tool_input: serde_json::Value, |
| 220 | registry: Option<&crate::tools::ToolRegistry>, |
| 221 | mcp_pool: Option<Arc<AsyncMutex<McpPool>>>, |
| 222 | context_override: Option<crate::tools::ToolContext>, |
| 223 | ) -> Result<ToolResult, ToolError> { |
| 224 | let _guard = if supports_parallel { |
| 225 | ToolExecGuard::Read(lock.read().await) |
| 226 | } else { |
| 227 | ToolExecGuard::Write(lock.write().await) |
| 228 | }; |
| 229 | |
| 230 | // RAII pause/resume: ensures `Event::ResumeEvents` always fires on |
| 231 | // drop, even if the tool future is cancelled mid-await. See |
| 232 | // `InteractiveTerminalGuard` doc-comment for the regression this |
| 233 | // closes (parent terminal scrollback hijacking the TUI after a |
| 234 | // cancelled interactive tool). |
| 235 | let _terminal = InteractiveTerminalGuard::engage(tx_event, interactive).await; |
| 236 | |
| 237 | if McpPool::is_mcp_tool(&tool_name) { |
| 238 | if let Some(pool) = mcp_pool { |
| 239 | Engine::execute_mcp_tool_with_pool(pool, &tool_name, tool_input).await |
| 240 | } else { |
| 241 | Err(ToolError::not_available(format!( |
| 242 | "tool '{tool_name}' is not registered" |
| 243 | ))) |
| 244 | } |
| 245 | } else if let Some(registry) = registry { |
| 246 | registry |
| 247 | .execute_full_with_context(&tool_name, tool_input, context_override.as_ref()) |
| 248 | .await |
| 249 | } else { |
| 250 | Err(ToolError::not_available(format!( |
| 251 | "tool '{tool_name}' is not registered" |
| 252 | ))) |
| 253 | } |
| 254 | } |
| 255 | } |
| 256 | |
| 257 | #[cfg(test)] |
| 258 | mod tests { |
| 259 | use super::*; |
| 260 | use serde_json::json; |
| 261 | use std::sync::Mutex; |
| 262 | |
| 263 | /// Tests in this module mutate `DEEPSEEK_TOOL_AUDIT_LOG` which is |
| 264 | /// process-global; serialise through this guard so the parallel |
| 265 | /// runner doesn't observe interleaved env mutations. |
| 266 | static AUDIT_TEST_GUARD: Mutex<()> = Mutex::new(()); |
| 267 | |
| 268 | fn audit_test_guard() -> std::sync::MutexGuard<'static, ()> { |
| 269 | AUDIT_TEST_GUARD.lock().unwrap_or_else(|e| e.into_inner()) |
| 270 | } |
| 271 | |
| 272 | #[test] |
| 273 | fn emit_tool_audit_writes_jsonl_line_when_env_var_set() { |
| 274 | let _g = audit_test_guard(); |
| 275 | let tmp = tempfile::tempdir().expect("tempdir"); |
| 276 | let path = tmp.path().join("audit.log"); |
| 277 | // SAFETY: serialised by the guard above. |
| 278 | unsafe { |
| 279 | std::env::set_var("DEEPSEEK_TOOL_AUDIT_LOG", &path); |
| 280 | } |
| 281 | |
| 282 | emit_tool_audit(json!({ |
| 283 | "event": "tool.spillover", |
| 284 | "tool_id": "call-abc", |
| 285 | "tool_name": "exec_shell", |
| 286 | "path": "/tmp/foo.txt", |
| 287 | })); |
| 288 | emit_tool_audit(json!({ |
| 289 | "event": "tool.result", |
| 290 | "tool_id": "call-xyz", |
| 291 | "success": true, |
| 292 | })); |
| 293 | |
| 294 | let body = std::fs::read_to_string(&path).expect("audit log written"); |
| 295 | let lines: Vec<&str> = body.lines().collect(); |
| 296 | assert_eq!(lines.len(), 2, "two emits → two lines"); |
| 297 | |
| 298 | // Each line round-trips as JSON, has the expected event key. |
| 299 | let first: serde_json::Value = serde_json::from_str(lines[0]).expect("first line is JSON"); |
| 300 | assert_eq!( |
| 301 | first.get("event").and_then(|v| v.as_str()), |
| 302 | Some("tool.spillover") |
| 303 | ); |
| 304 | assert_eq!( |
| 305 | first.get("tool_id").and_then(|v| v.as_str()), |
| 306 | Some("call-abc") |
| 307 | ); |
| 308 | |
| 309 | let second: serde_json::Value = |
| 310 | serde_json::from_str(lines[1]).expect("second line is JSON"); |
| 311 | assert_eq!( |
| 312 | second.get("event").and_then(|v| v.as_str()), |
| 313 | Some("tool.result") |
| 314 | ); |
| 315 | |
| 316 | // SAFETY: cleanup under the guard. |
| 317 | unsafe { |
| 318 | std::env::remove_var("DEEPSEEK_TOOL_AUDIT_LOG"); |
| 319 | } |
| 320 | } |
| 321 | |
| 322 | #[test] |
| 323 | fn emit_tool_audit_is_noop_when_env_var_unset() { |
| 324 | let _g = audit_test_guard(); |
| 325 | // SAFETY: serialised by the guard above. |
| 326 | unsafe { |
| 327 | std::env::remove_var("DEEPSEEK_TOOL_AUDIT_LOG"); |
| 328 | } |
| 329 | // Should not panic and should not create any file. We can't |
| 330 | // assert "no file written" without knowing where one might be |
| 331 | // written, but the contract is "do nothing", which we verify |
| 332 | // by ensuring the call returns without error. |
| 333 | emit_tool_audit(json!({"event": "noop", "x": 1})); |
| 334 | // Successful return is the assertion. |
| 335 | } |
| 336 | |
| 337 | #[test] |
| 338 | fn emit_tool_audit_creates_parent_directory() { |
| 339 | let _g = audit_test_guard(); |
| 340 | let tmp = tempfile::tempdir().expect("tempdir"); |
| 341 | // Path with a parent that doesn't exist yet — the writer |
| 342 | // should create it. |
| 343 | let nested = tmp.path().join("nested").join("dir").join("audit.log"); |
| 344 | // SAFETY: serialised by the guard above. |
| 345 | unsafe { |
| 346 | std::env::set_var("DEEPSEEK_TOOL_AUDIT_LOG", &nested); |
| 347 | } |
| 348 | emit_tool_audit(json!({"event": "test"})); |
| 349 | assert!(nested.exists(), "writer should mkdir -p the parent chain"); |
| 350 | |
| 351 | // SAFETY: cleanup under the guard. |
| 352 | unsafe { |
| 353 | std::env::remove_var("DEEPSEEK_TOOL_AUDIT_LOG"); |
| 354 | } |
| 355 | } |
| 356 | } |
| 357 |