Merge branch 'main' into feat/add-base-model

Author: mikeas1

.github/workflows/unit-tests.yml

@@ -45,6 +45,6 @@ jobs:
       - name: Install dependencies
         run: uv sync --dev --extra sql
       - name: Run tests and check coverage
-        run: uv run pytest --cov=a2a --cov-report=xml --cov-fail-under=90
+        run: uv run pytest --cov=a2a --cov-report=xml --cov-fail-under=89
       - name: Show coverage summary in log
         run: uv run coverage report

src/a2a/grpc/a2a_pb2.py

@@ -12,8 +12,10 @@
 from sse_starlette.sse import EventSourceResponse
 from starlette.applications import Starlette
 from starlette.authentication import BaseUser
+from starlette.exceptions import HTTPException
 from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
+from starlette.status import HTTP_413_REQUEST_ENTITY_TOO_LARGE
 
 from a2a.auth.user import UnauthenticatedUser
 from a2a.auth.user import User as A2AUser
@@ -45,6 +47,7 @@
 from a2a.utils.constants import (
     AGENT_CARD_WELL_KNOWN_PATH,
     DEFAULT_RPC_URL,
+    EXTENDED_AGENT_CARD_PATH,
 )
 from a2a.utils.errors import MethodNotImplementedError
 
@@ -176,7 +179,7 @@ def _generate_error_response(
             status_code=200,
         )
 
-    async def _handle_requests(self, request: Request) -> Response:
+    async def _handle_requests(self, request: Request) -> Response:  # noqa: PLR0911
         """Handles incoming POST requests to the main A2A endpoint.
 
         Parses the request body as JSON, validates it against A2A request types,
@@ -232,6 +235,15 @@ async def _handle_requests(self, request: Request) -> Response:
                 request_id,
                 A2AError(root=InvalidRequestError(data=json.loads(e.json()))),
             )
+        except HTTPException as e:
+            if e.status_code == HTTP_413_REQUEST_ENTITY_TOO_LARGE:
+                return self._generate_error_response(
+                    request_id,
+                    A2AError(
+                        root=InvalidRequestError(message='Payload too large')
+                    ),
+                )
+            raise e
         except Exception as e:
             logger.error(f'Unhandled exception: {e}')
             traceback.print_exc()
@@ -438,13 +450,16 @@ def build(
         self,
         agent_card_url: str = AGENT_CARD_WELL_KNOWN_PATH,
         rpc_url: str = DEFAULT_RPC_URL,
+        extended_agent_card_url: str = EXTENDED_AGENT_CARD_PATH,
         **kwargs: Any,
     ) -> FastAPI | Starlette:
         """Builds and returns the JSONRPC application instance.
 
         Args:
             agent_card_url: The URL for the agent card endpoint.
-            rpc_url: The URL for the A2A JSON-RPC endpoint
+            rpc_url: The URL for the A2A JSON-RPC endpoint.
+            extended_agent_card_url: The URL for the authenticated extended
+              agent card endpoint.
             **kwargs: Additional keyword arguments to pass to the FastAPI constructor.
 
         Returns:

src/a2a/server/apps/jsonrpc/jsonrpc_app.py

@@ -4,6 +4,8 @@
 
 from collections.abc import AsyncGenerator
 
