From 70464f0a30d8dedb6764bbae27d9626bb5746a0d Mon Sep 17 00:00:00 2001 From: Antriksh Jain Date: Thu, 9 Jul 2026 16:15:25 +0530 Subject: [PATCH 1/4] fix(agentserver-responses): move SSE keep-alive to a transport-layer wrapper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep-alive (periodic ': keep-alive' SSE comment frames) lived inside the event-pipeline orchestrator, which branched on the keep-alive interval. When enabled (hosted-platform default), background+stream ran through a merge path whose finally cancelled the handler on client disconnect, killing in-flight background runs so GET ?stream=true reconnect returned 200 with no new events. Keep-alive is a transport concern. Introduce with_keep_alive in streaming/_sse.py — a single-pump-task wrapper that interleaves heartbeats into any AsyncIterator[str] during idle gaps (passthrough when disabled), advancing the source in one task so its contextvars (request context, SSE sequence counter) are preserved. Apply it at the create-POST and reconnect-GET StreamingResponse sites in the endpoint handler. Remove all keep-alive knowledge from the orchestrator: _live_stream collapses from three keep-alive-aware paths to two lifecycle paths (non-bg simple iterator vs. bg+store shielded task). The shielded background path (which survives client disconnect, FR-012/FR-013) is now unconditional instead of gated on keep-alive being off. Reconnect GET streams now also get heartbeats. Add unit tests for with_keep_alive (passthrough, idle heartbeat, no drop/reorder, contextvar preservation, deterministic close, idle-from-start reconnect shape) and strengthen the sequence-number integrity assertion to catch contextvar regressions. --- .../responses/hosting/_endpoint_handler.py | 13 +- .../responses/hosting/_orchestrator.py | 191 ++++++------------ .../agentserver/responses/streaming/_sse.py | 68 ++++++- .../tests/contract/test_keep_alive.py | 9 +- .../tests/unit/test_keep_alive_wrapper.py | 118 +++++++++++ 5 files changed, 267 insertions(+), 132 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_keep_alive_wrapper.py 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..88ad26dd917c 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,90 @@ 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()) + # Keep-alive comment frames are injected at the transport layer + # (streaming._sse.with_keep_alive, applied in the endpoint handler); the + # orchestrator emits only real events. + 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: 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 - # --- 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 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 _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) + # 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 _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 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. + 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() + # 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 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..d18dab52fb1c 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,68 @@ 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: int | None, +) -> AsyncIterator[str]: + """Interleave SSE keep-alive comment frames into ``source`` during idle gaps. + + Passthrough when ``interval_seconds`` is falsy. A keep-alive comment frame is emitted + whenever no upstream item arrives within ``interval_seconds``; real items are never + dropped or reordered. Keep-alive is a transport concern, so this wraps the SSE byte + stream at the point it becomes an HTTP response rather than living in the event pipeline. + + The source is advanced by a single dedicated task, so any ``contextvars`` the source sets + once and relies on across its whole lifetime (the request/platform context and the SSE + sequence-number counter) persist exactly as when the source is iterated directly. + + :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: int | 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() + + async def _pump() -> None: + try: + async for item in source: + queue.put_nowait(item) + 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: + # Idle interval elapsed with no item; emit a heartbeat. The pending + # get_task is preserved (not cancelled) so no item is ever dropped. + yield encode_keep_alive_comment() + continue + item = get_task.result() + get_task = None + if item is sentinel: + break + yield item + finally: + # Deterministic cleanup — including "client disconnects right after a real item" + # (consumer suspended at ``yield item`` with the pump still running): cancelling the + # pump throws CancelledError into the source at its yield point, so the source's own + # finally (finalize / reset_request_context) runs and is awaited here. + 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..04510da6b50d 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,16 @@ 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. (A plain + # `== sorted(seq_nums)` check would pass on a corrupted all-zero stream, so assert + # uniqueness + strict monotonicity — this guards the contextvars-preservation + # requirement of the transport-layer keep-alive wrapper.) seq_nums = [e["data"].get("sequence_number") for e in events if "sequence_number" in e["data"]] + 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 seq_nums == sorted(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_keep_alive_wrapper.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_keep_alive_wrapper.py new file mode 100644 index 000000000000..d0c650877e2b --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_keep_alive_wrapper.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Unit tests for the transport-layer SSE keep-alive wrapper ``with_keep_alive``.""" + +from __future__ import annotations + +import asyncio +import contextvars +from typing import AsyncIterator, List + +import pytest + +from azure.ai.agentserver.responses.streaming._sse import encode_keep_alive_comment, with_keep_alive + +_KA = encode_keep_alive_comment() + + +async def _collect(agen: AsyncIterator[str]) -> List[str]: + return [item async for item in agen] + + +async def test_passthrough_when_interval_disabled() -> None: + """None / 0 interval => pure passthrough, no heartbeats, order preserved.""" + + 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_emits_heartbeat_on_idle_gap() -> None: + """A keep-alive frame is emitted when the source is idle longer than the interval.""" + + async def source() -> AsyncIterator[str]: + yield "a" + await asyncio.sleep(0.25) # idle gap >> interval + yield "b" + + out = await _collect(with_keep_alive(source(), 0.05)) + assert _KA in out + # Real items are preserved in order; heartbeats only appear between them. + assert [item for item in out if item != _KA] == ["a", "b"] + assert out.index("a") < out.index("b") + + +async def test_does_not_drop_or_reorder_items() -> None: + """Items arriving across multiple idle gaps are never dropped or reordered.""" + + async def source() -> AsyncIterator[str]: + for i in range(5): + await asyncio.sleep(0.03) + yield f"item-{i}" + + out = await _collect(with_keep_alive(source(), 0.01)) + assert [item for item in out if item != _KA] == [f"item-{i}" for i in range(5)] + + +async def test_preserves_contextvars_across_yields() -> None: + """The source is advanced by a single task, so a ContextVar it sets once persists + across all its yields. (Regression guard: a new-task-per-anext design loses it after + the first item and corrupts request context + SSE sequence numbers.)""" + + 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)) # fast drain, no heartbeats + assert out == ["a", "b"] + assert seen == ["SET", "SET", "SET"] + + +async def test_closes_source_deterministically_on_early_close() -> None: + """Closing the wrapper after one item (client disconnect mid-run) runs the source's + ``finally`` deterministically rather than deferring it to GC.""" + + finalized = asyncio.Event() + + async def source() -> AsyncIterator[str]: + try: + yield "a" + await asyncio.Event().wait() # block forever after the first item + yield "b" + finally: + finalized.set() + + agen = with_keep_alive(source(), 5) + first = await agen.__anext__() + assert first == "a" + await agen.aclose() + await asyncio.wait_for(finalized.wait(), timeout=1.0) + + +async def test_emits_heartbeat_while_source_idle_from_start() -> None: + """A source idle from the very start still gets keep-alive frames — the reconnect / + replay GET scenario, where the stream subscribes to an in-flight background run that + has not emitted anything since the cursor.""" + + async def idle_source() -> AsyncIterator[str]: + await asyncio.Event().wait() # never yields a real item + yield "never" # pragma: no cover + + agen = with_keep_alive(idle_source(), 0.05) + first = await asyncio.wait_for(agen.__anext__(), timeout=1.0) + assert first == _KA + await agen.aclose() + + +if __name__ == "__main__": # pragma: no cover + pytest.main([__file__, "-v"]) From 7459cb36b40170f595bd1962edb2535ce0dab81b Mon Sep 17 00:00:00 2001 From: Antriksh Jain Date: Thu, 9 Jul 2026 17:53:25 +0530 Subject: [PATCH 2/4] Consolidate keep-alive wrapper tests and trim comments Move the with_keep_alive unit tests into test_sse_writer.py (the unit home for _sse.py) and drop the separate test file. Reword source and test comments to describe current behavior. --- .../responses/hosting/_orchestrator.py | 31 ++--- .../agentserver/responses/streaming/_sse.py | 23 ++-- .../tests/contract/test_keep_alive.py | 5 +- .../tests/unit/test_keep_alive_wrapper.py | 118 ------------------ .../tests/unit/test_sse_writer.py | 100 +++++++++++++++ 5 files changed, 120 insertions(+), 157 deletions(-) delete mode 100644 sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_keep_alive_wrapper.py 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 88ad26dd917c..04bedb9cbc93 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 @@ -1499,9 +1499,8 @@ async def _live_stream(self, ctx: _ExecutionContext) -> AsyncIterator[str]: async def _finalize() -> None: await self._finalize_stream(ctx, state) - # Keep-alive comment frames are injected at the transport layer - # (streaming._sse.with_keep_alive, applied in the endpoint handler); the - # orchestrator emits only real events. + # The orchestrator emits only real events; keep-alive frames are added by the + # transport layer (streaming._sse.with_keep_alive in the endpoint handler). if not (ctx.background and ctx.store): # Simple path for non-background (or non-store) streaming. _stream_completed = False @@ -1515,20 +1514,16 @@ async def _finalize() -> None: 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 + # 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 - # Background+stream (store): 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. + # 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() @@ -1549,18 +1544,13 @@ async def _bg_producer_inner() -> None: ) 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. + # FR-013: Shield the producer so client-disconnect cancellation does not + # propagate into the handler. await asyncio.shield(_bg_producer_inner()) except asyncio.CancelledError: pass # outer task cancelled by scope; inner task continues @@ -1575,9 +1565,8 @@ async def _bg_producer() -> None: except Exception: # pylint: disable=broad-exception-caught pass # SSE connection dropped; bg_task continues independently 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. + # 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 bg_task 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 d18dab52fb1c..a8744113e109 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 @@ -186,14 +186,11 @@ async def with_keep_alive( ) -> AsyncIterator[str]: """Interleave SSE keep-alive comment frames into ``source`` during idle gaps. - Passthrough when ``interval_seconds`` is falsy. A keep-alive comment frame is emitted - whenever no upstream item arrives within ``interval_seconds``; real items are never - dropped or reordered. Keep-alive is a transport concern, so this wraps the SSE byte - stream at the point it becomes an HTTP response rather than living in the event pipeline. - - The source is advanced by a single dedicated task, so any ``contextvars`` the source sets - once and relies on across its whole lifetime (the request/platform context and the SSE - sequence-number counter) persist exactly as when the source is iterated directly. + 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] @@ -225,8 +222,8 @@ async def _pump() -> None: get_task = asyncio.ensure_future(queue.get()) done, _ = await asyncio.wait({get_task}, timeout=interval_seconds) if get_task not in done: - # Idle interval elapsed with no item; emit a heartbeat. The pending - # get_task is preserved (not cancelled) so no item is ever dropped. + # 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() @@ -235,10 +232,8 @@ async def _pump() -> None: break yield item finally: - # Deterministic cleanup — including "client disconnects right after a real item" - # (consumer suspended at ``yield item`` with the pump still running): cancelling the - # pump throws CancelledError into the source at its yield point, so the source's own - # finally (finalize / reset_request_context) runs and is awaited here. + # 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() 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 04510da6b50d..2c2181a4b32c 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,10 +178,7 @@ 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 must be present, unique, and strictly increasing. (A plain - # `== sorted(seq_nums)` check would pass on a corrupted all-zero stream, so assert - # uniqueness + strict monotonicity — this guards the contextvars-preservation - # requirement of the transport-layer keep-alive wrapper.) + # 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 len(seq_nums) >= 2 assert all(isinstance(n, int) for n in seq_nums) diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_keep_alive_wrapper.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_keep_alive_wrapper.py deleted file mode 100644 index d0c650877e2b..000000000000 --- a/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_keep_alive_wrapper.py +++ /dev/null @@ -1,118 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. -"""Unit tests for the transport-layer SSE keep-alive wrapper ``with_keep_alive``.""" - -from __future__ import annotations - -import asyncio -import contextvars -from typing import AsyncIterator, List - -import pytest - -from azure.ai.agentserver.responses.streaming._sse import encode_keep_alive_comment, with_keep_alive - -_KA = encode_keep_alive_comment() - - -async def _collect(agen: AsyncIterator[str]) -> List[str]: - return [item async for item in agen] - - -async def test_passthrough_when_interval_disabled() -> None: - """None / 0 interval => pure passthrough, no heartbeats, order preserved.""" - - 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_emits_heartbeat_on_idle_gap() -> None: - """A keep-alive frame is emitted when the source is idle longer than the interval.""" - - async def source() -> AsyncIterator[str]: - yield "a" - await asyncio.sleep(0.25) # idle gap >> interval - yield "b" - - out = await _collect(with_keep_alive(source(), 0.05)) - assert _KA in out - # Real items are preserved in order; heartbeats only appear between them. - assert [item for item in out if item != _KA] == ["a", "b"] - assert out.index("a") < out.index("b") - - -async def test_does_not_drop_or_reorder_items() -> None: - """Items arriving across multiple idle gaps are never dropped or reordered.""" - - async def source() -> AsyncIterator[str]: - for i in range(5): - await asyncio.sleep(0.03) - yield f"item-{i}" - - out = await _collect(with_keep_alive(source(), 0.01)) - assert [item for item in out if item != _KA] == [f"item-{i}" for i in range(5)] - - -async def test_preserves_contextvars_across_yields() -> None: - """The source is advanced by a single task, so a ContextVar it sets once persists - across all its yields. (Regression guard: a new-task-per-anext design loses it after - the first item and corrupts request context + SSE sequence numbers.)""" - - 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)) # fast drain, no heartbeats - assert out == ["a", "b"] - assert seen == ["SET", "SET", "SET"] - - -async def test_closes_source_deterministically_on_early_close() -> None: - """Closing the wrapper after one item (client disconnect mid-run) runs the source's - ``finally`` deterministically rather than deferring it to GC.""" - - finalized = asyncio.Event() - - async def source() -> AsyncIterator[str]: - try: - yield "a" - await asyncio.Event().wait() # block forever after the first item - yield "b" - finally: - finalized.set() - - agen = with_keep_alive(source(), 5) - first = await agen.__anext__() - assert first == "a" - await agen.aclose() - await asyncio.wait_for(finalized.wait(), timeout=1.0) - - -async def test_emits_heartbeat_while_source_idle_from_start() -> None: - """A source idle from the very start still gets keep-alive frames — the reconnect / - replay GET scenario, where the stream subscribes to an in-flight background run that - has not emitted anything since the cursor.""" - - async def idle_source() -> AsyncIterator[str]: - await asyncio.Event().wait() # never yields a real item - yield "never" # pragma: no cover - - agen = with_keep_alive(idle_source(), 0.05) - first = await asyncio.wait_for(agen.__anext__(), timeout=1.0) - assert first == _KA - await agen.aclose() - - -if __name__ == "__main__": # pragma: no cover - pytest.main([__file__, "-v"]) 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..8e43a1edcd5d 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,14 @@ from __future__ import annotations +import asyncio +import contextvars +from typing import AsyncIterator, List + 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 +66,96 @@ 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(agen: AsyncIterator[str]) -> List[str]: + return [item async for item in agen] + + +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_item_order_across_gaps() -> None: + """Items arriving across several idle gaps keep their order and count.""" + + async def source() -> AsyncIterator[str]: + for i in range(5): + await asyncio.sleep(0.03) + yield f"item-{i}" + + out = await _collect(with_keep_alive(source(), 0.01)) + assert [item for item in out if item != _KA] == [f"item-{i}" for i in range(5)] + + +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() + + agen = with_keep_alive(source(), 5) + assert await agen.__anext__() == "a" + await agen.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 + + agen = with_keep_alive(idle_source(), 0.05) + assert await asyncio.wait_for(agen.__anext__(), timeout=1.0) == _KA + await agen.aclose() From e25ea879db23b3362c84a47a456cfcfd4cdc5267 Mon Sep 17 00:00:00 2001 From: Antriksh Jain Date: Thu, 9 Jul 2026 17:56:29 +0530 Subject: [PATCH 3/4] Remove redundant keep-alive wrapper test The multi-gap order test covered the same branch and assertions as the single idle-gap heartbeat test. --- .../tests/unit/test_sse_writer.py | 12 ------------ 1 file changed, 12 deletions(-) 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 8e43a1edcd5d..b8792c758d52 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 @@ -97,18 +97,6 @@ async def source() -> AsyncIterator[str]: assert out.index("a") < out.index("b") -async def test_with_keep_alive__preserves_item_order_across_gaps() -> None: - """Items arriving across several idle gaps keep their order and count.""" - - async def source() -> AsyncIterator[str]: - for i in range(5): - await asyncio.sleep(0.03) - yield f"item-{i}" - - out = await _collect(with_keep_alive(source(), 0.01)) - assert [item for item in out if item != _KA] == [f"item-{i}" for i in range(5)] - - 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.""" From e868c9fe84ccb68faf2efd43a968b6f94c834bad Mon Sep 17 00:00:00 2001 From: Antriksh Jain Date: Fri, 10 Jul 2026 12:12:26 +0530 Subject: [PATCH 4/4] Address review feedback on keep-alive wrapper Widen interval_seconds to float, surface source errors to the consumer instead of swallowing them, drop a redundant assertion, and rename a test variable flagged by spell-check. --- .../responses/hosting/_orchestrator.py | 2 -- .../agentserver/responses/streaming/_sse.py | 10 ++++-- .../tests/contract/test_keep_alive.py | 1 - .../tests/unit/test_sse_writer.py | 33 ++++++++++++++----- 4 files changed, 33 insertions(+), 13 deletions(-) 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 04bedb9cbc93..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 @@ -1499,8 +1499,6 @@ async def _live_stream(self, ctx: _ExecutionContext) -> AsyncIterator[str]: async def _finalize() -> None: await self._finalize_stream(ctx, state) - # The orchestrator emits only real events; keep-alive frames are added by the - # transport layer (streaming._sse.with_keep_alive in the endpoint handler). if not (ctx.background and ctx.store): # Simple path for non-background (or non-store) streaming. _stream_completed = False 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 a8744113e109..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 @@ -182,7 +182,7 @@ def encode_keep_alive_comment(comment: str = "keep-alive") -> str: async def with_keep_alive( source: AsyncIterator[str], - interval_seconds: int | None, + interval_seconds: float | None, ) -> AsyncIterator[str]: """Interleave SSE keep-alive comment frames into ``source`` during idle gaps. @@ -195,7 +195,7 @@ async def with_keep_alive( :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: int | None + :type interval_seconds: float | None :returns: An async iterator that yields the source's items plus periodic keep-alive frames. :rtype: AsyncIterator[str] """ @@ -206,11 +206,15 @@ async def with_keep_alive( 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) @@ -231,6 +235,8 @@ async def _pump() -> 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. 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 2c2181a4b32c..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 @@ -183,7 +183,6 @@ def test_keep_alive__does_not_disrupt_event_stream_integrity() -> None: 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 seq_nums == sorted(seq_nums) assert all(b > a for a, b in zip(seq_nums, seq_nums[1:])), f"sequence numbers not strictly increasing: {seq_nums}" 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 b8792c758d52..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 @@ -8,6 +8,8 @@ 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 @@ -68,8 +70,8 @@ def _extract_sequence_number(encoded: str) -> int: assert seq_second == 1, f"second sequence_number must be 1 for a new stream, got {seq_second}" -async def _collect(agen: AsyncIterator[str]) -> List[str]: - return [item async for item in agen] +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: @@ -130,9 +132,9 @@ async def source() -> AsyncIterator[str]: finally: finalized.set() - agen = with_keep_alive(source(), 5) - assert await agen.__anext__() == "a" - await agen.aclose() + stream = with_keep_alive(source(), 5) + assert await stream.__anext__() == "a" + await stream.aclose() await asyncio.wait_for(finalized.wait(), timeout=1.0) @@ -144,6 +146,21 @@ async def idle_source() -> AsyncIterator[str]: await asyncio.Event().wait() yield "never" # pragma: no cover - agen = with_keep_alive(idle_source(), 0.05) - assert await asyncio.wait_for(agen.__anext__(), timeout=1.0) == _KA - await agen.aclose() + 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"]