diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index 6c1baca5bb42..de3343ebd8a7 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -1,5 +1,18 @@ # Release History +## 1.0.0 (2026-08-07) + +### Bugs Fixed + +- Added SSE keep-alive comments to idle `POST /invocations` event streams when + `SSE_KEEPALIVE_INTERVAL` is configured, preventing hosted proxy idle timeouts + from disconnecting clients before the agent emits its final events. + +### Other Changes + +- Updated the minimum `azure-ai-agentserver-core` dependency to the stable + `2.0.0` release. + ## 1.0.0b8 (2026-08-03) ### Samples diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index 6769c2b2ca16..dd2548c9ede8 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md @@ -313,6 +313,7 @@ The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`. | Environment variable | Default | Description | |---|---|---| +| `SSE_KEEPALIVE_INTERVAL` | unset (disabled) | Platform-injected interval, in seconds, for SSE comments on idle `POST /invocations` streams. `0` disables keep-alive. Surfaced on `app.config.sse_keepalive_interval`. | | `WS_KEEPALIVE_INTERVAL` | unset (disabled) | Platform-injected WebSocket Ping interval, in seconds. `0` disables keep-alive. Surfaced on `app.config.ws_ping_interval` and wired into Hypercorn's `websocket_ping_interval` by `AgentServerHost`. | ### Reference: close codes diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/api.metadata.yml b/sdk/agentserver/azure-ai-agentserver-invocations/api.metadata.yml index 570a1d2dada6..b3a77c616f75 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/api.metadata.yml +++ b/sdk/agentserver/azure-ai-agentserver-invocations/api.metadata.yml @@ -1,3 +1,3 @@ apiMdSha256: 28f40971c3ed93df127abaf7faf714a62065ef4440aaf511b9269c584811520c -parserVersion: 0.3.30 -pythonVersion: 3.11.9 +parserVersion: 0.3.31 +pythonVersion: 3.11.15 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py index e14184a6e2a1..0f36560fe2bd 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py @@ -38,6 +38,7 @@ from ._constants import InvocationConstants from ._invocation_ws import _WSHandlerMixin +from ._sse import _with_keep_alive logger = logging.getLogger("azure.ai.agentserver") @@ -480,7 +481,15 @@ async def _wrapped_body() -> AsyncIterator[Any]: if stream_ctx_token is not None: reset_request_context(stream_ctx_token) - response.body_iterator = _wrapped_body() + wrapped_body = _wrapped_body() + content_type = response.headers.get("content-type") or response.media_type or "" + if content_type.lower().startswith("text/event-stream"): + response.body_iterator = _with_keep_alive( + wrapped_body, + self.config.sse_keepalive_interval, + ) + else: + response.body_iterator = wrapped_body return response async def _create_invocation_endpoint(self, request: Request) -> Response: diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_sse.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_sse.py new file mode 100644 index 000000000000..557adbc51742 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_sse.py @@ -0,0 +1,77 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Internal server-sent events helpers for the invocations protocol.""" + +import asyncio # pylint: disable=do-not-import-asyncio +from collections.abc import AsyncIterator +from typing import Any + +_KEEP_ALIVE_COMMENT = ": keep-alive\n\n" + + +async def _with_keep_alive( + source: AsyncIterator[Any], + interval_seconds: float | None, +) -> AsyncIterator[Any]: + """Interleave SSE keep-alive comments while the source is idle. + + :param source: The source SSE stream. + :type source: ~collections.abc.AsyncIterator + :param interval_seconds: Seconds of inactivity between keep-alive comments, + or zero/``None`` to disable them. + :type interval_seconds: float or None + :return: The source chunks with keep-alive comments inserted during idle periods. + :rtype: ~collections.abc.AsyncIterator + """ + if not interval_seconds: + async for item in source: + yield item + return + + requests: asyncio.Queue[asyncio.Future[Any]] = asyncio.Queue(maxsize=1) + sentinel = object() + + async def _pump() -> None: + try: + while True: + result = await requests.get() + try: + item = await anext(source) + except StopAsyncIteration: + result.set_result(sentinel) + return + except BaseException as exc: # pylint: disable=broad-exception-caught + if not result.done(): + result.set_exception(exc) + return + result.set_result(item) + finally: + close = getattr(source, "aclose", None) + if close is not None: + await close() + + pump_task = asyncio.create_task(_pump()) + next_result: asyncio.Future[Any] | None = None + try: + while True: + if next_result is None: + next_result = asyncio.get_running_loop().create_future() + requests.put_nowait(next_result) + try: + item = await asyncio.wait_for( + asyncio.shield(next_result), + timeout=interval_seconds, + ) + except asyncio.TimeoutError: + yield _KEEP_ALIVE_COMMENT + continue + next_result = None + if item is sentinel: + break + yield item + finally: + if next_result is not None: + next_result.cancel() + pump_task.cancel() + await asyncio.gather(pump_task, return_exceptions=True) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py index 1b0058b20ee4..ed845cd7ee2a 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py @@ -2,4 +2,4 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -VERSION = "1.0.0b8" +VERSION = "1.0.0" diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/pyproject.toml b/sdk/agentserver/azure-ai-agentserver-invocations/pyproject.toml index 1d56a4ba2618..e8a41e49ac4c 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/pyproject.toml +++ b/sdk/agentserver/azure-ai-agentserver-invocations/pyproject.toml @@ -8,7 +8,7 @@ authors = [ ] license = "MIT" classifiers = [ - "Development Status :: 4 - Beta", + "Development Status :: 5 - Production/Stable", "Programming Language :: Python", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3", @@ -21,7 +21,7 @@ classifiers = [ keywords = ["azure", "azure sdk", "agent", "agentserver", "invocations"] dependencies = [ - "azure-ai-agentserver-core>=2.0.0b10", + "azure-ai-agentserver-core>=2.0.0", # Constraint on the transitive aiohttp: the `--pre` CI install otherwise # resolves the unbuildable aiohttp 4.0.0a1 alpha. Cap must be <4.0.0a0 # (<4.0.0 still admits 4.0.0a1 under PEP 440). @@ -72,8 +72,6 @@ mypy = true pyright = true verifytypes = false latestdependency = false -# azure-ai-agentserver-core>=2.0.0b8 is not yet on PyPI -mindependency = false pylint = true type_check_samples = false diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_sse_keep_alive.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_sse_keep_alive.py new file mode 100644 index 000000000000..a8331c217bf0 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_sse_keep_alive.py @@ -0,0 +1,218 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for invocations SSE keep-alive behavior.""" + +import asyncio +from collections.abc import AsyncIterator + +import pytest +from httpx import ASGITransport, AsyncClient +from starlette.requests import Request +from starlette.responses import StreamingResponse + +from azure.ai.agentserver.invocations import InvocationAgentServerHost +from azure.ai.agentserver.invocations._sse import _with_keep_alive + + +async def _collect(source: AsyncIterator[object]) -> list[object]: + return [item async for item in source] + + +@pytest.mark.asyncio +async def test_with_keep_alive_emits_comment_during_idle_gap() -> None: + async def source() -> AsyncIterator[str]: + yield "started" + await asyncio.sleep(0.1) + yield "finished" + + chunks = await _collect(_with_keep_alive(source(), 0.02)) + + assert ": keep-alive\n\n" in chunks + assert [chunk for chunk in chunks if chunk != ": keep-alive\n\n"] == [ + "started", + "finished", + ] + + +@pytest.mark.asyncio +async def test_with_keep_alive_passthrough_when_disabled() -> None: + async def source() -> AsyncIterator[bytes]: + yield b"first" + yield b"second" + + assert await _collect(_with_keep_alive(source(), 0)) == [b"first", b"second"] + + +@pytest.mark.asyncio +async def test_with_keep_alive_propagates_source_error() -> None: + async def source() -> AsyncIterator[str]: + yield "started" + raise RuntimeError("stream failed") + + stream = _with_keep_alive(source(), 0.02) + + assert await anext(stream) == "started" + with pytest.raises(RuntimeError, match="stream failed"): + await anext(stream) + + +@pytest.mark.asyncio +async def test_with_keep_alive_propagates_source_cancellation() -> None: + async def source() -> AsyncIterator[str]: + yield "started" + raise asyncio.CancelledError + + stream = _with_keep_alive(source(), 0.02) + + assert await anext(stream) == "started" + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(anext(stream), timeout=1.0) + + +@pytest.mark.asyncio +async def test_with_keep_alive_closes_source_when_consumer_stops() -> None: + finalized = asyncio.Event() + + async def source() -> AsyncIterator[str]: + try: + yield "started" + await asyncio.Event().wait() + finally: + finalized.set() + + stream = _with_keep_alive(source(), 0.02) + + assert await anext(stream) == "started" + await stream.aclose() + + await asyncio.wait_for(finalized.wait(), timeout=1.0) + + +@pytest.mark.asyncio +async def test_with_keep_alive_preserves_source_backpressure() -> None: + produced = 0 + + async def source() -> AsyncIterator[str]: + nonlocal produced + for index in range(100): + produced += 1 + yield str(index) + + stream = _with_keep_alive(source(), 0.02) + + assert await anext(stream) == "0" + await asyncio.sleep(0.05) + assert produced == 1 + + await stream.aclose() + + +@pytest.mark.asyncio +async def test_invocations_sse_stream_uses_configured_keep_alive(monkeypatch) -> None: + monkeypatch.setenv("SSE_KEEPALIVE_INTERVAL", "1") + app = InvocationAgentServerHost(configure_observability=None) + + @app.invoke_handler + async def handle(request: Request) -> StreamingResponse: + async def generate() -> AsyncIterator[str]: + yield "data: started\n\n" + await asyncio.sleep(1.1) + yield "data: finished\n\n" + + return StreamingResponse(generate(), media_type="text/event-stream") + + async with AsyncClient( + transport=ASGITransport(app=app), + base_url="http://testserver", + ) as client: + response = await client.post( + "/invocations", + headers={"accept": "text/event-stream"}, + ) + + assert response.status_code == 200 + assert ": keep-alive\n\n" in response.text + assert response.text.index("data: started") < response.text.index("data: finished") + + +@pytest.mark.asyncio +async def test_invocations_non_sse_stream_does_not_get_keep_alive(monkeypatch) -> None: + monkeypatch.setenv("SSE_KEEPALIVE_INTERVAL", "1") + app = InvocationAgentServerHost(configure_observability=None) + + @app.invoke_handler + async def handle(request: Request) -> StreamingResponse: + async def generate() -> AsyncIterator[str]: + yield '{"chunk": 1}\n' + await asyncio.sleep(1.1) + yield '{"chunk": 2}\n' + + return StreamingResponse(generate(), media_type="application/x-ndjson") + + async with AsyncClient( + transport=ASGITransport(app=app), + base_url="http://testserver", + ) as client: + response = await client.post("/invocations") + + assert response.status_code == 200 + assert ": keep-alive" not in response.text + + +@pytest.mark.asyncio +async def test_invocations_uses_final_content_type_header(monkeypatch) -> None: + monkeypatch.setenv("SSE_KEEPALIVE_INTERVAL", "1") + app = InvocationAgentServerHost(configure_observability=None) + + @app.invoke_handler + async def handle(request: Request) -> StreamingResponse: + async def generate() -> AsyncIterator[str]: + yield '{"chunk": 1}\n' + await asyncio.sleep(1.1) + yield '{"chunk": 2}\n' + + return StreamingResponse( + generate(), + media_type="text/event-stream", + headers={"content-type": "application/x-ndjson"}, + ) + + async with AsyncClient( + transport=ASGITransport(app=app), + base_url="http://testserver", + ) as client: + response = await client.post("/invocations") + + assert response.status_code == 200 + assert response.headers["content-type"] == "application/x-ndjson" + assert ": keep-alive" not in response.text + + +@pytest.mark.asyncio +async def test_invocations_adds_keep_alive_for_final_sse_header(monkeypatch) -> None: + monkeypatch.setenv("SSE_KEEPALIVE_INTERVAL", "1") + app = InvocationAgentServerHost(configure_observability=None) + + @app.invoke_handler + async def handle(request: Request) -> StreamingResponse: + async def generate() -> AsyncIterator[str]: + yield "data: started\n\n" + await asyncio.sleep(1.1) + yield "data: finished\n\n" + + return StreamingResponse( + generate(), + media_type="application/x-ndjson", + headers={"content-type": "text/event-stream"}, + ) + + async with AsyncClient( + transport=ASGITransport(app=app), + base_url="http://testserver", + ) as client: + response = await client.post("/invocations") + + assert response.status_code == 200 + assert response.headers["content-type"] == "text/event-stream" + assert ": keep-alive\n\n" in response.text