Skip to content

Commit f432ec7

Browse files
authored
Merge branch 'main' into doc/helper-utils
2 parents f27b1ea + 9da9ecc commit f432ec7

File tree

5 files changed

+26
-22
lines changed

5 files changed

+26
-22
lines changed

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -182,9 +182,9 @@ async def handle_authenticated_agent_card(
182182

183183
if self.extended_card_modifier:
184184
context = self._context_builder.build(request)
185-
# If no base extended card is provided, pass the public card to the modifier
186-
base_card = card_to_serve if card_to_serve else self.agent_card
187-
card_to_serve = self.extended_card_modifier(base_card, context)
185+
card_to_serve = self.extended_card_modifier(card_to_serve, context)
186+
elif self.card_modifier:
187+
card_to_serve = self.card_modifier(card_to_serve)
188188

189189
return card_to_serve.model_dump(mode='json', exclude_none=True)
190190

src/a2a/server/request_handlers/jsonrpc_handler.py

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ def __init__(
6666
[AgentCard, ServerCallContext], AgentCard
6767
]
6868
| None = None,
69+
card_modifier: Callable[[AgentCard], AgentCard] | None = None,
6970
):
7071
"""Initializes the JSONRPCHandler.
7172
@@ -76,11 +77,14 @@ def __init__(
7677
extended_card_modifier: An optional callback to dynamically modify
7778
the extended agent card before it is served. It receives the
7879
call context.
80+
card_modifier: An optional callback to dynamically modify the public
81+
agent card before it is served.
7982
"""
8083
self.agent_card = agent_card
8184
self.request_handler = request_handler
8285
self.extended_agent_card = extended_agent_card
8386
self.extended_card_modifier = extended_card_modifier
87+
self.card_modifier = card_modifier
8488

8589
async def on_message_send(
8690
self,
@@ -425,14 +429,10 @@ async def get_authenticated_extended_card(
425429
Returns:
426430
A `GetAuthenticatedExtendedCardResponse` object containing the config or a JSON-RPC error.
427431
"""
428-
if (
429-
self.extended_agent_card is None
430-
and self.extended_card_modifier is None
431-
):
432-
return GetAuthenticatedExtendedCardResponse(
433-
root=JSONRPCErrorResponse(
434-
id=request.id,
435-
error=AuthenticatedExtendedCardNotConfiguredError(),
432+
if not self.agent_card.supports_authenticated_extended_card:
433+
raise ServerError(
434+
error=AuthenticatedExtendedCardNotConfiguredError(
435+
message='Authenticated card not supported'
436436
)
437437
)
438438

@@ -443,6 +443,8 @@ async def get_authenticated_extended_card(
443443
card_to_serve = base_card
444444
if self.extended_card_modifier and context:
445445
card_to_serve = self.extended_card_modifier(base_card, context)
446+
elif self.card_modifier:
447+
card_to_serve = self.card_modifier(base_card)
446448

447449
return GetAuthenticatedExtendedCardResponse(
448450
root=GetAuthenticatedExtendedCardSuccessResponse(

src/a2a/utils/proto_utils.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,8 @@ def task_state(cls, state: types.TaskState) -> a2a_pb2.TaskState:
122122
return a2a_pb2.TaskState.TASK_STATE_FAILED
123123
case types.TaskState.input_required:
124124
return a2a_pb2.TaskState.TASK_STATE_INPUT_REQUIRED
125+
case types.TaskState.auth_required:
126+
return a2a_pb2.TaskState.TASK_STATE_AUTH_REQUIRED
125127
case _:
126128
return a2a_pb2.TaskState.TASK_STATE_UNSPECIFIED
127129

@@ -568,6 +570,8 @@ def task_state(cls, state: a2a_pb2.TaskState) -> types.TaskState:
568570
return types.TaskState.failed
569571
case a2a_pb2.TaskState.TASK_STATE_INPUT_REQUIRED:
570572
return types.TaskState.input_required
573+
case a2a_pb2.TaskState.TASK_STATE_AUTH_REQUIRED:
574+
return types.TaskState.auth_required
571575
case _:
572576
return types.TaskState.unknown
573577

tests/server/request_handlers/test_jsonrpc_handler.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import unittest
22
import unittest.async_case
3-
43
from collections.abc import AsyncGenerator
54
from typing import Any, NoReturn
65
from unittest.mock import AsyncMock, MagicMock, call, patch
@@ -75,7 +74,6 @@
7574
)
7675
from a2a.utils.errors import ServerError
7776

78-
7977
MINIMAL_TASK: dict[str, Any] = {
8078
'id': 'task_123',
8179
'contextId': 'session-xyz',
@@ -93,7 +91,9 @@ class TestJSONRPCtHandler(unittest.async_case.IsolatedAsyncioTestCase):
9391
@pytest.fixture(autouse=True)
9492
def init_fixtures(self) -> None:
9593
self.mock_agent_card = MagicMock(
96-
spec=AgentCard, url='http://agent.example.com/api'
94+
spec=AgentCard,
95+
url='http://agent.example.com/api',
96+
supports_authenticated_extended_card=True,
9797
)
9898

9999
async def test_on_get_task_success(self) -> None:
@@ -1233,6 +1233,7 @@ async def test_get_authenticated_extended_card_not_configured(self) -> None:
12331233
"""Test error when authenticated extended agent card is not configured."""
12341234
# Arrange
12351235
mock_request_handler = AsyncMock(spec=DefaultRequestHandler)
1236+
self.mock_agent_card.supports_extended_card = True
12361237
handler = JSONRPCHandler(
12371238
self.mock_agent_card,
12381239
mock_request_handler,
@@ -1248,11 +1249,12 @@ async def test_get_authenticated_extended_card_not_configured(self) -> None:
12481249
)
12491250

12501251
# Assert
1251-
self.assertIsInstance(response.root, JSONRPCErrorResponse)
1252-
self.assertEqual(response.root.id, 'ext-card-req-2')
1252+
# Authenticated Extended Card flag is set with no extended card,
1253+
# returns base card in this case.
12531254
self.assertIsInstance(
1254-
response.root.error, AuthenticatedExtendedCardNotConfiguredError
1255+
response.root, GetAuthenticatedExtendedCardSuccessResponse
12551256
)
1257+
self.assertEqual(response.root.id, 'ext-card-req-2')
12561258

12571259
async def test_get_authenticated_extended_card_with_modifier(self) -> None:
12581260
"""Test successful retrieval of a dynamically modified extended agent card."""

tests/utils/test_proto_utils.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -170,11 +170,7 @@ def test_enum_conversions(self):
170170
)
171171

172172
for state in types.TaskState:
173-
if state not in (
174-
types.TaskState.unknown,
175-
types.TaskState.rejected,
176-
types.TaskState.auth_required,
177-
):
173+
if state not in (types.TaskState.unknown, types.TaskState.rejected):
178174
proto_state = proto_utils.ToProto.task_state(state)
179175
assert proto_utils.FromProto.task_state(proto_state) == state
180176

0 commit comments

Comments
 (0)