Implement WebSocket support for InvocationAgentServerHost (#46841)

Author: yulin-li

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)
 

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]

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

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
 

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
 

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"

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,