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/// Maximum number of push notification configs allowed per task.
9///
10/// Each config costs memory and triggers an outbound request on every
11/// event, so an unbounded count is a resource-exhaustion vector.
12const MAX_PUSH_CONFIGS_PER_TASK: usize = 50;
13
14/// Interface for persisting push notification configurations.
15#[async_trait]
16pub trait PushConfigStore: Send + Sync + 'static {
17    /// Save a push config for a task. Generates an ID if none provided.
18    async fn save(
19        &self,
20        config: TaskPushNotificationConfig,
21    ) -> Result<TaskPushNotificationConfig, A2AError>;
22
23    /// Get a specific push config by task and config ID.
24    async fn get(
25        &self,
26        task_id: &str,
27        config_id: &str,
28    ) -> Result<TaskPushNotificationConfig, A2AError>;
29
30    /// List all push configs for a task.
31    async fn list(&self, task_id: &str) -> Result<Vec<TaskPushNotificationConfig>, A2AError>;
32
33    /// Delete a specific push config.
34    async fn delete(&self, task_id: &str, config_id: &str) -> Result<(), A2AError>;
35
36    /// Delete all push configs for a task.
37    async fn delete_all(&self, task_id: &str) -> Result<(), A2AError>;
38}
39
40/// In-memory push config store.
41pub struct InMemoryPushConfigStore {
42    configs: RwLock<HashMap<String, HashMap<String, TaskPushNotificationConfig>>>,
43}
44
45impl InMemoryPushConfigStore {
46    pub fn new() -> Self {
47        InMemoryPushConfigStore {
48            configs: RwLock::new(HashMap::new()),
49        }
50    }
51}
52
53impl Default for InMemoryPushConfigStore {
54    fn default() -> Self {
55        Self::new()
56    }
57}
58
59#[async_trait]
60impl PushConfigStore for InMemoryPushConfigStore {
61    async fn save(
62        &self,
63        mut config: TaskPushNotificationConfig,
64    ) -> Result<TaskPushNotificationConfig, A2AError> {
65        if config.url.is_empty() {
66            return Err(A2AError::invalid_params("push config URL cannot be empty"));
67        }
68        if config.task_id.is_empty() {
69            return Err(A2AError::invalid_params(
70                "push config task_id cannot be empty",
71            ));
72        }
73
74        if config.id.is_none() {
75            config.id = Some(uuid::Uuid::now_v7().to_string());
76        }
77
78        let mut store = self.configs.write().await;
79        let task_configs = store.entry(config.task_id.clone()).or_default();
80        let config_id = config.id.clone().unwrap();
81        // Enforce a per-task cap so a client cannot register unbounded
82        // configs, each consuming memory and triggering outbound requests.
83        // Replacing an existing config by ID is still allowed at the cap.
84        if !task_configs.contains_key(&config_id) && task_configs.len() >= MAX_PUSH_CONFIGS_PER_TASK
85        {
86            return Err(A2AError::invalid_params(format!(
87                "too many push notification configs for task (max {MAX_PUSH_CONFIGS_PER_TASK})"
88            )));
89        }
90        task_configs.insert(config_id, config.clone());
91        Ok(config)
92    }
93
94    async fn get(
95        &self,
96        task_id: &str,
97        config_id: &str,
98    ) -> Result<TaskPushNotificationConfig, A2AError> {
99        let store = self.configs.read().await;
100        store
101            .get(task_id)
102            .and_then(|configs| configs.get(config_id))
103            .cloned()
104            .ok_or_else(A2AError::push_notification_not_supported)
105    }
106
107    async fn list(&self, task_id: &str) -> Result<Vec<TaskPushNotificationConfig>, A2AError> {
108        let store = self.configs.read().await;
109        Ok(store
110            .get(task_id)
111            .map(|configs| configs.values().cloned().collect())
112            .unwrap_or_default())
113    }
114
115    async fn delete(&self, task_id: &str, config_id: &str) -> Result<(), A2AError> {
116        let mut store = self.configs.write().await;
117        if let Some(configs) = store.get_mut(task_id) {
118            configs.remove(config_id);
119        }
120        Ok(())
121    }
122
123    async fn delete_all(&self, task_id: &str) -> Result<(), A2AError> {
124        let mut store = self.configs.write().await;
125        store.remove(task_id);
126        Ok(())
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    fn make_config(task_id: &str, url: &str) -> TaskPushNotificationConfig {
135        TaskPushNotificationConfig {
136            url: url.to_string(),
137            id: None,
138            task_id: task_id.to_string(),
139            token: None,
140            authentication: None,
141            tenant: None,
142        }
143    }
144
145    #[tokio::test]
146    async fn test_save_and_get() {
147        let store = InMemoryPushConfigStore::new();
148        let config = store
149            .save(make_config("t1", "https://example.com/hook"))
150            .await
151            .unwrap();
152        assert!(config.id.is_some());
153
154        let got = store.get("t1", config.id.as_ref().unwrap()).await.unwrap();
155        assert_eq!(got.url, "https://example.com/hook");
156    }
157
158    #[tokio::test]
159    async fn test_save_empty_url() {
160        let store = InMemoryPushConfigStore::new();
161        let result = store.save(make_config("t1", "")).await;
162        assert!(result.is_err());
163    }
164
165    #[tokio::test]
166    async fn test_save_empty_task_id() {
167        let store = InMemoryPushConfigStore::new();
168        let result = store
169            .save(make_config("", "https://example.com/hook"))
170            .await;
171        assert!(result.is_err());
172    }
173
174    #[tokio::test]
175    async fn test_save_with_existing_id() {
176        let store = InMemoryPushConfigStore::new();
177        let mut config = make_config("t1", "https://example.com/hook");
178        config.id = Some("my-id".to_string());
179        let saved = store.save(config).await.unwrap();
180        assert_eq!(saved.id.as_deref(), Some("my-id"));
181    }
182
183    #[tokio::test]
184    async fn test_get_not_found() {
185        let store = InMemoryPushConfigStore::new();
186        let result = store.get("t0", "nonexistent").await;
187        assert!(result.is_err());
188    }
189
190    #[tokio::test]
191    async fn test_list() {
192        let store = InMemoryPushConfigStore::new();
193        store
194            .save(make_config("t1", "https://a.com"))
195            .await
196            .unwrap();
197        store
198            .save(make_config("t1", "https://b.com"))
199            .await
200            .unwrap();
201        store
202            .save(make_config("t2", "https://c.com"))
203            .await
204            .unwrap();
205
206        let configs = store.list("t1").await.unwrap();
207        assert_eq!(configs.len(), 2);
208
209        let configs = store.list("t2").await.unwrap();
210        assert_eq!(configs.len(), 1);
211
212        let configs = store.list("t3").await.unwrap();
213        assert_eq!(configs.len(), 0);
214    }
215
216    #[tokio::test]
217    async fn test_delete() {
218        let store = InMemoryPushConfigStore::new();
219        let config = store
220            .save(make_config("t1", "https://a.com"))
221            .await
222            .unwrap();
223        let config_id = config.id.unwrap();
224
225        store.delete("t1", &config_id).await.unwrap();
226        let result = store.get("t1", &config_id).await;
227        assert!(result.is_err());
228    }
229
230    #[tokio::test]
231    async fn test_delete_all() {
232        let store = InMemoryPushConfigStore::new();
233        store
234            .save(make_config("t1", "https://a.com"))
235            .await
236            .unwrap();
237        store
238            .save(make_config("t1", "https://b.com"))
239            .await
240            .unwrap();
241
242        store.delete_all("t1").await.unwrap();
243        let configs = store.list("t1").await.unwrap();
244        assert_eq!(configs.len(), 0);
245    }
246
247    #[tokio::test]
248    async fn test_save_caps_configs_per_task() {
249        let store = InMemoryPushConfigStore::new();
250        for i in 0..MAX_PUSH_CONFIGS_PER_TASK {
251            store
252                .save(make_config("t1", &format!("https://example.com/hook/{i}")))
253                .await
254                .unwrap();
255        }
256
257        // The (cap + 1)-th distinct config is rejected.
258        let result = store
259            .save(make_config("t1", "https://example.com/over-limit"))
260            .await;
261        assert!(result.is_err());
262        assert_eq!(result.unwrap_err().code, error_code::INVALID_PARAMS);
263
264        // Other tasks are not affected.
265        store
266            .save(make_config("t2", "https://example.com/hook"))
267            .await
268            .unwrap();
269    }
270
271    #[tokio::test]
272    async fn test_save_can_replace_existing_config_at_cap() {
273        let store = InMemoryPushConfigStore::new();
274        let mut first = make_config("t1", "https://example.com/hook/0");
275        first.id = Some("cfg-keep".to_string());
276        store.save(first).await.unwrap();
277        for i in 1..MAX_PUSH_CONFIGS_PER_TASK {
278            store
279                .save(make_config("t1", &format!("https://example.com/hook/{i}")))
280                .await
281                .unwrap();
282        }
283        assert_eq!(
284            store.list("t1").await.unwrap().len(),
285            MAX_PUSH_CONFIGS_PER_TASK
286        );
287
288        // Updating an existing config by ID is allowed even at the cap.
289        let mut config = make_config("t1", "https://example.com/replaced");
290        config.id = Some("cfg-keep".to_string());
291        let saved = store.save(config).await.unwrap();
292        assert_eq!(saved.id.as_deref(), Some("cfg-keep"));
293        assert_eq!(saved.url, "https://example.com/replaced");
294        assert_eq!(
295            store.list("t1").await.unwrap().len(),
296            MAX_PUSH_CONFIGS_PER_TASK
297        );
298    }
299}
a2a_server/push/store.rs

a2a_server/push/
store.rs
