Merge branch 'main' into fix-bug#367-client_hangs

Author: holtskinner

pyproject.toml

@@ -14,7 +14,7 @@ dependencies = [
   "pydantic>=2.11.3",
   "sse-starlette",
   "starlette",
-  "protobuf==5.29.5",
+  "protobuf>=5.29.5",
   "google-api-core>=1.26.0",
 ]
 

src/a2a/client/transports/grpc.py

@@ -87,7 +87,7 @@ async def send_message(
                 metadata=proto_utils.ToProto.metadata(request.metadata),
             )
         )
-        if response.task:
+        if response.HasField('task'):
             return proto_utils.FromProto.task(response.task)
         return proto_utils.FromProto.message(response.msg)
 

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

@@ -1,7 +1,5 @@
 import logging
 
-from collections.abc import AsyncIterator
-from contextlib import asynccontextmanager
 from typing import Any
 
 from fastapi import FastAPI
@@ -21,6 +19,28 @@
 logger = logging.getLogger(__name__)
 
 
+class A2AFastAPI(FastAPI):
+    """A FastAPI application that adds A2A-specific OpenAPI components."""
+
+    _a2a_components_added: bool = False
+
+    def openapi(self) -> dict[str, Any]:
+        """Generates the OpenAPI schema for the application."""
+        openapi_schema = super().openapi()
+        if not self._a2a_components_added:
+            a2a_request_schema = A2ARequest.model_json_schema(
+                ref_template='#/components/schemas/{model}'
+            )
+            defs = a2a_request_schema.pop('$defs', {})
+            component_schemas = openapi_schema.setdefault(
+                'components', {}
+            ).setdefault('schemas', {})
+            component_schemas.update(defs)
+            component_schemas['A2ARequest'] = a2a_request_schema
+            self._a2a_components_added = True
+        return openapi_schema
+
+
 class A2AFastAPIApplication(JSONRPCApplication):
     """A FastAPI application implementing the A2A protocol server endpoints.
 
@@ -92,23 +112,7 @@ def build(
         Returns:
             A configured FastAPI application instance.
         """
-
-        @asynccontextmanager
-        async def lifespan(app: FastAPI) -> AsyncIterator[None]:
-            a2a_request_schema = A2ARequest.model_json_schema(
-                ref_template='#/components/schemas/{model}'
-            )
-            defs = a2a_request_schema.pop('$defs', {})
-            openapi_schema = app.openapi()
-            component_schemas = openapi_schema.setdefault(
-                'components', {}
-            ).setdefault('schemas', {})
-            component_schemas.update(defs)
-            component_schemas['A2ARequest'] = a2a_request_schema
-
-            yield
-
-        app = FastAPI(lifespan=lifespan, **kwargs)
+        app = A2AFastAPI(**kwargs)
 
         self.add_routes_to_app(
             app, agent_card_url, rpc_url, extended_agent_card_url

tests/client/test_grpc_client.py

@@ -18,7 +18,7 @@
     TaskStatus,
     TextPart,
 )
-from a2a.utils import proto_utils
+from a2a.utils import get_text_parts, proto_utils
 
 
 # Fixtures
@@ -112,6 +112,28 @@ async def test_send_message_task_response(
     assert response.id == sample_task.id
 
 
+@pytest.mark.asyncio
+async def test_send_message_message_response(
+    grpc_transport: GrpcTransport,
+    mock_grpc_stub: AsyncMock,
+    sample_message_send_params: MessageSendParams,
+    sample_message: Message,
+):
+    """Test send_message that returns a Message."""
+    mock_grpc_stub.SendMessage.return_value = a2a_pb2.SendMessageResponse(
+        msg=proto_utils.ToProto.message(sample_message)
+    )
+
+    response = await grpc_transport.send_message(sample_message_send_params)
+
+    mock_grpc_stub.SendMessage.assert_awaited_once()
+    assert isinstance(response, Message)
+    assert response.message_id == sample_message.message_id
+    assert get_text_parts(response.parts) == get_text_parts(
+        sample_message.parts
+    )
+
+
 @pytest.mark.asyncio
 async def test_get_task(
     grpc_transport: GrpcTransport, mock_grpc_stub: AsyncMock, sample_task: Task

tests/server/apps/jsonrpc/test_serialization.py

@@ -1,6 +1,7 @@
 from unittest import mock
 
 import pytest
+from fastapi import FastAPI
 
 from pydantic import ValidationError
 from starlette.testclient import TestClient
@@ -183,3 +184,21 @@ def test_handle_unicode_characters(agent_card_with_api_key: AgentCard):
     data = response.json()
     assert 'error' not in data or data['error'] is None
     assert data['result']['parts'][0]['text'] == f'Received: {unicode_text}'
+
+
+def test_fastapi_sub_application(agent_card_with_api_key: AgentCard):
+    """
+    Tests that the A2AFastAPIApplication endpoint correctly passes the url in sub-application.
+    """
+    handler = mock.AsyncMock()
+    sub_app_instance = A2AFastAPIApplication(agent_card_with_api_key, handler)
+    app_instance = FastAPI()
+    app_instance.mount('/a2a', sub_app_instance.build())
+    client = TestClient(app_instance)
+
+    response = client.get('/a2a/openapi.json')
+    assert response.status_code == 200
+    response_data = response.json()
+
+    assert 'servers' in response_data
+    assert response_data['servers'] == [{'url': '/a2a'}]