make card modifier and extended card modifier async

Author: cbugwadia32

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] | AgentCard]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
         ]
         | None = None,
         max_content_length: int | None = 10 * 1024 * 1024,  # 10MB

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 isawaitable
 from typing import TYPE_CHECKING, Any
 
 from pydantic import ValidationError
@@ -178,9 +179,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] | AgentCard]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
         ]
         | None = None,
         max_content_length: int | None = 10 * 1024 * 1024,  # 10MB
@@ -576,7 +578,8 @@ 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)
+            result = self.card_modifier(card_to_serve)
+            card_to_serve = await result if isawaitable(result) else result
 
         return JSONResponse(
             card_to_serve.model_dump(
@@ -605,7 +608,8 @@ 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)
+            result = await self.extended_card_modifier(base_card, context)
+            card_to_serve = await result if isawaitable(result) else result
 
         if card_to_serve:
             return JSONResponse(

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] | AgentCard]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
         ]
         | None = None,
         max_content_length: int | None = 10 * 1024 * 1024,  # 10MB

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] | AgentCard]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
         ]
         | None = None,
     ):

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 isawaitable
 from typing import TYPE_CHECKING, Any
 
 
@@ -58,9 +59,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] | AgentCard]
+        | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
         ]
         | None = None,
     ):
@@ -150,7 +152,8 @@ 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)
+            result = self.card_modifier(card_to_serve)
+            card_to_serve = await result if isawaitable(result) else result
 
         return card_to_serve.model_dump(mode='json', exclude_none=True)
 
@@ -182,9 +185,11 @@ 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)
+            result = self.extended_card_modifier(card_to_serve, context)
+            card_to_serve = await result if isawaitable(result) else result
         elif self.card_modifier:
-            card_to_serve = self.card_modifier(card_to_serve)
+            result = self.card_modifier(card_to_serve)
+            card_to_serve = await result if isawaitable(result) else result
 
         return card_to_serve.model_dump(mode='json', exclude_none=True)
 

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 isawaitable
 
 
 try:
@@ -89,7 +90,8 @@ 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] | AgentCard]
+        | None = None,
     ):
         """Initializes the GrpcHandler.
 
@@ -339,7 +341,8 @@ 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)
+            result = self.card_modifier(card_to_serve)
+            card_to_serve = await result if isawaitable(result) else result
         return proto_utils.ToProto.agent_card(card_to_serve)
 
     async def abort_context(

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 isawaitable
 
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
@@ -63,10 +64,11 @@ def __init__(
         request_handler: RequestHandler,
         extended_agent_card: AgentCard | None = None,
         extended_card_modifier: Callable[
-            [AgentCard, ServerCallContext], AgentCard
+            [AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard
         ]
         | None = None,
-        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
+        card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
+        | None = None,
     ):
         """Initializes the JSONRPCHandler.
 
@@ -450,9 +452,11 @@ 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)
+            result = self.extended_card_modifier(base_card, context)
+            card_to_serve = await result if isawaitable(result) else result
         elif self.card_modifier:
-            card_to_serve = self.card_modifier(base_card)
+            result = self.card_modifier(base_card)
+            card_to_serve = await result if isawaitable(result) else result
 
         return GetAuthenticatedExtendedCardResponse(
             root=GetAuthenticatedExtendedCardSuccessResponse(

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'

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'