Skip to main content

soma_codemode/state/
workspace_files.rs

1use std::path::Path;
2
3use sha2::{Digest, Sha256};
4use tokio::io::{AsyncReadExt, AsyncWriteExt};
5
6use crate::path_safety::reject_existing_symlink_ancestors;
7use crate::ToolError;
8
9use super::path::VirtualPath;
10use super::workspace::{
11    internal_io, not_found_or_internal, serialize_error, DetectFileResult, HashFileResult,
12    JsonReadResult, MutationResult, ReadFileResult, StateWorkspace,
13};
14
15impl StateWorkspace {
16    pub async fn write_file(&self, path: &VirtualPath, content: &str) -> Result<(), ToolError> {
17        if content.len() > self.limits().max_file_bytes {
18            return Err(ToolError::InvalidParam {
19                message: format!(
20                    "state file content is {} bytes; maximum is {}",
21                    content.len(),
22                    self.limits().max_file_bytes
23                ),
24                param: "content".to_string(),
25            });
26        }
27        let destination = self.resolve(path);
28        self.ensure_path_allowed(&destination).await?;
29        self.check_write_quota(&destination, content.len() as u64)
30            .await?;
31        if let Some(parent) = destination.parent() {
32            tokio::fs::create_dir_all(parent)
33                .await
34                .map_err(internal_io("create state directory"))?;
35        }
36        self.ensure_path_allowed(&destination).await?;
37
38        let tmp = self.create_temp_path().await?;
39        let mut file = tokio::fs::OpenOptions::new()
40            .create_new(true)
41            .write(true)
42            .open(&tmp)
43            .await
44            .map_err(internal_io("create state temp file"))?;
45        file.write_all(content.as_bytes())
46            .await
47            .map_err(internal_io("write state temp file"))?;
48        file.flush()
49            .await
50            .map_err(internal_io("flush state temp file"))?;
51        drop(file);
52        let tmp_metadata = tokio::fs::symlink_metadata(&tmp)
53            .await
54            .map_err(internal_io("inspect state temp file"))?;
55        if !tmp_metadata.is_file() || tmp_metadata.file_type().is_symlink() {
56            drop(tokio::fs::remove_file(&tmp).await);
57            return Err(ToolError::Sdk {
58                sdk_kind: "permission_denied".to_string(),
59                message: "state temp path is not a regular file".to_string(),
60            });
61        }
62        reject_existing_symlink_ancestors(self.root(), &destination)?;
63        self.reject_existing_symlink_path(&destination).await?;
64        tokio::fs::rename(&tmp, &destination)
65            .await
66            .map_err(internal_io("move state temp file"))?;
67        Ok(())
68    }
69
70    pub async fn read_file(&self, path: &VirtualPath) -> Result<ReadFileResult, ToolError> {
71        let destination = self.resolve(path);
72        self.ensure_path_allowed(&destination).await?;
73        let file = tokio::fs::File::open(&destination)
74            .await
75            .map_err(not_found_or_internal("open state file"))?;
76        let mut content = String::new();
77        file.take(self.limits().max_result_bytes as u64 + 1)
78            .read_to_string(&mut content)
79            .await
80            .map_err(internal_io("read state file"))?;
81        if content.len() > self.limits().max_result_bytes {
82            return Err(ToolError::Sdk {
83                sdk_kind: "response_too_large".to_string(),
84                message: "state read result exceeded max result bytes".to_string(),
85            });
86        }
87        Ok(ReadFileResult {
88            path: path.as_str().to_string(),
89            bytes: content.len(),
90            content,
91        })
92    }
93
94    pub async fn append_file(
95        &self,
96        path: &VirtualPath,
97        content: &str,
98    ) -> Result<MutationResult, ToolError> {
99        let existing = match self.read_file(path).await {
100            Ok(file) => file.content,
101            Err(err) if err.kind() == "not_found" => String::new(),
102            Err(err) => return Err(err),
103        };
104        let next = format!("{existing}{content}");
105        self.write_file(path, &next).await?;
106        Ok(MutationResult {
107            ok: true,
108            path: path.as_str().to_string(),
109        })
110    }
111
112    pub async fn read_json(&self, path: &VirtualPath) -> Result<JsonReadResult, ToolError> {
113        let file = self.read_file(path).await?;
114        let value = serde_json::from_str(&file.content).map_err(|err| ToolError::InvalidParam {
115            message: format!("state file is not valid JSON: {err}"),
116            param: "path".to_string(),
117        })?;
118        Ok(JsonReadResult {
119            path: path.as_str().to_string(),
120            value,
121        })
122    }
123
124    pub async fn write_json(
125        &self,
126        path: &VirtualPath,
127        value: &serde_json::Value,
128        pretty: bool,
129    ) -> Result<(), ToolError> {
130        let mut content = if pretty {
131            serde_json::to_string_pretty(value).map_err(serialize_error)?
132        } else {
133            serde_json::to_string(value).map_err(serialize_error)?
134        };
135        content.push('\n');
136        self.write_file(path, &content).await
137    }
138
139    pub async fn hash_file(
140        &self,
141        path: &VirtualPath,
142        algorithm: &str,
143    ) -> Result<HashFileResult, ToolError> {
144        if algorithm != "sha256" {
145            return Err(ToolError::InvalidParam {
146                message: "state hashFile only supports sha256".to_string(),
147                param: "algorithm".to_string(),
148            });
149        }
150        let bytes = self.read_file_bytes(path).await?;
151        Ok(HashFileResult {
152            path: path.as_str().to_string(),
153            algorithm: algorithm.to_string(),
154            hex: hex::encode(Sha256::digest(&bytes)),
155            bytes: bytes.len(),
156        })
157    }
158
159    pub async fn detect_file(&self, path: &VirtualPath) -> Result<DetectFileResult, ToolError> {
160        let bytes = self.read_file_bytes(path).await?;
161        let extension = Path::new(path.as_str())
162            .extension()
163            .map(|value| value.to_string_lossy().to_ascii_lowercase())
164            .unwrap_or_default();
165        let text = std::str::from_utf8(&bytes).is_ok();
166        let json =
167            extension == "json" || serde_json::from_slice::<serde_json::Value>(&bytes).is_ok();
168        Ok(DetectFileResult {
169            path: path.as_str().to_string(),
170            extension,
171            text,
172            json,
173            bytes: bytes.len(),
174        })
175    }
176
177    pub(crate) async fn read_file_bytes(&self, path: &VirtualPath) -> Result<Vec<u8>, ToolError> {
178        let destination = self.resolve(path);
179        self.ensure_path_allowed(&destination).await?;
180        self.read_bounded_path(&destination, "open state file")
181            .await
182    }
183
184    pub(crate) async fn read_bounded_path(
185        &self,
186        path: &Path,
187        action: &'static str,
188    ) -> Result<Vec<u8>, ToolError> {
189        let file = tokio::fs::File::open(path)
190            .await
191            .map_err(not_found_or_internal(action))?;
192        let mut bytes = Vec::new();
193        file.take(self.limits().max_file_bytes as u64 + 1)
194            .read_to_end(&mut bytes)
195            .await
196            .map_err(internal_io("read state file"))?;
197        if bytes.len() > self.limits().max_file_bytes {
198            return Err(ToolError::Sdk {
199                sdk_kind: "response_too_large".to_string(),
200                message: "state file exceeded max readable bytes".to_string(),
201            });
202        }
203        Ok(bytes)
204    }
205}