diff --git a/src/select_ai/agent/a2a/server.py b/src/select_ai/agent/a2a/server.py index 6b28a26..1c60014 100644 --- a/src/select_ai/agent/a2a/server.py +++ b/src/select_ai/agent/a2a/server.py @@ -22,6 +22,7 @@ from a2a.server.routes import create_jsonrpc_routes from a2a.server.tasks import TaskUpdater from a2a.types import AgentCapabilities, AgentCard, AgentInterface, AgentSkill +from google.protobuf.json_format import ParseDict from starlette.applications import Starlette from starlette.responses import JSONResponse from starlette.routing import Route @@ -32,11 +33,9 @@ from select_ai.agent.a2a.task_store import OracleTaskStore from select_ai.version import __version__ -_A2UI_MIME_TYPE = "application/a2ui+json" - -def _a2ui_payload(result: str | None) -> dict | None: - """Return an A2UI response envelope, if ``RUN_TEAM`` returned one.""" +def _message_parts(result: str | None): + """Convert a serialized A2A message into its constituent parts.""" if not result: return None try: @@ -45,11 +44,24 @@ def _a2ui_payload(result: str | None) -> dict | None: return None if not isinstance(payload, dict): return None - if payload.get("metadata", {}).get("mimeType") != _A2UI_MIME_TYPE: - return None - if not isinstance(payload.get("data"), list): + message_parts = payload.get("parts") + if payload.get("kind") != "message" or not isinstance(message_parts, list): return None - return payload + parts = [] + for part in message_parts: + if not isinstance(part, dict): + return None + if part.get("kind") == "text" and isinstance(part.get("text"), str): + parts.append(new_text_part(part["text"])) + elif part.get("kind") == "data" and "data" in part: + output_part = new_data_part(part["data"]) + if isinstance(part.get("metadata"), dict): + ParseDict(part["metadata"], output_part.metadata) + parts.append(output_part) + else: + # Avoid silently discarding an unsupported part type. + return None + return parts or None class DatabaseTeamExecutor(AgentExecutor): @@ -81,11 +93,11 @@ async def execute(self, context, event_queue): prompt=context.get_user_input(), params={"conversation_id": conversation_id}, ) - a2ui_payload = _a2ui_payload(result) + message_parts = _message_parts(result) await updater.add_artifact( parts=( - [new_data_part(a2ui_payload)] - if a2ui_payload is not None + message_parts + if message_parts is not None else [new_text_part(result or "")] ), name="database-agent-result", diff --git a/tests/a2a/test_a2ui_parts.py b/tests/a2a/test_a2ui_parts.py new file mode 100644 index 0000000..82a1e40 --- /dev/null +++ b/tests/a2a/test_a2ui_parts.py @@ -0,0 +1,64 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2026, Oracle and/or its affiliates. +# +# Licensed under the Universal Permissive License v 1.0 as shown at +# https://oss.oracle.com/licenses/upl. +# ----------------------------------------------------------------------------- + +import json + +import pytest +from google.protobuf.json_format import MessageToDict + +pytest.importorskip("a2a") + +from select_ai.agent.a2a.server import _message_parts + + +def test_a2a_message_parts_are_forwarded_with_metadata(): + result = json.dumps( + { + "kind": "message", + "parts": [ + {"kind": "text", "text": "A2UI visualization ready."}, + { + "kind": "data", + "data": { + "version": "v0.9", + "createSurface": {"surfaceId": "smoke-test"}, + }, + "metadata": {"mimeType": "application/json+a2ui"}, + }, + ], + } + ) + + parts = _message_parts(result) + + assert [ + MessageToDict(part, preserving_proto_field_name=True) for part in parts + ] == [ + {"text": "A2UI visualization ready."}, + { + "data": { + "version": "v0.9", + "createSurface": {"surfaceId": "smoke-test"}, + }, + "metadata": {"mimeType": "application/json+a2ui"}, + }, + ] + + +def test_a2a_text_message_is_forwarded(): + result = json.dumps( + { + "kind": "message", + "parts": [{"kind": "text", "text": "ordinary response"}], + } + ) + + parts = _message_parts(result) + + assert [ + MessageToDict(part, preserving_proto_field_name=True) for part in parts + ] == [{"text": "ordinary response"}]