make card modifier and extended card modifier async

cbugwadia32 · cbugwadia32 · commit d543213262f7 · 2026-01-31T22:14:10.000-08:00
diff --git a/src/a2a/server/apps/jsonrpc/fastapi_app.py b/src/a2a/server/apps/jsonrpc/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
 
 
@@ -72,9 +72,10 @@ def __init__(  # noqa: PLR0913
         http_handler: RequestHandler,
         extended_agent_card: AgentCard | None = None,
         context_builder: CallContextBuilder | None = None,
-        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
+        card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard]
         ]
         | None = None,
         max_content_length: int | None = 10 * 1024 * 1024,  # 10MB
diff --git a/src/a2a/server/apps/jsonrpc/jsonrpc_app.py b/src/a2a/server/apps/jsonrpc/jsonrpc_app.py
@@ -4,7 +4,8 @@
 import traceback
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncGenerator, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable
+from inspect import iscoroutinefunction
 from typing import TYPE_CHECKING, Any
 
 from pydantic import ValidationError
@@ -178,10 +179,13 @@ def __init__(  # noqa: PLR0913
         http_handler: RequestHandler,
         extended_agent_card: AgentCard | None = None,
         context_builder: CallContextBuilder | None = None,
-        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
+        card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
+        | Callable[[AgentCard], AgentCard]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard]
         ]
+        | Callable[[AgentCard, ServerCallContext], AgentCard]
         | None = None,
         max_content_length: int | None = 10 * 1024 * 1024,  # 10MB
     ) -> None:
@@ -576,7 +580,10 @@ async def _handle_get_agent_card(self, request: Request) -> JSONResponse:
 
         card_to_serve = self.agent_card
         if self.card_modifier:
-            card_to_serve = self.card_modifier(card_to_serve)
+            if iscoroutinefunction(self.card_modifier):
+                card_to_serve = await self.card_modifier(card_to_serve)
+            else:
+                card_to_serve = self.card_modifier(card_to_serve)
 
         return JSONResponse(
             card_to_serve.model_dump(
@@ -605,7 +612,12 @@ async def _handle_get_authenticated_extended_agent_card(
             context = self._context_builder.build(request)
             # If no base extended card is provided, pass the public card to the modifier
             base_card = card_to_serve if card_to_serve else self.agent_card
-            card_to_serve = self.extended_card_modifier(base_card, context)
+            if iscoroutinefunction(self.extended_card_modifier):
+                card_to_serve = await self.extended_card_modifier(
+                    base_card, context
+                )
+            else:
+                card_to_serve = self.extended_card_modifier(base_card, context)
 
         if card_to_serve:
             return JSONResponse(
diff --git a/src/a2a/server/apps/jsonrpc/starlette_app.py b/src/a2a/server/apps/jsonrpc/starlette_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
 
 
@@ -54,9 +54,10 @@ def __init__(  # noqa: PLR0913
         http_handler: RequestHandler,
         extended_agent_card: AgentCard | None = None,
         context_builder: CallContextBuilder | None = None,
-        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
+        card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard]
         ]
         | None = None,
         max_content_length: int | None = 10 * 1024 * 1024,  # 10MB
diff --git a/src/a2a/server/apps/rest/fastapi_app.py b/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
 
 
@@ -49,9 +49,10 @@ def __init__(  # noqa: PLR0913
         http_handler: RequestHandler,
         extended_agent_card: AgentCard | None = None,
         context_builder: CallContextBuilder | None = None,
-        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
+        card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard]
         ]
         | None = None,
     ):
diff --git a/src/a2a/server/apps/rest/rest_adapter.py b/src/a2a/server/apps/rest/rest_adapter.py
@@ -2,6 +2,7 @@
 import logging
 
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
+from inspect import iscoroutinefunction
 from typing import TYPE_CHECKING, Any
 
 
@@ -58,10 +59,13 @@ def __init__(  # noqa: PLR0913
         http_handler: RequestHandler,
         extended_agent_card: AgentCard | None = None,
         context_builder: CallContextBuilder | None = None,
-        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
+        card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
+        | Callable[[AgentCard], AgentCard]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard]
         ]
+        | Callable[[AgentCard, ServerCallContext], AgentCard]
         | None = None,
     ):
         """Initializes the RESTApplication.
@@ -150,7 +154,10 @@ async def handle_get_agent_card(
         """
         card_to_serve = self.agent_card
         if self.card_modifier:
