fix: prevent first-owner write loss in in-memory stores

src/a2a/server/tasks/inmemory_push_notification_config_store.py

@@ -43,10 +43,8 @@ async def set_info(
     ) -> None:
         """Sets or updates the push notification configuration for a task in memory."""
         owner = self.owner_resolver(context)
-        if owner not in self._push_notification_infos:
-            self._push_notification_infos[owner] = {}
         with self.lock:
-            owner_infos = self._push_notification_infos[owner]
+            owner_infos = self._push_notification_infos.setdefault(owner, {})
             if task_id not in owner_infos:
                 owner_infos[task_id] = []
 

src/a2a/server/tasks/inmemory_task_store.py

@@ -38,11 +38,9 @@ def _get_owner_tasks(self, owner: str) -> dict[str, Task]:
     async def save(self, task: Task, context: ServerCallContext) -> None:
         """Saves or updates a task in the in-memory store for the resolved owner."""
         owner = self.owner_resolver(context)
-        if owner not in self.tasks:
-            self.tasks[owner] = {}
-
         with self.lock:
-            self.tasks[owner][task.id] = task
+            owner_tasks = self.tasks.setdefault(owner, {})
+            owner_tasks[task.id] = task
             logger.debug(
                 'Task %s for owner %s saved successfully.', task.id, owner
             )

tests/server/tasks/test_inmemory_push_notifications.py

@@ -1,3 +1,6 @@
+import asyncio
+import concurrent.futures
+import threading
 import unittest
 
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -63,6 +66,29 @@ def user_name(self) -> str:
 MINIMAL_CALL_CONTEXT = ServerCallContext(user=SampleUser(user_name='user'))
 
 
+def _lock_is_owned(lock: threading.RLock) -> bool:
+    is_owned = getattr(lock, '_is_owned', None)
+    return bool(is_owned()) if callable(is_owned) else False
+
+
+def _set_info_in_thread(
+    store: InMemoryPushNotificationConfigStore,
+    task_id: str,
+    config_id: str,
+    context: ServerCallContext,
+) -> None:
+    asyncio.run(
+        store.set_info(
+            task_id,
+            _create_sample_push_config(
+                url=f'http://example.com/{config_id}',
+                config_id=config_id,
+            ),
+            context,
+        )
+    )
+
+
 class TestInMemoryPushNotifier(unittest.IsolatedAsyncioTestCase):
     def setUp(self) -> None:
         self.mock_httpx_client = AsyncMock(spec=httpx.AsyncClient)
@@ -426,6 +452,66 @@ async def test_owner_resource_scoping(self) -> None:
         await self.config_store.delete_info('task1', context=context_user1)
         await self.config_store.delete_info('task1', context=context_user2)
 
