feat: allow configuring SSE shutdown grace periods

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"]

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,
+        )

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]]:

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)

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 [

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:

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,13 +60,15 @@ def create_rest_routes(
     dispatcher = RestDispatcher(
         request_handler=request_handler,
         context_builder=context_builder,
+        shutdown_grace_period=shutdown_grace_period,
     )
 
     routes: list[BaseRoute] = []
     if enable_v0_3_compat:
         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():

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)

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)

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'