From b9c25b50bf2c135d2730f9a8e179d4e7ab6d9873 Mon Sep 17 00:00:00 2001 From: Antriksh Jain Date: Fri, 3 Jul 2026 18:32:33 +0530 Subject: [PATCH 1/4] fix(agentserver-responses): keep background+stream handler alive on client disconnect when keep-alive is enabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The FR-012/FR-013 "background+stream survives client disconnect" behavior was only implemented on the no-keep-alive fast path. When SSE keep-alive is enabled (sse_keep_alive_interval_seconds set) — which is how the hosted platform runs the SDK — streaming took the keep-alive merge path, whose finally cancels the handler task on client disconnect. That killed in-flight background runs, so a later GET /responses/{id}?stream=true&starting_after=N reconnected to a dead run and returned 200 with no new events. Route background+store streams to the shielded, independent-task path regardless of keep-alive, and emit keep-alive comments from that path (fed to the same queue, stopped from both the producer and consumer finally so it can never outlive the run). The keep-alive merge path now only handles non-background streams, where cancelling on disconnect is correct. Add regression tests that run the disconnect scenario with keep-alive enabled. --- .../responses/hosting/_orchestrator.py | 47 ++++++- .../contract/test_bg_stream_disconnect.py | 119 ++++++++++++++++++ 2 files changed, 163 insertions(+), 3 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 999d5d641105..5d93a373e405 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,14 @@ 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: + # Fast path when keep-alive is disabled; ALSO always take this path for + # background+store streams so the handler runs as a shielded, independent task + # that survives client disconnect (FR-013) — keep-alive comments are still + # emitted below when configured. The hosted platform enables keep-alive, and the + # keep-alive merge path further down cancels the handler on client disconnect, + # which would kill an in-flight background run and break + # GET ?stream=true&starting_after=N reconnect. + if not self._runtime_options.sse_keep_alive_enabled or (ctx.background and ctx.store): if not (ctx.background and ctx.store): # Simple fast path for non-background streaming. _stream_completed = False @@ -1522,7 +1528,7 @@ async def _finalize() -> None: await _finalize() return - # Background+stream without keep-alive: run the handler as an independent + # 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 @@ -1530,6 +1536,10 @@ async def _finalize() -> None: # promptly), leaving GET-replay subscribers blocked on await q.get() forever. _SENTINEL_BG = object() bg_queue: asyncio.Queue[object] = asyncio.Queue() + # Optional keep-alive companion (assigned below when enabled). Declared here + # so the producer's own finally can stop it even if the consumer generator's + # finalizer is delayed — preventing a leaked task / unbounded bg_queue growth. + bg_keep_alive_task: "asyncio.Task[None] | None" = None async def _bg_producer_inner() -> None: try: @@ -1548,6 +1558,11 @@ async def _bg_producer_inner() -> None: ) state.captured_error = exc finally: + # Stop keep-alive the moment the run finishes — even if the consumer + # generator's finalizer is delayed — so it can't keep enqueueing into + # an unconsumed bg_queue. + if bg_keep_alive_task is not None and not bg_keep_alive_task.done(): + bg_keep_alive_task.cancel() # 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. @@ -1566,6 +1581,25 @@ async def _bg_producer() -> None: pass # outer task cancelled by scope; inner task continues bg_task = asyncio.create_task(_bg_producer()) + + # Emit periodic keep-alive comments while the client is still connected (only + # when configured) so an idle background handler doesn't get its live SSE + # connection dropped by an intermediary. Comments feed the same queue; the + # task is stopped from BOTH the producer finally (run done) and the consumer + # finally (client gone), whichever fires first, so it never outlives the run. + if self._runtime_options.sse_keep_alive_enabled: + + async def _bg_keep_alive(interval: int) -> None: + try: + while True: + await asyncio.sleep(interval) + await bg_queue.put(encode_keep_alive_comment()) + except asyncio.CancelledError: + return + + bg_keep_alive_task = asyncio.create_task( + _bg_keep_alive(self._runtime_options.sse_keep_alive_interval_seconds) # type: ignore[arg-type] + ) try: while True: item = await bg_queue.get() @@ -1575,6 +1609,13 @@ async def _bg_producer() -> None: except Exception: # pylint: disable=broad-exception-caught pass # SSE connection dropped; bg_task continues independently finally: + # Stop keep-alive first — there is no client left to keep alive. + if bg_keep_alive_task is not None and not bg_keep_alive_task.done(): + bg_keep_alive_task.cancel() + try: + await bg_keep_alive_task + except asyncio.CancelledError: + pass # 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. diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py index 036506cbe7a4..8ecc5601c870 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py @@ -21,6 +21,7 @@ from azure.ai.agentserver.responses import ResponsesAgentServerHost from azure.ai.agentserver.responses._id_generator import IdGenerator +from azure.ai.agentserver.responses._options import ResponsesServerOptions from azure.ai.agentserver.responses.streaming._event_stream import ResponseEventStream # ════════════════════════════════════════════════════════════ @@ -164,6 +165,16 @@ def _build_client(handler: Any) -> _AsyncAsgiClient: return _AsyncAsgiClient(app) +def _build_client_keep_alive(handler: Any, *, keep_alive_seconds: int = 1) -> _AsyncAsgiClient: + """Build a client with SSE keep-alive ENABLED — mirrors the hosted platform, which + runs with ``sse_keepalive_interval`` set and therefore routes streaming through the + keep-alive merge path.""" + options = ResponsesServerOptions(sse_keep_alive_interval_seconds=keep_alive_seconds) + app = ResponsesAgentServerHost(options=options) + app.response_handler(handler) + return _AsyncAsgiClient(app) + + async def _ensure_task_done(task: asyncio.Task[Any], handler: Any, timeout: float = 5.0) -> None: for attr in vars(handler): obj = getattr(handler, attr, None) @@ -441,3 +452,111 @@ async def test_bg_nostream_handler_continues_after_disconnect() -> None: get_resp = await client.get(f"/responses/{response_id}") assert get_resp.status_code == 200 assert get_resp.json()["status"] == "completed" + + +# ════════════════════════════════════════════════════════════ +# Keep-alive ENABLED (hosted-platform config): the same disconnect +# invariants must hold. Regression for the keep-alive merge path +# cancelling the background handler on client disconnect. +# ════════════════════════════════════════════════════════════ + + +@pytest.mark.asyncio +async def test_bg_stream_disconnect_handler_completes_all_events_with_keep_alive() -> None: + """FR-012 with keep-alive ENABLED — the hosted-platform configuration. + + The hosted platform runs with ``sse_keepalive_interval`` set, which routes streaming + through the keep-alive merge path. That path must NOT cancel a background handler on + client disconnect: the run must complete and GET must return 'completed' with all + output items, exactly like the keep-alive-disabled case (T036). + """ + total = 10 + disconnect_after = 3 + handler = _make_multi_output_handler(total, disconnect_after) + client = _build_client_keep_alive(handler, keep_alive_seconds=1) + response_id = IdGenerator.new_response_id() + + disconnect = asyncio.Event() + post_task = asyncio.create_task( + client.request_with_disconnect( + "POST", + "/responses", + json_body={ + "response_id": response_id, + "model": "test", + "background": True, + "stream": True, + "store": True, + }, + disconnect_event=disconnect, + ) + ) + try: + await asyncio.wait_for(handler.ready_for_disconnect.wait(), timeout=5.0) + + disconnect.set() + try: + await asyncio.wait_for(post_task, timeout=2.0) + except (asyncio.TimeoutError, Exception): + pass + + await asyncio.wait_for(handler.handler_completed.wait(), timeout=5.0) + await _wait_for_background_completion(client, response_id) + + get_resp = await client.get(f"/responses/{response_id}") + assert get_resp.status_code == 200 + doc = get_resp.json() + assert doc["status"] == "completed", ( + f"FR-012 (keep-alive): bg+stream handler should complete after disconnect, got '{doc['status']}'" + ) + finally: + await _ensure_task_done(post_task, handler) + + +@pytest.mark.asyncio +async def test_bg_stream_disconnect_does_not_cancel_handler_with_keep_alive() -> None: + """FR-013 with keep-alive ENABLED — the keep-alive merge path must not cancel the + background handler on client disconnect; it should complete normally.""" + handler = _make_cancellation_tracking_handler() + client = _build_client_keep_alive(handler, keep_alive_seconds=1) + response_id = IdGenerator.new_response_id() + + disconnect = asyncio.Event() + post_task = asyncio.create_task( + client.request_with_disconnect( + "POST", + "/responses", + json_body={ + "response_id": response_id, + "model": "test", + "background": True, + "stream": True, + "store": True, + }, + disconnect_event=disconnect, + ) + ) + try: + await asyncio.wait_for(handler.ready_for_disconnect.wait(), timeout=5.0) + + disconnect.set() + try: + await asyncio.wait_for(post_task, timeout=2.0) + except (asyncio.TimeoutError, Exception): + pass + + done_events = [handler.handler_completed, handler.handler_cancelled] + await asyncio.wait( + [asyncio.create_task(e.wait()) for e in done_events], + timeout=3.0, + return_when=asyncio.FIRST_COMPLETED, + ) + + assert handler.handler_completed.is_set(), ( + "FR-013 (keep-alive): handler should complete normally, not be cancelled by SSE disconnect" + ) + assert not handler.handler_cancelled.is_set(), ( + "FR-013 (keep-alive): handler CT should NOT be cancelled by SSE disconnect" + ) + finally: + await _ensure_task_done(post_task, handler) From 953323e8404602142f1539bddcacaeda6939f167 Mon Sep 17 00:00:00 2001 From: Antriksh Jain Date: Fri, 3 Jul 2026 19:02:45 +0530 Subject: [PATCH 2/4] address review comments - orchestrator: capture keep-alive interval in a local (drops type: ignore[arg-type]); narrow queue item with isinstance(str) instead of type: ignore[misc] - tests: correct keep-alive docstrings (background+store no longer uses the merge path); cancel/gather the event-wait tasks to avoid leaked pending tasks --- .../responses/hosting/_orchestrator.py | 10 +++---- .../contract/test_bg_stream_disconnect.py | 26 ++++++++++--------- 2 files changed, 19 insertions(+), 17 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 5d93a373e405..876153b5daf4 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 @@ -1587,7 +1587,8 @@ async def _bg_producer() -> None: # connection dropped by an intermediary. Comments feed the same queue; the # task is stopped from BOTH the producer finally (run done) and the consumer # finally (client gone), whichever fires first, so it never outlives the run. - if self._runtime_options.sse_keep_alive_enabled: + _ka_interval = self._runtime_options.sse_keep_alive_interval_seconds + if _ka_interval is not None: # keep-alive enabled async def _bg_keep_alive(interval: int) -> None: try: @@ -1597,15 +1598,14 @@ async def _bg_keep_alive(interval: int) -> None: except asyncio.CancelledError: return - bg_keep_alive_task = asyncio.create_task( - _bg_keep_alive(self._runtime_options.sse_keep_alive_interval_seconds) # type: ignore[arg-type] - ) + bg_keep_alive_task = asyncio.create_task(_bg_keep_alive(_ka_interval)) try: while True: item = await bg_queue.get() if item is _SENTINEL_BG: break - yield item # type: ignore[misc] + if isinstance(item, str): + yield item except Exception: # pylint: disable=broad-exception-caught pass # SSE connection dropped; bg_task continues independently finally: diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py index 8ecc5601c870..9ae79c567682 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py @@ -167,8 +167,7 @@ def _build_client(handler: Any) -> _AsyncAsgiClient: def _build_client_keep_alive(handler: Any, *, keep_alive_seconds: int = 1) -> _AsyncAsgiClient: """Build a client with SSE keep-alive ENABLED — mirrors the hosted platform, which - runs with ``sse_keepalive_interval`` set and therefore routes streaming through the - keep-alive merge path.""" + runs with ``sse_keepalive_interval`` set.""" options = ResponsesServerOptions(sse_keep_alive_interval_seconds=keep_alive_seconds) app = ResponsesAgentServerHost(options=options) app.response_handler(handler) @@ -465,10 +464,10 @@ async def test_bg_nostream_handler_continues_after_disconnect() -> None: async def test_bg_stream_disconnect_handler_completes_all_events_with_keep_alive() -> None: """FR-012 with keep-alive ENABLED — the hosted-platform configuration. - The hosted platform runs with ``sse_keepalive_interval`` set, which routes streaming - through the keep-alive merge path. That path must NOT cancel a background handler on - client disconnect: the run must complete and GET must return 'completed' with all - output items, exactly like the keep-alive-disabled case (T036). + Enabling keep-alive must not change background disconnect survival: the run must + complete and GET must return 'completed' with all output items, exactly like the + keep-alive-disabled case (T036). (Regression: keep-alive previously routed background + streams through the merge path, which cancelled the handler on disconnect.) """ total = 10 disconnect_after = 3 @@ -545,12 +544,15 @@ async def test_bg_stream_disconnect_does_not_cancel_handler_with_keep_alive() -> except (asyncio.TimeoutError, Exception): pass - done_events = [handler.handler_completed, handler.handler_cancelled] - await asyncio.wait( - [asyncio.create_task(e.wait()) for e in done_events], - timeout=3.0, - return_when=asyncio.FIRST_COMPLETED, - ) + wait_tasks = [ + asyncio.create_task(e.wait()) for e in (handler.handler_completed, handler.handler_cancelled) + ] + try: + await asyncio.wait(wait_tasks, timeout=3.0, return_when=asyncio.FIRST_COMPLETED) + finally: + for _t in wait_tasks: + _t.cancel() + await asyncio.gather(*wait_tasks, return_exceptions=True) assert handler.handler_completed.is_set(), ( "FR-013 (keep-alive): handler should complete normally, not be cancelled by SSE disconnect" From 2981130f984221c2f5113663b1fc5e3b68fda0af Mon Sep 17 00:00:00 2001 From: Antriksh Jain Date: Mon, 6 Jul 2026 15:55:19 +0530 Subject: [PATCH 3/4] refactor: extract background+stream path into _live_stream_background Fixes the Analyze/pylint CI failure: the inline keep-alive additions pushed _live_stream over pylint complexity limits (too-many-locals 27/25, too-many-branches 23/20). Move the background+store streaming branch (shielded independent-task lifecycle that survives client disconnect) into a dedicated _live_stream_background method. No behavior change. --- .../responses/hosting/_orchestrator.py | 219 ++++++++++-------- 1 file changed, 123 insertions(+), 96 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 876153b5daf4..3a0c1135a620 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 @@ -1528,102 +1528,10 @@ async def _finalize() -> None: 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. - _SENTINEL_BG = object() - bg_queue: asyncio.Queue[object] = asyncio.Queue() - # Optional keep-alive companion (assigned below when enabled). Declared here - # so the producer's own finally can stop it even if the consumer generator's - # finalizer is delayed — preventing a leaked task / unbounded bg_queue growth. - bg_keep_alive_task: "asyncio.Task[None] | None" = None - - 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: - # Stop keep-alive the moment the run finishes — even if the consumer - # generator's finalizer is delayed — so it can't keep enqueueing into - # an unconsumed bg_queue. - if bg_keep_alive_task is not None and not bg_keep_alive_task.done(): - bg_keep_alive_task.cancel() - # 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()) - - # Emit periodic keep-alive comments while the client is still connected (only - # when configured) so an idle background handler doesn't get its live SSE - # connection dropped by an intermediary. Comments feed the same queue; the - # task is stopped from BOTH the producer finally (run done) and the consumer - # finally (client gone), whichever fires first, so it never outlives the run. - _ka_interval = self._runtime_options.sse_keep_alive_interval_seconds - if _ka_interval is not None: # keep-alive enabled - - async def _bg_keep_alive(interval: int) -> None: - try: - while True: - await asyncio.sleep(interval) - await bg_queue.put(encode_keep_alive_comment()) - except asyncio.CancelledError: - return - - bg_keep_alive_task = asyncio.create_task(_bg_keep_alive(_ka_interval)) - try: - while True: - item = await bg_queue.get() - if item is _SENTINEL_BG: - break - if isinstance(item, str): - yield item - except Exception: # pylint: disable=broad-exception-caught - pass # SSE connection dropped; bg_task continues independently - finally: - # Stop keep-alive first — there is no client left to keep alive. - if bg_keep_alive_task is not None and not bg_keep_alive_task.done(): - bg_keep_alive_task.cancel() - try: - await bg_keep_alive_task - except asyncio.CancelledError: - pass - # 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 + # Background+stream (store): runs as a shielded, independent task that + # survives client disconnect (FR-013). See _live_stream_background. + async for chunk in self._live_stream_background(ctx, state, handler_iterator): + yield chunk return # --- Keep-alive path: merge handler events with periodic keep-alive comments --- @@ -1688,6 +1596,125 @@ async def _keep_alive_producer(interval: int) -> None: pass await _finalize() + async def _live_stream_background( + self, + ctx: _ExecutionContext, + state: _PipelineState, + handler_iterator: "AsyncIterator[generated_models.ResponseStreamEvent]", + ) -> AsyncIterator[str]: + """Drive a background+store SSE stream as a shielded, independent task. + + Extracted from :meth:`_live_stream` so the background lifecycle (which must + survive client disconnect per FR-013) is isolated and the parent method stays + within complexity limits. The handler runs in a shielded task decoupled from + the SSE consumer, so dropping the client connection never cancels the run. + + :param ctx: Current execution context. + :type ctx: _ExecutionContext + :param state: Mutable pipeline state for this invocation. + :type state: _PipelineState + :param handler_iterator: The handler's event async iterator. + :type handler_iterator: AsyncIterator[ResponseStreamEvent] + :returns: Async iterator of SSE-encoded strings. + :rtype: AsyncIterator[str] + """ + # 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() + # Optional keep-alive companion (assigned below when enabled). Declared here + # so the producer's own finally can stop it even if the consumer generator's + # finalizer is delayed — preventing a leaked task / unbounded bg_queue growth. + bg_keep_alive_task: "asyncio.Task[None] | None" = None + + 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: + # Stop keep-alive the moment the run finishes — even if the consumer + # generator's finalizer is delayed — so it can't keep enqueueing into + # an unconsumed bg_queue. + if bg_keep_alive_task is not None and not bg_keep_alive_task.done(): + bg_keep_alive_task.cancel() + # 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 self._finalize_stream(ctx, state) + 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()) + + # Emit periodic keep-alive comments while the client is still connected (only + # when configured) so an idle background handler doesn't get its live SSE + # connection dropped by an intermediary. Comments feed the same queue; the + # task is stopped from BOTH the producer finally (run done) and the consumer + # finally (client gone), whichever fires first, so it never outlives the run. + _ka_interval = self._runtime_options.sse_keep_alive_interval_seconds + if _ka_interval is not None: # keep-alive enabled + + async def _bg_keep_alive(interval: int) -> None: + try: + while True: + await asyncio.sleep(interval) + await bg_queue.put(encode_keep_alive_comment()) + except asyncio.CancelledError: + return + + bg_keep_alive_task = asyncio.create_task(_bg_keep_alive(_ka_interval)) + 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 + finally: + # Stop keep-alive first — there is no client left to keep alive. + if bg_keep_alive_task is not None and not bg_keep_alive_task.done(): + bg_keep_alive_task.cancel() + try: + await bg_keep_alive_task + except asyncio.CancelledError: + pass + # 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 + return + async def run_sync(self, ctx: _ExecutionContext) -> dict[str, Any]: """Execute a synchronous (non-stream, non-background) create-response request. From f6cf8211bfe79e96d8b04c952cbf18a04bacca03 Mon Sep 17 00:00:00 2001 From: Antriksh Jain Date: Mon, 6 Jul 2026 17:35:13 +0530 Subject: [PATCH 4/4] test: cover keep-alive comment emission on background stream Adds a test where a bg+stream handler stays idle longer than the keep-alive interval while the client remains connected, asserting ': keep-alive' comments are received. The existing keep-alive-enabled tests disconnect/finish before the interval elapses, so the emission path was never exercised. --- .../contract/test_bg_stream_disconnect.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py index 9ae79c567682..66773fea155c 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_bg_stream_disconnect.py @@ -293,6 +293,29 @@ async def _events(): return handler +def _make_idle_gap_handler(gap_seconds: float): + """Handler that stays idle between response.created and response.completed for + longer than the keep-alive interval, so keep-alive comments are emitted while + the client remains connected.""" + handler_completed = asyncio.Event() + + def handler(request: Any, context: Any, cancellation_signal: Any): + async def _events(): + stream = ResponseEventStream( + response_id=context.response_id, + model=getattr(request, "model", None), + ) + yield stream.emit_created() + await asyncio.sleep(gap_seconds) + yield stream.emit_completed() + handler_completed.set() + + return _events() + + handler.handler_completed = handler_completed + return handler + + # ════════════════════════════════════════════════════════════ # T036: bg+stream — client disconnects after 3 events, # handler produces 10 total → GET returns completed with all output @@ -562,3 +585,21 @@ async def test_bg_stream_disconnect_does_not_cancel_handler_with_keep_alive() -> ) finally: await _ensure_task_done(post_task, handler) + + +@pytest.mark.asyncio +async def test_bg_stream_emits_keep_alive_comments_while_connected() -> None: + """bg+stream with keep-alive enabled: while the client stays connected and the + handler is idle longer than the keep-alive interval, keep-alive comments are + emitted on the background stream (exercises the emission path, not just setup).""" + handler = _make_idle_gap_handler(gap_seconds=1.5) + client = _build_client_keep_alive(handler, keep_alive_seconds=1) + + response = await client.post( + "/responses", + json_body={"model": "test", "background": True, "stream": True, "store": True}, + ) + + assert response.status_code == 200 + body = response.body.decode() + assert ": keep-alive" in body, "background stream should emit keep-alive comments while the client is connected"