make extended_card_modifier async

Author: cbugwadia32

src/a2a/server/apps/rest/fastapi_app.py

@@ -1,9 +1,7 @@
 import logging
-
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any
 
-
 if TYPE_CHECKING:
     from fastapi import APIRouter, FastAPI, Request, Response
     from fastapi.responses import JSONResponse
@@ -31,7 +29,6 @@
 from a2a.types import AgentCard
 from a2a.utils.constants import AGENT_CARD_WELL_KNOWN_PATH
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -51,7 +48,7 @@ def __init__(  # noqa: PLR0913
         context_builder: CallContextBuilder | None = None,
         card_modifier: Callable[[AgentCard], AgentCard] | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard]
         ]
         | None = None,
     ):
@@ -68,7 +65,7 @@ def __init__(  # noqa: PLR0913
               ServerCallContext is passed.
             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
+            extended_card_modifier: An optional async callback to dynamically modify
               the extended agent card before it is served. It receives the
               call context.
         """

src/a2a/server/apps/rest/rest_adapter.py

@@ -60,7 +60,7 @@ def __init__(  # noqa: PLR0913
         context_builder: CallContextBuilder | None = None,
         card_modifier: Callable[[AgentCard], AgentCard] | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard]
         ]
         | None = None,
     ):
@@ -77,9 +77,9 @@ def __init__(  # noqa: PLR0913
               ServerCallContext is passed.
             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.
+            extended_card_modifier: An optional async callback to dynamically
+              modify the extended agent card before it is served. It receives
+              the call context.
         """
         if not _package_starlette_installed:
             raise ImportError(
@@ -182,7 +182,9 @@ async def handle_authenticated_agent_card(
 
         if self.extended_card_modifier:
             context = self._context_builder.build(request)
-            card_to_serve = self.extended_card_modifier(card_to_serve, context)
+            card_to_serve = await self.extended_card_modifier(
+                card_to_serve, context
+            )
         elif self.card_modifier:
             card_to_serve = self.card_modifier(card_to_serve)
 

tests/server/apps/rest/test_rest_fastapi_app.py

@@ -1,10 +1,9 @@
 import logging
-
+from collections.abc import Awaitable, Callable
 from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
-
 from fastapi import FastAPI
 from google.protobuf import json_format
 from httpx import ASGITransport, AsyncClient
@@ -13,6 +12,7 @@
 from a2a.server.apps.rest import fastapi_app, rest_adapter
 from a2a.server.apps.rest.fastapi_app import A2ARESTFastAPIApplication
 from a2a.server.apps.rest.rest_adapter import RESTAdapter
+from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.types import (
     AgentCard,
@@ -25,15 +25,14 @@
     TextPart,
 )
 
-
 logger = logging.getLogger(__name__)
 
 
 @pytest.fixture
 async def agent_card() -> AgentCard:
     mock_agent_card = MagicMock(spec=AgentCard)
     mock_agent_card.url = 'http://mockurl.com'
-    mock_agent_card.supports_authenticated_extended_card = False
+    mock_agent_card.supports_authenticated_extended_card = True
 
     # Mock the capabilities object with streaming disabled
     mock_capabilities = MagicMock()
@@ -58,6 +57,15 @@ async def streaming_agent_card() -> AgentCard:
     return mock_agent_card
 
 
+@pytest.fixture
+async def extended_card_modifier() -> Callable[
+    [AgentCard, ServerCallContext], Awaitable[AgentCard]
+]:
+    return MagicMock(
+        spec=Callable[[AgentCard, ServerCallContext], Awaitable[AgentCard]]
+    )
+
+
 @pytest.fixture
 async def request_handler() -> RequestHandler:
     return MagicMock(spec=RequestHandler)
@@ -84,13 +92,20 @@ async def streaming_client(streaming_app: FastAPI) -> AsyncClient:
 
 @pytest.fixture
 async def app(
-    agent_card: AgentCard, request_handler: RequestHandler
+    agent_card: AgentCard,
+    request_handler: RequestHandler,
+    extended_card_modifier: Callable[
+        [AgentCard, ServerCallContext], Awaitable[AgentCard]
+    ],
 ) -> FastAPI:
     """Builds the FastAPI application for testing."""
 
-    return A2ARESTFastAPIApplication(agent_card, request_handler).build(
-        agent_card_url='/well-known/agent.json', rpc_url=''
-    )
+    return A2ARESTFastAPIApplication(
+        agent_card,
+        request_handler,
+        extended_agent_card=agent_card,
+        extended_card_modifier=extended_card_modifier,
+    ).build(agent_card_url='/well-known/agent.json', rpc_url='')
 
 
 @pytest.fixture