fix(agentserver-responses): keep background+stream handler alive on client disconnect when keep-alive is enabled

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,67 +1528,10 @@ async def _finalize() -> None:
                     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())
-            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:
-                # 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 ---
@@ -1647,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.