diff --git a/src/a2a/server/apps/jsonrpc/fastapi_app.py b/src/a2a/server/apps/jsonrpc/fastapi_app.py index 20acfc575..0ec9d1ab2 100644 --- a/src/a2a/server/apps/jsonrpc/fastapi_app.py +++ b/src/a2a/server/apps/jsonrpc/fastapi_app.py @@ -1,5 +1,3 @@ -import importlib.resources -import json import logging from collections.abc import Awaitable, Callable @@ -36,43 +34,6 @@ 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.""" - if self.openapi_schema: - return self.openapi_schema - - # Try to use the a2a.json schema generated from the proto file - # if available, instead of generating one from the python types. - try: - from a2a import types # noqa: PLC0415 - - schema_file = importlib.resources.files(types).joinpath('a2a.json') - if schema_file.is_file(): - self.openapi_schema = json.loads( - schema_file.read_text(encoding='utf-8') - ) - if self.openapi_schema: - return self.openapi_schema - except Exception: # noqa: BLE001 - logger.warning( - "Could not load 'a2a.json' from 'a2a.types'. Falling back to auto-generation." - ) - - openapi_schema = super().openapi() - if not self._a2a_components_added: - # A2ARequest is now a Union type of proto messages, so we can't use - # model_json_schema. Instead, we just mark it as added without - # adding the schema since proto types don't have Pydantic schemas. - # The OpenAPI schema will still be functional for the endpoints. - self._a2a_components_added = True - return openapi_schema - - class A2AFastAPIApplication(JSONRPCApplication): """A FastAPI application implementing the A2A protocol server endpoints. @@ -180,7 +141,7 @@ def build( Returns: A configured FastAPI application instance. """ - app = A2AFastAPI(**kwargs) + app = FastAPI(**kwargs) self.add_routes_to_app(app, agent_card_url, rpc_url) diff --git a/src/a2a/server/apps/rest/fastapi_app.py b/src/a2a/server/apps/rest/fastapi_app.py index ea9a501b9..aac13b382 100644 --- a/src/a2a/server/apps/rest/fastapi_app.py +++ b/src/a2a/server/apps/rest/fastapi_app.py @@ -1,3 +1,5 @@ +import importlib.resources +import json import logging from collections.abc import Awaitable, Callable @@ -56,6 +58,45 @@ } +class A2AFastAPI(FastAPI): + """A FastAPI application that adds A2A-specific OpenAPI components.""" + + _a2a_components_added: bool = False + rpc_url: str = '' + + def openapi(self) -> dict[str, Any]: + """Generates the OpenAPI schema for the application.""" + if self.openapi_schema: + return self.openapi_schema + + # Try to use the a2a.json schema generated from the proto file + # if available, instead of generating one from the python types. + try: + from a2a import types # noqa: PLC0415 + + schema_file = importlib.resources.files(types).joinpath('a2a.json') + if schema_file.is_file(): + self.openapi_schema = json.loads( + schema_file.read_text(encoding='utf-8') + ) + if self.rpc_url and self.openapi_schema: + paths = self.openapi_schema.get('paths', {}) + self.openapi_schema['paths'] = { + f'{self.rpc_url}{k}': v for k, v in paths.items() + } + if self.openapi_schema: + return self.openapi_schema + except Exception: # noqa: BLE001 + logger.warning( + "Could not load 'a2a.json' from 'a2a.types'. Falling back to auto-generation." + ) + + openapi_schema = super().openapi() + if not self._a2a_components_added: + self._a2a_components_added = True + return openapi_schema + + class A2ARESTFastAPIApplication: """A FastAPI application implementing the A2A protocol server REST endpoints. @@ -140,7 +181,8 @@ def build( Returns: A configured FastAPI application instance. """ - app = FastAPI(**kwargs) + app = A2AFastAPI(**kwargs) + app.rpc_url = rpc_url @app.exception_handler(StarletteHTTPException) async def http_exception_handler( diff --git a/tck/sut_agent.py b/tck/sut_agent.py index 7196b828b..5af6309db 100644 --- a/tck/sut_agent.py +++ b/tck/sut_agent.py @@ -8,6 +8,7 @@ import grpc.aio import uvicorn +from fastapi import FastAPI from starlette.applications import Starlette import a2a.compat.v0_3.a2a_v0_3_pb2_grpc as a2a_v0_3_grpc @@ -19,6 +20,7 @@ from a2a.server.apps import ( A2ARESTFastAPIApplication, A2AStarletteApplication, + A2AFastAPIApplication, ) from a2a.server.events.event_queue import EventQueue from a2a.server.request_handlers.default_request_handler import ( @@ -196,22 +198,24 @@ def serve(task_store: TaskStore) -> None: task_store=task_store, ) - main_app = Starlette() + main_app = FastAPI() # JSONRPC - jsonrpc_server = A2AStarletteApplication( + jsonrpc_server = A2AFastAPIApplication( agent_card=agent_card, http_handler=request_handler, ) - jsonrpc_server.add_routes_to_app(main_app, rpc_url=JSONRPC_URL) + json_rpc_app = jsonrpc_server.build(rpc_url=JSONRPC_URL) # REST rest_server = A2ARESTFastAPIApplication( agent_card=agent_card, http_handler=request_handler, ) - rest_app = rest_server.build(rpc_url=REST_URL) - main_app.mount('', rest_app) + rest_app = rest_server.build() + + main_app.mount(REST_URL,rest_app) + main_app.mount('',json_rpc_app) config = uvicorn.Config( main_app, host='127.0.0.1', port=http_port, log_level='info'