-            card_to_serve = self.card_modifier(card_to_serve)
+            if iscoroutinefunction(self.card_modifier):
+                card_to_serve = await self.card_modifier(card_to_serve)
+            else:
+                card_to_serve = self.card_modifier(card_to_serve)
 
         return card_to_serve.model_dump(mode='json', exclude_none=True)
 
@@ -182,9 +189,19 @@ 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)
+            if iscoroutinefunction(self.extended_card_modifier):
+                card_to_serve = await self.extended_card_modifier(
+                    card_to_serve, context
+                )
+            else:
+                card_to_serve = self.extended_card_modifier(
+                    card_to_serve, context
+                )
         elif self.card_modifier:
-            card_to_serve = self.card_modifier(card_to_serve)
+            if iscoroutinefunction(self.extended_card_modifier):
+                card_to_serve = await self.card_modifier(card_to_serve)
+            else:
+                card_to_serve = self.card_modifier(card_to_serve)
 
         return card_to_serve.model_dump(mode='json', exclude_none=True)
 
diff --git a/src/a2a/server/request_handlers/grpc_handler.py b/src/a2a/server/request_handlers/grpc_handler.py
@@ -3,7 +3,8 @@
 import logging
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterable, Sequence
+from collections.abc import AsyncIterable, Awaitable, Sequence
+from inspect import iscoroutinefunction
 
 
 try:
@@ -89,7 +90,9 @@ def __init__(
         agent_card: AgentCard,
         request_handler: RequestHandler,
         context_builder: CallContextBuilder | None = None,
-        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
+        card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
+        | Callable[[AgentCard], AgentCard]
+        | None = None,
     ):
         """Initializes the GrpcHandler.
 
@@ -339,7 +342,10 @@ async def GetAgentCard(
         """Get the agent card for the agent served."""
         card_to_serve = self.agent_card
         if self.card_modifier:
-            card_to_serve = self.card_modifier(card_to_serve)
+            if iscoroutinefunction(self.card_modifier):
+                card_to_serve = await self.card_modifier(card_to_serve)
+            else:
+                card_to_serve = self.card_modifier(card_to_serve)
         return proto_utils.ToProto.agent_card(card_to_serve)
 
     async def abort_context(
diff --git a/src/a2a/server/request_handlers/jsonrpc_handler.py b/src/a2a/server/request_handlers/jsonrpc_handler.py
@@ -1,6 +1,7 @@
 import logging
 
-from collections.abc import AsyncIterable, Callable
+from collections.abc import AsyncIterable, Awaitable, Callable
+from inspect import iscoroutinefunction
 
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
@@ -63,10 +64,13 @@ def __init__(
         request_handler: RequestHandler,
         extended_agent_card: AgentCard | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard]
         ]
+        | Callable[[AgentCard, ServerCallContext], AgentCard]
+        | None = None,
+        card_modifier: Callable[[AgentCard], Awaitable[AgentCard]]
+        | Callable[[AgentCard], AgentCard]
         | None = None,
-        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
     ):
         """Initializes the JSONRPCHandler.
 
