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