| 1 | //! Tool dispatch — plan/execute helpers for the per-turn tool batch. |
| 2 | //! |
| 3 | //! Extracted from `core/engine.rs` (P1.3). The high-level ordering still |
| 4 | //! lives in `Engine::run_turn`; this module owns: |
| 5 | //! |
| 6 | //! * Streaming-buffer parsing into a finalized `serde_json::Value` tool input |
| 7 | //! (`final_tool_input`, `parse_tool_input`, fenced/JSON segment helpers). |
| 8 | //! * The `multi_tool_use.parallel` payload parser. |
| 9 | //! * Policy predicates the turn loop consults — when a batch can run in |
| 10 | //! parallel and the small set of read-only MCP tools that are safe to run |
| 11 | //! in parallel. |
| 12 | //! * The tool execution plan/outcome types the batch driver passes around. |
| 13 | //! |
| 14 | //! All items are `pub(super)`-only: the public engine surface (Op/Event, |
| 15 | //! `EngineHandle`, `spawn_engine`) stays in `core/engine.rs`. |
| 16 | |
| 17 | use serde_json::json; |
| 18 | |
| 19 | use crate::tools::spec::{ |
| 20 | ResourceClaim, ToolError, ToolExecutionOutcome, ToolResult, ToolResultContentBlock, |
| 21 | schedule_non_conflicting, |
| 22 | }; |
| 23 | use codewhale_models::{Tool, ToolCaller}; |
| 24 | |
| 25 | use super::ToolUseState; |
| 26 | |
| 27 | const MAX_SCHEMA_CONTAINER_REPAIR_BYTES: usize = 64 * 1024; |
| 28 | |
| 29 | // === Types ============================================================ |
| 30 | |
| 31 | #[allow(dead_code)] // `index` mirrors batch order for diagnostic ergonomics. |
| 32 | pub(super) struct ToolExecOutcome { |
| 33 | pub(super) index: usize, |
| 34 | pub(super) id: String, |
| 35 | pub(super) name: String, |
| 36 | pub(super) input: serde_json::Value, |
| 37 | pub(super) started_at: std::time::Instant, |
| 38 | pub(super) terminal: ToolExecutionOutcome, |
| 39 | pub(super) content_blocks: Vec<ToolResultContentBlock>, |
| 40 | /// Read-result bytes before spillover adds call-specific artifact paths. |
| 41 | pub(super) original_content_digest: Option<[u8; 32]>, |
| 42 | } |
| 43 | |
| 44 | /// Notice appended as a user-role message when the guard first asks the worker |
| 45 | /// to change strategy after repeated no-progress denials (#6015). |
| 46 | pub(crate) const FLEET_STRATEGY_SWITCH_NOTICE: &str = "Fleet strategy switch required: repeated permission denials produced no new evidence. The rejected action is held. Use another permitted tool from the current catalog to make progress, or report completed work and the blocker. Do not work around permissions or request the same approval again."; |
| 47 | |
| 48 | /// Notice appended when denials continue past the strategy switch: the next |
| 49 | /// response is report-only and its tool calls are admission-held (#6015). |
| 50 | pub(crate) const FLEET_FINAL_REPORT_NOTICE: &str = "Fleet no-progress final report: permission denials continued after the strategy switch without new evidence. Your next response is report-only; no tools will execute. Report what you completed, exact evidence, the permission blocker and remaining work. This is the last response unless the user changes direction or authority."; |
| 51 | |
| 52 | /// Terminal reason once the report-only response has been recorded (#6015). |
| 53 | pub(crate) const FLEET_NO_PROGRESS_STOP: &str = "Fleet worker stopped after repeated permission denials without new evidence. Work and tool results are retained in the transcript; review the blocker before resuming."; |
| 54 | |
| 55 | /// Progress observations for one provider response, independent of tool finish |
| 56 | /// order. Only typed permission denials contribute to the retry guard (#6015). |
| 57 | #[derive(Default)] |
| 58 | pub(crate) struct FleetDenialBatch { |
| 59 | denied: std::collections::HashSet<String>, |
| 60 | made_progress: bool, |
| 61 | } |
| 62 | |
| 63 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 64 | pub(crate) enum FleetDenialAction { |
| 65 | Continue, |
| 66 | SwitchStrategy, |
| 67 | FinalReport, |
| 68 | } |
| 69 | |
| 70 | /// Turn-local guard for an engine with a Fleet authority envelope. This is an |
| 71 | /// admission predicate and result accumulator, not another execution loop. |
| 72 | /// Three responses give the model two opportunities to use denial feedback; |
| 73 | /// after one strategy notice, three more denied responses request a report. |
| 74 | /// Fleet sub-agent workers run the same guard in their own loop (#6015). |
| 75 | #[derive(Default)] |
| 76 | pub(crate) struct FleetDenialGuard { |
| 77 | denied_rounds: std::collections::HashMap<String, u8>, |
| 78 | switch_requested: bool, |
| 79 | recovery_denied_rounds: u8, |
| 80 | denial_rounds_without_progress: u32, |
| 81 | report_only: bool, |
| 82 | // Last observed bytes per read request: paths alone cannot distinguish a |
| 83 | // changed file, and an unchanged read must not repeatedly reset denials. |
| 84 | // Coverage is bounded; an evicted observation is treated conservatively |
| 85 | // as new evidence. No raw arguments/bytes are kept. |
| 86 | reads: std::collections::VecDeque<([u8; 32], [u8; 32])>, |
| 87 | } |
| 88 | |
| 89 | impl FleetDenialGuard { |
| 90 | const REPEATED_DENIAL_ROUNDS: u8 = 3; |
| 91 | const MAX_OBSERVATIONS: usize = 32; |
| 92 | |
| 93 | pub(crate) fn reset(&mut self) { |
| 94 | *self = Self::default(); |
| 95 | } |
| 96 | |
| 97 | pub(crate) fn report_only(&self) -> bool { |
| 98 | self.report_only |
| 99 | } |
| 100 | |
| 101 | pub(crate) fn awaiting_strategy_change(&self) -> bool { |
| 102 | self.switch_requested |
| 103 | } |
| 104 | |
| 105 | pub(super) fn denial_rounds_without_progress(&self) -> u32 { |
| 106 | self.denial_rounds_without_progress |
| 107 | } |
| 108 | |
| 109 | pub(crate) fn original_content_digest( |
| 110 | name: &str, |
| 111 | input: &serde_json::Value, |
| 112 | output: &ToolResult, |
| 113 | ) -> Option<[u8; 32]> { |
| 114 | use sha2::{Digest, Sha256}; |
| 115 | let action = crate::tools::canonical_action::canonical_action_alias(name, input); |
| 116 | (output.success |
| 117 | && matches!( |
| 118 | action, |
| 119 | "read_file" | "list_dir" | "file_search" | "grep_files" |
| 120 | )) |
| 121 | .then(|| Sha256::digest(output.content.as_bytes()).into()) |
| 122 | } |
| 123 | |
| 124 | pub(crate) fn admission_error( |
| 125 | &self, |
| 126 | name: &str, |
| 127 | input: &serde_json::Value, |
| 128 | ) -> Option<ToolError> { |
| 129 | let action = crate::tools::canonical_action::canonical_action_alias(name, input); |
| 130 | if self.report_only { |
| 131 | Some(ToolError::permission_denied( |
| 132 | "Fleet no-progress final report: no tools may execute in this response. Report completed work, evidence and the remaining blocker; do not change permission mode or retry tools.", |
| 133 | )) |
| 134 | } else if self.switch_requested && self.denied_rounds.contains_key(action) { |
| 135 | Some(ToolError::permission_denied( |
| 136 | "Fleet permission-denial loop: this action is held until useful permitted work or an explicit authority change. Use another permitted tool or report the blocker; do not change permission mode or request permission again.", |
| 137 | )) |
| 138 | } else { |
| 139 | None |
| 140 | } |
| 141 | } |
| 142 | |
| 143 | pub(crate) fn observe( |
| 144 | &mut self, |
| 145 | batch: &mut FleetDenialBatch, |
| 146 | name: &str, |
| 147 | input: &serde_json::Value, |
| 148 | status: crate::tools::spec::ToolTerminalStatus, |
| 149 | result: &Result<ToolResult, ToolError>, |
| 150 | original_content_digest: Option<[u8; 32]>, |
| 151 | ) { |
| 152 | use crate::tools::spec::ToolTerminalStatus; |
| 153 | |
| 154 | let action = crate::tools::canonical_action::canonical_action_alias(name, input); |
| 155 | if status == ToolTerminalStatus::Denied |
| 156 | && matches!(result, Err(ToolError::PermissionDenied { .. })) |
| 157 | { |
| 158 | batch.denied.insert(action.to_owned()); |
| 159 | return; |
| 160 | } |
| 161 | let Ok(output) = result else { return }; |
| 162 | if status != ToolTerminalStatus::Succeeded |
| 163 | || !output.success |
| 164 | || output.metadata.as_ref().is_some_and(|metadata| { |
| 165 | metadata |
| 166 | .get("executed") |
| 167 | .and_then(serde_json::Value::as_bool) |
| 168 | == Some(false) |
| 169 | || metadata |
| 170 | .get("cancelled") |
| 171 | .and_then(serde_json::Value::as_bool) |
| 172 | == Some(true) |
| 173 | }) |
| 174 | { |
| 175 | return; |
| 176 | } |
| 177 | // Waiting is useful coordination, but its repeated success receipt is |
| 178 | // neither new evidence nor a failure. Its own timeouts still govern it. |
| 179 | if matches!( |
| 180 | action, |
| 181 | "exec_shell_wait" | "exec_wait" | "terminal/wait" | "wait_for_dev_server" | "sleep" |
| 182 | ) || name == "agent" |
| 183 | && input |
| 184 | .get("action") |
| 185 | .and_then(serde_json::Value::as_str) |
| 186 | .is_some_and(|action| matches!(action, "wait" | "status" | "list")) |
| 187 | { |
| 188 | return; |
| 189 | } |
| 190 | if matches!( |
| 191 | action, |
| 192 | "read_file" | "list_dir" | "file_search" | "grep_files" |
| 193 | ) { |
| 194 | use sha2::{Digest, Sha256}; |
| 195 | |
| 196 | let mut semantic_input = input.clone(); |
| 197 | if action != name |
| 198 | && let Some(object) = semantic_input.as_object_mut() |
| 199 | { |
| 200 | object.remove("action"); |
| 201 | } |
| 202 | let mut hasher = Sha256::new(); |
| 203 | hasher.update(action.as_bytes()); |
| 204 | hasher.update([0]); |
| 205 | // Tool JSON preserves insertion order; reordered equivalent keys |
| 206 | // must not manufacture a new read request. |
| 207 | hasher.update(crate::client::canonical_json(&semantic_input).as_bytes()); |
| 208 | let key: [u8; 32] = hasher.finalize().into(); |
| 209 | let contents = original_content_digest |
| 210 | .unwrap_or_else(|| Sha256::digest(output.content.as_bytes()).into()); |
| 211 | let previous = self |
| 212 | .reads |
| 213 | .iter() |
| 214 | .position(|(old_key, _)| *old_key == key) |
| 215 | .and_then(|index| self.reads.remove(index)); |
| 216 | batch.made_progress |= |
| 217 | previous.is_none_or(|(_, old_contents)| old_contents != contents); |
| 218 | self.reads.push_back((key, contents)); |
| 219 | if self.reads.len() > Self::MAX_OBSERVATIONS { |
| 220 | self.reads.pop_front(); |
| 221 | } |
| 222 | } else { |
| 223 | // A successful mutation or unfamiliar tool is useful work. Do not |
| 224 | // terminate it based on guesses about its content or side effects. |
| 225 | batch.made_progress = true; |
| 226 | } |
| 227 | } |
| 228 | |
| 229 | pub(crate) fn finish_batch(&mut self, batch: FleetDenialBatch) -> FleetDenialAction { |
| 230 | if self.report_only { |
| 231 | return FleetDenialAction::Continue; |
| 232 | } |
| 233 | if batch.made_progress { |
| 234 | self.denied_rounds.clear(); |
| 235 | self.switch_requested = false; |
| 236 | self.recovery_denied_rounds = 0; |
| 237 | self.denial_rounds_without_progress = 0; |
| 238 | return FleetDenialAction::Continue; |
| 239 | } |
| 240 | if batch.denied.is_empty() { |
| 241 | return FleetDenialAction::Continue; |
| 242 | } |
| 243 | self.denial_rounds_without_progress = self.denial_rounds_without_progress.saturating_add(1); |
| 244 | if self.switch_requested { |
| 245 | self.recovery_denied_rounds = self.recovery_denied_rounds.saturating_add(1); |
| 246 | if self.recovery_denied_rounds >= Self::REPEATED_DENIAL_ROUNDS { |
| 247 | self.report_only = true; |
| 248 | return FleetDenialAction::FinalReport; |
| 249 | } |
| 250 | return FleetDenialAction::Continue; |
| 251 | } |
| 252 | for action in batch.denied { |
| 253 | // The guard never retains payloads. Unknown families beyond this |
| 254 | // bounded window do not evict an already observed denial streak. |
| 255 | if self.denied_rounds.contains_key(&action) |
| 256 | || self.denied_rounds.len() < Self::MAX_OBSERVATIONS |
| 257 | { |
| 258 | let count = self.denied_rounds.entry(action).or_default(); |
| 259 | *count = count.saturating_add(1); |
| 260 | } |
| 261 | } |
| 262 | if self |
| 263 | .denied_rounds |
| 264 | .values() |
| 265 | .any(|count| *count >= Self::REPEATED_DENIAL_ROUNDS) |
| 266 | { |
| 267 | self.switch_requested = true; |
| 268 | FleetDenialAction::SwitchStrategy |
| 269 | } else { |
| 270 | FleetDenialAction::Continue |
| 271 | } |
| 272 | } |
| 273 | } |
| 274 | |
| 275 | #[derive(Debug, Clone)] |
| 276 | pub(super) struct ToolExecutionPlan { |
| 277 | pub(super) index: usize, |
| 278 | pub(super) id: String, |
| 279 | pub(super) name: String, |
| 280 | pub(super) input: serde_json::Value, |
| 281 | pub(super) caller: Option<ToolCaller>, |
| 282 | pub(super) interactive: bool, |
| 283 | pub(super) approval_required: bool, |
| 284 | pub(super) approval_description: String, |
| 285 | pub(super) approval_force_prompt: bool, |
| 286 | pub(super) supports_parallel: bool, |
| 287 | pub(super) read_only: bool, |
| 288 | pub(super) detached_start: bool, |
| 289 | pub(super) resources: Vec<ResourceClaim>, |
| 290 | pub(super) blocked_error: Option<ToolError>, |
| 291 | pub(super) guard_result: Option<ToolResult>, |
| 292 | } |
| 293 | |
| 294 | pub(super) enum ToolExecutionBatch { |
| 295 | Parallel(Vec<ToolExecutionPlan>), |
| 296 | Serial(Box<ToolExecutionPlan>), |
| 297 | } |
| 298 | |
| 299 | #[derive(Debug, serde::Serialize)] |
| 300 | pub(super) struct ParallelToolResultEntry { |
| 301 | pub(super) tool_name: String, |
| 302 | pub(super) success: bool, |
| 303 | pub(super) content: String, |
| 304 | #[serde(skip_serializing_if = "Option::is_none")] |
| 305 | pub(super) error: Option<String>, |
| 306 | } |
| 307 | |
| 308 | #[derive(Debug, serde::Serialize)] |
| 309 | pub(super) struct ParallelToolResult { |
| 310 | pub(super) results: Vec<ParallelToolResultEntry>, |
| 311 | } |
| 312 | |
| 313 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 314 | pub(super) enum ToolApprovalStamp { |
| 315 | ApprovedByUser, |
| 316 | ApprovedWithPolicy, |
| 317 | } |
| 318 | |
| 319 | impl ToolApprovalStamp { |
| 320 | fn decision(self) -> &'static str { |
| 321 | match self { |
| 322 | Self::ApprovedByUser => "approved_by_user", |
| 323 | Self::ApprovedWithPolicy => "approved_with_policy", |
| 324 | } |
| 325 | } |
| 326 | |
| 327 | fn model_visible_note(self) -> &'static str { |
| 328 | match self { |
| 329 | Self::ApprovedByUser => { |
| 330 | "[approval] This tool call required approval and was approved by the user before execution." |
| 331 | } |
| 332 | Self::ApprovedWithPolicy => { |
| 333 | "[approval] This tool call required approval and was approved by the user with an adjusted execution policy before execution." |
| 334 | } |
| 335 | } |
| 336 | } |
| 337 | } |
| 338 | |
| 339 | pub(super) fn stamp_tool_result_approval(result: &mut ToolResult, approval: ToolApprovalStamp) { |
| 340 | let approval_metadata = json!({ |
| 341 | "required": true, |
| 342 | "decision": approval.decision(), |
| 343 | "model_visible": true, |
| 344 | }); |
| 345 | let metadata = result.metadata.get_or_insert_with(|| json!({})); |
| 346 | if let Some(object) = metadata.as_object_mut() { |
| 347 | object.insert("approval".to_string(), approval_metadata); |
| 348 | } else { |
| 349 | let prior = std::mem::replace(metadata, json!({})); |
| 350 | if let Some(object) = metadata.as_object_mut() { |
| 351 | object.insert("_prior".to_string(), prior); |
| 352 | object.insert("approval".to_string(), approval_metadata); |
| 353 | } |
| 354 | } |
| 355 | |
| 356 | let note = approval.model_visible_note(); |
| 357 | if result.content.starts_with("[approval] ") { |
| 358 | return; |
| 359 | } |
| 360 | if result.content.is_empty() { |
| 361 | result.content = note.to_string(); |
| 362 | } else { |
| 363 | result.content = format!("{note}\n\n{}", result.content); |
| 364 | } |
| 365 | } |
| 366 | |
| 367 | // Hold the lock guard for the duration of a tool execution. |
| 368 | // The inner guards are held for RAII purposes (dropped when the guard is dropped). |
| 369 | pub(super) enum ToolExecGuard<'a> { |
| 370 | Read(#[allow(dead_code)] tokio::sync::RwLockReadGuard<'a, ()>), |
| 371 | Write(#[allow(dead_code)] tokio::sync::RwLockWriteGuard<'a, ()>), |
| 372 | } |
| 373 | |
| 374 | // === Caller policy and errors ======================================== |
| 375 | |
| 376 | pub(super) fn caller_type_for_tool_use(caller: Option<&ToolCaller>) -> &str { |
| 377 | caller.map_or("direct", |c| c.caller_type.as_str()) |
| 378 | } |
| 379 | |
| 380 | pub(super) fn caller_allowed_for_tool( |
| 381 | caller: Option<&ToolCaller>, |
| 382 | tool_def: Option<&Tool>, |
| 383 | ) -> bool { |
| 384 | let requested = caller_type_for_tool_use(caller); |
| 385 | if let Some(def) = tool_def |
| 386 | && let Some(allowed) = &def.allowed_callers |
| 387 | { |
| 388 | if allowed.is_empty() { |
| 389 | return requested == "direct"; |
| 390 | } |
| 391 | return allowed.iter().any(|item| item == requested); |
| 392 | } |
| 393 | requested == "direct" |
| 394 | } |
| 395 | |
| 396 | /// Whole-word check for "mode"/"modes" — a plain `contains("mode")` also |
| 397 | /// matched "model", letting provider model errors skip the actionable-hint |
| 398 | /// suffix (#3020). |
| 399 | fn mentions_mode_word(lower: &str) -> bool { |
| 400 | lower |
| 401 | .split(|ch: char| !ch.is_ascii_alphanumeric()) |
| 402 | .any(|word| word == "mode" || word == "modes") |
| 403 | } |
| 404 | |
| 405 | #[cfg(test)] |
| 406 | pub(super) fn format_tool_error(err: &ToolError, tool_name: &str) -> String { |
| 407 | format_tool_error_with_schema(err, tool_name, None) |
| 408 | } |
| 409 | |
| 410 | pub(super) fn format_tool_error_with_schema( |
| 411 | err: &ToolError, |
| 412 | tool_name: &str, |
| 413 | input_schema: Option<&serde_json::Value>, |
| 414 | ) -> String { |
| 415 | let message = match err { |
| 416 | ToolError::InvalidInput { message } => { |
| 417 | format!("Invalid input for tool '{tool_name}': {message}") |
| 418 | } |
| 419 | ToolError::MissingField { field } => { |
| 420 | format!("Tool '{tool_name}' is missing required field '{field}'") |
| 421 | } |
| 422 | ToolError::PathEscape { path } => format!( |
| 423 | "Path escapes workspace: {}. Use a workspace-relative path or enable trust mode.", |
| 424 | path.display() |
| 425 | ), |
| 426 | ToolError::ExecutionFailed { message } => message.clone(), |
| 427 | ToolError::Timeout { seconds } => format!( |
| 428 | "Tool '{tool_name}' timed out after {seconds}s. Try a narrower scope or a longer timeout." |
| 429 | ), |
| 430 | ToolError::Cancelled { message } => message.clone(), |
| 431 | ToolError::NotAvailable { message } => { |
| 432 | let lower = message.to_ascii_lowercase(); |
| 433 | // #3020: Pass through self-explanatory messages that already name the |
| 434 | // cause (mode switch, allow_shell, feature flag). Avoids appending a |
| 435 | // conflicting "Check mode, feature flags" suffix on top of |
| 436 | // "switch to Act mode" which already gives the recovery path. |
| 437 | if lower.contains("current tool catalog") |
| 438 | || lower.contains("did you mean:") |
| 439 | || mentions_mode_word(&lower) |
| 440 | || lower.contains("allow_shell") |
| 441 | || lower.contains("feature flag") |
| 442 | { |
| 443 | message.clone() |
| 444 | } else { |
| 445 | format!( |
| 446 | "Tool '{tool_name}' is not available: {message}. Check mode, feature flags, or tool name." |
| 447 | ) |
| 448 | } |
| 449 | } |
| 450 | ToolError::PermissionDenied { message } => { |
| 451 | let lower = message.to_ascii_lowercase(); |
| 452 | // #3020: Pass through messages that already name the denial cause. |
| 453 | if mentions_mode_word(&lower) |
| 454 | || lower.contains("allow_shell") |
| 455 | || lower.contains("denied by user") |
| 456 | { |
| 457 | message.clone() |
| 458 | } else { |
| 459 | format!( |
| 460 | "Tool '{tool_name}' was denied: {message}. Adjust approval mode or request permission." |
| 461 | ) |
| 462 | } |
| 463 | } |
| 464 | }; |
| 465 | |
| 466 | let (category, bad_field) = match err { |
| 467 | ToolError::InvalidInput { .. } => ("invalid_input", None), |
| 468 | ToolError::MissingField { field } => ("missing_field", Some(field.as_str())), |
| 469 | ToolError::PathEscape { .. } => ("path_escape", Some("path")), |
| 470 | ToolError::NotAvailable { .. } => ("tool_not_available", Some("tool_name")), |
| 471 | _ => return message, |
| 472 | }; |
| 473 | let valid_shape = input_schema.cloned().unwrap_or_else(|| { |
| 474 | serde_json::json!({ |
| 475 | "type": "object", |
| 476 | "guidance": format!("Use the advertised input schema for '{tool_name}'") |
| 477 | }) |
| 478 | }); |
| 479 | let feedback = serde_json::json!({ |
| 480 | "category": category, |
| 481 | "bad_field": bad_field, |
| 482 | "valid_shape": valid_shape, |
| 483 | "retryable": true, |
| 484 | "side_effect_status": "not_started" |
| 485 | }); |
| 486 | format!("{message}\nTool validation feedback: {feedback}") |
| 487 | } |
| 488 | |
| 489 | // === Streaming-buffer parsing ========================================= |
| 490 | |
| 491 | /// Promote a streaming `ToolUseState` to a finalized JSON input. |
| 492 | /// |
| 493 | /// Order of preference: |
| 494 | /// |
| 495 | /// 1. `input_buffer` (the raw streamed delta concatenation) — parsed as |
| 496 | /// JSON. This is the most authoritative because it's what the model |
| 497 | /// actually emitted. |
| 498 | /// 2. `input` (the per-delta best-effort parse mirror) — used when the |
| 499 | /// buffer is empty (pre-streaming tool calls take this path). |
| 500 | /// 3. `input_buffer` non-empty but unparseable → fall back to `input` |
| 501 | /// (the per-delta parser has already mirrored the most recent valid |
| 502 | /// partial parse into `tool_state.input`). |
| 503 | pub(super) fn final_tool_input(state: &ToolUseState) -> serde_json::Value { |
| 504 | if state.input_parse_error.is_some() { |
| 505 | return malformed_tool_arguments_input(&state.input_buffer); |
| 506 | } |
| 507 | if !state.input_buffer.trim().is_empty() |
| 508 | && let Some(parsed) = parse_tool_input(&state.input_buffer) |
| 509 | { |
| 510 | // Structure was synthesized to make this parse, so the argument text |
| 511 | // was cut off. Route it to the same malformed-arguments path as an |
| 512 | // outright parse failure rather than dispatching a completed guess. |
| 513 | if parsed.structure_synthesized { |
| 514 | return malformed_tool_arguments_input(&state.input_buffer); |
| 515 | } |
| 516 | return parsed.value; |
| 517 | } |
| 518 | state.input.clone() |
| 519 | } |
| 520 | |
| 521 | /// A parsed tool-argument buffer, plus whether the parse only succeeded |
| 522 | /// because the repair ladder synthesized structure (see |
| 523 | /// `crate::tools::arg_repair`). Mid-stream callers mirroring partial state |
| 524 | /// may ignore the flag; the caller making the final dispatch decision must |
| 525 | /// not, because synthesized structure means the argument text was cut off. |
| 526 | pub(super) struct ParsedToolInput { |
| 527 | pub(super) value: serde_json::Value, |
| 528 | pub(super) structure_synthesized: bool, |
| 529 | } |
| 530 | |
| 531 | pub(super) fn parse_tool_input(buffer: &str) -> Option<ParsedToolInput> { |
| 532 | let trimmed = buffer.trim(); |
| 533 | if trimmed.is_empty() { |
| 534 | return None; |
| 535 | } |
| 536 | // Try the deterministic arg-repair ladder first (handles trailing commas, |
| 537 | // unclosed braces, embedded control chars, etc.) |
| 538 | if let Ok(repaired) = crate::tools::arg_repair::repair(trimmed) { |
| 539 | return Some(ParsedToolInput { |
| 540 | value: repaired.value, |
| 541 | structure_synthesized: repaired.structure_synthesized, |
| 542 | }); |
| 543 | } |
| 544 | // Fall back to existing strategies for code-fenced, double-encoded, and |
| 545 | // segment-extraction patterns that the repair ladder doesn't cover. |
| 546 | if let Some(stripped) = strip_code_fences(trimmed) |
| 547 | && let Ok(value) = serde_json::from_str::<serde_json::Value>(&stripped) |
| 548 | { |
| 549 | return Some(ParsedToolInput { |
| 550 | value, |
| 551 | structure_synthesized: false, |
| 552 | }); |
| 553 | } |
| 554 | if let Ok(serde_json::Value::String(inner)) = serde_json::from_str::<serde_json::Value>(trimmed) |
| 555 | && let Ok(value) = serde_json::from_str::<serde_json::Value>(&inner) |
| 556 | { |
| 557 | return Some(ParsedToolInput { |
| 558 | value, |
| 559 | structure_synthesized: false, |
| 560 | }); |
| 561 | } |
| 562 | extract_json_segment(trimmed) |
| 563 | .and_then(|segment| serde_json::from_str::<serde_json::Value>(&segment).ok()) |
| 564 | .map(|value| ParsedToolInput { |
| 565 | value, |
| 566 | structure_synthesized: false, |
| 567 | }) |
| 568 | } |
| 569 | |
| 570 | /// Decode a JSON container that a provider encoded as a string when the tool |
| 571 | /// schema explicitly requires an object or array. |
| 572 | /// |
| 573 | /// This intentionally avoids general argument coercion: the string must be |
| 574 | /// bounded, parse as strict JSON, and decode to the declared container type. |
| 575 | /// Primitive strings are never coerced. |
| 576 | pub(super) fn normalize_schema_json_containers( |
| 577 | value: &mut serde_json::Value, |
| 578 | schema: &serde_json::Value, |
| 579 | ) -> usize { |
| 580 | let expected_container = if schema_declares_type(schema, "object") { |
| 581 | Some("object") |
| 582 | } else if schema_declares_type(schema, "array") { |
| 583 | Some("array") |
| 584 | } else { |
| 585 | None |
| 586 | }; |
| 587 | |
| 588 | if let (Some(expected), serde_json::Value::String(encoded)) = (expected_container, &*value) |
| 589 | && encoded.len() <= MAX_SCHEMA_CONTAINER_REPAIR_BYTES |
| 590 | && let Ok(decoded) = serde_json::from_str::<serde_json::Value>(encoded) |
| 591 | && ((expected == "object" && decoded.is_object()) |
| 592 | || (expected == "array" && decoded.is_array())) |
| 593 | { |
| 594 | *value = decoded; |
| 595 | return 1 + normalize_schema_json_containers(value, schema); |
| 596 | } |
| 597 | |
| 598 | match value { |
| 599 | serde_json::Value::Object(object) => { |
| 600 | let properties = schema |
| 601 | .get("properties") |
| 602 | .and_then(serde_json::Value::as_object); |
| 603 | object |
| 604 | .iter_mut() |
| 605 | .map(|(key, child)| { |
| 606 | properties |
| 607 | .and_then(|items| items.get(key)) |
| 608 | .map(|child_schema| normalize_schema_json_containers(child, child_schema)) |
| 609 | .unwrap_or(0) |
| 610 | }) |
| 611 | .sum() |
| 612 | } |
| 613 | serde_json::Value::Array(items) => schema |
| 614 | .get("items") |
| 615 | .map(|item_schema| { |
| 616 | items |
| 617 | .iter_mut() |
| 618 | .map(|item| normalize_schema_json_containers(item, item_schema)) |
| 619 | .sum() |
| 620 | }) |
| 621 | .unwrap_or(0), |
| 622 | _ => 0, |
| 623 | } |
| 624 | } |
| 625 | |
| 626 | fn schema_declares_type(schema: &serde_json::Value, expected: &str) -> bool { |
| 627 | match schema.get("type") { |
| 628 | Some(serde_json::Value::String(value)) => value == expected, |
| 629 | Some(serde_json::Value::Array(values)) => values.iter().any(|value| value == expected), |
| 630 | _ => false, |
| 631 | } |
| 632 | } |
| 633 | |
| 634 | pub(super) fn malformed_tool_arguments_input(buffer: &str) -> serde_json::Value { |
| 635 | json!({ "raw_arguments": buffer }) |
| 636 | } |
| 637 | |
| 638 | pub(super) fn malformed_tool_arguments_error(buffer: &str) -> String { |
| 639 | format!("malformed tool arguments from model: expected valid JSON, got {buffer:?}") |
| 640 | } |
| 641 | |
| 642 | fn strip_code_fences(text: &str) -> Option<String> { |
| 643 | if !text.contains("```") { |
| 644 | return None; |
| 645 | } |
| 646 | let line_count = text.lines().count(); |
| 647 | let mut lines = Vec::with_capacity(line_count); |
| 648 | for line in text.lines() { |
| 649 | if line.trim_start().starts_with("```") { |
| 650 | continue; |
| 651 | } |
| 652 | lines.push(line); |
| 653 | } |
| 654 | let stripped = lines.join("\n"); |
| 655 | let stripped = stripped.trim(); |
| 656 | if stripped.is_empty() { |
| 657 | None |
| 658 | } else { |
| 659 | Some(stripped.to_string()) |
| 660 | } |
| 661 | } |
| 662 | |
| 663 | fn extract_json_segment(text: &str) -> Option<String> { |
| 664 | extract_balanced_segment(text, '{', '}').or_else(|| extract_balanced_segment(text, '[', ']')) |
| 665 | } |
| 666 | |
| 667 | fn extract_balanced_segment(text: &str, open: char, close: char) -> Option<String> { |
| 668 | let start = text.find(open)?; |
| 669 | let mut depth = 0i32; |
| 670 | let mut end = None; |
| 671 | for (offset, ch) in text[start..].char_indices() { |
| 672 | if ch == open { |
| 673 | depth += 1; |
| 674 | } else if ch == close { |
| 675 | depth -= 1; |
| 676 | if depth == 0 { |
| 677 | end = Some(start + offset + ch.len_utf8()); |
| 678 | break; |
| 679 | } |
| 680 | } |
| 681 | } |
| 682 | end.map(|end_idx| text[start..end_idx].to_string()) |
| 683 | } |
| 684 | |
| 685 | fn normalize_parallel_tool_name(raw: &str) -> String { |
| 686 | let mut name = raw.trim(); |
| 687 | for prefix in ["functions.", "tools.", "tool."] { |
| 688 | if let Some(stripped) = name.strip_prefix(prefix) { |
| 689 | name = stripped; |
| 690 | break; |
| 691 | } |
| 692 | } |
| 693 | name.to_string() |
| 694 | } |
| 695 | |
| 696 | pub(super) fn parse_parallel_tool_calls( |
| 697 | input: &serde_json::Value, |
| 698 | ) -> Result<Vec<(String, serde_json::Value)>, ToolError> { |
| 699 | let tool_uses = input |
| 700 | .get("tool_uses") |
| 701 | .and_then(|v| v.as_array()) |
| 702 | .ok_or_else(|| ToolError::missing_field("tool_uses"))?; |
| 703 | if tool_uses.is_empty() { |
| 704 | return Err(ToolError::invalid_input( |
| 705 | "multi_tool_use.parallel requires at least one tool call", |
| 706 | )); |
| 707 | } |
| 708 | |
| 709 | let mut calls = Vec::with_capacity(tool_uses.len()); |
| 710 | for item in tool_uses { |
| 711 | let name = item |
| 712 | .get("recipient_name") |
| 713 | .or_else(|| item.get("tool_name")) |
| 714 | .or_else(|| item.get("name")) |
| 715 | .or_else(|| item.get("tool")) |
| 716 | .and_then(|v| v.as_str()) |
| 717 | .ok_or_else(|| ToolError::missing_field("recipient_name"))?; |
| 718 | let params = item |
| 719 | .get("parameters") |
| 720 | .or_else(|| item.get("input")) |
| 721 | .or_else(|| item.get("args")) |
| 722 | .or_else(|| item.get("arguments")) |
| 723 | .cloned() |
| 724 | .unwrap_or_else(|| json!({})); |
| 725 | calls.push((normalize_parallel_tool_name(name), params)); |
| 726 | } |
| 727 | |
| 728 | Ok(calls) |
| 729 | } |
| 730 | |
| 731 | // === Dispatch policy ================================================== |
| 732 | |
| 733 | #[cfg(test)] |
| 734 | pub(super) fn should_parallelize_tool_batch(plans: &[ToolExecutionPlan]) -> bool { |
| 735 | if plans.is_empty() || !plans.iter().all(tool_plan_can_join_parallel_batch) { |
| 736 | return false; |
| 737 | } |
| 738 | schedule_non_conflicting( |
| 739 | plans |
| 740 | .iter() |
| 741 | .map(|plan| ((), plan.resources.clone())) |
| 742 | .collect(), |
| 743 | ) |
| 744 | .len() |
| 745 | == 1 |
| 746 | } |
| 747 | |
| 748 | pub(super) fn tool_plan_is_parallel_safe(plan: &ToolExecutionPlan) -> bool { |
| 749 | plan.read_only && plan.supports_parallel && !plan.approval_required && !plan.interactive |
| 750 | } |
| 751 | |
| 752 | pub(super) fn tool_plan_can_join_parallel_batch(plan: &ToolExecutionPlan) -> bool { |
| 753 | plan.blocked_error.is_none() |
| 754 | && (tool_plan_is_parallel_safe(plan) |
| 755 | || (plan.detached_start && !plan.approval_required && !plan.interactive)) |
| 756 | } |
| 757 | |
| 758 | pub(super) fn plan_tool_execution_batches( |
| 759 | plans: Vec<ToolExecutionPlan>, |
| 760 | ) -> Vec<ToolExecutionBatch> { |
| 761 | let mut batches = Vec::new(); |
| 762 | let mut parallel_candidates = Vec::new(); |
| 763 | |
| 764 | let flush_parallel = |parallel_candidates: &mut Vec<_>, |
| 765 | batches: &mut Vec<ToolExecutionBatch>| { |
| 766 | for chunk in schedule_non_conflicting(std::mem::take(parallel_candidates)) { |
| 767 | batches.push(ToolExecutionBatch::Parallel(chunk)); |
| 768 | } |
| 769 | }; |
| 770 | |
| 771 | for plan in plans { |
| 772 | if tool_plan_can_join_parallel_batch(&plan) { |
| 773 | let resources = plan.resources.clone(); |
| 774 | parallel_candidates.push((plan, resources)); |
| 775 | continue; |
| 776 | } |
| 777 | |
| 778 | flush_parallel(&mut parallel_candidates, &mut batches); |
| 779 | batches.push(ToolExecutionBatch::Serial(Box::new(plan))); |
| 780 | } |
| 781 | |
| 782 | flush_parallel(&mut parallel_candidates, &mut batches); |
| 783 | |
| 784 | batches |
| 785 | } |
| 786 | |
| 787 | pub(super) fn mcp_tool_is_parallel_safe(name: &str) -> bool { |
| 788 | matches!( |
| 789 | name, |
| 790 | "list_mcp_resources" |
| 791 | | "list_mcp_resource_templates" |
| 792 | | "mcp_read_resource" |
| 793 | | "read_mcp_resource" |
| 794 | | "mcp_get_prompt" |
| 795 | ) |
| 796 | } |
| 797 | |
| 798 | pub(super) fn mcp_tool_is_read_only(name: &str) -> bool { |
| 799 | matches!( |
| 800 | name, |
| 801 | "list_mcp_resources" |
| 802 | | "list_mcp_resource_templates" |
| 803 | | "mcp_read_resource" |
| 804 | | "read_mcp_resource" |
| 805 | | "mcp_get_prompt" |
| 806 | ) |
| 807 | } |
| 808 | |
| 809 | pub(super) fn mcp_tool_approval_description(name: &str) -> String { |
| 810 | if mcp_tool_is_read_only(name) { |
| 811 | format!("Read-only MCP tool '{name}'") |
| 812 | } else { |
| 813 | format!("MCP tool '{name}' may have side effects") |
| 814 | } |
| 815 | } |
| 816 | |
| 817 | #[cfg(test)] |
| 818 | mod schema_json_container_tests { |
| 819 | use super::*; |
| 820 | use crate::tools::spec::ToolSpec; |
| 821 | use serde_json::json; |
| 822 | |
| 823 | #[test] |
| 824 | fn decodes_nested_containers_and_passes_tool_validation() { |
| 825 | let schema = crate::tools::user_input::RequestUserInputTool::default().input_schema(); |
| 826 | let encoded_options = serde_json::to_string(&json!([ |
| 827 | { "label": "Repository", "description": "Inspect the current repository" }, |
| 828 | { "label": "Workspace", "description": "Inspect the whole workspace" } |
| 829 | ])) |
| 830 | .expect("encode options"); |
| 831 | let encoded_questions = serde_json::to_string(&json!([{ |
| 832 | "header": "Scope", |
| 833 | "id": "scope", |
| 834 | "question": "Which scope should be inspected?", |
| 835 | "options": encoded_options |
| 836 | }])) |
| 837 | .expect("encode questions"); |
| 838 | let mut input = json!({ "questions": encoded_questions }); |
| 839 | |
| 840 | assert_eq!(normalize_schema_json_containers(&mut input, &schema), 2); |
| 841 | assert!(input["questions"].is_array()); |
| 842 | assert!(input["questions"][0]["options"].is_array()); |
| 843 | crate::tools::user_input::UserInputRequest::from_value(&input) |
| 844 | .expect("normalized input must still pass tool-specific validation"); |
| 845 | } |
| 846 | |
| 847 | #[test] |
| 848 | fn leaves_primitives_wrong_types_and_unbounded_strings_unchanged() { |
| 849 | let schema = json!({ |
| 850 | "type": "object", |
| 851 | "properties": { |
| 852 | "text": { "type": "string" }, |
| 853 | "count": { "type": "integer" }, |
| 854 | "items": { "type": "array" }, |
| 855 | "oversized": { "type": "array" } |
| 856 | } |
| 857 | }); |
| 858 | let oversized = format!("[\"{}\"]", "x".repeat(MAX_SCHEMA_CONTAINER_REPAIR_BYTES)); |
| 859 | let mut input = json!({ |
| 860 | "text": "[\"still text\"]", |
| 861 | "count": "10", |
| 862 | "items": "{\"wrong\":\"container\"}", |
| 863 | "oversized": oversized |
| 864 | }); |
| 865 | let before = input.clone(); |
| 866 | |
| 867 | assert_eq!(normalize_schema_json_containers(&mut input, &schema), 0); |
| 868 | assert_eq!(input, before); |
| 869 | } |
| 870 | } |
| 871 |