diff --git a/a2a_agents/python/a2ui_agent/src/a2ui/extension/send_a2ui_to_client_toolset.py b/a2a_agents/python/a2ui_agent/src/a2ui/extension/send_a2ui_to_client_toolset.py index baba339f4..7a0ffae57 100644 --- a/a2a_agents/python/a2ui_agent/src/a2ui/extension/send_a2ui_to_client_toolset.py +++ b/a2a_agents/python/a2ui_agent/src/a2ui/extension/send_a2ui_to_client_toolset.py @@ -80,6 +80,7 @@ async def get_schema(ctx: ReadonlyContext) -> dict[str, Any]: import inspect import json import logging +import re from typing import Any, Awaitable, Callable, Optional, TypeAlias, Union import jsonschema @@ -253,17 +254,40 @@ async def run_async( f" arg {self.A2UI_JSON_ARG_NAME} " ) - a2ui_json_payload = json.loads(a2ui_json) + a2ui_schema = await self.get_a2ui_schema(tool_context) - # Auto-wrap single object in list - if not isinstance(a2ui_json_payload, list): - logger.info( - "Received a single JSON object, wrapping in a list for validation." - ) - a2ui_json_payload = [a2ui_json_payload] + try: + # Attempt to parse and validate + a2ui_json_payload = json.loads(a2ui_json) - a2ui_schema = await self.get_a2ui_schema(tool_context) - validate_a2ui_json(a2ui_json_payload, a2ui_schema) + # Auto-wrap single object in list + if not isinstance(a2ui_json_payload, list): + logger.info("Received a single JSON object, wrapping in a list for validation.") + a2ui_json_payload = [a2ui_json_payload] + + jsonschema.validate(instance=a2ui_json_payload, schema=a2ui_schema) + + except (jsonschema.exceptions.ValidationError, json.JSONDecodeError) as e: + logger.warning(f"Initial A2UI JSON validation failed: {e}") + + # Run Fixer + fixed_a2ui_json = re.sub(r",(?=\s*[\]}])", "", a2ui_json) + + if fixed_a2ui_json != a2ui_json: + # Emit Warning + logger.warning("Detected trailing commas in LLM output; applied autofix.") + + # Re-parse and Re-validate + a2ui_json_payload = json.loads(fixed_a2ui_json) + + # Auto-wrap single object in list + if not isinstance(a2ui_json_payload, list): + logger.info("Received a single JSON object, wrapping in a list for validation.") + a2ui_json_payload = [a2ui_json_payload] + + jsonschema.validate(instance=a2ui_json_payload, schema=a2ui_schema) + else: + raise e logger.info( f"Validated call to tool {self.TOOL_NAME} with {self.A2UI_JSON_ARG_NAME}" @@ -329,4 +353,4 @@ def convert_send_a2ui_to_client_genai_part_to_a2a_part( converted_part = part_converter.convert_genai_part_to_a2a_part(part) logger.info(f"Returning converted part: {converted_part}") - return [converted_part] if converted_part else [] + return [converted_part] if converted_part else [] \ No newline at end of file diff --git a/a2a_agents/python/a2ui_agent/tests/extension/test_send_a2ui_to_client_toolset.py b/a2a_agents/python/a2ui_agent/tests/extension/test_send_a2ui_to_client_toolset.py index 03fd74748..c023fd780 100644 --- a/a2a_agents/python/a2ui_agent/tests/extension/test_send_a2ui_to_client_toolset.py +++ b/a2a_agents/python/a2ui_agent/tests/extension/test_send_a2ui_to_client_toolset.py @@ -243,6 +243,33 @@ async def test_send_tool_run_async_schema_validation_fail(): assert "'text' is a required property" in result["error"] +@pytest.mark.asyncio +async def test_send_tool_run_async_handles_trailing_comma(caplog): + """Tests that the tool's run_async can handle and fix a trailing comma in the JSON.""" + tool = SendA2uiToClientToolset._SendA2uiJsonToClientTool(TEST_A2UI_SCHEMA) + tool_context_mock = MagicMock(spec=ToolContext) + tool_context_mock.state = {} + tool_context_mock.actions = MagicMock(skip_summarization=False) + + # Malformed JSON with a trailing comma in the list + malformed_a2ui_str = '[{"type": "Text", "text": "Hello"},]' + + args = { + SendA2uiToClientToolset._SendA2uiJsonToClientTool.A2UI_JSON_ARG_NAME: malformed_a2ui_str + } + + result = await tool.run_async(args=args, tool_context=tool_context_mock) + + # Assert that the fix was successful and the result is correct + expected_a2ui = [{"type": "Text", "text": "Hello"}] + assert result == { + SendA2uiToClientToolset._SendA2uiJsonToClientTool.VALIDATED_A2UI_JSON_KEY: expected_a2ui + } + + # Assert that the warning was logged + assert "Detected trailing commas in LLM output; applied autofix." in caplog.text + + # endregion # region send_a2ui_to_client_part_converter Tests