feat: support extended_agent_card in A2AStarletteRouter

Signed-off-by: Shingo OKAWA <[email protected]>

Author: ognis1205

src/a2a/server/apps/starlette_app.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import traceback
+
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator
 from typing import Any
@@ -15,16 +16,29 @@
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.jsonrpc_handler import JSONRPCHandler
 from a2a.server.request_handlers.request_handler import RequestHandler
-from a2a.types import (A2AError, A2ARequest, AgentCard, CancelTaskRequest,
-                       GetTaskPushNotificationConfigRequest, GetTaskRequest,
-                       InternalError, InvalidRequestError, JSONParseError,
-                       JSONRPCError, JSONRPCErrorResponse, JSONRPCResponse,
-                       SendMessageRequest, SendStreamingMessageRequest,
-                       SendStreamingMessageResponse,
-                       SetTaskPushNotificationConfigRequest,
-                       TaskResubscriptionRequest, UnsupportedOperationError)
+from a2a.types import (
+    A2AError,
+    A2ARequest,
+    AgentCard,
+    CancelTaskRequest,
+    GetTaskPushNotificationConfigRequest,
+    GetTaskRequest,
+    InternalError,
+    InvalidRequestError,
+    JSONParseError,
+    JSONRPCError,
+    JSONRPCErrorResponse,
+    JSONRPCResponse,
+    SendMessageRequest,
+    SendStreamingMessageRequest,
+    SendStreamingMessageResponse,
+    SetTaskPushNotificationConfigRequest,
+    TaskResubscriptionRequest,
+    UnsupportedOperationError,
+)
 from a2a.utils.errors import MethodNotImplementedError
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -344,7 +358,9 @@ async def _handle_get_authenticated_extended_agent_card(
         # extended_agent_card was provided during server initialization,
         # return a 404
         return JSONResponse(
-            {'error': 'Authenticated extended agent card is supported but not configured on the server.'},
+            {
+                'error': 'Authenticated extended agent card is supported but not configured on the server.'
+            },
             status_code=404,
         )
 
@@ -437,24 +453,34 @@ def __init__(
         self,
         agent_card: AgentCard,
         http_handler: RequestHandler,
-        agent_card_path: str = '/agent.json',
-        rpc_path: str = '/',
+        extended_agent_card: AgentCard | None = None,
+        context_builder: CallContextBuilder | None = None,
     ):
         """Initializes the A2AStarletteRouter.
 
         Args:
             agent_card: The AgentCard describing the agent's capabilities.
             http_handler: The handler instance responsible for processing A2A
               requests via http.
-            agent_card_path: The URL path for the agent card endpoint.
-            rpc_path: The URL path for the A2A JSON-RPC endpoint (POST requests).
+            extended_agent_card: An optional, distinct AgentCard to be served
+              at the authenticated extended card endpoint.
+            context_builder: The CallContextBuilder used to construct the
+              ServerCallContext passed to the http_handler. If None, no
+              ServerCallContext is passed.
         """
         self.agent_card = agent_card
+        self.extended_agent_card = extended_agent_card
         self.handler = JSONRPCHandler(
             agent_card=agent_card, request_handler=http_handler
         )
-        self.agent_card_path = agent_card_path
-        self.rpc_path = rpc_path
+        if (
+            self.agent_card.supportsAuthenticatedExtendedCard
+            and self.extended_agent_card is None
+        ):
+            logger.error(
+                'AgentCard.supportsAuthenticatedExtendedCard is True, but no extended_agent_card was provided. The /agent/authenticatedExtendedCard endpoint will return 404.'
+            )
+        self._context_builder = context_builder
 
     def _generate_error_response(
         self, request_id: str | int | None, error: JSONRPCError | A2AError
@@ -474,7 +500,6 @@ def _generate_error_response(
             id=request_id,
             error=error if isinstance(error, JSONRPCError) else error.root,
         )
-
         log_level = (
             logging.ERROR
             if not isinstance(error, A2AError)
@@ -512,24 +537,25 @@ async def _handle_requests(self, request: Request) -> Response:
         """
         request_id = None
         body = None
-
         try:
             body = await request.json()
             a2a_request = A2ARequest.model_validate(body)
-
+            call_context = (
+                self._context_builder.build(request)
+                if self._context_builder
+                else None
+            )
             request_id = a2a_request.root.id
             request_obj = a2a_request.root
-
             if isinstance(
                 request_obj,
                 TaskResubscriptionRequest | SendStreamingMessageRequest,
             ):
                 return await self._process_streaming_request(
-                    request_id, a2a_request
+                    request_id, a2a_request, call_context
                 )
-
             return await self._process_non_streaming_request(
-                request_id, a2a_request
+                request_id, a2a_request, call_context
             )
         except MethodNotImplementedError:
             traceback.print_exc()
@@ -555,7 +581,10 @@ async def _handle_requests(self, request: Request) -> Response:
             )
 
     async def _process_streaming_request(
-        self, request_id: str | int | None, a2a_request: A2ARequest
+        self,
+        request_id: str | int | None,
+        a2a_request: A2ARequest,
+        context: ServerCallContext,
     ) -> Response:
         """Processes streaming requests (message/stream or tasks/resubscribe).
 
@@ -572,14 +601,20 @@ async def _process_streaming_request(
             request_obj,
             SendStreamingMessageRequest,
         ):
-            handler_result = self.handler.on_message_send_stream(request_obj)
+            handler_result = self.handler.on_message_send_stream(
+                request_obj, context
+            )
         elif isinstance(request_obj, TaskResubscriptionRequest):
-            handler_result = self.handler.on_resubscribe_to_task(request_obj)
-
+            handler_result = self.handler.on_resubscribe_to_task(
+                request_obj, context
+            )
         return self._create_response(handler_result)
 
     async def _process_non_streaming_request(
-        self, request_id: str | int | None, a2a_request: A2ARequest
+        self,
+        request_id: str | int | None,
+        a2a_request: A2ARequest,
+        context: ServerCallContext,
     ) -> Response:
         """Processes non-streaming requests (message/send, tasks/get, tasks/cancel, tasks/pushNotificationConfig/*).
 
@@ -594,18 +629,26 @@ async def _process_non_streaming_request(
         handler_result: Any = None
         match request_obj:
             case SendMessageRequest():
-                handler_result = await self.handler.on_message_send(request_obj)
+                handler_result = await self.handler.on_message_send(
+                    request_obj, context
+                )
             case CancelTaskRequest():
-                handler_result = await self.handler.on_cancel_task(request_obj)
+                handler_result = await self.handler.on_cancel_task(
+                    request_obj, context
+                )
             case GetTaskRequest():
-                handler_result = await self.handler.on_get_task(request_obj)
+                handler_result = await self.handler.on_get_task(
+                    request_obj, context
+                )
             case SetTaskPushNotificationConfigRequest():
                 handler_result = await self.handler.set_push_notification(
-                    request_obj
+                    request_obj,
+                    context,
                 )
             case GetTaskPushNotificationConfigRequest():
                 handler_result = await self.handler.get_push_notification(
-                    request_obj
+                    request_obj,
+                    context,
                 )
             case _:
                 logger.error(
@@ -617,7 +660,6 @@ async def _process_non_streaming_request(
                 handler_result = JSONRPCErrorResponse(
                     id=request_id, error=error
                 )
-
         return self._create_response(handler_result)
 
     def _create_response(
@@ -659,7 +701,6 @@ async def event_generator(
                     exclude_none=True,
                 )
             )
-
         return JSONResponse(
             handler_result.root.model_dump(mode='json', exclude_none=True)
         )
@@ -673,27 +714,75 @@ async def _handle_get_agent_card(self, request: Request) -> JSONResponse:
         Returns:
             A JSONResponse containing the agent card data.
         """
+        # The public agent card is a direct serialization of the agent_card
+        # provided at initialization.
         return JSONResponse(
             self.agent_card.model_dump(mode='json', exclude_none=True)
         )
 
-    def routes(self) -> list[Route]:
+    async def _handle_get_authenticated_extended_agent_card(
+        self, request: Request
+    ) -> JSONResponse:
+        """Handles GET requests for the authenticated extended agent card."""
+        if not self.agent_card.supportsAuthenticatedExtendedCard:
+            return JSONResponse(
+                {'error': 'Extended agent card not supported or not enabled.'},
+                status_code=404,
+            )
+        # If an explicit extended_agent_card is provided, serve that.
+        if self.extended_agent_card:
+            return JSONResponse(
+                self.extended_agent_card.model_dump(
+                    mode='json', exclude_none=True
+                )
+            )
+        # If supportsAuthenticatedExtendedCard is true, but no specific
+        # extended_agent_card was provided during server initialization,
+        # return a 404
+        return JSONResponse(
+            {
+                'error': 'Authenticated extended agent card is supported but not configured on the server.'
+            },
+            status_code=404,
+        )
+
+    def routes(
+        self,
+        agent_card_path: str = '/.well-known/agent.json',
+        extended_agent_card_path: str = '/agent/authenticatedExtendedCard',
+        rpc_path: str = '/',
+    ) -> list[Route]:
         """Returns the Starlette Routes for handling A2A requests.
 
+        Args:
+            agent_card_path: The URL path for the agent card endpoint.
+            rpc_path: The URL path for the A2A JSON-RPC endpoint (POST requests).
+            extended_agent_card_path: The URL path for the authenticated extended agent card endpoint.
+
         Returns:
             A list of Starlette Route objects.
         """
-        return [
+        routes = [
             Route(
-                self.rpc_path,
+                rpc_path,
                 self._handle_requests,
                 methods=['POST'],
                 name='a2a_handler',
             ),
             Route(
-                self.agent_card_path,
+                agent_card_path,
                 self._handle_get_agent_card,
                 methods=['GET'],
                 name='agent_card',
             ),
         ]
+        if self.agent_card.supportsAuthenticatedExtendedCard:
+            routes.append(
+                Route(
+                    extended_agent_card_path,
+                    self._handle_get_authenticated_extended_agent_card,
+                    methods=['GET'],
+                    name='authenticated_extended_agent_card',
+                )
+            )
+        return routes