Merge branch 'main' into task_store

Author: holtskinner

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

@@ -182,9 +182,9 @@ async def handle_authenticated_agent_card(
 
         if self.extended_card_modifier:
             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)
+            card_to_serve = self.extended_card_modifier(card_to_serve, context)
+        elif self.card_modifier:
+            card_to_serve = self.card_modifier(card_to_serve)
 
         return card_to_serve.model_dump(mode='json', exclude_none=True)
 

src/a2a/server/request_handlers/jsonrpc_handler.py

@@ -66,6 +66,7 @@ def __init__(
             [AgentCard, ServerCallContext], AgentCard
         ]
         | None = None,
+        card_modifier: Callable[[AgentCard], AgentCard] | None = None,
     ):
         """Initializes the JSONRPCHandler.
 
@@ -76,11 +77,14 @@ def __init__(
             extended_card_modifier: An optional callback to dynamically modify
               the extended agent card before it is served. It receives the
               call context.
+            card_modifier: An optional callback to dynamically modify the public
+              agent card before it is served.
         """
         self.agent_card = agent_card
         self.request_handler = request_handler
         self.extended_agent_card = extended_agent_card
         self.extended_card_modifier = extended_card_modifier
+        self.card_modifier = card_modifier
 
     async def on_message_send(
         self,
@@ -425,14 +429,10 @@ async def get_authenticated_extended_card(
         Returns:
             A `GetAuthenticatedExtendedCardResponse` object containing the config or a JSON-RPC error.
         """
-        if (
-            self.extended_agent_card is None
-            and self.extended_card_modifier is None
-        ):
-            return GetAuthenticatedExtendedCardResponse(
-                root=JSONRPCErrorResponse(
-                    id=request.id,
-                    error=AuthenticatedExtendedCardNotConfiguredError(),
+        if not self.agent_card.supports_authenticated_extended_card:
+            raise ServerError(
+                error=AuthenticatedExtendedCardNotConfiguredError(
+                    message='Authenticated card not supported'
                 )
             )
 
@@ -443,6 +443,8 @@ 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)
+        elif self.card_modifier:
+            card_to_serve = self.card_modifier(base_card)
 
         return GetAuthenticatedExtendedCardResponse(
             root=GetAuthenticatedExtendedCardSuccessResponse(

src/a2a/utils/proto_utils.py

@@ -122,6 +122,8 @@ def task_state(cls, state: types.TaskState) -> a2a_pb2.TaskState:
                 return a2a_pb2.TaskState.TASK_STATE_FAILED
             case types.TaskState.input_required:
                 return a2a_pb2.TaskState.TASK_STATE_INPUT_REQUIRED
+            case types.TaskState.auth_required:
+                return a2a_pb2.TaskState.TASK_STATE_AUTH_REQUIRED
             case _:
                 return a2a_pb2.TaskState.TASK_STATE_UNSPECIFIED
 
@@ -568,6 +570,8 @@ def task_state(cls, state: a2a_pb2.TaskState) -> types.TaskState:
                 return types.TaskState.failed
             case a2a_pb2.TaskState.TASK_STATE_INPUT_REQUIRED:
                 return types.TaskState.input_required
+            case a2a_pb2.TaskState.TASK_STATE_AUTH_REQUIRED:
+                return types.TaskState.auth_required
             case _:
                 return types.TaskState.unknown
 

tests/server/request_handlers/test_jsonrpc_handler.py

@@ -1,6 +1,5 @@
 import unittest
 import unittest.async_case
-
 from collections.abc import AsyncGenerator
 from typing import Any, NoReturn
 from unittest.mock import AsyncMock, MagicMock, call, patch
@@ -75,7 +74,6 @@
 )
 from a2a.utils.errors import ServerError
 
-
 MINIMAL_TASK: dict[str, Any] = {
     'id': 'task_123',
     'contextId': 'session-xyz',
@@ -93,7 +91,9 @@ class TestJSONRPCtHandler(unittest.async_case.IsolatedAsyncioTestCase):
     @pytest.fixture(autouse=True)
     def init_fixtures(self) -> None:
         self.mock_agent_card = MagicMock(
-            spec=AgentCard, url='http://agent.example.com/api'
+            spec=AgentCard,
+            url='http://agent.example.com/api',
+            supports_authenticated_extended_card=True,
         )
 
     async def test_on_get_task_success(self) -> None:
@@ -1235,6 +1235,7 @@ async def test_get_authenticated_extended_card_not_configured(self) -> None:
         """Test error when authenticated extended agent card is not configured."""
         # Arrange
         mock_request_handler = AsyncMock(spec=DefaultRequestHandler)
+        self.mock_agent_card.supports_extended_card = True
         handler = JSONRPCHandler(
             self.mock_agent_card,
             mock_request_handler,
@@ -1250,11 +1251,12 @@ async def test_get_authenticated_extended_card_not_configured(self) -> None:
         )
 
         # Assert
-        self.assertIsInstance(response.root, JSONRPCErrorResponse)
-        self.assertEqual(response.root.id, 'ext-card-req-2')
+        # Authenticated Extended Card flag is set with no extended card,
+        # returns base card in this case.
         self.assertIsInstance(
-            response.root.error, AuthenticatedExtendedCardNotConfiguredError
+            response.root, GetAuthenticatedExtendedCardSuccessResponse
         )
+        self.assertEqual(response.root.id, 'ext-card-req-2')
 
     async def test_get_authenticated_extended_card_with_modifier(self) -> None:
         """Test successful retrieval of a dynamically modified extended agent card."""

tests/utils/test_proto_utils.py

@@ -170,11 +170,7 @@ def test_enum_conversions(self):
         )
 
         for state in types.TaskState:
-            if state not in (
-                types.TaskState.unknown,
-                types.TaskState.rejected,
-                types.TaskState.auth_required,
-            ):
+            if state not in (types.TaskState.unknown, types.TaskState.rejected):
                 proto_state = proto_utils.ToProto.task_state(state)
                 assert proto_utils.FromProto.task_state(proto_state) == state