diff --git a/src/a2a/server/routes/common.py b/src/a2a/server/routes/common.py index 18b6865c5..c46847488 100644 --- a/src/a2a/server/routes/common.py +++ b/src/a2a/server/routes/common.py @@ -1,6 +1,10 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any +from google.protobuf.json_format import MessageToDict + +from a2a.types.a2a_pb2 import ListTasksResponse + if TYPE_CHECKING: from starlette.authentication import BaseUser @@ -21,6 +25,21 @@ from a2a.server.context import ServerCallContext +def serialize_list_tasks_response( + response: ListTasksResponse, include_artifacts: bool +) -> dict[str, Any]: + """Serializes a ListTasks response according to artifact semantics.""" + result = MessageToDict( + response, + preserving_proto_field_name=False, + always_print_fields_with_no_presence=True, + ) + if not include_artifacts: + for task in result['tasks']: + task.pop('artifacts', None) + return result + + class StarletteUser(User): """Adapts a Starlette BaseUser to the A2A User interface.""" diff --git a/src/a2a/server/routes/jsonrpc_dispatcher.py b/src/a2a/server/routes/jsonrpc_dispatcher.py index b59ed7551..228a57579 100644 --- a/src/a2a/server/routes/jsonrpc_dispatcher.py +++ b/src/a2a/server/routes/jsonrpc_dispatcher.py @@ -26,6 +26,7 @@ from a2a.server.routes.common import ( DefaultServerCallContextBuilder, ServerCallContextBuilder, + serialize_list_tasks_response, ) from a2a.types.a2a_pb2 import ( CancelTaskRequest, @@ -431,10 +432,8 @@ async def _handle_list_tasks( tasks_response = await self.request_handler.on_list_tasks( request_obj, context ) - return MessageToDict( - tasks_response, - preserving_proto_field_name=False, - always_print_fields_with_no_presence=True, + return serialize_list_tasks_response( + tasks_response, request_obj.include_artifacts ) async def _handle_create_task_push_notification_config( diff --git a/src/a2a/server/routes/rest_dispatcher.py b/src/a2a/server/routes/rest_dispatcher.py index 06c73c592..f2d81271a 100644 --- a/src/a2a/server/routes/rest_dispatcher.py +++ b/src/a2a/server/routes/rest_dispatcher.py @@ -10,6 +10,7 @@ from a2a.server.routes.common import ( DefaultServerCallContextBuilder, ServerCallContextBuilder, + serialize_list_tasks_response, ) from a2a.types import a2a_pb2 from a2a.types.a2a_pb2 import ( @@ -327,19 +328,19 @@ async def _handler( @rest_error_handler async def list_tasks(self, request: Request) -> Response: """Handles the 'tasks/list' REST method.""" + params = a2a_pb2.ListTasksRequest() @validate_version(constants.PROTOCOL_VERSION_1_0) async def _handler( context: ServerCallContext, ) -> a2a_pb2.ListTasksResponse: - params = a2a_pb2.ListTasksRequest() proto_utils.parse_params(request.query_params, params) return await self.request_handler.on_list_tasks(params, context) response = await self._handle_non_streaming(request, _handler) return JSONResponse( - content=MessageToDict( - response, always_print_fields_with_no_presence=True + content=serialize_list_tasks_response( + response, params.include_artifacts ) ) diff --git a/tests/server/routes/test_jsonrpc_dispatcher.py b/tests/server/routes/test_jsonrpc_dispatcher.py index 3bde4fc2e..ed55c476c 100644 --- a/tests/server/routes/test_jsonrpc_dispatcher.py +++ b/tests/server/routes/test_jsonrpc_dispatcher.py @@ -364,6 +364,32 @@ def test_list_tasks_routes_to_on_list_tasks(self, client, handler): call_context = handler.on_list_tasks.call_args[0][1] assert call_context.state['method'] == 'ListTasks' + @pytest.mark.parametrize( + ('params', 'artifacts_expected'), + [ + pytest.param({}, False, id='default'), + pytest.param({'includeArtifacts': True}, True, id='included'), + ], + ) + def test_list_tasks_artifact_presence( + self, + client: TestClient, + handler: AsyncMock, + params: dict[str, Any], + artifacts_expected: bool, + ) -> None: + handler.on_list_tasks.return_value = ListTasksResponse( + tasks=[Task(id='task1')] + ) + + response = client.post( + '/', json=_make_jsonrpc_request('ListTasks', params) + ) + response.raise_for_status() + + task = response.json()['result']['tasks'][0] + assert ('artifacts' in task) is artifacts_expected + def test_create_push_notification_config_routes_correctly( self, client, handler ): diff --git a/tests/server/routes/test_rest_dispatcher.py b/tests/server/routes/test_rest_dispatcher.py index 00cc45f7b..a062a10c1 100644 --- a/tests/server/routes/test_rest_dispatcher.py +++ b/tests/server/routes/test_rest_dispatcher.py @@ -73,7 +73,7 @@ def rest_dispatcher_instance(mock_handler): return RestDispatcher(request_handler=mock_handler) -from starlette.datastructures import Headers +from starlette.datastructures import Headers, QueryParams def make_mock_request( @@ -216,6 +216,34 @@ async def test_list_tasks(self, rest_dispatcher_instance, mock_handler): response = await rest_dispatcher_instance.list_tasks(req) assert response.status_code == 200 + @pytest.mark.parametrize( + ('query_params', 'artifacts_expected'), + [ + pytest.param(QueryParams(), False, id='default'), + pytest.param( + QueryParams({'includeArtifacts': 'true'}), + True, + id='included', + ), + ], + ) + async def test_list_tasks_artifact_presence( + self, + rest_dispatcher_instance: RestDispatcher, + mock_handler: AsyncMock, + query_params: QueryParams, + artifacts_expected: bool, + ) -> None: + mock_handler.on_list_tasks.return_value = ListTasksResponse( + tasks=[Task(id='test_task')] + ) + req = make_mock_request(method='GET', query_params=query_params) + + response = await rest_dispatcher_instance.list_tasks(req) + + task = json.loads(response.body)['tasks'][0] + assert ('artifacts' in task) is artifacts_expected + async def test_get_push_notification( self, rest_dispatcher_instance, mock_handler ):