diff --git a/src/a2a/server/routes/rest_dispatcher.py b/src/a2a/server/routes/rest_dispatcher.py index 80964e5bf..4c15a7c68 100644 --- a/src/a2a/server/routes/rest_dispatcher.py +++ b/src/a2a/server/routes/rest_dispatcher.py @@ -1,3 +1,4 @@ +import functools import logging from collections.abc import AsyncIterator, Awaitable, Callable @@ -65,6 +66,15 @@ class A2AJSONResponse(JSONResponse): logger = logging.getLogger(__name__) +# Spec Section 11.1: HTTP+JSON responses, errors included, use the A2A media +# type. The v0.3 compat adapter keeps the shared decorators' application/json. +_a2a_rest_error_handler = functools.partial( + rest_error_handler, media_type=constants.A2A_JSON_MEDIA_TYPE +) +_a2a_rest_stream_error_handler = functools.partial( + rest_stream_error_handler, media_type=constants.A2A_JSON_MEDIA_TYPE +) + TResponse = TypeVar('TResponse') @@ -170,7 +180,7 @@ async def event_generator() -> AsyncIterator[ServerSentEvent]: shutdown_grace_period=self._shutdown_grace_period, ) - @rest_error_handler + @_a2a_rest_error_handler async def on_message_send(self, request: Request) -> Response: """Handles the 'message/send' REST method.""" @@ -191,7 +201,7 @@ async def _handler( response = await self._handle_non_streaming(request, _handler) return A2AJSONResponse(content=MessageToDict(response)) - @rest_stream_error_handler + @_a2a_rest_stream_error_handler async def on_message_send_stream( self, request: Request ) -> EventSourceResponse: @@ -212,7 +222,7 @@ async def _handler( return await self._handle_streaming(request, _handler) - @rest_error_handler + @_a2a_rest_error_handler async def on_cancel_task(self, request: Request) -> Response: """Handles the 'tasks/cancel' REST method.""" @@ -231,7 +241,7 @@ async def _handler(context: ServerCallContext) -> a2a_pb2.Task: response = await self._handle_non_streaming(request, _handler) return A2AJSONResponse(content=MessageToDict(response)) - @rest_stream_error_handler + @_a2a_rest_stream_error_handler async def on_subscribe_to_task( self, request: Request ) -> EventSourceResponse: @@ -250,7 +260,7 @@ async def _handler( return await self._handle_streaming(request, _handler) - @rest_error_handler + @_a2a_rest_error_handler async def on_get_task(self, request: Request) -> Response: """Handles the 'tasks/{id}' REST method.""" @@ -267,7 +277,7 @@ async def _handler(context: ServerCallContext) -> a2a_pb2.Task: response = await self._handle_non_streaming(request, _handler) return A2AJSONResponse(content=MessageToDict(response)) - @rest_error_handler + @_a2a_rest_error_handler async def get_push_notification(self, request: Request) -> Response: """Handles the 'tasks/pushNotificationConfig/get' REST method.""" @@ -289,7 +299,7 @@ async def _handler( response = await self._handle_non_streaming(request, _handler) return A2AJSONResponse(content=MessageToDict(response)) - @rest_error_handler + @_a2a_rest_error_handler async def delete_push_notification(self, request: Request) -> Response: """Handles the 'tasks/pushNotificationConfig/delete' REST method.""" @@ -307,7 +317,7 @@ async def _handler(context: ServerCallContext) -> None: await self._handle_non_streaming(request, _handler) return A2AJSONResponse(content={}) - @rest_error_handler + @_a2a_rest_error_handler async def set_push_notification(self, request: Request) -> Response: """Handles the 'tasks/pushNotificationConfig/set' REST method.""" @@ -326,7 +336,7 @@ async def _handler( response = await self._handle_non_streaming(request, _handler) return A2AJSONResponse(content=MessageToDict(response)) - @rest_error_handler + @_a2a_rest_error_handler async def list_push_notifications(self, request: Request) -> Response: """Handles the 'tasks/pushNotificationConfig/list' REST method.""" @@ -344,7 +354,7 @@ async def _handler( response = await self._handle_non_streaming(request, _handler) return A2AJSONResponse(content=MessageToDict(response)) - @rest_error_handler + @_a2a_rest_error_handler async def list_tasks(self, request: Request) -> Response: """Handles the 'tasks/list' REST method.""" params = a2a_pb2.ListTasksRequest() @@ -363,7 +373,7 @@ async def _handler( ) ) - @rest_error_handler + @_a2a_rest_error_handler async def handle_authenticated_agent_card( self, request: Request ) -> Response: diff --git a/src/a2a/utils/error_handlers.py b/src/a2a/utils/error_handlers.py index ef8182940..0d1018918 100644 --- a/src/a2a/utils/error_handlers.py +++ b/src/a2a/utils/error_handlers.py @@ -35,6 +35,7 @@ ERROR_INFO_TYPE = 'type.googleapis.com/google.rpc.ErrorInfo' BAD_REQUEST_TYPE = 'type.googleapis.com/google.rpc.BadRequest' A2A_DOMAIN = 'a2a-protocol.org' +JSON_MEDIA_TYPE = 'application/json' def _error_info( @@ -131,7 +132,9 @@ def build_rest_error_payload(error: Exception) -> dict[str, Any]: ) -def _create_error_response(error: Exception) -> Response: +def _create_error_response( + error: Exception, media_type: str = JSON_MEDIA_TYPE +) -> Response: """Helper function to create a JSONResponse for an error.""" if isinstance(error, A2AError): log_level = ( @@ -157,29 +160,43 @@ def _create_error_response(error: Exception) -> Response: return JSONResponse( content=payload, status_code=http_code, - media_type='application/json', + media_type=media_type, ) def rest_error_handler( func: Callable[..., Awaitable[Response]], + *, + media_type: str = JSON_MEDIA_TYPE, ) -> Callable[..., Awaitable[Response]]: - """Decorator to catch A2AError and map it to an appropriate JSONResponse.""" + """Decorator to catch A2AError and map it to an appropriate JSONResponse. + + Args: + func: The handler to wrap. + media_type: The Content-Type of the error response. + """ @functools.wraps(func) async def wrapper(*args: Any, **kwargs: Any) -> Response: try: return await func(*args, **kwargs) except Exception as error: # noqa: BLE001 - return _create_error_response(error) + return _create_error_response(error, media_type) return wrapper def rest_stream_error_handler( func: Callable[..., Coroutine[Any, Any, Any]], + *, + media_type: str = JSON_MEDIA_TYPE, ) -> Callable[..., Coroutine[Any, Any, Any]]: - """Decorator to catch A2AError for a streaming method. Maps synchronous errors to JSONResponse and logs streaming errors.""" + """Decorator to catch A2AError for a streaming method. Maps synchronous errors to JSONResponse and logs streaming errors. + + Args: + func: The handler to wrap. + media_type: The Content-Type of a synchronous error response. + """ def _log_error(error: Exception) -> None: if isinstance(error, A2AError): @@ -223,7 +240,7 @@ async def error_catching_iterator() -> AsyncGenerator[ response.body_iterator = error_catching_iterator() except Exception as e: # noqa: BLE001 - return _create_error_response(e) + return _create_error_response(e, media_type) else: return response diff --git a/tests/server/routes/test_rest_routes.py b/tests/server/routes/test_rest_routes.py index a0450b2b0..b67a52a76 100644 --- a/tests/server/routes/test_rest_routes.py +++ b/tests/server/routes/test_rest_routes.py @@ -5,6 +5,7 @@ from a2a.server.request_handlers.request_handler import RequestHandler from a2a.server.routes.rest_routes import create_rest_routes from a2a.types.a2a_pb2 import AgentCard, ListTasksResponse, Task +from a2a.utils.errors import TaskNotFoundError from starlette.applications import Starlette from starlette.routing import BaseRoute from starlette.testclient import TestClient @@ -179,3 +180,28 @@ def test_list_tasks_response_uses_the_a2a_media_type(agent_card, mock_handler): response = client.get('/tasks', headers={'A2A-Version': '1.0'}) assert response.headers['content-type'].startswith('application/a2a+json') + + +def test_error_response_uses_the_a2a_media_type(agent_card, mock_handler): + """Error responses are labelled application/a2a+json too (spec 11.1).""" + mock_handler.on_get_task.side_effect = TaskNotFoundError() + client = TestClient(Starlette(routes=create_rest_routes(mock_handler))) + + response = client.get('/tasks/missing', headers={'A2A-Version': '1.0'}) + + assert response.status_code == 404 + assert response.headers['content-type'].startswith('application/a2a+json') + + +def test_stream_error_response_uses_the_a2a_media_type( + agent_card, mock_handler +): + """An error raised before the first stream event is labelled the same way.""" + client = TestClient(Starlette(routes=create_rest_routes(mock_handler))) + + response = client.post( + '/message:stream', json={}, headers={'A2A-Version': '9.9'} + ) + + assert response.status_code == 400 + assert response.headers['content-type'].startswith('application/a2a+json')