feat: Add GetExtendedAgentCard Support to RequestHandlers

Author: sokoliva

itk/main.py

@@ -292,7 +292,16 @@ async def main_async(http_port: int, grpc_port: int) -> None:
     handler = DefaultRequestHandler(
         agent_executor=V10AgentExecutor(),
         task_store=task_store,
+        agent_card=agent_card,
+        queue_manager=InMemoryQueueManager(),
+    )
+
+    handler_extended = DefaultRequestHandler(
+        agent_executor=V10AgentExecutor(),
+        task_store=task_store,
+        agent_card=agent_card,
         queue_manager=InMemoryQueueManager(),
+        extended_agent_card=agent_card,
     )
 
     app = FastAPI()
@@ -301,9 +310,7 @@ async def main_async(http_port: int, grpc_port: int) -> None:
         agent_card=agent_card, card_url='/.well-known/agent-card.json'
     )
     jsonrpc_routes = create_jsonrpc_routes(
-        agent_card=agent_card,
-        request_handler=handler,
-        extended_agent_card=agent_card,
+        request_handler=handler_extended,
         rpc_url='/',
         enable_v0_3_compat=True,
     )
@@ -313,7 +320,6 @@ async def main_async(http_port: int, grpc_port: int) -> None:
     )
 
     rest_routes = create_rest_routes(
-        agent_card=agent_card,
         request_handler=handler,
         enable_v0_3_compat=True,
     )
@@ -323,7 +329,7 @@ async def main_async(http_port: int, grpc_port: int) -> None:
 
     compat_servicer = CompatGrpcHandler(agent_card, handler)
     a2a_v0_3_pb2_grpc.add_A2AServiceServicer_to_server(compat_servicer, server)
-    servicer = GrpcHandler(agent_card, handler)
+    servicer = GrpcHandler(handler)
     a2a_pb2_grpc.add_A2AServiceServicer_to_server(servicer, server)
 
     server.add_insecure_port(f'127.0.0.1:{grpc_port}')

samples/hello_world_agent.py

