vtcode_core/tools/registry/
patch_recovery.rs

1//! One bounded fresh read shared by typed edit/patch mismatches and strict inspection.
2use std::collections::{HashMap, HashSet};
3use std::path::{Component, Path, PathBuf};
4
5use serde_json::Value;
6
7use super::{ToolExecutionError, ToolRegistry};
8use crate::config::constants::tools;
9
10#[derive(Default)]
11pub(super) struct PatchRecoveryReads {
12    canonical_paths_by_alias: HashMap<PathBuf, PathBuf>,
13    consumed_paths: HashSet<PathBuf>,
14}
15
16impl ToolRegistry {
17    /// Start a user/task turn without changing permissions or loop history.
18    pub fn begin_patch_recovery_turn(&self) {
19        *self.patch_recovery_reads.lock() = PatchRecoveryReads::default();
20    }
21
22    fn patch_recovery_key(&self, path: impl AsRef<Path>) -> Option<PathBuf> {
23        let path = self.workspace_root().join(path);
24        let mut normalized = PathBuf::new();
25        for component in path.components() {
26            match component {
27                Component::CurDir => {}
28                Component::ParentDir => {
29                    normalized.pop();
30                }
31                component => normalized.push(component.as_os_str()),
32            }
33        }
34        normalized.starts_with(self.workspace_root()).then_some(normalized)
35    }
36
37    pub(super) async fn grant_patch_recovery_read(&self, error: &ToolExecutionError) {
38        let Some(path) = error.patch_context_mismatch_path().or_else(|| {
39            error
40                .inspection_failure
41                .as_ref()
42                .and_then(|failure| failure.get("path"))
43                .and_then(Value::as_str)
44        }) else {
45            return;
46        };
47        // Repeat containment validation before creating an allowance. Invalid
48        // paths and symlink escapes never acquire a recovery read.
49        let Ok(canonical_path) = self.file_ops_tool().normalize_user_path(path).await else {
50            return;
51        };
52        if let Some(key) = self.patch_recovery_key(path) {
53            let mut reads = self.patch_recovery_reads.lock();
54            if !reads.consumed_paths.contains(&canonical_path) {
55                reads.canonical_paths_by_alias.insert(key, canonical_path.clone());
56                reads.canonical_paths_by_alias.insert(canonical_path.clone(), canonical_path);
57            }
58        }
59    }
60
61    fn patch_recovery_read_key(&self, tool_name: &str, args: &Value) -> Option<PathBuf> {
62        if crate::tools::tool_intent::is_command_run_tool_call(tool_name, args) {
63            let parts = crate::tools::command_args::command_words(args).ok()??;
64            let [command, option, script, path] = parts.as_slice() else {
65                return None;
66            };
67            if command != "sed" || option != "-n" {
68                return None;
69            }
70            let range = script.strip_suffix('p')?;
71            let (start, end) = range.split_once(',').unwrap_or((range, range));
72            let start = start.parse::<u64>().ok()?;
73            let end = end.parse::<u64>().ok()?;
74            if start == 0 || end < start || end.saturating_sub(start) >= 200 {
75                return None;
76            }
77            let directory = self
78                .pty_manager()
79                .working_dir_candidate(crate::tools::command_args::working_dir_text(args))
80                .ok()?;
81            return self.patch_recovery_key(directory.join(path));
82        }
83        if tool_name != tools::READ_FILE
84            && !(tool_name == tools::UNIFIED_FILE
85                && crate::tools::tool_intent::file_operation_action(args) == Some("read"))
86        {
87            return None;
88        }
89        let path = crate::tools::file_ops::bounded_line_read_path(args, 200)?;
90        self.patch_recovery_key(path)
91    }
92
93    /// Check admission without consuming the read before other guards run.
94    pub fn has_patch_recovery_read(&self, tool_name: &str, args: &Value) -> bool {
95        self.pending_patch_recovery_read_path(tool_name, args).is_some()
96    }
97
98    /// Canonical identity used to reserve a capped read before batch execution.
99    pub fn pending_patch_recovery_read_path(&self, tool_name: &str, args: &Value) -> Option<PathBuf> {
100        let key = self.patch_recovery_read_key(tool_name, args)?;
101        let reads = self.patch_recovery_reads.lock();
102        reads
103            .canonical_paths_by_alias
104            .get(&key)
105            .filter(|path| !reads.consumed_paths.contains(*path))
106            .cloned()
107    }
108
109    pub(super) fn consume_patch_recovery_read(&self, tool_name: &str, args: &Value) -> bool {
110        let Some(key) = self.patch_recovery_read_key(tool_name, args) else {
111            return false;
112        };
113        let mut reads = self.patch_recovery_reads.lock();
114        let Some(path) = reads.canonical_paths_by_alias.get(&key).cloned() else {
115            return false;
116        };
117        reads.consumed_paths.insert(path)
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use crate::tools::editing::PatchError;
125    use serde_json::json;
126
127    #[tokio::test]
128    async fn patch_recovery_requires_typed_context_and_cannot_replenish_within_turn() {
129        let temp = tempfile::tempdir().unwrap();
130        let registry = ToolRegistry::new(temp.path().to_path_buf()).await;
131        tokio::fs::write(temp.path().join("sample.txt"), "content\n").await.unwrap();
132        let read = json!({"path":"sample.txt", "limit":10});
133        for source in [
134            anyhow::anyhow!("permission denied"),
135            anyhow::Error::new(PatchError::InvalidFormat("bad".into())),
136            anyhow::Error::new(PatchError::InvalidPath {
137                operation: "update",
138                path: "../outside".into(),
139                reason: "outside workspace".into(),
140            }),
141        ] {
142            let error = ToolExecutionError::from_anyhow(tools::APPLY_PATCH, &source, 0, false, false, None);
143            registry.grant_patch_recovery_read(&error).await;
144            assert!(!registry.has_patch_recovery_read(tools::READ_FILE, &read));
145        }
146        let source = anyhow::Error::new(PatchError::ContextNotFound {
147            path: "sample.txt".into(),
148            context: "missing".into(),
149        })
150        .context("contextual wrapping");
151        let error = ToolExecutionError::from_anyhow(tools::APPLY_PATCH, &source, 0, false, false, None);
152        registry.grant_patch_recovery_read(&error).await;
153        assert!(!registry.has_patch_recovery_read(tools::READ_FILE, &json!({"path":"sample.txt"})));
154        assert!(!registry.has_patch_recovery_read(tools::READ_FILE, &json!({"path":"sample.txt", "limit":201})));
155        assert!(!registry.has_patch_recovery_read(tools::READ_FILE, &json!({"path":"other.txt", "limit":10})));
156        assert!(registry.has_patch_recovery_read(tools::EXEC_COMMAND, &json!({"cmd":"sed -n '2,4p' sample.txt"})));
157        for command in [
158            "sed -n '0,2p' sample.txt",
159            "sed -n '1,201p' sample.txt",
160            "sed -n '4,2p' sample.txt",
161            "sed -n '1,2p' sample.txt && touch ignored",
162        ] {
163            assert!(!registry.has_patch_recovery_read(tools::EXEC_COMMAND, &json!({"cmd":command})));
164        }
165        // The modern handler honors limit and ignores legacy max_lines.
166        assert!(
167            registry
168                .has_patch_recovery_read(tools::READ_FILE, &json!({"path":"sample.txt", "limit":10, "max_lines":1000}))
169        );
170        assert!(registry.has_patch_recovery_read(tools::READ_FILE, &read));
171        assert!(registry.consume_patch_recovery_read(tools::READ_FILE, &read));
172        registry.grant_patch_recovery_read(&error).await;
173        assert!(!registry.consume_patch_recovery_read(tools::READ_FILE, &read));
174        registry.begin_patch_recovery_turn();
175        assert!(!registry.has_patch_recovery_read(tools::READ_FILE, &read));
176        registry.grant_patch_recovery_read(&error).await;
177        assert!(registry.consume_patch_recovery_read(tools::READ_FILE, &read));
178    }
179    #[cfg(unix)]
180    #[tokio::test]
181    async fn patch_recovery_deduplicates_symlink_aliases_and_denies_escapes() {
182        let temp = tempfile::tempdir().unwrap();
183        let outside = tempfile::tempdir().unwrap();
184        tokio::fs::write(temp.path().join("sample.txt"), "content\n").await.unwrap();
185        tokio::fs::write(outside.path().join("outside.txt"), "outside\n").await.unwrap();
186        std::os::unix::fs::symlink(temp.path().join("sample.txt"), temp.path().join("alias.txt")).unwrap();
187        std::os::unix::fs::symlink(outside.path().join("outside.txt"), temp.path().join("escape.txt")).unwrap();
188        let registry = ToolRegistry::new(temp.path().to_path_buf()).await;
189        for path in ["sample.txt", "alias.txt", "escape.txt"] {
190            let source =
191                anyhow::Error::new(PatchError::ContextNotFound { path: path.into(), context: "missing".into() });
192            let error = ToolExecutionError::from_anyhow(tools::APPLY_PATCH, &source, 0, false, false, None);
193            registry.grant_patch_recovery_read(&error).await;
194            let read = json!({"path":path, "limit":10});
195            assert_eq!(registry.consume_patch_recovery_read(tools::READ_FILE, &read), path == "sample.txt");
196        }
197    }
198    #[tokio::test]
199    async fn patch_recovery_checks_the_effective_reader_range() {
200        let temp = tempfile::tempdir().unwrap();
201        let registry = ToolRegistry::new(temp.path().to_path_buf()).await;
202        let content = (1..=400).map(|line| format!("fixture-{line}\n")).collect::<String>();
203        tokio::fs::write(temp.path().join("sample.txt"), content).await.unwrap();
204        let source = anyhow::Error::new(PatchError::ContextNotFound {
205            path: "sample.txt".into(),
206            context: "missing".into(),
207        });
208        let error = ToolExecutionError::from_anyhow(tools::APPLY_PATCH, &source, 0, false, false, None);
209        registry.grant_patch_recovery_read(&error).await;
210        for args in [
211            json!({"path":"sample.txt", "limit_lines":1, "condense":false}),
212            json!({"path":"sample.txt", "max_lines":1}),
213            json!({"path":"sample.txt", "chunk_lines":1}),
214            json!({"path":"sample.txt", "limit":0}),
215            json!({"path":"sample.txt", "limit":201, "l":1}),
216            json!({"path":"sample.txt", "limit":1, "offset_bytes":0, "page_size_bytes":8192}),
217        ] {
218            assert!(!registry.has_patch_recovery_read(tools::READ_FILE, &args), "{args}");
219            assert!(!registry.consume_patch_recovery_read(tools::READ_FILE, &args), "{args}");
220        }
221        let narrow = json!({"path":"sample.txt", "limit":1, "condense":false, "max_tokens":4096});
222        assert!(registry.has_patch_recovery_read(tools::READ_FILE, &narrow));
223        let output = registry.file_ops_tool().read_file(narrow).await.unwrap();
224        let text = output["content"].as_str().unwrap();
225        assert!(text.contains("fixture-1"));
226        assert!(!text.contains("fixture-2"));
227        assert!(
228            crate::tools::file_ops::bounded_line_read_path(&json!({"path":"sample.txt", "limit":1}), 100).is_none()
229        );
230        for args in [
231            json!({"path":"sample.txt", "limit":200, "condense":false}),
232            json!({"path":"sample.txt", "limit":"200", "condense":false}),
233            json!({"path":"sample.txt", "l":200, "condense":false}),
234            json!({"path":"sample.txt", "page_size_lines":200, "condense":false}),
235            json!({"path":"sample.txt", "start_line":1, "end_line":200, "condense":false}),
236        ] {
237            assert!(registry.has_patch_recovery_read(tools::READ_FILE, &args), "{args}");
238            let output = registry.file_ops_tool().read_file(args.clone()).await.unwrap();
239            let text = output["content"].as_str().unwrap();
240            assert!(text.contains("fixture-200"), "{args}: {output}");
241            assert!(!text.contains("fixture-201"), "{args}: {output}");
242            assert!(!text.contains("fixture-400"), "{args}: {output}");
243        }
244    }
245
246    #[tokio::test]
247    async fn patch_recovery_resolves_shell_targets_in_the_execution_directory() {
248        let temp = tempfile::tempdir().unwrap();
249        tokio::fs::create_dir(temp.path().join("sub")).await.unwrap();
250        tokio::fs::write(temp.path().join("sample.txt"), "root-only\n").await.unwrap();
251        tokio::fs::write(temp.path().join("sub/sample.txt"), "nested-only\n")
252            .await
253            .unwrap();
254        let registry = ToolRegistry::new(temp.path().to_path_buf()).await;
255        registry.allow_all_tools().await.unwrap();
256        for path in ["sample.txt", "sub/sample.txt"] {
257            registry.begin_patch_recovery_turn();
258            let source =
259                anyhow::Error::new(PatchError::ContextNotFound { path: path.into(), context: "missing".into() });
260            let error = ToolExecutionError::from_anyhow(tools::APPLY_PATCH, &source, 0, false, false, None);
261            registry.grant_patch_recovery_read(&error).await;
262            for field in ["workdir", "cwd", "working_dir"] {
263                let args = json!({"cmd":"sed -n '1,2p' sample.txt", field:"sub"});
264                assert_eq!(
265                    registry.has_patch_recovery_read(tools::EXEC_COMMAND, &args),
266                    path == "sub/sample.txt",
267                    "{path}: {args}"
268                );
269            }
270            let args = json!({"cmd":"sed -n '1,2p' sample.txt", "workdir":"sub"});
271            if path == "sample.txt" {
272                assert!(!registry.consume_patch_recovery_read(tools::EXEC_COMMAND, &args));
273                assert!(
274                    registry.has_patch_recovery_read(tools::EXEC_COMMAND, &json!({"cmd":"sed -n '1,2p' sample.txt"}))
275                );
276            } else {
277                let output = registry.execute_tool(tools::EXEC_COMMAND, args.clone()).await.unwrap();
278                let text = output.to_string();
279                assert!(text.contains("nested-only"), "{output}");
280                assert!(!text.contains("root-only"), "{output}");
281                assert!(!registry.has_patch_recovery_read(tools::EXEC_COMMAND, &args));
282            }
283        }
284    }
285}
vtcode_core/tools/registry/patch_recovery.rs

vtcode_core/tools/registry/
patch_recovery.rs
