a2a_server/push/
store.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// SPDX-License-Identifier: Apache-2.0
3use a2a::*;
4use async_trait::async_trait;
5use std::collections::HashMap;
6use tokio::sync::RwLock;
7
8/// Interface for persisting push notification configurations.
9#[async_trait]
10pub trait PushConfigStore: Send + Sync + 'static {
11    /// Save a push config for a task. Generates an ID if none provided.
12    async fn save(
13        &self,
14        config: TaskPushNotificationConfig,
15    ) -> Result<TaskPushNotificationConfig, A2AError>;
16
17    /// Get a specific push config by task and config ID.
18    async fn get(
19        &self,
20        task_id: &str,
21        config_id: &str,
22    ) -> Result<TaskPushNotificationConfig, A2AError>;
23
24    /// List all push configs for a task.
25    async fn list(&self, task_id: &str) -> Result<Vec<TaskPushNotificationConfig>, A2AError>;
26
27    /// Delete a specific push config.
28    async fn delete(&self, task_id: &str, config_id: &str) -> Result<(), A2AError>;
29
30    /// Delete all push configs for a task.
31    async fn delete_all(&self, task_id: &str) -> Result<(), A2AError>;
32}
33
34/// In-memory push config store.
35pub struct InMemoryPushConfigStore {
36    configs: RwLock<HashMap<String, HashMap<String, TaskPushNotificationConfig>>>,
37}
38
39impl InMemoryPushConfigStore {
40    pub fn new() -> Self {
41        InMemoryPushConfigStore {
42            configs: RwLock::new(HashMap::new()),
43        }
44    }
45}
46
47impl Default for InMemoryPushConfigStore {
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53#[async_trait]
54impl PushConfigStore for InMemoryPushConfigStore {
55    async fn save(
56        &self,
57        mut config: TaskPushNotificationConfig,
58    ) -> Result<TaskPushNotificationConfig, A2AError> {
59        if config.url.is_empty() {
60            return Err(A2AError::invalid_params("push config URL cannot be empty"));
61        }
62        if config.task_id.is_empty() {
63            return Err(A2AError::invalid_params(
64                "push config task_id cannot be empty",
65            ));
66        }
67
68        if config.id.is_none() {
69            config.id = Some(uuid::Uuid::now_v7().to_string());
70        }
71
72        let mut store = self.configs.write().await;
73        let task_configs = store.entry(config.task_id.clone()).or_default();
74        let config_id = config.id.clone().unwrap();
75        task_configs.insert(config_id, config.clone());
76        Ok(config)
77    }
78
79    async fn get(
80        &self,
81        task_id: &str,
82        config_id: &str,
83    ) -> Result<TaskPushNotificationConfig, A2AError> {
84        let store = self.configs.read().await;
85        store
86            .get(task_id)
87            .and_then(|configs| configs.get(config_id))
88            .cloned()
89            .ok_or_else(A2AError::push_notification_not_supported)
90    }
91
92    async fn list(&self, task_id: &str) -> Result<Vec<TaskPushNotificationConfig>, A2AError> {
93        let store = self.configs.read().await;
94        Ok(store
95            .get(task_id)
96            .map(|configs| configs.values().cloned().collect())
97            .unwrap_or_default())
98    }
99
100    async fn delete(&self, task_id: &str, config_id: &str) -> Result<(), A2AError> {
101        let mut store = self.configs.write().await;
102        if let Some(configs) = store.get_mut(task_id) {
103            configs.remove(config_id);
104        }
105        Ok(())
106    }
107
108    async fn delete_all(&self, task_id: &str) -> Result<(), A2AError> {
109        let mut store = self.configs.write().await;
110        store.remove(task_id);
111        Ok(())
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    fn make_config(task_id: &str, url: &str) -> TaskPushNotificationConfig {
120        TaskPushNotificationConfig {
121            url: url.to_string(),
122            id: None,
123            task_id: task_id.to_string(),
124            token: None,
125            authentication: None,
126            tenant: None,
127        }
128    }
129
130    #[tokio::test]
131    async fn test_save_and_get() {
132        let store = InMemoryPushConfigStore::new();
133        let config = store
134            .save(make_config("t1", "https://example.com/hook"))
135            .await
136            .unwrap();
137        assert!(config.id.is_some());
138
139        let got = store.get("t1", config.id.as_ref().unwrap()).await.unwrap();
140        assert_eq!(got.url, "https://example.com/hook");
141    }
142
143    #[tokio::test]
144    async fn test_save_empty_url() {
145        let store = InMemoryPushConfigStore::new();
146        let result = store.save(make_config("t1", "")).await;
147        assert!(result.is_err());
148    }
149
150    #[tokio::test]
151    async fn test_save_empty_task_id() {
152        let store = InMemoryPushConfigStore::new();
153        let result = store
154            .save(make_config("", "https://example.com/hook"))
155            .await;
156        assert!(result.is_err());
157    }
158
159    #[tokio::test]
160    async fn test_save_with_existing_id() {
161        let store = InMemoryPushConfigStore::new();
162        let mut config = make_config("t1", "https://example.com/hook");
163        config.id = Some("my-id".to_string());
164        let saved = store.save(config).await.unwrap();
165        assert_eq!(saved.id.as_deref(), Some("my-id"));
166    }
167
168    #[tokio::test]
169    async fn test_get_not_found() {
170        let store = InMemoryPushConfigStore::new();
171        let result = store.get("t0", "nonexistent").await;
172        assert!(result.is_err());
173    }
174
175    #[tokio::test]
176    async fn test_list() {
177        let store = InMemoryPushConfigStore::new();
178        store
179            .save(make_config("t1", "https://a.com"))
180            .await
181            .unwrap();
182        store
183            .save(make_config("t1", "https://b.com"))
184            .await
185            .unwrap();
186        store
187            .save(make_config("t2", "https://c.com"))
188            .await
189            .unwrap();
190
191        let configs = store.list("t1").await.unwrap();
192        assert_eq!(configs.len(), 2);
193
194        let configs = store.list("t2").await.unwrap();
195        assert_eq!(configs.len(), 1);
196
197        let configs = store.list("t3").await.unwrap();
198        assert_eq!(configs.len(), 0);
199    }
200
201    #[tokio::test]
202    async fn test_delete() {
203        let store = InMemoryPushConfigStore::new();
204        let config = store
205            .save(make_config("t1", "https://a.com"))
206            .await
207            .unwrap();
208        let config_id = config.id.unwrap();
209
210        store.delete("t1", &config_id).await.unwrap();
211        let result = store.get("t1", &config_id).await;
212        assert!(result.is_err());
213    }
214
215    #[tokio::test]
216    async fn test_delete_all() {
217        let store = InMemoryPushConfigStore::new();
218        store
219            .save(make_config("t1", "https://a.com"))
220            .await
221            .unwrap();
222        store
223            .save(make_config("t1", "https://b.com"))
224            .await
225            .unwrap();
226
227        store.delete_all("t1").await.unwrap();
228        let configs = store.list("t1").await.unwrap();
229        assert_eq!(configs.len(), 0);
230    }
231}
a2a_server/push/store.rs

a2a_server/push/
store.rs
