Skip to content

Commit b261ceb

Browse files
bartek-wbartek-w
authored
feat: A2A Version Header validation on server side. (#865)
1 parent 24f5f1e commit b261ceb

16 files changed

Lines changed: 861 additions & 146 deletions

File tree

‎src/a2a/compat/v0_3/jsonrpc_adapter.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,9 @@
3838
from a2a.server.jsonrpc_models import (
3939
JSONRPCError as CoreJSONRPCError,
4040
)
41+
from a2a.utils import constants
4142
from a2a.utils.errors import ExtendedAgentCardNotConfiguredError
42-
from a2a.utils.helpers import maybe_await
43+
from a2a.utils.helpers import maybe_await, validate_version
4344

4445

4546
logger = logging.getLogger(__name__)
@@ -152,6 +153,7 @@ async def handle_request(
152153
request_id, CoreInternalError(message=str(e))
153154
)
154155

156+
@validate_version(constants.PROTOCOL_VERSION_0_3)
155157
async def _process_non_streaming_request(
156158
self,
157159
request_id: 'str | int | None',
@@ -266,6 +268,7 @@ async def get_authenticated_extended_card(
266268

267269
return conversions.to_compat_agent_card(card_to_serve)
268270

271+
@validate_version(constants.PROTOCOL_VERSION_0_3)
269272
async def _process_streaming_request(
270273
self,
271274
request_id: 'str | int | None',

‎src/a2a/compat/v0_3/rest_handler.py‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,12 @@
2828
from a2a.compat.v0_3 import types as types_v03
2929
from a2a.compat.v0_3.request_handler import RequestHandler03
3030
from a2a.server.context import ServerCallContext
31-
from a2a.utils.helpers import validate, validate_async_generator
31+
from a2a.utils import constants
32+
from a2a.utils.helpers import (
33+
validate,
34+
validate_async_generator,
35+
validate_version,
36+
)
3237
from a2a.utils.telemetry import SpanKind, trace_class
3338

3439

@@ -53,6 +58,7 @@ def __init__(
5358
self.agent_card = agent_card
5459
self.handler03 = RequestHandler03(request_handler=request_handler)
5560

61+
@validate_version(constants.PROTOCOL_VERSION_0_3)
5662
async def on_message_send(
5763
self,
5864
request: Request,
@@ -78,6 +84,7 @@ async def on_message_send(
7884
pb2_v03_resp = proto_utils.ToProto.task_or_message(v03_resp)
7985
return MessageToDict(pb2_v03_resp)
8086

87+
@validate_version(constants.PROTOCOL_VERSION_0_3)
8188
@validate_async_generator(
8289
lambda self: self.agent_card.capabilities.streaming,
8390
'Streaming is not supported by the agent',
@@ -110,6 +117,7 @@ async def on_message_send_stream(
110117
)
111118
yield MessageToDict(v03_pb_resp)
112119

120+
@validate_version(constants.PROTOCOL_VERSION_0_3)
113121
async def on_cancel_task(
114122
self,
115123
request: Request,
@@ -134,6 +142,7 @@ async def on_cancel_task(
134142
pb2_v03_task = proto_utils.ToProto.task(v03_resp)
135143
return MessageToDict(pb2_v03_task)
136144

145+
@validate_version(constants.PROTOCOL_VERSION_0_3)
137146
@validate_async_generator(
138147
lambda self: self.agent_card.capabilities.streaming,
139148
'Streaming is not supported by the agent',
@@ -166,6 +175,7 @@ async def on_subscribe_to_task(
166175
)
167176
yield MessageToDict(v03_pb_resp)
168177

178+
@validate_version(constants.PROTOCOL_VERSION_0_3)
169179
async def get_push_notification(
170180
self,
171181
request: Request,
@@ -198,6 +208,7 @@ async def get_push_notification(
198208
)
199209
return MessageToDict(pb2_v03_config)
200210

211+
@validate_version(constants.PROTOCOL_VERSION_0_3)
201212
@validate(
202213
lambda self: self.agent_card.capabilities.push_notifications,
203214
'Push notifications are not supported by the agent',
@@ -242,6 +253,7 @@ async def set_push_notification(
242253
)
243254
return MessageToDict(pb2_v03_config)
244255

256+
@validate_version(constants.PROTOCOL_VERSION_0_3)
245257
async def on_get_task(
246258
self,
247259
request: Request,
@@ -271,6 +283,7 @@ async def on_get_task(
271283
pb2_v03_task = proto_utils.ToProto.task(v03_resp)
272284
return MessageToDict(pb2_v03_task)
273285

286+
@validate_version(constants.PROTOCOL_VERSION_0_3)
274287
async def list_push_notifications(
275288
self,
276289
request: Request,
@@ -297,6 +310,7 @@ async def list_push_notifications(
297310

298311
return MessageToDict(pb2_v03_resp)
299312

313+
@validate_version(constants.PROTOCOL_VERSION_0_3)
300314
async def list_tasks(
301315
self,
302316
request: Request,

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -254,7 +254,7 @@ def __init__( # noqa: PLR0913
254254
agent_card=agent_card,
255255
http_handler=http_handler,
256256
extended_agent_card=extended_agent_card,
257-
context_builder=context_builder,
257+
context_builder=self._context_builder,
258258
card_modifier=card_modifier,
259259
extended_card_modifier=extended_card_modifier,
260260
)
@@ -444,6 +444,8 @@ async def _handle_requests(self, request: Request) -> Response: # noqa: PLR0911
444444
InvalidRequestError(message='Payload too large'),
445445
)
446446
raise e
447+
except A2AError as e:
448+
return self._generate_error_response(request_id, e)
447449
except Exception as e:
448450
logger.exception('Unhandled exception')
449451
return self._generate_error_response(

‎src/a2a/server/request_handlers/jsonrpc_handler.py‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
Task,
3232
TaskPushNotificationConfig,
3333
)
34-
from a2a.utils import proto_utils
34+
from a2a.utils import constants, proto_utils
3535
from a2a.utils.errors import (
3636
JSON_RPC_ERROR_CODE_MAP,
3737
A2AError,
@@ -49,7 +49,12 @@
4949
UnsupportedOperationError,
5050
VersionNotSupportedError,
5151
)
52-
from a2a.utils.helpers import maybe_await, validate, validate_async_generator
52+
from a2a.utils.helpers import (
53+
maybe_await,
54+
validate,
55+
validate_async_generator,
56+
validate_version,
57+
)
5358
from a2a.utils.telemetry import SpanKind, trace_class
5459

5560

@@ -142,6 +147,7 @@ def _get_request_id(
142147
return None
143148
return context.state.get('request_id')
144149

150+
@validate_version(constants.PROTOCOL_VERSION_1_0)
145151
async def on_message_send(
146152
self,
147153
request: SendMessageRequest,
@@ -171,6 +177,7 @@ async def on_message_send(
171177
except A2AError as e:
172178
return _build_error_response(request_id, e)
173179

180+
@validate_version(constants.PROTOCOL_VERSION_1_0)
174181
@validate_async_generator(
175182
lambda self: self.agent_card.capabilities.streaming,
176183
'Streaming is not supported by the agent',
@@ -209,6 +216,7 @@ async def on_message_send_stream(
209216
e,
210217
)
211218

219+
@validate_version(constants.PROTOCOL_VERSION_1_0)
212220
async def on_cancel_task(
213221
self,
214222
request: CancelTaskRequest,
@@ -235,6 +243,7 @@ async def on_cancel_task(
235243

236244
return _build_error_response(request_id, TaskNotFoundError())
237245

246+
@validate_version(constants.PROTOCOL_VERSION_1_0)
238247
@validate_async_generator(
239248
lambda self: self.agent_card.capabilities.streaming,
240249
'Streaming is not supported by the agent',
@@ -273,6 +282,7 @@ async def on_subscribe_to_task(
273282
e,
274283
)
275284

285+
@validate_version(constants.PROTOCOL_VERSION_1_0)
276286
async def get_push_notification_config(
277287
self,
278288
request: GetTaskPushNotificationConfigRequest,
@@ -299,6 +309,7 @@ async def get_push_notification_config(
299309
except A2AError as e:
300310
return _build_error_response(request_id, e)
301311

312+
@validate_version(constants.PROTOCOL_VERSION_1_0)
302313
@validate(
303314
lambda self: self.agent_card.capabilities.push_notifications,
304315
'Push notifications are not supported by the agent',
@@ -336,6 +347,7 @@ async def set_push_notification_config(
336347
except A2AError as e:
337348
return _build_error_response(request_id, e)
338349

350+
@validate_version(constants.PROTOCOL_VERSION_1_0)
339351
async def on_get_task(
340352
self,
341353
request: GetTaskRequest,
@@ -362,6 +374,7 @@ async def on_get_task(
362374

363375
return _build_error_response(request_id, TaskNotFoundError())
364376

377+
@validate_version(constants.PROTOCOL_VERSION_1_0)
365378
async def list_tasks(
366379
self,
367380
request: ListTasksRequest,
@@ -390,6 +403,7 @@ async def list_tasks(
390403
except A2AError as e:
391404
return _build_error_response(request_id, e)
392405

406+
@validate_version(constants.PROTOCOL_VERSION_1_0)
393407
async def list_push_notification_configs(
394408
self,
395409
request: ListTaskPushNotificationConfigsRequest,
@@ -415,6 +429,7 @@ async def list_push_notification_configs(
415429
except A2AError as e:
416430
return _build_error_response(request_id, e)
417431

432+
@validate_version(constants.PROTOCOL_VERSION_1_0)
418433
async def delete_push_notification_config(
419434
self,
420435
request: DeleteTaskPushNotificationConfigRequest,

‎src/a2a/server/request_handlers/rest_handler.py‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,13 @@
2727
GetTaskPushNotificationConfigRequest,
2828
SubscribeToTaskRequest,
2929
)
30-
from a2a.utils import proto_utils
30+
from a2a.utils import constants, proto_utils
3131
from a2a.utils.errors import TaskNotFoundError
32-
from a2a.utils.helpers import validate, validate_async_generator
32+
from a2a.utils.helpers import (
33+
validate,
34+
validate_async_generator,
35+
validate_version,
36+
)
3337
from a2a.utils.telemetry import SpanKind, trace_class
3438

3539

@@ -61,6 +65,7 @@ def __init__(
6165
self.agent_card = agent_card
6266
self.request_handler = request_handler
6367

68+
@validate_version(constants.PROTOCOL_VERSION_1_0)
6469
async def on_message_send(
6570
self,
6671
request: Request,
@@ -87,6 +92,7 @@ async def on_message_send(
8792
response = a2a_pb2.SendMessageResponse(message=task_or_message)
8893
return MessageToDict(response)
8994

95+
@validate_version(constants.PROTOCOL_VERSION_1_0)
9096
@validate_async_generator(
9197
lambda self: self.agent_card.capabilities.streaming,
9298
'Streaming is not supported by the agent',
@@ -117,6 +123,7 @@ async def on_message_send_stream(
117123
response = proto_utils.to_stream_response(event)
118124
yield MessageToDict(response)
119125

126+
@validate_version(constants.PROTOCOL_VERSION_1_0)
120127
async def on_cancel_task(
121128
self,
122129
request: Request,
@@ -139,6 +146,7 @@ async def on_cancel_task(
139146
return MessageToDict(task)
140147
raise TaskNotFoundError
141148

149+
@validate_version(constants.PROTOCOL_VERSION_1_0)
142150
@validate_async_generator(
143151
lambda self: self.agent_card.capabilities.streaming,
144152
'Streaming is not supported by the agent',
@@ -165,6 +173,7 @@ async def on_subscribe_to_task(
165173
):
166174
yield MessageToDict(proto_utils.to_stream_response(event))
167175

176+
@validate_version(constants.PROTOCOL_VERSION_1_0)
168177
async def get_push_notification(
169178
self,
170179
request: Request,
@@ -192,6 +201,7 @@ async def get_push_notification(
192201
)
193202
return MessageToDict(config)
194203

204+
@validate_version(constants.PROTOCOL_VERSION_1_0)
195205
@validate(
196206
lambda self: self.agent_card.capabilities.push_notifications,
197207
'Push notifications are not supported by the agent',
@@ -229,6 +239,7 @@ async def set_push_notification(
229239
)
230240
return MessageToDict(config)
231241

242+
@validate_version(constants.PROTOCOL_VERSION_1_0)
232243
async def on_get_task(
233244
self,
234245
request: Request,
@@ -251,6 +262,7 @@ async def on_get_task(
251262
return MessageToDict(task)
252263
raise TaskNotFoundError
253264

265+
@validate_version(constants.PROTOCOL_VERSION_1_0)
254266
async def delete_push_notification(
255267
self,
256268
request: Request,
@@ -275,6 +287,7 @@ async def delete_push_notification(
275287
)
276288
return {}
277289

290+
@validate_version(constants.PROTOCOL_VERSION_1_0)
278291
async def list_tasks(
279292
self,
280293
request: Request,
@@ -295,6 +308,7 @@ async def list_tasks(
295308
result = await self.request_handler.on_list_tasks(params, context)
296309
return MessageToDict(result, always_print_fields_with_no_presence=True)
297310

311+
@validate_version(constants.PROTOCOL_VERSION_1_0)
298312
async def list_push_notifications(
299313
self,
300314
request: Request,

0 commit comments

Comments
 (0)