Skip to main content

rs_adk/artifacts/
in_memory.rs

1//! In-memory artifact service using DashMap.
2
3use async_trait::async_trait;
4use dashmap::DashMap;
5
6use super::{now_secs, Artifact, ArtifactError, ArtifactMetadata, ArtifactService};
7
8/// In-memory artifact service backed by [`DashMap`] for lock-free concurrent access.
9///
10/// Each artifact name can have multiple versions. Suitable for testing,
11/// prototyping, and single-process deployments. Data is lost on process restart.
12pub struct InMemoryArtifactService {
13    /// Maps session ID to artifact name to version history.
14    store: DashMap<String, DashMap<String, Vec<Artifact>>>,
15}
16
17impl InMemoryArtifactService {
18    /// Create a new in-memory artifact service.
19    pub fn new() -> Self {
20        Self {
21            store: DashMap::new(),
22        }
23    }
24}
25
26impl Default for InMemoryArtifactService {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32#[async_trait]
33impl ArtifactService for InMemoryArtifactService {
34    async fn save(
35        &self,
36        session_id: &str,
37        mut artifact: Artifact,
38    ) -> Result<ArtifactMetadata, ArtifactError> {
39        let session = self.store.entry(session_id.to_string()).or_default();
40
41        let mut versions = session.entry(artifact.metadata.name.clone()).or_default();
42
43        let version = versions.len() as u32 + 1;
44        artifact.metadata.version = version;
45        artifact.metadata.updated_at = now_secs();
46        if version == 1 {
47            artifact.metadata.created_at = artifact.metadata.updated_at;
48        }
49
50        let metadata = artifact.metadata.clone();
51        versions.push(artifact);
52        Ok(metadata)
53    }
54
55    async fn load(&self, session_id: &str, name: &str) -> Result<Option<Artifact>, ArtifactError> {
56        let result = self.store.get(session_id).and_then(|session| {
57            session
58                .get(name)
59                .and_then(|versions| versions.last().cloned())
60        });
61        Ok(result)
62    }
63
64    async fn load_version(
65        &self,
66        session_id: &str,
67        name: &str,
68        version: u32,
69    ) -> Result<Option<Artifact>, ArtifactError> {
70        let result = self.store.get(session_id).and_then(|session| {
71            session.get(name).and_then(|versions| {
72                if version == 0 || version as usize > versions.len() {
73                    None
74                } else {
75                    Some(versions[(version - 1) as usize].clone())
76                }
77            })
78        });
79        Ok(result)
80    }
81
82    async fn list(&self, session_id: &str) -> Result<Vec<ArtifactMetadata>, ArtifactError> {
83        let result = self
84            .store
85            .get(session_id)
86            .map(|session| {
87                session
88                    .iter()
89                    .filter_map(|entry| entry.value().last().map(|a| a.metadata.clone()))
90                    .collect()
91            })
92            .unwrap_or_default();
93        Ok(result)
94    }
95
96    async fn delete(&self, session_id: &str, name: &str) -> Result<(), ArtifactError> {
97        if let Some(session) = self.store.get(session_id) {
98            session.remove(name);
99        }
100        Ok(())
101    }
102}