+    async def test_set_info_creates_owner_bucket_under_lock(self) -> None:
+        """Creating the first owner bucket must happen while the RLock is held."""
+        store = InMemoryPushNotificationConfigStore()
+        lock_held: list[bool] = []
+
+        class _LockHeldOwnerMap(
+            dict[str, dict[str, list[TaskPushNotificationConfig]]]
+        ):
+            def setdefault(
+                self,
+                key: str,
+                default: (
+                    dict[str, list[TaskPushNotificationConfig]] | None
+                ) = None,
+            ) -> dict[str, list[TaskPushNotificationConfig]]:
+                lock_held.append(_lock_is_owned(store.lock))
+                if default is None:
+                    default = {}
+                return super().setdefault(key, default)
+
+            def __setitem__(
+                self,
+                key: str,
+                value: dict[str, list[TaskPushNotificationConfig]],
+            ) -> None:
+                lock_held.append(_lock_is_owned(store.lock))
+                super().__setitem__(key, value)
+
+        store._push_notification_infos = _LockHeldOwnerMap()
+        await store.set_info(
+            'task-a',
+            _create_sample_push_config(config_id='cfg-a'),
+            MINIMAL_CALL_CONTEXT,
+        )
+        self.assertTrue(lock_held)
+        self.assertTrue(all(lock_held))
+
+    async def test_concurrent_first_owner_set_info_keeps_both(
+        self,
+    ) -> None:
+        """Concurrent first configs for a new owner must both persist."""
+        store = InMemoryPushNotificationConfigStore()
+        context = ServerCallContext(user=SampleUser(user_name='race-owner'))
+        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
+            futures = [
+                pool.submit(
+                    _set_info_in_thread, store, 'task-a', 'cfg-a', context
+                ),
+                pool.submit(
+                    _set_info_in_thread, store, 'task-b', 'cfg-b', context
+                ),
+            ]
+            for future in futures:
+                future.result(timeout=10)
+
+        configs_a = await store.get_info('task-a', context)
+        configs_b = await store.get_info('task-b', context)
+        self.assertEqual([config.id for config in configs_a], ['cfg-a'])
+        self.assertEqual([config.id for config in configs_b], ['cfg-b'])
+
 
 class TestPushNotificationDispatchAcrossOwners(
     unittest.IsolatedAsyncioTestCase

tests/server/tasks/test_inmemory_task_store.py

@@ -1,3 +1,7 @@
+import asyncio
+import concurrent.futures
+import threading
+
 from datetime import datetime, timezone
 
 import pytest
@@ -368,3 +372,61 @@ async def test_inmemory_task_store_copying_behavior(use_copying: bool):
     else:
         assert retrieved_task_2.status.state == TaskState.TASK_STATE_COMPLETED
         assert retrieved_task_2 is retrieved_task
+
+
+def _lock_is_owned(lock: threading.RLock) -> bool:
+    is_owned = getattr(lock, '_is_owned', None)
+    return bool(is_owned()) if callable(is_owned) else False
+
+
+def _save_task_in_thread(
+    store: InMemoryTaskStore,
+    task_id: str,
+    context: ServerCallContext,
+) -> None:
+    asyncio.run(store.save(create_minimal_task(task_id=task_id), context))
+
+
+def test_save_creates_owner_bucket_under_lock() -> None:
+    """Creating the first owner bucket must happen while the RLock is held."""
+    store = InMemoryTaskStore(use_copying=False)
+    impl = store._impl
+    lock_held: list[bool] = []
+
+    class _LockHeldOwnerMap(dict[str, dict[str, Task]]):
+        def setdefault(
+            self,
+            key: str,
+            default: dict[str, Task] | None = None,
+        ) -> dict[str, Task]:
+            lock_held.append(_lock_is_owned(impl.lock))
+            if default is None:
+                default = {}
+            return super().setdefault(key, default)
+
+        def __setitem__(self, key: str, value: dict[str, Task]) -> None:
+            lock_held.append(_lock_is_owned(impl.lock))
+            super().__setitem__(key, value)
+
+    impl.tasks = _LockHeldOwnerMap()
+    asyncio.run(store.save(create_minimal_task(), TEST_CONTEXT))
+    assert lock_held
+    assert all(lock_held)
+
+
+def test_concurrent_first_owner_saves_keep_both_tasks() -> None:
+    """Concurrent first saves for a new owner must keep both tasks."""
+    store = InMemoryTaskStore(use_copying=False)
+    context = ServerCallContext(user=SampleUser('race-owner'))
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
+        futures = [
+            pool.submit(_save_task_in_thread, store, 'task-a', context),
+            pool.submit(_save_task_in_thread, store, 'task-b', context),
+        ]
+        for future in futures:
+            future.result(timeout=10)
+
+    page = asyncio.run(store.list(ListTasksRequest(), context))
+    assert {task.id for task in page.tasks} == {'task-a', 'task-b'}
+    assert asyncio.run(store.get('task-a', context)) is not None
+    assert asyncio.run(store.get('task-b', context)) is not None