soma_codemode/artifacts/
store.rs

1use std::path::PathBuf;
2use std::sync::{Arc, Mutex};
3
4use serde::{Deserialize, Serialize};
5use sha2::{Digest, Sha256};
6use tokio::io::AsyncWriteExt;
7
8use crate::ToolError;
9
10use super::path::{artifact_root, safe_artifact_path};
11
12const DEFAULT_MAX_FILE_BYTES: usize = 8 * 1024 * 1024;
13const DEFAULT_MAX_RUN_BYTES: usize = 64 * 1024 * 1024;
14const DEFAULT_MAX_FILES: usize = 128;
15
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
17pub struct ArtifactReceipt {
18    pub path: String,
19    pub absolute_path: String,
20    pub content_type: String,
21    pub bytes: usize,
22    pub sha256: String,
23}
24
25#[derive(Debug, Clone)]
26pub struct ArtifactStore {
27    run_id: String,
28    max_bytes: usize,
29    max_run_bytes: usize,
30    max_files: usize,
31    usage: Arc<Mutex<ArtifactUsage>>,
32}
33
34#[derive(Debug, Default)]
35struct ArtifactUsage {
36    bytes: usize,
37    files: usize,
38}
39
40impl ArtifactStore {
41    pub fn new(run_id: impl Into<String>) -> Result<Self, ToolError> {
42        let run_id = validate_run_id(run_id.into())?;
43        Ok(Self {
44            run_id,
45            max_bytes: DEFAULT_MAX_FILE_BYTES,
46            max_run_bytes: DEFAULT_MAX_RUN_BYTES,
47            max_files: DEFAULT_MAX_FILES,
48            usage: Arc::new(Mutex::new(ArtifactUsage::default())),
49        })
50    }
51
52    pub fn with_max_bytes(mut self, max_bytes: usize) -> Self {
53        self.max_bytes = max_bytes.max(1);
54        self
55    }
56
57    pub fn with_run_limits(mut self, max_run_bytes: usize, max_files: usize) -> Self {
58        self.max_run_bytes = max_run_bytes.max(1);
59        self.max_files = max_files.max(1);
60        self
61    }
62
63    pub fn root(&self) -> PathBuf {
64        artifact_root(&self.run_id)
65    }
66
67    pub async fn write_text(
68        &self,
69        rel_path: &str,
70        content: &str,
71        content_type: Option<&str>,
72    ) -> Result<ArtifactReceipt, ToolError> {
73        if content.len() > self.max_bytes {
74            return Err(ToolError::InvalidParam {
75                message: "artifact content exceeded size limit".to_string(),
76                param: "content".to_string(),
77            });
78        }
79        self.check_run_quota(content.len())?;
80        let root = self.root();
81        let target = safe_artifact_path(&root, rel_path)?;
82        if let Some(parent) = target.parent() {
83            tokio::fs::create_dir_all(parent).await.map_err(|err| {
84                ToolError::internal_message(format!("create artifact dir: {err}"))
85            })?;
86        }
87        let mut file = tokio::fs::File::create(&target)
88            .await
89            .map_err(|err| ToolError::internal_message(format!("create artifact: {err}")))?;
90        file.write_all(content.as_bytes())
91            .await
92            .map_err(|err| ToolError::internal_message(format!("write artifact: {err}")))?;
93        self.record_write(content.len())?;
94        Ok(ArtifactReceipt {
95            path: rel_path.to_string(),
96            absolute_path: target.display().to_string(),
97            content_type: content_type.unwrap_or("text/plain").to_string(),
98            bytes: content.len(),
99            sha256: hex::encode(Sha256::digest(content.as_bytes())),
100        })
101    }
102
103    fn check_run_quota(&self, next_bytes: usize) -> Result<(), ToolError> {
104        let usage = self
105            .usage
106            .lock()
107            .map_err(|_| ToolError::internal_message("artifact usage lock poisoned"))?;
108        if usage.files >= self.max_files
109            || usage.bytes.saturating_add(next_bytes) > self.max_run_bytes
110        {
111            return Err(ToolError::InvalidParam {
112                message: "artifact run quota exceeded".to_string(),
113                param: "content".to_string(),
114            });
115        }
116        Ok(())
117    }
118
119    fn record_write(&self, bytes: usize) -> Result<(), ToolError> {
120        let mut usage = self
121            .usage
122            .lock()
123            .map_err(|_| ToolError::internal_message("artifact usage lock poisoned"))?;
124        usage.files = usage.files.saturating_add(1);
125        usage.bytes = usage.bytes.saturating_add(bytes);
126        Ok(())
127    }
128}
129
130fn validate_run_id(run_id: String) -> Result<String, ToolError> {
131    let trimmed = run_id.trim();
132    let safe = !trimmed.is_empty()
133        && trimmed != "."
134        && trimmed != ".."
135        && trimmed.len() <= 128
136        && trimmed
137            .chars()
138            .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.'));
139    if safe {
140        Ok(trimmed.to_string())
141    } else {
142        Err(ToolError::InvalidParam {
143            message: "artifact run id must be a safe single path segment".to_string(),
144            param: "execution_id".to_string(),
145        })
146    }
147}
soma_codemode/artifacts/store.rs

soma_codemode/artifacts/
store.rs
