fix(server)!: deliver push notifications to all registrars across owners

Author: sokoliva

itk/main.py

@@ -31,7 +31,6 @@
     InMemoryPushNotificationConfigStore,
 )
 from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
-from a2a.server.context import ServerCallContext
 from a2a.types import a2a_pb2_grpc
 from a2a.types.a2a_pb2 import (
     AgentCapabilities,
@@ -339,7 +338,6 @@ async def main_async(http_port: int, grpc_port: int) -> None:
     push_sender = BasePushNotificationSender(
         httpx_client=httpx.AsyncClient(),
         config_store=push_config_store,
-        context=ServerCallContext(),
     )
 
     handler = DefaultRequestHandler(

src/a2a/server/tasks/base_push_notification_sender.py

@@ -5,7 +5,6 @@
 
 from google.protobuf.json_format import MessageToDict
 
-from a2a.server.context import ServerCallContext
 from a2a.server.tasks.push_notification_config_store import (
     PushNotificationConfigStore,
 )
@@ -27,26 +26,21 @@ def __init__(
         self,
         httpx_client: httpx.AsyncClient,
         config_store: PushNotificationConfigStore,
-        context: ServerCallContext,
     ) -> None:
         """Initializes the BasePushNotificationSender.
 
         Args:
             httpx_client: An async HTTP client instance to send notifications.
             config_store: A PushNotificationConfigStore instance to retrieve configurations.
-            context: The `ServerCallContext` that this push notification is produced under.
         """
         self._client = httpx_client
         self._config_store = config_store
-        self._call_context: ServerCallContext = context
 
     async def send_notification(
         self, task_id: str, event: PushNotificationEvent
     ) -> None:
         """Sends a push notification for an event if configuration exists."""
-        push_configs = await self._config_store.get_info(
-            task_id, self._call_context
-        )
+        push_configs = await self._config_store.get_info_for_dispatch(task_id)
         if not push_configs:
             return
 

src/a2a/server/tasks/database_push_notification_config_store.py

@@ -309,7 +309,10 @@ async def get_info(
         task_id: str,
         context: ServerCallContext,
     ) -> list[TaskPushNotificationConfig]:
-        """Retrieves all push notification configurations for a task, for the given owner."""
+        """Retrieves all push notification configurations for a task, for the given owner.
+
+        Used by the user-callable read endpoints.
+        """
         await self._ensure_initialized()
         owner = self.owner_resolver(context)
         async with self.async_session_maker() as session:
@@ -335,6 +338,35 @@ async def get_info(
                     )
             return configs
 
+    async def get_info_for_dispatch(
+        self,
+        task_id: str,
+    ) -> list[TaskPushNotificationConfig]:
+        """Retrieves all push notification configurations for a task, across all owners.
+
+        Used by the push-notification dispatch path.
+        """
+        await self._ensure_initialized()
+        async with self.async_session_maker() as session:
+            stmt = select(self.config_model).where(
+                self.config_model.task_id == task_id
+            )
+            result = await session.execute(stmt)
+            models = result.scalars().all()
+
+            configs = []
+            for model in models:
+                try:
+                    configs.append(self._from_orm(model))
+                except ValueError:  # noqa: PERF203
+                    logger.exception(
+                        'Could not deserialize push notification config for task %s, config %s, owner %s',
+                        model.task_id,
+                        model.config_id,
+                        model.owner,
+                    )
+            return configs
+
     async def delete_info(
         self,
         task_id: str,

src/a2a/server/tasks/inmemory_push_notification_config_store.py

@@ -72,12 +72,29 @@ async def get_info(
         task_id: str,
         context: ServerCallContext,
     ) -> list[TaskPushNotificationConfig]:
-        """Retrieves all push notification configurations for a task from memory, for the given owner."""
+        """Retrieves all push notification configurations for a task from memory, for the given owner.
+
+        Used by the user-callable read endpoints.
+        """
         owner = self.owner_resolver(context)
         async with self.lock:
             owner_infos = self._get_owner_push_notification_infos(owner)
             return list(owner_infos.get(task_id, []))
 
+    async def get_info_for_dispatch(
+        self,
+        task_id: str,
+    ) -> list[TaskPushNotificationConfig]:
+        """Retrieves all push notification configurations for a task across all owners.
+
+        Used by the push-notification dispatch path.
+        """
+        async with self.lock:
+            results: list[TaskPushNotificationConfig] = []
+            for all_configs in self._push_notification_infos.values():
+                results.extend(all_configs.get(task_id, []))
+            return results
+
     async def delete_info(
         self,
         task_id: str,

src/a2a/server/tasks/push_notification_config_store.py

@@ -22,7 +22,26 @@ async def get_info(
         task_id: str,
         context: ServerCallContext,
     ) -> list[TaskPushNotificationConfig]:
-        """Retrieves the push notification configuration for a task."""
+        """Retrieves push notification configurations for a task, scoped to the caller.
+
+        This is the user-callable read path. Implementations MUST return
+        only configurations owned by the caller (as resolved from
+        context).
+        """
+
+    @abstractmethod
+    async def get_info_for_dispatch(
+        self,
+        task_id: str,
+    ) -> list[TaskPushNotificationConfig]:
+        """Retrieves all push notification configurations for a task, across all owners.
+
+        This is the internal read path used by the push-notification
+        dispatch loop. Implementations MUST return every configuration
+        registered for task_id regardless of which user registered
+        it. Authorization already happened at registration time and
+        the dispatch path fires every registered webhook for the task.
+        """
 
     @abstractmethod
     async def delete_info(

tests/e2e/push_notifications/agent_app.py

@@ -1,14 +1,17 @@
 import httpx
 
 from fastapi import FastAPI
+from starlette.applications import Starlette
+from starlette.requests import Request
 
+from a2a.auth.user import UnauthenticatedUser, User
 from a2a.server.agent_execution import AgentExecutor, RequestContext
 from a2a.server.context import ServerCallContext
 from a2a.server.events import EventQueue
-from starlette.applications import Starlette
-from a2a.server.routes.rest_routes import create_rest_routes
-from a2a.server.routes import create_agent_card_routes
 from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.routes import create_agent_card_routes
+from a2a.server.routes.common import DefaultServerCallContextBuilder
+from a2a.server.routes.rest_routes import create_rest_routes
 from a2a.server.tasks import (
     BasePushNotificationSender,
     InMemoryPushNotificationConfigStore,
@@ -30,6 +33,9 @@
 )
 
 
+_TEST_USER_HEADER = 'x-test-user'
+
+
 def test_agent_card(url: str) -> AgentCard:
     """Returns an agent card for the test agent."""
     return AgentCard(
@@ -151,11 +157,85 @@ def create_agent_app(
         push_sender=BasePushNotificationSender(
             httpx_client=notification_client,
             config_store=push_config_store,
-            context=ServerCallContext(),
         ),
     )
     rest_routes = create_rest_routes(request_handler=handler)
     agent_card_routes = create_agent_card_routes(
         agent_card=card, card_url='/.well-known/agent-card.json'
     )
     return Starlette(routes=[*rest_routes, *agent_card_routes])
+
+
+class _NamedTestUser(User):
+    """Authenticated test user identified by ``user_name``."""
+
+    def __init__(self, user_name: str) -> None:
+        self._user_name = user_name
+
+    @property
+    def is_authenticated(self) -> bool:
+        return True
+
+    @property
+    def user_name(self) -> str:
+        return self._user_name
+
+
+class _HeaderUserContextBuilder(DefaultServerCallContextBuilder):
+    """Builds a ServerCallContext whose user is read from a request header."""
+
+    def build_user(self, request: Request) -> User:
+        user_name = request.headers.get(_TEST_USER_HEADER)
+        if user_name:
+            return _NamedTestUser(user_name)
+        return UnauthenticatedUser()
+
+
+def create_multi_user_agent_app(
+    url: str, notification_client: httpx.AsyncClient
+) -> Starlette:
+    """Creates a multi-user variant of the test agent app.
+
+    Differences from create_agent_app:
+
+    - Identity is read from the x-test-user header on each request
+      via _HeaderUserContextBuilder. Multiple authenticated
+      users (e.g. alice, bob) can therefore call the same
+      server.
+    - The InMemoryTaskStore uses a constant owner resolver, so
+      every authenticated user has access to every task.
+    - The InMemoryPushNotificationConfigStore keeps the default
+      per-user owner resolver, so each registrar's configs live in their
+      own owner partition; this exercises cross-owner aggregation in
+      get_info_for_dispatch.
+    """
+    # Shared task visibility: any authenticated user can see any task.
+    task_store = InMemoryTaskStore(owner_resolver=lambda _ctx: 'shared')
+
+    # Per-user push-config partitioning (the default).
+    push_config_store = InMemoryPushNotificationConfigStore()
+
+    card = test_agent_card(url)
+    extended_card = test_agent_card(url)
+    extended_card.name = 'Test Agent Extended'
+
+    handler = DefaultRequestHandler(
+        agent_executor=TestAgentExecutor(),
+        task_store=task_store,
+        agent_card=card,
+        extended_agent_card=extended_card,
+        push_config_store=push_config_store,
+        push_sender=BasePushNotificationSender(
+            httpx_client=notification_client,
+            config_store=push_config_store,
+        ),
+    )
+
+    rest_routes = create_rest_routes(
+        request_handler=handler,
+        context_builder=_HeaderUserContextBuilder(),
+    )
+    agent_card_routes = create_agent_card_routes(
+        agent_card=card, card_url='/.well-known/agent-card.json'
+    )
+    return Starlette(routes=[*rest_routes, *agent_card_routes])