make extended_card_modifier async

Author: cbugwadia32

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

@@ -1,6 +1,6 @@
 import logging
 
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any
 
 
@@ -51,7 +51,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 +68,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)
 

src/a2a/utils/signing.py

@@ -1,14 +1,11 @@
 import json
-
 from collections.abc import Callable
 from typing import Any, TypedDict
 
 from a2a.utils.helpers import canonicalize_agent_card
 
-
 try:
     import jwt
-
     from jwt.api_jwk import PyJWK
     from jwt.exceptions import PyJWTError
     from jwt.utils import base64url_decode, base64url_encode
@@ -66,7 +63,7 @@ def create_agent_card_signer(
         A callable that takes an AgentCard and returns the modified AgentCard with a signature.
     """
 
-    def agent_card_signer(agent_card: AgentCard) -> AgentCard:
+    async def agent_card_signer(agent_card: AgentCard) -> AgentCard:
         """Signs agent card."""
         canonical_payload = canonicalize_agent_card(agent_card)
         payload_dict = json.loads(canonical_payload)

tests/server/apps/rest/test_rest_fastapi_app.py

@@ -1,5 +1,6 @@
 import logging
 
+from collections.abc import Awaitable, Callable
 from typing import Any
 from unittest.mock import MagicMock
 
@@ -13,6 +14,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,
@@ -33,7 +35,7 @@
 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 +60,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 +95,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