Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions src/a2a/server/routes/common.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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."""

Expand Down
7 changes: 3 additions & 4 deletions src/a2a/server/routes/jsonrpc_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from a2a.server.routes.common import (
DefaultServerCallContextBuilder,
ServerCallContextBuilder,
serialize_list_tasks_response,
)
from a2a.types.a2a_pb2 import (
CancelTaskRequest,
Expand Down Expand Up @@ -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(
Expand Down
7 changes: 4 additions & 3 deletions src/a2a/server/routes/rest_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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
)
)

Expand Down
26 changes: 26 additions & 0 deletions tests/server/routes/test_jsonrpc_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
):
Expand Down
30 changes: 29 additions & 1 deletion tests/server/routes/test_rest_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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
):
Expand Down
Loading