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..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 @@ -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. 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..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 @@ -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,15 @@ 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.""" + 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) @@ -283,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 @@ -441,3 +474,132 @@ 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. + + 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 + 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 + + 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" + ) + 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) + + +@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"