@@ -450,9 +454,17 @@ async def get_authenticated_extended_card(
 
         card_to_serve = base_card
         if self.extended_card_modifier and context:
-            card_to_serve = self.extended_card_modifier(base_card, context)
+            if iscoroutinefunction(self.extended_card_modifier):
+                card_to_serve = await self.extended_card_modifier(
+                    base_card, context
+                )
+            else:
+                card_to_serve = self.extended_card_modifier(base_card, context)
         elif self.card_modifier:
-            card_to_serve = self.card_modifier(base_card)
+            if iscoroutinefunction(self.card_modifier):
+                card_to_serve = await self.card_modifier(base_card)
+            else:
+                card_to_serve = self.card_modifier(base_card)
 
         return GetAuthenticatedExtendedCardResponse(
             root=GetAuthenticatedExtendedCardSuccessResponse(
diff --git a/src/a2a/utils/signing.py b/src/a2a/utils/signing.py
@@ -1,6 +1,6 @@
 import json
 
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from typing import Any, TypedDict
 
 from a2a.utils.helpers import canonicalize_agent_card
@@ -54,7 +54,7 @@ def create_agent_card_signer(
     signing_key: PyJWK | str | bytes,
     protected_header: ProtectedHeader,
     header: dict[str, Any] | None = None,
-) -> Callable[[AgentCard], AgentCard]:
+) -> Callable[[AgentCard], Awaitable[AgentCard]]:
     """Creates a function that signs an AgentCard and adds the signature.
 
     Args:
@@ -66,7 +66,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)
diff --git a/tests/server/request_handlers/test_grpc_handler.py b/tests/server/request_handlers/test_grpc_handler.py
@@ -209,6 +209,34 @@ async def test_get_agent_card_with_modifier(
 ) -> None:
     """Test GetAgentCard call with a card_modifier."""
 
+    async def modifier(card: types.AgentCard) -> types.AgentCard:
+        modified_card = card.model_copy(deep=True)
+        modified_card.name = 'Modified gRPC Agent'
+        return modified_card
+
+    grpc_handler_modified = GrpcHandler(
+        agent_card=sample_agent_card,
+        request_handler=mock_request_handler,
+        card_modifier=modifier,
+    )
+
+    request_proto = a2a_pb2.GetAgentCardRequest()
+    response = await grpc_handler_modified.GetAgentCard(
+        request_proto, mock_grpc_context
+    )
+
+    assert response.name == 'Modified gRPC Agent'
+    assert response.version == sample_agent_card.version
+
+
+@pytest.mark.asyncio
+async def test_get_agent_card_with_modifier_sync(
+    mock_request_handler: AsyncMock,
+    sample_agent_card: types.AgentCard,
+    mock_grpc_context: AsyncMock,
+) -> None:
+    """Test GetAgentCard call with a synchronous card_modifier."""
+
     def modifier(card: types.AgentCard) -> types.AgentCard:
         modified_card = card.model_copy(deep=True)
         modified_card.name = 'Modified gRPC Agent'
diff --git a/tests/server/request_handlers/test_jsonrpc_handler.py b/tests/server/request_handlers/test_jsonrpc_handler.py
@@ -1295,6 +1295,57 @@ async def test_get_authenticated_extended_card_with_modifier(self) -> None:
             skills=[],
         )
 
+        async def modifier(
+            card: AgentCard, context: ServerCallContext
+        ) -> AgentCard:
+            modified_card = card.model_copy(deep=True)
+            modified_card.name = 'Modified Card'
+            modified_card.description = (
+                f'Modified for context: {context.state.get("foo")}'
+            )
+            return modified_card
+
+        handler = JSONRPCHandler(
+            self.mock_agent_card,
+            mock_request_handler,
+            extended_agent_card=mock_base_card,
+            extended_card_modifier=modifier,
+        )
+        request = GetAuthenticatedExtendedCardRequest(id='ext-card-req-mod')
+        call_context = ServerCallContext(state={'foo': 'bar'})
+
+        # Act
+        response: GetAuthenticatedExtendedCardResponse = (
+            await handler.get_authenticated_extended_card(request, call_context)
+        )
+
+        # Assert
+        self.assertIsInstance(
+            response.root, GetAuthenticatedExtendedCardSuccessResponse
+        )
+        self.assertEqual(response.root.id, 'ext-card-req-mod')
+        modified_card = response.root.result
+        self.assertEqual(modified_card.name, 'Modified Card')
+        self.assertEqual(modified_card.description, 'Modified for context: bar')
+        self.assertEqual(modified_card.version, '1.0')
+
+    async def test_get_authenticated_extended_card_with_modifier_sync(
+        self,
+    ) -> None:
+        """Test successful retrieval of a synchronously dynamically modified extended agent card."""
+        # Arrange
+        mock_request_handler = AsyncMock(spec=DefaultRequestHandler)
+        mock_base_card = AgentCard(
+            name='Base Card',
+            description='Base details',
+            url='http://agent.example.com/api',
+            version='1.0',
+            capabilities=AgentCapabilities(),
+            default_input_modes=['text/plain'],
+            default_output_modes=['application/json'],
+            skills=[],
+        )
+
         def modifier(card: AgentCard, context: ServerCallContext) -> AgentCard:
             modified_card = card.model_copy(deep=True)
             modified_card.name = 'Modified Card'
diff --git a/tests/server/test_integration.py b/tests/server/test_integration.py
diff --git a/tests/utils/test_signing.py b/tests/utils/test_signing.py
