Merge branch 'main' into refactor-apps2

Author: holtskinner

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

@@ -1,12 +1,16 @@
 import logging
 
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from typing import Any
 
 from fastapi import FastAPI
 
 from a2a.server.apps.jsonrpc.jsonrpc_app import (
     JSONRPCApplication,
 )
+from a2a.server.request_handlers.jsonrpc_handler import RequestHandler
+from a2a.types import A2ARequest
 from a2a.utils.constants import (
     AGENT_CARD_WELL_KNOWN_PATH,
     DEFAULT_RPC_URL,
@@ -40,7 +44,22 @@ def add_routes_to_app(
             rpc_url: The URL for the A2A JSON-RPC endpoint.
             extended_agent_card_url: The URL for the authenticated extended agent card endpoint.
         """
-        app.post(rpc_url)(self._handle_requests)
+        app.post(
+            rpc_url,
+            openapi_extra={
+                'requestBody': {
+                    'content': {
+                        'application/json': {
+                            'schema': {
+                                '$ref': '#/components/schemas/A2ARequest'
+                            }
+                        }
+                    },
+                    'required': True,
+                    'description': 'A2ARequest',
+                }
+            },
+        )(self._handle_requests)
         app.get(agent_card_url)(self._handle_get_agent_card)
 
         if self.agent_card.supportsAuthenticatedExtendedCard:
@@ -66,7 +85,23 @@ def build(
         Returns:
             A configured FastAPI application instance.
         """
-        app = FastAPI(**kwargs)
+
+        @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)
 
         self.add_routes_to_app(
             app, agent_card_url, rpc_url, extended_agent_card_url