+from pydantic import ValidationError
+
 from a2a.server.events.event_queue import Event, EventQueue
 from a2a.types import (
     InternalError,
@@ -138,6 +140,9 @@ async def consume_all(self) -> AsyncGenerator[Event]:
                 # python 3.12 and get a queue empty error on an open queue
                 if self.queue.is_closed():
                     break
+            except ValidationError as e:
+                logger.error(f'Invalid event format received: {e}')
+                continue
             except Exception as e:
                 logger.error(
                     f'Stopping event consumption due to exception: {e}'

src/a2a/server/events/event_consumer.py

@@ -1,6 +1,8 @@
 from unittest import mock
 
 import pytest
+
+from pydantic import ValidationError
 from starlette.testclient import TestClient
 
 from a2a.server.apps import A2AFastAPIApplication, A2AStarletteApplication
@@ -9,9 +11,14 @@
     AgentCapabilities,
     AgentCard,
     In,
+    InvalidRequestError,
+    JSONParseError,
+    Message,
+    Part,
+    Role,
     SecurityScheme,
+    TextPart,
 )
-from pydantic import ValidationError
 
 
 @pytest.fixture
@@ -92,3 +99,88 @@ def test_fastapi_agent_card_with_api_key_scheme_alias(
     assert 'in' in security_scheme_json
     assert 'in_' not in security_scheme_json
     assert security_scheme_json['in'] == 'header'
+
+
+def test_handle_invalid_json(agent_card_with_api_key: AgentCard):
+    """Test handling of malformed JSON."""
+    handler = mock.AsyncMock()
+    app_instance = A2AStarletteApplication(agent_card_with_api_key, handler)
+    client = TestClient(app_instance.build())
+
+    response = client.post(
+        '/',
+        content='{ "jsonrpc": "2.0", "method": "test", "id": 1, "params": { "key": "value" }',
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert data['error']['code'] == JSONParseError().code
+
+
+def test_handle_oversized_payload(agent_card_with_api_key: AgentCard):
+    """Test handling of oversized JSON payloads."""
+    handler = mock.AsyncMock()
+    app_instance = A2AStarletteApplication(agent_card_with_api_key, handler)
+    client = TestClient(app_instance.build())
+
+    large_string = 'a' * 2_000_000  # 2MB string
+    payload = {
+        'jsonrpc': '2.0',
+        'method': 'test',
+        'id': 1,
+        'params': {'data': large_string},
+    }
+
+    # Starlette/FastAPI's default max request size is around 1MB.
+    # This test will likely fail with a 413 Payload Too Large if the default is not increased.
+    # If the application is expected to handle larger payloads, the server configuration needs to be adjusted.
+    # For this test, we expect a 413 or a graceful JSON-RPC error if the app handles it.
+
+    try:
+        response = client.post('/', json=payload)
+        # If the app handles it gracefully and returns a JSON-RPC error
+        if response.status_code == 200:
+            data = response.json()
+            assert data['error']['code'] == InvalidRequestError().code
+        else:
+            assert response.status_code == 413
+    except Exception as e:
+        # Depending on server setup, it might just drop the connection for very large payloads
+        assert isinstance(e, (ConnectionResetError, RuntimeError))
+
+
+def test_handle_unicode_characters(agent_card_with_api_key: AgentCard):
+    """Test handling of unicode characters in JSON payload."""
+    handler = mock.AsyncMock()
+    app_instance = A2AStarletteApplication(agent_card_with_api_key, handler)
+    client = TestClient(app_instance.build())
+
+    unicode_text = 'こんにちは世界'  # "Hello world" in Japanese
+    unicode_payload = {
+        'jsonrpc': '2.0',
+        'method': 'message/send',
+        'id': 'unicode_test',
+        'params': {
+            'message': {
+                'role': 'user',
+                'parts': [{'kind': 'text', 'text': unicode_text}],
+                'messageId': 'msg-unicode',
+            }
+        },
+    }
+
+    # Mock a handler for this method
+    handler.on_message_send.return_value = Message(
+        role=Role.agent,
+        parts=[Part(root=TextPart(text=f'Received: {unicode_text}'))],
+        messageId='response-unicode',
+    )
+
+    response = client.post('/', json=unicode_payload)
+
+    # We are not testing the handler logic here, just that the server can correctly
+    # deserialize the unicode payload without errors. A 200 response with any valid
+    # JSON-RPC response indicates success.
+    assert response.status_code == 200
+    data = response.json()
+    assert 'error' not in data or data['error'] is None
+    assert data['result']['parts'][0]['text'] == f'Received: {unicode_text}'

tests/server/apps/jsonrpc/test_serialization.py

@@ -5,6 +5,8 @@
 
 import pytest
 
+from pydantic import ValidationError
+
 from a2a.server.events.event_consumer import EventConsumer, QueueClosed
 from a2a.server.events.event_queue import EventQueue
 from a2a.types import (
@@ -343,3 +345,28 @@ def test_agent_task_callback_no_exception(event_consumer: EventConsumer):
 
     assert event_consumer._exception is None  # Should remain None
     mock_task.exception.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_consume_all_handles_validation_error(
+    event_consumer: EventConsumer, mock_event_queue: AsyncMock
+):
+    """Test that consume_all gracefully handles a pydantic.ValidationError."""
+    # Simulate dequeue_event raising a ValidationError
+    mock_event_queue.dequeue_event.side_effect = [
+        ValidationError.from_exception_data(title='Test Error', line_errors=[]),
+        asyncio.CancelledError,  # To stop the loop for the test
+    ]
+
+    with patch(
+        'a2a.server.events.event_consumer.logger.error'
+    ) as logger_error_mock:
+        with pytest.raises(asyncio.CancelledError):
+            async for _ in event_consumer.consume_all():
+                pass
+
+        # Check that the specific error was logged and the consumer continued
+        logger_error_mock.assert_called_once()
+        assert (
+            'Invalid event format received' in logger_error_mock.call_args[0][0]
+        )