Skip to main content

soma_codemode/state/
workspace_meta.rs

1use regex::Regex;
2use serde::Serialize;
3
4use crate::path_safety::{reject_existing_symlink_ancestors, rel_to_unix_string};
5use crate::ToolError;
6
7use super::path::{is_reserved_metadata_path, VirtualPath};
8use super::workspace::{
9    internal_io, not_found_or_internal, serialize_error, ExistsResult, GlobResult, ListResult,
10    MutationResult, ReplaceInFilesResult, SearchFilesResult, SearchMatch, StatResult,
11    StateWorkspace, WalkEntry, WalkTreeResult,
12};
13
14impl StateWorkspace {
15    pub async fn exists(&self, path: &VirtualPath) -> Result<ExistsResult, ToolError> {
16        let destination = self.resolve(path);
17        self.ensure_path_allowed(&destination).await?;
18        let exists = match tokio::fs::metadata(&destination).await {
19            Ok(_) => true,
20            Err(err) if err.kind() == std::io::ErrorKind::NotFound => false,
21            Err(err) => return Err(internal_io("read state path metadata")(err)),
22        };
23        Ok(ExistsResult {
24            path: path.as_str().to_string(),
25            exists,
26        })
27    }
28
29    pub async fn stat(&self, path: &VirtualPath) -> Result<StatResult, ToolError> {
30        let destination = self.resolve(path);
31        self.ensure_path_allowed(&destination).await?;
32        let metadata = tokio::fs::metadata(&destination)
33            .await
34            .map_err(not_found_or_internal("read state path metadata"))?;
35        let kind = if metadata.is_file() {
36            "file"
37        } else if metadata.is_dir() {
38            "directory"
39        } else {
40            return Err(ToolError::Sdk {
41                sdk_kind: "permission_denied".to_string(),
42                message: "state path kind is not supported".to_string(),
43            });
44        };
45        Ok(StatResult {
46            path: path.as_str().to_string(),
47            kind: kind.to_string(),
48            bytes: metadata.len(),
49        })
50    }
51
52    pub async fn mkdir(&self, path: &VirtualPath) -> Result<MutationResult, ToolError> {
53        let destination = self.resolve(path);
54        self.ensure_path_allowed(&destination).await?;
55        self.check_entry_quota_for_path(&destination).await?;
56        tokio::fs::create_dir_all(&destination)
57            .await
58            .map_err(internal_io("create state directory"))?;
59        self.ensure_path_allowed(&destination).await?;
60        Ok(MutationResult {
61            ok: true,
62            path: path.as_str().to_string(),
63        })
64    }
65
66    pub async fn remove(
67        &self,
68        path: &VirtualPath,
69        recursive: bool,
70    ) -> Result<MutationResult, ToolError> {
71        if is_reserved_metadata_path(path.as_str()) {
72            return Err(ToolError::Sdk {
73                sdk_kind: "permission_denied".to_string(),
74                message: "state metadata paths cannot be removed".to_string(),
75            });
76        }
77        let destination = self.resolve(path);
78        self.ensure_path_allowed(&destination).await?;
79        let metadata = tokio::fs::metadata(&destination)
80            .await
81            .map_err(not_found_or_internal("read state path metadata"))?;
82        if metadata.is_file() {
83            tokio::fs::remove_file(&destination)
84                .await
85                .map_err(internal_io("remove state file"))?;
86        } else if metadata.is_dir() {
87            if recursive {
88                tokio::fs::remove_dir_all(&destination)
89                    .await
90                    .map_err(internal_io("remove state directory tree"))?;
91            } else {
92                tokio::fs::remove_dir(&destination)
93                    .await
94                    .map_err(internal_io("remove state directory"))?;
95            }
96        } else {
97            return Err(ToolError::Sdk {
98                sdk_kind: "permission_denied".to_string(),
99                message: "state path kind is not supported".to_string(),
100            });
101        }
102        Ok(MutationResult {
103            ok: true,
104            path: path.as_str().to_string(),
105        })
106    }
107
108    pub async fn copy(
109        &self,
110        from: &VirtualPath,
111        to: &VirtualPath,
112    ) -> Result<MutationResult, ToolError> {
113        let source = self.read_file(from).await?;
114        self.write_file(to, &source.content).await?;
115        Ok(MutationResult {
116            ok: true,
117            path: to.as_str().to_string(),
118        })
119    }
120
121    pub async fn move_path(
122        &self,
123        from: &VirtualPath,
124        to: &VirtualPath,
125    ) -> Result<MutationResult, ToolError> {
126        let source = self.resolve(from);
127        let destination = self.resolve(to);
128        self.ensure_path_allowed(&source).await?;
129        self.ensure_path_allowed(&destination).await?;
130        self.check_entry_quota_for_path(&destination).await?;
131        if let Some(parent) = destination.parent() {
132            tokio::fs::create_dir_all(parent)
133                .await
134                .map_err(internal_io("create state move directory"))?;
135        }
136        reject_existing_symlink_ancestors(self.root(), &destination)?;
137        tokio::fs::rename(&source, &destination)
138            .await
139            .map_err(not_found_or_internal("move state path"))?;
140        Ok(MutationResult {
141            ok: true,
142            path: to.as_str().to_string(),
143        })
144    }
145
146    pub async fn walk_tree(
147        &self,
148        path: &VirtualPath,
149        limit: usize,
150    ) -> Result<WalkTreeResult, ToolError> {
151        let limit = normalize_limit(limit);
152        let start = self.resolve(path);
153        self.ensure_path_allowed(&start).await?;
154        let mut entries = Vec::new();
155        let mut stack = vec![start];
156        while let Some(dir) = stack.pop() {
157            let mut read_dir = tokio::fs::read_dir(&dir)
158                .await
159                .map_err(not_found_or_internal("read state directory"))?;
160            while let Some(entry) = read_dir
161                .next_entry()
162                .await
163                .map_err(internal_io("read state directory entry"))?
164            {
165                let path = entry.path();
166                let relative = match path.strip_prefix(self.root()) {
167                    Ok(relative) => relative,
168                    Err(_) => continue,
169                };
170                let virtual_path = rel_to_unix_string(relative);
171                if is_reserved_metadata_path(&virtual_path) {
172                    continue;
173                }
174                let metadata = tokio::fs::symlink_metadata(&path)
175                    .await
176                    .map_err(internal_io("read state workspace metadata"))?;
177                if metadata.file_type().is_symlink() {
178                    return Err(ToolError::Sdk {
179                        sdk_kind: "symlink_rejected".to_string(),
180                        message: "state walk rejected a symlink".to_string(),
181                    });
182                }
183                let kind = if metadata.is_dir() {
184                    stack.push(path);
185                    "directory"
186                } else if metadata.is_file() {
187                    "file"
188                } else {
189                    return Err(ToolError::Sdk {
190                        sdk_kind: "permission_denied".to_string(),
191                        message: "state path kind is not supported".to_string(),
192                    });
193                };
194                entries.push(WalkEntry {
195                    path: virtual_path,
196                    kind: kind.to_string(),
197                    bytes: metadata.len(),
198                });
199                if entries.len() > limit {
200                    entries.sort_by(|left, right| left.path.cmp(&right.path));
201                    entries.truncate(limit);
202                    return Ok(WalkTreeResult {
203                        entries,
204                        truncated: true,
205                    });
206                }
207            }
208        }
209        entries.sort_by(|left, right| left.path.cmp(&right.path));
210        Ok(WalkTreeResult {
211            entries,
212            truncated: false,
213        })
214    }
215
216    pub async fn list(&self, path: &VirtualPath) -> Result<ListResult, ToolError> {
217        let dir = self.resolve(path);
218        self.ensure_path_allowed(&dir).await?;
219        let mut read_dir = tokio::fs::read_dir(&dir)
220            .await
221            .map_err(not_found_or_internal("read state directory"))?;
222        let mut entries = Vec::new();
223        while let Some(entry) = read_dir
224            .next_entry()
225            .await
226            .map_err(internal_io("read state directory entry"))?
227        {
228            let name = entry.file_name().to_string_lossy().to_string();
229            let child_path = if path.as_str().is_empty() {
230                name.clone()
231            } else {
232                format!("{}/{}", path.as_str(), name)
233            };
234            if is_reserved_metadata_path(&child_path) {
235                continue;
236            }
237            entries.push(name);
238            if entries.len() as u64 > self.limits().max_entries {
239                return Err(ToolError::Sdk {
240                    sdk_kind: "response_too_large".to_string(),
241                    message: "state list exceeded max entries".to_string(),
242                });
243            }
244        }
245        entries.sort();
246        Ok(ListResult { entries })
247    }
248
249    pub async fn glob(&self, pattern: &str, limit: usize) -> Result<GlobResult, ToolError> {
250        let limit = normalize_limit(limit);
251        let matcher = glob_pattern_regex(pattern)?;
252        let walked = self.walk_files(self.limits().max_entries as usize).await?;
253        let mut matches = Vec::new();
254        for file in walked.files {
255            if matcher.is_match(&file) {
256                matches.push(file);
257                if matches.len() > limit {
258                    matches.truncate(limit);
259                    return Ok(GlobResult {
260                        matches,
261                        truncated: true,
262                    });
263                }
264            }
265        }
266        Ok(GlobResult {
267            matches,
268            truncated: walked.truncated,
269        })
270    }
271
272    pub async fn search_files(
273        &self,
274        pattern: &str,
275        query: &str,
276        limit: usize,
277    ) -> Result<SearchFilesResult, ToolError> {
278        if query.is_empty() {
279            return Err(ToolError::InvalidParam {
280                message: "state search query must not be empty".to_string(),
281                param: "query".to_string(),
282            });
283        }
284        let limit = normalize_limit(limit);
285        let glob = self
286            .glob(pattern, self.limits().max_entries as usize)
287            .await?;
288        let mut matches = Vec::new();
289        for path in glob.matches {
290            let virtual_path = VirtualPath::parse(&path)?;
291            let file = self.read_file(&virtual_path).await?;
292            for (index, line) in file.content.lines().enumerate() {
293                if line.contains(query) {
294                    matches.push(SearchMatch {
295                        path: path.clone(),
296                        line: index + 1,
297                        text: cap_line_preview(line),
298                    });
299                    if matches.len() > limit {
300                        matches.truncate(limit);
301                        return Ok(SearchFilesResult {
302                            matches,
303                            truncated: true,
304                        });
305                    }
306                    ensure_serialized_result_fits(&matches, self.limits().max_result_bytes)?;
307                }
308            }
309        }
310        Ok(SearchFilesResult {
311            matches,
312            truncated: glob.truncated,
313        })
314    }
315
316    pub async fn replace_in_files(
317        &self,
318        pattern: &str,
319        search: &str,
320        replace: &str,
321        dry_run: bool,
322    ) -> Result<ReplaceInFilesResult, ToolError> {
323        if search.is_empty() {
324            return Err(ToolError::InvalidParam {
325                message: "state replace search must not be empty".to_string(),
326                param: "search".to_string(),
327            });
328        }
329        let changed_paths = self.replace_targets(pattern, search).await?;
330        let changed = changed_paths
331            .iter()
332            .map(|path| path.as_str().to_string())
333            .collect::<Vec<_>>();
334        if dry_run {
335            return Ok(ReplaceInFilesResult { changed, dry_run });
336        }
337        let mut originals = Vec::new();
338        for path in changed_paths {
339            let file = self.read_file(&path).await?;
340            if !file.content.contains(search) {
341                return Err(self
342                    .restore_originals_after_failure(
343                        &originals,
344                        ToolError::Sdk {
345                            sdk_kind: "edit_conflict".to_string(),
346                            message: format!(
347                                "state replace input no longer matches `{}`",
348                                path.as_str()
349                            ),
350                        },
351                    )
352                    .await);
353            }
354            let original = file.content;
355            let next = original.replace(search, replace);
356            if let Err(err) = self.write_file(&path, &next).await {
357                return Err(self.restore_originals_after_failure(&originals, err).await);
358            }
359            originals.push((path, original));
360        }
361        Ok(ReplaceInFilesResult { changed, dry_run })
362    }
363
364    async fn replace_targets(
365        &self,
366        pattern: &str,
367        search: &str,
368    ) -> Result<Vec<VirtualPath>, ToolError> {
369        let glob = self
370            .glob(pattern, self.limits().max_entries as usize)
371            .await?;
372        if glob.truncated {
373            return Err(ToolError::Sdk {
374                sdk_kind: "response_too_large".to_string(),
375                message: "state replace input exceeded max entries".to_string(),
376            });
377        }
378        let mut changed_paths = Vec::new();
379        for path in glob.matches {
380            let virtual_path = VirtualPath::parse(&path)?;
381            let file = self.read_file(&virtual_path).await?;
382            if file.content.contains(search) {
383                changed_paths.push(virtual_path);
384            }
385        }
386        Ok(changed_paths)
387    }
388}
389
390fn normalize_limit(limit: usize) -> usize {
391    limit.clamp(1, 10_000)
392}
393
394fn glob_pattern_regex(pattern: &str) -> Result<Regex, ToolError> {
395    if pattern.trim().is_empty() {
396        return Err(ToolError::InvalidParam {
397            message: "state glob pattern must not be empty".to_string(),
398            param: "pattern".to_string(),
399        });
400    }
401    if pattern.contains("..") || pattern.starts_with('/') || pattern.contains(':') {
402        return Err(ToolError::Sdk {
403            sdk_kind: "path_traversal".to_string(),
404            message: "state glob pattern must stay inside the workspace".to_string(),
405        });
406    }
407    let mut regex = String::from("^");
408    let chars = pattern.chars().collect::<Vec<_>>();
409    let mut i = 0;
410    while i < chars.len() {
411        match chars[i] {
412            '*' if chars.get(i + 1) == Some(&'*') && chars.get(i + 2) == Some(&'/') => {
413                regex.push_str("(?:.*/)?");
414                i += 3;
415            }
416            '*' if chars.get(i + 1) == Some(&'*') => {
417                regex.push_str(".*");
418                i += 2;
419            }
420            '*' => {
421                regex.push_str("[^/]*");
422                i += 1;
423            }
424            '?' => {
425                regex.push_str("[^/]");
426                i += 1;
427            }
428            ch => {
429                regex.push_str(&regex::escape(&ch.to_string()));
430                i += 1;
431            }
432        }
433    }
434    regex.push('$');
435    Regex::new(&regex).map_err(|err| ToolError::InvalidParam {
436        message: format!("invalid state glob pattern: {err}"),
437        param: "pattern".to_string(),
438    })
439}
440
441fn cap_line_preview(line: &str) -> String {
442    line.chars().take(512).collect()
443}
444
445fn ensure_serialized_result_fits<T: Serialize>(value: &T, max: usize) -> Result<(), ToolError> {
446    let len = serde_json::to_vec(value).map_err(serialize_error)?.len();
447    if len > max {
448        return Err(ToolError::Sdk {
449            sdk_kind: "response_too_large".to_string(),
450            message: "state search result exceeded max result bytes".to_string(),
451        });
452    }
453    Ok(())
454}