Skip to main content

soma_codemode/state/
workspace_edit.rs

1use sha2::{Digest, Sha256};
2
3use crate::ToolError;
4
5use super::path::VirtualPath;
6use super::workspace::{
7    cleanup_file_after_quota_error, internal_io, serialize_error, ApplyEditPlanResult,
8    EditPlanResult, FileEdit, StateWorkspace,
9};
10
11impl StateWorkspace {
12    pub async fn plan_edits(&self, edits: Vec<FileEdit>) -> Result<EditPlanResult, ToolError> {
13        let edits = normalize_edits(edits)?;
14        let canonical = serde_json::to_vec(&edits).map_err(serialize_error)?;
15        if canonical.len() > self.limits().max_file_bytes {
16            return Err(ToolError::Sdk {
17                sdk_kind: "response_too_large".to_string(),
18                message: "state edit plan exceeded max file bytes".to_string(),
19            });
20        }
21        let plan_id = hex::encode(Sha256::digest(&canonical));
22        let plan_path = self.plan_path(&plan_id);
23        if let Some(parent) = plan_path.parent() {
24            self.check_entry_quota_for_path(parent).await?;
25            tokio::fs::create_dir_all(parent)
26                .await
27                .map_err(internal_io("create state edit plan directory"))?;
28        }
29        self.check_write_quota(&plan_path, canonical.len() as u64)
30            .await?;
31        tokio::fs::write(&plan_path, canonical)
32            .await
33            .map_err(internal_io("write state edit plan"))?;
34        if let Err(err) = self.enforce_total_limits().await {
35            cleanup_file_after_quota_error(&plan_path, err, "state edit plan").await?;
36        }
37        Ok(EditPlanResult { plan_id, edits })
38    }
39
40    pub async fn apply_edit_plan(&self, plan_id: &str) -> Result<ApplyEditPlanResult, ToolError> {
41        validate_plan_id(plan_id)?;
42        let plan_path = self.plan_path(plan_id);
43        let plan = self
44            .read_bounded_path(&plan_path, "read state edit plan")
45            .await?;
46        let edits: Vec<FileEdit> = serde_json::from_slice(&plan).map_err(|err| ToolError::Sdk {
47            sdk_kind: "internal_error".to_string(),
48            message: format!("failed to parse state edit plan: {err}"),
49        })?;
50        let mut planned = Vec::new();
51        for edit in edits {
52            let path = VirtualPath::parse(&edit.path)?;
53            let original = self.read_file(&path).await?;
54            if !original.content.contains(&edit.search) {
55                return Err(ToolError::Sdk {
56                    sdk_kind: "edit_conflict".to_string(),
57                    message: format!("state edit plan no longer matches `{}`", path.as_str()),
58                });
59            }
60            let next = original.content.replace(&edit.search, &edit.replace);
61            planned.push((path, original.content, next));
62        }
63        let mut changed = Vec::new();
64        let mut originals = Vec::new();
65        for (path, original, next) in planned {
66            if let Err(err) = self.write_file(&path, &next).await {
67                return Err(self.restore_originals_after_failure(&originals, err).await);
68            }
69            originals.push((path.clone(), original));
70            changed.push(path.as_str().to_string());
71        }
72        Ok(ApplyEditPlanResult { ok: true, changed })
73    }
74
75    pub(super) async fn restore_originals_after_failure(
76        &self,
77        originals: &[(VirtualPath, String)],
78        original_error: ToolError,
79    ) -> ToolError {
80        for (path, content) in originals.iter().rev() {
81            if let Err(rollback_error) = self.write_file(path, content).await {
82                return ToolError::Sdk {
83                    sdk_kind: "rollback_failed".to_string(),
84                    message: format!(
85                        "state batch mutation failed with `{}` and rollback of `{}` failed with `{}`",
86                        original_error.kind(),
87                        path.as_str(),
88                        rollback_error.kind()
89                    ),
90                };
91            }
92        }
93        original_error
94    }
95}
96
97fn normalize_edits(edits: Vec<FileEdit>) -> Result<Vec<FileEdit>, ToolError> {
98    if edits.is_empty() {
99        return Err(ToolError::InvalidParam {
100            message: "state edit plan must include at least one edit".to_string(),
101            param: "edits".to_string(),
102        });
103    }
104    edits
105        .into_iter()
106        .map(|edit| {
107            if edit.search.is_empty() {
108                return Err(ToolError::InvalidParam {
109                    message: "state edit search must not be empty".to_string(),
110                    param: "search".to_string(),
111                });
112            }
113            let path = VirtualPath::parse(&edit.path)?.as_str().to_string();
114            Ok(FileEdit {
115                path,
116                search: edit.search,
117                replace: edit.replace,
118            })
119        })
120        .collect()
121}
122
123fn validate_plan_id(plan_id: &str) -> Result<(), ToolError> {
124    let valid = plan_id.len() == 64 && plan_id.chars().all(|ch| ch.is_ascii_hexdigit());
125    if valid {
126        Ok(())
127    } else {
128        Err(ToolError::InvalidParam {
129            message: "state edit plan id must be a sha256 hex string".to_string(),
130            param: "planId".to_string(),
131        })
132    }
133}