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..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 @@ -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(), ) @@ -322,3 +329,40 @@ 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() -> float: + """Resolve the WebSocket Ping/Pong keep-alive interval from the env var. + + 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``). + + :return: The resolved interval in seconds (``0`` means disabled). + :rtype: float + :raises ValueError: If the env var is set but not parseable as a + non-negative finite number. + """ + import math # local import — only used here + + 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"Invalid value for {_ENV_WS_KEEPALIVE_INTERVAL}: " + f"{env_raw!r} (expected a non-negative finite number)" + ) + return resolved diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index 5c629cbd2388..ce89988e351f 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -5,6 +5,10 @@ ### 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 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. +- New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages). ### Other Changes diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index 5e9dfe515657..913e95d7329e 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,69 @@ 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 — 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`. + +### Per-connection tracing + +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. + +### 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 + +| Environment variable | Default | Description | +|---|---|---| +| `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 + +| 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 ### Reporting issues @@ -196,6 +264,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: 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/_constants.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py index 49ed1853ae68..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 @@ -21,3 +21,25 @@ 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" + + # 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_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 4bea0b550b36..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 @@ -39,6 +39,7 @@ ) from ._constants import InvocationConstants +from ._invocation_ws import _WSHandlerMixin logger = logging.getLogger("azure.ai.agentserver") @@ -147,13 +148,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): @@ -162,13 +168,19 @@ 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 @@ -189,8 +201,13 @@ def __init__( self._cancel_invocation_fn: Optional[Callable] = None self._openapi_spec = openapi_spec + # 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 = [ + invocation_routes: list[Any] = [ Route( "/invocations/docs/openapi.json", self._get_openapi_spec_endpoint, 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..86246eaf540e --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py @@ -0,0 +1,476 @@ +# --------------------------------------------------------- +# 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 (disabled by default; + enable via the ``WS_KEEPALIVE_INTERVAL`` environment variable surfaced + on ``AgentConfig.ws_ping_interval``) so idle connections can survive + 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 + ``azure.ai.agentserver.invocations_ws.session_id``, + ``azure.ai.agentserver.invocations_ws.close_code``, and + ``azure.ai.agentserver.invocations_ws.duration_ms``. +""" + +import inspect +import logging +import time +import uuid +from collections.abc import Awaitable, Callable +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, +) + +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") + + +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(_MixinBase): + """Pure mixin that adds the ``@app.ws_handler`` decorator and ``/invocations_ws`` route. + + 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__. + _ws_fn: Optional[WSHandler] + + def _init_ws_state(self) -> None: + """Initialize WS handler slots. + + 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 + + # ------------------------------------------------------------------ + # Public configuration accessor + # ------------------------------------------------------------------ + + @property + 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 float(self.config.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, or its + signature cannot be invoked with a single positional argument + (the WebSocket). + """ + 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 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 + # 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 + + def _ensure_ws_route_registered(self) -> None: + """Append the ``/invocations_ws`` WebSocketRoute to the router (idempotent). + + 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 + + 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", + ) + ) + + # ------------------------------------------------------------------ + # 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. 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 + # 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}) + + # 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: + 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: 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, + ) -> 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]] + :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: + raise RuntimeError("_invoke_user_handler called with no registered ws_handler") + try: + await 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 _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, + ) + + # 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 + # ------------------------------------------------------------------ + + @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 ``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. + + :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) + + # 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={ + 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/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..7d5b9ac42588 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py @@ -0,0 +1,276 @@ +"""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 can + ask Hypercorn to send WebSocket protocol-level Ping frames + (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) +------------------------------------- + +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:: + + 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"} +""" +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("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", + " 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 — leading underscore + signals "intentionally ignored" to linters). + :type _text: str + :return: An async generator of token strings. + :rtype: AsyncGenerator[str, None] + """ + 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 + # 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( + 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 + # ``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), + ) + + 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() + # 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 + + 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 ``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"}) + + 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..30b9fa1fbb66 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py @@ -0,0 +1,59 @@ +"""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) + pip install websockets + 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 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. + """ + 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..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") # --------------------------------------------------------------------------- @@ -199,6 +200,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, "azure.ai.agentserver.invocations_ws.session_id") and hasattr(r, "azure.ai.agentserver.invocations_ws.close_code") + ] + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- 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"}} 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..5fd135ab875d --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py @@ -0,0 +1,162 @@ +# --------------------------------------------------------- +# 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 (``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. +""" +import logging + +import pytest +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 ( + _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, "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 + assert isinstance(duration_ms, int) + assert duration_ms >= 0 + + +def test_ws_close_event_duration_is_non_negative(caplog): + """``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) + + 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], "azure.ai.agentserver.invocations_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], "azure.ai.agentserver.invocations_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, "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 new file mode 100644 index 000000000000..835e21c1f5e1 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py @@ -0,0 +1,115 @@ +# --------------------------------------------------------- +# 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 + + +# --------------------------------------------------------------------------- +# 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(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 + + 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="single 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="single 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_rejects_upgrade(): + """If no @ws_handler is registered the route is absent and the upgrade is rejected.""" + app = InvocationAgentServerHost() + client = TestClient(app) + + # 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 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..90fc3183e455 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py @@ -0,0 +1,132 @@ +# --------------------------------------------------------- +# 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, "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. + 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, 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. + 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] + # 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], "azure.ai.agentserver.invocations_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], "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 new file mode 100644 index 000000000000..7fa992566d76 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py @@ -0,0 +1,147 @@ +# --------------------------------------------------------- +# 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) + # 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") + 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..0dbeb9e2a14c --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py @@ -0,0 +1,79 @@ +# --------------------------------------------------------- +# 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. +""" +import pytest +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 + + +@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() + 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..2a1f5171dcf5 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py @@ -0,0 +1,125 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tests for the ``ws_ping_interval`` configuration knob. + +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 + + +_ENV_NAME = "WS_KEEPALIVE_INTERVAL" + + +# --------------------------------------------------------------------------- +# Default / accepted values (env-driven) +# --------------------------------------------------------------------------- + +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(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(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(monkeypatch): + """Fractional intervals are coerced to ``float``.""" + 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 + + +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 + + +# --------------------------------------------------------------------------- +# Rejected values (validation surfaces at AgentConfig.from_env) +# --------------------------------------------------------------------------- + +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() + + +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 (delegated to core's _build_hypercorn_config) +# --------------------------------------------------------------------------- + +def test_ws_ping_interval_propagates_to_hypercorn_config(monkeypatch): + """The configured interval lands on the Hypercorn server config.""" + 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(monkeypatch): + """Zero explicitly sets Hypercorn's ``websocket_ping_interval`` to ``None``.""" + 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] + + +# --------------------------------------------------------------------------- +# Property surface +# --------------------------------------------------------------------------- + +def test_ws_ping_interval_property_is_read_only(monkeypatch): + """``ws_ping_interval`` is exposed only as a property (no setter).""" + monkeypatch.setenv(_ENV_NAME, "20") + app = InvocationAgentServerHost() + with pytest.raises(AttributeError): + app.ws_ping_interval = 10 # type: ignore[misc] + + +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 == app.config.ws_ping_interval == 7.5 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..8629dc6078ad --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py @@ -0,0 +1,104 @@ +# --------------------------------------------------------- +# 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_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() + 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): + 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): + 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..2163e8cd9286 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py @@ -0,0 +1,101 @@ +# --------------------------------------------------------- +# 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 ``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)] + + +# --------------------------------------------------------------------------- +# 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) + # Uniqueness is the meaningful invariant — implies len == 5 too. + 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]) + + +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 new file mode 100644 index 000000000000..a83dcfe18690 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py @@ -0,0 +1,463 @@ +# --------------------------------------------------------- +# 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("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) + 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("azure.ai.agentserver.invocations_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("azure.ai.agentserver.invocations_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("azure.ai.agentserver.invocations_ws.close_code") == 1011 + + +def test_ws_span_records_duration_ms(): + """``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: + ws.send_text("hi") + ws.receive_text() + + spans = _ws_session_spans(_get_spans()) + attrs = dict(spans[0].attributes) + duration_ms = attrs.get("azure.ai.agentserver.invocations_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("azure.ai.agentserver.invocations_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 + # 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_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): + with client.websocket_connect("/invocations_ws"): + pass + + spans = _ws_session_spans(_get_spans()) + assert spans == [] + + +# --------------------------------------------------------------------------- +# 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(monkeypatch): + """Agent name without version still appears in span name.""" + 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: + 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