diff --git a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py index 0f9cbfe39ee6..4701681f9c41 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py +++ b/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, diff --git a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py index 999d5d641105..69ec63dbc08a 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py +++ b/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. diff --git a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/streaming/_sse.py b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/streaming/_sse.py index 9152500afa10..8dad60a5e661 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/streaming/_sse.py +++ b/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) diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_keep_alive.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_keep_alive.py index f9dbf63a91d0..7b423ba1ef37 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_keep_alive.py +++ b/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: diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_sse_writer.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_sse_writer.py index 259063f82960..10c7d3dfe366 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_sse_writer.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_sse_writer.py @@ -4,7 +4,16 @@ from __future__ import annotations +import asyncio +import contextvars +from typing import AsyncIterator, List + +import pytest + from azure.ai.agentserver.responses.streaming import _sse +from azure.ai.agentserver.responses.streaming._sse import encode_keep_alive_comment, with_keep_alive + +_KA = encode_keep_alive_comment() class _FakeEvent: @@ -59,3 +68,99 @@ def _extract_sequence_number(encoded: str) -> int: assert seq_first == 0, f"first sequence_number must be 0 for a new stream, got {seq_first}" assert seq_second == 1, f"second sequence_number must be 1 for a new stream, got {seq_second}" + + +async def _collect(stream: AsyncIterator[str]) -> List[str]: + return [item async for item in stream] + + +async def test_with_keep_alive__passthrough_when_disabled() -> None: + """A None or 0 interval yields the source unchanged with no keep-alive frames.""" + + async def source() -> AsyncIterator[str]: + yield "a" + yield "b" + + assert await _collect(with_keep_alive(source(), None)) == ["a", "b"] + assert await _collect(with_keep_alive(source(), 0)) == ["a", "b"] + + +async def test_with_keep_alive__emits_heartbeat_on_idle_gap() -> None: + """A keep-alive frame is emitted when the source stays idle past the interval.""" + + async def source() -> AsyncIterator[str]: + yield "a" + await asyncio.sleep(0.25) + yield "b" + + out = await _collect(with_keep_alive(source(), 0.05)) + assert _KA in out + assert [item for item in out if item != _KA] == ["a", "b"] + assert out.index("a") < out.index("b") + + +async def test_with_keep_alive__preserves_contextvars_across_yields() -> None: + """A ContextVar the source sets once stays set across all of its yields, because the + source is advanced by a single task.""" + + cvar: contextvars.ContextVar[str] = contextvars.ContextVar("cvar", default="UNSET") + seen: List[str] = [] + + async def source() -> AsyncIterator[str]: + cvar.set("SET") + seen.append(cvar.get()) + yield "a" + seen.append(cvar.get()) + yield "b" + seen.append(cvar.get()) + + out = await _collect(with_keep_alive(source(), 5)) + assert out == ["a", "b"] + assert seen == ["SET", "SET", "SET"] + + +async def test_with_keep_alive__runs_source_finally_on_early_close() -> None: + """Closing the wrapper after one item runs the source's finally before returning.""" + + finalized = asyncio.Event() + + async def source() -> AsyncIterator[str]: + try: + yield "a" + await asyncio.Event().wait() + yield "b" + finally: + finalized.set() + + stream = with_keep_alive(source(), 5) + assert await stream.__anext__() == "a" + await stream.aclose() + await asyncio.wait_for(finalized.wait(), timeout=1.0) + + +async def test_with_keep_alive__emits_heartbeat_while_source_idle_from_start() -> None: + """A source idle from the start still emits keep-alive frames, matching a reconnect + GET subscribed to an in-flight run that has not emitted since the cursor.""" + + async def idle_source() -> AsyncIterator[str]: + await asyncio.Event().wait() + yield "never" # pragma: no cover + + stream = with_keep_alive(idle_source(), 0.05) + assert await asyncio.wait_for(stream.__anext__(), timeout=1.0) == _KA + await stream.aclose() + + +async def test_with_keep_alive__propagates_source_error() -> None: + """An error raised by the source surfaces to the consumer after buffered items.""" + + async def source() -> AsyncIterator[str]: + yield "a" + raise RuntimeError("boom") + + collected: List[str] = [] + with pytest.raises(RuntimeError, match="boom"): + async for item in with_keep_alive(source(), 5): + collected.append(item) + + assert collected == ["a"]