Move SSE keep-alive to a transport-layer wrapper

sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py

@@ -48,7 +48,7 @@
 from ..store._base import ResponseProviderProtocol, ResponseStreamProviderProtocol
 from ..store._foundry_errors import FoundryApiError, FoundryBadRequestError, FoundryResourceNotFoundError
 from ..streaming._helpers import _encode_sse
-from ..streaming._sse import encode_sse_any_event
+from ..streaming._sse import encode_sse_any_event, with_keep_alive
 from ..streaming._state_machine import _normalize_lifecycle_events
 from ._execution_context import _ExecutionContext
 from ._observability import (
@@ -696,7 +696,10 @@ async def _iter_with_context():  # type: ignore[return]
                             disconnect_task.cancel()
 
                 sse_response = StreamingResponse(
-                    _iter_with_context(),
+                    with_keep_alive(
+                        _iter_with_context(),
+                        self._runtime_options.sse_keep_alive_interval_seconds,
+                    ),
                     media_type="text/event-stream",
                     headers={**self._sse_headers, **self._session_headers(agent_session_id)},
                 )
@@ -1046,7 +1049,11 @@ async def _stream_from_subject():
             async for event in record.subject.subscribe(cursor=_cursor):  # type: ignore[union-attr]
                 yield encode_sse_any_event(event)
 
-        return StreamingResponse(_stream_from_subject(), media_type="text/event-stream", headers=merged_headers)
+        return StreamingResponse(
+            with_keep_alive(_stream_from_subject(), self._runtime_options.sse_keep_alive_interval_seconds),
+            media_type="text/event-stream",
+            headers=merged_headers,
+        )
 
     async def _try_replay_persisted_stream(
         self,

sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py

@@ -41,7 +41,7 @@
     _extract_response_snapshot_from_events,
 )
 from ..streaming._internals import construct_event_model
-from ..streaming._sse import encode_keep_alive_comment, encode_sse_any_event, new_stream_counter
+from ..streaming._sse import encode_sse_any_event, new_stream_counter
 from ..streaming._state_machine import EventStreamValidator
 from ._event_subject import _ResponseEventSubject
 from ._execution_context import _ExecutionContext
@@ -1461,7 +1461,6 @@ def run_stream(self, ctx: _ExecutionContext) -> AsyncIterator[str]:
         - Empty handler (fallback synthesised events).
         - Mid-stream handler errors (``response.failed`` SSE event, S-035).
         - Cancellation terminal events.
-        - Optional SSE keep-alive comments.
 
         :param ctx: Current execution context.
         :type ctx: _ExecutionContext
@@ -1476,7 +1475,8 @@ async def _live_stream(self, ctx: _ExecutionContext) -> AsyncIterator[str]:
         Delegates all event processing (first-event handling, normalisation,
         bg record registration, S-035 / S-015 / B11 terminal events) to
         :meth:`_process_handler_events`.  This method only encodes each event
-        dict to SSE and handles keep-alive comment injection.
+        dict to SSE. Keep-alive comment injection is handled at the transport
+        layer (:func:`streaming._sse.with_keep_alive`) in the endpoint handler.
 
         :param ctx: Current execution context.
         :type ctx: _ExecutionContext
@@ -1499,153 +1499,77 @@ async def _live_stream(self, ctx: _ExecutionContext) -> AsyncIterator[str]:
         async def _finalize() -> None:
             await self._finalize_stream(ctx, state)
 
-        # --- Fast path: no keep-alive ---
-        if not self._runtime_options.sse_keep_alive_enabled:
-            if not (ctx.background and ctx.store):
-                # Simple fast path for non-background streaming.
-                _stream_completed = False
-                try:
-                    async for event in self._process_handler_events(ctx, state, handler_iterator):
-                        yield encode_sse_any_event(event)
-                    _stream_completed = True
-                    # Persist-then-yield: resolve the buffered terminal event
-                    if state.pending_terminal is not None:
-                        record = state.bg_record or _make_ephemeral_record(ctx, state)
-                        resolved = await self._persist_and_resolve_terminal(ctx, state, record)
-                        yield encode_sse_any_event(resolved)
-                finally:
-                    # B17: If the stream did not complete naturally (e.g. client
-                    # disconnect → CancelledError), mark it as interrupted so
-                    # _finalize_stream skips persistence for non-bg streams.
-                    if not _stream_completed:
-                        state.stream_interrupted = True
-                    await _finalize()
-                return
-
-            # Background+stream without keep-alive: run the handler as an independent
-            # asyncio.Task so that finalization (including subject.complete()) is
-            # guaranteed to run even when the original SSE connection is dropped before
-            # all events are delivered.  Without this, _live_stream can be abandoned
-            # mid-iteration by Starlette (the async-generator finalizer may not fire
-            # promptly), leaving GET-replay subscribers blocked on await q.get() forever.
-            _SENTINEL_BG = object()
-            bg_queue: asyncio.Queue[object] = asyncio.Queue()
-
-            async def _bg_producer_inner() -> None:
-                try:
-                    async for event in self._process_handler_events(ctx, state, handler_iterator):
-                        await bg_queue.put(encode_sse_any_event(event))
-                    # Persist-then-yield: resolve the buffered terminal event
-                    if state.pending_terminal is not None:
-                        record = state.bg_record or _make_ephemeral_record(ctx, state)
-                        resolved = await self._persist_and_resolve_terminal(ctx, state, record)
-                        await bg_queue.put(encode_sse_any_event(resolved))
-                except Exception as exc:  # pylint: disable=broad-exception-caught
-                    logger.error(
-                        "Background stream producer failed (response_id=%s)",
-                        ctx.response_id,
-                        exc_info=exc,
-                    )
-                    state.captured_error = exc
-                finally:
-                    # Always finalize (includes subject.complete()) — this runs even if
-                    # the original POST SSE connection was dropped and _live_stream is
-                    # never properly closed by Starlette.
-                    await _finalize()
-                    await bg_queue.put(_SENTINEL_BG)
-
-            async def _bg_producer() -> None:
-                try:
-                    # FR-013: Shield the inner producer via asyncio.shield so
-                    # that Starlette's anyio cancel-scope cancellation (triggered
-                    # by client disconnect) does NOT propagate into the handler.
-                    # asyncio.shield() creates a new inner Task whose cancellation
-                    # is independent of the outer task.
-                    await asyncio.shield(_bg_producer_inner())
-                except asyncio.CancelledError:
-                    pass  # outer task cancelled by scope; inner task continues
-
-            bg_task = asyncio.create_task(_bg_producer())
+        if not (ctx.background and ctx.store):
+            # Simple path for non-background (or non-store) streaming.
+            _stream_completed = False
             try:
-                while True:
-                    item = await bg_queue.get()
-                    if item is _SENTINEL_BG:
-                        break
-                    yield item  # type: ignore[misc]
-            except Exception:  # pylint: disable=broad-exception-caught
-                pass  # SSE connection dropped; bg_task continues independently
+                async for event in self._process_handler_events(ctx, state, handler_iterator):
+                    yield encode_sse_any_event(event)
+                _stream_completed = True
+                # Persist-then-yield: resolve the buffered terminal event
+                if state.pending_terminal is not None:
+                    record = state.bg_record or _make_ephemeral_record(ctx, state)
+                    resolved = await self._persist_and_resolve_terminal(ctx, state, record)
+                    yield encode_sse_any_event(resolved)
             finally:
-                # Wait for the handler task so _finalize() has run before we exit.
-                # Do NOT cancel it — background+stream must reach a terminal state
-                # regardless of client connectivity.
-                if not bg_task.done():
-                    try:
-                        await bg_task
-                    except Exception:  # pylint: disable=broad-exception-caught
-                        pass
+                # B17: mark an interrupted (not naturally completed) stream so
+                # _finalize_stream skips persistence for non-bg streams.
+                if not _stream_completed:
+                    state.stream_interrupted = True
+                await _finalize()
             return
 
-        # --- Keep-alive path: merge handler events with periodic keep-alive comments ---
-        # via a shared asyncio.Queue so comments are sent even while the handler is idle.
-        _SENTINEL = object()
-        merge_queue: asyncio.Queue[str | object] = asyncio.Queue()
+        # Background+stream (store): run the handler as an independent task so finalization
+        # (including subject.complete()) runs even when the SSE connection is dropped before
+        # all events are delivered.
+        _SENTINEL_BG = object()
+        bg_queue: asyncio.Queue[object] = asyncio.Queue()
 
-        async def _handler_producer() -> None:
+        async def _bg_producer_inner() -> None:
             try:
                 async for event in self._process_handler_events(ctx, state, handler_iterator):
-                    await merge_queue.put(encode_sse_any_event(event))
+                    await bg_queue.put(encode_sse_any_event(event))
                 # Persist-then-yield: resolve the buffered terminal event
                 if state.pending_terminal is not None:
                     record = state.bg_record or _make_ephemeral_record(ctx, state)
                     resolved = await self._persist_and_resolve_terminal(ctx, state, record)
-                    await merge_queue.put(encode_sse_any_event(resolved))
+                    await bg_queue.put(encode_sse_any_event(resolved))
+            except Exception as exc:  # pylint: disable=broad-exception-caught
+                logger.error(
+                    "Background stream producer failed (response_id=%s)",
+                    ctx.response_id,
+                    exc_info=exc,
+                )
+                state.captured_error = exc
             finally:
-                await merge_queue.put(_SENTINEL)
+                await _finalize()
+                await bg_queue.put(_SENTINEL_BG)
 
-        async def _keep_alive_producer(interval: int) -> None:
+        async def _bg_producer() -> None:
             try:
-                while True:
-                    await asyncio.sleep(interval)
-                    await merge_queue.put(encode_keep_alive_comment())
+                # FR-013: Shield the producer so client-disconnect cancellation does not
+                # propagate into the handler.
+                await asyncio.shield(_bg_producer_inner())
             except asyncio.CancelledError:
-                return
-
-        handler_task = asyncio.create_task(_handler_producer())
-        keep_alive_task = asyncio.create_task(
-            _keep_alive_producer(self._runtime_options.sse_keep_alive_interval_seconds)  # type: ignore[arg-type]
-        )
+                pass  # outer task cancelled by scope; inner task continues
 
-        _ka_stream_completed = False
+        bg_task = asyncio.create_task(_bg_producer())
         try:
             while True:
-                item = await merge_queue.get()
-                if item is _SENTINEL:
-                    _ka_stream_completed = True
+                item = await bg_queue.get()
+                if item is _SENTINEL_BG:
                     break
                 yield item  # type: ignore[misc]
-        except Exception as exc:  # pylint: disable=broad-exception-caught
-            logger.error(
-                "Stream consumer failed (response_id=%s)",
-                ctx.response_id,
-                exc_info=exc,
-            )
-            state.captured_error = exc
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass  # SSE connection dropped; bg_task continues independently
         finally:
-            if not _ka_stream_completed:
-                state.stream_interrupted = True
-            keep_alive_task.cancel()
-            try:
-                await keep_alive_task
-            except asyncio.CancelledError:
-                pass
-            # Ensure the handler task has finished before finalising
-            if not handler_task.done():
-                handler_task.cancel()
+            # Await the producer (do not cancel it) so it reaches a terminal state and
+            # _finalize() has run before returning.
+            if not bg_task.done():
                 try:
-                    await handler_task
-                except asyncio.CancelledError:
+                    await bg_task
+                except Exception:  # pylint: disable=broad-exception-caught
                     pass
-            await _finalize()
 
     async def run_sync(self, ctx: _ExecutionContext) -> dict[str, Any]:
         """Execute a synchronous (non-stream, non-background) create-response request.

sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/streaming/_sse.py

@@ -4,11 +4,12 @@
 
 from __future__ import annotations
 
+import asyncio  # pylint: disable=do-not-import-asyncio
 import itertools
 import json
 from contextvars import ContextVar
 from datetime import date, datetime, time, timedelta
-from typing import Any, Mapping
+from typing import Any, AsyncIterator, Mapping
 
 from ..models._generated import ResponseStreamEvent
 
@@ -177,3 +178,69 @@ def encode_keep_alive_comment(comment: str = "keep-alive") -> str:
     :rtype: str
     """
     return f": {comment}\n\n"
+
+
+async def with_keep_alive(
+    source: AsyncIterator[str],
+    interval_seconds: float | None,
+) -> AsyncIterator[str]:
+    """Interleave SSE keep-alive comment frames into ``source`` during idle gaps.
+
+    Yields the source unchanged when ``interval_seconds`` is falsy. Otherwise a keep-alive
+    comment frame is emitted whenever no upstream item arrives within ``interval_seconds``;
+    real items are never dropped or reordered. The source is advanced by a single task, so
+    any ``contextvars`` it sets (request context, SSE sequence counter) persist across the
+    whole stream.
+
+    :param source: The upstream async iterator of SSE-encoded strings.
+    :type source: AsyncIterator[str]
+    :param interval_seconds: Idle interval before emitting a keep-alive frame, or ``None`` to disable.
+    :type interval_seconds: float | None
+    :returns: An async iterator that yields the source's items plus periodic keep-alive frames.
+    :rtype: AsyncIterator[str]
+    """
+    if not interval_seconds:
+        async for item in source:
+            yield item
+        return
+
+    queue: asyncio.Queue[Any] = asyncio.Queue()
+    sentinel = object()
+    pump_error: BaseException | None = None
+
+    async def _pump() -> None:
+        nonlocal pump_error
+        try:
+            async for item in source:
+                queue.put_nowait(item)
+        except Exception as exc:  # pylint: disable=broad-exception-caught
+            pump_error = exc
+        finally:
+            queue.put_nowait(sentinel)
+
+    pump_task = asyncio.ensure_future(_pump())
+    get_task: "asyncio.Future[Any] | None" = None
+    try:
+        while True:
+            if get_task is None:
+                get_task = asyncio.ensure_future(queue.get())
+            done, _ = await asyncio.wait({get_task}, timeout=interval_seconds)
+            if get_task not in done:
+                # No item within the interval; emit a heartbeat and keep the pending
+                # get_task so the next item is not dropped.
+                yield encode_keep_alive_comment()
+                continue
+            item = get_task.result()
+            get_task = None
+            if item is sentinel:
+                break
+            yield item
+        if pump_error is not None:
+            raise pump_error
+    finally:
+        # Stop the pump and any pending get, and await them so the source's finally
+        # (finalize, request-context reset) runs before returning.
+        pending = [task for task in (pump_task, get_task) if task is not None]
+        for task in pending:
+            task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)

sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_keep_alive.py

@@ -178,9 +178,12 @@ def test_keep_alive__does_not_disrupt_event_stream_integrity() -> None:
 
     event_types = [e["type"] for e in events]
     assert "response.created" in event_types
-    # Sequence numbers should still be monotonically increasing
+    # Sequence numbers must be present, unique, and strictly increasing.
     seq_nums = [e["data"].get("sequence_number") for e in events if "sequence_number" in e["data"]]
-    assert seq_nums == sorted(seq_nums)
+    assert len(seq_nums) >= 2
+    assert all(isinstance(n, int) for n in seq_nums)
+    assert len(set(seq_nums)) == len(seq_nums), f"sequence numbers not unique: {seq_nums}"
+    assert all(b > a for a, b in zip(seq_nums, seq_nums[1:])), f"sequence numbers not strictly increasing: {seq_nums}"
 
 
 def test_keep_alive__no_comments_after_stream_ends() -> None: