[agentserver-invocations] Address PR review comments 18-52 for invocations_ws

sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md

@@ -4,19 +4,10 @@
 
 ### Features Added
 
-- `invocations_ws` (WebSocket) protocol support on `InvocationAgentServerHost`.
-  Register a handler with the new `@app.ws_handler` decorator to host a
-  full-duplex WebSocket endpoint at `/invocations_ws` on the same host that
-  serves `POST /invocations`. The SDK calls `await websocket.accept()` before
-  invoking the handler, runs WebSocket Ping/Pong keep-alive in the background
-  (default 30 s; configurable via the new `ws_ping_interval` constructor
-  argument), closes the connection cleanly on handler return, and maps
-  uncaught exceptions to RFC 6455 close code `1011`. Each connection emits a
-  structured close-event log line carrying `ws.session_id`, `ws.close_code`,
-  and `ws.duration_ms`, and the same fields are recorded as OpenTelemetry
-  span attributes. `/readiness`, OTEL export, graceful shutdown, and the
-  `x-platform-server` identity header continue to be inherited from
-  `azure-ai-agentserver-core`.
+- WebSocket protocol support — `InvocationAgentServerHost` now hosts `/invocations_ws` alongside `POST /invocations`. Register the handler with the new `@app.ws_handler` decorator. The route is registered lazily on first decoration, so hosts without a registered handler return HTTP 404.
+- WebSocket Ping/Pong keep-alive — disabled by default; enable by passing `InvocationAgentServerHost(ws_ping_interval=<seconds>)` or by setting the `WS_KEEPALIVE_INTERVAL` env var (auto-injected by AgentService into hosted-agent containers). The constructor argument takes precedence; `0` (or unset) disables keep-alive. Wired through to Hypercorn's `websocket_ping_interval`.
+- WebSocket telemetry — structured close-event log line and OpenTelemetry span attributes `azure.ai.agentserver.invocations_ws.{session_id,close_code,duration_ms}`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation.
+- New samples: `samples/ws_invoke_agent/` (echo) and `samples/ws_bidirectional_streaming_agent/` (concurrent token streaming with cancel/bye control messages).
 
 ### Breaking Changes
 

sdk/agentserver/azure-ai-agentserver-invocations/README.md

@@ -221,9 +221,9 @@ app.run()
 
 - Registers `/invocations_ws` on the same Starlette host as `/invocations` and `/readiness`.
 - Calls `await websocket.accept()` before invoking your handler.
-- Runs WebSocket Ping/Pong keep-alive in the background — default 30 s, configurable via `InvocationAgentServerHost(ws_ping_interval=...)`. Set `ws_ping_interval=0` to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic.
+- Runs WebSocket Ping/Pong keep-alive in the background — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers) or by passing `InvocationAgentServerHost(ws_ping_interval=...)`. Set the value to `0` (in any source) to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic.
 - Closes the connection cleanly on handler return (close code `1000`) or maps an uncaught handler exception to close code `1011`.
-- Emits a structured close-event log line carrying `ws.session_id`, `ws.close_code`, and `ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end.
+- Emits a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `azure.ai.agentserver.invocations_ws.close_code`, and `azure.ai.agentserver.invocations_ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end.
 - Inherits `/readiness`, OpenTelemetry export, graceful shutdown, and the `x-platform-server` identity header from `azure-ai-agentserver-core`.
 
 ### Handler signature
