1use anyhow::{Context, Result};
4use serde::{Deserialize, Serialize};
5use std::collections::{HashMap, VecDeque};
6use std::path::Path;
7use std::sync::{Arc, RwLock};
8
9const DEFAULT_MAX_ARTIFACTS: usize = 256;
10const DEFAULT_MAX_BYTES: usize = 16 * 1024 * 1024;
11
12#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
13pub struct ToolArtifact {
14 pub artifact_id: String,
15 pub artifact_uri: String,
16 pub tool_name: String,
17 pub content: String,
18 pub original_bytes: usize,
19 pub shown_bytes: usize,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23struct ArtifactStoreSnapshot {
24 artifacts: Vec<ToolArtifact>,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub struct ArtifactStoreLimits {
29 pub max_artifacts: usize,
30 pub max_bytes: usize,
31}
32
33impl Default for ArtifactStoreLimits {
34 fn default() -> Self {
35 Self {
36 max_artifacts: DEFAULT_MAX_ARTIFACTS,
37 max_bytes: DEFAULT_MAX_BYTES,
38 }
39 }
40}
41
42#[derive(Debug, Default)]
43struct ArtifactStoreState {
44 artifacts: HashMap<String, ToolArtifact>,
45 insertion_order: VecDeque<String>,
46 total_bytes: usize,
47}
48
49#[derive(Debug, Clone)]
50pub struct ArtifactStore {
51 inner: Arc<RwLock<ArtifactStoreState>>,
52 limits: ArtifactStoreLimits,
53}
54
55impl ArtifactStore {
56 pub fn new() -> Self {
57 Self::with_limits(ArtifactStoreLimits::default())
58 }
59
60 pub fn with_limits(limits: ArtifactStoreLimits) -> Self {
61 Self {
62 inner: Arc::new(RwLock::new(ArtifactStoreState::default())),
63 limits,
64 }
65 }
66
67 pub fn put(&self, artifact: ToolArtifact) {
68 let mut state = self.inner.write().unwrap();
69 let artifact_uri = artifact.artifact_uri.clone();
70 if let Some(existing) = state.artifacts.remove(&artifact_uri) {
71 state.total_bytes = state.total_bytes.saturating_sub(existing.content.len());
72 state.insertion_order.retain(|uri| uri != &artifact_uri);
73 }
74
75 state.total_bytes += artifact.content.len();
76 state.insertion_order.push_back(artifact_uri.clone());
77 state.artifacts.insert(artifact_uri, artifact);
78
79 self.enforce_limits(&mut state);
80 }
81
82 pub fn get(&self, artifact_uri: &str) -> Option<ToolArtifact> {
83 self.inner
84 .read()
85 .unwrap()
86 .artifacts
87 .get(artifact_uri)
88 .cloned()
89 }
90
91 pub fn len(&self) -> usize {
92 self.inner.read().unwrap().artifacts.len()
93 }
94
95 pub fn is_empty(&self) -> bool {
96 self.len() == 0
97 }
98
99 pub fn total_bytes(&self) -> usize {
100 self.inner.read().unwrap().total_bytes
101 }
102
103 pub fn limits(&self) -> ArtifactStoreLimits {
104 self.limits
105 }
106
107 pub fn artifacts(&self) -> Vec<ToolArtifact> {
108 self.ordered_artifacts()
109 }
110
111 pub fn save_to_dir(&self, dir: impl AsRef<Path>) -> Result<()> {
112 let dir = dir.as_ref();
113 std::fs::create_dir_all(dir)
114 .with_context(|| format!("failed to create artifact directory '{}'", dir.display()))?;
115 let snapshot = ArtifactStoreSnapshot {
116 artifacts: self.ordered_artifacts(),
117 };
118 let json = serde_json::to_string_pretty(&snapshot)
119 .context("failed to serialize artifact store snapshot")?;
120 let path = artifact_manifest_path(dir);
121 std::fs::write(&path, json)
122 .with_context(|| format!("failed to write artifact manifest '{}'", path.display()))?;
123 Ok(())
124 }
125
126 pub fn load_from_dir(dir: impl AsRef<Path>) -> Result<Self> {
127 Self::load_from_dir_with_limits(dir, ArtifactStoreLimits::default())
128 }
129
130 pub fn load_from_dir_with_limits(
131 dir: impl AsRef<Path>,
132 limits: ArtifactStoreLimits,
133 ) -> Result<Self> {
134 let path = artifact_manifest_path(dir.as_ref());
135 if !path.exists() {
136 return Ok(Self::with_limits(limits));
137 }
138
139 let json = std::fs::read_to_string(&path)
140 .with_context(|| format!("failed to read artifact manifest '{}'", path.display()))?;
141 let snapshot: ArtifactStoreSnapshot =
142 serde_json::from_str(&json).context("failed to parse artifact store snapshot")?;
143 let store = Self::with_limits(limits);
144 for artifact in snapshot.artifacts {
145 store.put(artifact);
146 }
147 Ok(store)
148 }
149
150 fn enforce_limits(&self, state: &mut ArtifactStoreState) {
151 while state.artifacts.len() > self.limits.max_artifacts
152 || state.total_bytes > self.limits.max_bytes
153 {
154 let Some(uri) = state.insertion_order.pop_front() else {
155 break;
156 };
157 if let Some(removed) = state.artifacts.remove(&uri) {
158 state.total_bytes = state.total_bytes.saturating_sub(removed.content.len());
159 }
160 }
161 }
162
163 fn ordered_artifacts(&self) -> Vec<ToolArtifact> {
164 let state = self.inner.read().unwrap();
165 state
166 .insertion_order
167 .iter()
168 .filter_map(|uri| state.artifacts.get(uri).cloned())
169 .collect()
170 }
171}
172
173impl Default for ArtifactStore {
174 fn default() -> Self {
175 Self::new()
176 }
177}
178
179fn artifact_manifest_path(dir: &Path) -> std::path::PathBuf {
180 dir.join("artifacts.json")
181}
182
183#[cfg(test)]
184mod tests {
185 use super::*;
186
187 #[test]
188 fn test_artifact_store_put_and_get() {
189 let store = ArtifactStore::new();
190 let artifact = ToolArtifact {
191 artifact_id: "tool-output:test:abc".to_string(),
192 artifact_uri: "a3s://tool-output/test/abc".to_string(),
193 tool_name: "test".to_string(),
194 content: "full output".to_string(),
195 original_bytes: 11,
196 shown_bytes: 4,
197 };
198
199 store.put(artifact.clone());
200
201 assert_eq!(store.len(), 1);
202 assert_eq!(store.get("a3s://tool-output/test/abc"), Some(artifact));
203 }
204
205 #[test]
206 fn test_artifact_store_missing_uri() {
207 let store = ArtifactStore::new();
208
209 assert!(store.is_empty());
210 assert!(store.get("a3s://missing").is_none());
211 }
212
213 #[test]
214 fn test_artifact_store_evicts_oldest_by_count() {
215 let store = ArtifactStore::with_limits(ArtifactStoreLimits {
216 max_artifacts: 2,
217 max_bytes: 1024,
218 });
219
220 for index in 0..3 {
221 store.put(ToolArtifact {
222 artifact_id: format!("tool-output:test:{index}"),
223 artifact_uri: format!("a3s://tool-output/test/{index}"),
224 tool_name: "test".to_string(),
225 content: format!("artifact {index}"),
226 original_bytes: 10,
227 shown_bytes: 4,
228 });
229 }
230
231 assert_eq!(store.len(), 2);
232 assert!(store.get("a3s://tool-output/test/0").is_none());
233 assert!(store.get("a3s://tool-output/test/1").is_some());
234 assert!(store.get("a3s://tool-output/test/2").is_some());
235 }
236
237 #[test]
238 fn test_artifact_store_evicts_oldest_by_bytes() {
239 let store = ArtifactStore::with_limits(ArtifactStoreLimits {
240 max_artifacts: 10,
241 max_bytes: 8,
242 });
243
244 store.put(ToolArtifact {
245 artifact_id: "tool-output:test:a".to_string(),
246 artifact_uri: "a3s://tool-output/test/a".to_string(),
247 tool_name: "test".to_string(),
248 content: "aaaa".to_string(),
249 original_bytes: 4,
250 shown_bytes: 2,
251 });
252 store.put(ToolArtifact {
253 artifact_id: "tool-output:test:b".to_string(),
254 artifact_uri: "a3s://tool-output/test/b".to_string(),
255 tool_name: "test".to_string(),
256 content: "bbbbb".to_string(),
257 original_bytes: 5,
258 shown_bytes: 2,
259 });
260
261 assert_eq!(store.len(), 1);
262 assert_eq!(store.total_bytes(), 5);
263 assert!(store.get("a3s://tool-output/test/a").is_none());
264 assert!(store.get("a3s://tool-output/test/b").is_some());
265 }
266
267 #[test]
268 fn test_artifact_store_replacing_artifact_updates_order_and_bytes() {
269 let store = ArtifactStore::with_limits(ArtifactStoreLimits {
270 max_artifacts: 2,
271 max_bytes: 1024,
272 });
273
274 store.put(ToolArtifact {
275 artifact_id: "tool-output:test:a".to_string(),
276 artifact_uri: "a3s://tool-output/test/a".to_string(),
277 tool_name: "test".to_string(),
278 content: "a".to_string(),
279 original_bytes: 1,
280 shown_bytes: 1,
281 });
282 store.put(ToolArtifact {
283 artifact_id: "tool-output:test:b".to_string(),
284 artifact_uri: "a3s://tool-output/test/b".to_string(),
285 tool_name: "test".to_string(),
286 content: "bb".to_string(),
287 original_bytes: 2,
288 shown_bytes: 1,
289 });
290 store.put(ToolArtifact {
291 artifact_id: "tool-output:test:a".to_string(),
292 artifact_uri: "a3s://tool-output/test/a".to_string(),
293 tool_name: "test".to_string(),
294 content: "aaaa".to_string(),
295 original_bytes: 4,
296 shown_bytes: 1,
297 });
298
299 assert_eq!(store.len(), 2);
300 assert_eq!(store.total_bytes(), 6);
301 assert_eq!(
302 store.get("a3s://tool-output/test/a").unwrap().content,
303 "aaaa"
304 );
305 }
306
307 #[test]
308 fn test_artifact_store_saves_and_loads_manifest() {
309 let dir = tempfile::tempdir().unwrap();
310 let store = ArtifactStore::with_limits(ArtifactStoreLimits {
311 max_artifacts: 10,
312 max_bytes: 1024,
313 });
314 store.put(ToolArtifact {
315 artifact_id: "tool-output:test:a".to_string(),
316 artifact_uri: "a3s://tool-output/test/a".to_string(),
317 tool_name: "test".to_string(),
318 content: "artifact content".to_string(),
319 original_bytes: 16,
320 shown_bytes: 4,
321 });
322
323 store.save_to_dir(dir.path()).unwrap();
324 let loaded = ArtifactStore::load_from_dir_with_limits(
325 dir.path(),
326 ArtifactStoreLimits {
327 max_artifacts: 10,
328 max_bytes: 1024,
329 },
330 )
331 .unwrap();
332
333 assert_eq!(loaded.len(), 1);
334 assert_eq!(
335 loaded
336 .get("a3s://tool-output/test/a")
337 .expect("artifact")
338 .content,
339 "artifact content"
340 );
341 }
342
343 #[test]
344 fn test_artifact_store_load_missing_manifest_returns_empty_store() {
345 let dir = tempfile::tempdir().unwrap();
346
347 let loaded = ArtifactStore::load_from_dir(dir.path()).unwrap();
348
349 assert!(loaded.is_empty());
350 }
351}