address feedback

Author: cbugwadia32

src/a2a/server/apps/jsonrpc/jsonrpc_app.py

@@ -51,7 +51,7 @@
     PREV_AGENT_CARD_WELL_KNOWN_PATH,
 )
 from a2a.utils.errors import MethodNotImplementedError
-from a2a.utils.helpers import apply_optional_awaitable
+from a2a.utils.helpers import maybe_await
 
 
 logger = logging.getLogger(__name__)
@@ -578,9 +578,7 @@ async def _handle_get_agent_card(self, request: Request) -> JSONResponse:
 
         card_to_serve = self.agent_card
         if self.card_modifier:
-            card_to_serve = await apply_optional_awaitable(
-                self.card_modifier, card_to_serve
-            )
+            card_to_serve = await maybe_await(self.card_modifier(card_to_serve))
 
         return JSONResponse(
             card_to_serve.model_dump(
@@ -609,8 +607,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 = await apply_optional_awaitable(
-                self.extended_card_modifier, base_card, context
+            card_to_serve = await maybe_await(
+                self.extended_card_modifier(base_card, context)
             )
 
         if card_to_serve:

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

@@ -4,7 +4,7 @@
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
 from typing import TYPE_CHECKING, Any
 
-from a2a.utils.helpers import apply_optional_awaitable
+from a2a.utils.helpers import maybe_await
 
 
 if TYPE_CHECKING:
@@ -153,9 +153,7 @@ async def handle_get_agent_card(
         """
         card_to_serve = self.agent_card
         if self.card_modifier:
-            card_to_serve = await apply_optional_awaitable(
-                self.card_modifier, card_to_serve
-            )
+            card_to_serve = await maybe_await(self.card_modifier(card_to_serve))
 
         return card_to_serve.model_dump(mode='json', exclude_none=True)
 
@@ -187,13 +185,11 @@ async def handle_authenticated_agent_card(
 
         if self.extended_card_modifier:
             context = self._context_builder.build(request)
-            card_to_serve = await apply_optional_awaitable(
-                self.extended_card_modifier, card_to_serve, context
+            card_to_serve = await maybe_await(
+                self.extended_card_modifier(card_to_serve, context)
             )
         elif self.card_modifier:
-            card_to_serve = await apply_optional_awaitable(
-                self.card_modifier, card_to_serve
-            )
+            card_to_serve = await maybe_await(self.card_modifier(card_to_serve))
 
         return card_to_serve.model_dump(mode='json', exclude_none=True)
 

src/a2a/server/request_handlers/grpc_handler.py

@@ -34,11 +34,7 @@
 from a2a.types import AgentCard, TaskNotFoundError
 from a2a.utils import proto_utils
 from a2a.utils.errors import ServerError
-from a2a.utils.helpers import (
-    apply_optional_awaitable,
-    validate,
-    validate_async_generator,
-)
+from a2a.utils.helpers import maybe_await, validate, validate_async_generator
 
 
 logger = logging.getLogger(__name__)
@@ -344,9 +340,7 @@ async def GetAgentCard(
         """Get the agent card for the agent served."""
         card_to_serve = self.agent_card
         if self.card_modifier:
-            card_to_serve = await apply_optional_awaitable(
-                self.card_modifier, card_to_serve
-            )
+            card_to_serve = await maybe_await(self.card_modifier(card_to_serve))
         return proto_utils.ToProto.agent_card(card_to_serve)
 
     async def abort_context(

src/a2a/server/request_handlers/jsonrpc_handler.py

@@ -46,7 +46,7 @@
     TaskStatusUpdateEvent,
 )
 from a2a.utils.errors import ServerError
-from a2a.utils.helpers import apply_optional_awaitable, validate
+from a2a.utils.helpers import maybe_await, validate
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -451,13 +451,11 @@ async def get_authenticated_extended_card(
 
         card_to_serve = base_card
         if self.extended_card_modifier and context:
-            card_to_serve = await apply_optional_awaitable(
-                self.extended_card_modifier, base_card, context
+            card_to_serve = await maybe_await(
+                self.extended_card_modifier(base_card, context)
             )
         elif self.card_modifier:
-            card_to_serve = await apply_optional_awaitable(
-                self.card_modifier, base_card
-            )
+            card_to_serve = await maybe_await(self.card_modifier(base_card))
 
         return GetAuthenticatedExtendedCardResponse(
             root=GetAuthenticatedExtendedCardSuccessResponse(

src/a2a/utils/helpers.py

@@ -6,8 +6,7 @@
 import logging
 
 from collections.abc import Awaitable, Callable
-from inspect import isawaitable
-from typing import Any, ParamSpec, TypeVar
+from typing import Any, TypeVar
 from uuid import uuid4
 
 from a2a.types import (
@@ -26,7 +25,6 @@
 
 
 T = TypeVar('T')
-P = ParamSpec('P')
 
 
 logger = logging.getLogger(__name__)
@@ -375,11 +373,8 @@ def canonicalize_agent_card(agent_card: AgentCard) -> str:
     return json.dumps(cleaned_dict, separators=(',', ':'), sort_keys=True)
 
 
-async def apply_optional_awaitable(
-    func: Callable[P, Awaitable[T] | T], *args: P.args, **kwargs: P.kwargs
-) -> T:
-    """Applies a function that may be sync or async and returns the result."""
-    result = func(*args, **kwargs)
-    if isawaitable(result):
-        return await result
-    return result
+async def maybe_await(value: T | Awaitable[T]) -> T:
+    """Awaits a value if it's awaitable, otherwise simply provides it back."""
+    if inspect.isawaitable(value):
+        return await value
+    return value