a2a_server/push/
store.rs

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

a2a_server/push/
store.rs
