diff --git a/pyproject.toml b/pyproject.toml index f5b1432f6..a0fe9ef6a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ classifiers = [ ] [project.optional-dependencies] -http-server = ["sse-starlette", "starlette"] +http-server = ["sse-starlette>=3.3.0", "starlette"] fastapi = ["a2a-sdk[http-server]", "fastapi>=0.115.2"] encryption = ["cryptography>=43.0.0"] grpc = ["grpcio>=1.60", "grpcio-tools>=1.60", "grpcio_reflection>=1.7.0"] diff --git a/src/a2a/compat/v0_3/jsonrpc_adapter.py b/src/a2a/compat/v0_3/jsonrpc_adapter.py index 580034e9b..653281304 100644 --- a/src/a2a/compat/v0_3/jsonrpc_adapter.py +++ b/src/a2a/compat/v0_3/jsonrpc_adapter.py @@ -67,13 +67,15 @@ def __init__( self, http_handler: 'RequestHandler', context_builder: 'ServerCallContextBuilder | None' = None, - ): + shutdown_grace_period: float = 0, + ) -> None: self.handler = RequestHandler03( request_handler=http_handler, ) self._context_builder = V03ServerCallContextBuilder( context_builder or DefaultServerCallContextBuilder() ) + self._shutdown_grace_period = shutdown_grace_period def supports_method(self, method: str) -> bool: """Returns True if the v0.3 adapter supports the given method name.""" @@ -277,4 +279,7 @@ async def event_generator( ) } - return EventSourceResponse(event_generator(stream_gen)) + return EventSourceResponse( + event_generator(stream_gen), + shutdown_grace_period=self._shutdown_grace_period, + ) diff --git a/src/a2a/compat/v0_3/rest_adapter.py b/src/a2a/compat/v0_3/rest_adapter.py index d0c4c82c1..9d5fb7786 100644 --- a/src/a2a/compat/v0_3/rest_adapter.py +++ b/src/a2a/compat/v0_3/rest_adapter.py @@ -59,11 +59,13 @@ def __init__( self, http_handler: 'RequestHandler', context_builder: 'ServerCallContextBuilder | None' = None, - ): + shutdown_grace_period: float = 0, + ) -> None: self.handler = REST03Handler(request_handler=http_handler) self._context_builder = V03ServerCallContextBuilder( context_builder or DefaultServerCallContextBuilder() ) + self._shutdown_grace_period = shutdown_grace_period @rest_error_handler async def _handle_request( @@ -97,7 +99,8 @@ async def event_generator( yield json_utils.dumps(item) return EventSourceResponse( - event_generator(method(request, call_context)) + event_generator(method(request, call_context)), + shutdown_grace_period=self._shutdown_grace_period, ) def routes(self) -> dict[tuple[str, str], Callable[[Request], Any]]: diff --git a/src/a2a/server/routes/jsonrpc_dispatcher.py b/src/a2a/server/routes/jsonrpc_dispatcher.py index 228a57579..9cfbf95a2 100644 --- a/src/a2a/server/routes/jsonrpc_dispatcher.py +++ b/src/a2a/server/routes/jsonrpc_dispatcher.py @@ -5,7 +5,7 @@ import traceback from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from google.protobuf.json_format import MessageToDict, ParseDict from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response @@ -130,6 +130,7 @@ def __init__( request_handler: RequestHandler, context_builder: ServerCallContextBuilder | None = None, enable_v0_3_compat: bool = False, + shutdown_grace_period: float = 0, ) -> None: """Initializes the JsonRpcDispatcher. @@ -140,6 +141,8 @@ def __init__( ServerCallContext passed to the request_handler. If None the DefaultServerCallContextBuilder is used. enable_v0_3_compat: Whether to enable v0.3 backward compatibility on the same endpoint. + shutdown_grace_period: Seconds to allow active SSE streams to + finish before force-cancellation during shutdown. """ if not _package_starlette_installed: raise ImportError( @@ -153,12 +156,14 @@ def __init__( context_builder or DefaultServerCallContextBuilder() ) self.enable_v0_3_compat = enable_v0_3_compat + self._shutdown_grace_period = shutdown_grace_period self._v03_adapter: JSONRPC03Adapter | None = None if self.enable_v0_3_compat: self._v03_adapter = JSONRPC03Adapter( http_handler=request_handler, context_builder=self._context_builder, + shutdown_grace_period=shutdown_grace_period, ) def _generate_error_response( @@ -594,7 +599,11 @@ async def event_generator( 'data': json_utils.dumps(error_response), } - return EventSourceResponse(event_generator(handler_result)) # ty:ignore[invalid-argument-type] + stream = cast('AsyncGenerator[dict[str, Any]]', handler_result) + return EventSourceResponse( + event_generator(stream), + shutdown_grace_period=self._shutdown_grace_period, + ) # handler_result is a dict (JSON-RPC response) return JSONResponse(handler_result) diff --git a/src/a2a/server/routes/jsonrpc_routes.py b/src/a2a/server/routes/jsonrpc_routes.py index a94d513ae..648e9dfa2 100644 --- a/src/a2a/server/routes/jsonrpc_routes.py +++ b/src/a2a/server/routes/jsonrpc_routes.py @@ -30,6 +30,7 @@ def create_jsonrpc_routes( rpc_url: str, context_builder: ServerCallContextBuilder | None = None, enable_v0_3_compat: bool = False, + shutdown_grace_period: float = 0, ) -> list['Route']: """Creates the Starlette Route for the A2A protocol JSON-RPC endpoint. @@ -45,6 +46,10 @@ def create_jsonrpc_routes( ServerCallContext passed to the request_handler. If None the DefaultServerCallContextBuilder is used. enable_v0_3_compat: Whether to enable v0.3 backward compatibility on the same endpoint. + shutdown_grace_period: Seconds to allow active SSE streams to finish + before force-cancellation during shutdown. This value should be less + than the ASGI server's graceful shutdown timeout. Defaults to 0, + matching the ``sse-starlette`` default behavior. """ if not _package_starlette_installed: raise ImportError( @@ -57,6 +62,7 @@ def create_jsonrpc_routes( request_handler=request_handler, context_builder=context_builder, enable_v0_3_compat=enable_v0_3_compat, + shutdown_grace_period=shutdown_grace_period, ) return [ diff --git a/src/a2a/server/routes/rest_dispatcher.py b/src/a2a/server/routes/rest_dispatcher.py index f2d81271a..afcf2e515 100644 --- a/src/a2a/server/routes/rest_dispatcher.py +++ b/src/a2a/server/routes/rest_dispatcher.py @@ -72,6 +72,7 @@ def __init__( self, request_handler: RequestHandler, context_builder: ServerCallContextBuilder | None = None, + shutdown_grace_period: float = 0, ) -> None: """Initializes the RestDispatcher. @@ -80,6 +81,8 @@ def __init__( context_builder: The ServerCallContextBuilder used to construct the ServerCallContext passed to the request_handler. If None the DefaultServerCallContextBuilder is used. + shutdown_grace_period: Seconds to allow active SSE streams to + finish before force-cancellation during shutdown. """ if not _package_starlette_installed: raise ImportError( @@ -92,6 +95,7 @@ def __init__( context_builder or DefaultServerCallContextBuilder() ) self.request_handler = request_handler + self._shutdown_grace_period = shutdown_grace_period def _build_call_context(self, request: Request) -> ServerCallContext: call_context = self._context_builder.build(request) @@ -137,7 +141,10 @@ async def _handle_streaming( try: first_item = await anext(stream) except StopAsyncIteration: - return EventSourceResponse(iter([])) + return EventSourceResponse( + iter([]), + shutdown_grace_period=self._shutdown_grace_period, + ) async def event_generator() -> AsyncIterator[ServerSentEvent]: yield ServerSentEvent(data=json_utils.dumps(first_item)) @@ -151,7 +158,10 @@ async def event_generator() -> AsyncIterator[ServerSentEvent]: event='error', ) - return EventSourceResponse(event_generator()) + return EventSourceResponse( + event_generator(), + shutdown_grace_period=self._shutdown_grace_period, + ) @rest_error_handler async def on_message_send(self, request: Request) -> Response: diff --git a/src/a2a/server/routes/rest_routes.py b/src/a2a/server/routes/rest_routes.py index 2ba8cecfc..b74141189 100644 --- a/src/a2a/server/routes/rest_routes.py +++ b/src/a2a/server/routes/rest_routes.py @@ -32,6 +32,7 @@ def create_rest_routes( context_builder: ServerCallContextBuilder | None = None, enable_v0_3_compat: bool = False, path_prefix: str = '', + shutdown_grace_period: float = 0, ) -> list['BaseRoute']: """Creates the Starlette Routes for the A2A protocol REST endpoint. @@ -44,6 +45,10 @@ def create_rest_routes( enable_v0_3_compat: If True, mounts backward-compatible v0.3 protocol endpoints using REST03Adapter. path_prefix: The URL prefix for the REST endpoints. + shutdown_grace_period: Seconds to allow active SSE streams to finish + before force-cancellation during shutdown. This value should be less + than the ASGI server's graceful shutdown timeout. Defaults to 0, + matching the ``sse-starlette`` default behavior. """ if not _package_starlette_installed: raise ImportError( @@ -55,6 +60,7 @@ def create_rest_routes( dispatcher = RestDispatcher( request_handler=request_handler, context_builder=context_builder, + shutdown_grace_period=shutdown_grace_period, ) routes: list[BaseRoute] = [] @@ -62,6 +68,7 @@ def create_rest_routes( v03_adapter = REST03Adapter( http_handler=request_handler, context_builder=context_builder, + shutdown_grace_period=shutdown_grace_period, ) v03_routes = v03_adapter.routes() for (path, method), endpoint in v03_routes.items(): diff --git a/tests/compat/v0_3/test_jsonrpc_app_compat.py b/tests/compat/v0_3/test_jsonrpc_app_compat.py index 4da4091c5..5679b1934 100644 --- a/tests/compat/v0_3/test_jsonrpc_app_compat.py +++ b/tests/compat/v0_3/test_jsonrpc_app_compat.py @@ -1,9 +1,13 @@ import logging -from unittest.mock import AsyncMock +from collections.abc import AsyncIterator +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from a2a.compat.v0_3 import jsonrpc_adapter +from a2a.compat.v0_3.jsonrpc_adapter import JSONRPC03Adapter +from a2a.server.context import ServerCallContext from a2a.server.request_handlers.request_handler import RequestHandler from a2a.server.routes import create_jsonrpc_routes from a2a.types.a2a_pb2 import ( @@ -23,6 +27,33 @@ logger = logging.getLogger(__name__) +@pytest.mark.asyncio +async def test_shutdown_grace_period_is_passed_to_event_source_response( + mock_handler: AsyncMock, +) -> None: + async def stream_generator() -> AsyncIterator[MagicMock]: + yield MagicMock() + + adapter = JSONRPC03Adapter( + http_handler=mock_handler, + shutdown_grace_period=30.0, + ) + adapter.handler.on_message_send_stream = MagicMock( + return_value=stream_generator() + ) + request_obj = MagicMock(method='message/stream') + + with patch.object(jsonrpc_adapter, 'EventSourceResponse') as response_class: + await adapter._process_streaming_request( + request_id='1', + request_obj=request_obj, + context=ServerCallContext(), + ) + + response_class.assert_called_once() + assert response_class.call_args.kwargs['shutdown_grace_period'] == 30.0 + + @pytest.fixture def mock_handler(): handler = AsyncMock(spec=RequestHandler) diff --git a/tests/compat/v0_3/test_rest_routes_compat.py b/tests/compat/v0_3/test_rest_routes_compat.py index fd6583432..5e16df65d 100644 --- a/tests/compat/v0_3/test_rest_routes_compat.py +++ b/tests/compat/v0_3/test_rest_routes_compat.py @@ -1,11 +1,11 @@ import logging from collections.abc import AsyncIterator -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from a2a.compat.v0_3 import a2a_v0_3_pb2 +from a2a.compat.v0_3 import a2a_v0_3_pb2, rest_adapter from a2a.compat.v0_3.rest_adapter import REST03Adapter from a2a.server.request_handlers.request_handler import RequestHandler from a2a.server.routes import create_agent_card_routes @@ -30,6 +30,33 @@ logger = logging.getLogger(__name__) +@pytest.mark.anyio +async def test_shutdown_grace_period_is_passed_to_event_source_response( + request_handler: RequestHandler, +) -> None: + async def stream( + request: Request, context: object + ) -> AsyncIterator[dict[str, str]]: + yield {'result': 'value'} + + adapter = REST03Adapter( + http_handler=request_handler, + shutdown_grace_period=30.0, + ) + mock_req = MagicMock(spec=Request) + mock_req.body = AsyncMock(return_value=b'{}') + mock_req.headers = Headers({'a2a-version': '0.3'}) + mock_req.user = MagicMock(is_authenticated=False) + mock_req.auth = None + mock_req.scope = {} + + with patch.object(rest_adapter, 'EventSourceResponse') as response_class: + await adapter._handle_streaming_request(stream, mock_req) + + response_class.assert_called_once() + assert response_class.call_args.kwargs['shutdown_grace_period'] == 30.0 + + @pytest.fixture async def agent_card() -> AgentCard: mock_agent_card = MagicMock(spec=AgentCard) diff --git a/tests/server/routes/test_jsonrpc_dispatcher.py b/tests/server/routes/test_jsonrpc_dispatcher.py index ed55c476c..4bd649ba2 100644 --- a/tests/server/routes/test_jsonrpc_dispatcher.py +++ b/tests/server/routes/test_jsonrpc_dispatcher.py @@ -1,5 +1,6 @@ import asyncio +from collections.abc import AsyncGenerator from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -126,6 +127,25 @@ def test_create_dispatcher_with_missing_deps_raises_importerror( JsonRpcDispatcher(**mock_app_params) +class TestJsonRpcDispatcherStreamingResponse: + def test_shutdown_grace_period_is_passed_to_event_source_response( + self, mock_handler + ) -> None: + async def stream_generator() -> AsyncGenerator[dict[str, Any]]: + yield {'result': {}} + + dispatcher = JsonRpcDispatcher( + request_handler=mock_handler, + shutdown_grace_period=30.0, + ) + + response = dispatcher._create_response( + ServerCallContext(), stream_generator() + ) + + assert getattr(response, '_shutdown_grace_period') == 30.0 + + class TestJsonRpcDispatcherExtensions: def test_request_with_single_extension(self, client, mock_handler): headers = {HTTP_EXTENSION_HEADER: 'foo'} @@ -196,6 +216,22 @@ def test_no_tenant_extraction(self, client, mock_handler): class TestJsonRpcDispatcherV03Compat: + def test_shutdown_grace_period_is_forwarded_to_adapter( + self, mock_handler + ) -> None: + with patch.object( + jsonrpc_dispatcher, 'JSONRPC03Adapter' + ) as adapter_class: + JsonRpcDispatcher( + request_handler=mock_handler, + enable_v0_3_compat=True, + shutdown_grace_period=30.0, + ) + + adapter_class.assert_called_once() + assert adapter_class.call_args.kwargs['http_handler'] is mock_handler + assert adapter_class.call_args.kwargs['shutdown_grace_period'] == 30.0 + def test_v0_3_compat_flag_routes_to_adapter(self, mock_handler): mock_agent_card = MagicMock(spec=AgentCard) mock_agent_card.url = 'http://mockurl.com' diff --git a/tests/server/routes/test_jsonrpc_routes.py b/tests/server/routes/test_jsonrpc_routes.py index a9e166f69..4c25873a9 100644 --- a/tests/server/routes/test_jsonrpc_routes.py +++ b/tests/server/routes/test_jsonrpc_routes.py @@ -1,4 +1,4 @@ -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch import pytest @@ -56,3 +56,34 @@ def test_jsonrpc_custom_url(agent_card, mock_handler): assert 'error' in resp_json # Method not found error from dispatcher assert resp_json['error']['code'] == -32601 + + +def test_shutdown_grace_period_is_forwarded_to_dispatcher(mock_handler) -> None: + """Tests that the route factory configures SSE cooperative shutdown.""" + with patch( + 'a2a.server.routes.jsonrpc_routes.JsonRpcDispatcher' + ) as dispatcher_class: + create_jsonrpc_routes( + request_handler=mock_handler, + rpc_url='/a2a/jsonrpc', + ) + dispatcher_class.assert_called_once_with( + request_handler=mock_handler, + context_builder=None, + enable_v0_3_compat=False, + shutdown_grace_period=0, + ) + + dispatcher_class.reset_mock() + create_jsonrpc_routes( + request_handler=mock_handler, + rpc_url='/a2a/jsonrpc', + shutdown_grace_period=30.0, + ) + + dispatcher_class.assert_called_once_with( + request_handler=mock_handler, + context_builder=None, + enable_v0_3_compat=False, + shutdown_grace_period=30.0, + ) diff --git a/tests/server/routes/test_rest_dispatcher.py b/tests/server/routes/test_rest_dispatcher.py index a062a10c1..e0d72c4f2 100644 --- a/tests/server/routes/test_rest_dispatcher.py +++ b/tests/server/routes/test_rest_dispatcher.py @@ -276,6 +276,32 @@ async def test_handle_authenticated_agent_card( @pytest.mark.asyncio class TestRestDispatcherStreaming: + async def test_shutdown_grace_period_is_passed_to_streaming_responses( + self, mock_handler + ) -> None: + async def stream( + context: ServerCallContext, + ) -> AsyncIterator[dict[str, str]]: + yield {'result': 'value'} + + async def empty_stream( + context: ServerCallContext, + ) -> AsyncIterator[dict[str, str]]: + for item in []: + yield item + + dispatcher = RestDispatcher( + request_handler=mock_handler, + shutdown_grace_period=30.0, + ) + req = make_mock_request(method='POST') + + response = await dispatcher._handle_streaming(req, stream) + empty_response = await dispatcher._handle_streaming(req, empty_stream) + + assert getattr(response, '_shutdown_grace_period') == 30.0 + assert getattr(empty_response, '_shutdown_grace_period') == 30.0 + async def test_on_message_send_stream_success( self, rest_dispatcher_instance ): diff --git a/tests/server/routes/test_rest_routes.py b/tests/server/routes/test_rest_routes.py index 7d37762c4..1839041cd 100644 --- a/tests/server/routes/test_rest_routes.py +++ b/tests/server/routes/test_rest_routes.py @@ -1,4 +1,4 @@ -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch import pytest @@ -92,3 +92,47 @@ def test_rest_list_tasks(agent_card, mock_handler): response = client.get('/tasks', headers={'A2A-Version': '1.0'}) assert response.status_code == 200 assert mock_handler.on_list_tasks.called + + +def test_shutdown_grace_period_is_forwarded_to_dispatcher(mock_handler) -> None: + """Tests that the route factory configures SSE cooperative shutdown.""" + with patch( + 'a2a.server.routes.rest_routes.RestDispatcher' + ) as dispatcher_class: + create_rest_routes(request_handler=mock_handler) + dispatcher_class.assert_called_once_with( + request_handler=mock_handler, + context_builder=None, + shutdown_grace_period=0, + ) + + dispatcher_class.reset_mock() + create_rest_routes( + request_handler=mock_handler, + shutdown_grace_period=30.0, + ) + + dispatcher_class.assert_called_once_with( + request_handler=mock_handler, + context_builder=None, + shutdown_grace_period=30.0, + ) + + +def test_shutdown_grace_period_is_forwarded_to_v03_adapter( + mock_handler, +) -> None: + """Tests that v0.3 compatibility routes share the shutdown grace period.""" + with patch('a2a.server.routes.rest_routes.REST03Adapter') as adapter_class: + adapter_class.return_value.routes.return_value = {} + create_rest_routes( + request_handler=mock_handler, + enable_v0_3_compat=True, + shutdown_grace_period=30.0, + ) + + adapter_class.assert_called_once_with( + http_handler=mock_handler, + context_builder=None, + shutdown_grace_period=30.0, + ) diff --git a/uv.lock b/uv.lock index 3c6f3ccde..c27d94965 100644 --- a/uv.lock +++ b/uv.lock @@ -143,9 +143,9 @@ requires-dist = [ { name = "sqlalchemy", extras = ["asyncio", "postgresql-asyncpg"], marker = "extra == 'all'", specifier = ">=2.0.0" }, { name = "sqlalchemy", extras = ["asyncio", "postgresql-asyncpg"], marker = "extra == 'postgresql'", specifier = ">=2.0.0" }, { name = "sqlalchemy", extras = ["asyncio", "postgresql-asyncpg"], marker = "extra == 'sql'", specifier = ">=2.0.0" }, - { name = "sse-starlette", marker = "extra == 'all'" }, - { name = "sse-starlette", marker = "extra == 'fastapi'" }, - { name = "sse-starlette", marker = "extra == 'http-server'" }, + { name = "sse-starlette", marker = "extra == 'all'", specifier = ">=3.3.0" }, + { name = "sse-starlette", marker = "extra == 'fastapi'", specifier = ">=3.3.0" }, + { name = "sse-starlette", marker = "extra == 'http-server'", specifier = ">=3.3.0" }, { name = "starlette", marker = "extra == 'all'" }, { name = "starlette", marker = "extra == 'fastapi'" }, { name = "starlette", marker = "extra == 'http-server'" },