diff --git a/src/a2a/server/tasks/inmemory_push_notification_config_store.py b/src/a2a/server/tasks/inmemory_push_notification_config_store.py index f8b0b151b..9b257c0e5 100644 --- a/src/a2a/server/tasks/inmemory_push_notification_config_store.py +++ b/src/a2a/server/tasks/inmemory_push_notification_config_store.py @@ -12,6 +12,14 @@ logger = logging.getLogger(__name__) +def _copy_config( + notification_config: TaskPushNotificationConfig, +) -> TaskPushNotificationConfig: + stored = TaskPushNotificationConfig() + stored.CopyFrom(notification_config) + return stored + + class InMemoryPushNotificationConfigStore(PushNotificationConfigStore): """In-memory implementation of PushNotificationConfigStore interface. @@ -48,16 +56,18 @@ async def set_info( if task_id not in owner_infos: owner_infos[task_id] = [] - if not notification_config.id: - notification_config.id = task_id + stored = TaskPushNotificationConfig() + stored.CopyFrom(notification_config) + if not stored.id: + stored.id = task_id # Remove existing config with the same ID for config in owner_infos[task_id]: - if config.id == notification_config.id: + if config.id == stored.id: owner_infos[task_id].remove(config) break - owner_infos[task_id].append(notification_config) + owner_infos[task_id].append(stored) logger.debug( 'Push notification config for task %s with config id %s for owner %s saved/updated.', task_id, @@ -77,7 +87,7 @@ async def get_info( owner = self.owner_resolver(context) with self.lock: owner_infos = self._get_owner_push_notification_infos(owner) - return list(owner_infos.get(task_id, [])) + return [_copy_config(item) for item in owner_infos.get(task_id, [])] async def get_info_for_dispatch( self, @@ -90,7 +100,9 @@ async def get_info_for_dispatch( with self.lock: results: list[TaskPushNotificationConfig] = [] for all_configs in self._push_notification_infos.values(): - results.extend(all_configs.get(task_id, [])) + results.extend( + _copy_config(item) for item in all_configs.get(task_id, []) + ) return results async def delete_info( diff --git a/tests/server/tasks/test_inmemory_push_notifications.py b/tests/server/tasks/test_inmemory_push_notifications.py index 0a53352f8..74aa2b154 100644 --- a/tests/server/tasks/test_inmemory_push_notifications.py +++ b/tests/server/tasks/test_inmemory_push_notifications.py @@ -112,6 +112,21 @@ async def test_set_info_adds_new_config(self) -> None: ) self.assertEqual(retrieved, [config]) + async def test_set_info_copies_caller_and_returned_config(self) -> None: + task_id = 'task_copy' + config = _create_sample_push_config(url='http://orig.url/callback') + await self.config_store.set_info(task_id, config, MINIMAL_CALL_CONTEXT) + + config.url = 'http://mutated.url/callback' + retrieved = await self.config_store.get_info( + task_id, MINIMAL_CALL_CONTEXT + ) + self.assertEqual(retrieved[0].url, 'http://orig.url/callback') + + retrieved[0].url = 'http://got.url/callback' + again = await self.config_store.get_info(task_id, MINIMAL_CALL_CONTEXT) + self.assertEqual(again[0].url, 'http://orig.url/callback') + async def test_set_info_appends_to_existing_config(self) -> None: task_id = 'task_update' initial_config = _create_sample_push_config(