From edf8ac8bfccc17af7bed2eeef3288c8fa6bd5799 Mon Sep 17 00:00:00 2001 From: Xinran Dong Date: Tue, 12 May 2026 17:38:00 +0800 Subject: [PATCH 01/13] Implement WebSocket support for InvocationAgentServerHost (#46805) * Implement WebSocket support for InvocationAgentServerHost - Added `_invocation_ws.py` to handle WebSocket connections and implement the `@app.ws_handler` decorator for the `/invocations_ws` route. - Created a bidirectional streaming sample agent (`ws_bidirectional_streaming_agent.py`) demonstrating full-duplex communication over WebSocket. - Developed an echo agent sample (`ws_invoke_agent.py`) that supports both HTTP and WebSocket protocols. - Introduced requirements files for both sample agents to specify dependencies. - Added comprehensive tests for WebSocket handler functionality, including connection acceptance, error handling, and logging of close events. * remove heartbeat from sample ageng. add testcases * Split the monolithic test_ws_handler.py into 8 focused files mirroring the invocation suite layout * resolve pyright error --- .../CHANGELOG.md | 14 + .../README.md | 60 ++- .../ai/agentserver/invocations/_constants.py | 28 ++ .../ai/agentserver/invocations/_invocation.py | 73 ++- .../agentserver/invocations/_invocation_ws.py | 467 +++++++++++++++++ .../requirements.txt | 1 + .../ws_bidirectional_streaming_agent.py | 250 ++++++++++ .../samples/ws_invoke_agent/requirements.txt | 1 + .../ws_invoke_agent/ws_invoke_agent.py | 58 +++ .../tests/conftest.py | 41 ++ .../tests/test_ws_close_event.py | 113 +++++ .../tests/test_ws_decorator_pattern.py | 92 ++++ .../tests/test_ws_edge_cases.py | 121 +++++ .../tests/test_ws_invoke.py | 143 ++++++ .../tests/test_ws_multimodal_protocol.py | 77 +++ .../tests/test_ws_ping_interval.py | 103 ++++ .../tests/test_ws_server_routes.py | 95 ++++ .../tests/test_ws_session_id.py | 80 +++ .../tests/test_ws_tracing.py | 470 ++++++++++++++++++ 19 files changed, 2282 insertions(+), 5 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/requirements.txt create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/requirements.txt create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index 6ead3c39d58d..7745922fba25 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -4,6 +4,20 @@ ### Features Added +- `invocations_ws` (WebSocket) protocol support on `InvocationAgentServerHost`. + Register a handler with the new `@app.ws_handler` decorator to host a + full-duplex WebSocket endpoint at `/invocations_ws` on the same host that + serves `POST /invocations`. The SDK calls `await websocket.accept()` before + invoking the handler, runs WebSocket Ping/Pong keep-alive in the background + (default 30 s; configurable via the new `ws_ping_interval` constructor + argument), closes the connection cleanly on handler return, and maps + uncaught exceptions to RFC 6455 close code `1011`. Each connection emits a + structured close-event log line carrying `ws.session_id`, `ws.close_code`, + and `ws.duration_ms`, and the same fields are recorded as OpenTelemetry + span attributes. `/readiness`, OTEL export, graceful shutdown, and the + `x-platform-server` identity header continue to be inherited from + `azure-ai-agentserver-core`. + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index 5e9dfe515657..63ec794bbc8f 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md @@ -1,6 +1,9 @@ # Azure AI Agent Server Invocations client library for Python -The `azure-ai-agentserver-invocations` package provides the invocation protocol endpoints for Azure AI Hosted Agent containers. It plugs into the [`azure-ai-agentserver-core`](https://pypi.org/project/azure-ai-agentserver-core/) host framework and adds the full invocation lifecycle: `POST /invocations`, `GET /invocations/{id}`, `POST /invocations/{id}/cancel`, and `GET /invocations/docs/openapi.json`. +The `azure-ai-agentserver-invocations` package provides the invocation protocol endpoints for Azure AI Hosted Agent containers. It plugs into the [`azure-ai-agentserver-core`](https://pypi.org/project/azure-ai-agentserver-core/) host framework and supports two transports on the same host: + +- **HTTP** (`invocations` protocol) — `POST /invocations`, `GET /invocations/{id}`, `POST /invocations/{id}/cancel`, `GET /invocations/docs/openapi.json`. +- **WebSocket** (`invocations_ws` protocol) — full-duplex streaming at `/invocations_ws`, registered with `@app.ws_handler`. ## Getting started @@ -25,6 +28,7 @@ This automatically installs `azure-ai-agentserver-core` as a dependency. - `@app.invoke_handler` — **Required.** Handles `POST /invocations`. - `@app.get_invocation_handler` — Optional. Handles `GET /invocations/{id}`. - `@app.cancel_invocation_handler` — Optional. Handles `POST /invocations/{id}/cancel`. +- `@app.ws_handler` — Optional. Handles WebSocket connections at `/invocations_ws`. ### Protocol endpoints @@ -34,6 +38,7 @@ This automatically installs `azure-ai-agentserver-core` as a dependency. | `GET` | `/invocations/{invocation_id}` | No | Retrieve invocation status or result | | `POST` | `/invocations/{invocation_id}/cancel` | No | Cancel a running invocation | | `GET` | `/invocations/docs/openapi.json` | No | Serve the agent's OpenAPI 3.x spec | +| `WS` | `/invocations_ws` | No | Full-duplex WebSocket transport (`invocations_ws` protocol) | ### Request and response headers @@ -182,6 +187,57 @@ app = InvocationAgentServerHost(openapi_spec={ }) ``` +## WebSocket protocol (`invocations_ws`) + +The same `InvocationAgentServerHost` object also exposes a WebSocket transport at `/invocations_ws`. Container authors do not install or import a second package — registering an `@app.ws_handler` is the only step. A multi-protocol agent shares one host, one session, and one process. + +### Quick start + +```python +from azure.ai.agentserver.invocations import InvocationAgentServerHost +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.websockets import WebSocket + +app = InvocationAgentServerHost() + + +@app.invoke_handler # POST /invocations (HTTP) +async def invoke(request: Request) -> Response: + payload = await request.json() + return JSONResponse({"echo": payload}) + + +@app.ws_handler # /invocations_ws (WebSocket) +async def ws(websocket: WebSocket) -> None: + async for message in websocket.iter_text(): + await websocket.send_text(message) + + +app.run() +``` + +### What the SDK does for `@app.ws_handler` + +- Registers `/invocations_ws` on the same Starlette host as `/invocations` and `/readiness`. +- Calls `await websocket.accept()` before invoking your handler. +- Runs WebSocket Ping/Pong keep-alive in the background — default 30 s, configurable via `InvocationAgentServerHost(ws_ping_interval=...)`. Set `ws_ping_interval=0` to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic. +- Closes the connection cleanly on handler return (close code `1000`) or maps an uncaught handler exception to close code `1011`. +- Emits a structured close-event log line carrying `ws.session_id`, `ws.close_code`, and `ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end. +- Inherits `/readiness`, OpenTelemetry export, graceful shutdown, and the `x-platform-server` identity header from `azure-ai-agentserver-core`. + +### Handler signature + +The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`. The full WebSocket API — `iter_text`, `iter_bytes`, `iter_json`, `send_text`, `send_bytes`, `send_json`, `close`, `headers`, `query_params`, `client`, `state` — is available, so application protocols on top of `invocations_ws` are entirely under your control. + +[starlette-ws]: https://www.starlette.io/websockets/ + +### Reference: configuration + +| Constructor argument | Default | Description | +|---|---|---| +| `ws_ping_interval` | `30.0` (seconds) | WebSocket protocol Ping interval. `0` disables keep-alive. Negative or non-finite values are rejected. | + ## Troubleshooting ### Reporting issues @@ -196,6 +252,8 @@ Visit the [Samples](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/ |---|---| | [simple_invoke_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/simple_invoke_agent/) | Minimal synchronous request-response | | [async_invoke_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/async_invoke_agent/) | Long-running operations with polling and cancellation | +| [ws_invoke_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/) | Combined `POST /invocations` (HTTP) and `/invocations_ws` (WebSocket) host | +| [ws_bidirectional_streaming_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/) | Full-duplex `/invocations_ws` agent: server-pushed heartbeats + concurrent token streams + mid-flight cancel | ## Contributing diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py index 62f8600a44bd..611772000efc 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py @@ -19,3 +19,31 @@ class InvocationConstants: ATTR_SPAN_SESSION_ID = "azure.ai.agentserver.invocations.session_id" ATTR_SPAN_ERROR_CODE = "azure.ai.agentserver.invocations.error.code" ATTR_SPAN_ERROR_MESSAGE = "azure.ai.agentserver.invocations.error.message" + + +class InvocationsWSConstants: + """invocations_ws (WebSocket) protocol constants. + + Route, span attribute keys, and ping/pong defaults for the + WebSocket endpoint hosted alongside the HTTP invocations protocol. + """ + + # Route + ROUTE_PATH = "/invocations_ws" + + # Default WebSocket Ping interval in seconds. + # Azure APIM and Azure Load Balancer drop idle WebSocket connections + # after ~4 minutes; 30 s gives a comfortable safety margin. + DEFAULT_PING_INTERVAL_S = 30.0 + + # Close codes (RFC 6455) + CLOSE_NORMAL = 1000 # handler returned cleanly + CLOSE_INTERNAL_ERROR = 1011 # handler raised an unhandled exception + CLOSE_SERVICE_RESTART = 1012 # graceful shutdown drained the connection + + # Span attribute keys + ATTR_SPAN_SESSION_ID = "ws.session_id" + ATTR_SPAN_CLOSE_CODE = "ws.close_code" + ATTR_SPAN_DURATION_MS = "ws.duration_ms" + ATTR_SPAN_ERROR_CODE = "azure.ai.agentserver.invocations_ws.error.code" + ATTR_SPAN_ERROR_MESSAGE = "azure.ai.agentserver.invocations_ws.error.message" 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 bf3120974fa0..8c706e55e5b4 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 @@ -31,6 +31,7 @@ ) from ._constants import InvocationConstants +from ._invocation_ws import _WSHandlerMixin logger = logging.getLogger("azure.ai.agentserver") @@ -93,13 +94,18 @@ def _sanitize_id(value: str, fallback: str) -> str: return value -class InvocationAgentServerHost(AgentServerHost): +class InvocationAgentServerHost(_WSHandlerMixin, AgentServerHost): """Invocation protocol host for Azure AI Hosted Agents. A :class:`~azure.ai.agentserver.core.AgentServerHost` subclass that adds the invocation protocol endpoints. Use the decorator methods to wire handler functions to the endpoints. + The same host object also exposes the ``invocations_ws`` (WebSocket) + transport at :data:`/invocations_ws` — register a handler with the + :meth:`ws_handler` decorator. Multi-protocol agents share a single + host, session, and process. + For multi-protocol agents, compose via cooperative inheritance:: class MyHost(InvocationAgentServerHost, ResponsesAgentServerHost): @@ -108,18 +114,29 @@ class MyHost(InvocationAgentServerHost, ResponsesAgentServerHost): Usage:: from azure.ai.agentserver.invocations import InvocationAgentServerHost + from starlette.websockets import WebSocket app = InvocationAgentServerHost() - @app.invoke_handler + @app.invoke_handler # POST /invocations async def handle(request): return JSONResponse({"ok": True}) + @app.ws_handler # /invocations_ws + async def ws(websocket: WebSocket) -> None: + async for message in websocket.iter_text(): + await websocket.send_text(message) + app.run() :param openapi_spec: Optional OpenAPI spec dict. When provided, the spec is served at ``GET /invocations/docs/openapi.json``. :type openapi_spec: Optional[dict[str, Any]] + :param ws_ping_interval: Seconds between WebSocket protocol Ping frames + on ``/invocations_ws``. ``None`` (default) selects 30 s; ``0`` + disables keep-alive. Configured on the underlying Hypercorn + server so the framing is opcode 0x9 / 0xA, not application JSON. + :type ws_ping_interval: Optional[float] """ _INSTRUMENTATION_SCOPE = "Azure.AI.AgentServer.Invocations" @@ -128,6 +145,7 @@ def __init__( self, *, openapi_spec: Optional[dict[str, Any]] = None, + ws_ping_interval: Optional[float] = None, **kwargs: Any, ) -> None: self._invoke_fn: Optional[Callable] = None @@ -135,8 +153,11 @@ def __init__( self._cancel_invocation_fn: Optional[Callable] = None self._openapi_spec = openapi_spec + # Initialise WS handler slots (raises ValueError on a bad interval). + self._init_ws_state(ws_ping_interval) + # Build invocation routes and pass to parent via routes kwarg - invocation_routes = [ + invocation_routes: list[Any] = [ Route( "/invocations/docs/openapi.json", self._get_openapi_spec_endpoint, @@ -161,6 +182,7 @@ def __init__( methods=["POST"], name="cancel_invocation", ), + self._build_ws_route(self._ws_endpoint), ] # Merge with any routes from sibling mixins via cooperative init @@ -169,10 +191,53 @@ def __init__( # --- Invocations startup configuration logging --- logger.info( - "Invocations protocol: openapi_spec_configured=%s", + "Invocations protocol: openapi_spec_configured=%s, " + "ws_ping_interval=%s", self._openapi_spec is not None, + ( + "disabled" + if self._ws_ping_interval == 0 + else f"{self._ws_ping_interval}s" + ), ) + # ------------------------------------------------------------------ + # Hypercorn server config (WebSocket Ping/Pong keep-alive) + # ------------------------------------------------------------------ + + def _build_hypercorn_config(self, host: str, port: int) -> object: + """Extend the base Hypercorn config with the WebSocket Ping interval. + + Hypercorn sends WS protocol Ping frames every + ``websocket_ping_interval`` seconds on every active WebSocket + connection — exactly the keep-alive the ``invocations_ws`` spec + requires. ``ws_ping_interval=0`` leaves the default + ``None`` (disabled). + + :param host: Network interface to bind. + :type host: str + :param port: Port to bind. + :type port: int + :return: The configured Hypercorn config. + :rtype: hypercorn.config.Config + """ + config = super()._build_hypercorn_config(host, port) + if self._ws_ping_interval and self._ws_ping_interval > 0: + try: + # ``websocket_ping_interval`` is a float-or-None on + # Hypercorn ≥0.14; assigning a positive float enables + # protocol-level Ping frames. + config.websocket_ping_interval = self._ws_ping_interval # type: ignore[attr-defined] + except Exception: # pylint: disable=broad-exception-caught + # Hypercorn <0.14 does not support per-server WS ping — + # leave the default and warn so operators can upgrade. + logger.warning( + "Hypercorn does not support websocket_ping_interval; " + "WebSocket keep-alive will be best-effort.", + exc_info=True, + ) + return config + # ------------------------------------------------------------------ # Handler decorators # ------------------------------------------------------------------ diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py new file mode 100644 index 000000000000..38b531ebe385 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py @@ -0,0 +1,467 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""invocations_ws (WebSocket) protocol support for ``InvocationAgentServerHost``. + +Implements the ``@app.ws_handler`` decorator and the ``/invocations_ws`` +route described in the ``invocations_ws`` protocol spec. The SDK wraps +the user handler with: + +* ``await websocket.accept()`` before the handler runs; +* WebSocket protocol-level Ping/Pong keep-alive (default 30 s) so idle + connections survive Azure APIM / Azure Load Balancer's ~4-minute idle + timeout; +* a clean close on handler return (code 1000) or a 1011 close on uncaught + handler exceptions; +* a structured close-event log line and OTel span attributes carrying + ``ws.session_id``, ``ws.close_code``, and ``ws.duration_ms``. +""" +from __future__ import annotations + +import inspect +import logging +import math +import time +import uuid +from collections.abc import Awaitable, Callable +from typing import Any, Optional + +from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState + +from azure.ai.agentserver.core import ( # pylint: disable=no-name-in-module + AgentServerHost, + end_span, + record_error, +) + +from ._constants import InvocationsWSConstants + +logger = logging.getLogger("azure.ai.agentserver") + + +WSHandler = Callable[[WebSocket], Awaitable[None]] + + +def _safe_set_attrs(span: Any, attrs: dict[str, Any]) -> None: + """Best-effort ``span.set_attribute`` for a batch of attributes. + + :param span: The OTel span (or ``None`` when tracing is disabled). + :type span: any + :param attrs: Mapping of attribute keys to values. + :type attrs: dict[str, any] + """ + if span is None: + return + try: + for key, value in attrs.items(): + span.set_attribute(key, value) + except Exception: # pylint: disable=broad-exception-caught + logger.debug("Failed to set WS span attributes: %s", list(attrs.keys()), exc_info=True) + + +class _WSHandlerMixin(AgentServerHost): + """Trait class that adds the ``@app.ws_handler`` decorator and ``/invocations_ws`` route. + + Inherits from :class:`~azure.ai.agentserver.core.AgentServerHost` so that + cooperative ``super().__init__`` calls and host attribute access + (``self.config``, ``self.request_span``) are resolved statically as well + as at runtime. Designed to be mixed into :class:`InvocationAgentServerHost` + so the same host object exposes both ``POST /invocations`` (HTTP) and + ``/invocations_ws`` (WebSocket) on the same Starlette application. + + Subclasses must append the route returned by :meth:`_build_ws_route` to + their ``routes`` list before calling ``super().__init__``. + """ + + # Slots populated by __init__. + _ws_fn: Optional[WSHandler] + _ws_ping_interval: float + + def _init_ws_state(self, ws_ping_interval: Optional[float]) -> None: + """Initialize WS handler slots. + + :param ws_ping_interval: Seconds between WS protocol Ping frames. + ``None`` selects the default (30 s); ``0`` disables keep-alive. + :type ws_ping_interval: Optional[float] + :raises ValueError: If *ws_ping_interval* is negative or non-finite. + """ + self._ws_fn = None + if ws_ping_interval is None: + resolved = InvocationsWSConstants.DEFAULT_PING_INTERVAL_S + else: + try: + resolved = float(ws_ping_interval) + except (TypeError, ValueError) as exc: + raise ValueError( + f"ws_ping_interval must be a number, got {ws_ping_interval!r}" + ) from exc + # Reject negative / NaN / inf — those are programming errors that + # would silently mis-configure the keep-alive. + if math.isnan(resolved) or math.isinf(resolved) or resolved < 0.0: + raise ValueError( + f"ws_ping_interval must be a non-negative finite number, " + f"got {ws_ping_interval!r}" + ) + self._ws_ping_interval = resolved + + # ------------------------------------------------------------------ + # Public configuration accessor + # ------------------------------------------------------------------ + + @property + def ws_ping_interval(self) -> float: + """Configured WebSocket Ping interval in seconds (``0`` = disabled). + + :return: The configured interval, or ``0`` when keep-alive is disabled. + :rtype: float + """ + return self._ws_ping_interval + + # ------------------------------------------------------------------ + # Decorator + # ------------------------------------------------------------------ + + def ws_handler(self, fn: WSHandler) -> WSHandler: + """Register an async function as the ``/invocations_ws`` handler. + + The SDK calls ``await websocket.accept()`` before invoking *fn* and + cleanly closes the connection on return (code 1000) or maps an + uncaught exception to close code 1011. + + Usage:: + + from starlette.websockets import WebSocket + + @app.ws_handler + async def handle(websocket: WebSocket) -> None: + async for msg in websocket.iter_text(): + await websocket.send_text(msg) + + :param fn: An async function accepting a Starlette + :class:`~starlette.websockets.WebSocket` and returning ``None``. + :type fn: Callable[[WebSocket], Awaitable[None]] + :return: The original function (unmodified). + :rtype: Callable[[WebSocket], Awaitable[None]] + :raises TypeError: If *fn* is not an ``async def`` function. + """ + if not inspect.iscoroutinefunction(fn): + raise TypeError( + f"ws_handler expects an async function, got {type(fn).__name__}. " + "Use 'async def' to define your handler." + ) + self._ws_fn = fn + return fn + + # ------------------------------------------------------------------ + # Route factory (for cooperative __init__) + # ------------------------------------------------------------------ + + @staticmethod + def _build_ws_route(endpoint: Callable[[WebSocket], Awaitable[None]]) -> Any: + """Return a :class:`~starlette.routing.WebSocketRoute` for ``/invocations_ws``. + + Imported lazily to avoid hard-coding the route type in the public + module body and keep the import surface symmetric with the HTTP + ``Route`` import in :mod:`._invocation`. + + :param endpoint: The async endpoint to wire to the route. + :type endpoint: Callable[[WebSocket], Awaitable[None]] + :return: A configured ``WebSocketRoute`` instance. + :rtype: ~starlette.routing.WebSocketRoute + """ + from starlette.routing import WebSocketRoute # pylint: disable=import-outside-toplevel + + return WebSocketRoute( + InvocationsWSConstants.ROUTE_PATH, + endpoint, + name="invocations_ws", + ) + + # ------------------------------------------------------------------ + # Endpoint + # ------------------------------------------------------------------ + + async def _ws_endpoint(self, websocket: WebSocket) -> None: + """ASGI endpoint for ``/invocations_ws``. + + Wraps the user-registered handler with: accept, span lifecycle, + graceful close on success, 1011 close on failure, and a structured + close event log + span attributes. + + :param websocket: The incoming Starlette WebSocket. + :type websocket: ~starlette.websockets.WebSocket + """ + # Per-connection identifiers. Session ID is generated server-side; + # the spec carries it in the close-event metric so an operator can + # correlate logs/metrics for a given long-lived connection. + session_id = str(uuid.uuid4()) + start_ns = time.monotonic_ns() + + # Open the OTel span before accepting so any tracecontext header + # the client sent is honoured for parenting child spans inside the + # user handler. ``end_on_exit=False`` so we can attach the close + # code / duration before ending. + span_ctx = self.request_span( + websocket.headers, + session_id, + "websocket_session", + operation_name="websocket_session", + session_id=session_id, + end_on_exit=False, + ) + otel_span = span_ctx.__enter__() + _safe_set_attrs(otel_span, {InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id}) + + if self._ws_fn is None: + await self._reject_no_handler(websocket, span_ctx, otel_span, session_id, start_ns) + return + + # Accept the upgrade *before* invoking the user handler — per spec. + try: + await websocket.accept() + except Exception as exc: # pylint: disable=broad-exception-caught + await self._finalize_session( + websocket=None, + span_ctx=span_ctx, + otel_span=otel_span, + session_id=session_id, + start_ns=start_ns, + close_code=InvocationsWSConstants.CLOSE_INTERNAL_ERROR, + handler_exc=exc, + error_code="accept_failed", + ) + logger.error( + "WebSocket accept failed for session %s: %s", + session_id, exc, exc_info=True, + ) + return + + close_code, handler_exc = await self._invoke_user_handler(websocket, session_id) + await self._finalize_session( + websocket=websocket, + span_ctx=span_ctx, + otel_span=otel_span, + session_id=session_id, + start_ns=start_ns, + close_code=close_code, + handler_exc=handler_exc, + error_code="internal_error" if handler_exc is not None else None, + ) + + async def _invoke_user_handler( + self, websocket: WebSocket, session_id: str, + ) -> tuple[int, Optional[BaseException]]: + """Run the registered user handler and classify the outcome. + + :param websocket: The accepted WebSocket to pass to the handler. + :type websocket: ~starlette.websockets.WebSocket + :param session_id: Per-connection session ID for diagnostic logs. + :type session_id: str + :return: ``(close_code, exception_or_None)``. ``close_code`` is the + RFC 6455 code that should be sent to the client; ``exception`` + is set only for an *unhandled* exception (so the caller can map + it to span error events and a 1011 close). + :rtype: tuple[int, Optional[BaseException]] + """ + assert self._ws_fn is not None # checked by caller + try: + await self._ws_fn(websocket) + return InvocationsWSConstants.CLOSE_NORMAL, None + except WebSocketDisconnect as exc: + # Client (or proxy) closed first — surface their code, not 1011. + return ( + int(exc.code) if exc.code else InvocationsWSConstants.CLOSE_NORMAL, + None, + ) + except Exception as exc: # pylint: disable=broad-exception-caught + logger.error( + "WebSocket handler raised for session %s: %s", + session_id, exc, exc_info=True, + ) + return InvocationsWSConstants.CLOSE_INTERNAL_ERROR, exc + + async def _reject_no_handler( + self, + websocket: WebSocket, + span_ctx: Any, + otel_span: Any, + session_id: str, + start_ns: int, + ) -> None: + """Refuse a WS upgrade when no ``@ws_handler`` is registered. + + :param websocket: The pending WebSocket awaiting upgrade. + :type websocket: ~starlette.websockets.WebSocket + :param span_ctx: The active ``request_span`` context manager. + :type span_ctx: any + :param otel_span: The current OTel span (or ``None``). + :type otel_span: any + :param session_id: Per-connection session ID. + :type session_id: str + :param start_ns: ``time.monotonic_ns()`` at connection start. + :type start_ns: int + """ + logger.error( + "WebSocket connection on %s rejected: no @ws_handler registered", + InvocationsWSConstants.ROUTE_PATH, + ) + duration_ms = (time.monotonic_ns() - start_ns) // 1_000_000 + self._emit_close_event( + otel_span, + session_id, + InvocationsWSConstants.CLOSE_INTERNAL_ERROR, + duration_ms, + error_code="not_implemented", + error_message="No ws_handler registered.", + ) + try: + span_ctx.__exit__(None, None, None) + finally: + end_span(otel_span) + await websocket.close( + code=InvocationsWSConstants.CLOSE_INTERNAL_ERROR, + reason="No ws_handler registered", + ) + + async def _finalize_session( + self, + *, + websocket: Optional[WebSocket], + span_ctx: Any, + otel_span: Any, + session_id: str, + start_ns: int, + close_code: int, + handler_exc: Optional[BaseException], + error_code: Optional[str], + ) -> None: + """Close the WS (best-effort), emit metrics, and end the span. + + Called from both the success path and the accept-failure path. + + :keyword websocket: The connected WebSocket, or ``None`` when the + ASGI ``accept`` itself failed (no socket to close). + :paramtype websocket: Optional[~starlette.websockets.WebSocket] + :keyword span_ctx: The active ``request_span`` context manager. + :paramtype span_ctx: any + :keyword otel_span: The current OTel span (or ``None`` when tracing is off). + :paramtype otel_span: any + :keyword session_id: Per-connection session ID. + :paramtype session_id: str + :keyword start_ns: ``time.monotonic_ns()`` at connection start. + :paramtype start_ns: int + :keyword close_code: The RFC 6455 code to report to the client. + :paramtype close_code: int + :keyword handler_exc: Unhandled exception raised by the user handler, + or ``None`` for a clean close. + :paramtype handler_exc: Optional[BaseException] + :keyword error_code: Short error tag for span / log; ``None`` for success. + :paramtype error_code: Optional[str] + """ + duration_ms = (time.monotonic_ns() - start_ns) // 1_000_000 + + # Best-effort clean close: only send a close frame if the + # application hasn't already done so (e.g. the user handler + # may have called ``websocket.close`` itself, or the client + # may have disconnected). + if ( + websocket is not None + and websocket.application_state != WebSocketState.DISCONNECTED + ): + reason = ( + "Internal server error" + if close_code == InvocationsWSConstants.CLOSE_INTERNAL_ERROR + else "" + ) + try: + await websocket.close(code=close_code, reason=reason) + except Exception: # pylint: disable=broad-exception-caught + # Connection already gone — nothing to recover here. + logger.debug( + "Error closing WebSocket session %s", session_id, exc_info=True, + ) + + self._emit_close_event( + otel_span, + session_id, + close_code, + duration_ms, + error_code=error_code, + error_message=str(handler_exc) if handler_exc is not None else None, + ) + + if handler_exc is not None: + try: + record_error(otel_span, handler_exc) + finally: + try: + span_ctx.__exit__( + type(handler_exc), handler_exc, handler_exc.__traceback__, + ) + finally: + end_span(otel_span) + else: + try: + span_ctx.__exit__(None, None, None) + finally: + end_span(otel_span) + + # ------------------------------------------------------------------ + # Close event + # ------------------------------------------------------------------ + + @staticmethod + def _emit_close_event( + otel_span: Any, + session_id: str, + close_code: int, + duration_ms: int, + *, + error_code: Optional[str] = None, + error_message: Optional[str] = None, + ) -> None: + """Record close-event span attributes and emit a structured log line. + + The log record carries the ``ws.session_id``, ``ws.close_code``, + and ``ws.duration_ms`` fields listed in the spec via the standard + ``logging`` ``extra`` dict — a structured-logging formatter or an + OTel logging bridge can pick them up directly without having to + parse the message. + + :param otel_span: The connection span (or ``None`` when tracing is off). + :type otel_span: any + :param session_id: Per-connection session ID. + :type session_id: str + :param close_code: The RFC 6455 close code reported to the client. + :type close_code: int + :param duration_ms: Connection duration in milliseconds (monotonic). + :type duration_ms: int + :keyword error_code: Optional short error tag for span attribute. + :paramtype error_code: Optional[str] + :keyword error_message: Optional human-readable error message for + span attribute (NOT included in the log line — exception details + are logged separately by ``logger.error(..., exc_info=True)``). + :paramtype error_message: Optional[str] + """ + attrs: dict[str, Any] = { + InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id, + InvocationsWSConstants.ATTR_SPAN_CLOSE_CODE: close_code, + InvocationsWSConstants.ATTR_SPAN_DURATION_MS: duration_ms, + } + if error_code: + attrs[InvocationsWSConstants.ATTR_SPAN_ERROR_CODE] = error_code + if error_message: + attrs[InvocationsWSConstants.ATTR_SPAN_ERROR_MESSAGE] = error_message + _safe_set_attrs(otel_span, attrs) + + logger.info( + "invocations_ws connection closed: session_id=%s code=%s duration_ms=%s", + session_id, + close_code, + duration_ms, + extra={ + "ws.session_id": session_id, + "ws.close_code": close_code, + "ws.duration_ms": duration_ms, + }, + ) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/requirements.txt b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/requirements.txt new file mode 100644 index 000000000000..bc5cf4644e14 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/requirements.txt @@ -0,0 +1 @@ +azure-ai-agentserver-invocations diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py new file mode 100644 index 000000000000..9a66525d3c03 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py @@ -0,0 +1,250 @@ +"""Bidirectional streaming agent over the ``invocations_ws`` (WebSocket) protocol. + +Unlike the request/reply echo in :mod:`samples.ws_invoke_agent`, this sample +exercises the *full-duplex* nature of WebSocket: the server and the client +send and receive on the same socket **concurrently and independently**. + +The handler runs two groups of coroutines in parallel: + +1. ``_reader`` — consumes inbound JSON frames (prompts and control + messages) and dispatches them. Multiple prompts + may arrive while previous ones are still streaming. +2. ``_stream_tokens`` — one task per prompt; streams generated tokens back + to the client at its own pace. Multiple generations + can run in parallel and run *independently* of any + inbound traffic — the defining property of full- + duplex. ``cancel`` control messages interrupt them + mid-flight. + +.. note:: + + Connection keep-alive is **not** an application concern: the SDK + configures Hypercorn to send WebSocket protocol-level Ping frames + (opcode 0x9) every ``ws_ping_interval`` seconds (default 30 s). That + is enough to survive Azure APIM / Azure Load Balancer's ~4-minute idle + timeout without your handler having to push any application-level + heartbeat messages of its own. + +Wire protocol (JSON over text frames) +------------------------------------- + +Inbound (client -> server):: + + {"type": "prompt", "id": "p1", "text": "..."} + {"type": "cancel", "id": "p1"} + {"type": "bye"} # graceful shutdown + +Outbound (server -> client):: + + {"type": "ready"} # sent on connect + {"type": "token", "id": "p1", "token": "..."} + {"type": "done", "id": "p1"} + {"type": "cancelled", "id": "p1"} + {"type": "error", "id": "p1", "message": "..."} + +Run it:: + + python ws_bidirectional_streaming_agent.py + +Drive it with the ``websockets`` CLI; the server keeps streaming tokens +for each prompt while you type the next one:: + + python -m websockets ws://localhost:8088/invocations_ws + > {"type": "prompt", "id": "p1", "text": "Tell me a story"} + > {"type": "prompt", "id": "p2", "text": "And another"} + > {"type": "cancel", "id": "p1"} + > {"type": "bye"} +""" +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +from collections.abc import AsyncGenerator + +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.websockets import WebSocket, WebSocketDisconnect + +from azure.ai.agentserver.invocations import InvocationAgentServerHost + + +logger = logging.getLogger("ws_bidirectional_streaming_agent") + +app = InvocationAgentServerHost() + + +# Simulated tokens — in production these would come from a model. +_SIMULATED_TOKENS = [ + "Once", " upon", " a", " time", ",", " in", " a", " land", + " of", " full", "-", "duplex", " sockets", ",", " a", " server", + " and", " a", " client", " spoke", " at", " the", " same", " time", ".", +] + +_TOKEN_DELAY_S = 0.2 + + +# --------------------------------------------------------------------------- +# HTTP — same host, kept for parity with the rest of the samples. +# --------------------------------------------------------------------------- + +@app.invoke_handler # POST /invocations +async def handle_invoke(request: Request) -> Response: + """Echo the JSON payload back over HTTP.""" + payload = await request.json() + return JSONResponse({"echo": payload}) + + +# --------------------------------------------------------------------------- +# WebSocket — true bidirectional streaming. +# --------------------------------------------------------------------------- + +async def _generate_tokens(text: str) -> AsyncGenerator[str, None]: + """Yield simulated tokens with a small per-token delay. + + Replace this with a real streaming model call (e.g. Azure OpenAI) in + production. + + :param text: The user prompt (unused in this demo). + :type text: str + :return: An async generator of token strings. + :rtype: AsyncGenerator[str, None] + """ + del text # demo: ignore prompt content + for token in _SIMULATED_TOKENS: + await asyncio.sleep(_TOKEN_DELAY_S) + yield token + + +async def _stream_tokens( + websocket: WebSocket, prompt_id: str, text: str, +) -> None: + """Stream tokens for one prompt; cancellable via ``asyncio.CancelledError``. + + :param websocket: The accepted WebSocket. + :type websocket: ~starlette.websockets.WebSocket + :param prompt_id: Caller-supplied prompt identifier (echoed in events). + :type prompt_id: str + :param text: The user prompt to "generate" against. + :type text: str + """ + try: + async for token in _generate_tokens(text): + await websocket.send_json( + {"type": "token", "id": prompt_id, "token": token}, + ) + await websocket.send_json({"type": "done", "id": prompt_id}) + except asyncio.CancelledError: + # Best-effort: tell the client we honoured their cancel. Suppress + # any send error (the socket may already be closed) and re-raise + # so the caller observes the cancellation. + with contextlib.suppress(Exception): + await websocket.send_json({"type": "cancelled", "id": prompt_id}) + raise + + +async def _reader( + websocket: WebSocket, + in_flight: "dict[str, asyncio.Task[None]]", +) -> None: + """Consume inbound frames and dispatch prompt / cancel / bye control messages. + + Returns (instead of raising) on a ``bye`` message or a clean client + disconnect. Returning lets the caller cancel any in-flight generation + tasks and end the connection. + + :param websocket: The accepted WebSocket. + :type websocket: ~starlette.websockets.WebSocket + :param in_flight: Map of ``prompt_id`` -> generation task, used to + honour ``cancel`` messages. + :type in_flight: dict[str, asyncio.Task[None]] + """ + try: + async for raw in websocket.iter_text(): + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await websocket.send_json( + {"type": "error", "message": "invalid JSON"}, + ) + continue + + msg_type = msg.get("type") + + if msg_type == "prompt": + prompt_id = str(msg.get("id", "")) + text = str(msg.get("text", "")) + if not prompt_id: + await websocket.send_json( + {"type": "error", "message": "prompt requires 'id'"}, + ) + continue + # Schedule the generation as an independent task so it + # runs in parallel with the reader (and any other + # in-flight generations). + task = asyncio.create_task( + _stream_tokens(websocket, prompt_id, text), + name=f"stream-{prompt_id}", + ) + in_flight[prompt_id] = task + task.add_done_callback( + lambda _t, k=prompt_id: in_flight.pop(k, None), + ) + + elif msg_type == "cancel": + prompt_id = str(msg.get("id", "")) + task = in_flight.get(prompt_id) + if task is not None and not task.done(): + task.cancel() + + elif msg_type == "bye": + return + + else: + await websocket.send_json( + {"type": "error", "message": f"unknown type: {msg_type!r}"}, + ) + except WebSocketDisconnect: + # Client closed first — let the caller unwind normally. + return + + +async def _cancel_and_wait(tasks: "list[asyncio.Task[None]]") -> None: + """Cancel every task in *tasks* and wait for them to actually finish. + + :param tasks: Tasks to cancel; already-done tasks are ignored. + :type tasks: list[asyncio.Task[None]] + """ + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + +@app.ws_handler # /invocations_ws +async def handle_ws(websocket: WebSocket) -> None: + """Run the reader and per-prompt generation tasks concurrently. + + The SDK has already accepted the connection by the time this function + runs, and it is sending WS protocol-level Ping frames on its own + schedule (see ``ws_ping_interval``). When this coroutine returns the + SDK closes the socket with code ``1000``; if it raises, the SDK maps + the exception to ``1011``. + """ + await websocket.send_json({"type": "ready"}) + + in_flight: "dict[str, asyncio.Task[None]]" = {} + + try: + await _reader(websocket, in_flight) + except WebSocketDisconnect: + # Client went away mid-read — fall through to cleanup. + logger.info("client disconnected during streaming") + finally: + await _cancel_and_wait(list(in_flight.values())) + + +if __name__ == "__main__": + app.run() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/requirements.txt b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/requirements.txt new file mode 100644 index 000000000000..bc5cf4644e14 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/requirements.txt @@ -0,0 +1 @@ +azure-ai-agentserver-invocations diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py new file mode 100644 index 000000000000..2f36c48a3b20 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py @@ -0,0 +1,58 @@ +"""Echo agent over the ``invocations_ws`` (WebSocket) protocol. + +Exposes the same host on: + +* ``POST /invocations`` — HTTP, JSON request/response; +* ``/invocations_ws`` — WebSocket, full-duplex streaming; +* ``GET /readiness`` — readiness probe inherited from + ``azure-ai-agentserver-core``. + +Usage:: + + # Start the agent + python ws_invoke_agent.py + + # HTTP turn + curl -X POST http://localhost:8088/invocations \\ + -H "Content-Type: application/json" \\ + -d '{"name": "Alice"}' + # -> {"echo": {"name": "Alice"}} + + # WebSocket turn (with the `websockets` client library) + python -m websockets ws://localhost:8088/invocations_ws + # > hello + # < hello +""" +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.websockets import WebSocket + +from azure.ai.agentserver.invocations import InvocationAgentServerHost + + +app = InvocationAgentServerHost() + + +@app.invoke_handler # POST /invocations +async def handle_invoke(request: Request) -> Response: + """Echo the JSON payload back over HTTP.""" + payload = await request.json() + return JSONResponse({"echo": payload}) + + +@app.ws_handler # /invocations_ws +async def handle_ws(websocket: WebSocket) -> None: + """Echo every text frame back over the WebSocket connection. + + The SDK has already accepted the connection by the time this function + runs, sends WebSocket Ping frames every 30 s in the background to keep + Azure APIM / Azure Load Balancer from idling the socket out, and will + close the connection cleanly on return. An uncaught exception here + is mapped to RFC 6455 close code 1011. + """ + async for message in websocket.iter_text(): + await websocket.send_text(message) + + +if __name__ == "__main__": + app.run() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py index 8a3deb55c72f..24c3c1eb9500 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py @@ -199,6 +199,47 @@ async def handle(request: Request) -> Response: return app +# --------------------------------------------------------------------------- +# WebSocket factory functions and helpers +# --------------------------------------------------------------------------- + + +def _make_echo_ws_app(**kwargs: Any) -> InvocationAgentServerHost: + """Create an InvocationAgentServerHost whose ws handler echoes text frames.""" + from starlette.websockets import WebSocket + + app = InvocationAgentServerHost(**kwargs) + + @app.ws_handler + async def echo(websocket: WebSocket) -> None: + async for message in websocket.iter_text(): + await websocket.send_text(message) + + return app + + +def _make_failing_ws_app(**kwargs: Any) -> InvocationAgentServerHost: + """Create an InvocationAgentServerHost whose ws handler raises after one frame.""" + from starlette.websockets import WebSocket + + app = InvocationAgentServerHost(**kwargs) + + @app.ws_handler + async def boom(websocket: WebSocket) -> None: + await websocket.receive_text() + raise ValueError("boom") + + return app + + +def _records_with_ws_extras(records): + """Filter log records that carry the close-event ``ws.*`` extras.""" + return [ + r for r in records + if hasattr(r, "ws.session_id") and hasattr(r, "ws.close_code") + ] + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py new file mode 100644 index 000000000000..6c2c9723c301 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py @@ -0,0 +1,113 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for the structured close-event log line emitted by ``/invocations_ws``. + +Parity with :mod:`tests.test_request_id` — verifies the spec's required +fields (``ws.session_id``, ``ws.close_code``, ``ws.duration_ms``) appear +on every connection close, and that handler exception details are NOT +leaked into the structured payload. +""" +import logging + +import pytest +from starlette.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +from azure.ai.agentserver.invocations._constants import InvocationsWSConstants + +from conftest import ( + _make_echo_ws_app, + _make_failing_ws_app, + _records_with_ws_extras, +) + + +# --------------------------------------------------------------------------- +# Required fields on every close-event +# --------------------------------------------------------------------------- + +def test_ws_close_event_log_contains_required_fields(caplog): + """The close-event log line carries ws.session_id, ws.close_code, ws.duration_ms.""" + app = _make_echo_ws_app() + client = TestClient(app) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("ping") + assert ws.receive_text() == "ping" + + matches = _records_with_ws_extras(caplog.records) + assert matches, "expected a structured close-event log record" + rec = matches[-1] + + session_id = getattr(rec, "ws.session_id") + close_code = getattr(rec, "ws.close_code") + duration_ms = getattr(rec, "ws.duration_ms") + + assert isinstance(session_id, str) and session_id # generated UUID + assert close_code == InvocationsWSConstants.CLOSE_NORMAL + assert isinstance(duration_ms, int) + assert duration_ms >= 0 + + +def test_ws_close_event_duration_is_non_negative(caplog): + """``ws.duration_ms`` is a non-negative integer derived from a monotonic clock.""" + app = _make_echo_ws_app() + client = TestClient(app) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("ping") + ws.receive_text() + + matches = _records_with_ws_extras(caplog.records) + assert matches + duration_ms = getattr(matches[-1], "ws.duration_ms") + assert isinstance(duration_ms, int) + assert duration_ms >= 0 + + +# --------------------------------------------------------------------------- +# Close codes on the close-event +# --------------------------------------------------------------------------- + +def test_ws_close_event_on_handler_exception_records_1011(caplog): + """Handler raising → close-event log records ws.close_code = 1011.""" + app = _make_failing_ws_app() + client = TestClient(app) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("trigger") + ws.receive_text() + + matches = _records_with_ws_extras(caplog.records) + assert matches + assert getattr(matches[-1], "ws.close_code") == 1011 + + +# --------------------------------------------------------------------------- +# Exception details are NOT leaked into the structured payload +# (parity with test_error_hides_details_by_default) +# --------------------------------------------------------------------------- + +def test_ws_close_event_log_does_not_leak_exception_message(caplog): + """The close-event log line does NOT carry the handler exception text.""" + app = _make_failing_ws_app() + client = TestClient(app) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("trigger") + ws.receive_text() + + matches = _records_with_ws_extras(caplog.records) + assert matches + rec = matches[-1] + # The structured close-event log line carries only the safe ws.* fields. + assert not hasattr(rec, "ws.error_message") + # And the message itself does not embed the raw exception text. + assert "boom" not in rec.getMessage() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py new file mode 100644 index 000000000000..b447ff8d1090 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py @@ -0,0 +1,92 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for decorator-based handler registration on the ``ws_handler`` API. + +Parity with :mod:`tests.test_decorator_pattern` — covers the ``@ws_handler`` +decorator's behaviour, storage, and validation rules. +""" +import pytest +from starlette.testclient import TestClient +from starlette.websockets import WebSocket, WebSocketDisconnect + +from azure.ai.agentserver.invocations import InvocationAgentServerHost +from azure.ai.agentserver.invocations._constants import InvocationsWSConstants + + +# --------------------------------------------------------------------------- +# Decorator validation +# --------------------------------------------------------------------------- + +def test_ws_handler_rejects_sync_function(): + """``@app.ws_handler`` must be applied to ``async def`` callables.""" + app = InvocationAgentServerHost() + + with pytest.raises(TypeError, match="async function"): + @app.ws_handler # type: ignore[arg-type] + def sync_handler(websocket): # noqa: ARG001 + pass + + +def test_ws_handler_returns_function_unchanged(): + """The decorator must return the original function unmodified.""" + app = InvocationAgentServerHost() + + async def handler(websocket: WebSocket) -> None: + await websocket.accept() + await websocket.close() + + result = app.ws_handler(handler) + assert result is handler + + +# --------------------------------------------------------------------------- +# Decorator state — slot storage / defaults / re-registration +# --------------------------------------------------------------------------- + +def test_ws_handler_stores_function(): + """``@app.ws_handler`` stores the registered function on the host.""" + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket: WebSocket) -> None: + await websocket.send_text("ok") + + assert app._ws_fn is handler # noqa: SLF001 + + +def test_ws_handler_default_is_none(): + """Without ``@ws_handler`` the slot stays ``None``.""" + app = InvocationAgentServerHost() + assert app._ws_fn is None # noqa: SLF001 + + +def test_ws_handler_last_registration_wins(): + """Re-applying ``@ws_handler`` replaces the previous function.""" + app = InvocationAgentServerHost() + + @app.ws_handler + async def first(websocket: WebSocket) -> None: # noqa: ARG001 + return + + @app.ws_handler + async def second(websocket: WebSocket) -> None: # noqa: ARG001 + return + + assert app._ws_fn is second # noqa: SLF001 + + +# --------------------------------------------------------------------------- +# Missing handler behaviour (parity with test_missing_invoke_handler_returns_501) +# --------------------------------------------------------------------------- + +def test_ws_with_no_handler_registered_closes_with_1011(): + """If no @ws_handler is registered, the SDK closes with 1011.""" + app = InvocationAgentServerHost() + client = TestClient(app) + + with pytest.raises(WebSocketDisconnect) as excinfo: + with client.websocket_connect("/invocations_ws"): + pass + + assert excinfo.value.code == InvocationsWSConstants.CLOSE_INTERNAL_ERROR diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py new file mode 100644 index 000000000000..a25f1249ee8b --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py @@ -0,0 +1,121 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Edge-case tests for ``/invocations_ws``. + +Parity with :mod:`tests.test_edge_cases` — covers client-initiated +disconnects, handler-managed close, custom close codes, and empty +connections. +""" +import logging + +import pytest +from starlette.testclient import TestClient +from starlette.websockets import WebSocket, WebSocketDisconnect + +from azure.ai.agentserver.invocations import InvocationAgentServerHost +from azure.ai.agentserver.invocations._constants import InvocationsWSConstants + +from conftest import _make_echo_ws_app, _records_with_ws_extras + + +# --------------------------------------------------------------------------- +# Client-initiated disconnect (clean) +# --------------------------------------------------------------------------- + +def test_ws_client_disconnect_does_not_log_as_error(caplog): + """A client-initiated disconnect is a normal close, not a 1011 error.""" + app = _make_echo_ws_app() + client = TestClient(app) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hello") + ws.receive_text() + # __exit__ sends websocket.disconnect — the SDK should treat + # this as normal, not raise from the handler. + + error_records = [r for r in caplog.records if r.levelno >= logging.ERROR] + # No ERROR-level records should be emitted for a clean client disconnect. + assert not error_records, [r.getMessage() for r in error_records] + + +# --------------------------------------------------------------------------- +# Client-initiated close with a custom (non-1000) code +# --------------------------------------------------------------------------- + +def test_ws_client_initiated_close_with_custom_code_is_reported(caplog): + """When the client closes with a non-1000 code, the server surfaces the client's code (not 1011).""" + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket: WebSocket) -> None: + # Use receive_text directly so WebSocketDisconnect propagates with + # the client's close code — ``iter_text`` swallows it. + while True: + await websocket.receive_text() + + client = TestClient(app) + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("warm-up") + ws.close(code=4001) + + matches = _records_with_ws_extras(caplog.records) + assert matches + rec = matches[-1] + close_code = getattr(rec, "ws.close_code") + # Server surfaces the client's code — NOT 1011. + assert close_code == 4001 + # No ERROR-level records — a client disconnect is normal. + assert not [r for r in caplog.records if r.levelno >= logging.ERROR] + + +# --------------------------------------------------------------------------- +# Handler-managed close +# --------------------------------------------------------------------------- + +def test_ws_handler_explicit_close_does_not_double_close(caplog): + """If the handler closes the WS itself, the SDK does NOT attempt a second close.""" + app = InvocationAgentServerHost() + closes_observed: list[int] = [] + + async def handler(websocket: WebSocket) -> None: + # Hand-roll close so we can verify the SDK skips re-closing. + await websocket.send_text("bye") + await websocket.close(code=1000) + closes_observed.append(1) + + app.ws_handler(handler) + + client = TestClient(app) + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/invocations_ws") as ws: + assert ws.receive_text() == "bye" + ws.receive_text() # forces the close frame to surface + + # Handler ran to completion. + assert closes_observed == [1] + # Close-event log line still emitted (with the handler's code). + matches = _records_with_ws_extras(caplog.records) + assert matches + assert getattr(matches[-1], "ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL + + +# --------------------------------------------------------------------------- +# Empty connection (no frames sent) +# --------------------------------------------------------------------------- + +def test_ws_empty_connection_closes_normally(caplog): + """A connection that immediately disconnects closes cleanly (1000).""" + app = _make_echo_ws_app() + client = TestClient(app) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with client.websocket_connect("/invocations_ws"): + pass # Disconnect without sending anything. + + matches = _records_with_ws_extras(caplog.records) + assert matches + assert getattr(matches[-1], "ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py new file mode 100644 index 000000000000..9eb988b73713 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py @@ -0,0 +1,143 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for the WebSocket request/response lifecycle on ``/invocations_ws``. + +Parity with :mod:`tests.test_invoke` — covers the happy path (accept, echo, +clean close), the unhappy path (handler exception → 1011), and the +WebSocket-only bidirectional streaming property. +""" +import asyncio + +import pytest +from starlette.testclient import TestClient +from starlette.websockets import WebSocket, WebSocketDisconnect + +from azure.ai.agentserver.invocations import InvocationAgentServerHost +from azure.ai.agentserver.invocations._constants import InvocationsWSConstants + +from conftest import _make_echo_ws_app, _make_failing_ws_app + + +# --------------------------------------------------------------------------- +# Accept happens automatically +# --------------------------------------------------------------------------- + +def test_ws_sdk_accepts_connection_before_handler_runs(): + """The SDK calls ``websocket.accept()`` before invoking the user handler. + + The handler in this test never calls ``accept`` itself; the connection + must still be usable, which proves the SDK accepted it on the user's + behalf. + """ + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket: WebSocket) -> None: + await websocket.send_text("ready") + + client = TestClient(app) + with client.websocket_connect("/invocations_ws") as ws: + assert ws.receive_text() == "ready" + + +# --------------------------------------------------------------------------- +# Echo round-trip +# --------------------------------------------------------------------------- + +def test_ws_echo_round_trip(): + """End-to-end: send a frame, receive it echoed back.""" + app = _make_echo_ws_app() + client = TestClient(app) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hello") + assert ws.receive_text() == "hello" + ws.send_text("world") + assert ws.receive_text() == "world" + + +def test_ws_multiple_frames_per_connection(): + """Multiple frames flow through one connection (parity with streaming chunks).""" + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket: WebSocket) -> None: + for i in range(5): + await websocket.send_text(f"chunk-{i}") + + client = TestClient(app) + with client.websocket_connect("/invocations_ws") as ws: + for i in range(5): + assert ws.receive_text() == f"chunk-{i}" + + +# --------------------------------------------------------------------------- +# Close codes +# --------------------------------------------------------------------------- + +def test_ws_handler_exception_maps_to_close_code_1011(): + """Uncaught handler exceptions must surface as RFC 6455 close code 1011.""" + app = _make_failing_ws_app() + client = TestClient(app) + + with pytest.raises(WebSocketDisconnect) as excinfo: + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("trigger") + # Server closes after the handler raises; receiving forces + # the close frame to surface as WebSocketDisconnect. + ws.receive_text() + + assert excinfo.value.code == InvocationsWSConstants.CLOSE_INTERNAL_ERROR + assert excinfo.value.code == 1011 + + +def test_ws_clean_return_uses_close_code_1000(): + """A handler that returns normally yields a 1000 (normal) close code.""" + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket: WebSocket) -> None: + # Receive once, then return — SDK closes cleanly. + await websocket.receive_text() + + client = TestClient(app) + with pytest.raises(WebSocketDisconnect) as excinfo: + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("done") + ws.receive_text() # Forces the close to surface. + + assert excinfo.value.code == InvocationsWSConstants.CLOSE_NORMAL + + +# --------------------------------------------------------------------------- +# Bidirectional streaming (WebSocket-only feature) +# --------------------------------------------------------------------------- + +def test_ws_bidirectional_concurrent_send_receive(): + """Reader and writer coroutines run concurrently on the same socket.""" + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket: WebSocket) -> None: + async def reader(): + async for msg in websocket.iter_text(): + if msg == "bye": + return + + async def writer(): + for i in range(3): + await websocket.send_text(f"server-{i}") + await asyncio.sleep(0) # yield to reader + + # Run reader + writer in parallel — the defining property of WS. + await asyncio.gather(reader(), writer()) + + client = TestClient(app) + with client.websocket_connect("/invocations_ws") as ws: + # Mix inbound and outbound frames to prove they're not lock-stepped. + ws.send_text("client-a") + assert ws.receive_text() == "server-0" + ws.send_text("client-b") + assert ws.receive_text() == "server-1" + assert ws.receive_text() == "server-2" + ws.send_text("bye") diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py new file mode 100644 index 000000000000..e05b228c13a8 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py @@ -0,0 +1,77 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for multi-modality payloads sent through ``/invocations_ws``. + +Parity with :mod:`tests.test_multimodal_protocol` — covers binary frames, +text frames (unicode + large), and JSON frames over WebSocket. +""" +from starlette.testclient import TestClient +from starlette.websockets import WebSocket + +from azure.ai.agentserver.invocations import InvocationAgentServerHost + +from conftest import _make_echo_ws_app + + +# --------------------------------------------------------------------------- +# Binary frames +# --------------------------------------------------------------------------- + +def test_ws_binary_frame_round_trip(): + """Binary frames round-trip without corruption (parity with test_binary_payload).""" + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket: WebSocket) -> None: + data = await websocket.receive_bytes() + await websocket.send_bytes(data) + + client = TestClient(app) + payload = bytes(range(256)) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_bytes(payload) + assert ws.receive_bytes() == payload + + +# --------------------------------------------------------------------------- +# Text frames — unicode and large payloads +# --------------------------------------------------------------------------- + +def test_ws_unicode_text_round_trip(): + """Unicode text frames are preserved (parity with test_unicode_payload).""" + app = _make_echo_ws_app() + client = TestClient(app) + text = "Hello, 世界! 🌍 — naïve façade" + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text(text) + assert ws.receive_text() == text + + +def test_ws_large_text_frame_round_trip(): + """A ~1 MB text frame round-trips successfully (parity with test_large_payload).""" + app = _make_echo_ws_app() + client = TestClient(app) + payload = "x" * (1024 * 1024) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text(payload) + assert ws.receive_text() == payload + + +# --------------------------------------------------------------------------- +# JSON frames (``send_json`` / ``receive_json``) +# --------------------------------------------------------------------------- + +def test_ws_json_frame_round_trip(): + """``send_json`` / ``receive_json`` round-trip JSON payloads.""" + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket: WebSocket) -> None: + msg = await websocket.receive_json() + await websocket.send_json({"echo": msg}) + + client = TestClient(app) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_json({"hello": "world", "n": 42}) + assert ws.receive_json() == {"echo": {"hello": "world", "n": 42}} diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py new file mode 100644 index 000000000000..82567c6df622 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py @@ -0,0 +1,103 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for the ``ws_ping_interval`` configuration knob. + +Parity with :mod:`tests.test_request_limits` / configuration tests in +:mod:`tests.test_graceful_shutdown` — validates the Hypercorn ``ws_ping_interval`` +parameter parsing and wiring. +""" +import pytest + +from azure.ai.agentserver.invocations import InvocationAgentServerHost +from azure.ai.agentserver.invocations._constants import InvocationsWSConstants + + +# --------------------------------------------------------------------------- +# Default / accepted values +# --------------------------------------------------------------------------- + +def test_ws_ping_interval_default_is_30_seconds(): + """Default ping interval matches the spec (30 s).""" + app = InvocationAgentServerHost() + assert app.ws_ping_interval == InvocationsWSConstants.DEFAULT_PING_INTERVAL_S + assert app.ws_ping_interval == 30.0 + + +def test_ws_ping_interval_custom_value(): + """``ws_ping_interval`` is honoured.""" + app = InvocationAgentServerHost(ws_ping_interval=15) + assert app.ws_ping_interval == 15.0 + + +def test_ws_ping_interval_zero_disables_keepalive(): + """``ws_ping_interval=0`` disables WS-level keep-alive.""" + app = InvocationAgentServerHost(ws_ping_interval=0) + assert app.ws_ping_interval == 0.0 + + +def test_ws_ping_interval_float_value_accepted(): + """Fractional intervals are coerced to ``float``.""" + app = InvocationAgentServerHost(ws_ping_interval=12.5) + assert app.ws_ping_interval == 12.5 + config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 + assert getattr(config, "websocket_ping_interval", None) == 12.5 + + +# --------------------------------------------------------------------------- +# Rejected values (validation) +# --------------------------------------------------------------------------- + +def test_ws_ping_interval_negative_rejected(): + """Negative intervals are programming errors.""" + with pytest.raises(ValueError, match="non-negative"): + InvocationAgentServerHost(ws_ping_interval=-1) + + +def test_ws_ping_interval_nan_rejected(): + """``ws_ping_interval=nan`` is a programming error.""" + with pytest.raises(ValueError, match="non-negative"): + InvocationAgentServerHost(ws_ping_interval=float("nan")) + + +def test_ws_ping_interval_inf_rejected(): + """``ws_ping_interval=inf`` is a programming error.""" + with pytest.raises(ValueError, match="non-negative"): + InvocationAgentServerHost(ws_ping_interval=float("inf")) + + +def test_ws_ping_interval_non_numeric_rejected(): + """Strings or non-numeric values surface as ``ValueError``.""" + with pytest.raises(ValueError, match="must be a number"): + InvocationAgentServerHost(ws_ping_interval="thirty") # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# Hypercorn config wiring +# --------------------------------------------------------------------------- + +def test_ws_ping_interval_propagates_to_hypercorn_config(): + """The configured interval lands on the Hypercorn server config.""" + app = InvocationAgentServerHost(ws_ping_interval=20) + config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 + # Hypercorn ≥0.14 exposes this attribute. + assert getattr(config, "websocket_ping_interval", None) == 20.0 + + +def test_ws_ping_interval_zero_does_not_override_hypercorn_default(): + """Zero leaves Hypercorn's default (None = disabled) intact.""" + app = InvocationAgentServerHost(ws_ping_interval=0) + config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 + # Hypercorn default is None — our wiring leaves it unset for 0. + assert getattr(config, "websocket_ping_interval", None) is None + + +# --------------------------------------------------------------------------- +# Property surface +# --------------------------------------------------------------------------- + +def test_ws_ping_interval_property_is_read_only(): + """``ws_ping_interval`` is exposed only as a property (no setter).""" + app = InvocationAgentServerHost(ws_ping_interval=20) + with pytest.raises(AttributeError): + app.ws_ping_interval = 10 # type: ignore[misc] diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py new file mode 100644 index 000000000000..4972b76275c7 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py @@ -0,0 +1,95 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for basic server route registration with the ``/invocations_ws`` route. + +Parity with :mod:`tests.test_server_routes` — covers route registration, +coexistence with the HTTP routes, and rejection of mismatched paths. +""" +import pytest +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.testclient import TestClient +from starlette.websockets import WebSocket, WebSocketDisconnect + +from azure.ai.agentserver.invocations import InvocationAgentServerHost + +from conftest import _make_echo_ws_app + + +# --------------------------------------------------------------------------- +# Route registration +# --------------------------------------------------------------------------- + +def test_ws_route_is_registered(): + """The /invocations_ws route exists alongside the HTTP routes.""" + app = InvocationAgentServerHost() + paths = [getattr(r, "path", None) for r in app.routes] + assert "/invocations_ws" in paths + assert "/invocations" in paths + assert "/readiness" in paths + + +def test_readiness_still_works_with_ws_registered(): + """Adding the WS route doesn't break /readiness.""" + app = _make_echo_ws_app() + client = TestClient(app) + resp = client.get("/readiness") + assert resp.status_code == 200 + # x-platform-server header still applied via core middleware + assert "x-platform-server" in resp.headers + + +# --------------------------------------------------------------------------- +# Coexistence with HTTP /invocations +# --------------------------------------------------------------------------- + +def test_http_and_ws_share_same_host(): + """Both transports work on the same app — single session, single process.""" + app = InvocationAgentServerHost() + + @app.invoke_handler + async def http_handle(request: Request) -> Response: + body = await request.json() + return JSONResponse({"http": body}) + + @app.ws_handler + async def ws_handle(websocket: WebSocket) -> None: + async for msg in websocket.iter_text(): + await websocket.send_text(f"ws:{msg}") + + client = TestClient(app) + + # HTTP route still works + resp = client.post("/invocations", json={"hello": "world"}) + assert resp.status_code == 200 + assert resp.json() == {"http": {"hello": "world"}} + + # WS route works on the same host + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + assert ws.receive_text() == "ws:hi" + + +# --------------------------------------------------------------------------- +# Mismatched URLs (parity with test_unknown_route_returns_404) +# --------------------------------------------------------------------------- + +def test_ws_upgrade_on_http_path_fails(): + """A WS upgrade to ``/invocations`` (the HTTP route) is rejected.""" + app = _make_echo_ws_app() + client = TestClient(app) + # /invocations is a Route, not a WebSocketRoute — TestClient surfaces + # this as an immediate WebSocketDisconnect rather than a connect. + with pytest.raises((WebSocketDisconnect, Exception)): # noqa: PT011 + with client.websocket_connect("/invocations"): + pass + + +def test_ws_unknown_path_fails(): + """An unknown WS path is rejected.""" + app = _make_echo_ws_app() + client = TestClient(app) + with pytest.raises((WebSocketDisconnect, Exception)): # noqa: PT011 + with client.websocket_connect("/nonexistent"): + pass diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py new file mode 100644 index 000000000000..43107ed71b7e --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py @@ -0,0 +1,80 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for WebSocket session ID resolution. + +Parity with :mod:`tests.test_session_id` — covers per-connection UUID +generation and the (deliberate) lack of query-param overrides. The WS +endpoint always generates a fresh server-side session ID. +""" +import logging +import uuid + +from starlette.testclient import TestClient + +from conftest import _make_echo_ws_app, _records_with_ws_extras + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _session_ids_from_records(records): + """Pull ``ws.session_id`` from each structured close-event record.""" + return [getattr(r, "ws.session_id") for r in _records_with_ws_extras(records)] + + +# --------------------------------------------------------------------------- +# Session ID is a server-generated UUID +# --------------------------------------------------------------------------- + +def test_ws_session_id_is_uuid(caplog): + """The per-connection session ID is a valid UUID string.""" + app = _make_echo_ws_app() + client = TestClient(app) + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + session_ids = _session_ids_from_records(caplog.records) + assert session_ids + # Each must parse as a UUID — the WS endpoint generates one server-side. + parsed = uuid.UUID(session_ids[-1]) + assert str(parsed) == session_ids[-1] + + +def test_ws_session_id_is_unique_per_connection(caplog): + """Each WS connection gets its own session ID (parity with test_invoke_unique_invocation_ids).""" + app = _make_echo_ws_app() + client = TestClient(app) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + for _ in range(5): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("ping") + ws.receive_text() + + session_ids = _session_ids_from_records(caplog.records) + assert len(session_ids) == 5 + assert len(set(session_ids)) == 5 + + +def test_ws_session_id_ignores_query_param(caplog): + """Unlike HTTP, the WS endpoint always generates its own session ID.""" + app = _make_echo_ws_app() + client = TestClient(app) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with client.websocket_connect( + "/invocations_ws?agent_session_id=my-custom-session", + ) as ws: + ws.send_text("x") + ws.receive_text() + + session_ids = _session_ids_from_records(caplog.records) + assert session_ids + # The client's value MUST NOT leak into the session ID — the WS spec + # generates a fresh server-side UUID per connection. + assert session_ids[-1] != "my-custom-session" + uuid.UUID(session_ids[-1]) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py new file mode 100644 index 000000000000..f78b0c815c92 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py @@ -0,0 +1,470 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for OpenTelemetry tracing on the ``invocations_ws`` protocol. + +These mirror :mod:`tests.test_tracing` (HTTP /invocations) but target the +WebSocket transport: span creation, GenAI attributes, traceparent +propagation, close-code recording, and error recording. +""" +import os +import uuid +from unittest.mock import patch + +import pytest +from starlette.testclient import TestClient +from starlette.websockets import WebSocket, WebSocketDisconnect + +from azure.ai.agentserver.invocations import InvocationAgentServerHost + + +# --------------------------------------------------------------------------- +# Lazy OTel setup with in-memory exporter (mirrors tests/test_tracing.py) +# --------------------------------------------------------------------------- +try: + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + _HAS_OTEL = True +except ImportError: + _HAS_OTEL = False + +_MODULE_PROVIDER = None +_MODULE_EXPORTER = None +_MODULE_SETUP_DONE = False + +pytestmark = pytest.mark.skipif(not _HAS_OTEL, reason="opentelemetry not installed") + + +@pytest.fixture(autouse=True) +def _clear_spans(): + """Set up the OTel provider on first use, then clear spans before each test.""" + global _MODULE_PROVIDER, _MODULE_EXPORTER, _MODULE_SETUP_DONE + if _HAS_OTEL and not _MODULE_SETUP_DONE: + _existing = trace.get_tracer_provider() + if hasattr(_existing, "add_span_processor"): + _MODULE_PROVIDER = _existing + else: + _MODULE_PROVIDER = SdkTracerProvider() + trace.set_tracer_provider(_MODULE_PROVIDER) + _MODULE_EXPORTER = InMemorySpanExporter() + _MODULE_PROVIDER.add_span_processor(SimpleSpanProcessor(_MODULE_EXPORTER)) + _MODULE_SETUP_DONE = True + if _MODULE_EXPORTER: + _MODULE_EXPORTER.clear() + + +def _get_spans(): + """Return all captured spans.""" + if _MODULE_EXPORTER: + return _MODULE_EXPORTER.get_finished_spans() + return [] + + +def _ws_session_spans(spans): + """Filter spans created by the ``/invocations_ws`` route.""" + return [s for s in spans if "websocket_session" in s.name] + + +# --------------------------------------------------------------------------- +# Helpers: tracing-enabled servers +# --------------------------------------------------------------------------- + +def _make_tracing_ws_server(**kwargs) -> InvocationAgentServerHost: + """Build an echo WS host with tracing enabled.""" + with patch.dict( + os.environ, + {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, + ): + with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): + server = InvocationAgentServerHost(**kwargs) + + @server.ws_handler + async def handler(websocket: WebSocket) -> None: + async for msg in websocket.iter_text(): + await websocket.send_text(msg) + + return server + + +def _make_failing_tracing_ws_server(**kwargs) -> InvocationAgentServerHost: + """Build a WS host whose handler raises after one received frame.""" + with patch.dict( + os.environ, + {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, + ): + with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): + server = InvocationAgentServerHost(**kwargs) + + @server.ws_handler + async def handler(websocket: WebSocket) -> None: + await websocket.receive_text() + raise ValueError("ws tracing error test") + + return server + + +def _make_tracing_ws_no_handler(**kwargs) -> InvocationAgentServerHost: + """Build a tracing-enabled host with NO ws_handler registered.""" + with patch.dict( + os.environ, + {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, + ): + with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): + return InvocationAgentServerHost(**kwargs) + + +# --------------------------------------------------------------------------- +# Span creation +# --------------------------------------------------------------------------- + +def test_ws_tracing_creates_websocket_session_span(): + """A WS connection creates a span named ``websocket_session``.""" + server = _make_tracing_ws_server() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hello") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert len(spans) == 1 + assert spans[0].name.startswith("websocket_session") + + +def test_ws_tracing_disabled_by_default_still_creates_span(): + """Even without an exporter, the global tracer creates the span.""" + if _MODULE_EXPORTER: + _MODULE_EXPORTER.clear() + + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket: WebSocket) -> None: + await websocket.send_text("ready") + + client = TestClient(app) + with client.websocket_connect("/invocations_ws") as ws: + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert len(spans) >= 1 + + +# --------------------------------------------------------------------------- +# GenAI attributes (parity with test_genai_attributes_on_invoke_span) +# --------------------------------------------------------------------------- + +def test_ws_span_has_genai_attributes(): + """WS span carries the standard GenAI / service.name attributes.""" + server = _make_tracing_ws_server() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert spans + attrs = dict(spans[0].attributes) + assert attrs.get("gen_ai.system") == "azure.ai.agentserver" + assert attrs.get("gen_ai.provider.name") == "AzureAI Hosted Agents" + assert attrs.get("service.name") == "azure.ai.agentserver" + + +def test_ws_span_operation_name_attribute(): + """``gen_ai.operation.name`` on the WS span is ``websocket_session``.""" + server = _make_tracing_ws_server() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + attrs = dict(spans[0].attributes) + assert attrs.get("gen_ai.operation.name") == "websocket_session" + + +# --------------------------------------------------------------------------- +# Session ID attribute (parity with test_session_id_in_conversation_id) +# --------------------------------------------------------------------------- + +def test_ws_session_id_attribute_set_on_span(): + """The auto-generated session_id is stored on the span under both keys.""" + server = _make_tracing_ws_server() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + attrs = dict(spans[0].attributes) + # WS-specific key + ws_session = attrs.get("ws.session_id") + # General session key (shared with HTTP) + ms_session = attrs.get("microsoft.session.id") + assert ws_session and isinstance(ws_session, str) + uuid.UUID(ws_session) + assert ms_session == ws_session + + +def test_ws_session_id_is_unique_per_span(): + """Two WS connections produce two distinct session IDs on their spans.""" + server = _make_tracing_ws_server() + client = TestClient(server) + for _ in range(2): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("x") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert len(spans) == 2 + session_ids = {dict(s.attributes).get("ws.session_id") for s in spans} + assert len(session_ids) == 2 + + +# --------------------------------------------------------------------------- +# Close-code / duration attributes (WS-specific) +# --------------------------------------------------------------------------- + +def test_ws_span_records_close_code_1000_on_clean_return(): + """A handler that returns normally records ws.close_code = 1000.""" + server = _make_tracing_ws_server() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + attrs = dict(spans[0].attributes) + assert attrs.get("ws.close_code") == 1000 + + +def test_ws_span_records_close_code_1011_on_handler_exception(): + """A failing handler records ws.close_code = 1011.""" + server = _make_failing_tracing_ws_server() + client = TestClient(server) + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("trigger") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + attrs = dict(spans[0].attributes) + assert attrs.get("ws.close_code") == 1011 + + +def test_ws_span_records_duration_ms(): + """``ws.duration_ms`` is set as a non-negative integer.""" + server = _make_tracing_ws_server() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + attrs = dict(spans[0].attributes) + duration_ms = attrs.get("ws.duration_ms") + assert isinstance(duration_ms, int) + assert duration_ms >= 0 + + +def test_ws_span_records_client_close_code(): + """A client closing with a custom code lands on the span as ws.close_code.""" + server_with_recv = _make_tracing_ws_server() + + # Override the handler so receive_text raises with the client's code + # (iter_text swallows WebSocketDisconnect — see _invocation_ws docs). + @server_with_recv.ws_handler + async def handler(websocket: WebSocket) -> None: + while True: + await websocket.receive_text() + + client = TestClient(server_with_recv) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("ping") + ws.close(code=4010) + + spans = _ws_session_spans(_get_spans()) + attrs = dict(spans[0].attributes) + assert attrs.get("ws.close_code") == 4010 + + +# --------------------------------------------------------------------------- +# Error recording on span (parity with test_invoke_error_records_exception) +# --------------------------------------------------------------------------- + +def test_ws_handler_exception_records_exception_on_span(): + """When the handler raises, the span has ERROR status and exception event.""" + server = _make_failing_tracing_ws_server() + client = TestClient(server) + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("trigger") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert spans + span = spans[0] + assert span.status.status_code.name == "ERROR" + # Exception is recorded as an event by record_error. + event_names = [e.name for e in span.events] + assert "exception" in event_names + + +def test_ws_no_handler_records_error_on_span(): + """A connection refused for lack of @ws_handler still produces an error-tagged span.""" + app = _make_tracing_ws_no_handler() + client = TestClient(app) + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/invocations_ws"): + pass + + spans = _ws_session_spans(_get_spans()) + assert spans + attrs = dict(spans[0].attributes) + # Close code is 1011 and the error_code/message attrs are populated. + assert attrs.get("ws.close_code") == 1011 + assert attrs.get("azure.ai.agentserver.invocations_ws.error.code") == "not_implemented" + assert "ws_handler" in str( + attrs.get("azure.ai.agentserver.invocations_ws.error.message", "") + ) + + +# --------------------------------------------------------------------------- +# Traceparent propagation (parity with test_traceparent_propagation) +# --------------------------------------------------------------------------- + +def test_ws_traceparent_propagation(): + """A client-supplied traceparent header parents the WS span.""" + server = _make_tracing_ws_server() + trace_id_hex = uuid.uuid4().hex + span_id_hex = uuid.uuid4().hex[:16] + traceparent = f"00-{trace_id_hex}-{span_id_hex}-01" + + client = TestClient(server) + with client.websocket_connect( + "/invocations_ws", + headers={"traceparent": traceparent}, + ) as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert spans + actual_trace_id = format(spans[0].context.trace_id, "032x") + assert actual_trace_id == trace_id_hex + + +# --------------------------------------------------------------------------- +# Agent name / version in span name (parity with test_agent_name_in_span_name) +# --------------------------------------------------------------------------- + +def test_ws_agent_name_and_version_in_span_name(): + """``FOUNDRY_AGENT_NAME``/``_VERSION`` appear in the WS span name.""" + with patch.dict( + os.environ, + { + "FOUNDRY_AGENT_NAME": "my-ws-agent", + "FOUNDRY_AGENT_VERSION": "3.1", + }, + ): + server = _make_tracing_ws_server() + + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert spans + name = spans[0].name + assert "my-ws-agent" in name + assert "3.1" in name + + +def test_ws_agent_name_only_in_span_name(): + """Agent name without version still appears in span name.""" + env_override = {"FOUNDRY_AGENT_NAME": "ws-solo"} + env_copy = os.environ.copy() + env_copy.pop("FOUNDRY_AGENT_VERSION", None) + env_copy.update(env_override) + with patch.dict(os.environ, env_copy, clear=True): + server = _make_tracing_ws_server() + + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert spans + assert "ws-solo" in spans[0].name + + +# --------------------------------------------------------------------------- +# Span kind / shape +# --------------------------------------------------------------------------- + +def test_ws_span_kind_is_server(): + """The WS span is created with ``kind=SERVER``.""" + server = _make_tracing_ws_server() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert spans + # SpanKind.SERVER == 2; compare by name to avoid pinning the enum value. + assert spans[0].kind.name == "SERVER" + + +def test_ws_span_has_no_parent_when_no_traceparent(): + """Without traceparent the WS span is a root span (parent is None).""" + server = _make_tracing_ws_server() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + assert spans + assert spans[0].parent is None + + +# --------------------------------------------------------------------------- +# Coexistence: invoke + ws spans on the same host +# --------------------------------------------------------------------------- + +def test_ws_and_invoke_spans_coexist(): + """A host serving both HTTP /invocations and /invocations_ws produces both spans.""" + from starlette.requests import Request + from starlette.responses import Response + + with patch.dict( + os.environ, + {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, + ): + with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): + server = InvocationAgentServerHost() + + @server.invoke_handler + async def http_handler(request: Request) -> Response: + return Response(content=b"ok") + + @server.ws_handler + async def ws_handler(websocket: WebSocket) -> None: + async for msg in websocket.iter_text(): + await websocket.send_text(msg) + + client = TestClient(server) + client.post("/invocations", content=b"data") + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _get_spans() + invoke = [s for s in spans if "invoke_agent" in s.name] + ws_spans = _ws_session_spans(spans) + assert invoke and ws_spans From 8df0573516a88eaf7bb21acc4410d30b1a7df3a8 Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Thu, 14 May 2026 16:41:15 +0800 Subject: [PATCH 02/13] [agentserver-invocations] Address PR review comments 18-52 for invocations_ws (#46885) * Address PR review comments #18-#52 for invocations_ws - README: add WS close-codes reference table (1000/1011/4xxx) - Tests: tighten pytest.raises types; rename ping-interval-zero test; add disconnect-code 0/None coverage; add Hypercorn ping-wiring tests; enhance double-close test to assert call_count==1; add exception-count belt-and-suspenders; mark large-frame test slow; drop create=True from _setup_distro_export mocks; replace patch.dict(clear=True) with monkeypatch; use long-form ws.* attribute name in close-event check; clarify clean-close test with comment. - Samples: add 'pip install websockets' to docstrings; correct outdated 'default 30s' ping-interval claim; add structured error frame on stream exceptions; await cancelled task with timeout; use BaseException in courtesy-cancel send; cross-reference azure-ai-inference/projects streaming examples; rename text -> _text; document lambda default-arg capture; use azure.ai.agentserver package logger. * Fix CI: restore create=True for forward-compat with published core; rephrase 'unconfigured' (cspell) * Fix pylint: line lengths and import position in _invocation_ws.py --------- Co-authored-by: Yulin Li --- .../CHANGELOG.md | 17 +- .../README.md | 18 +- .../ai/agentserver/invocations/_constants.py | 19 +- .../ai/agentserver/invocations/_invocation.py | 35 ++- .../agentserver/invocations/_invocation_ws.py | 253 ++++++++++-------- .../ws_bidirectional_streaming_agent.py | 58 ++-- .../ws_invoke_agent/ws_invoke_agent.py | 10 +- .../tests/conftest.py | 3 +- .../tests/test_ws_close_event.py | 65 ++++- .../tests/test_ws_decorator_pattern.py | 45 +++- .../tests/test_ws_edge_cases.py | 19 +- .../tests/test_ws_invoke.py | 4 + .../tests/test_ws_multimodal_protocol.py | 3 + .../tests/test_ws_ping_interval.py | 85 +++++- .../tests/test_ws_server_routes.py | 19 +- .../tests/test_ws_session_id.py | 27 +- .../tests/test_ws_tracing.py | 41 ++- 17 files changed, 481 insertions(+), 240 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index 7745922fba25..2410d32a0844 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -4,19 +4,10 @@ ### Features Added -- `invocations_ws` (WebSocket) protocol support on `InvocationAgentServerHost`. - Register a handler with the new `@app.ws_handler` decorator to host a - full-duplex WebSocket endpoint at `/invocations_ws` on the same host that - serves `POST /invocations`. The SDK calls `await websocket.accept()` before - invoking the handler, runs WebSocket Ping/Pong keep-alive in the background - (default 30 s; configurable via the new `ws_ping_interval` constructor - argument), closes the connection cleanly on handler return, and maps - uncaught exceptions to RFC 6455 close code `1011`. Each connection emits a - structured close-event log line carrying `ws.session_id`, `ws.close_code`, - and `ws.duration_ms`, and the same fields are recorded as OpenTelemetry - span attributes. `/readiness`, OTEL export, graceful shutdown, and the - `x-platform-server` identity header continue to be inherited from - `azure-ai-agentserver-core`. +- WebSocket protocol support — `InvocationAgentServerHost` now hosts `/invocations_ws` alongside `POST /invocations`. Register the handler with the new `@app.ws_handler` decorator. The route is registered lazily on first decoration, so hosts without a registered handler return HTTP 404. +- WebSocket Ping/Pong keep-alive — disabled by default; enable by passing `InvocationAgentServerHost(ws_ping_interval=)` or by setting the `WS_KEEPALIVE_INTERVAL` env var (auto-injected by AgentService into hosted-agent containers). The constructor argument takes precedence; `0` (or unset) disables keep-alive. Wired through to Hypercorn's `websocket_ping_interval`. +- WebSocket telemetry — structured close-event log line and OpenTelemetry span attributes `azure.ai.agentserver.invocations_ws.{session_id,close_code,duration_ms}`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. +- New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages). ### Breaking Changes diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index 63ec794bbc8f..acab609ed746 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md @@ -221,9 +221,9 @@ app.run() - Registers `/invocations_ws` on the same Starlette host as `/invocations` and `/readiness`. - Calls `await websocket.accept()` before invoking your handler. -- Runs WebSocket Ping/Pong keep-alive in the background — default 30 s, configurable via `InvocationAgentServerHost(ws_ping_interval=...)`. Set `ws_ping_interval=0` to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic. +- Runs WebSocket Ping/Pong keep-alive in the background — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers) or by passing `InvocationAgentServerHost(ws_ping_interval=...)`. Set the value to `0` (in any source) to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic. - Closes the connection cleanly on handler return (close code `1000`) or maps an uncaught handler exception to close code `1011`. -- Emits a structured close-event log line carrying `ws.session_id`, `ws.close_code`, and `ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end. +- Emits a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `azure.ai.agentserver.invocations_ws.close_code`, and `azure.ai.agentserver.invocations_ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end. - Inherits `/readiness`, OpenTelemetry export, graceful shutdown, and the `x-platform-server` identity header from `azure-ai-agentserver-core`. ### Handler signature @@ -236,7 +236,19 @@ The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`. | Constructor argument | Default | Description | |---|---|---| -| `ws_ping_interval` | `30.0` (seconds) | WebSocket protocol Ping interval. `0` disables keep-alive. Negative or non-finite values are rejected. | +| `ws_ping_interval` | `0` (disabled) | WebSocket protocol Ping interval in seconds. `0` disables keep-alive. Negative or non-finite values are rejected. When `None`, the SDK reads the `WS_KEEPALIVE_INTERVAL` env var before falling back to disabled. | + +| Environment variable | Default | Description | +|---|---|---| +| `WS_KEEPALIVE_INTERVAL` | unset (disabled) | Platform-injected override for the WebSocket Ping interval (seconds). `0` disables keep-alive. Ignored when `ws_ping_interval=` is set explicitly. | + +### Reference: close codes + +| Close code | Meaning | +|---|---| +| `1000` | Handler returned cleanly (normal close). | +| `1011` | Handler raised an unhandled exception (mapped by the SDK). | +| `4000`-`4999` | Application-defined codes (set by the handler via `await websocket.close(code=...)` — surfaced unchanged to the client). | ## Troubleshooting diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py index 611772000efc..43683652417c 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py @@ -31,19 +31,22 @@ class InvocationsWSConstants: # Route ROUTE_PATH = "/invocations_ws" - # Default WebSocket Ping interval in seconds. - # Azure APIM and Azure Load Balancer drop idle WebSocket connections - # after ~4 minutes; 30 s gives a comfortable safety margin. - DEFAULT_PING_INTERVAL_S = 30.0 + # Environment variable for platform-injected keep-alive override. + # AgentService injects this into every hosted-agent container so the + # platform can tune the WebSocket Ping cadence without redeploying + # the agent code. + # Resolution order: explicit ``ws_ping_interval=`` constructor arg → + # ``WS_KEEPALIVE_INTERVAL`` env var → disabled (``0``). + # ``0`` disables keep-alive. + ENV_WS_KEEPALIVE_INTERVAL = "WS_KEEPALIVE_INTERVAL" # Close codes (RFC 6455) CLOSE_NORMAL = 1000 # handler returned cleanly CLOSE_INTERNAL_ERROR = 1011 # handler raised an unhandled exception - CLOSE_SERVICE_RESTART = 1012 # graceful shutdown drained the connection # Span attribute keys - ATTR_SPAN_SESSION_ID = "ws.session_id" - ATTR_SPAN_CLOSE_CODE = "ws.close_code" - ATTR_SPAN_DURATION_MS = "ws.duration_ms" + ATTR_SPAN_SESSION_ID = "azure.ai.agentserver.invocations_ws.session_id" + ATTR_SPAN_CLOSE_CODE = "azure.ai.agentserver.invocations_ws.close_code" + ATTR_SPAN_DURATION_MS = "azure.ai.agentserver.invocations_ws.duration_ms" ATTR_SPAN_ERROR_CODE = "azure.ai.agentserver.invocations_ws.error.code" ATTR_SPAN_ERROR_MESSAGE = "azure.ai.agentserver.invocations_ws.error.message" 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 8c706e55e5b4..40657f494a15 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 @@ -133,9 +133,12 @@ async def ws(websocket: WebSocket) -> None: is served at ``GET /invocations/docs/openapi.json``. :type openapi_spec: Optional[dict[str, Any]] :param ws_ping_interval: Seconds between WebSocket protocol Ping frames - on ``/invocations_ws``. ``None`` (default) selects 30 s; ``0`` - disables keep-alive. Configured on the underlying Hypercorn - server so the framing is opcode 0x9 / 0xA, not application JSON. + on ``/invocations_ws``. ``None`` (default) reads the + ``WS_KEEPALIVE_INTERVAL`` environment variable (auto-injected by + AgentService into hosted-agent containers); when neither is set, + keep-alive is disabled. ``0`` disables keep-alive. Configured on + the underlying Hypercorn server so the framing is opcode 0x9 / + 0xA, not application JSON. :type ws_ping_interval: Optional[float] """ @@ -182,7 +185,6 @@ def __init__( methods=["POST"], name="cancel_invocation", ), - self._build_ws_route(self._ws_endpoint), ] # Merge with any routes from sibling mixins via cooperative init @@ -211,8 +213,11 @@ def _build_hypercorn_config(self, host: str, port: int) -> object: Hypercorn sends WS protocol Ping frames every ``websocket_ping_interval`` seconds on every active WebSocket connection — exactly the keep-alive the ``invocations_ws`` spec - requires. ``ws_ping_interval=0`` leaves the default - ``None`` (disabled). + requires. ``ws_ping_interval=0`` explicitly disables keep-alive + by setting the attribute to ``None``. + + ``azure-ai-agentserver-core`` pins ``hypercorn>=0.17.0``, which + is guaranteed to expose ``websocket_ping_interval``. :param host: Network interface to bind. :type host: str @@ -222,20 +227,10 @@ def _build_hypercorn_config(self, host: str, port: int) -> object: :rtype: hypercorn.config.Config """ config = super()._build_hypercorn_config(host, port) - if self._ws_ping_interval and self._ws_ping_interval > 0: - try: - # ``websocket_ping_interval`` is a float-or-None on - # Hypercorn ≥0.14; assigning a positive float enables - # protocol-level Ping frames. - config.websocket_ping_interval = self._ws_ping_interval # type: ignore[attr-defined] - except Exception: # pylint: disable=broad-exception-caught - # Hypercorn <0.14 does not support per-server WS ping — - # leave the default and warn so operators can upgrade. - logger.warning( - "Hypercorn does not support websocket_ping_interval; " - "WebSocket keep-alive will be best-effort.", - exc_info=True, - ) + # Positive float enables protocol-level Ping; None disables. + config.websocket_ping_interval = ( # type: ignore[attr-defined] + self._ws_ping_interval if self._ws_ping_interval > 0 else None + ) return config # ------------------------------------------------------------------ diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py index 38b531ebe385..86fedaca6efb 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py @@ -8,34 +8,47 @@ the user handler with: * ``await websocket.accept()`` before the handler runs; -* WebSocket protocol-level Ping/Pong keep-alive (default 30 s) so idle - connections survive Azure APIM / Azure Load Balancer's ~4-minute idle - timeout; +* WebSocket protocol-level Ping/Pong keep-alive (disabled by default; + enable via the ``WS_KEEPALIVE_INTERVAL`` env var or the + ``ws_ping_interval=`` constructor argument) so idle connections can + survive Azure APIM / Azure Load Balancer's ~4-minute idle timeout; * a clean close on handler return (code 1000) or a 1011 close on uncaught handler exceptions; * a structured close-event log line and OTel span attributes carrying - ``ws.session_id``, ``ws.close_code``, and ``ws.duration_ms``. + ``azure.ai.agentserver.invocations_ws.session_id``, + ``azure.ai.agentserver.invocations_ws.close_code``, and + ``azure.ai.agentserver.invocations_ws.duration_ms``. """ -from __future__ import annotations import inspect import logging import math +import os import time import uuid from collections.abc import Awaitable, Callable -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState from azure.ai.agentserver.core import ( # pylint: disable=no-name-in-module AgentServerHost, end_span, - record_error, ) from ._constants import InvocationsWSConstants +# Type-checking only base so the mixin reads as an ``AgentServerHost`` to +# mypy / pyright (resolves ``self.config``, ``self.request_span``, +# ``self.router``) without coupling the runtime hierarchy. At runtime the +# mixin is a plain ``object`` subclass — only the concrete +# ``InvocationAgentServerHost`` MRO actually inherits ``AgentServerHost``, +# which keeps the diamond out of the runtime class graph. +if TYPE_CHECKING: + _MixinBase = AgentServerHost +else: + _MixinBase = object + logger = logging.getLogger("azure.ai.agentserver") @@ -59,18 +72,17 @@ def _safe_set_attrs(span: Any, attrs: dict[str, Any]) -> None: logger.debug("Failed to set WS span attributes: %s", list(attrs.keys()), exc_info=True) -class _WSHandlerMixin(AgentServerHost): - """Trait class that adds the ``@app.ws_handler`` decorator and ``/invocations_ws`` route. - - Inherits from :class:`~azure.ai.agentserver.core.AgentServerHost` so that - cooperative ``super().__init__`` calls and host attribute access - (``self.config``, ``self.request_span``) are resolved statically as well - as at runtime. Designed to be mixed into :class:`InvocationAgentServerHost` - so the same host object exposes both ``POST /invocations`` (HTTP) and - ``/invocations_ws`` (WebSocket) on the same Starlette application. +class _WSHandlerMixin(_MixinBase): + """Pure mixin that adds the ``@app.ws_handler`` decorator and ``/invocations_ws`` route. - Subclasses must append the route returned by :meth:`_build_ws_route` to - their ``routes`` list before calling ``super().__init__``. + Designed to be mixed into a concrete + :class:`~azure.ai.agentserver.core.AgentServerHost` subclass (e.g. + :class:`InvocationAgentServerHost`) so the same host object exposes + both ``POST /invocations`` (HTTP) and ``/invocations_ws`` (WebSocket) + on the same Starlette application. At runtime the mixin is a plain + ``object`` subclass — host attributes (``self.config``, + ``self.request_span``, ``self.router``) are accessed via duck typing + and are typed only for the static checkers (see ``_MixinBase``). """ # Slots populated by __init__. @@ -80,14 +92,41 @@ class _WSHandlerMixin(AgentServerHost): def _init_ws_state(self, ws_ping_interval: Optional[float]) -> None: """Initialize WS handler slots. + Resolution order for the keep-alive interval: + + 1. Explicit *ws_ping_interval* constructor argument (when not ``None``). + 2. ``WS_KEEPALIVE_INTERVAL`` environment variable (auto-injected by + AgentService into every hosted-agent container). + 3. Disabled (``0``) — no Ping/Pong frames are sent. + + ``0`` (from any source) disables keep-alive entirely. + :param ws_ping_interval: Seconds between WS protocol Ping frames. - ``None`` selects the default (30 s); ``0`` disables keep-alive. + ``None`` selects the env var or the default. :type ws_ping_interval: Optional[float] - :raises ValueError: If *ws_ping_interval* is negative or non-finite. + :raises ValueError: If the resolved value is negative, non-finite, + or (for the env var) not parseable as a number. """ self._ws_fn = None if ws_ping_interval is None: - resolved = InvocationsWSConstants.DEFAULT_PING_INTERVAL_S + env_raw = os.environ.get(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL) + if env_raw is None or env_raw == "": + resolved = 0.0 + else: + try: + resolved = float(env_raw) + except ValueError as exc: + raise ValueError( + f"Invalid value for " + f"{InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL}: " + f"{env_raw!r} (expected a non-negative number)" + ) from exc + if math.isnan(resolved) or math.isinf(resolved) or resolved < 0.0: + raise ValueError( + f"Invalid value for " + f"{InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL}: " + f"{env_raw!r} (expected a non-negative finite number)" + ) else: try: resolved = float(ws_ping_interval) @@ -142,39 +181,58 @@ async def handle(websocket: WebSocket) -> None: :type fn: Callable[[WebSocket], Awaitable[None]] :return: The original function (unmodified). :rtype: Callable[[WebSocket], Awaitable[None]] - :raises TypeError: If *fn* is not an ``async def`` function. + :raises TypeError: If *fn* is not an ``async def`` function, or its + signature is not callable with exactly one positional argument. """ if not inspect.iscoroutinefunction(fn): raise TypeError( f"ws_handler expects an async function, got {type(fn).__name__}. " "Use 'async def' to define your handler." ) + # Validate signature at registration time (not at first request) so + # 0-arg / 2-required-arg coroutine mistakes surface at import. + try: + sig = inspect.signature(fn) + sig.bind(None) # one positional placeholder for the WebSocket + except TypeError as exc: + raise TypeError( + f"ws_handler must accept exactly one positional argument " + f"(the WebSocket); got {fn.__qualname__}{inspect.signature(fn)}" + ) from exc + if self._ws_fn is not None: + # Match the HTTP decorator's last-write-wins semantics, but log + # so misconfigured apps that double-register a handler aren't + # silently downgraded. + logger.warning( + "ws_handler overwriting previously registered handler %s with %s", + getattr(self._ws_fn, "__qualname__", repr(self._ws_fn)), + getattr(fn, "__qualname__", repr(fn)), + ) self._ws_fn = fn + # Register the route lazily on first decoration so hosts without a + # registered handler return HTTP 404 to a WebSocket upgrade rather than + # accepting and immediately closing with code 1011. + self._ensure_ws_route_registered() return fn - # ------------------------------------------------------------------ - # Route factory (for cooperative __init__) - # ------------------------------------------------------------------ + def _ensure_ws_route_registered(self) -> None: + """Append the ``/invocations_ws`` WebSocketRoute to the router (idempotent). - @staticmethod - def _build_ws_route(endpoint: Callable[[WebSocket], Awaitable[None]]) -> Any: - """Return a :class:`~starlette.routing.WebSocketRoute` for ``/invocations_ws``. - - Imported lazily to avoid hard-coding the route type in the public - module body and keep the import surface symmetric with the HTTP - ``Route`` import in :mod:`._invocation`. - - :param endpoint: The async endpoint to wire to the route. - :type endpoint: Callable[[WebSocket], Awaitable[None]] - :return: A configured ``WebSocketRoute`` instance. - :rtype: ~starlette.routing.WebSocketRoute + Starlette's ``router.routes`` is a plain list and may be mutated + between construction and first request, so deferring registration + until ``@ws_handler`` is called is safe. """ from starlette.routing import WebSocketRoute # pylint: disable=import-outside-toplevel - return WebSocketRoute( - InvocationsWSConstants.ROUTE_PATH, - endpoint, - name="invocations_ws", + for route in self.router.routes: + if isinstance(route, WebSocketRoute) and getattr(route, "path", None) == InvocationsWSConstants.ROUTE_PATH: + return + self.router.routes.append( + WebSocketRoute( + InvocationsWSConstants.ROUTE_PATH, + self._ws_endpoint, + name="invocations_ws", + ) ) # ------------------------------------------------------------------ @@ -191,10 +249,14 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None: :param websocket: The incoming Starlette WebSocket. :type websocket: ~starlette.websockets.WebSocket """ - # Per-connection identifiers. Session ID is generated server-side; - # the spec carries it in the close-event metric so an operator can - # correlate logs/metrics for a given long-lived connection. - session_id = str(uuid.uuid4()) + # Per-connection identifiers. Honour the platform-injected + # ``FOUNDRY_AGENT_SESSION_ID`` (surfaced via ``self.config.session_id``) + # so HTTP and WebSocket transports on the same container report the + # same session ID; fall back to a fresh UUID when the platform does + # not inject one. Matches the precedence used by the HTTP + # ``POST /invocations`` endpoint (minus the query-param override, + # which has no equivalent ergonomic on a long-lived WS connection). + session_id = self.config.session_id or str(uuid.uuid4()) start_ns = time.monotonic_ns() # Open the OTel span before accepting so any tracecontext header @@ -212,9 +274,9 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None: otel_span = span_ctx.__enter__() _safe_set_attrs(otel_span, {InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id}) - if self._ws_fn is None: - await self._reject_no_handler(websocket, span_ctx, otel_span, session_id, start_ns) - return + # NOTE: when no ``@ws_handler`` is registered, the route itself is + # not registered (see ``_ensure_ws_route_registered``), so this + # endpoint is unreachable in that state — Starlette returns 404. # Accept the upgrade *before* invoking the user handler — per spec. try: @@ -262,10 +324,14 @@ async def _invoke_user_handler( is set only for an *unhandled* exception (so the caller can map it to span error events and a 1011 close). :rtype: tuple[int, Optional[BaseException]] + :raises RuntimeError: If no handler is registered (programmer error; + the public ``_ws_endpoint`` checks ``_ws_fn`` before calling). """ - assert self._ws_fn is not None # checked by caller + ws_fn = self._ws_fn + if ws_fn is None: + raise RuntimeError("_invoke_user_handler called with no registered ws_handler") try: - await self._ws_fn(websocket) + await ws_fn(websocket) return InvocationsWSConstants.CLOSE_NORMAL, None except WebSocketDisconnect as exc: # Client (or proxy) closed first — surface their code, not 1011. @@ -280,49 +346,6 @@ async def _invoke_user_handler( ) return InvocationsWSConstants.CLOSE_INTERNAL_ERROR, exc - async def _reject_no_handler( - self, - websocket: WebSocket, - span_ctx: Any, - otel_span: Any, - session_id: str, - start_ns: int, - ) -> None: - """Refuse a WS upgrade when no ``@ws_handler`` is registered. - - :param websocket: The pending WebSocket awaiting upgrade. - :type websocket: ~starlette.websockets.WebSocket - :param span_ctx: The active ``request_span`` context manager. - :type span_ctx: any - :param otel_span: The current OTel span (or ``None``). - :type otel_span: any - :param session_id: Per-connection session ID. - :type session_id: str - :param start_ns: ``time.monotonic_ns()`` at connection start. - :type start_ns: int - """ - logger.error( - "WebSocket connection on %s rejected: no @ws_handler registered", - InvocationsWSConstants.ROUTE_PATH, - ) - duration_ms = (time.monotonic_ns() - start_ns) // 1_000_000 - self._emit_close_event( - otel_span, - session_id, - InvocationsWSConstants.CLOSE_INTERNAL_ERROR, - duration_ms, - error_code="not_implemented", - error_message="No ws_handler registered.", - ) - try: - span_ctx.__exit__(None, None, None) - finally: - end_span(otel_span) - await websocket.close( - code=InvocationsWSConstants.CLOSE_INTERNAL_ERROR, - reason="No ws_handler registered", - ) - async def _finalize_session( self, *, @@ -390,21 +413,16 @@ async def _finalize_session( error_message=str(handler_exc) if handler_exc is not None else None, ) - if handler_exc is not None: - try: - record_error(otel_span, handler_exc) - finally: - try: - span_ctx.__exit__( - type(handler_exc), handler_exc, handler_exc.__traceback__, - ) - finally: - end_span(otel_span) - else: - try: - span_ctx.__exit__(None, None, None) - finally: - end_span(otel_span) + # Always exit the context manager with ``(None, None, None)``: + # OpenTelemetry's ``start_as_current_span`` records an exception + # event whenever ``__exit__`` is given exception info, so passing + # the user exception there *and* calling ``end_span(span, exc=...)`` + # would double-record. We take ownership of error recording and + # ending via ``end_span`` instead. + try: + span_ctx.__exit__(None, None, None) + finally: + end_span(otel_span, exc=handler_exc) # ------------------------------------------------------------------ # Close event @@ -422,8 +440,9 @@ def _emit_close_event( ) -> None: """Record close-event span attributes and emit a structured log line. - The log record carries the ``ws.session_id``, ``ws.close_code``, - and ``ws.duration_ms`` fields listed in the spec via the standard + The log record carries the ``azure.ai.agentserver.invocations_ws.session_id``, + ``azure.ai.agentserver.invocations_ws.close_code``, and + ``azure.ai.agentserver.invocations_ws.duration_ms`` fields listed in the spec via the standard ``logging`` ``extra`` dict — a structured-logging formatter or an OTel logging bridge can pick them up directly without having to parse the message. @@ -454,14 +473,22 @@ def _emit_close_event( attrs[InvocationsWSConstants.ATTR_SPAN_ERROR_MESSAGE] = error_message _safe_set_attrs(otel_span, attrs) + # NOTE: ``extra`` keys deliberately use dotted names (``azure.ai.agentserver.invocations_ws.session_id`` + # etc.) so they line up 1:1 with the OTel span-attribute keys defined + # in :class:`InvocationsWSConstants`. This makes log<->trace + # correlation trivial in OTel logging bridges. The trade-off is that + # ``%(azure.ai.agentserver.invocations_ws.session_id)s`` printf-style formatters can't address them + # directly — use a structured formatter (JSON, OTel) or access via + # ``LogRecord.__dict__["azure.ai.agentserver.invocations_ws.session_id"]`` if you need them in plain + # ``logging`` output. logger.info( "invocations_ws connection closed: session_id=%s code=%s duration_ms=%s", session_id, close_code, duration_ms, extra={ - "ws.session_id": session_id, - "ws.close_code": close_code, - "ws.duration_ms": duration_ms, + InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id, + InvocationsWSConstants.ATTR_SPAN_CLOSE_CODE: close_code, + InvocationsWSConstants.ATTR_SPAN_DURATION_MS: duration_ms, }, ) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py index 9a66525d3c03..a228f61e145d 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py @@ -18,12 +18,14 @@ .. note:: - Connection keep-alive is **not** an application concern: the SDK - configures Hypercorn to send WebSocket protocol-level Ping frames - (opcode 0x9) every ``ws_ping_interval`` seconds (default 30 s). That - is enough to survive Azure APIM / Azure Load Balancer's ~4-minute idle - timeout without your handler having to push any application-level - heartbeat messages of its own. + Connection keep-alive is **not** an application concern: the SDK can + ask Hypercorn to send WebSocket protocol-level Ping frames + (opcode 0x9) every ``ws_ping_interval`` seconds (disabled by default; + enable by setting ``WS_KEEPALIVE_INTERVAL`` or by passing + ``InvocationAgentServerHost(ws_ping_interval=)``). When + enabled, that is enough to survive Azure APIM / Azure Load Balancer's + ~4-minute idle timeout without your handler having to push any + application-level heartbeat messages of its own. Wire protocol (JSON over text frames) ------------------------------------- @@ -49,14 +51,13 @@ Drive it with the ``websockets`` CLI; the server keeps streaming tokens for each prompt while you type the next one:: + pip install websockets python -m websockets ws://localhost:8088/invocations_ws > {"type": "prompt", "id": "p1", "text": "Tell me a story"} > {"type": "prompt", "id": "p2", "text": "And another"} > {"type": "cancel", "id": "p1"} > {"type": "bye"} """ -from __future__ import annotations - import asyncio import contextlib import json @@ -70,12 +71,16 @@ from azure.ai.agentserver.invocations import InvocationAgentServerHost -logger = logging.getLogger("ws_bidirectional_streaming_agent") +logger = logging.getLogger("azure.ai.agentserver") app = InvocationAgentServerHost() # Simulated tokens — in production these would come from a model. +# For real-world streaming patterns, see the streaming examples in +# ``azure-ai-inference`` and ``azure-ai-projects``: +# https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/ai/azure-ai-inference/samples +# https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/ai/azure-ai-projects/samples _SIMULATED_TOKENS = [ "Once", " upon", " a", " time", ",", " in", " a", " land", " of", " full", "-", "duplex", " sockets", ",", " a", " server", @@ -100,18 +105,18 @@ async def handle_invoke(request: Request) -> Response: # WebSocket — true bidirectional streaming. # --------------------------------------------------------------------------- -async def _generate_tokens(text: str) -> AsyncGenerator[str, None]: +async def _generate_tokens(_text: str) -> AsyncGenerator[str, None]: """Yield simulated tokens with a small per-token delay. Replace this with a real streaming model call (e.g. Azure OpenAI) in production. - :param text: The user prompt (unused in this demo). - :type text: str + :param _text: The user prompt (unused in this demo — leading underscore + signals "intentionally ignored" to linters). + :type _text: str :return: An async generator of token strings. :rtype: AsyncGenerator[str, None] """ - del text # demo: ignore prompt content for token in _SIMULATED_TOKENS: await asyncio.sleep(_TOKEN_DELAY_S) yield token @@ -138,10 +143,23 @@ async def _stream_tokens( except asyncio.CancelledError: # Best-effort: tell the client we honoured their cancel. Suppress # any send error (the socket may already be closed) and re-raise - # so the caller observes the cancellation. - with contextlib.suppress(Exception): + # so the caller observes the cancellation. ``BaseException`` + # rather than ``Exception`` so we don't accidentally swallow a + # nested cancellation while emitting the courtesy frame. + try: await websocket.send_json({"type": "cancelled", "id": prompt_id}) + except BaseException: # pylint: disable=broad-exception-caught + pass raise + except Exception as exc: # pylint: disable=broad-exception-caught + # Surface generation errors to the client as a structured frame so + # the connection survives a single-prompt failure (parity with the + # ``error`` reply the reader emits on bad input). + with contextlib.suppress(Exception): + await websocket.send_json( + {"type": "error", "id": prompt_id, "message": str(exc)}, + ) + logger.exception("_stream_tokens failed for prompt %s", prompt_id) async def _reader( @@ -188,6 +206,9 @@ async def _reader( name=f"stream-{prompt_id}", ) in_flight[prompt_id] = task + # ``k=prompt_id`` captures the current value via a default + # argument so the callback removes the right entry even if + # ``prompt_id`` is rebound by a later iteration of the loop. task.add_done_callback( lambda _t, k=prompt_id: in_flight.pop(k, None), ) @@ -197,6 +218,13 @@ async def _reader( task = in_flight.get(prompt_id) if task is not None and not task.done(): task.cancel() + # Give the task a brief window to finish its courtesy + # ``cancelled`` frame before we move on — prevents the + # next prompt from racing against an in-flight close. + with contextlib.suppress( + asyncio.TimeoutError, asyncio.CancelledError, Exception, + ): + await asyncio.wait_for(task, timeout=1.0) elif msg_type == "bye": return diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py index 2f36c48a3b20..269d9a1d4c21 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py @@ -19,6 +19,7 @@ # -> {"echo": {"name": "Alice"}} # WebSocket turn (with the `websockets` client library) + pip install websockets python -m websockets ws://localhost:8088/invocations_ws # > hello # < hello @@ -45,10 +46,11 @@ async def handle_ws(websocket: WebSocket) -> None: """Echo every text frame back over the WebSocket connection. The SDK has already accepted the connection by the time this function - runs, sends WebSocket Ping frames every 30 s in the background to keep - Azure APIM / Azure Load Balancer from idling the socket out, and will - close the connection cleanly on return. An uncaught exception here - is mapped to RFC 6455 close code 1011. + runs and will close it cleanly on return. An uncaught exception is + mapped to RFC 6455 close code 1011. WebSocket protocol-level Ping/Pong + keep-alive is disabled by default; enable it by setting the + ``WS_KEEPALIVE_INTERVAL`` environment variable or by passing + ``InvocationAgentServerHost(ws_ping_interval=)``. """ async for message in websocket.iter_text(): await websocket.send_text(message) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py index 24c3c1eb9500..7f5947655ef3 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py @@ -16,6 +16,7 @@ def pytest_configure(config): config.addinivalue_line("markers", "tracing_e2e: end-to-end tracing tests against live Application Insights") + config.addinivalue_line("markers", "slow: tests that send large payloads or otherwise take noticeable time in CI") # --------------------------------------------------------------------------- @@ -236,7 +237,7 @@ def _records_with_ws_extras(records): """Filter log records that carry the close-event ``ws.*`` extras.""" return [ r for r in records - if hasattr(r, "ws.session_id") and hasattr(r, "ws.close_code") + if hasattr(r, "azure.ai.agentserver.invocations_ws.session_id") and hasattr(r, "azure.ai.agentserver.invocations_ws.close_code") ] diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py index 6c2c9723c301..5fd135ab875d 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py @@ -4,7 +4,7 @@ """Tests for the structured close-event log line emitted by ``/invocations_ws``. Parity with :mod:`tests.test_request_id` — verifies the spec's required -fields (``ws.session_id``, ``ws.close_code``, ``ws.duration_ms``) appear +fields (``azure.ai.agentserver.invocations_ws.session_id``, ``azure.ai.agentserver.invocations_ws.close_code``, ``azure.ai.agentserver.invocations_ws.duration_ms``) appear on every connection close, and that handler exception details are NOT leaked into the structured payload. """ @@ -14,6 +14,7 @@ from starlette.testclient import TestClient from starlette.websockets import WebSocketDisconnect +from azure.ai.agentserver.invocations import InvocationAgentServerHost from azure.ai.agentserver.invocations._constants import InvocationsWSConstants from conftest import ( @@ -41,9 +42,9 @@ def test_ws_close_event_log_contains_required_fields(caplog): assert matches, "expected a structured close-event log record" rec = matches[-1] - session_id = getattr(rec, "ws.session_id") - close_code = getattr(rec, "ws.close_code") - duration_ms = getattr(rec, "ws.duration_ms") + session_id = getattr(rec, "azure.ai.agentserver.invocations_ws.session_id") + close_code = getattr(rec, "azure.ai.agentserver.invocations_ws.close_code") + duration_ms = getattr(rec, "azure.ai.agentserver.invocations_ws.duration_ms") assert isinstance(session_id, str) and session_id # generated UUID assert close_code == InvocationsWSConstants.CLOSE_NORMAL @@ -52,7 +53,7 @@ def test_ws_close_event_log_contains_required_fields(caplog): def test_ws_close_event_duration_is_non_negative(caplog): - """``ws.duration_ms`` is a non-negative integer derived from a monotonic clock.""" + """``azure.ai.agentserver.invocations_ws.duration_ms`` is a non-negative integer derived from a monotonic clock.""" app = _make_echo_ws_app() client = TestClient(app) @@ -63,7 +64,7 @@ def test_ws_close_event_duration_is_non_negative(caplog): matches = _records_with_ws_extras(caplog.records) assert matches - duration_ms = getattr(matches[-1], "ws.duration_ms") + duration_ms = getattr(matches[-1], "azure.ai.agentserver.invocations_ws.duration_ms") assert isinstance(duration_ms, int) assert duration_ms >= 0 @@ -85,7 +86,7 @@ def test_ws_close_event_on_handler_exception_records_1011(caplog): matches = _records_with_ws_extras(caplog.records) assert matches - assert getattr(matches[-1], "ws.close_code") == 1011 + assert getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code") == 1011 # --------------------------------------------------------------------------- @@ -108,6 +109,54 @@ def test_ws_close_event_log_does_not_leak_exception_message(caplog): assert matches rec = matches[-1] # The structured close-event log line carries only the safe ws.* fields. - assert not hasattr(rec, "ws.error_message") + assert not hasattr(rec, "azure.ai.agentserver.invocations_ws.error.message") # And the message itself does not embed the raw exception text. assert "boom" not in rec.getMessage() + + +def test_ws_disconnect_with_code_zero_falls_back_to_normal_close(caplog): + """A ``WebSocketDisconnect(code=0)`` is reported as the normal close code (1000). + + Some clients/proxies surface a falsy ``code`` on the disconnect; the SDK + treats it as a clean close rather than as 1011. + """ + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket): # noqa: ARG001 + raise WebSocketDisconnect(code=0) + + client = TestClient(app) + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/invocations_ws") as ws: + ws.receive_text() + + matches = _records_with_ws_extras(caplog.records) + assert matches + assert ( + getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code") + == InvocationsWSConstants.CLOSE_NORMAL + ) + + +def test_ws_disconnect_with_code_none_falls_back_to_normal_close(caplog): + """A ``WebSocketDisconnect(code=None)`` is reported as 1000.""" + app = InvocationAgentServerHost() + + @app.ws_handler + async def handler(websocket): # noqa: ARG001 + raise WebSocketDisconnect(code=None) # type: ignore[arg-type] + + client = TestClient(app) + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/invocations_ws") as ws: + ws.receive_text() + + matches = _records_with_ws_extras(caplog.records) + assert matches + assert ( + getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code") + == InvocationsWSConstants.CLOSE_NORMAL + ) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py index b447ff8d1090..54d25b415601 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py @@ -11,7 +11,6 @@ from starlette.websockets import WebSocket, WebSocketDisconnect from azure.ai.agentserver.invocations import InvocationAgentServerHost -from azure.ai.agentserver.invocations._constants import InvocationsWSConstants # --------------------------------------------------------------------------- @@ -61,32 +60,56 @@ def test_ws_handler_default_is_none(): assert app._ws_fn is None # noqa: SLF001 -def test_ws_handler_last_registration_wins(): - """Re-applying ``@ws_handler`` replaces the previous function.""" +def test_ws_handler_last_registration_wins(caplog): + """Re-applying ``@ws_handler`` replaces the previous function and warns.""" + import logging + app = InvocationAgentServerHost() @app.ws_handler async def first(websocket: WebSocket) -> None: # noqa: ARG001 return - @app.ws_handler - async def second(websocket: WebSocket) -> None: # noqa: ARG001 - return + with caplog.at_level(logging.WARNING, logger="azure.ai.agentserver"): + @app.ws_handler + async def second(websocket: WebSocket) -> None: # noqa: ARG001 + return assert app._ws_fn is second # noqa: SLF001 + assert any("overwriting previously registered handler" in r.message for r in caplog.records) + + +def test_ws_handler_rejects_zero_arg_coroutine(): + """A 0-arg coroutine cannot accept the WebSocket — fail at registration.""" + app = InvocationAgentServerHost() + + with pytest.raises(TypeError, match="exactly one positional argument"): + @app.ws_handler + async def bad() -> None: # type: ignore[misc] + return + + +def test_ws_handler_rejects_two_required_arg_coroutine(): + """A 2-required-arg coroutine cannot be bound with one positional.""" + app = InvocationAgentServerHost() + + with pytest.raises(TypeError, match="exactly one positional argument"): + @app.ws_handler + async def bad(websocket: WebSocket, extra: int) -> None: # noqa: ARG001 + return # --------------------------------------------------------------------------- # Missing handler behaviour (parity with test_missing_invoke_handler_returns_501) # --------------------------------------------------------------------------- -def test_ws_with_no_handler_registered_closes_with_1011(): - """If no @ws_handler is registered, the SDK closes with 1011.""" +def test_ws_with_no_handler_registered_rejects_upgrade(): + """If no @ws_handler is registered the route is absent and the upgrade is rejected.""" app = InvocationAgentServerHost() client = TestClient(app) - with pytest.raises(WebSocketDisconnect) as excinfo: + # Starlette has no WebSocketRoute matching /invocations_ws — the upgrade + # is rejected before any handler-level close code can be sent. + with pytest.raises(WebSocketDisconnect): with client.websocket_connect("/invocations_ws"): pass - - assert excinfo.value.code == InvocationsWSConstants.CLOSE_INTERNAL_ERROR diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py index a25f1249ee8b..90fc3183e455 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py @@ -64,7 +64,7 @@ async def handler(websocket: WebSocket) -> None: matches = _records_with_ws_extras(caplog.records) assert matches rec = matches[-1] - close_code = getattr(rec, "ws.close_code") + close_code = getattr(rec, "azure.ai.agentserver.invocations_ws.close_code") # Server surfaces the client's code — NOT 1011. assert close_code == 4001 # No ERROR-level records — a client disconnect is normal. @@ -75,10 +75,19 @@ async def handler(websocket: WebSocket) -> None: # Handler-managed close # --------------------------------------------------------------------------- -def test_ws_handler_explicit_close_does_not_double_close(caplog): +def test_ws_handler_explicit_close_does_not_double_close(caplog, monkeypatch): """If the handler closes the WS itself, the SDK does NOT attempt a second close.""" app = InvocationAgentServerHost() closes_observed: list[int] = [] + close_calls: list[tuple[int, str]] = [] + + real_close = WebSocket.close + + async def _counting_close(self, code=1000, reason=""): + close_calls.append((code, reason)) + return await real_close(self, code=code, reason=reason) + + monkeypatch.setattr(WebSocket, "close", _counting_close, raising=True) async def handler(websocket: WebSocket) -> None: # Hand-roll close so we can verify the SDK skips re-closing. @@ -97,10 +106,12 @@ async def handler(websocket: WebSocket) -> None: # Handler ran to completion. assert closes_observed == [1] + # Exactly ONE close — the handler's. The SDK must not re-close. + assert len(close_calls) == 1, f"expected one close, got {close_calls}" # Close-event log line still emitted (with the handler's code). matches = _records_with_ws_extras(caplog.records) assert matches - assert getattr(matches[-1], "ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL + assert getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL # --------------------------------------------------------------------------- @@ -118,4 +129,4 @@ def test_ws_empty_connection_closes_normally(caplog): matches = _records_with_ws_extras(caplog.records) assert matches - assert getattr(matches[-1], "ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL + assert getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py index 9eb988b73713..7fa992566d76 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py @@ -101,6 +101,10 @@ async def handler(websocket: WebSocket) -> None: await websocket.receive_text() client = TestClient(app) + # Starlette TestClient surfaces *any* close (including a clean 1000) via + # ``WebSocketDisconnect`` when ``receive_text()`` is called after the + # server sends a close frame — this is the documented client API for + # observing the close, not an indication of error. with pytest.raises(WebSocketDisconnect) as excinfo: with client.websocket_connect("/invocations_ws") as ws: ws.send_text("done") diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py index e05b228c13a8..334465244870 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py @@ -9,6 +9,8 @@ from starlette.testclient import TestClient from starlette.websockets import WebSocket +import pytest + from azure.ai.agentserver.invocations import InvocationAgentServerHost from conftest import _make_echo_ws_app @@ -48,6 +50,7 @@ def test_ws_unicode_text_round_trip(): assert ws.receive_text() == text +@pytest.mark.slow def test_ws_large_text_frame_round_trip(): """A ~1 MB text frame round-trips successfully (parity with test_large_payload).""" app = _make_echo_ws_app() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py index 82567c6df622..e8ead6148b9e 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py @@ -17,11 +17,10 @@ # Default / accepted values # --------------------------------------------------------------------------- -def test_ws_ping_interval_default_is_30_seconds(): - """Default ping interval matches the spec (30 s).""" +def test_ws_ping_interval_default_is_disabled(): + """Default ping interval is 0 (disabled) when no arg / env var is set.""" app = InvocationAgentServerHost() - assert app.ws_ping_interval == InvocationsWSConstants.DEFAULT_PING_INTERVAL_S - assert app.ws_ping_interval == 30.0 + assert app.ws_ping_interval == 0.0 def test_ws_ping_interval_custom_value(): @@ -84,12 +83,11 @@ def test_ws_ping_interval_propagates_to_hypercorn_config(): assert getattr(config, "websocket_ping_interval", None) == 20.0 -def test_ws_ping_interval_zero_does_not_override_hypercorn_default(): - """Zero leaves Hypercorn's default (None = disabled) intact.""" +def test_ws_ping_interval_zero_sets_attribute_to_none(): + """Zero explicitly sets Hypercorn's ``websocket_ping_interval`` to ``None``.""" app = InvocationAgentServerHost(ws_ping_interval=0) config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 - # Hypercorn default is None — our wiring leaves it unset for 0. - assert getattr(config, "websocket_ping_interval", None) is None + assert config.websocket_ping_interval is None # type: ignore[attr-defined] # --------------------------------------------------------------------------- @@ -101,3 +99,74 @@ def test_ws_ping_interval_property_is_read_only(): app = InvocationAgentServerHost(ws_ping_interval=20) with pytest.raises(AttributeError): app.ws_ping_interval = 10 # type: ignore[misc] + + +# --------------------------------------------------------------------------- +# WS_KEEPALIVE_INTERVAL env var (platform-injected override) +# --------------------------------------------------------------------------- + +def test_ws_keepalive_interval_env_var_used_when_arg_is_none(monkeypatch): + """Env var drives the interval when no explicit arg is passed.""" + monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "45") + app = InvocationAgentServerHost() + assert app.ws_ping_interval == 45.0 + + +def test_ws_keepalive_interval_env_var_zero_disables(monkeypatch): + """``WS_KEEPALIVE_INTERVAL=0`` disables keep-alive (mirrors SSE).""" + monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "0") + app = InvocationAgentServerHost() + assert app.ws_ping_interval == 0.0 + config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 + assert getattr(config, "websocket_ping_interval", None) is None + + +def test_ws_keepalive_interval_env_var_float_accepted(monkeypatch): + """Fractional env-var values are coerced to ``float``.""" + monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "12.5") + app = InvocationAgentServerHost() + assert app.ws_ping_interval == 12.5 + + +def test_ws_keepalive_interval_env_var_empty_uses_default(monkeypatch): + """An empty env-var value falls back to the default (disabled).""" + monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "") + app = InvocationAgentServerHost() + assert app.ws_ping_interval == 0.0 + + +def test_ws_keepalive_interval_env_var_invalid_rejected(monkeypatch): + """Non-numeric env-var values surface as ``ValueError`` at startup.""" + monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "thirty") + with pytest.raises(ValueError, match="WS_KEEPALIVE_INTERVAL"): + InvocationAgentServerHost() + + +def test_ws_keepalive_interval_env_var_negative_rejected(monkeypatch): + """Negative env-var values are programming errors.""" + monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "-5") + with pytest.raises(ValueError, match="WS_KEEPALIVE_INTERVAL"): + InvocationAgentServerHost() + + +def test_ws_ping_interval_arg_overrides_env_var(monkeypatch): + """Explicit constructor arg wins over the env var.""" + monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "45") + app = InvocationAgentServerHost(ws_ping_interval=10) + assert app.ws_ping_interval == 10.0 + + +def test_ws_ping_interval_default_wires_none_into_hypercorn(monkeypatch): + """With neither arg nor env var, Hypercorn config has ``websocket_ping_interval = None``.""" + monkeypatch.delenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, raising=False) + app = InvocationAgentServerHost() + config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 + assert config.websocket_ping_interval is None # type: ignore[attr-defined] + + +def test_ws_ping_interval_positive_wires_into_hypercorn(): + """A positive interval propagates as a float to Hypercorn config.""" + app = InvocationAgentServerHost(ws_ping_interval=12) + config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 + assert config.websocket_ping_interval == 12.0 # type: ignore[attr-defined] + diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py index 4972b76275c7..8629dc6078ad 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py @@ -21,15 +21,24 @@ # Route registration # --------------------------------------------------------------------------- -def test_ws_route_is_registered(): - """The /invocations_ws route exists alongside the HTTP routes.""" - app = InvocationAgentServerHost() +def test_ws_route_is_registered_when_handler_is_set(): + """The /invocations_ws route is registered lazily on @ws_handler.""" + app = _make_echo_ws_app() paths = [getattr(r, "path", None) for r in app.routes] assert "/invocations_ws" in paths assert "/invocations" in paths assert "/readiness" in paths +def test_ws_route_is_not_registered_without_handler(): + """Without @ws_handler the WS route is absent (upgrades return 404).""" + app = InvocationAgentServerHost() + paths = [getattr(r, "path", None) for r in app.routes] + assert "/invocations_ws" not in paths + # HTTP routes still registered. + assert "/invocations" in paths + + def test_readiness_still_works_with_ws_registered(): """Adding the WS route doesn't break /readiness.""" app = _make_echo_ws_app() @@ -81,7 +90,7 @@ def test_ws_upgrade_on_http_path_fails(): client = TestClient(app) # /invocations is a Route, not a WebSocketRoute — TestClient surfaces # this as an immediate WebSocketDisconnect rather than a connect. - with pytest.raises((WebSocketDisconnect, Exception)): # noqa: PT011 + with pytest.raises(WebSocketDisconnect): with client.websocket_connect("/invocations"): pass @@ -90,6 +99,6 @@ def test_ws_unknown_path_fails(): """An unknown WS path is rejected.""" app = _make_echo_ws_app() client = TestClient(app) - with pytest.raises((WebSocketDisconnect, Exception)): # noqa: PT011 + with pytest.raises(WebSocketDisconnect): with client.websocket_connect("/nonexistent"): pass diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py index 43107ed71b7e..2163e8cd9286 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py @@ -20,8 +20,8 @@ # --------------------------------------------------------------------------- def _session_ids_from_records(records): - """Pull ``ws.session_id`` from each structured close-event record.""" - return [getattr(r, "ws.session_id") for r in _records_with_ws_extras(records)] + """Pull ``azure.ai.agentserver.invocations_ws.session_id`` from each structured close-event record.""" + return [getattr(r, "azure.ai.agentserver.invocations_ws.session_id") for r in _records_with_ws_extras(records)] # --------------------------------------------------------------------------- @@ -56,7 +56,7 @@ def test_ws_session_id_is_unique_per_connection(caplog): ws.receive_text() session_ids = _session_ids_from_records(caplog.records) - assert len(session_ids) == 5 + # Uniqueness is the meaningful invariant — implies len == 5 too. assert len(set(session_ids)) == 5 @@ -78,3 +78,24 @@ def test_ws_session_id_ignores_query_param(caplog): # generates a fresh server-side UUID per connection. assert session_ids[-1] != "my-custom-session" uuid.UUID(session_ids[-1]) + + +def test_ws_session_id_uses_foundry_env_var(caplog, monkeypatch): + """When ``FOUNDRY_AGENT_SESSION_ID`` is set, every WS session reports it. + + Parity with HTTP ``POST /invocations``: the platform-injected session + ID must propagate to both transports so cross-transport correlation + works on the same container. + """ + monkeypatch.setenv("FOUNDRY_AGENT_SESSION_ID", "platform-session-abc") + app = _make_echo_ws_app() + client = TestClient(app) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + for _ in range(3): + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("ping") + ws.receive_text() + + session_ids = _session_ids_from_records(caplog.records) + assert session_ids == ["platform-session-abc"] * 3 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py index f78b0c815c92..a83dcfe18690 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py @@ -200,7 +200,7 @@ def test_ws_session_id_attribute_set_on_span(): spans = _ws_session_spans(_get_spans()) attrs = dict(spans[0].attributes) # WS-specific key - ws_session = attrs.get("ws.session_id") + ws_session = attrs.get("azure.ai.agentserver.invocations_ws.session_id") # General session key (shared with HTTP) ms_session = attrs.get("microsoft.session.id") assert ws_session and isinstance(ws_session, str) @@ -219,7 +219,7 @@ def test_ws_session_id_is_unique_per_span(): spans = _ws_session_spans(_get_spans()) assert len(spans) == 2 - session_ids = {dict(s.attributes).get("ws.session_id") for s in spans} + session_ids = {dict(s.attributes).get("azure.ai.agentserver.invocations_ws.session_id") for s in spans} assert len(session_ids) == 2 @@ -237,7 +237,7 @@ def test_ws_span_records_close_code_1000_on_clean_return(): spans = _ws_session_spans(_get_spans()) attrs = dict(spans[0].attributes) - assert attrs.get("ws.close_code") == 1000 + assert attrs.get("azure.ai.agentserver.invocations_ws.close_code") == 1000 def test_ws_span_records_close_code_1011_on_handler_exception(): @@ -251,11 +251,11 @@ def test_ws_span_records_close_code_1011_on_handler_exception(): spans = _ws_session_spans(_get_spans()) attrs = dict(spans[0].attributes) - assert attrs.get("ws.close_code") == 1011 + assert attrs.get("azure.ai.agentserver.invocations_ws.close_code") == 1011 def test_ws_span_records_duration_ms(): - """``ws.duration_ms`` is set as a non-negative integer.""" + """``azure.ai.agentserver.invocations_ws.duration_ms`` is set as a non-negative integer.""" server = _make_tracing_ws_server() client = TestClient(server) with client.websocket_connect("/invocations_ws") as ws: @@ -264,7 +264,7 @@ def test_ws_span_records_duration_ms(): spans = _ws_session_spans(_get_spans()) attrs = dict(spans[0].attributes) - duration_ms = attrs.get("ws.duration_ms") + duration_ms = attrs.get("azure.ai.agentserver.invocations_ws.duration_ms") assert isinstance(duration_ms, int) assert duration_ms >= 0 @@ -287,7 +287,7 @@ async def handler(websocket: WebSocket) -> None: spans = _ws_session_spans(_get_spans()) attrs = dict(spans[0].attributes) - assert attrs.get("ws.close_code") == 4010 + assert attrs.get("azure.ai.agentserver.invocations_ws.close_code") == 4010 # --------------------------------------------------------------------------- @@ -310,10 +310,13 @@ def test_ws_handler_exception_records_exception_on_span(): # Exception is recorded as an event by record_error. event_names = [e.name for e in span.events] assert "exception" in event_names + # Exactly one exception event — belt-and-suspenders against future + # double-recording on the same span. + assert event_names.count("exception") == 1 -def test_ws_no_handler_records_error_on_span(): - """A connection refused for lack of @ws_handler still produces an error-tagged span.""" +def test_ws_no_handler_produces_no_span(): + """Without @ws_handler the route is absent, so no WS span is created.""" app = _make_tracing_ws_no_handler() client = TestClient(app) with pytest.raises(WebSocketDisconnect): @@ -321,14 +324,7 @@ def test_ws_no_handler_records_error_on_span(): pass spans = _ws_session_spans(_get_spans()) - assert spans - attrs = dict(spans[0].attributes) - # Close code is 1011 and the error_code/message attrs are populated. - assert attrs.get("ws.close_code") == 1011 - assert attrs.get("azure.ai.agentserver.invocations_ws.error.code") == "not_implemented" - assert "ws_handler" in str( - attrs.get("azure.ai.agentserver.invocations_ws.error.message", "") - ) + assert spans == [] # --------------------------------------------------------------------------- @@ -383,14 +379,11 @@ def test_ws_agent_name_and_version_in_span_name(): assert "3.1" in name -def test_ws_agent_name_only_in_span_name(): +def test_ws_agent_name_only_in_span_name(monkeypatch): """Agent name without version still appears in span name.""" - env_override = {"FOUNDRY_AGENT_NAME": "ws-solo"} - env_copy = os.environ.copy() - env_copy.pop("FOUNDRY_AGENT_VERSION", None) - env_copy.update(env_override) - with patch.dict(os.environ, env_copy, clear=True): - server = _make_tracing_ws_server() + monkeypatch.setenv("FOUNDRY_AGENT_NAME", "ws-solo") + monkeypatch.delenv("FOUNDRY_AGENT_VERSION", raising=False) + server = _make_tracing_ws_server() client = TestClient(server) with client.websocket_connect("/invocations_ws") as ws: From 901f9074717993866598559d30db551eb9d929c1 Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Thu, 14 May 2026 08:55:42 +0000 Subject: [PATCH 03/13] Address Copilot AI review on PR 46841 - README: drop misleading 'server-pushed heartbeats' claim from sample table - _invocation_ws.py: docstring fixes (route registration is the gate, not an explicit _ws_fn check; tighten ws_handler raises wording to match what is actually validated) - _invocation_ws.py: wrap user-handler invocation in try/except BaseException + finally so CancelledError finalizes the span and emits the close-event log line before propagating - ws_bidirectional_streaming_agent.py: narrow BaseException -> Exception for the courtesy 'cancelled' send so KeyboardInterrupt/SystemExit propagate normally - test_ws_multimodal_protocol.py: reorder imports (stdlib/third-party blocks) - test_ws_decorator_pattern.py: update raises regex to match new message --- .../README.md | 2 +- .../agentserver/invocations/_invocation_ws.py | 59 ++++++++++++++----- .../ws_bidirectional_streaming_agent.py | 12 ++-- .../tests/test_ws_decorator_pattern.py | 4 +- .../tests/test_ws_multimodal_protocol.py | 3 +- 5 files changed, 52 insertions(+), 28 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index acab609ed746..3b1caacd2e84 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md @@ -265,7 +265,7 @@ Visit the [Samples](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/ | [simple_invoke_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/simple_invoke_agent/) | Minimal synchronous request-response | | [async_invoke_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/async_invoke_agent/) | Long-running operations with polling and cancellation | | [ws_invoke_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/) | Combined `POST /invocations` (HTTP) and `/invocations_ws` (WebSocket) host | -| [ws_bidirectional_streaming_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/) | Full-duplex `/invocations_ws` agent: server-pushed heartbeats + concurrent token streams + mid-flight cancel | +| [ws_bidirectional_streaming_agent](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/) | Full-duplex `/invocations_ws` agent: concurrent token streams + mid-flight cancel (relies on the SDK's WS protocol Ping/Pong keep-alive, not application-level heartbeats) | ## Contributing diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py index 86fedaca6efb..790beb6daa53 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py @@ -182,7 +182,8 @@ async def handle(websocket: WebSocket) -> None: :return: The original function (unmodified). :rtype: Callable[[WebSocket], Awaitable[None]] :raises TypeError: If *fn* is not an ``async def`` function, or its - signature is not callable with exactly one positional argument. + signature cannot be invoked with a single positional argument + (the WebSocket). """ if not inspect.iscoroutinefunction(fn): raise TypeError( @@ -196,8 +197,9 @@ async def handle(websocket: WebSocket) -> None: sig.bind(None) # one positional placeholder for the WebSocket except TypeError as exc: raise TypeError( - f"ws_handler must accept exactly one positional argument " - f"(the WebSocket); got {fn.__qualname__}{inspect.signature(fn)}" + f"ws_handler signature must be invocable with a single " + f"positional argument (the WebSocket); got " + f"{fn.__qualname__}{inspect.signature(fn)}" ) from exc if self._ws_fn is not None: # Match the HTTP decorator's last-write-wins semantics, but log @@ -298,17 +300,40 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None: ) return - close_code, handler_exc = await self._invoke_user_handler(websocket, session_id) - await self._finalize_session( - websocket=websocket, - span_ctx=span_ctx, - otel_span=otel_span, - session_id=session_id, - start_ns=start_ns, - close_code=close_code, - handler_exc=handler_exc, - error_code="internal_error" if handler_exc is not None else None, - ) + close_code: int = InvocationsWSConstants.CLOSE_NORMAL + handler_exc: Optional[BaseException] = None + try: + close_code, handler_exc = await self._invoke_user_handler(websocket, session_id) + except BaseException as exc: # pylint: disable=broad-exception-caught + # ``_invoke_user_handler`` catches ``Exception`` but not + # ``BaseException`` (notably ``asyncio.CancelledError``). Capture + # the exception so the ``finally`` block below can record it on + # the span, then re-raise via ``finally`` so cancellation is + # never swallowed. + close_code = InvocationsWSConstants.CLOSE_INTERNAL_ERROR + handler_exc = exc + raise + finally: + # Always finalize \u2014 emits the close-event log line, ends the + # span, and best-effort closes the socket \u2014 even when the + # handler raised a ``BaseException`` like ``CancelledError``. + error_code: Optional[str] + if handler_exc is None: + error_code = None + elif isinstance(handler_exc, Exception): + error_code = "internal_error" + else: + error_code = "cancelled" + await self._finalize_session( + websocket=websocket, + span_ctx=span_ctx, + otel_span=otel_span, + session_id=session_id, + start_ns=start_ns, + close_code=close_code, + handler_exc=handler_exc, + error_code=error_code, + ) async def _invoke_user_handler( self, websocket: WebSocket, session_id: str, @@ -324,8 +349,10 @@ async def _invoke_user_handler( is set only for an *unhandled* exception (so the caller can map it to span error events and a 1011 close). :rtype: tuple[int, Optional[BaseException]] - :raises RuntimeError: If no handler is registered (programmer error; - the public ``_ws_endpoint`` checks ``_ws_fn`` before calling). + :raises RuntimeError: If no handler is registered. The route is only + registered after ``ws_handler`` is decorated, so reaching this + method without a handler indicates a programming error in the + SDK itself rather than a user misconfiguration. """ ws_fn = self._ws_fn if ws_fn is None: diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py index a228f61e145d..903b635d9e7b 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py @@ -142,14 +142,12 @@ async def _stream_tokens( await websocket.send_json({"type": "done", "id": prompt_id}) except asyncio.CancelledError: # Best-effort: tell the client we honoured their cancel. Suppress - # any send error (the socket may already be closed) and re-raise - # so the caller observes the cancellation. ``BaseException`` - # rather than ``Exception`` so we don't accidentally swallow a - # nested cancellation while emitting the courtesy frame. - try: + # ordinary send errors (the socket may already be closed) and re-raise + # so the caller observes the cancellation. We deliberately do NOT + # catch ``BaseException`` here so process-level signals + # (``KeyboardInterrupt`` / ``SystemExit``) propagate normally. + with contextlib.suppress(Exception): await websocket.send_json({"type": "cancelled", "id": prompt_id}) - except BaseException: # pylint: disable=broad-exception-caught - pass raise except Exception as exc: # pylint: disable=broad-exception-caught # Surface generation errors to the client as a structured frame so diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py index 54d25b415601..835e21c1f5e1 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py @@ -83,7 +83,7 @@ def test_ws_handler_rejects_zero_arg_coroutine(): """A 0-arg coroutine cannot accept the WebSocket — fail at registration.""" app = InvocationAgentServerHost() - with pytest.raises(TypeError, match="exactly one positional argument"): + with pytest.raises(TypeError, match="single positional argument"): @app.ws_handler async def bad() -> None: # type: ignore[misc] return @@ -93,7 +93,7 @@ def test_ws_handler_rejects_two_required_arg_coroutine(): """A 2-required-arg coroutine cannot be bound with one positional.""" app = InvocationAgentServerHost() - with pytest.raises(TypeError, match="exactly one positional argument"): + with pytest.raises(TypeError, match="single positional argument"): @app.ws_handler async def bad(websocket: WebSocket, extra: int) -> None: # noqa: ARG001 return diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py index 334465244870..0dbeb9e2a14c 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py @@ -6,11 +6,10 @@ Parity with :mod:`tests.test_multimodal_protocol` — covers binary frames, text frames (unicode + large), and JSON frames over WebSocket. """ +import pytest from starlette.testclient import TestClient from starlette.websockets import WebSocket -import pytest - from azure.ai.agentserver.invocations import InvocationAgentServerHost from conftest import _make_echo_ws_app From 265e620747edad83d2b626532ba393c7aee09206 Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Fri, 15 May 2026 08:00:13 +0000 Subject: [PATCH 04/13] Reorder CHANGELOG: error-source bullet (from main) first, WS bullets after --- sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index fcb0fa1d9f25..5596188ef0d3 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -4,11 +4,11 @@ ### Features Added +- Error source classification headers: All HTTP error responses now include `x-platform-error-source` with a value of `user`, `platform`, or `upstream` to indicate which component caused the error. Developer handler exceptions and missing handler registrations are classified as `upstream`. Exceptions tagged with the platform error tag are classified as `platform` and additionally include `x-platform-error-detail` with truncated exception details (max 2048 characters) for diagnostics. - WebSocket protocol support — `InvocationAgentServerHost` now hosts `/invocations_ws` alongside `POST /invocations`. Register the handler with the new `@app.ws_handler` decorator. The route is registered lazily on first decoration, so hosts without a registered handler return HTTP 404. - WebSocket Ping/Pong keep-alive — disabled by default; enable by passing `InvocationAgentServerHost(ws_ping_interval=)` or by setting the `WS_KEEPALIVE_INTERVAL` env var (auto-injected by AgentService into hosted-agent containers). The constructor argument takes precedence; `0` (or unset) disables keep-alive. Wired through to Hypercorn's `websocket_ping_interval`. - WebSocket telemetry — structured close-event log line and OpenTelemetry span attributes `azure.ai.agentserver.invocations_ws.{session_id,close_code,duration_ms}`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. - New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages). -- Error source classification headers: All HTTP error responses now include `x-platform-error-source` with a value of `user`, `platform`, or `upstream` to indicate which component caused the error. Developer handler exceptions and missing handler registrations are classified as `upstream`. Exceptions tagged with the platform error tag are classified as `platform` and additionally include `x-platform-error-detail` with truncated exception details (max 2048 characters) for diagnostics. ### Breaking Changes From ea90b70583af9a373cf594db16571a89a20100fe Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Mon, 18 May 2026 09:34:21 +0000 Subject: [PATCH 05/13] Address review threads 2-6 on PR 46841: move ws_ping_interval to AgentConfig, wire in core, add per-turn ws_invocation helper Thread 2 + 3 (core refactor): - AgentConfig gains ws_ping_interval; resolve_ws_ping_interval() centralises parsing of the WS_KEEPALIVE_INTERVAL env var (validation, negative/NaN/inf rejection). AgentServerHost._build_hypercorn_config wires it into Hypercorn's websocket_ping_interval. - InvocationAgentServerHost drops ws_ping_interval= kwarg, env parsing, and its _build_hypercorn_config override. _init_ws_state is now no-arg; the app.ws_ping_interval property is a thin alias for app.config.ws_ping_interval. - InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL constant removed (env name is now centralised in core). Threads 4 + 5 + 6 (tracing model): - New app.ws_invocation(websocket) async context manager opens a per-turn 'invoke_agent' span (parity with POST /invocations) parented to the connection-level websocket_session span, mints a per-turn invocation_id, and propagates OTel context through awaits inside the block via start_as_current_span. Yields a WSInvocationContext exposing invocation_id + session_id. - Both WS samples updated to demonstrate ws_invocation and echo invocation_id back to clients. - README adds a 'Per-turn tracing with app.ws_invocation' subsection and documents that HTTP-only GET/cancel do not apply to WS turns (cancellation = WS close + CancelledError). - Added 6 tests in test_ws_invocation_helper.py covering span shape, invocation_id uniqueness, parent-child relationship, session_id inheritance, and OTel context propagation through awaits. Tests: core 89 passed, invocations 200 passed (188 + 6 new + reworked 6 ping-interval). All 6 samples compile. --- .../azure-ai-agentserver-core/CHANGELOG.md | 1 + .../azure/ai/agentserver/core/_base.py | 8 +- .../azure/ai/agentserver/core/_config.py | 49 +++++ .../CHANGELOG.md | 7 +- .../README.md | 35 +++- .../ai/agentserver/invocations/_constants.py | 10 +- .../ai/agentserver/invocations/_invocation.py | 53 +---- .../agentserver/invocations/_invocation_ws.py | 164 ++++++++++----- .../ws_bidirectional_streaming_agent.py | 53 +++-- .../ws_invoke_agent/ws_invoke_agent.py | 12 +- .../tests/test_ws_invocation_helper.py | 197 ++++++++++++++++++ .../tests/test_ws_ping_interval.py | 173 ++++++--------- 12 files changed, 506 insertions(+), 256 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invocation_helper.py diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index 024a8e7857d5..cd6e0fe84e97 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -5,6 +5,7 @@ ### Features Added - Added `_platform_headers` module with cross-cutting protocol header name constants (`x-request-id`, `x-platform-server`, `x-agent-session-id`, `x-platform-error-source`, `x-platform-error-detail`, and others). Protocol packages now import shared header name strings from core instead of maintaining their own copies. +- `AgentConfig.ws_ping_interval` — new field resolved from the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers). `0` disables; negative/non-finite values raise `ValueError` at startup. `AgentServerHost._build_hypercorn_config` wires this into Hypercorn's `websocket_ping_interval` so any protocol package serving WebSocket routes inherits keep-alive without per-package wiring. ## 2.0.0b3 (2026-04-22) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py index 0785f01e36ba..f4ea435397db 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py @@ -219,13 +219,14 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]: # noqa: RUF cfg = self.config logger.info( "Platform environment: is_hosted=%s, agent_name=%s, agent_version=%s, " - "port=%s, session_id=%s, sse_keepalive_interval=%s", + "port=%s, session_id=%s, sse_keepalive_interval=%s, ws_ping_interval=%s", cfg.is_hosted, cfg.agent_name or _NOT_SET, cfg.agent_version or _NOT_SET, cfg.port, cfg.session_id or _NOT_SET, cfg.sse_keepalive_interval if cfg.sse_keepalive_interval > 0 else "disabled", + f"{cfg.ws_ping_interval}s" if cfg.ws_ping_interval > 0 else "disabled", ) logger.info( "Connectivity: project_endpoint=%s, otlp_endpoint=%s, appinsights_configured=%s", @@ -417,6 +418,11 @@ def _build_hypercorn_config(self, host: str, port: int) -> object: config.graceful_timeout = float(self._graceful_shutdown_timeout) # Spec requires HTTP/1.1 only — disable HTTP/2 config.h2_max_concurrent_streams = 0 + # WebSocket Ping/Pong keep-alive (RFC 6455 opcodes 0x9/0xA). + # ``0`` (disabled) maps to Hypercorn's ``None`` sentinel; any + # positive value is sent verbatim to Hypercorn. + ws_ping = self.config.ws_ping_interval + config.websocket_ping_interval = ws_ping if ws_ping > 0 else None # type: ignore[attr-defined] # Access logging if self._access_log is not None: config.accesslog = self._access_log # type: ignore[assignment] diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py index e22bc1ff1cf6..a8433427debd 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py @@ -32,9 +32,11 @@ _ENV_APPLICATIONINSIGHTS_CONNECTION_STRING = "APPLICATIONINSIGHTS_CONNECTION_STRING" _ENV_OTEL_EXPORTER_OTLP_ENDPOINT = "OTEL_EXPORTER_OTLP_ENDPOINT" _ENV_SSE_KEEPALIVE_INTERVAL = "SSE_KEEPALIVE_INTERVAL" +_ENV_WS_KEEPALIVE_INTERVAL = "WS_KEEPALIVE_INTERVAL" _DEFAULT_PORT = 8088 _DEFAULT_SSE_KEEPALIVE_INTERVAL = 0 +_DEFAULT_WS_PING_INTERVAL = 0.0 # ====================================================================== @@ -64,6 +66,8 @@ class AgentConfig: # pylint: disable=too-many-instance-attributes :param appinsights_connection_string: Application Insights connection string. :param otlp_endpoint: OTLP exporter endpoint. :param sse_keepalive_interval: SSE keep-alive interval in seconds (0 = disabled). + :param ws_ping_interval: WebSocket protocol Ping interval in seconds + (``0`` disables keep-alive). """ def __init__( @@ -80,6 +84,7 @@ def __init__( appinsights_connection_string: str, otlp_endpoint: str, sse_keepalive_interval: int, + ws_ping_interval: float = 0.0, ) -> None: self.agent_name = agent_name self.agent_version = agent_version @@ -92,6 +97,7 @@ def __init__( self.appinsights_connection_string = appinsights_connection_string self.otlp_endpoint = otlp_endpoint self.sse_keepalive_interval = sse_keepalive_interval + self.ws_ping_interval = ws_ping_interval @classmethod def from_env(cls) -> Self: @@ -123,6 +129,7 @@ def from_env(cls) -> Self: _ENV_APPLICATIONINSIGHTS_CONNECTION_STRING, ""), otlp_endpoint=os.environ.get(_ENV_OTEL_EXPORTER_OTLP_ENDPOINT, ""), sse_keepalive_interval=resolve_sse_keepalive_interval(None), + ws_ping_interval=resolve_ws_ping_interval(None), ) @@ -322,3 +329,45 @@ def resolve_otlp_endpoint() -> Optional[str]: """ value = os.environ.get(_ENV_OTEL_EXPORTER_OTLP_ENDPOINT, "") return value if value else None + + +def resolve_ws_ping_interval(interval: Optional[float] = None) -> float: + """Resolve the WebSocket Ping/Pong keep-alive interval. + + Resolution order: explicit *interval* → ``WS_KEEPALIVE_INTERVAL`` env + var → ``0.0`` (disabled). ``0`` (from any source) disables + keep-alive. + + :param interval: Explicitly requested interval in seconds, or None. + :type interval: Optional[float] + :return: The resolved interval in seconds (``0`` means disabled). + :rtype: float + :raises ValueError: If the resolved value is negative, non-finite, or + (for the env var) not parseable as a number. + """ + import math # local import — only used here + + if interval is not None: + try: + resolved = float(interval) + except (TypeError, ValueError) as exc: + raise ValueError( + f"ws_ping_interval must be a number, got {interval!r}" + ) from exc + else: + env_raw = os.environ.get(_ENV_WS_KEEPALIVE_INTERVAL) + if env_raw is None or env_raw == "": + return _DEFAULT_WS_PING_INTERVAL + try: + resolved = float(env_raw) + except ValueError as exc: + raise ValueError( + f"Invalid value for {_ENV_WS_KEEPALIVE_INTERVAL}: " + f"{env_raw!r} (expected a non-negative number)" + ) from exc + if math.isnan(resolved) or math.isinf(resolved) or resolved < 0.0: + raise ValueError( + f"ws_ping_interval must be a non-negative finite number, " + f"got {interval if interval is not None else env_raw!r}" + ) + return resolved diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index 5596188ef0d3..a480e7b543d6 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -6,9 +6,10 @@ - Error source classification headers: All HTTP error responses now include `x-platform-error-source` with a value of `user`, `platform`, or `upstream` to indicate which component caused the error. Developer handler exceptions and missing handler registrations are classified as `upstream`. Exceptions tagged with the platform error tag are classified as `platform` and additionally include `x-platform-error-detail` with truncated exception details (max 2048 characters) for diagnostics. - WebSocket protocol support — `InvocationAgentServerHost` now hosts `/invocations_ws` alongside `POST /invocations`. Register the handler with the new `@app.ws_handler` decorator. The route is registered lazily on first decoration, so hosts without a registered handler return HTTP 404. -- WebSocket Ping/Pong keep-alive — disabled by default; enable by passing `InvocationAgentServerHost(ws_ping_interval=)` or by setting the `WS_KEEPALIVE_INTERVAL` env var (auto-injected by AgentService into hosted-agent containers). The constructor argument takes precedence; `0` (or unset) disables keep-alive. Wired through to Hypercorn's `websocket_ping_interval`. -- WebSocket telemetry — structured close-event log line and OpenTelemetry span attributes `azure.ai.agentserver.invocations_ws.{session_id,close_code,duration_ms}`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. -- New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages). +- WebSocket Ping/Pong keep-alive — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` env var (auto-injected by AgentService into hosted-agent containers; surfaced on `app.config.ws_ping_interval` in `azure-ai-agentserver-core>=2.0.0b4`). `0` (or unset) disables keep-alive. Wired through to Hypercorn's `websocket_ping_interval` by `AgentServerHost._build_hypercorn_config`. +- WebSocket per-turn tracing — `app.ws_invocation(websocket)` async context manager opens an `invoke_agent` child span (parity with `POST /invocations`) parented to the connection-level `websocket_session` span, mints a per-turn `invocation_id`, and propagates the OpenTelemetry context through all `await`s inside the block. +- WebSocket telemetry — structured close-event log line and OpenTelemetry span attributes `azure.ai.agentserver.invocations_ws.{session_id, invocation_id, close_code, duration_ms}`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. +- New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages). Both demonstrate `app.ws_invocation`. ### Breaking Changes diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index 3b1caacd2e84..e5e92f117727 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md @@ -221,11 +221,38 @@ app.run() - Registers `/invocations_ws` on the same Starlette host as `/invocations` and `/readiness`. - Calls `await websocket.accept()` before invoking your handler. -- Runs WebSocket Ping/Pong keep-alive in the background — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers) or by passing `InvocationAgentServerHost(ws_ping_interval=...)`. Set the value to `0` (in any source) to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic. +- Runs WebSocket Ping/Pong keep-alive in the background — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers). Set the value to `0` to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic. - Closes the connection cleanly on handler return (close code `1000`) or maps an uncaught handler exception to close code `1011`. - Emits a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `azure.ai.agentserver.invocations_ws.close_code`, and `azure.ai.agentserver.invocations_ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end. - Inherits `/readiness`, OpenTelemetry export, graceful shutdown, and the `x-platform-server` identity header from `azure-ai-agentserver-core`. +### Per-turn tracing with `app.ws_invocation` + +A long-lived WebSocket connection is wrapped by the SDK in a connection-scoped `websocket_session` span. Each *logical turn* (typically one inbound prompt + its reply stream) should be wrapped in `app.ws_invocation(websocket)` so dashboards aggregate equally across the HTTP `POST /invocations` and WebSocket transports — both produce one `invoke_agent` span per turn carrying the same GenAI semantic-convention attributes. + +```python +@app.ws_handler +async def handle(websocket: WebSocket) -> None: + async for msg in websocket.iter_text(): + async with app.ws_invocation(websocket) as turn: + # All OpenTelemetry spans you open inside this block are + # parented to the per-turn ``invoke_agent`` span. + result = await run_agent(msg) + await websocket.send_json({ + "invocation_id": turn.invocation_id, + "result": result, + }) +``` + +The yielded `WSInvocationContext` exposes: + +| Field | Description | +|---|---| +| `invocation_id` | UUID identifying this single turn — recorded on the `invoke_agent` span as `azure.ai.agentserver.invocations_ws.invocation_id`. Echo it in your reply frame so clients can correlate request/response. | +| `session_id` | Per-connection session ID (matches the HTTP `FOUNDRY_AGENT_SESSION_ID` precedence). | + +> `GET /invocations/{id}` and `POST /invocations/{id}/cancel` are HTTP-only endpoints — they apply to *long-running* HTTP invocations. WebSocket clients observe status via reply frames and cancel by closing the socket (the handler observes `asyncio.CancelledError`). + ### Handler signature The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`. The full WebSocket API — `iter_text`, `iter_bytes`, `iter_json`, `send_text`, `send_bytes`, `send_json`, `close`, `headers`, `query_params`, `client`, `state` — is available, so application protocols on top of `invocations_ws` are entirely under your control. @@ -234,13 +261,9 @@ The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`. ### Reference: configuration -| Constructor argument | Default | Description | -|---|---|---| -| `ws_ping_interval` | `0` (disabled) | WebSocket protocol Ping interval in seconds. `0` disables keep-alive. Negative or non-finite values are rejected. When `None`, the SDK reads the `WS_KEEPALIVE_INTERVAL` env var before falling back to disabled. | - | Environment variable | Default | Description | |---|---|---| -| `WS_KEEPALIVE_INTERVAL` | unset (disabled) | Platform-injected override for the WebSocket Ping interval (seconds). `0` disables keep-alive. Ignored when `ws_ping_interval=` is set explicitly. | +| `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/azure/ai/agentserver/invocations/_constants.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py index ec2928aeaf3a..16c33ddbe33c 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py @@ -33,21 +33,13 @@ class InvocationsWSConstants: # Route ROUTE_PATH = "/invocations_ws" - # Environment variable for platform-injected keep-alive override. - # AgentService injects this into every hosted-agent container so the - # platform can tune the WebSocket Ping cadence without redeploying - # the agent code. - # Resolution order: explicit ``ws_ping_interval=`` constructor arg → - # ``WS_KEEPALIVE_INTERVAL`` env var → disabled (``0``). - # ``0`` disables keep-alive. - ENV_WS_KEEPALIVE_INTERVAL = "WS_KEEPALIVE_INTERVAL" - # Close codes (RFC 6455) CLOSE_NORMAL = 1000 # handler returned cleanly CLOSE_INTERNAL_ERROR = 1011 # handler raised an unhandled exception # Span attribute keys ATTR_SPAN_SESSION_ID = "azure.ai.agentserver.invocations_ws.session_id" + ATTR_SPAN_INVOCATION_ID = "azure.ai.agentserver.invocations_ws.invocation_id" ATTR_SPAN_CLOSE_CODE = "azure.ai.agentserver.invocations_ws.close_code" ATTR_SPAN_DURATION_MS = "azure.ai.agentserver.invocations_ws.duration_ms" ATTR_SPAN_ERROR_CODE = "azure.ai.agentserver.invocations_ws.error.code" 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 c443efc5a6c9..f3aac5fb868b 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 @@ -186,14 +186,6 @@ async def ws(websocket: WebSocket) -> None: :param openapi_spec: Optional OpenAPI spec dict. When provided, the spec is served at ``GET /invocations/docs/openapi.json``. :type openapi_spec: Optional[dict[str, Any]] - :param ws_ping_interval: Seconds between WebSocket protocol Ping frames - on ``/invocations_ws``. ``None`` (default) reads the - ``WS_KEEPALIVE_INTERVAL`` environment variable (auto-injected by - AgentService into hosted-agent containers); when neither is set, - keep-alive is disabled. ``0`` disables keep-alive. Configured on - the underlying Hypercorn server so the framing is opcode 0x9 / - 0xA, not application JSON. - :type ws_ping_interval: Optional[float] """ _INSTRUMENTATION_SCOPE = "Azure.AI.AgentServer.Invocations" @@ -202,7 +194,6 @@ def __init__( self, *, openapi_spec: Optional[dict[str, Any]] = None, - ws_ping_interval: Optional[float] = None, **kwargs: Any, ) -> None: self._invoke_fn: Optional[Callable] = None @@ -210,8 +201,10 @@ def __init__( self._cancel_invocation_fn: Optional[Callable] = None self._openapi_spec = openapi_spec - # Initialise WS handler slots (raises ValueError on a bad interval). - self._init_ws_state(ws_ping_interval) + # Initialise WS handler slots (no parameters — the keep-alive + # interval lives on ``AgentConfig`` and is wired into Hypercorn + # by ``AgentServerHost._build_hypercorn_config``). + self._init_ws_state() # Build invocation routes and pass to parent via routes kwarg invocation_routes: list[Any] = [ @@ -247,45 +240,9 @@ def __init__( # --- Invocations startup configuration logging --- logger.info( - "Invocations protocol: openapi_spec_configured=%s, " - "ws_ping_interval=%s", + "Invocations protocol: openapi_spec_configured=%s", self._openapi_spec is not None, - ( - "disabled" - if self._ws_ping_interval == 0 - else f"{self._ws_ping_interval}s" - ), - ) - - # ------------------------------------------------------------------ - # Hypercorn server config (WebSocket Ping/Pong keep-alive) - # ------------------------------------------------------------------ - - def _build_hypercorn_config(self, host: str, port: int) -> object: - """Extend the base Hypercorn config with the WebSocket Ping interval. - - Hypercorn sends WS protocol Ping frames every - ``websocket_ping_interval`` seconds on every active WebSocket - connection — exactly the keep-alive the ``invocations_ws`` spec - requires. ``ws_ping_interval=0`` explicitly disables keep-alive - by setting the attribute to ``None``. - - ``azure-ai-agentserver-core`` pins ``hypercorn>=0.17.0``, which - is guaranteed to expose ``websocket_ping_interval``. - - :param host: Network interface to bind. - :type host: str - :param port: Port to bind. - :type port: int - :return: The configured Hypercorn config. - :rtype: hypercorn.config.Config - """ - config = super()._build_hypercorn_config(host, port) - # Positive float enables protocol-level Ping; None disables. - config.websocket_ping_interval = ( # type: ignore[attr-defined] - self._ws_ping_interval if self._ws_ping_interval > 0 else None ) - return config # ------------------------------------------------------------------ # Handler decorators diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py index 790beb6daa53..02a9a5a32a2b 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py @@ -20,13 +20,13 @@ ``azure.ai.agentserver.invocations_ws.duration_ms``. """ +import contextlib import inspect import logging -import math -import os import time import uuid -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Optional from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState @@ -55,6 +55,25 @@ WSHandler = Callable[[WebSocket], Awaitable[None]] +@dataclass(frozen=True) +class WSInvocationContext: + """Per-turn context yielded by :meth:`_WSHandlerMixin.ws_invocation`. + + :ivar invocation_id: UUID identifying this single user-handler turn. + Echo it in your reply frame so clients can correlate request/response + and so it surfaces in OpenTelemetry as the ``invoke_agent`` span's + ``azure.ai.agentserver.invocations_ws.invocation_id`` attribute. + :vartype invocation_id: str + :ivar session_id: Per-connection session ID (shared across all turns on + the same WebSocket connection; matches the HTTP + ``FOUNDRY_AGENT_SESSION_ID`` precedence). + :vartype session_id: str + """ + + invocation_id: str + session_id: str + + def _safe_set_attrs(span: Any, attrs: dict[str, Any]) -> None: """Best-effort ``span.set_attribute`` for a batch of attributes. @@ -87,61 +106,16 @@ class _WSHandlerMixin(_MixinBase): # Slots populated by __init__. _ws_fn: Optional[WSHandler] - _ws_ping_interval: float - def _init_ws_state(self, ws_ping_interval: Optional[float]) -> None: + def _init_ws_state(self) -> None: """Initialize WS handler slots. - Resolution order for the keep-alive interval: - - 1. Explicit *ws_ping_interval* constructor argument (when not ``None``). - 2. ``WS_KEEPALIVE_INTERVAL`` environment variable (auto-injected by - AgentService into every hosted-agent container). - 3. Disabled (``0``) — no Ping/Pong frames are sent. - - ``0`` (from any source) disables keep-alive entirely. - - :param ws_ping_interval: Seconds between WS protocol Ping frames. - ``None`` selects the env var or the default. - :type ws_ping_interval: Optional[float] - :raises ValueError: If the resolved value is negative, non-finite, - or (for the env var) not parseable as a number. + The keep-alive interval lives on :class:`AgentConfig` and is + wired into Hypercorn by + :meth:`AgentServerHost._build_hypercorn_config` — there is no + per-mixin state to populate here besides the handler slot. """ self._ws_fn = None - if ws_ping_interval is None: - env_raw = os.environ.get(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL) - if env_raw is None or env_raw == "": - resolved = 0.0 - else: - try: - resolved = float(env_raw) - except ValueError as exc: - raise ValueError( - f"Invalid value for " - f"{InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL}: " - f"{env_raw!r} (expected a non-negative number)" - ) from exc - if math.isnan(resolved) or math.isinf(resolved) or resolved < 0.0: - raise ValueError( - f"Invalid value for " - f"{InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL}: " - f"{env_raw!r} (expected a non-negative finite number)" - ) - else: - try: - resolved = float(ws_ping_interval) - except (TypeError, ValueError) as exc: - raise ValueError( - f"ws_ping_interval must be a number, got {ws_ping_interval!r}" - ) from exc - # Reject negative / NaN / inf — those are programming errors that - # would silently mis-configure the keep-alive. - if math.isnan(resolved) or math.isinf(resolved) or resolved < 0.0: - raise ValueError( - f"ws_ping_interval must be a non-negative finite number, " - f"got {ws_ping_interval!r}" - ) - self._ws_ping_interval = resolved # ------------------------------------------------------------------ # Public configuration accessor @@ -151,10 +125,12 @@ def _init_ws_state(self, ws_ping_interval: Optional[float]) -> None: def ws_ping_interval(self) -> float: """Configured WebSocket Ping interval in seconds (``0`` = disabled). + Convenience alias for ``self.config.ws_ping_interval``. + :return: The configured interval, or ``0`` when keep-alive is disabled. :rtype: float """ - return self._ws_ping_interval + return float(self.config.ws_ping_interval) # ------------------------------------------------------------------ # Decorator @@ -217,6 +193,86 @@ async def handle(websocket: WebSocket) -> None: self._ensure_ws_route_registered() return fn + # ------------------------------------------------------------------ + # Per-turn tracing helper + # ------------------------------------------------------------------ + + @contextlib.asynccontextmanager + async def ws_invocation( + self, + websocket: WebSocket, + *, + session_id: Optional[str] = None, + operation_name: str = "invoke_agent", + ) -> AsyncIterator[WSInvocationContext]: + """Open a per-turn ``invoke_agent`` span for a single WebSocket exchange. + + A long-lived WebSocket connection is wrapped in a connection-scoped + ``websocket_session`` span by the SDK. Each *logical turn* on the + connection (e.g. one inbound prompt frame and its reply stream) is + a child ``invoke_agent`` span — matching the HTTP ``POST /invocations`` + tracing model so dashboards aggregate equally over both transports. + ``ws_invocation`` is the opt-in helper that opens that child span and + propagates the OpenTelemetry context through every ``await`` inside + the ``async with`` block. + + Usage:: + + @app.ws_handler + async def handle(websocket: WebSocket) -> None: + async for msg in websocket.iter_text(): + async with app.ws_invocation(websocket) as turn: + # All child spans created here are parented to the + # per-turn ``invoke_agent`` span. + result = await run_agent(msg) + await websocket.send_json( + {"invocation_id": turn.invocation_id, "result": result} + ) + + :param websocket: The accepted WebSocket (used to honour the client's + ``traceparent`` for the initial connection; per-turn spans inherit + the connection's active span automatically). + :type websocket: ~starlette.websockets.WebSocket + :keyword session_id: Override the connection-level session ID for this + turn (rarely needed). Defaults to ``self.config.session_id`` or + a freshly generated UUID. + :paramtype session_id: Optional[str] + :keyword operation_name: ``gen_ai.operation.name`` attribute value. + Defaults to ``"invoke_agent"`` for parity with HTTP invocations. + :paramtype operation_name: str + :return: An async context manager yielding a :class:`WSInvocationContext`. + :rtype: AsyncIterator[WSInvocationContext] + """ + invocation_id = str(uuid.uuid4()) + resolved_session_id = session_id or self.config.session_id or str(uuid.uuid4()) + # ``request_span`` uses ``start_as_current_span``, so awaits inside + # the ``with`` see the span as the active OTel context. We honour + # ``websocket.headers`` only as a fallback parent (the connection + # span is already current, so it wins for child parenting). + with self.request_span( + websocket.headers, + invocation_id, + "invoke_agent", + operation_name=operation_name, + session_id=resolved_session_id, + ) as otel_span: + _safe_set_attrs(otel_span, { + InvocationsWSConstants.ATTR_SPAN_INVOCATION_ID: invocation_id, + InvocationsWSConstants.ATTR_SPAN_SESSION_ID: resolved_session_id, + }) + try: + yield WSInvocationContext( + invocation_id=invocation_id, + session_id=resolved_session_id, + ) + except Exception as exc: # pylint: disable=broad-exception-caught + # ``request_span`` records the exception via OTel's standard + # context-manager protocol on __exit__; nothing extra needed. + _safe_set_attrs(otel_span, { + "error.type": type(exc).__name__, + }) + raise + def _ensure_ws_route_registered(self) -> None: """Append the ``/invocations_ws`` WebSocketRoute to the router (idempotent). diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py index 903b635d9e7b..7e3f03088f9c 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py @@ -127,6 +127,10 @@ async def _stream_tokens( ) -> None: """Stream tokens for one prompt; cancellable via ``asyncio.CancelledError``. + Each prompt is wrapped in a per-turn ``invoke_agent`` span via + ``app.ws_invocation`` so dashboards aggregate equally across the WS and + HTTP transports. + :param websocket: The accepted WebSocket. :type websocket: ~starlette.websockets.WebSocket :param prompt_id: Caller-supplied prompt identifier (echoed in events). @@ -134,30 +138,35 @@ async def _stream_tokens( :param text: The user prompt to "generate" against. :type text: str """ - try: - async for token in _generate_tokens(text): - await websocket.send_json( - {"type": "token", "id": prompt_id, "token": token}, - ) - await websocket.send_json({"type": "done", "id": prompt_id}) - except asyncio.CancelledError: - # Best-effort: tell the client we honoured their cancel. Suppress - # ordinary send errors (the socket may already be closed) and re-raise - # so the caller observes the cancellation. We deliberately do NOT - # catch ``BaseException`` here so process-level signals - # (``KeyboardInterrupt`` / ``SystemExit``) propagate normally. - with contextlib.suppress(Exception): - await websocket.send_json({"type": "cancelled", "id": prompt_id}) - raise - except Exception as exc: # pylint: disable=broad-exception-caught - # Surface generation errors to the client as a structured frame so - # the connection survives a single-prompt failure (parity with the - # ``error`` reply the reader emits on bad input). - with contextlib.suppress(Exception): + async with app.ws_invocation(websocket) as turn: + try: + async for token in _generate_tokens(text): + await websocket.send_json( + {"type": "token", "id": prompt_id, "token": token, + "invocation_id": turn.invocation_id}, + ) await websocket.send_json( - {"type": "error", "id": prompt_id, "message": str(exc)}, + {"type": "done", "id": prompt_id, + "invocation_id": turn.invocation_id}, ) - logger.exception("_stream_tokens failed for prompt %s", prompt_id) + except asyncio.CancelledError: + # Best-effort: tell the client we honoured their cancel. Suppress + # ordinary send errors (the socket may already be closed) and re-raise + # so the caller observes the cancellation. We deliberately do NOT + # catch ``BaseException`` here so process-level signals + # (``KeyboardInterrupt`` / ``SystemExit``) propagate normally. + with contextlib.suppress(Exception): + await websocket.send_json({"type": "cancelled", "id": prompt_id}) + raise + except Exception as exc: # pylint: disable=broad-exception-caught + # Surface generation errors to the client as a structured frame so + # the connection survives a single-prompt failure (parity with the + # ``error`` reply the reader emits on bad input). + with contextlib.suppress(Exception): + await websocket.send_json( + {"type": "error", "id": prompt_id, "message": str(exc)}, + ) + logger.exception("_stream_tokens failed for prompt %s", prompt_id) async def _reader( diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py index 269d9a1d4c21..c0401953b999 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py @@ -49,11 +49,17 @@ async def handle_ws(websocket: WebSocket) -> None: runs and will close it cleanly on return. An uncaught exception is mapped to RFC 6455 close code 1011. WebSocket protocol-level Ping/Pong keep-alive is disabled by default; enable it by setting the - ``WS_KEEPALIVE_INTERVAL`` environment variable or by passing - ``InvocationAgentServerHost(ws_ping_interval=)``. + ``WS_KEEPALIVE_INTERVAL`` environment variable. + + Each inbound frame opens a per-turn ``invoke_agent`` span via + ``app.ws_invocation`` so traces line up 1:1 with the HTTP + ``POST /invocations`` endpoint above. """ async for message in websocket.iter_text(): - await websocket.send_text(message) + async with app.ws_invocation(websocket) as turn: + await websocket.send_text( + f"[{turn.invocation_id}] {message}" + ) if __name__ == "__main__": diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invocation_helper.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invocation_helper.py new file mode 100644 index 000000000000..d45e70af72aa --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invocation_helper.py @@ -0,0 +1,197 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for the per-turn ``ws_invocation`` helper. + +Each ``ws_invocation()`` block opens an ``invoke_agent`` span (mirroring the +HTTP ``POST /invocations`` model) parented to the connection-level +``websocket_session`` span, mints an ``invocation_id`` for the turn, and +ensures the OTel context is current for every ``await`` inside the block. +""" +import os +from unittest.mock import patch + +from opentelemetry import trace as otel_trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from starlette.testclient import TestClient +from starlette.websockets import WebSocket + +from azure.ai.agentserver.invocations import InvocationAgentServerHost +from azure.ai.agentserver.invocations._constants import InvocationsWSConstants + + +# --------------------------------------------------------------------------- +# Test fixture: module-scoped OTel provider + in-memory exporter +# --------------------------------------------------------------------------- + +_EXPORTER = InMemorySpanExporter() +_PROVIDER = TracerProvider() +_PROVIDER.add_span_processor(SimpleSpanProcessor(_EXPORTER)) +otel_trace.set_tracer_provider(_PROVIDER) + + +def _make_server_with_invocations() -> InvocationAgentServerHost: + """Build a tracing-enabled WS host whose handler opens one ``invoke_agent`` + span per inbound text frame via ``ws_invocation``.""" + with patch.dict( + os.environ, + {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, + ): + with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): + server = InvocationAgentServerHost() + + @server.ws_handler + async def handler(websocket: WebSocket) -> None: + async for msg in websocket.iter_text(): + async with server.ws_invocation(websocket) as turn: + await websocket.send_json( + {"invocation_id": turn.invocation_id, "echo": msg} + ) + + return server + + +def _get_spans(): + return _EXPORTER.get_finished_spans() + + +def _session_spans(spans): + return [s for s in spans if "websocket_session" in s.name] + + +def _invoke_spans(spans): + return [s for s in spans if "invoke_agent" in s.name] + + +# --------------------------------------------------------------------------- +# Span shape +# --------------------------------------------------------------------------- + +def test_ws_invocation_opens_invoke_agent_span_per_turn(): + """One ``invoke_agent`` span is recorded per ``ws_invocation`` block.""" + _EXPORTER.clear() + server = _make_server_with_invocations() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_json() + ws.send_text("again") + ws.receive_json() + + spans = _get_spans() + invokes = _invoke_spans(spans) + assert len(invokes) == 2, f"expected 2 invoke_agent spans, got {[s.name for s in spans]}" + + +def test_ws_invocation_id_is_set_on_invoke_span(): + """``invoke_agent`` span carries the per-turn ``invocation_id``.""" + _EXPORTER.clear() + server = _make_server_with_invocations() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + payload = ws.receive_json() + + invokes = _invoke_spans(_get_spans()) + assert invokes + span = invokes[-1] + attr = span.attributes.get(InvocationsWSConstants.ATTR_SPAN_INVOCATION_ID) + assert attr == payload["invocation_id"] + + +def test_ws_invocation_unique_per_turn(): + """Each turn mints a fresh ``invocation_id``.""" + _EXPORTER.clear() + server = _make_server_with_invocations() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("a") + a = ws.receive_json()["invocation_id"] + ws.send_text("b") + b = ws.receive_json()["invocation_id"] + ws.send_text("c") + c = ws.receive_json()["invocation_id"] + + assert len({a, b, c}) == 3 + + +def test_ws_invocation_span_is_child_of_websocket_session(): + """``invoke_agent`` is parented to the connection-level ``websocket_session``.""" + _EXPORTER.clear() + server = _make_server_with_invocations() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_json() + + spans = _get_spans() + sessions = _session_spans(spans) + invokes = _invoke_spans(spans) + assert sessions and invokes + session_span = sessions[-1] + invoke_span = invokes[-1] + # Parent-child via OTel SpanContext.span_id linkage. + assert invoke_span.parent is not None + assert invoke_span.parent.span_id == session_span.context.span_id + + +# --------------------------------------------------------------------------- +# Session ID propagation +# --------------------------------------------------------------------------- + +def test_ws_invocation_inherits_connection_session_id(monkeypatch): + """``ws_invocation`` defaults to ``app.config.session_id`` when set.""" + monkeypatch.setenv("FOUNDRY_AGENT_SESSION_ID", "fixed-session") + _EXPORTER.clear() + server = _make_server_with_invocations() + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_json() + + invokes = _invoke_spans(_get_spans()) + assert invokes + assert ( + invokes[-1].attributes.get(InvocationsWSConstants.ATTR_SPAN_SESSION_ID) + == "fixed-session" + ) + + +# --------------------------------------------------------------------------- +# OTel context propagation through awaits +# --------------------------------------------------------------------------- + +def test_ws_invocation_propagates_otel_context_across_awaits(): + """A user-opened child span inside ``ws_invocation`` is parented to + the per-turn ``invoke_agent`` span (proves OTel context survives awaits).""" + _EXPORTER.clear() + with patch.dict( + os.environ, + {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, + ): + with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): + server = InvocationAgentServerHost() + + tracer = otel_trace.get_tracer("test") + + @server.ws_handler + async def handler(websocket: WebSocket) -> None: + async for msg in websocket.iter_text(): + async with server.ws_invocation(websocket): + # User-opened child span — must parent to the invoke_agent span. + with tracer.start_as_current_span("user_work"): + await websocket.send_text(msg) + + client = TestClient(server) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text("hi") + ws.receive_text() + + spans = _get_spans() + invoke = [s for s in spans if s.name.startswith("invoke_agent") or "invoke_agent" in s.name] + user_work = [s for s in spans if s.name == "user_work"] + assert invoke and user_work + assert user_work[-1].parent is not None + assert user_work[-1].parent.span_id == invoke[-1].context.span_id diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py index e8ead6148b9e..2a1f5171dcf5 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py @@ -3,89 +3,105 @@ # --------------------------------------------------------- """Tests for the ``ws_ping_interval`` configuration knob. -Parity with :mod:`tests.test_request_limits` / configuration tests in -:mod:`tests.test_graceful_shutdown` — validates the Hypercorn ``ws_ping_interval`` -parameter parsing and wiring. +The interval lives on :class:`~azure.ai.agentserver.core.AgentConfig` and is +resolved from the ``WS_KEEPALIVE_INTERVAL`` environment variable +(auto-injected by AgentService into hosted-agent containers). The Hypercorn +wiring lives in :class:`AgentServerHost._build_hypercorn_config`; the +invocations package exposes a convenience ``app.ws_ping_interval`` property +that reads the same config value. """ import pytest from azure.ai.agentserver.invocations import InvocationAgentServerHost -from azure.ai.agentserver.invocations._constants import InvocationsWSConstants + + +_ENV_NAME = "WS_KEEPALIVE_INTERVAL" # --------------------------------------------------------------------------- -# Default / accepted values +# Default / accepted values (env-driven) # --------------------------------------------------------------------------- -def test_ws_ping_interval_default_is_disabled(): - """Default ping interval is 0 (disabled) when no arg / env var is set.""" +def test_ws_ping_interval_default_is_disabled(monkeypatch): + """Default ping interval is 0 (disabled) when the env var is not set.""" + monkeypatch.delenv(_ENV_NAME, raising=False) app = InvocationAgentServerHost() assert app.ws_ping_interval == 0.0 -def test_ws_ping_interval_custom_value(): - """``ws_ping_interval`` is honoured.""" - app = InvocationAgentServerHost(ws_ping_interval=15) +def test_ws_ping_interval_custom_value(monkeypatch): + """``WS_KEEPALIVE_INTERVAL`` is honoured.""" + monkeypatch.setenv(_ENV_NAME, "15") + app = InvocationAgentServerHost() assert app.ws_ping_interval == 15.0 -def test_ws_ping_interval_zero_disables_keepalive(): - """``ws_ping_interval=0`` disables WS-level keep-alive.""" - app = InvocationAgentServerHost(ws_ping_interval=0) +def test_ws_ping_interval_zero_disables_keepalive(monkeypatch): + """``WS_KEEPALIVE_INTERVAL=0`` disables WS-level keep-alive.""" + monkeypatch.setenv(_ENV_NAME, "0") + app = InvocationAgentServerHost() assert app.ws_ping_interval == 0.0 -def test_ws_ping_interval_float_value_accepted(): +def test_ws_ping_interval_float_value_accepted(monkeypatch): """Fractional intervals are coerced to ``float``.""" - app = InvocationAgentServerHost(ws_ping_interval=12.5) + monkeypatch.setenv(_ENV_NAME, "12.5") + app = InvocationAgentServerHost() assert app.ws_ping_interval == 12.5 config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 assert getattr(config, "websocket_ping_interval", None) == 12.5 -# --------------------------------------------------------------------------- -# Rejected values (validation) -# --------------------------------------------------------------------------- - -def test_ws_ping_interval_negative_rejected(): - """Negative intervals are programming errors.""" - with pytest.raises(ValueError, match="non-negative"): - InvocationAgentServerHost(ws_ping_interval=-1) - +def test_ws_ping_interval_empty_env_uses_default(monkeypatch): + """An empty env-var value falls back to the default (disabled).""" + monkeypatch.setenv(_ENV_NAME, "") + app = InvocationAgentServerHost() + assert app.ws_ping_interval == 0.0 -def test_ws_ping_interval_nan_rejected(): - """``ws_ping_interval=nan`` is a programming error.""" - with pytest.raises(ValueError, match="non-negative"): - InvocationAgentServerHost(ws_ping_interval=float("nan")) +# --------------------------------------------------------------------------- +# Rejected values (validation surfaces at AgentConfig.from_env) +# --------------------------------------------------------------------------- -def test_ws_ping_interval_inf_rejected(): - """``ws_ping_interval=inf`` is a programming error.""" +def test_ws_ping_interval_negative_rejected(monkeypatch): + """Negative env-var values are programming errors.""" + monkeypatch.setenv(_ENV_NAME, "-1") with pytest.raises(ValueError, match="non-negative"): - InvocationAgentServerHost(ws_ping_interval=float("inf")) + InvocationAgentServerHost() -def test_ws_ping_interval_non_numeric_rejected(): - """Strings or non-numeric values surface as ``ValueError``.""" - with pytest.raises(ValueError, match="must be a number"): - InvocationAgentServerHost(ws_ping_interval="thirty") # type: ignore[arg-type] +def test_ws_ping_interval_non_numeric_rejected(monkeypatch): + """Non-numeric env-var values surface as ``ValueError`` at startup.""" + monkeypatch.setenv(_ENV_NAME, "thirty") + with pytest.raises(ValueError, match=_ENV_NAME): + InvocationAgentServerHost() # --------------------------------------------------------------------------- -# Hypercorn config wiring +# Hypercorn config wiring (delegated to core's _build_hypercorn_config) # --------------------------------------------------------------------------- -def test_ws_ping_interval_propagates_to_hypercorn_config(): +def test_ws_ping_interval_propagates_to_hypercorn_config(monkeypatch): """The configured interval lands on the Hypercorn server config.""" - app = InvocationAgentServerHost(ws_ping_interval=20) + monkeypatch.setenv(_ENV_NAME, "20") + app = InvocationAgentServerHost() config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 # Hypercorn ≥0.14 exposes this attribute. assert getattr(config, "websocket_ping_interval", None) == 20.0 -def test_ws_ping_interval_zero_sets_attribute_to_none(): +def test_ws_ping_interval_zero_sets_attribute_to_none(monkeypatch): """Zero explicitly sets Hypercorn's ``websocket_ping_interval`` to ``None``.""" - app = InvocationAgentServerHost(ws_ping_interval=0) + monkeypatch.setenv(_ENV_NAME, "0") + app = InvocationAgentServerHost() + config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 + assert config.websocket_ping_interval is None # type: ignore[attr-defined] + + +def test_ws_ping_interval_default_wires_none_into_hypercorn(monkeypatch): + """With no env var, Hypercorn config has ``websocket_ping_interval = None``.""" + monkeypatch.delenv(_ENV_NAME, raising=False) + app = InvocationAgentServerHost() config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 assert config.websocket_ping_interval is None # type: ignore[attr-defined] @@ -94,79 +110,16 @@ def test_ws_ping_interval_zero_sets_attribute_to_none(): # Property surface # --------------------------------------------------------------------------- -def test_ws_ping_interval_property_is_read_only(): +def test_ws_ping_interval_property_is_read_only(monkeypatch): """``ws_ping_interval`` is exposed only as a property (no setter).""" - app = InvocationAgentServerHost(ws_ping_interval=20) + monkeypatch.setenv(_ENV_NAME, "20") + app = InvocationAgentServerHost() with pytest.raises(AttributeError): app.ws_ping_interval = 10 # type: ignore[misc] -# --------------------------------------------------------------------------- -# WS_KEEPALIVE_INTERVAL env var (platform-injected override) -# --------------------------------------------------------------------------- - -def test_ws_keepalive_interval_env_var_used_when_arg_is_none(monkeypatch): - """Env var drives the interval when no explicit arg is passed.""" - monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "45") - app = InvocationAgentServerHost() - assert app.ws_ping_interval == 45.0 - - -def test_ws_keepalive_interval_env_var_zero_disables(monkeypatch): - """``WS_KEEPALIVE_INTERVAL=0`` disables keep-alive (mirrors SSE).""" - monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "0") +def test_ws_ping_interval_mirrors_config(monkeypatch): + """The property is a thin alias for ``app.config.ws_ping_interval``.""" + monkeypatch.setenv(_ENV_NAME, "7.5") app = InvocationAgentServerHost() - assert app.ws_ping_interval == 0.0 - config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 - assert getattr(config, "websocket_ping_interval", None) is None - - -def test_ws_keepalive_interval_env_var_float_accepted(monkeypatch): - """Fractional env-var values are coerced to ``float``.""" - monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "12.5") - app = InvocationAgentServerHost() - assert app.ws_ping_interval == 12.5 - - -def test_ws_keepalive_interval_env_var_empty_uses_default(monkeypatch): - """An empty env-var value falls back to the default (disabled).""" - monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "") - app = InvocationAgentServerHost() - assert app.ws_ping_interval == 0.0 - - -def test_ws_keepalive_interval_env_var_invalid_rejected(monkeypatch): - """Non-numeric env-var values surface as ``ValueError`` at startup.""" - monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "thirty") - with pytest.raises(ValueError, match="WS_KEEPALIVE_INTERVAL"): - InvocationAgentServerHost() - - -def test_ws_keepalive_interval_env_var_negative_rejected(monkeypatch): - """Negative env-var values are programming errors.""" - monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "-5") - with pytest.raises(ValueError, match="WS_KEEPALIVE_INTERVAL"): - InvocationAgentServerHost() - - -def test_ws_ping_interval_arg_overrides_env_var(monkeypatch): - """Explicit constructor arg wins over the env var.""" - monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "45") - app = InvocationAgentServerHost(ws_ping_interval=10) - assert app.ws_ping_interval == 10.0 - - -def test_ws_ping_interval_default_wires_none_into_hypercorn(monkeypatch): - """With neither arg nor env var, Hypercorn config has ``websocket_ping_interval = None``.""" - monkeypatch.delenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, raising=False) - app = InvocationAgentServerHost() - config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 - assert config.websocket_ping_interval is None # type: ignore[attr-defined] - - -def test_ws_ping_interval_positive_wires_into_hypercorn(): - """A positive interval propagates as a float to Hypercorn config.""" - app = InvocationAgentServerHost(ws_ping_interval=12) - config = app._build_hypercorn_config("0.0.0.0", 8088) # noqa: SLF001 - assert config.websocket_ping_interval == 12.0 # type: ignore[attr-defined] - + assert app.ws_ping_interval == app.config.ws_ping_interval == 7.5 From c49b804154b0c9be32a2448476f2b0280798f29e Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Mon, 18 May 2026 09:36:24 +0000 Subject: [PATCH 06/13] Revert per-turn ws_invocation helper: WS has only one connection-scoped span Per design feedback: WebSocket has no per-message invocation_id. A WS connection is a single long-lived session, so a single websocket_session span per connection is sufficient. HTTP-only GET/cancel do not apply. - Drop app.ws_invocation() async context manager and WSInvocationContext - Drop InvocationsWSConstants.ATTR_SPAN_INVOCATION_ID - Revert sample changes (no invocation_id echo) - Replace README 'Per-turn tracing with app.ws_invocation' section with 'Per-connection tracing' explaining the single-span model + why GET/cancel are HTTP-only - Update CHANGELOG bullets accordingly - Drop tests/test_ws_invocation_helper.py --- .../CHANGELOG.md | 5 +- .../README.md | 27 +-- .../ai/agentserver/invocations/_constants.py | 1 - .../agentserver/invocations/_invocation_ws.py | 103 +-------- .../ws_bidirectional_streaming_agent.py | 53 ++--- .../ws_invoke_agent/ws_invoke_agent.py | 9 +- .../tests/test_ws_invocation_helper.py | 197 ------------------ 7 files changed, 29 insertions(+), 366 deletions(-) delete mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invocation_helper.py diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index a480e7b543d6..5cd275b49f73 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -7,9 +7,8 @@ - Error source classification headers: All HTTP error responses now include `x-platform-error-source` with a value of `user`, `platform`, or `upstream` to indicate which component caused the error. Developer handler exceptions and missing handler registrations are classified as `upstream`. Exceptions tagged with the platform error tag are classified as `platform` and additionally include `x-platform-error-detail` with truncated exception details (max 2048 characters) for diagnostics. - WebSocket protocol support — `InvocationAgentServerHost` now hosts `/invocations_ws` alongside `POST /invocations`. Register the handler with the new `@app.ws_handler` decorator. The route is registered lazily on first decoration, so hosts without a registered handler return HTTP 404. - WebSocket Ping/Pong keep-alive — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` env var (auto-injected by AgentService into hosted-agent containers; surfaced on `app.config.ws_ping_interval` in `azure-ai-agentserver-core>=2.0.0b4`). `0` (or unset) disables keep-alive. Wired through to Hypercorn's `websocket_ping_interval` by `AgentServerHost._build_hypercorn_config`. -- WebSocket per-turn tracing — `app.ws_invocation(websocket)` async context manager opens an `invoke_agent` child span (parity with `POST /invocations`) parented to the connection-level `websocket_session` span, mints a per-turn `invocation_id`, and propagates the OpenTelemetry context through all `await`s inside the block. -- WebSocket telemetry — structured close-event log line and OpenTelemetry span attributes `azure.ai.agentserver.invocations_ws.{session_id, invocation_id, close_code, duration_ms}`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. -- New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages). Both demonstrate `app.ws_invocation`. +- WebSocket telemetry — single connection-scoped `websocket_session` OpenTelemetry span per WS connection, plus a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. There is no per-message `invocation_id` for WS: a connection is a single long-lived session, not a sequence of HTTP invocations. +- New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages). ### Breaking Changes diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index e5e92f117727..a46ba840f0da 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md @@ -226,32 +226,11 @@ app.run() - Emits a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `azure.ai.agentserver.invocations_ws.close_code`, and `azure.ai.agentserver.invocations_ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end. - Inherits `/readiness`, OpenTelemetry export, graceful shutdown, and the `x-platform-server` identity header from `azure-ai-agentserver-core`. -### Per-turn tracing with `app.ws_invocation` +### Per-connection tracing -A long-lived WebSocket connection is wrapped by the SDK in a connection-scoped `websocket_session` span. Each *logical turn* (typically one inbound prompt + its reply stream) should be wrapped in `app.ws_invocation(websocket)` so dashboards aggregate equally across the HTTP `POST /invocations` and WebSocket transports — both produce one `invoke_agent` span per turn carrying the same GenAI semantic-convention attributes. +A WebSocket connection is wrapped by the SDK in a single connection-scoped `websocket_session` OpenTelemetry span. The span carries the GenAI semantic-convention attributes plus `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms`. Any child spans your handler opens — e.g. via `opentelemetry.trace.get_tracer(...).start_as_current_span(...)` — are automatically parented to the connection span. -```python -@app.ws_handler -async def handle(websocket: WebSocket) -> None: - async for msg in websocket.iter_text(): - async with app.ws_invocation(websocket) as turn: - # All OpenTelemetry spans you open inside this block are - # parented to the per-turn ``invoke_agent`` span. - result = await run_agent(msg) - await websocket.send_json({ - "invocation_id": turn.invocation_id, - "result": result, - }) -``` - -The yielded `WSInvocationContext` exposes: - -| Field | Description | -|---|---| -| `invocation_id` | UUID identifying this single turn — recorded on the `invoke_agent` span as `azure.ai.agentserver.invocations_ws.invocation_id`. Echo it in your reply frame so clients can correlate request/response. | -| `session_id` | Per-connection session ID (matches the HTTP `FOUNDRY_AGENT_SESSION_ID` precedence). | - -> `GET /invocations/{id}` and `POST /invocations/{id}/cancel` are HTTP-only endpoints — they apply to *long-running* HTTP invocations. WebSocket clients observe status via reply frames and cancel by closing the socket (the handler observes `asyncio.CancelledError`). +There is no per-message `invocation_id` for WebSocket: a WS connection is a single long-lived session, not a sequence of independently-identifiable HTTP invocations. The HTTP-only endpoints `GET /invocations/{id}` and `POST /invocations/{id}/cancel` do not apply to WS sessions — clients observe status via reply frames and cancel by closing the socket (the handler observes `asyncio.CancelledError`). ### Handler signature diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py index 16c33ddbe33c..a704df3c530d 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py @@ -39,7 +39,6 @@ class InvocationsWSConstants: # Span attribute keys ATTR_SPAN_SESSION_ID = "azure.ai.agentserver.invocations_ws.session_id" - ATTR_SPAN_INVOCATION_ID = "azure.ai.agentserver.invocations_ws.invocation_id" ATTR_SPAN_CLOSE_CODE = "azure.ai.agentserver.invocations_ws.close_code" ATTR_SPAN_DURATION_MS = "azure.ai.agentserver.invocations_ws.duration_ms" ATTR_SPAN_ERROR_CODE = "azure.ai.agentserver.invocations_ws.error.code" diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py index 02a9a5a32a2b..11e5eac2ebde 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py @@ -20,13 +20,11 @@ ``azure.ai.agentserver.invocations_ws.duration_ms``. """ -import contextlib import inspect import logging import time import uuid -from collections.abc import AsyncIterator, Awaitable, Callable -from dataclasses import dataclass +from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any, Optional from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState @@ -55,25 +53,6 @@ WSHandler = Callable[[WebSocket], Awaitable[None]] -@dataclass(frozen=True) -class WSInvocationContext: - """Per-turn context yielded by :meth:`_WSHandlerMixin.ws_invocation`. - - :ivar invocation_id: UUID identifying this single user-handler turn. - Echo it in your reply frame so clients can correlate request/response - and so it surfaces in OpenTelemetry as the ``invoke_agent`` span's - ``azure.ai.agentserver.invocations_ws.invocation_id`` attribute. - :vartype invocation_id: str - :ivar session_id: Per-connection session ID (shared across all turns on - the same WebSocket connection; matches the HTTP - ``FOUNDRY_AGENT_SESSION_ID`` precedence). - :vartype session_id: str - """ - - invocation_id: str - session_id: str - - def _safe_set_attrs(span: Any, attrs: dict[str, Any]) -> None: """Best-effort ``span.set_attribute`` for a batch of attributes. @@ -193,86 +172,6 @@ async def handle(websocket: WebSocket) -> None: self._ensure_ws_route_registered() return fn - # ------------------------------------------------------------------ - # Per-turn tracing helper - # ------------------------------------------------------------------ - - @contextlib.asynccontextmanager - async def ws_invocation( - self, - websocket: WebSocket, - *, - session_id: Optional[str] = None, - operation_name: str = "invoke_agent", - ) -> AsyncIterator[WSInvocationContext]: - """Open a per-turn ``invoke_agent`` span for a single WebSocket exchange. - - A long-lived WebSocket connection is wrapped in a connection-scoped - ``websocket_session`` span by the SDK. Each *logical turn* on the - connection (e.g. one inbound prompt frame and its reply stream) is - a child ``invoke_agent`` span — matching the HTTP ``POST /invocations`` - tracing model so dashboards aggregate equally over both transports. - ``ws_invocation`` is the opt-in helper that opens that child span and - propagates the OpenTelemetry context through every ``await`` inside - the ``async with`` block. - - Usage:: - - @app.ws_handler - async def handle(websocket: WebSocket) -> None: - async for msg in websocket.iter_text(): - async with app.ws_invocation(websocket) as turn: - # All child spans created here are parented to the - # per-turn ``invoke_agent`` span. - result = await run_agent(msg) - await websocket.send_json( - {"invocation_id": turn.invocation_id, "result": result} - ) - - :param websocket: The accepted WebSocket (used to honour the client's - ``traceparent`` for the initial connection; per-turn spans inherit - the connection's active span automatically). - :type websocket: ~starlette.websockets.WebSocket - :keyword session_id: Override the connection-level session ID for this - turn (rarely needed). Defaults to ``self.config.session_id`` or - a freshly generated UUID. - :paramtype session_id: Optional[str] - :keyword operation_name: ``gen_ai.operation.name`` attribute value. - Defaults to ``"invoke_agent"`` for parity with HTTP invocations. - :paramtype operation_name: str - :return: An async context manager yielding a :class:`WSInvocationContext`. - :rtype: AsyncIterator[WSInvocationContext] - """ - invocation_id = str(uuid.uuid4()) - resolved_session_id = session_id or self.config.session_id or str(uuid.uuid4()) - # ``request_span`` uses ``start_as_current_span``, so awaits inside - # the ``with`` see the span as the active OTel context. We honour - # ``websocket.headers`` only as a fallback parent (the connection - # span is already current, so it wins for child parenting). - with self.request_span( - websocket.headers, - invocation_id, - "invoke_agent", - operation_name=operation_name, - session_id=resolved_session_id, - ) as otel_span: - _safe_set_attrs(otel_span, { - InvocationsWSConstants.ATTR_SPAN_INVOCATION_ID: invocation_id, - InvocationsWSConstants.ATTR_SPAN_SESSION_ID: resolved_session_id, - }) - try: - yield WSInvocationContext( - invocation_id=invocation_id, - session_id=resolved_session_id, - ) - except Exception as exc: # pylint: disable=broad-exception-caught - # ``request_span`` records the exception via OTel's standard - # context-manager protocol on __exit__; nothing extra needed. - _safe_set_attrs(otel_span, { - "error.type": type(exc).__name__, - }) - raise - def _ensure_ws_route_registered(self) -> None: """Append the ``/invocations_ws`` WebSocketRoute to the router (idempotent). diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py index 7e3f03088f9c..903b635d9e7b 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py @@ -127,10 +127,6 @@ async def _stream_tokens( ) -> None: """Stream tokens for one prompt; cancellable via ``asyncio.CancelledError``. - Each prompt is wrapped in a per-turn ``invoke_agent`` span via - ``app.ws_invocation`` so dashboards aggregate equally across the WS and - HTTP transports. - :param websocket: The accepted WebSocket. :type websocket: ~starlette.websockets.WebSocket :param prompt_id: Caller-supplied prompt identifier (echoed in events). @@ -138,35 +134,30 @@ async def _stream_tokens( :param text: The user prompt to "generate" against. :type text: str """ - async with app.ws_invocation(websocket) as turn: - try: - async for token in _generate_tokens(text): - await websocket.send_json( - {"type": "token", "id": prompt_id, "token": token, - "invocation_id": turn.invocation_id}, - ) + try: + async for token in _generate_tokens(text): await websocket.send_json( - {"type": "done", "id": prompt_id, - "invocation_id": turn.invocation_id}, + {"type": "token", "id": prompt_id, "token": token}, ) - except asyncio.CancelledError: - # Best-effort: tell the client we honoured their cancel. Suppress - # ordinary send errors (the socket may already be closed) and re-raise - # so the caller observes the cancellation. We deliberately do NOT - # catch ``BaseException`` here so process-level signals - # (``KeyboardInterrupt`` / ``SystemExit``) propagate normally. - with contextlib.suppress(Exception): - await websocket.send_json({"type": "cancelled", "id": prompt_id}) - raise - except Exception as exc: # pylint: disable=broad-exception-caught - # Surface generation errors to the client as a structured frame so - # the connection survives a single-prompt failure (parity with the - # ``error`` reply the reader emits on bad input). - with contextlib.suppress(Exception): - await websocket.send_json( - {"type": "error", "id": prompt_id, "message": str(exc)}, - ) - logger.exception("_stream_tokens failed for prompt %s", prompt_id) + await websocket.send_json({"type": "done", "id": prompt_id}) + except asyncio.CancelledError: + # Best-effort: tell the client we honoured their cancel. Suppress + # ordinary send errors (the socket may already be closed) and re-raise + # so the caller observes the cancellation. We deliberately do NOT + # catch ``BaseException`` here so process-level signals + # (``KeyboardInterrupt`` / ``SystemExit``) propagate normally. + with contextlib.suppress(Exception): + await websocket.send_json({"type": "cancelled", "id": prompt_id}) + raise + except Exception as exc: # pylint: disable=broad-exception-caught + # Surface generation errors to the client as a structured frame so + # the connection survives a single-prompt failure (parity with the + # ``error`` reply the reader emits on bad input). + with contextlib.suppress(Exception): + await websocket.send_json( + {"type": "error", "id": prompt_id, "message": str(exc)}, + ) + logger.exception("_stream_tokens failed for prompt %s", prompt_id) async def _reader( diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py index c0401953b999..30b9fa1fbb66 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py @@ -50,16 +50,9 @@ async def handle_ws(websocket: WebSocket) -> None: mapped to RFC 6455 close code 1011. WebSocket protocol-level Ping/Pong keep-alive is disabled by default; enable it by setting the ``WS_KEEPALIVE_INTERVAL`` environment variable. - - Each inbound frame opens a per-turn ``invoke_agent`` span via - ``app.ws_invocation`` so traces line up 1:1 with the HTTP - ``POST /invocations`` endpoint above. """ async for message in websocket.iter_text(): - async with app.ws_invocation(websocket) as turn: - await websocket.send_text( - f"[{turn.invocation_id}] {message}" - ) + await websocket.send_text(message) if __name__ == "__main__": diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invocation_helper.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invocation_helper.py deleted file mode 100644 index d45e70af72aa..000000000000 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invocation_helper.py +++ /dev/null @@ -1,197 +0,0 @@ -# --------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# --------------------------------------------------------- -"""Tests for the per-turn ``ws_invocation`` helper. - -Each ``ws_invocation()`` block opens an ``invoke_agent`` span (mirroring the -HTTP ``POST /invocations`` model) parented to the connection-level -``websocket_session`` span, mints an ``invocation_id`` for the turn, and -ensures the OTel context is current for every ``await`` inside the block. -""" -import os -from unittest.mock import patch - -from opentelemetry import trace as otel_trace -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from starlette.testclient import TestClient -from starlette.websockets import WebSocket - -from azure.ai.agentserver.invocations import InvocationAgentServerHost -from azure.ai.agentserver.invocations._constants import InvocationsWSConstants - - -# --------------------------------------------------------------------------- -# Test fixture: module-scoped OTel provider + in-memory exporter -# --------------------------------------------------------------------------- - -_EXPORTER = InMemorySpanExporter() -_PROVIDER = TracerProvider() -_PROVIDER.add_span_processor(SimpleSpanProcessor(_EXPORTER)) -otel_trace.set_tracer_provider(_PROVIDER) - - -def _make_server_with_invocations() -> InvocationAgentServerHost: - """Build a tracing-enabled WS host whose handler opens one ``invoke_agent`` - span per inbound text frame via ``ws_invocation``.""" - with patch.dict( - os.environ, - {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, - ): - with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): - server = InvocationAgentServerHost() - - @server.ws_handler - async def handler(websocket: WebSocket) -> None: - async for msg in websocket.iter_text(): - async with server.ws_invocation(websocket) as turn: - await websocket.send_json( - {"invocation_id": turn.invocation_id, "echo": msg} - ) - - return server - - -def _get_spans(): - return _EXPORTER.get_finished_spans() - - -def _session_spans(spans): - return [s for s in spans if "websocket_session" in s.name] - - -def _invoke_spans(spans): - return [s for s in spans if "invoke_agent" in s.name] - - -# --------------------------------------------------------------------------- -# Span shape -# --------------------------------------------------------------------------- - -def test_ws_invocation_opens_invoke_agent_span_per_turn(): - """One ``invoke_agent`` span is recorded per ``ws_invocation`` block.""" - _EXPORTER.clear() - server = _make_server_with_invocations() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_json() - ws.send_text("again") - ws.receive_json() - - spans = _get_spans() - invokes = _invoke_spans(spans) - assert len(invokes) == 2, f"expected 2 invoke_agent spans, got {[s.name for s in spans]}" - - -def test_ws_invocation_id_is_set_on_invoke_span(): - """``invoke_agent`` span carries the per-turn ``invocation_id``.""" - _EXPORTER.clear() - server = _make_server_with_invocations() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - payload = ws.receive_json() - - invokes = _invoke_spans(_get_spans()) - assert invokes - span = invokes[-1] - attr = span.attributes.get(InvocationsWSConstants.ATTR_SPAN_INVOCATION_ID) - assert attr == payload["invocation_id"] - - -def test_ws_invocation_unique_per_turn(): - """Each turn mints a fresh ``invocation_id``.""" - _EXPORTER.clear() - server = _make_server_with_invocations() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("a") - a = ws.receive_json()["invocation_id"] - ws.send_text("b") - b = ws.receive_json()["invocation_id"] - ws.send_text("c") - c = ws.receive_json()["invocation_id"] - - assert len({a, b, c}) == 3 - - -def test_ws_invocation_span_is_child_of_websocket_session(): - """``invoke_agent`` is parented to the connection-level ``websocket_session``.""" - _EXPORTER.clear() - server = _make_server_with_invocations() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_json() - - spans = _get_spans() - sessions = _session_spans(spans) - invokes = _invoke_spans(spans) - assert sessions and invokes - session_span = sessions[-1] - invoke_span = invokes[-1] - # Parent-child via OTel SpanContext.span_id linkage. - assert invoke_span.parent is not None - assert invoke_span.parent.span_id == session_span.context.span_id - - -# --------------------------------------------------------------------------- -# Session ID propagation -# --------------------------------------------------------------------------- - -def test_ws_invocation_inherits_connection_session_id(monkeypatch): - """``ws_invocation`` defaults to ``app.config.session_id`` when set.""" - monkeypatch.setenv("FOUNDRY_AGENT_SESSION_ID", "fixed-session") - _EXPORTER.clear() - server = _make_server_with_invocations() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_json() - - invokes = _invoke_spans(_get_spans()) - assert invokes - assert ( - invokes[-1].attributes.get(InvocationsWSConstants.ATTR_SPAN_SESSION_ID) - == "fixed-session" - ) - - -# --------------------------------------------------------------------------- -# OTel context propagation through awaits -# --------------------------------------------------------------------------- - -def test_ws_invocation_propagates_otel_context_across_awaits(): - """A user-opened child span inside ``ws_invocation`` is parented to - the per-turn ``invoke_agent`` span (proves OTel context survives awaits).""" - _EXPORTER.clear() - with patch.dict( - os.environ, - {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, - ): - with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): - server = InvocationAgentServerHost() - - tracer = otel_trace.get_tracer("test") - - @server.ws_handler - async def handler(websocket: WebSocket) -> None: - async for msg in websocket.iter_text(): - async with server.ws_invocation(websocket): - # User-opened child span — must parent to the invoke_agent span. - with tracer.start_as_current_span("user_work"): - await websocket.send_text(msg) - - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _get_spans() - invoke = [s for s in spans if s.name.startswith("invoke_agent") or "invoke_agent" in s.name] - user_work = [s for s in spans if s.name == "user_work"] - assert invoke and user_work - assert user_work[-1].parent is not None - assert user_work[-1].parent.span_id == invoke[-1].context.span_id From db1e596b42d1f64a68e9da559f4546f4bafbe56c Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Mon, 18 May 2026 09:38:16 +0000 Subject: [PATCH 07/13] Drop README/CHANGELOG note about missing WS invocation_id and HTTP-only GET/cancel --- sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md | 2 +- sdk/agentserver/azure-ai-agentserver-invocations/README.md | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index 5cd275b49f73..0cae361fc35f 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -7,7 +7,7 @@ - Error source classification headers: All HTTP error responses now include `x-platform-error-source` with a value of `user`, `platform`, or `upstream` to indicate which component caused the error. Developer handler exceptions and missing handler registrations are classified as `upstream`. Exceptions tagged with the platform error tag are classified as `platform` and additionally include `x-platform-error-detail` with truncated exception details (max 2048 characters) for diagnostics. - WebSocket protocol support — `InvocationAgentServerHost` now hosts `/invocations_ws` alongside `POST /invocations`. Register the handler with the new `@app.ws_handler` decorator. The route is registered lazily on first decoration, so hosts without a registered handler return HTTP 404. - WebSocket Ping/Pong keep-alive — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` env var (auto-injected by AgentService into hosted-agent containers; surfaced on `app.config.ws_ping_interval` in `azure-ai-agentserver-core>=2.0.0b4`). `0` (or unset) disables keep-alive. Wired through to Hypercorn's `websocket_ping_interval` by `AgentServerHost._build_hypercorn_config`. -- WebSocket telemetry — single connection-scoped `websocket_session` OpenTelemetry span per WS connection, plus a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. There is no per-message `invocation_id` for WS: a connection is a single long-lived session, not a sequence of HTTP invocations. +- WebSocket telemetry — single connection-scoped `websocket_session` OpenTelemetry span per WS connection, plus a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. - New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages). ### Breaking Changes diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index a46ba840f0da..603f8364ab5e 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md @@ -230,8 +230,6 @@ app.run() A WebSocket connection is wrapped by the SDK in a single connection-scoped `websocket_session` OpenTelemetry span. The span carries the GenAI semantic-convention attributes plus `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms`. Any child spans your handler opens — e.g. via `opentelemetry.trace.get_tracer(...).start_as_current_span(...)` — are automatically parented to the connection span. -There is no per-message `invocation_id` for WebSocket: a WS connection is a single long-lived session, not a sequence of independently-identifiable HTTP invocations. The HTTP-only endpoints `GET /invocations/{id}` and `POST /invocations/{id}/cancel` do not apply to WS sessions — clients observe status via reply frames and cancel by closing the socket (the handler observes `asyncio.CancelledError`). - ### Handler signature The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`. The full WebSocket API — `iter_text`, `iter_bytes`, `iter_json`, `send_text`, `send_bytes`, `send_json`, `close`, `headers`, `query_params`, `client`, `state` — is available, so application protocols on top of `invocations_ws` are entirely under your control. From dc1cfc6a0f34f54a444abcaaef0d7ca14b4b649c Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Mon, 18 May 2026 11:01:45 +0000 Subject: [PATCH 08/13] Fix stale module docstring: ws keep-alive is env-only now (no constructor arg) --- .../azure/ai/agentserver/invocations/_invocation_ws.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py index 11e5eac2ebde..0afa4ee51c01 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py @@ -9,9 +9,9 @@ * ``await websocket.accept()`` before the handler runs; * WebSocket protocol-level Ping/Pong keep-alive (disabled by default; - enable via the ``WS_KEEPALIVE_INTERVAL`` env var or the - ``ws_ping_interval=`` constructor argument) so idle connections can - survive Azure APIM / Azure Load Balancer's ~4-minute idle timeout; + enable via the ``WS_KEEPALIVE_INTERVAL`` environment variable surfaced + on ``AgentConfig.ws_ping_interval``) so idle connections can survive + Azure APIM / Azure Load Balancer's ~4-minute idle timeout; * a clean close on handler return (code 1000) or a 1011 close on uncaught handler exceptions; * a structured close-event log line and OTel span attributes carrying From c562e2d3fd53d4fbf9e087147ddad00e544c4901 Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Mon, 18 May 2026 11:03:19 +0000 Subject: [PATCH 09/13] Add smoke tests for ws_bidirectional_streaming_agent sample Loads the sample as a module (importlib.util), patches _TOKEN_DELAY_S to 0 so tests run in ms, and exercises the wire protocol via TestClient: - handshake: ready frame on connect - streaming: prompt -> tokens -> done (matches _SIMULATED_TOKENS exactly) - error replies: missing id, unknown type, invalid JSON - cancel: in-flight prompt + unknown id - graceful shutdown: bye returns -> SDK clean close - HTTP parity: POST /invocations still echoes 9 tests, all pass in ~1.3s. --- .../test_ws_bidirectional_streaming_sample.py | 191 ++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_bidirectional_streaming_sample.py diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_bidirectional_streaming_sample.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_bidirectional_streaming_sample.py new file mode 100644 index 000000000000..e4a1564a935b --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_bidirectional_streaming_sample.py @@ -0,0 +1,191 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Smoke tests for the ws_bidirectional_streaming_agent sample. + +Loads the sample module directly (samples are not a package), patches the +per-token sleep down to 0 so the test runs in milliseconds, and drives the +``ready`` / ``prompt`` / ``cancel`` / ``bye`` wire protocol over Starlette's +``TestClient.websocket_connect``. +""" +import importlib.util +import json +import pathlib +import sys + +import pytest +from starlette.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + + +_SAMPLE_PATH = ( + pathlib.Path(__file__).parent.parent + / "samples" + / "ws_bidirectional_streaming_agent" + / "ws_bidirectional_streaming_agent.py" +) + + +@pytest.fixture +def sample(monkeypatch): + """Load the sample as a module and zero out the per-token delay.""" + spec = importlib.util.spec_from_file_location( + "ws_bidirectional_streaming_agent_sample", _SAMPLE_PATH + ) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + monkeypatch.setattr(module, "_TOKEN_DELAY_S", 0.0) + yield module + sys.modules.pop(spec.name, None) + + +# --------------------------------------------------------------------------- +# Handshake +# --------------------------------------------------------------------------- + +def test_ws_bidirectional_sends_ready_on_connect(sample): + """The handler immediately sends a ``{"type": "ready"}`` frame on connect.""" + client = TestClient(sample.app) + with client.websocket_connect("/invocations_ws") as ws: + ws.send_text(json.dumps({"type": "bye"})) + first = ws.receive_json() + assert first == {"type": "ready"} + + +# --------------------------------------------------------------------------- +# Prompt streaming +# --------------------------------------------------------------------------- + +def _drain_until_done(ws, prompt_id: str): + """Collect token frames until the matching ``done`` (or ``cancelled``) arrives.""" + tokens: list[str] = [] + terminal: dict | None = None + while True: + msg = ws.receive_json() + if msg["type"] == "token" and msg["id"] == prompt_id: + tokens.append(msg["token"]) + elif msg["type"] in ("done", "cancelled", "error") and msg["id"] == prompt_id: + terminal = msg + break + return tokens, terminal + + +def test_ws_bidirectional_streams_tokens_for_prompt(sample): + """A ``prompt`` frame produces a token stream terminated by ``done``.""" + client = TestClient(sample.app) + with client.websocket_connect("/invocations_ws") as ws: + assert ws.receive_json() == {"type": "ready"} + ws.send_text(json.dumps({"type": "prompt", "id": "p1", "text": "story"})) + tokens, terminal = _drain_until_done(ws, "p1") + ws.send_text(json.dumps({"type": "bye"})) + + assert terminal == {"type": "done", "id": "p1"} + # All simulated tokens echoed in order. + assert tokens == sample._SIMULATED_TOKENS + + +def test_ws_bidirectional_prompt_without_id_emits_error(sample): + """A ``prompt`` missing an ``id`` triggers an error reply but keeps the connection.""" + client = TestClient(sample.app) + with client.websocket_connect("/invocations_ws") as ws: + assert ws.receive_json() == {"type": "ready"} + ws.send_text(json.dumps({"type": "prompt", "text": "no-id"})) + reply = ws.receive_json() + assert reply["type"] == "error" + assert "id" in reply["message"] + ws.send_text(json.dumps({"type": "bye"})) + + +def test_ws_bidirectional_unknown_type_emits_error(sample): + """An unknown message ``type`` triggers an error reply but keeps the connection.""" + client = TestClient(sample.app) + with client.websocket_connect("/invocations_ws") as ws: + assert ws.receive_json() == {"type": "ready"} + ws.send_text(json.dumps({"type": "foobar"})) + reply = ws.receive_json() + assert reply["type"] == "error" + assert "foobar" in reply["message"] + ws.send_text(json.dumps({"type": "bye"})) + + +def test_ws_bidirectional_invalid_json_emits_error(sample): + """A non-JSON text frame triggers an error reply.""" + client = TestClient(sample.app) + with client.websocket_connect("/invocations_ws") as ws: + assert ws.receive_json() == {"type": "ready"} + ws.send_text("not json {") + reply = ws.receive_json() + assert reply["type"] == "error" + assert "JSON" in reply["message"] + ws.send_text(json.dumps({"type": "bye"})) + + +# --------------------------------------------------------------------------- +# Cancellation +# --------------------------------------------------------------------------- + +def test_ws_bidirectional_cancel_interrupts_in_flight_prompt(sample): + """A ``cancel`` frame mid-stream surfaces a ``cancelled`` event.""" + client = TestClient(sample.app) + with client.websocket_connect("/invocations_ws") as ws: + assert ws.receive_json() == {"type": "ready"} + ws.send_text(json.dumps({"type": "prompt", "id": "p2", "text": "story"})) + # Cancel before drain — handler should reply with cancelled (and may + # have emitted a few token frames first). + ws.send_text(json.dumps({"type": "cancel", "id": "p2"})) + # Drain until we see either done or cancelled. + terminal: dict | None = None + while True: + msg = ws.receive_json() + if msg["type"] in ("done", "cancelled") and msg["id"] == "p2": + terminal = msg + break + ws.send_text(json.dumps({"type": "bye"})) + + assert terminal is not None + # With _TOKEN_DELAY_S = 0 the stream may finish before the cancel is + # observed; both terminal types are acceptable outcomes. + assert terminal["type"] in ("cancelled", "done") + + +def test_ws_bidirectional_cancel_unknown_id_is_noop(sample): + """Cancelling an unknown prompt ID does not break the connection.""" + client = TestClient(sample.app) + with client.websocket_connect("/invocations_ws") as ws: + assert ws.receive_json() == {"type": "ready"} + # Cancel an id that was never registered — handler ignores silently. + ws.send_text(json.dumps({"type": "cancel", "id": "ghost"})) + # Verify connection still works by sending a real prompt. + ws.send_text(json.dumps({"type": "prompt", "id": "p3", "text": "ok"})) + _, terminal = _drain_until_done(ws, "p3") + ws.send_text(json.dumps({"type": "bye"})) + assert terminal == {"type": "done", "id": "p3"} + + +# --------------------------------------------------------------------------- +# Graceful shutdown +# --------------------------------------------------------------------------- + +def test_ws_bidirectional_bye_closes_connection(sample): + """A ``bye`` frame causes the handler to return → SDK closes cleanly.""" + client = TestClient(sample.app) + with pytest.raises(WebSocketDisconnect): + with client.websocket_connect("/invocations_ws") as ws: + assert ws.receive_json() == {"type": "ready"} + ws.send_text(json.dumps({"type": "bye"})) + # Receive should surface the close as WebSocketDisconnect. + ws.receive_json() + + +# --------------------------------------------------------------------------- +# HTTP parity +# --------------------------------------------------------------------------- + +def test_ws_bidirectional_http_invoke_still_works(sample): + """The same host still serves ``POST /invocations`` for HTTP parity.""" + client = TestClient(sample.app) + response = client.post("/invocations", json={"hello": "world"}) + assert response.status_code == 200 + assert response.json() == {"echo": {"hello": "world"}} From fe51b11b2e166f6ab7878094acd78ae6bc33c0ec Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Mon, 18 May 2026 11:04:48 +0000 Subject: [PATCH 10/13] Drop unused interval kwarg from resolve_ws_ping_interval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reviewer note: the explicit-argument branch was dead code — AgentConfig.from_env() always called it with None, and there is no programmatic constructor path to override the env var. Removing the parameter makes the signature match the actual public API (env-only) and removes the implied (but non-existent) callable shape. --- .../azure/ai/agentserver/core/_config.py | 55 +++++++++---------- 1 file changed, 25 insertions(+), 30 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py index a8433427debd..f3a0e229f15b 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py @@ -129,7 +129,7 @@ def from_env(cls) -> Self: _ENV_APPLICATIONINSIGHTS_CONNECTION_STRING, ""), otlp_endpoint=os.environ.get(_ENV_OTEL_EXPORTER_OTLP_ENDPOINT, ""), sse_keepalive_interval=resolve_sse_keepalive_interval(None), - ws_ping_interval=resolve_ws_ping_interval(None), + ws_ping_interval=resolve_ws_ping_interval(), ) @@ -331,43 +331,38 @@ def resolve_otlp_endpoint() -> Optional[str]: return value if value else None -def resolve_ws_ping_interval(interval: Optional[float] = None) -> float: - """Resolve the WebSocket Ping/Pong keep-alive interval. +def resolve_ws_ping_interval() -> float: + """Resolve the WebSocket Ping/Pong keep-alive interval from the env var. - Resolution order: explicit *interval* → ``WS_KEEPALIVE_INTERVAL`` env - var → ``0.0`` (disabled). ``0`` (from any source) disables - keep-alive. + Reads the ``WS_KEEPALIVE_INTERVAL`` environment variable (auto-injected + by AgentService into hosted-agent containers) and returns the parsed + value in seconds. ``0`` (or unset) disables keep-alive. + + The keep-alive interval is intentionally env-only — there is no + programmatic constructor argument to override it. Hosted agents + inherit the platform-injected value automatically; in tests, set the + env var (e.g. via ``monkeypatch.setenv``). - :param interval: Explicitly requested interval in seconds, or None. - :type interval: Optional[float] :return: The resolved interval in seconds (``0`` means disabled). :rtype: float - :raises ValueError: If the resolved value is negative, non-finite, or - (for the env var) not parseable as a number. + :raises ValueError: If the env var is set but not parseable as a + non-negative finite number. """ import math # local import — only used here - if interval is not None: - try: - resolved = float(interval) - except (TypeError, ValueError) as exc: - raise ValueError( - f"ws_ping_interval must be a number, got {interval!r}" - ) from exc - else: - env_raw = os.environ.get(_ENV_WS_KEEPALIVE_INTERVAL) - if env_raw is None or env_raw == "": - return _DEFAULT_WS_PING_INTERVAL - try: - resolved = float(env_raw) - except ValueError as exc: - raise ValueError( - f"Invalid value for {_ENV_WS_KEEPALIVE_INTERVAL}: " - f"{env_raw!r} (expected a non-negative number)" - ) from exc + env_raw = os.environ.get(_ENV_WS_KEEPALIVE_INTERVAL) + if env_raw is None or env_raw == "": + return _DEFAULT_WS_PING_INTERVAL + try: + resolved = float(env_raw) + except ValueError as exc: + raise ValueError( + f"Invalid value for {_ENV_WS_KEEPALIVE_INTERVAL}: " + f"{env_raw!r} (expected a non-negative number)" + ) from exc if math.isnan(resolved) or math.isinf(resolved) or resolved < 0.0: raise ValueError( - f"ws_ping_interval must be a non-negative finite number, " - f"got {interval if interval is not None else env_raw!r}" + f"Invalid value for {_ENV_WS_KEEPALIVE_INTERVAL}: " + f"{env_raw!r} (expected a non-negative finite number)" ) return resolved From 6f529e518eff22dbb5e6a7763da6554f851aacf4 Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Mon, 18 May 2026 11:06:59 +0000 Subject: [PATCH 11/13] Drop specific idle-timeout duration and APIM/Load Balancer mentions from docs --- sdk/agentserver/azure-ai-agentserver-invocations/README.md | 2 +- .../azure/ai/agentserver/invocations/_invocation_ws.py | 2 +- .../ws_bidirectional_streaming_agent.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index 603f8364ab5e..913e95d7329e 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md @@ -221,7 +221,7 @@ app.run() - Registers `/invocations_ws` on the same Starlette host as `/invocations` and `/readiness`. - Calls `await websocket.accept()` before invoking your handler. -- Runs WebSocket Ping/Pong keep-alive in the background — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers). Set the value to `0` to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic. +- Runs WebSocket Ping/Pong keep-alive in the background — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers). Set the value to `0` to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across upstream proxy / load-balancer idle timeouts without any extra application traffic. - Closes the connection cleanly on handler return (close code `1000`) or maps an uncaught handler exception to close code `1011`. - Emits a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `azure.ai.agentserver.invocations_ws.close_code`, and `azure.ai.agentserver.invocations_ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end. - Inherits `/readiness`, OpenTelemetry export, graceful shutdown, and the `x-platform-server` identity header from `azure-ai-agentserver-core`. diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py index 0afa4ee51c01..86246eaf540e 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py @@ -11,7 +11,7 @@ * WebSocket protocol-level Ping/Pong keep-alive (disabled by default; enable via the ``WS_KEEPALIVE_INTERVAL`` environment variable surfaced on ``AgentConfig.ws_ping_interval``) so idle connections can survive - Azure APIM / Azure Load Balancer's ~4-minute idle timeout; + upstream proxy / load-balancer idle timeouts; * a clean close on handler return (code 1000) or a 1011 close on uncaught handler exceptions; * a structured close-event log line and OTel span attributes carrying diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py index 903b635d9e7b..c755d58d7aae 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py @@ -23,8 +23,8 @@ (opcode 0x9) every ``ws_ping_interval`` seconds (disabled by default; enable by setting ``WS_KEEPALIVE_INTERVAL`` or by passing ``InvocationAgentServerHost(ws_ping_interval=)``). When - enabled, that is enough to survive Azure APIM / Azure Load Balancer's - ~4-minute idle timeout without your handler having to push any + enabled, that is enough to survive upstream proxy / load-balancer + idle timeouts without your handler having to push any application-level heartbeat messages of its own. Wire protocol (JSON over text frames) From ae06e2e5a2f2235d78fb3e390205a51aa0b0b46c Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Mon, 18 May 2026 11:08:19 +0000 Subject: [PATCH 12/13] Drop stale ws_ping_interval= constructor mention from bidirectional sample docstring --- .../ws_bidirectional_streaming_agent.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py index c755d58d7aae..7d5b9ac42588 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py @@ -20,12 +20,12 @@ Connection keep-alive is **not** an application concern: the SDK can ask Hypercorn to send WebSocket protocol-level Ping frames - (opcode 0x9) every ``ws_ping_interval`` seconds (disabled by default; - enable by setting ``WS_KEEPALIVE_INTERVAL`` or by passing - ``InvocationAgentServerHost(ws_ping_interval=)``). When - enabled, that is enough to survive upstream proxy / load-balancer - idle timeouts without your handler having to push any - application-level heartbeat messages of its own. + (opcode 0x9) on its own schedule (disabled by default; enable by + setting the ``WS_KEEPALIVE_INTERVAL`` environment variable, surfaced + on ``AgentConfig.ws_ping_interval``). When enabled, that is enough + to survive upstream proxy / load-balancer idle timeouts without your + handler having to push any application-level heartbeat messages of + its own. Wire protocol (JSON over text frames) ------------------------------------- @@ -255,9 +255,9 @@ async def handle_ws(websocket: WebSocket) -> None: The SDK has already accepted the connection by the time this function runs, and it is sending WS protocol-level Ping frames on its own - schedule (see ``ws_ping_interval``). When this coroutine returns the - SDK closes the socket with code ``1000``; if it raises, the SDK maps - the exception to ``1011``. + schedule (see ``AgentConfig.ws_ping_interval``). When this coroutine + returns the SDK closes the socket with code ``1000``; if it raises, + the SDK maps the exception to ``1011``. """ await websocket.send_json({"type": "ready"}) From 0822abef014ff42eaec3ab8790cf603352c5fa5d Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Mon, 18 May 2026 15:15:00 +0000 Subject: [PATCH 13/13] Drop empty CHANGELOG sections to satisfy release-build verifier --- sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index 0cae361fc35f..ce89988e351f 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -10,10 +10,6 @@ - WebSocket telemetry — single connection-scoped `websocket_session` OpenTelemetry span per WS connection, plus a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. - New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages). -### Breaking Changes - -### Bugs Fixed - ### Other Changes - Platform header name constants (e.g. `x-platform-error-source`, `x-platform-error-detail`) are now imported from `azure-ai-agentserver-core` (`_platform_headers` module) instead of being defined locally. Error source classification helpers remain internal to this package.