@@ -236,7 +236,19 @@ The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`.
 
 | Constructor argument | Default | Description |
 |---|---|---|
-| `ws_ping_interval` | `30.0` (seconds) | WebSocket protocol Ping interval. `0` disables keep-alive. Negative or non-finite values are rejected. |
+| `ws_ping_interval` | `0` (disabled) | WebSocket protocol Ping interval in seconds. `0` disables keep-alive. Negative or non-finite values are rejected. When `None`, the SDK reads the `WS_KEEPALIVE_INTERVAL` env var before falling back to disabled. |
+
+| Environment variable | Default | Description |
+|---|---|---|
+| `WS_KEEPALIVE_INTERVAL` | unset (disabled) | Platform-injected override for the WebSocket Ping interval (seconds). `0` disables keep-alive. Ignored when `ws_ping_interval=` is set explicitly. |
+
+### Reference: close codes
+
+| Close code | Meaning |
+|---|---|
+| `1000` | Handler returned cleanly (normal close). |
+| `1011` | Handler raised an unhandled exception (mapped by the SDK). |
+| `4000`-`4999` | Application-defined codes (set by the handler via `await websocket.close(code=...)` — surfaced unchanged to the client). |
 
 ## Troubleshooting
 

sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py

@@ -31,19 +31,22 @@ class InvocationsWSConstants:
     # Route
     ROUTE_PATH = "/invocations_ws"
 
-    # Default WebSocket Ping interval in seconds.
-    # Azure APIM and Azure Load Balancer drop idle WebSocket connections
-    # after ~4 minutes; 30 s gives a comfortable safety margin.
-    DEFAULT_PING_INTERVAL_S = 30.0
+    # Environment variable for platform-injected keep-alive override.
+    # AgentService injects this into every hosted-agent container so the
+    # platform can tune the WebSocket Ping cadence without redeploying
+    # the agent code.
+    # Resolution order: explicit ``ws_ping_interval=`` constructor arg →
+    # ``WS_KEEPALIVE_INTERVAL`` env var → disabled (``0``).
+    # ``0`` disables keep-alive.
+    ENV_WS_KEEPALIVE_INTERVAL = "WS_KEEPALIVE_INTERVAL"
 
     # Close codes (RFC 6455)
     CLOSE_NORMAL = 1000  # handler returned cleanly
     CLOSE_INTERNAL_ERROR = 1011  # handler raised an unhandled exception
-    CLOSE_SERVICE_RESTART = 1012  # graceful shutdown drained the connection
 
     # Span attribute keys
-    ATTR_SPAN_SESSION_ID = "ws.session_id"
-    ATTR_SPAN_CLOSE_CODE = "ws.close_code"
-    ATTR_SPAN_DURATION_MS = "ws.duration_ms"
+    ATTR_SPAN_SESSION_ID = "azure.ai.agentserver.invocations_ws.session_id"
+    ATTR_SPAN_CLOSE_CODE = "azure.ai.agentserver.invocations_ws.close_code"
+    ATTR_SPAN_DURATION_MS = "azure.ai.agentserver.invocations_ws.duration_ms"
     ATTR_SPAN_ERROR_CODE = "azure.ai.agentserver.invocations_ws.error.code"
     ATTR_SPAN_ERROR_MESSAGE = "azure.ai.agentserver.invocations_ws.error.message"

sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py

@@ -133,9 +133,12 @@ async def ws(websocket: WebSocket) -> None:
         is served at ``GET /invocations/docs/openapi.json``.
     :type openapi_spec: Optional[dict[str, Any]]
     :param ws_ping_interval: Seconds between WebSocket protocol Ping frames
-        on ``/invocations_ws``.  ``None`` (default) selects 30 s; ``0``
-        disables keep-alive.  Configured on the underlying Hypercorn
-        server so the framing is opcode 0x9 / 0xA, not application JSON.
+        on ``/invocations_ws``. ``None`` (default) reads the
+        ``WS_KEEPALIVE_INTERVAL`` environment variable (auto-injected by
+        AgentService into hosted-agent containers); when neither is set,
+        keep-alive is disabled. ``0`` disables keep-alive. Configured on
+        the underlying Hypercorn server so the framing is opcode 0x9 /
+        0xA, not application JSON.
     :type ws_ping_interval: Optional[float]
     """
 
@@ -182,7 +185,6 @@ def __init__(
                 methods=["POST"],
                 name="cancel_invocation",
             ),
-            self._build_ws_route(self._ws_endpoint),
         ]
 
         # Merge with any routes from sibling mixins via cooperative init
@@ -211,8 +213,11 @@ def _build_hypercorn_config(self, host: str, port: int) -> object:
         Hypercorn sends WS protocol Ping frames every
         ``websocket_ping_interval`` seconds on every active WebSocket
         connection — exactly the keep-alive the ``invocations_ws`` spec
-        requires.  ``ws_ping_interval=0`` leaves the default
-        ``None`` (disabled).
+        requires.  ``ws_ping_interval=0`` explicitly disables keep-alive
+        by setting the attribute to ``None``.
+
+        ``azure-ai-agentserver-core`` pins ``hypercorn>=0.17.0``, which
+        is guaranteed to expose ``websocket_ping_interval``.
 
         :param host: Network interface to bind.
         :type host: str
@@ -222,20 +227,10 @@ def _build_hypercorn_config(self, host: str, port: int) -> object:
         :rtype: hypercorn.config.Config
         """
         config = super()._build_hypercorn_config(host, port)
-        if self._ws_ping_interval and self._ws_ping_interval > 0:
-            try:
-                # ``websocket_ping_interval`` is a float-or-None on
-                # Hypercorn ≥0.14; assigning a positive float enables
-                # protocol-level Ping frames.
-                config.websocket_ping_interval = self._ws_ping_interval  # type: ignore[attr-defined]
-            except Exception:  # pylint: disable=broad-exception-caught
-                # Hypercorn <0.14 does not support per-server WS ping —
-                # leave the default and warn so operators can upgrade.
-                logger.warning(
-                    "Hypercorn does not support websocket_ping_interval; "
-                    "WebSocket keep-alive will be best-effort.",
-                    exc_info=True,
-                )
+        # Positive float enables protocol-level Ping; None disables.
+        config.websocket_ping_interval = (  # type: ignore[attr-defined]
+            self._ws_ping_interval if self._ws_ping_interval > 0 else None
+        )
         return config
 
     # ------------------------------------------------------------------