@@ -191,17 +191,17 @@ async def serve(
 
     task_store = InMemoryTaskStore()
     request_handler = DefaultRequestHandler(
-        agent_executor=SampleAgentExecutor(), task_store=task_store
+        agent_executor=SampleAgentExecutor(),
+        task_store=task_store,
+        agent_card=agent_card,
     )
 
     rest_routes = create_rest_routes(
-        agent_card=agent_card,
         request_handler=request_handler,
         path_prefix='/a2a/rest',
         enable_v0_3_compat=True,
     )
     jsonrpc_routes = create_jsonrpc_routes(
-        agent_card=agent_card,
         request_handler=request_handler,
         rpc_url='/a2a/jsonrpc',
         enable_v0_3_compat=True,
@@ -216,7 +216,7 @@ async def serve(
 
     grpc_server = grpc.aio.server()
     grpc_server.add_insecure_port(f'{host}:{grpc_port}')
-    servicer = GrpcHandler(agent_card, request_handler)
+    servicer = GrpcHandler(request_handler)
     a2a_pb2_grpc.add_A2AServiceServicer_to_server(servicer, grpc_server)
 
     compat_grpc_server = grpc.aio.server()

src/a2a/server/request_handlers/default_request_handler.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from typing import cast
 
 from a2a.server.agent_execution import (
@@ -31,8 +31,10 @@
     TaskStore,
 )
 from a2a.types.a2a_pb2 import (
+    AgentCard,
     CancelTaskRequest,
     DeleteTaskPushNotificationConfigRequest,
+    GetExtendedAgentCardRequest,
     GetTaskPushNotificationConfigRequest,
     GetTaskRequest,
     ListTaskPushNotificationConfigsRequest,
@@ -47,13 +49,16 @@
     TaskState,
 )
 from a2a.utils.errors import (
+    ExtendedAgentCardNotConfiguredError,
     InternalError,
     InvalidParamsError,
+    PushNotificationConfigNotFoundError,
     PushNotificationNotSupportedError,
     TaskNotCancelableError,
     TaskNotFoundError,
     UnsupportedOperationError,
 )
+from a2a.utils.helpers import maybe_await, validate
 from a2a.utils.task import (
     apply_history_length,
     validate_history_length,
@@ -88,27 +93,43 @@ def __init__(  # noqa: PLR0913
         self,
         agent_executor: AgentExecutor,
         task_store: TaskStore,
+        agent_card: AgentCard,
         queue_manager: QueueManager | None = None,
         push_config_store: PushNotificationConfigStore | None = None,
         push_sender: PushNotificationSender | None = None,
         request_context_builder: RequestContextBuilder | None = None,
+        extended_agent_card: AgentCard | None = None,
+        card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
+        | None = None,
+        extended_card_modifier: Callable[
+            [AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
+        ]
+        | None = None,
     ) -> None:
         """Initializes the DefaultRequestHandler.
 
         Args:
             agent_executor: The `AgentExecutor` instance to run agent logic.
             task_store: The `TaskStore` instance to manage task persistence.
+            agent_card: The AgentCard describing the agent's capabilities.
             queue_manager: The `QueueManager` instance to manage event queues. Defaults to `InMemoryQueueManager`.
             push_config_store: The `PushNotificationConfigStore` instance for managing push notification configurations. Defaults to None.
             push_sender: The `PushNotificationSender` instance for sending push notifications. Defaults to None.
             request_context_builder: The `RequestContextBuilder` instance used
               to build request contexts. Defaults to `SimpleRequestContextBuilder`.
+            extended_agent_card: An optional, distinct AgentCard to be served at the authenticated extended card endpoint.
+            card_modifier: An optional callback to dynamically modify the public agent card before it is served.
+            extended_card_modifier: An optional callback to dynamically modify the extended agent card before it is served. It receives the call context.
         """
         self.agent_executor = agent_executor
         self.task_store = task_store
+        self._agent_card = agent_card
         self._queue_manager = queue_manager or InMemoryQueueManager()
         self._push_config_store = push_config_store
         self._push_sender = push_sender
+        self.extended_agent_card = extended_agent_card
+        self.card_modifier = card_modifier
+        self.extended_card_modifier = extended_card_modifier
         self._request_context_builder = (
             request_context_builder
             or SimpleRequestContextBuilder(
@@ -122,6 +143,11 @@ def __init__(  # noqa: PLR0913
         # asyncio tasks and to surface unexpected exceptions.
         self._background_tasks = set()
 
+    @property
+    def agent_card(self) -> AgentCard:
+        """The agent card to be served by default."""
+        return self._agent_card
+
     @validate_request_params
     async def on_get_task(
         self,
@@ -396,6 +422,10 @@ async def push_notification_callback(event: Event) -> None:
         return result
 
     @validate_request_params
+    @validate(
+        lambda self: self.agent_card.capabilities.streaming,
+        'Streaming is not supported by the agent',
+    )
     async def on_message_send_stream(
         self,
         params: SendMessageRequest,
@@ -485,6 +515,14 @@ async def _cleanup_producer(
             self._running_agents.pop(task_id, None)
 
     @validate_request_params
+    @validate(
+        lambda self: (
+            self.agent_card.capabilities.push_notifications
+            and self._push_config_store
+        ),
+        error_message='Push notifications are not supported by the agent',
+        error_type=PushNotificationNotSupportedError,
+    )
     async def on_create_task_push_notification_config(
         self,
         params: TaskPushNotificationConfig,
@@ -494,15 +532,12 @@ async def on_create_task_push_notification_config(
 
         Requires a `PushNotifier` to be configured.
         """
-        if not self._push_config_store:
-            raise PushNotificationNotSupportedError
-
         task_id = params.task_id
         task: Task | None = await self.task_store.get(task_id, context)
         if not task:
             raise TaskNotFoundError
 
-        await self._push_config_store.set_info(
+        await self._push_config_store.set_info(  # type: ignore[union-attr]
             task_id,
             params,
             context,
@@ -511,6 +546,14 @@ async def on_create_task_push_notification_config(
         return params
 
     @validate_request_params
+    @validate(
+        lambda self: (
+            self.agent_card.capabilities.push_notifications
+            and self._push_config_store
+        ),
+        error_message='Push notifications are not supported by the agent',
+        error_type=PushNotificationNotSupportedError,
+    )
     async def on_get_task_push_notification_config(
         self,
         params: GetTaskPushNotificationConfigRequest,
@@ -520,26 +563,27 @@ async def on_get_task_push_notification_config(
 
         Requires a `PushConfigStore` to be configured.
         """
-        if not self._push_config_store:
-            raise PushNotificationNotSupportedError
-
         task_id = params.task_id
         config_id = params.id
         task: Task | None = await self.task_store.get(task_id, context)
         if not task:
             raise TaskNotFoundError
 
         push_notification_configs: list[TaskPushNotificationConfig] = (
-            await self._push_config_store.get_info(task_id, context) or []
+            await self._push_config_store.get_info(task_id, context) or []  # type: ignore[union-attr]
         )
 
         for config in push_notification_configs:
             if config.id == config_id:
                 return config
 
-        raise InternalError(message='Push notification config not found')
+        raise PushNotificationConfigNotFoundError
 
     @validate_request_params
+    @validate(
+        lambda self: self.agent_card.capabilities.streaming,
+        'Streaming is not supported by the agent',
+    )
     async def on_subscribe_to_task(
         self,
         params: SubscribeToTaskRequest,
@@ -583,6 +627,14 @@ async def on_subscribe_to_task(
             yield event
 
     @validate_request_params
+    @validate(
+        lambda self: (
+            self.agent_card.capabilities.push_notifications
+            and self._push_config_store
+        ),
+        error_message='Push notifications are not supported by the agent',
+        error_type=PushNotificationNotSupportedError,
+    )
     async def on_list_task_push_notification_configs(
         self,
         params: ListTaskPushNotificationConfigsRequest,
@@ -592,15 +644,12 @@ async def on_list_task_push_notification_configs(
 
         Requires a `PushConfigStore` to be configured.
         """
-        if not self._push_config_store:
-            raise PushNotificationNotSupportedError
-
         task_id = params.task_id
         task: Task | None = await self.task_store.get(task_id, context)
         if not task:
             raise TaskNotFoundError
 
-        push_notification_config_list = await self._push_config_store.get_info(
+        push_notification_config_list = await self._push_config_store.get_info(  # type: ignore[union-attr]
             task_id, context
         )
 
@@ -609,6 +658,14 @@ async def on_list_task_push_notification_configs(
         )
 
     @validate_request_params
+    @validate(
+        lambda self: (
+            self.agent_card.capabilities.push_notifications
+            and self._push_config_store
+        ),
+        error_message='Push notifications are not supported by the agent',
+        error_type=PushNotificationNotSupportedError,
+    )
     async def on_delete_task_push_notification_config(
         self,
         params: DeleteTaskPushNotificationConfigRequest,
@@ -618,13 +675,35 @@ async def on_delete_task_push_notification_config(
 
         Requires a `PushConfigStore` to be configured.
         """
-        if not self._push_config_store:
-            raise PushNotificationNotSupportedError
-
         task_id = params.task_id
         config_id = params.id
         task: Task | None = await self.task_store.get(task_id, context)
         if not task:
             raise TaskNotFoundError
 
-        await self._push_config_store.delete_info(task_id, context, config_id)
+        await self._push_config_store.delete_info(task_id, context, config_id)  # type: ignore[union-attr]
+
+    @validate_request_params
+    @validate(
+        lambda self: self.agent_card.capabilities.extended_agent_card,
+        error_message='The agent does not have an extended agent card configured',
+        error_type=ExtendedAgentCardNotConfiguredError,
+    )
+    async def on_get_extended_agent_card(
+        self,
+        params: GetExtendedAgentCardRequest,
+        context: ServerCallContext,
+    ) -> AgentCard:
+        """Default handler for 'GetExtendedAgentCard'.
+
+        Requires `capabilities.extended_agent_card` to be true.
+        """
+        card = self.extended_agent_card or self.agent_card
+
+        if self.extended_card_modifier:
+            return await maybe_await(self.extended_card_modifier(card, context))
+
+        if self.card_modifier:
+            return await maybe_await(self.card_modifier(card))
+
+        return card