Skip to main content

soma_codemode/state/
workspace.rs

1use std::path::{Component, Path, PathBuf};
2
3use serde::{Deserialize, Serialize};
4
5use crate::path_safety::{reject_existing_symlink_ancestors, rel_to_unix_string};
6use crate::ToolError;
7
8use super::path::{is_reserved_metadata_path, VirtualPath};
9use super::quota::StateWorkspaceLimits;
10
11#[derive(Debug, Clone)]
12pub struct StateWorkspace {
13    root: PathBuf,
14    limits: StateWorkspaceLimits,
15}
16
17#[derive(Debug, Clone, Serialize)]
18pub struct ReadFileResult {
19    pub path: String,
20    pub content: String,
21    pub bytes: usize,
22}
23
24#[derive(Debug, Clone, Serialize)]
25pub struct ListResult {
26    pub entries: Vec<String>,
27}
28
29#[derive(Debug, Clone, Serialize)]
30pub struct MutationResult {
31    pub ok: bool,
32    pub path: String,
33}
34
35#[derive(Debug, Clone, Serialize)]
36pub struct ExistsResult {
37    pub path: String,
38    pub exists: bool,
39}
40
41#[derive(Debug, Clone, Serialize)]
42pub struct StatResult {
43    pub path: String,
44    pub kind: String,
45    pub bytes: u64,
46}
47
48#[derive(Debug, Clone, Serialize)]
49pub struct WalkEntry {
50    pub path: String,
51    pub kind: String,
52    pub bytes: u64,
53}
54
55#[derive(Debug, Clone, Serialize)]
56pub struct WalkTreeResult {
57    pub entries: Vec<WalkEntry>,
58    pub truncated: bool,
59}
60
61#[derive(Debug, Clone, Serialize)]
62pub struct JsonReadResult {
63    pub path: String,
64    pub value: serde_json::Value,
65}
66
67#[derive(Debug, Clone, Serialize)]
68pub struct HashFileResult {
69    pub path: String,
70    pub algorithm: String,
71    pub hex: String,
72    pub bytes: usize,
73}
74
75#[derive(Debug, Clone, Serialize)]
76pub struct DetectFileResult {
77    pub path: String,
78    pub extension: String,
79    pub text: bool,
80    pub json: bool,
81    pub bytes: usize,
82}
83
84#[derive(Debug, Clone, Serialize)]
85pub struct ArchiveCreateResult {
86    pub ok: bool,
87    pub destination: String,
88    pub entries: usize,
89}
90
91#[derive(Debug, Clone, Serialize)]
92pub struct ArchiveListResult {
93    pub path: String,
94    pub entries: Vec<String>,
95    pub truncated: bool,
96}
97
98#[derive(Debug, Clone, Serialize)]
99pub struct GlobResult {
100    pub matches: Vec<String>,
101    pub truncated: bool,
102}
103
104#[derive(Debug, Clone, Serialize)]
105pub struct SearchMatch {
106    pub path: String,
107    pub line: usize,
108    pub text: String,
109}
110
111#[derive(Debug, Clone, Serialize)]
112pub struct SearchFilesResult {
113    pub matches: Vec<SearchMatch>,
114    pub truncated: bool,
115}
116
117#[derive(Debug, Clone, Serialize)]
118pub struct ReplaceInFilesResult {
119    pub changed: Vec<String>,
120    pub dry_run: bool,
121}
122
123#[derive(Debug, Clone, Deserialize, Serialize)]
124pub struct FileEdit {
125    pub path: String,
126    pub search: String,
127    pub replace: String,
128}
129
130#[derive(Debug, Clone, Serialize)]
131pub struct EditPlanResult {
132    pub plan_id: String,
133    pub edits: Vec<FileEdit>,
134}
135
136#[derive(Debug, Clone, Serialize)]
137pub struct ApplyEditPlanResult {
138    pub ok: bool,
139    pub changed: Vec<String>,
140}
141
142pub(crate) struct WalkFilesResult {
143    pub(crate) files: Vec<String>,
144    pub(crate) truncated: bool,
145}
146
147struct WorkspaceUsage {
148    bytes: u64,
149    entries: u64,
150}
151
152pub fn default_search_limit() -> usize {
153    200
154}
155
156pub fn default_true() -> bool {
157    true
158}
159
160impl StateWorkspace {
161    pub fn new(root: impl Into<PathBuf>) -> Self {
162        Self {
163            root: root.into(),
164            limits: StateWorkspaceLimits::default(),
165        }
166    }
167
168    pub fn with_limits(root: impl Into<PathBuf>, limits: StateWorkspaceLimits) -> Self {
169        Self {
170            root: root.into(),
171            limits,
172        }
173    }
174
175    pub fn root(&self) -> &Path {
176        &self.root
177    }
178
179    pub fn root_path(&self) -> &PathBuf {
180        &self.root
181    }
182
183    pub(crate) fn limits(&self) -> StateWorkspaceLimits {
184        self.limits
185    }
186
187    pub(crate) fn resolve(&self, path: &VirtualPath) -> PathBuf {
188        self.root.join(path.as_str())
189    }
190
191    pub(crate) async fn ensure_path_allowed(&self, path: &Path) -> Result<(), ToolError> {
192        reject_existing_symlink_ancestors(&self.root, path)?;
193        self.reject_existing_symlink_path(path).await
194    }
195
196    pub(crate) async fn reject_existing_symlink_path(&self, path: &Path) -> Result<(), ToolError> {
197        match tokio::fs::symlink_metadata(path).await {
198            Ok(metadata) if metadata.file_type().is_symlink() => Err(ToolError::Sdk {
199                sdk_kind: "symlink_rejected".to_string(),
200                message: "state path is denied because it is a symlink".to_string(),
201            }),
202            Ok(_) => Ok(()),
203            Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
204            Err(err) => Err(internal_io("read state path metadata")(err)),
205        }
206    }
207
208    pub(crate) async fn create_temp_path(&self) -> Result<PathBuf, ToolError> {
209        let dir = self.root.join(".soma-state").join("tmp");
210        tokio::fs::create_dir_all(&dir)
211            .await
212            .map_err(internal_io("create state temp directory"))?;
213        reject_existing_symlink_ancestors(&self.root, &dir)?;
214        let metadata = tokio::fs::symlink_metadata(&dir)
215            .await
216            .map_err(internal_io("inspect state temp directory"))?;
217        if !metadata.is_dir() || metadata.file_type().is_symlink() {
218            return Err(ToolError::Sdk {
219                sdk_kind: "permission_denied".to_string(),
220                message: "state temp directory is not a directory".to_string(),
221            });
222        }
223        Ok(dir.join(format!("{}.tmp", ulid::Ulid::generate())))
224    }
225
226    pub(crate) async fn check_write_quota(
227        &self,
228        destination: &Path,
229        next_file_bytes: u64,
230    ) -> Result<(), ToolError> {
231        let current_file_bytes = match tokio::fs::metadata(destination).await {
232            Ok(metadata) if metadata.is_file() => metadata.len(),
233            Ok(_) | Err(_) => 0,
234        };
235        let usage = workspace_usage(&self.root).await?;
236        let projected = usage
237            .bytes
238            .saturating_sub(current_file_bytes)
239            .saturating_add(next_file_bytes);
240        if projected > self.limits.max_total_bytes {
241            return Err(quota_error(format!(
242                "state workspace would be {projected} bytes; maximum is {}",
243                self.limits.max_total_bytes
244            )));
245        }
246        let projected_entries = usage
247            .entries
248            .saturating_add(missing_entry_count(&self.root, destination).await?);
249        if projected_entries > self.limits.max_entries {
250            return Err(quota_error(format!(
251                "state workspace would have {projected_entries} entries; maximum is {}",
252                self.limits.max_entries
253            )));
254        }
255        Ok(())
256    }
257
258    pub async fn enforce_total_bytes(&self) -> Result<(), ToolError> {
259        self.enforce_total_limits().await
260    }
261
262    pub(crate) async fn enforce_total_limits(&self) -> Result<(), ToolError> {
263        let usage = workspace_usage(&self.root).await?;
264        if usage.bytes > self.limits.max_total_bytes {
265            return Err(quota_error(format!(
266                "state workspace is {} bytes; maximum is {}",
267                usage.bytes, self.limits.max_total_bytes
268            )));
269        }
270        if usage.entries > self.limits.max_entries {
271            return Err(quota_error(format!(
272                "state workspace has {} entries; maximum is {}",
273                usage.entries, self.limits.max_entries
274            )));
275        }
276        Ok(())
277    }
278
279    pub(crate) async fn check_entry_quota_for_path(
280        &self,
281        destination: &Path,
282    ) -> Result<(), ToolError> {
283        let usage = workspace_usage(&self.root).await?;
284        let projected = usage
285            .entries
286            .saturating_add(missing_entry_count(&self.root, destination).await?);
287        if projected > self.limits.max_entries {
288            return Err(quota_error(format!(
289                "state workspace would have {projected} entries; maximum is {}",
290                self.limits.max_entries
291            )));
292        }
293        Ok(())
294    }
295
296    pub(crate) fn plan_path(&self, plan_id: &str) -> PathBuf {
297        self.root
298            .join(".soma-state")
299            .join("plans")
300            .join(format!("{plan_id}.json"))
301    }
302
303    pub(crate) async fn walk_files(&self, limit: usize) -> Result<WalkFilesResult, ToolError> {
304        let mut files = Vec::new();
305        let mut truncated = false;
306        let mut stack = vec![self.root.clone()];
307        while let Some(dir) = stack.pop() {
308            let mut read_dir = match tokio::fs::read_dir(&dir).await {
309                Ok(read_dir) => read_dir,
310                Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
311                Err(error) => return Err(internal_io("walk state workspace")(error)),
312            };
313            while let Some(entry) = read_dir
314                .next_entry()
315                .await
316                .map_err(internal_io("walk state workspace entry"))?
317            {
318                let path = entry.path();
319                let relative = match path.strip_prefix(&self.root) {
320                    Ok(relative) => relative,
321                    Err(_) => continue,
322                };
323                let virtual_path = rel_to_unix_string(relative);
324                if is_reserved_metadata_path(&virtual_path) {
325                    continue;
326                }
327                let metadata = tokio::fs::symlink_metadata(&path)
328                    .await
329                    .map_err(internal_io("read state workspace metadata"))?;
330                if metadata.file_type().is_symlink() {
331                    continue;
332                }
333                if metadata.is_dir() {
334                    stack.push(path);
335                } else if metadata.is_file() {
336                    files.push(virtual_path);
337                    if files.len() > limit {
338                        truncated = true;
339                        break;
340                    }
341                }
342            }
343            if truncated {
344                break;
345            }
346        }
347        files.sort();
348        if truncated {
349            files.truncate(limit);
350        }
351        Ok(WalkFilesResult { files, truncated })
352    }
353}
354
355pub(crate) fn internal_io(action: &'static str) -> impl FnOnce(std::io::Error) -> ToolError {
356    move |err| ToolError::Sdk {
357        sdk_kind: "internal_error".to_string(),
358        message: format!("failed to {action}: {err}"),
359    }
360}
361
362pub(crate) fn not_found_or_internal(
363    action: &'static str,
364) -> impl FnOnce(std::io::Error) -> ToolError {
365    move |err| ToolError::Sdk {
366        sdk_kind: if err.kind() == std::io::ErrorKind::NotFound {
367            "not_found"
368        } else {
369            "internal_error"
370        }
371        .to_string(),
372        message: format!("failed to {action}: {err}"),
373    }
374}
375
376pub(crate) fn serialize_error(err: serde_json::Error) -> ToolError {
377    ToolError::Sdk {
378        sdk_kind: "internal_error".to_string(),
379        message: format!("failed to serialize state value: {err}"),
380    }
381}
382
383pub(crate) async fn cleanup_file_after_quota_error(
384    path: &Path,
385    original: ToolError,
386    label: &str,
387) -> Result<(), ToolError> {
388    match tokio::fs::remove_file(path).await {
389        Ok(()) => Err(original),
390        Err(err) if err.kind() == std::io::ErrorKind::NotFound => Err(original),
391        Err(err) => Err(ToolError::Sdk {
392            sdk_kind: "quota_cleanup_failed".to_string(),
393            message: format!("{label} exceeded quota and cleanup failed: {err}"),
394        }),
395    }
396}
397
398fn quota_error(message: String) -> ToolError {
399    ToolError::Sdk {
400        sdk_kind: "quota_exceeded".to_string(),
401        message,
402    }
403}
404
405async fn workspace_usage(root: &Path) -> Result<WorkspaceUsage, ToolError> {
406    let mut usage = WorkspaceUsage {
407        bytes: 0,
408        entries: 0,
409    };
410    let mut stack = vec![root.to_path_buf()];
411    while let Some(dir) = stack.pop() {
412        let mut read_dir = match tokio::fs::read_dir(&dir).await {
413            Ok(read_dir) => read_dir,
414            Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
415            Err(error) => return Err(internal_io("scan state workspace")(error)),
416        };
417        while let Some(entry) = read_dir
418            .next_entry()
419            .await
420            .map_err(internal_io("scan state workspace entry"))?
421        {
422            let path = entry.path();
423            let metadata = tokio::fs::symlink_metadata(&path)
424                .await
425                .map_err(internal_io("read state workspace metadata"))?;
426            if metadata.file_type().is_symlink() {
427                continue;
428            }
429            if path.strip_prefix(root).is_ok() {
430                usage.entries = usage.entries.saturating_add(1);
431            }
432            if metadata.is_dir() {
433                stack.push(path);
434            } else if metadata.is_file() {
435                usage.bytes = usage.bytes.saturating_add(metadata.len());
436            }
437        }
438    }
439    Ok(usage)
440}
441
442async fn missing_entry_count(root: &Path, destination: &Path) -> Result<u64, ToolError> {
443    let relative = destination.strip_prefix(root).map_err(|_| ToolError::Sdk {
444        sdk_kind: "path_traversal".to_string(),
445        message: "state path escapes the workspace".to_string(),
446    })?;
447    let virtual_path = rel_to_unix_string(relative);
448    if is_reserved_metadata_path(&virtual_path) {
449        return Ok(0);
450    }
451    let parts = relative
452        .components()
453        .filter_map(|component| match component {
454            Component::Normal(part) => Some(part.to_owned()),
455            _ => None,
456        })
457        .collect::<Vec<_>>();
458    let mut current = root.to_path_buf();
459    for (index, part) in parts.iter().enumerate() {
460        current.push(part);
461        match tokio::fs::symlink_metadata(&current).await {
462            Ok(_) => {}
463            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
464                return Ok((parts.len() - index) as u64);
465            }
466            Err(err) => return Err(internal_io("read state path metadata")(err)),
467        }
468    }
469    Ok(0)
470}