返回 CodeWhale
tool_preparation.rs
根目录 / crates / tui / src / core / engine / tool_preparation.rs
1 //! Side-effect-free preparation of concrete tool inputs.
2 //!
3 //! The turn loop remains the authority orchestrator. This module only makes
4 //! the input-specific policy decision inspectable and reusable, including a
5 //! mandatory second preparation after a hook rewrites input.
6
7 use std::collections::BTreeSet;
8 use std::path::PathBuf;
9
10 use serde_json::Value;
11
12 use crate::mcp::McpPool;
13 use crate::tools::ToolRegistry;
14 use crate::tools::spec::{ApprovalRequirement, PreparedToolCall, ResourceClaim, ToolError};
15
16 use super::dispatch::{
17 mcp_tool_approval_description, mcp_tool_is_parallel_safe, mcp_tool_is_read_only,
18 };
19 use super::tool_catalog::{CODE_EXECUTION_TOOL_NAME, JS_EXECUTION_TOOL_NAME, is_tool_search_tool};
20
21 #[derive(Debug, Clone, PartialEq)]
22 pub(super) struct PreparedToolPolicy {
23 pub(super) call: PreparedToolCall,
24 pub(super) auto_approve: bool,
25 }
26
27 /// Prepare a concrete call without mutating external state.
28 pub(super) fn prepare_tool_call(
29 name: &str,
30 input: Value,
31 registry: Option<&ToolRegistry>,
32 session_auto_approve: bool,
33 ) -> Result<PreparedToolPolicy, ToolError> {
34 if McpPool::is_mcp_tool(name) {
35 let read_only = mcp_tool_is_read_only(name);
36 if !read_only
37 && let Some(authority) =
38 registry.and_then(|registry| registry.context().tool_authority.as_ref())
39 {
40 return Err(ToolError::permission_denied(format!(
41 "worker '{}' cannot run mutating MCP tool {name}: it has no bounded file target under the machine-readable authority envelope",
42 authority.owner
43 )));
44 }
45 return Ok(PreparedToolPolicy {
46 call: PreparedToolCall {
47 name: name.to_string(),
48 input,
49 description: mcp_tool_approval_description(name),
50 read_only,
51 supports_parallel: mcp_tool_is_parallel_safe(name),
52 starts_detached: false,
53 approval: if read_only {
54 ApprovalRequirement::Auto
55 } else {
56 ApprovalRequirement::Suggest
57 },
58 resources: vec![ResourceClaim::GlobalExclusive],
59 },
60 auto_approve: session_auto_approve,
61 });
62 }
63
64 if let Some(registry) = registry
65 && let Some(spec) = registry.get(name)
66 {
67 let mut call = spec.prepare(input, registry.context())?;
68 call.resources = registered_resource_claims(name, &call.input, registry.context())?;
69 return Ok(PreparedToolPolicy {
70 call,
71 auto_approve: registry.context().auto_approve,
72 });
73 }
74
75 if name == CODE_EXECUTION_TOOL_NAME {
76 reject_unbounded_execution_under_authority(name, registry)?;
77 return Ok(conservative_execution_policy(
78 name,
79 input,
80 "Run model-provided Python code in local execution sandbox",
81 session_auto_approve,
82 ));
83 }
84
85 if name == JS_EXECUTION_TOOL_NAME {
86 reject_unbounded_execution_under_authority(name, registry)?;
87 return Ok(conservative_execution_policy(
88 name,
89 input,
90 "Run model-provided JavaScript code in local Node.js execution sandbox",
91 session_auto_approve,
92 ));
93 }
94
95 if is_tool_search_tool(name) {
96 return Ok(PreparedToolPolicy {
97 call: PreparedToolCall {
98 name: name.to_string(),
99 input,
100 description: "Search tool catalog".to_string(),
101 read_only: true,
102 supports_parallel: false,
103 starts_detached: false,
104 approval: ApprovalRequirement::Auto,
105 resources: Vec::new(),
106 },
107 auto_approve: session_auto_approve,
108 });
109 }
110
111 Err(ToolError::not_available(format!(
112 "tool '{name}' has no preparation path"
113 )))
114 }
115
116 fn reject_unbounded_execution_under_authority(
117 name: &str,
118 registry: Option<&ToolRegistry>,
119 ) -> Result<(), ToolError> {
120 let Some(authority) = registry.and_then(|registry| registry.context().tool_authority.as_ref())
121 else {
122 return Ok(());
123 };
124 Err(ToolError::permission_denied(format!(
125 "worker '{}' cannot run {name}: arbitrary code execution cannot prove a bounded file target under the machine-readable authority envelope",
126 authority.owner
127 )))
128 }
129
130 /// Re-run preparation from the rewritten input rather than patching any
131 /// previously derived field.
132 pub(super) fn reprepare_tool_call_after_hook(
133 name: &str,
134 updated_input: Value,
135 registry: Option<&ToolRegistry>,
136 session_auto_approve: bool,
137 ) -> Result<PreparedToolPolicy, ToolError> {
138 prepare_tool_call(name, updated_input, registry, session_auto_approve)
139 }
140
141 fn conservative_execution_policy(
142 name: &str,
143 input: Value,
144 description: &str,
145 auto_approve: bool,
146 ) -> PreparedToolPolicy {
147 PreparedToolPolicy {
148 call: PreparedToolCall {
149 name: name.to_string(),
150 input,
151 description: description.to_string(),
152 read_only: false,
153 supports_parallel: false,
154 starts_detached: false,
155 approval: ApprovalRequirement::Suggest,
156 resources: vec![ResourceClaim::GlobalExclusive],
157 },
158 auto_approve,
159 }
160 }
161
162 fn registered_resource_claims(
163 name: &str,
164 input: &Value,
165 context: &crate::tools::ToolContext,
166 ) -> Result<Vec<ResourceClaim>, ToolError> {
167 match name {
168 "read_file" => path_claim(input, "path", None, context, ResourceClaim::ReadPath),
169 "write_file" | "edit_file" => {
170 path_claim(input, "path", None, context, ResourceClaim::WritePath)
171 }
172 "list_dir" | "grep_files" | "file_search" => {
173 path_claim(input, "path", Some("."), context, ResourceClaim::ReadTree)
174 }
175 "apply_patch" => apply_patch_resource_claims(input, context),
176 "terminal/run" => Ok(terminal_claim(input, "session", Some("term-1"))),
177 "terminal/send" | "terminal/wait" | "terminal/cancel" | "terminal/reset" => {
178 Ok(terminal_claim(input, "session", None))
179 }
180 "exec_shell_wait"
181 | "exec_wait"
182 | "exec_shell_interact"
183 | "exec_interact"
184 | "exec_shell_cancel" => Ok(terminal_claim(input, "task_id", None)),
185 _ => Ok(global_exclusive_claim()),
186 }
187 }
188
189 fn path_claim(
190 input: &Value,
191 key: &str,
192 default: Option<&str>,
193 context: &crate::tools::ToolContext,
194 build: fn(PathBuf) -> ResourceClaim,
195 ) -> Result<Vec<ResourceClaim>, ToolError> {
196 let raw = input
197 .get(key)
198 .and_then(Value::as_str)
199 .map(str::trim)
200 .filter(|path| !path.is_empty())
201 .or(default);
202 let Some(raw) = raw else {
203 return Ok(global_exclusive_claim());
204 };
205 Ok(context
206 .resolve_path(raw)
207 .map_or_else(|_| global_exclusive_claim(), |path| vec![build(path)]))
208 }
209
210 fn apply_patch_resource_claims(
211 input: &Value,
212 context: &crate::tools::ToolContext,
213 ) -> Result<Vec<ResourceClaim>, ToolError> {
214 let Ok(preflight) = crate::tools::apply_patch::preflight_apply_patch(input) else {
215 return Ok(global_exclusive_claim());
216 };
217 if preflight.touched_files.is_empty() {
218 return Ok(global_exclusive_claim());
219 }
220
221 let mut claims = BTreeSet::new();
222 for path in preflight.touched_files {
223 let Ok(path) = context.resolve_path(&path) else {
224 return Ok(global_exclusive_claim());
225 };
226 claims.insert(ResourceClaim::WritePath(path));
227 }
228 Ok(claims.into_iter().collect())
229 }
230
231 fn terminal_claim(input: &Value, key: &str, default: Option<&str>) -> Vec<ResourceClaim> {
232 input
233 .get(key)
234 .and_then(Value::as_str)
235 .map(str::trim)
236 .filter(|id| !id.is_empty())
237 .or(default)
238 .map_or_else(global_exclusive_claim, |id| {
239 vec![ResourceClaim::Terminal(id.to_string())]
240 })
241 }
242
243 fn global_exclusive_claim() -> Vec<ResourceClaim> {
244 vec![ResourceClaim::GlobalExclusive]
245 }
246
247 #[cfg(test)]
248 mod tests {
249 use std::sync::Arc;
250
251 use async_trait::async_trait;
252 use serde_json::json;
253 use tempfile::tempdir;
254
255 use crate::tools::spec::{ToolCapability, ToolContext, ToolResult, ToolSpec};
256
257 use super::*;
258
259 struct InputDependentTool;
260
261 #[async_trait]
262 impl ToolSpec for InputDependentTool {
263 fn name(&self) -> &str {
264 "input_dependent"
265 }
266
267 fn description(&self) -> &str {
268 "characterization tool"
269 }
270
271 fn input_schema(&self) -> Value {
272 json!({"type": "object"})
273 }
274
275 fn capabilities(&self) -> Vec<ToolCapability> {
276 vec![ToolCapability::WritesFiles]
277 }
278
279 fn approval_requirement_for(&self, input: &Value) -> ApprovalRequirement {
280 if input.get("safe").and_then(Value::as_bool) == Some(true) {
281 ApprovalRequirement::Auto
282 } else {
283 ApprovalRequirement::Required
284 }
285 }
286
287 fn is_read_only_for(&self, input: &Value) -> bool {
288 input.get("safe").and_then(Value::as_bool) == Some(true)
289 }
290
291 fn supports_parallel_for(&self, input: &Value) -> bool {
292 self.is_read_only_for(input)
293 }
294
295 fn starts_detached_for(&self, input: &Value) -> bool {
296 input.get("detached").and_then(Value::as_bool) == Some(true)
297 }
298
299 async fn execute(
300 &self,
301 _input: Value,
302 _context: &ToolContext,
303 ) -> Result<ToolResult, ToolError> {
304 unreachable!("preparation must not execute the tool")
305 }
306 }
307
308 fn registry() -> (tempfile::TempDir, ToolRegistry) {
309 let root = tempdir().expect("tempdir");
310 let mut context = ToolContext::new(root.path().to_path_buf());
311 context.auto_approve = true;
312 let mut registry = ToolRegistry::new(context);
313 registry.register(Arc::new(InputDependentTool));
314 (root, registry)
315 }
316
317 #[test]
318 fn prepared_policy_matches_existing_input_specific_decisions() {
319 let (_root, registry) = registry();
320 let spec = registry.get("input_dependent").expect("registered tool");
321
322 for input in [
323 json!({"safe": true, "detached": false}),
324 json!({"safe": false, "detached": true}),
325 ] {
326 let prepared =
327 prepare_tool_call("input_dependent", input.clone(), Some(&registry), false)
328 .expect("prepare");
329
330 assert_eq!(
331 prepared.call.approval,
332 spec.approval_requirement_for(&input)
333 );
334 assert_eq!(prepared.call.read_only, spec.is_read_only_for(&input));
335 assert_eq!(
336 prepared.call.supports_parallel,
337 spec.supports_parallel_for(&input)
338 );
339 assert_eq!(
340 prepared.call.starts_detached,
341 spec.starts_detached_for(&input)
342 );
343 assert!(prepared.auto_approve);
344 }
345 }
346
347 #[test]
348 fn hook_rewrite_discards_every_original_prepared_decision() {
349 let (_root, registry) = registry();
350 let original = prepare_tool_call(
351 "input_dependent",
352 json!({"safe": true, "detached": false}),
353 Some(&registry),
354 false,
355 )
356 .expect("prepare original");
357 let rewritten = reprepare_tool_call_after_hook(
358 "input_dependent",
359 json!({"safe": false, "detached": true}),
360 Some(&registry),
361 false,
362 )
363 .expect("reprepare rewritten input");
364
365 assert_eq!(original.call.approval, ApprovalRequirement::Auto);
366 assert!(original.call.read_only);
367 assert!(original.call.supports_parallel);
368 assert!(!original.call.starts_detached);
369
370 assert_eq!(rewritten.call.approval, ApprovalRequirement::Required);
371 assert!(!rewritten.call.read_only);
372 assert!(!rewritten.call.supports_parallel);
373 assert!(rewritten.call.starts_detached);
374 assert_eq!(
375 rewritten.call.input,
376 json!({"safe": false, "detached": true})
377 );
378 }
379
380 #[test]
381 fn bypass_preparation_preserves_legacy_policy_table() {
382 struct Expected {
383 name: &'static str,
384 approval: ApprovalRequirement,
385 read_only: bool,
386 supports_parallel: bool,
387 global_exclusive: bool,
388 }
389
390 for expected in [
391 Expected {
392 name: "read_mcp_resource",
393 approval: ApprovalRequirement::Auto,
394 read_only: true,
395 supports_parallel: true,
396 global_exclusive: true,
397 },
398 Expected {
399 name: "mcp_filesystem_write",
400 approval: ApprovalRequirement::Suggest,
401 read_only: false,
402 supports_parallel: false,
403 global_exclusive: true,
404 },
405 Expected {
406 name: CODE_EXECUTION_TOOL_NAME,
407 approval: ApprovalRequirement::Suggest,
408 read_only: false,
409 supports_parallel: false,
410 global_exclusive: true,
411 },
412 Expected {
413 name: JS_EXECUTION_TOOL_NAME,
414 approval: ApprovalRequirement::Suggest,
415 read_only: false,
416 supports_parallel: false,
417 global_exclusive: true,
418 },
419 Expected {
420 name: "tool_search",
421 approval: ApprovalRequirement::Auto,
422 read_only: true,
423 supports_parallel: false,
424 global_exclusive: false,
425 },
426 ] {
427 let prepared = prepare_tool_call(expected.name, json!({}), None, false)
428 .unwrap_or_else(|error| panic!("prepare {}: {error}", expected.name));
429 assert_eq!(
430 prepared.call.approval, expected.approval,
431 "{}",
432 expected.name
433 );
434 assert_eq!(
435 prepared.call.read_only, expected.read_only,
436 "{}",
437 expected.name
438 );
439 assert_eq!(
440 prepared.call.supports_parallel, expected.supports_parallel,
441 "{}",
442 expected.name
443 );
444 assert_eq!(
445 prepared.call.resources == vec![ResourceClaim::GlobalExclusive],
446 expected.global_exclusive,
447 "{}",
448 expected.name
449 );
450 assert!(!prepared.call.starts_detached, "{}", expected.name);
451 assert!(!prepared.auto_approve, "{}", expected.name);
452 }
453 }
454
455 #[test]
456 fn mcp_write_preparation_respects_session_auto_approval() {
457 let prepared = prepare_tool_call("mcp_filesystem_write", json!({}), None, true)
458 .expect("prepare MCP write tool with session auto-approval");
459
460 assert_eq!(prepared.call.approval, ApprovalRequirement::Suggest);
461 assert!(!prepared.call.read_only);
462 assert!(!prepared.call.supports_parallel);
463 assert_eq!(
464 prepared.call.resources,
465 vec![ResourceClaim::GlobalExclusive]
466 );
467 assert!(prepared.auto_approve);
468 assert!(!super::super::turn_loop::registered_tool_approval_required(
469 &prepared.call.name,
470 prepared.call.approval,
471 prepared.auto_approve,
472 ));
473 }
474
475 #[test]
476 fn hook_rewrite_reprepares_resource_claims_from_final_input() {
477 let root = tempdir().expect("tempdir");
478 let context = ToolContext::new(root.path().to_path_buf());
479 let original_path = context.resolve_path("before.rs").expect("original path");
480 let rewritten_path = context.resolve_path("after.rs").expect("rewritten path");
481 let mut registry = ToolRegistry::new(context);
482 registry.register(Arc::new(crate::tools::file::ReadFileTool));
483
484 let original = prepare_tool_call(
485 "read_file",
486 json!({"path": "before.rs"}),
487 Some(&registry),
488 false,
489 )
490 .expect("prepare original read");
491 let rewritten = reprepare_tool_call_after_hook(
492 "read_file",
493 json!({"path": "after.rs"}),
494 Some(&registry),
495 false,
496 )
497 .expect("reprepare rewritten read");
498
499 assert_eq!(
500 original.call.resources,
501 vec![ResourceClaim::ReadPath(original_path)]
502 );
503 assert_eq!(
504 rewritten.call.resources,
505 vec![ResourceClaim::ReadPath(rewritten_path)]
506 );
507 }
508
509 #[test]
510 fn registered_file_claims_are_canonical_and_input_specific() {
511 let root = tempdir().expect("tempdir");
512 let context = ToolContext::new(root.path().to_path_buf());
513 let exact = context.resolve_path("src/lib.rs").expect("exact path");
514 let tree = context.resolve_path("src").expect("tree path");
515
516 assert_eq!(
517 registered_resource_claims("read_file", &json!({"path": "src/lib.rs"}), &context,)
518 .expect("read claim"),
519 vec![ResourceClaim::ReadPath(exact.clone())]
520 );
521 assert_eq!(
522 registered_resource_claims("edit_file", &json!({"path": "src/lib.rs"}), &context,)
523 .expect("write claim"),
524 vec![ResourceClaim::WritePath(exact)]
525 );
526 assert_eq!(
527 registered_resource_claims("grep_files", &json!({"path": "src"}), &context)
528 .expect("tree claim"),
529 vec![ResourceClaim::ReadTree(tree)]
530 );
531 assert_eq!(
532 registered_resource_claims("read_file", &json!({"path": "../../outside"}), &context,)
533 .expect("path escape must fall back conservatively"),
534 vec![ResourceClaim::GlobalExclusive]
535 );
536 }
537
538 #[cfg(unix)]
539 #[test]
540 fn symlink_aliases_resolve_to_the_same_file_claim() {
541 use std::os::unix::fs::symlink;
542
543 let root = tempdir().expect("tempdir");
544 let real_dir = root.path().join("real");
545 std::fs::create_dir(&real_dir).expect("create real directory");
546 std::fs::write(real_dir.join("lib.rs"), "fn main() {}\n").expect("write real file");
547 symlink("real", root.path().join("alias")).expect("create directory symlink");
548 let context = ToolContext::new(root.path().to_path_buf());
549 let canonical_real = real_dir
550 .join("lib.rs")
551 .canonicalize()
552 .expect("canonical file");
553
554 let read_alias =
555 registered_resource_claims("read_file", &json!({"path": "alias/lib.rs"}), &context)
556 .expect("alias claim");
557 let write_real =
558 registered_resource_claims("write_file", &json!({"path": "real/lib.rs"}), &context)
559 .expect("real claim");
560
561 assert_eq!(read_alias, vec![ResourceClaim::ReadPath(canonical_real)]);
562 assert!(read_alias[0].conflicts_with(&write_real[0]));
563 }
564
565 #[test]
566 fn apply_patch_claims_every_resolved_target_or_falls_back_global() {
567 let root = tempdir().expect("tempdir");
568 let context = ToolContext::new(root.path().to_path_buf());
569 let a = context.resolve_path("a.rs").expect("a path");
570 let b = context.resolve_path("b.rs").expect("b path");
571
572 let claims = registered_resource_claims(
573 "apply_patch",
574 &json!({
575 "replace": [
576 {"path": "b.rs", "content": "b"},
577 {"path": "a.rs", "content": "a"}
578 ]
579 }),
580 &context,
581 )
582 .expect("patch claims");
583 assert_eq!(
584 claims,
585 vec![ResourceClaim::WritePath(a), ResourceClaim::WritePath(b)]
586 );
587
588 assert_eq!(
589 registered_resource_claims(
590 "apply_patch",
591 &json!({"patch": "not a unified diff"}),
592 &context,
593 )
594 .expect("fallback claim"),
595 vec![ResourceClaim::GlobalExclusive]
596 );
597 assert_eq!(
598 registered_resource_claims(
599 "apply_patch",
600 &json!({"replace": [{"path": "../../outside", "content": "nope"}]}),
601 &context,
602 )
603 .expect("escaped target fallback"),
604 vec![ResourceClaim::GlobalExclusive]
605 );
606 }
607
608 #[test]
609 fn terminal_and_unknown_tools_keep_conservative_claims() {
610 let root = tempdir().expect("tempdir");
611 let context = ToolContext::new(root.path().to_path_buf());
612
613 assert_eq!(
614 registered_resource_claims("terminal/run", &json!({}), &context)
615 .expect("default terminal"),
616 vec![ResourceClaim::Terminal("term-1".to_string())]
617 );
618 assert_eq!(
619 registered_resource_claims(
620 "exec_shell_interact",
621 &json!({"task_id": "task-7"}),
622 &context,
623 )
624 .expect("task terminal"),
625 vec![ResourceClaim::Terminal("task-7".to_string())]
626 );
627 assert_eq!(
628 registered_resource_claims("plugin_tool", &json!({}), &context).expect("unknown tool"),
629 vec![ResourceClaim::GlobalExclusive]
630 );
631 }
632 }
633
633 lines RUST