fix: openapi working in sub-app

Author: racinmat

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

@@ -1,11 +1,16 @@
 import logging
+import sys
 
-from collections.abc import AsyncIterator
-from contextlib import asynccontextmanager
 from typing import Any
 
 from fastapi import FastAPI
 
+
+if sys.version_info < (3, 12):  # pragma: no cover
+    from typing_extensions import override
+else:  # pragma: no cover
+    from typing import override
+
 from a2a.server.apps.jsonrpc.jsonrpc_app import (
     CallContextBuilder,
     JSONRPCApplication,
@@ -22,6 +27,28 @@
 logger = logging.getLogger(__name__)
 
 
+class A2AFastAPI(FastAPI):
+    """A FastAPI application that adds A2A-specific OpenAPI components."""
+
+    a2a_components_added: bool = False
+
+    @override
+    def openapi(self) -> dict[str, Any]:
+        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.
 
@@ -112,23 +139,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