a3s_code_core/tools/
artifacts.rs

1//! In-memory artifact storage for large tool observations.
2
3use anyhow::{Context, Result};
4use serde::{Deserialize, Serialize};
5use std::collections::{HashMap, HashSet, VecDeque};
6use std::io::{Read, Write};
7use std::path::Path;
8use std::sync::{Arc, RwLock};
9use thiserror::Error;
10
11const DEFAULT_MAX_ARTIFACTS: usize = 256;
12const DEFAULT_MAX_BYTES: usize = 16 * 1024 * 1024;
13/// Hard safety boundary for the on-disk artifact manifest.
14///
15/// Artifact contents are already bounded by [`ArtifactStoreLimits`], but the
16/// manifest is read before those limits can be applied. Keep that read
17/// bounded so a corrupt or untrusted session directory cannot force an
18/// unbounded allocation during recovery.
19const MAX_ARTIFACT_MANIFEST_BYTES: u64 = 256 * 1024 * 1024;
20
21#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
22pub struct ToolArtifact {
23    pub artifact_id: String,
24    pub artifact_uri: String,
25    pub tool_name: String,
26    pub content: String,
27    pub original_bytes: usize,
28    pub shown_bytes: usize,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32struct ArtifactStoreSnapshot {
33    artifacts: Vec<ToolArtifact>,
34    /// Host-supplied retention roots that must survive eviction/GC.
35    #[serde(default)]
36    retained_uris: Vec<String>,
37}
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub struct ArtifactStoreLimits {
41    pub max_artifacts: usize,
42    pub max_bytes: usize,
43}
44
45/// Conflict raised when a content-addressed URI is reused for different
46/// artifact bytes or metadata.  The existing [`ArtifactStore::put`] method
47/// remains available for mutable cache callers; research and replay paths
48/// should use [`ArtifactStore::put_content_addressed`] instead.
49#[derive(Debug, Clone, Eq, PartialEq, Error)]
50pub enum ArtifactStoreError {
51    #[error("artifact URI '{artifact_uri}' is already bound to different content")]
52    Conflict { artifact_uri: String },
53}
54
55impl Default for ArtifactStoreLimits {
56    fn default() -> Self {
57        Self {
58            max_artifacts: DEFAULT_MAX_ARTIFACTS,
59            max_bytes: DEFAULT_MAX_BYTES,
60        }
61    }
62}
63
64#[derive(Debug, Default)]
65struct ArtifactStoreState {
66    artifacts: HashMap<String, ToolArtifact>,
67    insertion_order: VecDeque<String>,
68    total_bytes: usize,
69    /// URIs that must survive limit eviction and explicit GC until unpinned.
70    retained_uris: HashSet<String>,
71}
72
73#[derive(Debug, Clone)]
74pub struct ArtifactStore {
75    inner: Arc<RwLock<ArtifactStoreState>>,
76    limits: ArtifactStoreLimits,
77}
78
79impl ArtifactStore {
80    pub fn new() -> Self {
81        Self::with_limits(ArtifactStoreLimits::default())
82    }
83
84    pub fn with_limits(limits: ArtifactStoreLimits) -> Self {
85        Self {
86            inner: Arc::new(RwLock::new(ArtifactStoreState::default())),
87            limits,
88        }
89    }
90
91    pub fn put(&self, artifact: ToolArtifact) {
92        let mut state = self.inner.write().unwrap();
93        let artifact_uri = artifact.artifact_uri.clone();
94        if let Some(existing) = state.artifacts.remove(&artifact_uri) {
95            state.total_bytes = state.total_bytes.saturating_sub(existing.content.len());
96            state.insertion_order.retain(|uri| uri != &artifact_uri);
97        }
98
99        // Accounting is part of the eviction invariant. Saturating here
100        // keeps a malformed/oversized value from wrapping the counter and
101        // bypassing the byte limit in release builds.
102        state.total_bytes = state.total_bytes.saturating_add(artifact.content.len());
103        state.insertion_order.push_back(artifact_uri.clone());
104        state.artifacts.insert(artifact_uri, artifact);
105
106        self.enforce_limits(&mut state);
107    }
108
109    /// Insert an artifact without allowing an existing URI to be overwritten.
110    ///
111    /// Exact replay is idempotent and returns `Ok(false)`.  A URI collision
112    /// with different bytes or metadata fails closed and leaves the retained
113    /// artifact untouched.  Retention eviction remains explicit store policy:
114    /// after an object is evicted, a later create-only write may reinsert its
115    /// URI because the store no longer owns that historical object.
116    pub fn put_content_addressed(
117        &self,
118        artifact: ToolArtifact,
119    ) -> Result<bool, ArtifactStoreError> {
120        let mut state = self.inner.write().unwrap();
121        let artifact_uri = artifact.artifact_uri.clone();
122        if let Some(existing) = state.artifacts.get(&artifact_uri) {
123            if existing == &artifact {
124                return Ok(false);
125            }
126            return Err(ArtifactStoreError::Conflict { artifact_uri });
127        }
128
129        state.total_bytes += artifact.content.len();
130        state.insertion_order.push_back(artifact_uri.clone());
131        state.artifacts.insert(artifact_uri, artifact);
132        self.enforce_limits(&mut state);
133        Ok(true)
134    }
135
136    pub fn get(&self, artifact_uri: &str) -> Option<ToolArtifact> {
137        self.inner
138            .read()
139            .unwrap()
140            .artifacts
141            .get(artifact_uri)
142            .cloned()
143    }
144
145    pub fn len(&self) -> usize {
146        self.inner.read().unwrap().artifacts.len()
147    }
148
149    pub fn is_empty(&self) -> bool {
150        self.len() == 0
151    }
152
153    pub fn total_bytes(&self) -> usize {
154        self.inner.read().unwrap().total_bytes
155    }
156
157    pub fn limits(&self) -> ArtifactStoreLimits {
158        self.limits
159    }
160
161    pub fn artifacts(&self) -> Vec<ToolArtifact> {
162        self.ordered_artifacts()
163    }
164
165    /// Pin artifact URIs so limit eviction and unreferenced GC cannot remove
166    /// content still reachable from a retained identity (STORE-GC1).
167    pub fn pin_uris(&self, uris: impl IntoIterator<Item = impl Into<String>>) {
168        let mut state = self.inner.write().unwrap();
169        for uri in uris {
170            let uri = uri.into();
171            if !uri.is_empty() {
172                state.retained_uris.insert(uri);
173            }
174        }
175    }
176
177    /// Replace the retained URI set with exactly the supplied roots.
178    pub fn set_retained_uris(&self, uris: impl IntoIterator<Item = impl Into<String>>) {
179        let mut state = self.inner.write().unwrap();
180        state.retained_uris.clear();
181        for uri in uris {
182            let uri = uri.into();
183            if !uri.is_empty() {
184                state.retained_uris.insert(uri);
185            }
186        }
187    }
188
189    pub fn retained_uris(&self) -> HashSet<String> {
190        self.inner.read().unwrap().retained_uris.clone()
191    }
192
193    /// Remove artifacts that are not in the retained root set.
194    ///
195    /// Returns the number of removed objects. Retained roots that are absent
196    /// from the store are ignored; present retained objects are never removed.
197    pub fn gc_unreferenced(&self) -> usize {
198        let mut state = self.inner.write().unwrap();
199        let removable: Vec<String> = state
200            .insertion_order
201            .iter()
202            .filter(|uri| !state.retained_uris.contains(uri.as_str()))
203            .cloned()
204            .collect();
205        let mut removed = 0usize;
206        for uri in removable {
207            if let Some(artifact) = state.artifacts.remove(&uri) {
208                state.total_bytes = state.total_bytes.saturating_sub(artifact.content.len());
209                state.insertion_order.retain(|queued| queued != &uri);
210                removed += 1;
211            }
212        }
213        removed
214    }
215
216    pub fn save_to_dir(&self, dir: impl AsRef<Path>) -> Result<()> {
217        let dir = dir.as_ref();
218        std::fs::create_dir_all(dir)
219            .with_context(|| format!("failed to create artifact directory '{}'", dir.display()))?;
220        let mut retained: Vec<String> = self.retained_uris().into_iter().collect();
221        retained.sort();
222        let snapshot = ArtifactStoreSnapshot {
223            artifacts: self.ordered_artifacts(),
224            retained_uris: retained,
225        };
226        let json = serde_json::to_string_pretty(&snapshot)
227            .context("failed to serialize artifact store snapshot")?;
228        let path = artifact_manifest_path(dir);
229        if json.len() as u64 > MAX_ARTIFACT_MANIFEST_BYTES {
230            anyhow::bail!(
231                "refusing to write artifact manifest '{}': {} bytes exceeds the {} byte limit",
232                path.display(),
233                json.len(),
234                MAX_ARTIFACT_MANIFEST_BYTES
235            );
236        }
237
238        // The manifest is the legacy fragment-store boundary. Publish it as
239        // one generation so a reader can never observe a partially written
240        // JSON document after a crash or concurrent session load.
241        let mut temp = tempfile::NamedTempFile::new_in(dir).with_context(|| {
242            format!(
243                "failed to create temporary artifact manifest in '{}'",
244                dir.display()
245            )
246        })?;
247        temp.write_all(json.as_bytes())
248            .context("failed to write temporary artifact manifest")?;
249        temp.flush()
250            .context("failed to flush temporary artifact manifest")?;
251        temp.as_file()
252            .sync_all()
253            .context("failed to sync temporary artifact manifest")?;
254        temp.persist(&path)
255            .map_err(|error| error.error)
256            .with_context(|| {
257                format!(
258                    "failed to atomically replace artifact manifest '{}'",
259                    path.display()
260                )
261            })?;
262        Ok(())
263    }
264
265    pub fn load_from_dir(dir: impl AsRef<Path>) -> Result<Self> {
266        Self::load_from_dir_with_limits(dir, ArtifactStoreLimits::default())
267    }
268
269    pub fn load_from_manifest_bytes(bytes: &[u8]) -> Result<Self> {
270        Self::load_from_manifest_bytes_with_limits(bytes, ArtifactStoreLimits::default())
271    }
272
273    pub fn load_from_manifest_bytes_with_limits(
274        bytes: &[u8],
275        limits: ArtifactStoreLimits,
276    ) -> Result<Self> {
277        if bytes.len() as u64 > MAX_ARTIFACT_MANIFEST_BYTES {
278            anyhow::bail!(
279                "refusing to parse artifact manifest: {} bytes exceeds the {} byte limit",
280                bytes.len(),
281                MAX_ARTIFACT_MANIFEST_BYTES
282            );
283        }
284        let snapshot: ArtifactStoreSnapshot =
285            serde_json::from_slice(bytes).context("failed to parse artifact store snapshot")?;
286        let store = Self::with_limits(limits);
287        for artifact in snapshot.artifacts {
288            store
289                .put_content_addressed(artifact)
290                .map_err(|error| anyhow::anyhow!("invalid artifact manifest: {error}"))?;
291        }
292        store.set_retained_uris(snapshot.retained_uris);
293        Ok(store)
294    }
295
296    pub fn load_from_dir_with_limits(
297        dir: impl AsRef<Path>,
298        limits: ArtifactStoreLimits,
299    ) -> Result<Self> {
300        let path = artifact_manifest_path(dir.as_ref());
301        if !path.exists() {
302            return Ok(Self::with_limits(limits));
303        }
304
305        let json = read_manifest(&path)?;
306        Self::load_from_manifest_bytes_with_limits(&json, limits)
307    }
308
309    fn enforce_limits(&self, state: &mut ArtifactStoreState) {
310        while state.artifacts.len() > self.limits.max_artifacts
311            || state.total_bytes > self.limits.max_bytes
312        {
313            let Some(uri) = pop_oldest_evictable(state) else {
314                break;
315            };
316            if let Some(removed) = state.artifacts.remove(&uri) {
317                state.total_bytes = state.total_bytes.saturating_sub(removed.content.len());
318            }
319        }
320    }
321
322    fn ordered_artifacts(&self) -> Vec<ToolArtifact> {
323        let state = self.inner.read().unwrap();
324        state
325            .insertion_order
326            .iter()
327            .filter_map(|uri| state.artifacts.get(uri).cloned())
328            .collect()
329    }
330}
331
332fn pop_oldest_evictable(state: &mut ArtifactStoreState) -> Option<String> {
333    let mut skipped = VecDeque::new();
334    let mut evicted = None;
335    while let Some(uri) = state.insertion_order.pop_front() {
336        if state.retained_uris.contains(&uri) {
337            skipped.push_back(uri);
338            continue;
339        }
340        evicted = Some(uri);
341        break;
342    }
343    while let Some(uri) = skipped.pop_back() {
344        state.insertion_order.push_front(uri);
345    }
346    evicted
347}
348
349impl Default for ArtifactStore {
350    fn default() -> Self {
351        Self::new()
352    }
353}
354
355fn artifact_manifest_path(dir: &Path) -> std::path::PathBuf {
356    dir.join("artifacts.json")
357}
358
359fn read_manifest(path: &Path) -> Result<Vec<u8>> {
360    let file = std::fs::File::open(path)
361        .with_context(|| format!("failed to open artifact manifest '{}'", path.display()))?;
362    let declared_len = file
363        .metadata()
364        .with_context(|| format!("failed to inspect artifact manifest '{}'", path.display()))?
365        .len();
366    if declared_len > MAX_ARTIFACT_MANIFEST_BYTES {
367        anyhow::bail!(
368            "refusing to read artifact manifest '{}': {} bytes exceeds the {} byte limit",
369            path.display(),
370            declared_len,
371            MAX_ARTIFACT_MANIFEST_BYTES
372        );
373    }
374
375    // Re-check through a bounded reader so a file that grows after metadata
376    // inspection cannot race past the allocation boundary.
377    let mut reader = file.take(MAX_ARTIFACT_MANIFEST_BYTES + 1);
378    let mut bytes = Vec::with_capacity(
379        usize::try_from(declared_len)
380            .unwrap_or(usize::MAX)
381            .min(1024 * 1024),
382    );
383    reader
384        .read_to_end(&mut bytes)
385        .with_context(|| format!("failed to read artifact manifest '{}'", path.display()))?;
386    if bytes.len() as u64 > MAX_ARTIFACT_MANIFEST_BYTES {
387        anyhow::bail!(
388            "refusing to read artifact manifest '{}': document exceeds the {} byte limit",
389            path.display(),
390            MAX_ARTIFACT_MANIFEST_BYTES
391        );
392    }
393    Ok(bytes)
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399
400    #[test]
401    fn test_artifact_store_put_and_get() {
402        let store = ArtifactStore::new();
403        let artifact = ToolArtifact {
404            artifact_id: "tool-output:test:abc".to_string(),
405            artifact_uri: "a3s://tool-output/test/abc".to_string(),
406            tool_name: "test".to_string(),
407            content: "full output".to_string(),
408            original_bytes: 11,
409            shown_bytes: 4,
410        };
411
412        store.put(artifact.clone());
413
414        assert_eq!(store.len(), 1);
415        assert_eq!(store.get("a3s://tool-output/test/abc"), Some(artifact));
416    }
417
418    #[test]
419    fn test_content_addressed_put_is_idempotent_and_conflict_safe() {
420        let store = ArtifactStore::new();
421        let artifact = ToolArtifact {
422            artifact_id: "tool-output:test:immutable".to_string(),
423            artifact_uri: "a3s://tool-output/test/immutable".to_string(),
424            tool_name: "test".to_string(),
425            content: "immutable output".to_string(),
426            original_bytes: 16,
427            shown_bytes: 8,
428        };
429
430        assert!(store.put_content_addressed(artifact.clone()).unwrap());
431        assert!(!store.put_content_addressed(artifact.clone()).unwrap());
432        assert_eq!(store.len(), 1);
433
434        let mut conflicting = artifact.clone();
435        conflicting.content = "different output".to_string();
436        assert_eq!(
437            store.put_content_addressed(conflicting),
438            Err(ArtifactStoreError::Conflict {
439                artifact_uri: artifact.artifact_uri.clone()
440            })
441        );
442        assert_eq!(store.get(&artifact.artifact_uri), Some(artifact));
443    }
444
445    #[test]
446    fn test_artifact_store_missing_uri() {
447        let store = ArtifactStore::new();
448
449        assert!(store.is_empty());
450        assert!(store.get("a3s://missing").is_none());
451    }
452
453    #[test]
454    fn test_artifact_store_evicts_oldest_by_count() {
455        let store = ArtifactStore::with_limits(ArtifactStoreLimits {
456            max_artifacts: 2,
457            max_bytes: 1024,
458        });
459
460        for index in 0..3 {
461            store.put(ToolArtifact {
462                artifact_id: format!("tool-output:test:{index}"),
463                artifact_uri: format!("a3s://tool-output/test/{index}"),
464                tool_name: "test".to_string(),
465                content: format!("artifact {index}"),
466                original_bytes: 10,
467                shown_bytes: 4,
468            });
469        }
470
471        assert_eq!(store.len(), 2);
472        assert!(store.get("a3s://tool-output/test/0").is_none());
473        assert!(store.get("a3s://tool-output/test/1").is_some());
474        assert!(store.get("a3s://tool-output/test/2").is_some());
475    }
476
477    #[test]
478    fn test_artifact_store_evicts_oldest_by_bytes() {
479        let store = ArtifactStore::with_limits(ArtifactStoreLimits {
480            max_artifacts: 10,
481            max_bytes: 8,
482        });
483
484        store.put(ToolArtifact {
485            artifact_id: "tool-output:test:a".to_string(),
486            artifact_uri: "a3s://tool-output/test/a".to_string(),
487            tool_name: "test".to_string(),
488            content: "aaaa".to_string(),
489            original_bytes: 4,
490            shown_bytes: 2,
491        });
492        store.put(ToolArtifact {
493            artifact_id: "tool-output:test:b".to_string(),
494            artifact_uri: "a3s://tool-output/test/b".to_string(),
495            tool_name: "test".to_string(),
496            content: "bbbbb".to_string(),
497            original_bytes: 5,
498            shown_bytes: 2,
499        });
500
501        assert_eq!(store.len(), 1);
502        assert_eq!(store.total_bytes(), 5);
503        assert!(store.get("a3s://tool-output/test/a").is_none());
504        assert!(store.get("a3s://tool-output/test/b").is_some());
505    }
506
507    #[test]
508    fn test_artifact_store_replacing_artifact_updates_order_and_bytes() {
509        let store = ArtifactStore::with_limits(ArtifactStoreLimits {
510            max_artifacts: 2,
511            max_bytes: 1024,
512        });
513
514        store.put(ToolArtifact {
515            artifact_id: "tool-output:test:a".to_string(),
516            artifact_uri: "a3s://tool-output/test/a".to_string(),
517            tool_name: "test".to_string(),
518            content: "a".to_string(),
519            original_bytes: 1,
520            shown_bytes: 1,
521        });
522        store.put(ToolArtifact {
523            artifact_id: "tool-output:test:b".to_string(),
524            artifact_uri: "a3s://tool-output/test/b".to_string(),
525            tool_name: "test".to_string(),
526            content: "bb".to_string(),
527            original_bytes: 2,
528            shown_bytes: 1,
529        });
530        store.put(ToolArtifact {
531            artifact_id: "tool-output:test:a".to_string(),
532            artifact_uri: "a3s://tool-output/test/a".to_string(),
533            tool_name: "test".to_string(),
534            content: "aaaa".to_string(),
535            original_bytes: 4,
536            shown_bytes: 1,
537        });
538
539        assert_eq!(store.len(), 2);
540        assert_eq!(store.total_bytes(), 6);
541        assert_eq!(
542            store.get("a3s://tool-output/test/a").unwrap().content,
543            "aaaa"
544        );
545    }
546
547    #[test]
548    fn test_artifact_store_saves_and_loads_manifest() {
549        let dir = tempfile::tempdir().unwrap();
550        let store = ArtifactStore::with_limits(ArtifactStoreLimits {
551            max_artifacts: 10,
552            max_bytes: 1024,
553        });
554        store.put(ToolArtifact {
555            artifact_id: "tool-output:test:a".to_string(),
556            artifact_uri: "a3s://tool-output/test/a".to_string(),
557            tool_name: "test".to_string(),
558            content: "artifact content".to_string(),
559            original_bytes: 16,
560            shown_bytes: 4,
561        });
562
563        store.save_to_dir(dir.path()).unwrap();
564        let loaded = ArtifactStore::load_from_dir_with_limits(
565            dir.path(),
566            ArtifactStoreLimits {
567                max_artifacts: 10,
568                max_bytes: 1024,
569            },
570        )
571        .unwrap();
572
573        assert_eq!(loaded.len(), 1);
574        assert_eq!(
575            loaded
576                .get("a3s://tool-output/test/a")
577                .expect("artifact")
578                .content,
579            "artifact content"
580        );
581    }
582
583    #[test]
584    fn test_artifact_store_load_missing_manifest_returns_empty_store() {
585        let dir = tempfile::tempdir().unwrap();
586
587        let loaded = ArtifactStore::load_from_dir(dir.path()).unwrap();
588
589        assert!(loaded.is_empty());
590    }
591
592    #[test]
593    fn test_artifact_store_rejects_oversized_manifest_before_reading() {
594        let dir = tempfile::tempdir().unwrap();
595        let path = dir.path().join("artifacts.json");
596        let file = std::fs::File::create(&path).unwrap();
597        file.set_len(MAX_ARTIFACT_MANIFEST_BYTES + 1).unwrap();
598
599        let error = ArtifactStore::load_from_dir(dir.path()).unwrap_err();
600
601        assert!(error.to_string().contains("exceeds the"));
602        assert!(error.to_string().contains("artifact manifest"));
603    }
604
605    #[test]
606    fn retained_uris_survive_limit_eviction() {
607        let store = ArtifactStore::with_limits(ArtifactStoreLimits {
608            max_artifacts: 2,
609            max_bytes: 1024,
610        });
611        store.put(ToolArtifact {
612            artifact_id: "tool-output:test:pinned".to_string(),
613            artifact_uri: "a3s://tool-output/test/pinned".to_string(),
614            tool_name: "test".to_string(),
615            content: "pinned".to_string(),
616            original_bytes: 6,
617            shown_bytes: 3,
618        });
619        store.put(ToolArtifact {
620            artifact_id: "tool-output:test:kept".to_string(),
621            artifact_uri: "a3s://tool-output/test/kept".to_string(),
622            tool_name: "test".to_string(),
623            content: "kept".to_string(),
624            original_bytes: 4,
625            shown_bytes: 2,
626        });
627        store.set_retained_uris([
628            "a3s://tool-output/test/pinned",
629            "a3s://tool-output/test/kept",
630        ]);
631        store.put(ToolArtifact {
632            artifact_id: "tool-output:test:third".to_string(),
633            artifact_uri: "a3s://tool-output/test/third".to_string(),
634            tool_name: "test".to_string(),
635            content: "third".to_string(),
636            original_bytes: 5,
637            shown_bytes: 2,
638        });
639        assert!(store.get("a3s://tool-output/test/pinned").is_some());
640        assert!(store.get("a3s://tool-output/test/kept").is_some());
641        assert!(
642            store.get("a3s://tool-output/test/third").is_none(),
643            "unpinned overflow should still be evicted"
644        );
645        assert_eq!(store.len(), 2);
646    }
647
648    #[test]
649    fn gc_unreferenced_keeps_only_retained_roots() {
650        let store = ArtifactStore::new();
651        for index in 0..3 {
652            store.put(ToolArtifact {
653                artifact_id: format!("tool-output:test:{index}"),
654                artifact_uri: format!("a3s://tool-output/test/{index}"),
655                tool_name: "test".to_string(),
656                content: format!("artifact {index}"),
657                original_bytes: 10,
658                shown_bytes: 4,
659            });
660        }
661        store.set_retained_uris(["a3s://tool-output/test/1"]);
662        assert_eq!(store.gc_unreferenced(), 2);
663        assert!(store.get("a3s://tool-output/test/0").is_none());
664        assert!(store.get("a3s://tool-output/test/1").is_some());
665        assert!(store.get("a3s://tool-output/test/2").is_none());
666        assert_eq!(store.retained_uris().len(), 1);
667    }
668}
a3s_code_core/tools/artifacts.rs

a3s_code_core/tools/
artifacts.rs
