| 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::{ |
| 8 | fs::OpenOptions, |
| 9 | io::Write, |
| 10 | path::{Path, PathBuf}, |
| 11 | sync::Arc, |
| 12 | time::Duration, |
| 13 | }; |
| 14 | |
| 15 | use super::*; |
| 16 | |
| 17 | const TOOL_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); |
| 18 | |
| 19 | fn inherited_interactive_shell_refusal(tool_name: &str, interactive: bool) -> Option<ToolError> { |
| 20 | if !interactive || !matches!(tool_name, "bash" | "Bash" | "exec_shell") { |
| 21 | return None; |
| 22 | } |
| 23 | crate::tools::shell::inherited_interactive_terminal_refusal() |
| 24 | .map(|message| ToolError::execution_failed(message.to_string())) |
| 25 | } |
| 26 | |
| 27 | #[cfg(all(test, unix))] |
| 28 | thread_local! { |
| 29 | static REPLAY_SPAN_SEQ: std::cell::Cell<Option<u64>> = const { std::cell::Cell::new(None) }; |
| 30 | } |
| 31 | |
| 32 | /// Number operation spans on this thread from 1 until the guard drops, the |
| 33 | /// way an otherwise idle process would (see [`OperationSpanGuard::span_id`]). |
| 34 | #[cfg(all(test, unix))] |
| 35 | pub(crate) fn pin_replay_span_sequence() -> ReplaySpanSequenceGuard { |
| 36 | REPLAY_SPAN_SEQ.with(|cell| cell.set(Some(1))); |
| 37 | ReplaySpanSequenceGuard |
| 38 | } |
| 39 | |
| 40 | #[cfg(all(test, unix))] |
| 41 | pub(crate) struct ReplaySpanSequenceGuard; |
| 42 | |
| 43 | #[cfg(all(test, unix))] |
| 44 | impl Drop for ReplaySpanSequenceGuard { |
| 45 | fn drop(&mut self) { |
| 46 | REPLAY_SPAN_SEQ.with(|cell| cell.set(None)); |
| 47 | } |
| 48 | } |
| 49 | |
| 50 | /// Pairs an observed `OperationActivityStarted` with at most one completion. |
| 51 | /// |
| 52 | /// The turn loop drops an in-flight tool future when the user cancels |
| 53 | /// (`tokio::select!` on the cancel token, or `drop(tool_tasks)` for a parallel |
| 54 | /// batch), so a Completed sent inline after the await would never be sent. |
| 55 | /// Dropping an unfinished span sends `Completed { Cancelled }` with |
| 56 | /// `try_send`: best effort, like the other guards here, because `Drop` cannot |
| 57 | /// await a full channel. Cancellation may release a completion parked on a |
| 58 | /// full queue; the turn's reserved terminal observation settles its lifecycle. |
| 59 | pub(super) struct OperationSpanGuard { |
| 60 | tx: mpsc::Sender<Event>, |
| 61 | span: Option<(String, codewhale_protocol::engine_owner::OwnerActivityKind)>, |
| 62 | cancel: Option<CancellationToken>, |
| 63 | } |
| 64 | |
| 65 | impl OperationSpanGuard { |
| 66 | /// A process-unique span id. The model's tool-call id is not unique: |
| 67 | /// gateways that elide ids fall back to `call_{block_index}`, which |
| 68 | /// repeats every step, and a consumer deduplicates completed spans. |
| 69 | fn span_id(call_id: &str) -> String { |
| 70 | static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); |
| 71 | // A conformance replay numbers its own spans. The process-wide |
| 72 | // counter also counts every span another test starts in a shared |
| 73 | // process, so a golden recorded in isolation drifted there (#6698). |
| 74 | #[cfg(all(test, unix))] |
| 75 | if let Some(seq) = REPLAY_SPAN_SEQ.with(|cell| { |
| 76 | let seq = cell.get()?; |
| 77 | cell.set(Some(seq + 1)); |
| 78 | Some(seq) |
| 79 | }) { |
| 80 | return format!("{call_id}#{seq}"); |
| 81 | } |
| 82 | let seq = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); |
| 83 | format!("{call_id}#{seq}") |
| 84 | } |
| 85 | |
| 86 | pub(super) async fn start( |
| 87 | tx: mpsc::Sender<Event>, |
| 88 | call_id: &str, |
| 89 | activity_kind: codewhale_protocol::engine_owner::OwnerActivityKind, |
| 90 | cancel: Option<CancellationToken>, |
| 91 | ) -> Self { |
| 92 | let span_id = Self::span_id(call_id); |
| 93 | // Reservation is cancel safe: if this future is dropped the event |
| 94 | // was either sent or not, and the guard is armed only once it was. |
| 95 | let sent = match super::streaming::reserve_event_capacity( |
| 96 | &tx, |
| 97 | cancel.as_ref(), |
| 98 | super::streaming::EventReservationPolicy::Strict, |
| 99 | ) |
| 100 | .await |
| 101 | { |
| 102 | Ok(permit) => { |
| 103 | permit.send(Event::OperationActivityStarted { |
| 104 | span_id: span_id.clone(), |
| 105 | activity_kind, |
| 106 | }); |
| 107 | true |
| 108 | } |
| 109 | Err(_) => false, |
| 110 | }; |
| 111 | Self { |
| 112 | tx, |
| 113 | span: sent.then_some((span_id, activity_kind)), |
| 114 | cancel, |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | async fn complete(mut self, outcome: codewhale_protocol::engine_owner::OwnerOperationOutcome) { |
| 119 | if let Some((span_id, activity_kind)) = self.span.take() |
| 120 | && let Ok(permit) = super::streaming::reserve_event_capacity( |
| 121 | &self.tx, |
| 122 | self.cancel.as_ref(), |
| 123 | super::streaming::EventReservationPolicy::Receipt, |
| 124 | ) |
| 125 | .await |
| 126 | { |
| 127 | permit.send(Event::OperationActivityCompleted { |
| 128 | span_id, |
| 129 | activity_kind, |
| 130 | outcome, |
| 131 | }); |
| 132 | } |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | impl Drop for OperationSpanGuard { |
| 137 | fn drop(&mut self) { |
| 138 | if let Some((span_id, activity_kind)) = self.span.take() { |
| 139 | let _ = self.tx.try_send(Event::OperationActivityCompleted { |
| 140 | span_id, |
| 141 | activity_kind, |
| 142 | outcome: codewhale_protocol::engine_owner::OwnerOperationOutcome::Cancelled, |
| 143 | }); |
| 144 | } |
| 145 | } |
| 146 | } |
| 147 | |
| 148 | /// Emits delayed, best-effort liveness pulses for one running tool. |
| 149 | /// |
| 150 | /// Keep the ticker in its own task instead of embedding `tokio::time::Interval` |
| 151 | /// in the already-large engine turn future. Besides keeping the turn future |
| 152 | /// compact, this leaves pre-execution MCP discovery and approval scheduling |
| 153 | /// untouched. Dropping the guard cancels and aborts the ticker synchronously. |
| 154 | struct ToolHeartbeatGuard { |
| 155 | cancel: tokio_util::sync::CancellationToken, |
| 156 | task: tokio::task::JoinHandle<()>, |
| 157 | } |
| 158 | |
| 159 | impl ToolHeartbeatGuard { |
| 160 | fn start(tx_event: mpsc::Sender<Event>, interval: Duration) -> Self { |
| 161 | let cancel = tokio_util::sync::CancellationToken::new(); |
| 162 | let task_cancel = cancel.clone(); |
| 163 | let task = tokio::spawn(async move { |
| 164 | let mut ticker = tokio::time::interval(interval); |
| 165 | ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); |
| 166 | // Tokio intervals tick immediately once. Consume that tick so fast |
| 167 | // tools do not produce a pulse and the first heartbeat is delayed. |
| 168 | ticker.tick().await; |
| 169 | |
| 170 | loop { |
| 171 | tokio::select! { |
| 172 | biased; |
| 173 | |
| 174 | () = task_cancel.cancelled() => break, |
| 175 | _ = ticker.tick() => { |
| 176 | match tx_event.try_send(Event::ToolCallHeartbeat) { |
| 177 | Ok(()) | Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {} |
| 178 | Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => break, |
| 179 | } |
| 180 | } |
| 181 | } |
| 182 | } |
| 183 | }); |
| 184 | Self { cancel, task } |
| 185 | } |
| 186 | } |
| 187 | |
| 188 | impl Drop for ToolHeartbeatGuard { |
| 189 | fn drop(&mut self) { |
| 190 | self.cancel.cancel(); |
| 191 | self.task.abort(); |
| 192 | } |
| 193 | } |
| 194 | |
| 195 | /// RAII guard that pauses the TUI's terminal-state ownership for the duration |
| 196 | /// of an interactive tool, then restores it on drop. |
| 197 | /// |
| 198 | /// Background: interactive tools (anything that needs the raw TTY — external |
| 199 | /// editor, `exec_shell` with stdin, etc.) need the TUI to leave alt-screen, |
| 200 | /// disable raw mode, and release mouse capture so the child sees a normal |
| 201 | /// terminal. The TUI listens for `Event::PauseEvents` / `Event::ResumeEvents` |
| 202 | /// and runs `pause_terminal` / `resume_terminal` in response. |
| 203 | /// |
| 204 | /// Earlier code sent `PauseEvents` before tool execution and `ResumeEvents` |
| 205 | /// after. That worked on the happy path, but if the tool's future was dropped |
| 206 | /// — Ctrl+C cancellation, sub-agent abort, parent task cancelled while the |
| 207 | /// tool was awaiting — the second `await` never reached and `ResumeEvents` |
| 208 | /// was never sent. It also let interactive children start before the UI had |
| 209 | /// actually left alt-screen/raw mode. Both failures strand the TUI in a |
| 210 | /// regular shell scrollback: the parent shell scrollbar takes over, mouse |
| 211 | /// wheel scrolls the host terminal instead of the transcript, and the TUI |
| 212 | /// renders at the bottom of cooked-mode output. |
| 213 | /// |
| 214 | /// Reserve the matching resume before pausing. `Drop` consumes that permit |
| 215 | /// synchronously, including on cancellation and a full channel. No detached |
| 216 | /// sender can outlive the tool or compete with the turn's terminal event. |
| 217 | pub(super) struct InteractiveTerminalGuard { |
| 218 | resume: Option<(mpsc::Sender<Event>, mpsc::OwnedPermit<Event>)>, |
| 219 | } |
| 220 | |
| 221 | impl InteractiveTerminalGuard { |
| 222 | /// Send `PauseEvents` and arm the guard. If `interactive` is false the |
| 223 | /// guard is a no-op — `Drop` will skip the resume. |
| 224 | pub(super) async fn engage( |
| 225 | tx: mpsc::Sender<Event>, |
| 226 | interactive: bool, |
| 227 | cancel: Option<CancellationToken>, |
| 228 | ) -> Result<Self, ToolError> { |
| 229 | if !interactive { |
| 230 | return Ok(Self { resume: None }); |
| 231 | } |
| 232 | let resume = super::streaming::reserve_event_capacity( |
| 233 | &tx, |
| 234 | cancel.as_ref(), |
| 235 | super::streaming::EventReservationPolicy::Strict, |
| 236 | ) |
| 237 | .await |
| 238 | .map_err(terminal_handoff_send_error)?; |
| 239 | let ack = Arc::new(tokio::sync::Notify::new()); |
| 240 | let pause = super::streaming::reserve_event_capacity( |
| 241 | &tx, |
| 242 | cancel.as_ref(), |
| 243 | super::streaming::EventReservationPolicy::Strict, |
| 244 | ) |
| 245 | .await |
| 246 | .map_err(terminal_handoff_send_error)?; |
| 247 | // No await separates the pause send and guard installation. A |
| 248 | // cancelled reservation above never paused and needs no resume. |
| 249 | pause.send(Event::PauseEvents { |
| 250 | ack: Some(ack.clone()), |
| 251 | }); |
| 252 | let guard = Self { |
| 253 | resume: Some((tx, resume)), |
| 254 | }; |
| 255 | let acknowledged = match cancel.as_ref() { |
| 256 | Some(cancel) => tokio::select! { |
| 257 | biased; |
| 258 | () = cancel.cancelled() => { |
| 259 | return Err(ToolError::cancelled("Terminal handoff cancelled; interactive tool was not launched.")); |
| 260 | } |
| 261 | result = tokio::time::timeout(Duration::from_millis(750), ack.notified()) => result, |
| 262 | }, |
| 263 | None => tokio::time::timeout(Duration::from_millis(750), ack.notified()).await, |
| 264 | }; |
| 265 | if acknowledged.is_err() { |
| 266 | return Err(ToolError::execution_failed( |
| 267 | "Terminal handoff was not acknowledged; interactive tool was not launched.", |
| 268 | )); |
| 269 | } |
| 270 | Ok(guard) |
| 271 | } |
| 272 | } |
| 273 | |
| 274 | fn terminal_handoff_send_error(reason: super::streaming::EventSendError) -> ToolError { |
| 275 | match reason { |
| 276 | super::streaming::EventSendError::Cancelled => { |
| 277 | ToolError::cancelled("Terminal handoff cancelled; interactive tool was not launched.") |
| 278 | } |
| 279 | super::streaming::EventSendError::Closed => ToolError::execution_failed( |
| 280 | "Terminal handoff channel closed; interactive tool was not launched.", |
| 281 | ), |
| 282 | } |
| 283 | } |
| 284 | |
| 285 | impl Drop for InteractiveTerminalGuard { |
| 286 | fn drop(&mut self) { |
| 287 | if let Some((tx, resume)) = self.resume.take() |
| 288 | && !tx.is_closed() |
| 289 | { |
| 290 | resume.send(Event::ResumeEvents); |
| 291 | } |
| 292 | } |
| 293 | } |
| 294 | |
| 295 | pub(crate) fn emit_tool_audit(event: serde_json::Value) { |
| 296 | let Some(path) = std::env::var_os("CODEWHALE_TOOL_AUDIT_LOG") |
| 297 | .or_else(|| std::env::var_os("DEEPSEEK_TOOL_AUDIT_LOG")) |
| 298 | else { |
| 299 | return; |
| 300 | }; |
| 301 | emit_tool_audit_to_path(&PathBuf::from(path), event); |
| 302 | } |
| 303 | |
| 304 | fn emit_tool_audit_to_path(path: &Path, event: serde_json::Value) { |
| 305 | let line = match serde_json::to_string(&event) { |
| 306 | Ok(line) => line, |
| 307 | Err(e) => { |
| 308 | tracing::error!("Failed to serialize tool audit event: {e}"); |
| 309 | return; |
| 310 | } |
| 311 | }; |
| 312 | if let Some(parent) = path.parent() |
| 313 | && let Err(e) = std::fs::create_dir_all(parent) |
| 314 | { |
| 315 | tracing::error!( |
| 316 | "Failed to create audit log directory {}: {e}", |
| 317 | parent.display() |
| 318 | ); |
| 319 | return; |
| 320 | } |
| 321 | match OpenOptions::new().create(true).append(true).open(path) { |
| 322 | Ok(mut file) => { |
| 323 | if let Err(e) = writeln!(file, "{line}") { |
| 324 | tracing::error!("Failed to write to audit log {}: {e}", path.display()); |
| 325 | } |
| 326 | } |
| 327 | Err(e) => { |
| 328 | tracing::error!("Failed to open audit log {}: {e}", path.display()); |
| 329 | } |
| 330 | } |
| 331 | } |
| 332 | |
| 333 | impl Engine { |
| 334 | pub(crate) async fn execute_mcp_tool_with_pool( |
| 335 | pool: Arc<AsyncMutex<McpPool>>, |
| 336 | tx_event: &mpsc::Sender<Event>, |
| 337 | name: &str, |
| 338 | input: serde_json::Value, |
| 339 | disallowed_tools: &[String], |
| 340 | decision: Option<&super::HumanDecision>, |
| 341 | ) -> Result<RichToolResult, ToolError> { |
| 342 | McpPool::authorize_call(disallowed_tools, name, &input) |
| 343 | .map_err(|error| ToolError::not_available(error.to_string()))?; |
| 344 | // A synthetic `mcp_<server>_authenticate` call runs the shared OAuth |
| 345 | // login flow with the pool lock released during the browser wait, so |
| 346 | // parallel MCP tools and the `/mcp` manager keep working while the |
| 347 | // user signs in. On success it changes the callable MCP surface (the |
| 348 | // server's real tools replace the synthetic one); flag that so the |
| 349 | // turn loop merges the refreshed catalog before the next model |
| 350 | // request instead of leaving the model with names it cannot legally |
| 351 | // call yet. |
| 352 | let auth_target = pool.lock().await.authenticate_tool_target(name); |
| 353 | if let Some(server) = auth_target { |
| 354 | let mut notice_task = None; |
| 355 | let mut result = crate::mcp::authenticate_tool_via_pool(&pool, &server, |url| { |
| 356 | // The model cannot relay the URL until the call returns, and |
| 357 | // the call returns only after the sign-in completes — so the |
| 358 | // user must see it now. This status is the only copy of the |
| 359 | // URL: `try_send` drops it on a full channel, so the send |
| 360 | // waits for room on its own task instead of failing the |
| 361 | // login. |
| 362 | let server = server.clone(); |
| 363 | let url = url.to_string(); |
| 364 | let tx = tx_event.clone(); |
| 365 | notice_task = Some(super::turn_heartbeat::AbortOnDrop(tokio::spawn(async move { |
| 366 | if let Ok(permit) = super::streaming::reserve_event_capacity( |
| 367 | &tx, |
| 368 | None, |
| 369 | super::streaming::EventReservationPolicy::Receipt, |
| 370 | ) |
| 371 | .await |
| 372 | { |
| 373 | permit.send(Event::status(format!( |
| 374 | "◆ auth required: sign in to MCP server '{server}' in your browser — {url}" |
| 375 | ))); |
| 376 | } |
| 377 | }))); |
| 378 | }) |
| 379 | .await |
| 380 | .map_err(|e| ToolError::execution_failed(format!("MCP tool failed: {e}")))?; |
| 381 | if let Some(task) = notice_task.as_mut() { |
| 382 | let _ = (&mut task.0).await; |
| 383 | } |
| 384 | McpPool::filter_authenticate_result(&mut result, disallowed_tools); |
| 385 | let mut rich = crate::tools::registry::mcp_result_to_bounded_rich_tool_result(result); |
| 386 | if rich.result.success { |
| 387 | rich.result.metadata = Some(serde_json::json!({ "mcp_catalog_changed": true })); |
| 388 | } |
| 389 | return Ok(rich); |
| 390 | } |
| 391 | let needs_auth_generation_before = pool.lock().await.needs_auth_generation(); |
| 392 | let result = pool |
| 393 | .lock() |
| 394 | .await |
| 395 | .call_tool_with_disallowed(name, input, disallowed_tools, decision) |
| 396 | .await; |
| 397 | match result { |
| 398 | Ok(result) => { |
| 399 | Ok(crate::tools::registry::mcp_result_to_bounded_rich_tool_result(result)) |
| 400 | } |
| 401 | Err(error) => { |
| 402 | // A credential the server stopped accepting mid-session |
| 403 | // flips it into the typed needs-auth state and drops its |
| 404 | // connection — the callable surface changed. For THAT |
| 405 | // transition only, return the failure as a result (not an |
| 406 | // Err) carrying `mcp_catalog_changed`, so the turn loop |
| 407 | // replaces the pool's catalog slice: dead tools leave, the |
| 408 | // synthetic login tool arrives, and the error's own hint |
| 409 | // stays model-readable. Every other failure keeps the Err |
| 410 | // contract. |
| 411 | let auth_surface_changed = { |
| 412 | let pool = pool.lock().await; |
| 413 | pool.needs_auth_generation() != needs_auth_generation_before |
| 414 | }; |
| 415 | if !auth_surface_changed { |
| 416 | return Err(ToolError::execution_failed(format!( |
| 417 | "MCP tool failed: {error}" |
| 418 | ))); |
| 419 | } |
| 420 | let tool_result = |
| 421 | crate::tools::spec::ToolResult::error(format!("MCP tool failed: {error}")) |
| 422 | .with_metadata(serde_json::json!({ "mcp_catalog_changed": true })); |
| 423 | Ok(RichToolResult::plain(tool_result)) |
| 424 | } |
| 425 | } |
| 426 | } |
| 427 | |
| 428 | #[allow(clippy::too_many_arguments)] |
| 429 | pub(super) async fn execute_tool_with_lock( |
| 430 | lock: Arc<RwLock<()>>, |
| 431 | supports_parallel: bool, |
| 432 | interactive: bool, |
| 433 | tx_event: mpsc::Sender<Event>, |
| 434 | cancel_token: Option<CancellationToken>, |
| 435 | tool_name: String, |
| 436 | activity_call_id: Option<String>, |
| 437 | tool_input: serde_json::Value, |
| 438 | workspace: PathBuf, |
| 439 | registry: Option<&crate::tools::ToolRegistry>, |
| 440 | mcp_pool: Option<Arc<AsyncMutex<McpPool>>>, |
| 441 | context_override: Option<crate::tools::ToolContext>, |
| 442 | ) -> Result<RichToolResult, ToolError> { |
| 443 | let child = context_override |
| 444 | .as_ref() |
| 445 | .or_else(|| registry.map(|registry| registry.context())) |
| 446 | .and_then(|context| context.child_host.clone()); |
| 447 | let execution = Self::execute_tool_admitted( |
| 448 | lock, |
| 449 | supports_parallel, |
| 450 | interactive, |
| 451 | tx_event, |
| 452 | cancel_token, |
| 453 | tool_name, |
| 454 | activity_call_id, |
| 455 | tool_input, |
| 456 | workspace, |
| 457 | registry, |
| 458 | mcp_pool, |
| 459 | context_override, |
| 460 | ); |
| 461 | match child { |
| 462 | Some(child) => child.run_tool_bounded(execution).await, |
| 463 | None => execution.await, |
| 464 | } |
| 465 | } |
| 466 | |
| 467 | #[allow(clippy::too_many_arguments)] |
| 468 | async fn execute_tool_admitted( |
| 469 | lock: Arc<RwLock<()>>, |
| 470 | supports_parallel: bool, |
| 471 | interactive: bool, |
| 472 | tx_event: mpsc::Sender<Event>, |
| 473 | cancel_token: Option<CancellationToken>, |
| 474 | tool_name: String, |
| 475 | activity_call_id: Option<String>, |
| 476 | tool_input: serde_json::Value, |
| 477 | workspace: PathBuf, |
| 478 | registry: Option<&crate::tools::ToolRegistry>, |
| 479 | mcp_pool: Option<Arc<AsyncMutex<McpPool>>>, |
| 480 | context_override: Option<crate::tools::ToolContext>, |
| 481 | ) -> Result<RichToolResult, ToolError> { |
| 482 | if cancel_token |
| 483 | .as_ref() |
| 484 | .is_some_and(CancellationToken::is_cancelled) |
| 485 | { |
| 486 | return Err(ToolError::permission_denied( |
| 487 | "Turn stopped by user. Tool call blocked.", |
| 488 | )); |
| 489 | } |
| 490 | // Unix inherited-terminal shell calls are impossible without a full |
| 491 | // POSIX job-control lease. Refuse them before the terminal guard so a |
| 492 | // known-invalid call cannot flash host scrollback or drain input. |
| 493 | if let Some(error) = inherited_interactive_shell_refusal(&tool_name, interactive) { |
| 494 | return Err(error); |
| 495 | } |
| 496 | // This guard starts before lock acquisition, so contention as well as |
| 497 | // registry/MCP/interpreter execution remains visibly live. |
| 498 | let _heartbeat = ToolHeartbeatGuard::start(tx_event.clone(), TOOL_HEARTBEAT_INTERVAL); |
| 499 | let started_at = std::time::Instant::now(); |
| 500 | let dispatch = if McpPool::is_mcp_tool(&tool_name) { |
| 501 | "mcp" |
| 502 | } else if matches!( |
| 503 | tool_name.as_str(), |
| 504 | CODE_EXECUTION_TOOL_NAME | JS_EXECUTION_TOOL_NAME | EXECUTE_TOOLS_TOOL_NAME |
| 505 | ) { |
| 506 | "interpreter" |
| 507 | } else if registry.is_some() { |
| 508 | "registry" |
| 509 | } else { |
| 510 | "missing" |
| 511 | }; |
| 512 | let input_bytes = serde_json::to_string(&tool_input) |
| 513 | .map(|s| s.len()) |
| 514 | .unwrap_or(0); |
| 515 | tracing::debug!( |
| 516 | target: "engine.tool_execution", |
| 517 | tool = %tool_name, |
| 518 | dispatch, |
| 519 | interactive, |
| 520 | supports_parallel, |
| 521 | input_bytes, |
| 522 | "tool.exec.start", |
| 523 | ); |
| 524 | |
| 525 | let acquire = async { |
| 526 | if supports_parallel { |
| 527 | ToolExecGuard::Read(lock.read().await) |
| 528 | } else { |
| 529 | ToolExecGuard::Write(lock.write().await) |
| 530 | } |
| 531 | }; |
| 532 | let _guard = match cancel_token.as_ref() { |
| 533 | Some(cancel) => tokio::select! { |
| 534 | biased; |
| 535 | () = cancel.cancelled() => return Err(ToolError::cancelled("Tool lock wait cancelled.")), |
| 536 | guard = acquire => guard, |
| 537 | }, |
| 538 | None => acquire.await, |
| 539 | }; |
| 540 | |
| 541 | // RAII pause/resume: ensures `Event::ResumeEvents` always fires on |
| 542 | // drop, even if the tool future is cancelled mid-await. See |
| 543 | // `InteractiveTerminalGuard` doc-comment for the regression this |
| 544 | // closes (parent terminal scrollback hijacking the TUI after a |
| 545 | // cancelled interactive tool). |
| 546 | let _terminal = |
| 547 | InteractiveTerminalGuard::engage(tx_event.clone(), interactive, cancel_token.clone()) |
| 548 | .await?; |
| 549 | |
| 550 | if cancel_token |
| 551 | .as_ref() |
| 552 | .is_some_and(CancellationToken::is_cancelled) |
| 553 | { |
| 554 | return Err(ToolError::permission_denied( |
| 555 | "Turn stopped by user. Tool call blocked.", |
| 556 | )); |
| 557 | } |
| 558 | |
| 559 | if let Some(context) = context_override |
| 560 | .as_ref() |
| 561 | .or_else(|| registry.map(|registry| registry.context())) |
| 562 | { |
| 563 | super::tool_catalog::enforce_tool_denial(context, &tool_name, &tool_input)?; |
| 564 | } |
| 565 | |
| 566 | let tool_authority = context_override |
| 567 | .as_ref() |
| 568 | .and_then(|context| context.tool_authority.as_ref()) |
| 569 | .or_else(|| registry.and_then(|registry| registry.context().tool_authority.as_ref())); |
| 570 | if let Some(authority) = tool_authority { |
| 571 | if McpPool::is_mcp_tool(&tool_name) |
| 572 | && !super::dispatch::mcp_tool_is_read_only(&tool_name) |
| 573 | { |
| 574 | return Err(ToolError::permission_denied(format!( |
| 575 | "worker '{}' cannot run mutating MCP tool {tool_name}: it has no authorized file target", |
| 576 | authority.owner |
| 577 | ))); |
| 578 | } |
| 579 | if matches!( |
| 580 | tool_name.as_str(), |
| 581 | CODE_EXECUTION_TOOL_NAME | JS_EXECUTION_TOOL_NAME | EXECUTE_TOOLS_TOOL_NAME |
| 582 | ) { |
| 583 | return Err(ToolError::permission_denied(format!( |
| 584 | "worker '{}' cannot run {tool_name}: arbitrary code execution is outside its machine-readable authority envelope", |
| 585 | authority.owner |
| 586 | ))); |
| 587 | } |
| 588 | } |
| 589 | |
| 590 | // Typed owner activity: classified only after every gate above has |
| 591 | // passed, from the same authority that dispatches the call (the MCP |
| 592 | // pool's resolved server map, the interpreter, or the registry plus |
| 593 | // the canonical action alias). Names and arguments never leave here. |
| 594 | let activity_kind = match activity_call_id.as_ref() { |
| 595 | None => None, |
| 596 | Some(_) if McpPool::is_mcp_tool(&tool_name) => match mcp_pool.as_ref() { |
| 597 | Some(pool) => pool |
| 598 | .lock() |
| 599 | .await |
| 600 | .resolved_tool_servers() |
| 601 | .get(&tool_name) |
| 602 | .map(|server| crate::tools::activity::mcp_activity_kind(server)), |
| 603 | None => None, |
| 604 | }, |
| 605 | Some(_) |
| 606 | if matches!( |
| 607 | tool_name.as_str(), |
| 608 | CODE_EXECUTION_TOOL_NAME | JS_EXECUTION_TOOL_NAME |
| 609 | ) => |
| 610 | { |
| 611 | Some(codewhale_protocol::engine_owner::OwnerActivityKind::Executing) |
| 612 | } |
| 613 | // The code-mode wrapper is not itself an operation. |
| 614 | Some(_) if tool_name == EXECUTE_TOOLS_TOOL_NAME => None, |
| 615 | Some(_) => registry |
| 616 | .filter(|registry| registry.get(&tool_name).is_some()) |
| 617 | .and_then(|_| { |
| 618 | crate::tools::activity::registry_activity_kind(&tool_name, &tool_input) |
| 619 | }), |
| 620 | }; |
| 621 | let operation_span = match (activity_call_id.as_deref(), activity_kind) { |
| 622 | (Some(call_id), Some(activity_kind)) => Some( |
| 623 | OperationSpanGuard::start( |
| 624 | tx_event.clone(), |
| 625 | call_id, |
| 626 | activity_kind, |
| 627 | cancel_token.clone(), |
| 628 | ) |
| 629 | .await, |
| 630 | ), |
| 631 | _ => None, |
| 632 | }; |
| 633 | if cancel_token |
| 634 | .as_ref() |
| 635 | .is_some_and(CancellationToken::is_cancelled) |
| 636 | { |
| 637 | return Err(ToolError::cancelled("Tool activity admission cancelled.")); |
| 638 | } |
| 639 | |
| 640 | let child_mcp_call = if McpPool::is_mcp_tool(&tool_name) { |
| 641 | let context = context_override |
| 642 | .as_ref() |
| 643 | .or_else(|| registry.map(|registry| registry.context())); |
| 644 | if context.is_some_and(|context| context.child_host.is_some()) { |
| 645 | let registry = registry.ok_or_else(|| { |
| 646 | ToolError::permission_denied( |
| 647 | "child MCP call has no canonical registered capability", |
| 648 | ) |
| 649 | })?; |
| 650 | registry |
| 651 | .admit_child_call( |
| 652 | &tool_name, |
| 653 | &tool_input, |
| 654 | context.expect("captured child context"), |
| 655 | ) |
| 656 | .await? |
| 657 | } else { |
| 658 | None |
| 659 | } |
| 660 | } else { |
| 661 | None |
| 662 | }; |
| 663 | if let Some(context) = context_override |
| 664 | .as_ref() |
| 665 | .or_else(|| registry.map(|registry| registry.context())) |
| 666 | && (context.acp_host.is_some() || context.child_host.is_some()) |
| 667 | { |
| 668 | let spec = registry |
| 669 | .and_then(|registry| registry.get(&tool_name)) |
| 670 | .ok_or_else(|| { |
| 671 | ToolError::not_available("ACP call has no admitted registered tool") |
| 672 | })?; |
| 673 | crate::tools::registry::enforce_tool_authority( |
| 674 | &tool_name, |
| 675 | &tool_input, |
| 676 | spec.as_ref(), |
| 677 | context, |
| 678 | )?; |
| 679 | crate::extension_host::validate_caller_plugins(context.plugin_registry.as_deref()) |
| 680 | .map_err(ToolError::not_available)?; |
| 681 | if let Some(id) = activity_call_id.as_ref() { |
| 682 | if let Ok(permit) = super::streaming::reserve_event_capacity( |
| 683 | &tx_event, |
| 684 | cancel_token.as_ref(), |
| 685 | super::streaming::EventReservationPolicy::Receipt, |
| 686 | ) |
| 687 | .await |
| 688 | { |
| 689 | // Event backpressure is an await boundary. Recheck the |
| 690 | // exact captured authority before claiming dispatch. |
| 691 | if cancel_token |
| 692 | .as_ref() |
| 693 | .is_some_and(CancellationToken::is_cancelled) |
| 694 | { |
| 695 | return Err(ToolError::cancelled("ACP dispatch admission cancelled")); |
| 696 | } |
| 697 | crate::tools::registry::enforce_tool_authority( |
| 698 | &tool_name, |
| 699 | &tool_input, |
| 700 | spec.as_ref(), |
| 701 | context, |
| 702 | )?; |
| 703 | crate::extension_host::validate_caller_plugins( |
| 704 | context.plugin_registry.as_deref(), |
| 705 | ) |
| 706 | .map_err(ToolError::not_available)?; |
| 707 | permit.send(Event::ToolExecutionStarted { id: id.clone() }); |
| 708 | } else { |
| 709 | return Err(ToolError::cancelled("ACP dispatch observation unavailable")); |
| 710 | } |
| 711 | } |
| 712 | } |
| 713 | let outcome: Result<RichToolResult, ToolError> = if McpPool::is_mcp_tool(&tool_name) { |
| 714 | if let Some(pool) = mcp_pool { |
| 715 | let disallowed_tools = context_override |
| 716 | .as_ref() |
| 717 | .or_else(|| registry.map(|registry| registry.context())) |
| 718 | .map(|context| context.disallowed_tools.as_slice()) |
| 719 | .unwrap_or_default(); |
| 720 | // Only a per-call override can carry a person's decision; the |
| 721 | // registry's shared context never does. |
| 722 | let decision = context_override |
| 723 | .as_ref() |
| 724 | .and_then(|context| context.human_decision.as_ref()); |
| 725 | Engine::execute_mcp_tool_with_pool( |
| 726 | pool, |
| 727 | &tx_event, |
| 728 | &tool_name, |
| 729 | tool_input, |
| 730 | disallowed_tools, |
| 731 | decision, |
| 732 | ) |
| 733 | .await |
| 734 | } else { |
| 735 | Err(ToolError::not_available(format!( |
| 736 | "tool '{tool_name}' is not registered" |
| 737 | ))) |
| 738 | } |
| 739 | } else if matches!( |
| 740 | tool_name.as_str(), |
| 741 | CODE_EXECUTION_TOOL_NAME | JS_EXECUTION_TOOL_NAME |
| 742 | ) { |
| 743 | if let Some(context) = context_override |
| 744 | .as_ref() |
| 745 | .or_else(|| registry.map(|registry| registry.context())) |
| 746 | { |
| 747 | let result = if tool_name == CODE_EXECUTION_TOOL_NAME { |
| 748 | execute_code_execution_tool(&tool_input, &workspace, context).await |
| 749 | } else { |
| 750 | execute_js_execution_tool(&tool_input, &workspace, context).await |
| 751 | }; |
| 752 | result.map(RichToolResult::plain) |
| 753 | } else { |
| 754 | Err(ToolError::not_available( |
| 755 | "local code execution requires an effective tool context", |
| 756 | )) |
| 757 | } |
| 758 | } else if tool_name == EXECUTE_TOOLS_TOOL_NAME { |
| 759 | if let Some(registry) = registry { |
| 760 | let context = context_override |
| 761 | .as_ref() |
| 762 | .cloned() |
| 763 | .unwrap_or_else(|| registry.context().clone()); |
| 764 | crate::tools::codemode::execute_tools_tool(&tool_input, registry, &context) |
| 765 | .await |
| 766 | .map(RichToolResult::plain) |
| 767 | } else { |
| 768 | Err(ToolError::not_available(format!( |
| 769 | "tool '{tool_name}' is not registered" |
| 770 | ))) |
| 771 | } |
| 772 | } else if let Some(registry) = registry { |
| 773 | registry |
| 774 | .execute_rich_full_with_context(&tool_name, tool_input, context_override.as_ref()) |
| 775 | .await |
| 776 | } else { |
| 777 | Err(ToolError::not_available(format!( |
| 778 | "tool '{tool_name}' is not registered" |
| 779 | ))) |
| 780 | }; |
| 781 | |
| 782 | if let Some(operation_span) = operation_span { |
| 783 | let cancelled = cancel_token |
| 784 | .as_ref() |
| 785 | .is_some_and(CancellationToken::is_cancelled); |
| 786 | operation_span |
| 787 | .complete(crate::tools::activity::operation_outcome( |
| 788 | &outcome, cancelled, |
| 789 | )) |
| 790 | .await; |
| 791 | } |
| 792 | |
| 793 | if outcome.as_ref().is_ok_and(|result| result.result.success) |
| 794 | && let Some((authority, writes)) = child_mcp_call |
| 795 | { |
| 796 | authority.record_settled_writes(writes).await; |
| 797 | } |
| 798 | let duration_ms = started_at.elapsed().as_millis() as u64; |
| 799 | // The surface-agnostic choke point for every tool call, so this one |
| 800 | // bump covers exec and the CLI as well as the TUI. `memory_search` is |
| 801 | // counted here for the same reason — one site, not one per tool. |
| 802 | let telemetry = codewhale_telemetry::session_counters(); |
| 803 | telemetry.bump(codewhale_telemetry::Counter::ToolCalls); |
| 804 | if tool_name == "memory_search" { |
| 805 | telemetry.bump(codewhale_telemetry::Counter::MemorySearch); |
| 806 | } |
| 807 | match &outcome { |
| 808 | Ok(result) => { |
| 809 | tracing::debug!( |
| 810 | target: "engine.tool_execution", |
| 811 | tool = %tool_name, |
| 812 | dispatch, |
| 813 | duration_ms, |
| 814 | success = result.result.success, |
| 815 | output_bytes = result.result.content.len(), |
| 816 | "tool.exec.end", |
| 817 | ); |
| 818 | } |
| 819 | Err(err) => { |
| 820 | let kind = match err { |
| 821 | ToolError::InvalidInput { .. } => "invalid_input", |
| 822 | ToolError::MissingField { .. } => "missing_field", |
| 823 | ToolError::PathEscape { .. } => "path_escape", |
| 824 | ToolError::ExecutionFailed { .. } => "execution_failed", |
| 825 | ToolError::Timeout { .. } => "timeout", |
| 826 | ToolError::Cancelled { .. } => "cancelled", |
| 827 | ToolError::NotAvailable { .. } => "not_available", |
| 828 | ToolError::PermissionDenied { .. } => "permission_denied", |
| 829 | }; |
| 830 | // The discriminant and nothing else. `ToolError::PathEscape`'s |
| 831 | // `Display` *is* an absolute path, and several sibling |
| 832 | // variants render a literal source fragment the model emitted. |
| 833 | match err { |
| 834 | ToolError::PermissionDenied { .. } => { |
| 835 | telemetry.bump_error(codewhale_telemetry::ErrorCounter::ToolDeniedByPolicy) |
| 836 | } |
| 837 | ToolError::Timeout { .. } => { |
| 838 | telemetry.bump_error(codewhale_telemetry::ErrorCounter::ToolTimeout); |
| 839 | } |
| 840 | _ => {} |
| 841 | } |
| 842 | tracing::warn!( |
| 843 | target: "engine.tool_execution", |
| 844 | tool = %tool_name, |
| 845 | dispatch, |
| 846 | duration_ms, |
| 847 | error_kind = kind, |
| 848 | error = %err, |
| 849 | "tool.exec.end", |
| 850 | ); |
| 851 | } |
| 852 | } |
| 853 | outcome |
| 854 | } |
| 855 | } |
| 856 | |
| 857 | #[cfg(test)] |
| 858 | mod tests { |
| 859 | use super::*; |
| 860 | use serde_json::json; |
| 861 | use std::time::Duration; |
| 862 | |
| 863 | const TEST_HEARTBEAT_INTERVAL: Duration = Duration::from_millis(10); |
| 864 | |
| 865 | #[tokio::test] |
| 866 | async fn tool_heartbeat_emits_for_slow_tool() { |
| 867 | let (tx, mut rx) = mpsc::channel(4); |
| 868 | let guard = ToolHeartbeatGuard::start(tx, TEST_HEARTBEAT_INTERVAL); |
| 869 | |
| 870 | let event = tokio::time::timeout(Duration::from_secs(1), rx.recv()) |
| 871 | .await |
| 872 | .expect("heartbeat before slow tool completes") |
| 873 | .expect("event channel stays open"); |
| 874 | |
| 875 | assert!(matches!(event, Event::ToolCallHeartbeat)); |
| 876 | drop(guard); |
| 877 | } |
| 878 | |
| 879 | #[tokio::test] |
| 880 | async fn event_capacity_cancelled_activity_preserves_exact_outcomes_and_releases_full_queue() { |
| 881 | use codewhale_protocol::engine_owner::{OwnerActivityKind, OwnerOperationOutcome}; |
| 882 | for outcome in [ |
| 883 | OwnerOperationOutcome::Succeeded, |
| 884 | OwnerOperationOutcome::Failed, |
| 885 | OwnerOperationOutcome::Denied, |
| 886 | OwnerOperationOutcome::Cancelled, |
| 887 | ] { |
| 888 | let (tx, mut rx) = mpsc::channel(2); |
| 889 | let cancel = CancellationToken::new(); |
| 890 | let span = OperationSpanGuard::start( |
| 891 | tx, |
| 892 | "completed-call", |
| 893 | OwnerActivityKind::Tool, |
| 894 | Some(cancel.clone()), |
| 895 | ) |
| 896 | .await; |
| 897 | cancel.cancel(); |
| 898 | tokio::time::timeout(Duration::from_millis(100), span.complete(outcome)) |
| 899 | .await |
| 900 | .expect("completed activity preserves available capacity after cancellation"); |
| 901 | let Event::OperationActivityStarted { |
| 902 | span_id, |
| 903 | activity_kind, |
| 904 | } = rx.try_recv().unwrap() |
| 905 | else { |
| 906 | panic!("activity start"); |
| 907 | }; |
| 908 | let Event::OperationActivityCompleted { |
| 909 | span_id: completed_span, |
| 910 | activity_kind: completed_kind, |
| 911 | outcome: completed_outcome, |
| 912 | } = rx |
| 913 | .try_recv() |
| 914 | .expect("cancellation must not discard the completion") |
| 915 | else { |
| 916 | panic!("activity completion"); |
| 917 | }; |
| 918 | assert_eq!(completed_span, span_id); |
| 919 | assert_eq!(completed_kind, activity_kind); |
| 920 | assert_eq!( |
| 921 | completed_outcome, outcome, |
| 922 | "retain the observed outcome exactly" |
| 923 | ); |
| 924 | assert!(rx.try_recv().is_err(), "exactly one completion"); |
| 925 | } |
| 926 | |
| 927 | let (tx, mut rx) = mpsc::channel(1); |
| 928 | let cancel = CancellationToken::new(); |
| 929 | let span = OperationSpanGuard::start( |
| 930 | tx, |
| 931 | "full-call", |
| 932 | OwnerActivityKind::Tool, |
| 933 | Some(cancel.clone()), |
| 934 | ) |
| 935 | .await; |
| 936 | cancel.cancel(); |
| 937 | tokio::time::timeout( |
| 938 | Duration::from_millis(100), |
| 939 | span.complete(OwnerOperationOutcome::Succeeded), |
| 940 | ) |
| 941 | .await |
| 942 | .expect("a full queue cannot retain a cancelled activity sender"); |
| 943 | assert!(matches!( |
| 944 | rx.try_recv().unwrap(), |
| 945 | Event::OperationActivityStarted { .. } |
| 946 | )); |
| 947 | assert!( |
| 948 | rx.try_recv().is_err(), |
| 949 | "no sender survives to publish after draining" |
| 950 | ); |
| 951 | } |
| 952 | |
| 953 | #[tokio::test] |
| 954 | async fn tool_heartbeat_is_delayed_for_fast_tool() { |
| 955 | let (tx, mut rx) = mpsc::channel(4); |
| 956 | |
| 957 | let guard = ToolHeartbeatGuard::start(tx, TEST_HEARTBEAT_INTERVAL); |
| 958 | drop(guard); |
| 959 | tokio::time::sleep(TEST_HEARTBEAT_INTERVAL * 2).await; |
| 960 | |
| 961 | assert!(rx.try_recv().is_err(), "fast tool emitted a heartbeat"); |
| 962 | } |
| 963 | |
| 964 | #[tokio::test] |
| 965 | async fn tool_heartbeat_stops_after_tool_completes() { |
| 966 | let (tx, mut rx) = mpsc::channel(8); |
| 967 | let guard = ToolHeartbeatGuard::start(tx, TEST_HEARTBEAT_INTERVAL); |
| 968 | |
| 969 | let event = tokio::time::timeout(Duration::from_secs(1), rx.recv()) |
| 970 | .await |
| 971 | .expect("heartbeat before slow tool completes") |
| 972 | .expect("event channel stays open"); |
| 973 | assert!(matches!(event, Event::ToolCallHeartbeat)); |
| 974 | |
| 975 | drop(guard); |
| 976 | tokio::task::yield_now().await; |
| 977 | while rx.try_recv().is_ok() {} |
| 978 | tokio::time::sleep(TEST_HEARTBEAT_INTERVAL * 2).await; |
| 979 | assert!( |
| 980 | rx.try_recv().is_err(), |
| 981 | "heartbeat continued after tool completion" |
| 982 | ); |
| 983 | } |
| 984 | |
| 985 | #[tokio::test] |
| 986 | async fn full_event_channel_never_blocks_tool_heartbeat() { |
| 987 | let (tx, mut rx) = mpsc::channel(1); |
| 988 | tx.try_send(Event::status("filler")).expect("fill channel"); |
| 989 | |
| 990 | let result = tokio::time::timeout(Duration::from_secs(1), async { |
| 991 | let guard = ToolHeartbeatGuard::start(tx, TEST_HEARTBEAT_INTERVAL); |
| 992 | tokio::time::sleep(TEST_HEARTBEAT_INTERVAL * 3).await; |
| 993 | drop(guard); |
| 994 | "done" |
| 995 | }) |
| 996 | .await |
| 997 | .expect("full event channel must not block tool completion"); |
| 998 | |
| 999 | assert_eq!(result, "done"); |
| 1000 | assert!(matches!(rx.recv().await, Some(Event::Status { .. }))); |
| 1001 | assert!(rx.try_recv().is_err(), "heartbeat displaced queued event"); |
| 1002 | } |
| 1003 | |
| 1004 | #[tokio::test] |
| 1005 | async fn terminal_guard_queues_resume_when_event_channel_is_full() { |
| 1006 | let (tx, mut rx) = mpsc::channel(2); |
| 1007 | let resume = tx.clone().reserve_owned().await.expect("resume capacity"); |
| 1008 | tx.try_send(Event::status("filler")).expect("fill channel"); |
| 1009 | |
| 1010 | drop(InteractiveTerminalGuard { |
| 1011 | resume: Some((tx, resume)), |
| 1012 | }); |
| 1013 | |
| 1014 | assert!(matches!(rx.recv().await, Some(Event::Status { .. }))); |
| 1015 | let resumed = tokio::time::timeout(Duration::from_secs(1), rx.recv()) |
| 1016 | .await |
| 1017 | .expect("queued resume event") |
| 1018 | .expect("event channel still open"); |
| 1019 | assert!(matches!(resumed, Event::ResumeEvents)); |
| 1020 | assert!(rx.try_recv().is_err(), "restoration is exactly once"); |
| 1021 | } |
| 1022 | |
| 1023 | #[tokio::test] |
| 1024 | async fn terminal_guard_waits_for_pause_ack_before_returning() { |
| 1025 | let (tx, mut rx) = mpsc::channel(4); |
| 1026 | let task = tokio::spawn(InteractiveTerminalGuard::engage(tx, true, None)); |
| 1027 | |
| 1028 | let event = tokio::time::timeout(Duration::from_secs(1), rx.recv()) |
| 1029 | .await |
| 1030 | .expect("pause event") |
| 1031 | .expect("event channel still open"); |
| 1032 | let ack = match event { |
| 1033 | Event::PauseEvents { ack: Some(ack) } => ack, |
| 1034 | other => panic!("expected PauseEvents with ack, got {other:?}"), |
| 1035 | }; |
| 1036 | |
| 1037 | tokio::task::yield_now().await; |
| 1038 | assert!(!task.is_finished(), "guard returned before pause ack"); |
| 1039 | |
| 1040 | ack.notify_one(); |
| 1041 | let guard = tokio::time::timeout(Duration::from_secs(1), task) |
| 1042 | .await |
| 1043 | .expect("guard returned after ack") |
| 1044 | .expect("guard task joined") |
| 1045 | .expect("terminal handoff acknowledged"); |
| 1046 | |
| 1047 | drop(guard); |
| 1048 | let resumed = tokio::time::timeout(Duration::from_secs(1), rx.recv()) |
| 1049 | .await |
| 1050 | .expect("resume event") |
| 1051 | .expect("event channel still open"); |
| 1052 | assert!(matches!(resumed, Event::ResumeEvents)); |
| 1053 | } |
| 1054 | |
| 1055 | #[tokio::test] |
| 1056 | async fn terminal_guard_refuses_child_and_queues_resume_when_pause_is_not_acknowledged() { |
| 1057 | let (tx, mut rx) = mpsc::channel(4); |
| 1058 | let task = tokio::spawn(InteractiveTerminalGuard::engage(tx, true, None)); |
| 1059 | |
| 1060 | let event = tokio::time::timeout(Duration::from_secs(1), rx.recv()) |
| 1061 | .await |
| 1062 | .expect("pause event") |
| 1063 | .expect("event channel still open"); |
| 1064 | let _unacknowledged_pause = match event { |
| 1065 | Event::PauseEvents { ack: Some(ack) } => ack, |
| 1066 | other => panic!("expected PauseEvents with ack, got {other:?}"), |
| 1067 | }; |
| 1068 | |
| 1069 | let handoff = tokio::time::timeout(Duration::from_secs(2), task) |
| 1070 | .await |
| 1071 | .expect("guard refused child after pause timeout") |
| 1072 | .expect("guard task joined"); |
| 1073 | let err = match handoff { |
| 1074 | Ok(_) => panic!("unacknowledged terminal handoff must fail closed"), |
| 1075 | Err(err) => err, |
| 1076 | }; |
| 1077 | assert!( |
| 1078 | err.to_string().contains("was not acknowledged"), |
| 1079 | "unexpected handoff error: {err}" |
| 1080 | ); |
| 1081 | |
| 1082 | let resumed = tokio::time::timeout(Duration::from_secs(1), rx.recv()) |
| 1083 | .await |
| 1084 | .expect("queued resume event") |
| 1085 | .expect("event channel still open"); |
| 1086 | assert!(matches!(resumed, Event::ResumeEvents)); |
| 1087 | } |
| 1088 | |
| 1089 | #[tokio::test] |
| 1090 | async fn terminal_guard_cancellation_during_pause_ack_still_queues_resume() { |
| 1091 | let (tx, mut rx) = mpsc::channel(4); |
| 1092 | let task = tokio::spawn(InteractiveTerminalGuard::engage(tx, true, None)); |
| 1093 | |
| 1094 | let event = tokio::time::timeout(Duration::from_secs(1), rx.recv()) |
| 1095 | .await |
| 1096 | .expect("pause event") |
| 1097 | .expect("event channel still open"); |
| 1098 | let _unacknowledged_pause = match event { |
| 1099 | Event::PauseEvents { ack: Some(ack) } => ack, |
| 1100 | other => panic!("expected PauseEvents with ack, got {other:?}"), |
| 1101 | }; |
| 1102 | |
| 1103 | task.abort(); |
| 1104 | let cancelled = match task.await { |
| 1105 | Ok(_) => panic!("engage future should be cancelled"), |
| 1106 | Err(cancelled) => cancelled, |
| 1107 | }; |
| 1108 | assert!(cancelled.is_cancelled()); |
| 1109 | |
| 1110 | let resumed = tokio::time::timeout(Duration::from_secs(1), rx.recv()) |
| 1111 | .await |
| 1112 | .expect("queued resume event") |
| 1113 | .expect("event channel still open"); |
| 1114 | assert!(matches!(resumed, Event::ResumeEvents)); |
| 1115 | } |
| 1116 | |
| 1117 | #[tokio::test] |
| 1118 | async fn terminal_guard_cancelled_capacity_reservations_never_pause_or_restore() { |
| 1119 | for capacity in [1, 2] { |
| 1120 | let (tx, mut rx) = mpsc::channel(capacity); |
| 1121 | tx.try_send(Event::status("occupied")).unwrap(); |
| 1122 | let cancel = CancellationToken::new(); |
| 1123 | let mut engage = Box::pin(InteractiveTerminalGuard::engage( |
| 1124 | tx.clone(), |
| 1125 | true, |
| 1126 | Some(cancel.clone()), |
| 1127 | )); |
| 1128 | assert!( |
| 1129 | tokio::time::timeout(Duration::from_millis(20), &mut engage) |
| 1130 | .await |
| 1131 | .is_err() |
| 1132 | ); |
| 1133 | // Capacity 1 stalls the resume reservation; capacity 2 stalls |
| 1134 | // the pause reservation while the first permit is held. |
| 1135 | cancel.cancel(); |
| 1136 | assert!( |
| 1137 | tokio::time::timeout(Duration::from_secs(1), &mut engage) |
| 1138 | .await |
| 1139 | .unwrap() |
| 1140 | .is_err() |
| 1141 | ); |
| 1142 | drop(engage); |
| 1143 | assert!(matches!(rx.try_recv(), Ok(Event::Status { .. }))); |
| 1144 | assert!( |
| 1145 | rx.try_recv().is_err(), |
| 1146 | "a pause that never entered needs no resume" |
| 1147 | ); |
| 1148 | assert_eq!( |
| 1149 | tx.capacity(), |
| 1150 | capacity, |
| 1151 | "cancelled waits release every permit" |
| 1152 | ); |
| 1153 | } |
| 1154 | } |
| 1155 | |
| 1156 | #[tokio::test] |
| 1157 | async fn terminal_guard_cancelled_ack_restores_once_without_detached_sender() { |
| 1158 | let (tx, mut rx) = mpsc::channel(3); |
| 1159 | let cancel = CancellationToken::new(); |
| 1160 | let task = tokio::spawn(InteractiveTerminalGuard::engage( |
| 1161 | tx.clone(), |
| 1162 | true, |
| 1163 | Some(cancel.clone()), |
| 1164 | )); |
| 1165 | assert!(matches!(rx.recv().await, Some(Event::PauseEvents { .. }))); |
| 1166 | tx.try_send(Event::status("fill after pause")).unwrap(); |
| 1167 | cancel.cancel(); |
| 1168 | assert!( |
| 1169 | tokio::time::timeout(Duration::from_secs(1), task) |
| 1170 | .await |
| 1171 | .unwrap() |
| 1172 | .unwrap() |
| 1173 | .is_err() |
| 1174 | ); |
| 1175 | assert!(matches!(rx.try_recv(), Ok(Event::Status { .. }))); |
| 1176 | assert!(matches!(rx.try_recv(), Ok(Event::ResumeEvents))); |
| 1177 | assert!(rx.try_recv().is_err()); |
| 1178 | assert_eq!(tx.capacity(), 3); |
| 1179 | } |
| 1180 | |
| 1181 | #[tokio::test] |
| 1182 | async fn terminal_guard_closed_receiver_refuses_or_releases_held_restoration() { |
| 1183 | let (tx, rx) = mpsc::channel(3); |
| 1184 | drop(rx); |
| 1185 | assert!( |
| 1186 | InteractiveTerminalGuard::engage(tx, true, None) |
| 1187 | .await |
| 1188 | .is_err() |
| 1189 | ); |
| 1190 | |
| 1191 | let (tx, mut rx) = mpsc::channel(3); |
| 1192 | let task = tokio::spawn(InteractiveTerminalGuard::engage(tx.clone(), true, None)); |
| 1193 | let Some(Event::PauseEvents { ack: Some(ack) }) = rx.recv().await else { |
| 1194 | panic!("pause with acknowledgement"); |
| 1195 | }; |
| 1196 | ack.notify_one(); |
| 1197 | let guard = task.await.unwrap().unwrap(); |
| 1198 | drop(rx); |
| 1199 | drop(guard); |
| 1200 | assert!(tx.is_closed()); |
| 1201 | assert_eq!( |
| 1202 | tx.capacity(), |
| 1203 | 3, |
| 1204 | "receiver closure cannot strand restoration capacity" |
| 1205 | ); |
| 1206 | } |
| 1207 | |
| 1208 | #[cfg(unix)] |
| 1209 | #[test] |
| 1210 | fn inherited_interactive_shell_is_refused_before_terminal_handoff() { |
| 1211 | for tool_name in ["bash", "Bash", "exec_shell"] { |
| 1212 | let err = inherited_interactive_shell_refusal(tool_name, true) |
| 1213 | .expect("Unix inherited-interactive shell must fail preflight"); |
| 1214 | assert!( |
| 1215 | err.to_string().contains("foreground TTY ownership"), |
| 1216 | "{tool_name}: {err}" |
| 1217 | ); |
| 1218 | } |
| 1219 | assert!(inherited_interactive_shell_refusal("Bash", false).is_none()); |
| 1220 | assert!( |
| 1221 | inherited_interactive_shell_refusal(REQUEST_USER_INPUT_NAME, true).is_none(), |
| 1222 | "user-input modal keeps its own terminal handoff" |
| 1223 | ); |
| 1224 | } |
| 1225 | |
| 1226 | #[test] |
| 1227 | fn emit_tool_audit_to_path_writes_jsonl_lines() { |
| 1228 | let tmp = tempfile::tempdir().expect("tempdir"); |
| 1229 | let path = tmp.path().join("audit.log"); |
| 1230 | let marker = path.display().to_string(); |
| 1231 | |
| 1232 | emit_tool_audit_to_path( |
| 1233 | &path, |
| 1234 | json!({ |
| 1235 | "event": "tool.spillover", |
| 1236 | "test_marker": marker, |
| 1237 | "tool_id": "call-abc", |
| 1238 | "tool_name": "exec_shell", |
| 1239 | "path": "/tmp/foo.txt", |
| 1240 | }), |
| 1241 | ); |
| 1242 | emit_tool_audit_to_path( |
| 1243 | &path, |
| 1244 | json!({ |
| 1245 | "event": "tool.result", |
| 1246 | "test_marker": marker, |
| 1247 | "tool_id": "call-xyz", |
| 1248 | "success": true, |
| 1249 | }), |
| 1250 | ); |
| 1251 | |
| 1252 | let body = std::fs::read_to_string(&path).expect("audit log written"); |
| 1253 | let entries: Vec<serde_json::Value> = body |
| 1254 | .lines() |
| 1255 | .map(|line| serde_json::from_str(line).expect("audit line is JSON")) |
| 1256 | .filter(|entry: &serde_json::Value| { |
| 1257 | entry.get("test_marker").and_then(|v| v.as_str()) == Some(marker.as_str()) |
| 1258 | }) |
| 1259 | .collect(); |
| 1260 | assert_eq!(entries.len(), 2, "two marked emits -> two lines"); |
| 1261 | |
| 1262 | // Each line round-trips as JSON, has the expected event key. |
| 1263 | let first = &entries[0]; |
| 1264 | assert_eq!( |
| 1265 | first.get("event").and_then(|v| v.as_str()), |
| 1266 | Some("tool.spillover") |
| 1267 | ); |
| 1268 | assert_eq!( |
| 1269 | first.get("tool_id").and_then(|v| v.as_str()), |
| 1270 | Some("call-abc") |
| 1271 | ); |
| 1272 | |
| 1273 | let second = &entries[1]; |
| 1274 | assert_eq!( |
| 1275 | second.get("event").and_then(|v| v.as_str()), |
| 1276 | Some("tool.result") |
| 1277 | ); |
| 1278 | } |
| 1279 | |
| 1280 | #[test] |
| 1281 | fn emit_tool_audit_creates_parent_directory() { |
| 1282 | let tmp = tempfile::tempdir().expect("tempdir"); |
| 1283 | // Path with a parent that doesn't exist yet — the writer |
| 1284 | // should create it. |
| 1285 | let nested = tmp.path().join("nested").join("dir").join("audit.log"); |
| 1286 | emit_tool_audit_to_path(&nested, json!({"event": "test"})); |
| 1287 | assert!(nested.exists(), "writer should mkdir -p the parent chain"); |
| 1288 | } |
| 1289 | } |
| 1290 |