diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md
index 7745922fba25..2410d32a0844 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md
@@ -4,19 +4,10 @@
 
 ### Features Added
 
-- `invocations_ws` (WebSocket) protocol support on `InvocationAgentServerHost`.
-  Register a handler with the new `@app.ws_handler` decorator to host a
-  full-duplex WebSocket endpoint at `/invocations_ws` on the same host that
-  serves `POST /invocations`. The SDK calls `await websocket.accept()` before
-  invoking the handler, runs WebSocket Ping/Pong keep-alive in the background
-  (default 30 s; configurable via the new `ws_ping_interval` constructor
-  argument), closes the connection cleanly on handler return, and maps
-  uncaught exceptions to RFC 6455 close code `1011`. Each connection emits a
-  structured close-event log line carrying `ws.session_id`, `ws.close_code`,
-  and `ws.duration_ms`, and the same fields are recorded as OpenTelemetry
-  span attributes. `/readiness`, OTEL export, graceful shutdown, and the
-  `x-platform-server` identity header continue to be inherited from
-  `azure-ai-agentserver-core`.
+- WebSocket protocol support — `InvocationAgentServerHost` now hosts `/invocations_ws` alongside `POST /invocations`. Register the handler with the new `@app.ws_handler` decorator. The route is registered lazily on first decoration, so hosts without a registered handler return HTTP 404.
+- WebSocket Ping/Pong keep-alive — disabled by default; enable by passing `InvocationAgentServerHost(ws_ping_interval=<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
 
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md
index 63ec794bbc8f..acab609ed746 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md
@@ -221,9 +221,9 @@ app.run()
 
 - Registers `/invocations_ws` on the same Starlette host as `/invocations` and `/readiness`.
 - Calls `await websocket.accept()` before invoking your handler.
-- Runs WebSocket Ping/Pong keep-alive in the background — default 30 s, configurable via `InvocationAgentServerHost(ws_ping_interval=...)`. Set `ws_ping_interval=0` to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic.
+- Runs WebSocket Ping/Pong keep-alive in the background — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` environment variable (auto-injected by AgentService into hosted-agent containers) or by passing `InvocationAgentServerHost(ws_ping_interval=...)`. Set the value to `0` (in any source) to disable. Frames are sent at the WebSocket protocol layer (RFC 6455 opcode `0x9`/`0xA`) by the underlying Hypercorn server, which keeps the connection alive across Azure APIM and Azure Load Balancer's ~4 minute idle timeout without any extra application traffic.
 - Closes the connection cleanly on handler return (close code `1000`) or maps an uncaught handler exception to close code `1011`.
-- Emits a structured close-event log line carrying `ws.session_id`, `ws.close_code`, and `ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end.
+- Emits a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `azure.ai.agentserver.invocations_ws.close_code`, and `azure.ai.agentserver.invocations_ws.duration_ms`. The same fields are recorded as OpenTelemetry span attributes so the connection lifetime is visible end-to-end.
 - Inherits `/readiness`, OpenTelemetry export, graceful shutdown, and the `x-platform-server` identity header from `azure-ai-agentserver-core`.
 
 ### Handler signature
@@ -236,7 +236,19 @@ The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`.
 
 | Constructor argument | Default | Description |
 |---|---|---|
-| `ws_ping_interval` | `30.0` (seconds) | WebSocket protocol Ping interval. `0` disables keep-alive. Negative or non-finite values are rejected. |
+| `ws_ping_interval` | `0` (disabled) | WebSocket protocol Ping interval in seconds. `0` disables keep-alive. Negative or non-finite values are rejected. When `None`, the SDK reads the `WS_KEEPALIVE_INTERVAL` env var before falling back to disabled. |
+
+| Environment variable | Default | Description |
+|---|---|---|
+| `WS_KEEPALIVE_INTERVAL` | unset (disabled) | Platform-injected override for the WebSocket Ping interval (seconds). `0` disables keep-alive. Ignored when `ws_ping_interval=` is set explicitly. |
+
+### Reference: close codes
+
+| Close code | Meaning |
+|---|---|
+| `1000` | Handler returned cleanly (normal close). |
+| `1011` | Handler raised an unhandled exception (mapped by the SDK). |
+| `4000`-`4999` | Application-defined codes (set by the handler via `await websocket.close(code=...)` — surfaced unchanged to the client). |
 
 ## Troubleshooting
 
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py
index 611772000efc..43683652417c 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py
@@ -31,19 +31,22 @@ class InvocationsWSConstants:
     # Route
     ROUTE_PATH = "/invocations_ws"
 
-    # Default WebSocket Ping interval in seconds.
-    # Azure APIM and Azure Load Balancer drop idle WebSocket connections
-    # after ~4 minutes; 30 s gives a comfortable safety margin.
-    DEFAULT_PING_INTERVAL_S = 30.0
+    # Environment variable for platform-injected keep-alive override.
+    # AgentService injects this into every hosted-agent container so the
+    # platform can tune the WebSocket Ping cadence without redeploying
+    # the agent code.
+    # Resolution order: explicit ``ws_ping_interval=`` constructor arg →
+    # ``WS_KEEPALIVE_INTERVAL`` env var → disabled (``0``).
+    # ``0`` disables keep-alive.
+    ENV_WS_KEEPALIVE_INTERVAL = "WS_KEEPALIVE_INTERVAL"
 
     # Close codes (RFC 6455)
     CLOSE_NORMAL = 1000  # handler returned cleanly
     CLOSE_INTERNAL_ERROR = 1011  # handler raised an unhandled exception
-    CLOSE_SERVICE_RESTART = 1012  # graceful shutdown drained the connection
 
     # Span attribute keys
-    ATTR_SPAN_SESSION_ID = "ws.session_id"
-    ATTR_SPAN_CLOSE_CODE = "ws.close_code"
-    ATTR_SPAN_DURATION_MS = "ws.duration_ms"
+    ATTR_SPAN_SESSION_ID = "azure.ai.agentserver.invocations_ws.session_id"
+    ATTR_SPAN_CLOSE_CODE = "azure.ai.agentserver.invocations_ws.close_code"
+    ATTR_SPAN_DURATION_MS = "azure.ai.agentserver.invocations_ws.duration_ms"
     ATTR_SPAN_ERROR_CODE = "azure.ai.agentserver.invocations_ws.error.code"
     ATTR_SPAN_ERROR_MESSAGE = "azure.ai.agentserver.invocations_ws.error.message"
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py
index 8c706e55e5b4..40657f494a15 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py
@@ -133,9 +133,12 @@ async def ws(websocket: WebSocket) -> None:
         is served at ``GET /invocations/docs/openapi.json``.
     :type openapi_spec: Optional[dict[str, Any]]
     :param ws_ping_interval: Seconds between WebSocket protocol Ping frames
-        on ``/invocations_ws``.  ``None`` (default) selects 30 s; ``0``
-        disables keep-alive.  Configured on the underlying Hypercorn
-        server so the framing is opcode 0x9 / 0xA, not application JSON.
+        on ``/invocations_ws``. ``None`` (default) reads the
+        ``WS_KEEPALIVE_INTERVAL`` environment variable (auto-injected by
+        AgentService into hosted-agent containers); when neither is set,
+        keep-alive is disabled. ``0`` disables keep-alive. Configured on
+        the underlying Hypercorn server so the framing is opcode 0x9 /
+        0xA, not application JSON.
     :type ws_ping_interval: Optional[float]
     """
 
@@ -182,7 +185,6 @@ def __init__(
                 methods=["POST"],
                 name="cancel_invocation",
             ),
-            self._build_ws_route(self._ws_endpoint),
         ]
 
         # Merge with any routes from sibling mixins via cooperative init
@@ -211,8 +213,11 @@ def _build_hypercorn_config(self, host: str, port: int) -> object:
         Hypercorn sends WS protocol Ping frames every
         ``websocket_ping_interval`` seconds on every active WebSocket
         connection — exactly the keep-alive the ``invocations_ws`` spec
-        requires.  ``ws_ping_interval=0`` leaves the default
-        ``None`` (disabled).
+        requires.  ``ws_ping_interval=0`` explicitly disables keep-alive
+        by setting the attribute to ``None``.
+
+        ``azure-ai-agentserver-core`` pins ``hypercorn>=0.17.0``, which
+        is guaranteed to expose ``websocket_ping_interval``.
 
         :param host: Network interface to bind.
         :type host: str
@@ -222,20 +227,10 @@ def _build_hypercorn_config(self, host: str, port: int) -> object:
         :rtype: hypercorn.config.Config
         """
         config = super()._build_hypercorn_config(host, port)
-        if self._ws_ping_interval and self._ws_ping_interval > 0:
-            try:
-                # ``websocket_ping_interval`` is a float-or-None on
-                # Hypercorn ≥0.14; assigning a positive float enables
-                # protocol-level Ping frames.
-                config.websocket_ping_interval = self._ws_ping_interval  # type: ignore[attr-defined]
-            except Exception:  # pylint: disable=broad-exception-caught
-                # Hypercorn <0.14 does not support per-server WS ping —
-                # leave the default and warn so operators can upgrade.
-                logger.warning(
-                    "Hypercorn does not support websocket_ping_interval; "
-                    "WebSocket keep-alive will be best-effort.",
-                    exc_info=True,
-                )
+        # Positive float enables protocol-level Ping; None disables.
+        config.websocket_ping_interval = (  # type: ignore[attr-defined]
+            self._ws_ping_interval if self._ws_ping_interval > 0 else None
+        )
         return config
 
     # ------------------------------------------------------------------
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py
index 38b531ebe385..86fedaca6efb 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py
@@ -8,34 +8,47 @@
 the user handler with:
 
 * ``await websocket.accept()`` before the handler runs;
-* WebSocket protocol-level Ping/Pong keep-alive (default 30 s) so idle
-  connections survive Azure APIM / Azure Load Balancer's ~4-minute idle
-  timeout;
+* WebSocket protocol-level Ping/Pong keep-alive (disabled by default;
+  enable via the ``WS_KEEPALIVE_INTERVAL`` env var or the
+  ``ws_ping_interval=`` constructor argument) so idle connections can
+  survive Azure APIM / Azure Load Balancer's ~4-minute idle timeout;
 * a clean close on handler return (code 1000) or a 1011 close on uncaught
   handler exceptions;
 * a structured close-event log line and OTel span attributes carrying
-  ``ws.session_id``, ``ws.close_code``, and ``ws.duration_ms``.
+  ``azure.ai.agentserver.invocations_ws.session_id``,
+  ``azure.ai.agentserver.invocations_ws.close_code``, and
+  ``azure.ai.agentserver.invocations_ws.duration_ms``.
 """
-from __future__ import annotations
 
 import inspect
 import logging
 import math
+import os
 import time
 import uuid
 from collections.abc import Awaitable, Callable
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
 
 from azure.ai.agentserver.core import (  # pylint: disable=no-name-in-module
     AgentServerHost,
     end_span,
-    record_error,
 )
 
 from ._constants import InvocationsWSConstants
 
+# Type-checking only base so the mixin reads as an ``AgentServerHost`` to
+# mypy / pyright (resolves ``self.config``, ``self.request_span``,
+# ``self.router``) without coupling the runtime hierarchy.  At runtime the
+# mixin is a plain ``object`` subclass — only the concrete
+# ``InvocationAgentServerHost`` MRO actually inherits ``AgentServerHost``,
+# which keeps the diamond out of the runtime class graph.
+if TYPE_CHECKING:
+    _MixinBase = AgentServerHost
+else:
+    _MixinBase = object
+
 logger = logging.getLogger("azure.ai.agentserver")
 
 
@@ -59,18 +72,17 @@ def _safe_set_attrs(span: Any, attrs: dict[str, Any]) -> None:
         logger.debug("Failed to set WS span attributes: %s", list(attrs.keys()), exc_info=True)
 
 
-class _WSHandlerMixin(AgentServerHost):
-    """Trait class that adds the ``@app.ws_handler`` decorator and ``/invocations_ws`` route.
-
-    Inherits from :class:`~azure.ai.agentserver.core.AgentServerHost` so that
-    cooperative ``super().__init__`` calls and host attribute access
-    (``self.config``, ``self.request_span``) are resolved statically as well
-    as at runtime.  Designed to be mixed into :class:`InvocationAgentServerHost`
-    so the same host object exposes both ``POST /invocations`` (HTTP) and
-    ``/invocations_ws`` (WebSocket) on the same Starlette application.
+class _WSHandlerMixin(_MixinBase):
+    """Pure mixin that adds the ``@app.ws_handler`` decorator and ``/invocations_ws`` route.
 
-    Subclasses must append the route returned by :meth:`_build_ws_route` to
-    their ``routes`` list before calling ``super().__init__``.
+    Designed to be mixed into a concrete
+    :class:`~azure.ai.agentserver.core.AgentServerHost` subclass (e.g.
+    :class:`InvocationAgentServerHost`) so the same host object exposes
+    both ``POST /invocations`` (HTTP) and ``/invocations_ws`` (WebSocket)
+    on the same Starlette application.  At runtime the mixin is a plain
+    ``object`` subclass — host attributes (``self.config``,
+    ``self.request_span``, ``self.router``) are accessed via duck typing
+    and are typed only for the static checkers (see ``_MixinBase``).
     """
 
     # Slots populated by __init__.
@@ -80,14 +92,41 @@ class _WSHandlerMixin(AgentServerHost):
     def _init_ws_state(self, ws_ping_interval: Optional[float]) -> None:
         """Initialize WS handler slots.
 
+        Resolution order for the keep-alive interval:
+
+        1. Explicit *ws_ping_interval* constructor argument (when not ``None``).
+        2. ``WS_KEEPALIVE_INTERVAL`` environment variable (auto-injected by
+           AgentService into every hosted-agent container).
+        3. Disabled (``0``) — no Ping/Pong frames are sent.
+
+        ``0`` (from any source) disables keep-alive entirely.
+
         :param ws_ping_interval: Seconds between WS protocol Ping frames.
-            ``None`` selects the default (30 s); ``0`` disables keep-alive.
+            ``None`` selects the env var or the default.
         :type ws_ping_interval: Optional[float]
-        :raises ValueError: If *ws_ping_interval* is negative or non-finite.
+        :raises ValueError: If the resolved value is negative, non-finite,
+            or (for the env var) not parseable as a number.
         """
         self._ws_fn = None
         if ws_ping_interval is None:
-            resolved = InvocationsWSConstants.DEFAULT_PING_INTERVAL_S
+            env_raw = os.environ.get(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL)
+            if env_raw is None or env_raw == "":
+                resolved = 0.0
+            else:
+                try:
+                    resolved = float(env_raw)
+                except ValueError as exc:
+                    raise ValueError(
+                        f"Invalid value for "
+                        f"{InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL}: "
+                        f"{env_raw!r} (expected a non-negative number)"
+                    ) from exc
+                if math.isnan(resolved) or math.isinf(resolved) or resolved < 0.0:
+                    raise ValueError(
+                        f"Invalid value for "
+                        f"{InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL}: "
+                        f"{env_raw!r} (expected a non-negative finite number)"
+                    )
         else:
             try:
                 resolved = float(ws_ping_interval)
@@ -142,39 +181,58 @@ async def handle(websocket: WebSocket) -> None:
         :type fn: Callable[[WebSocket], Awaitable[None]]
         :return: The original function (unmodified).
         :rtype: Callable[[WebSocket], Awaitable[None]]
-        :raises TypeError: If *fn* is not an ``async def`` function.
+        :raises TypeError: If *fn* is not an ``async def`` function, or its
+            signature is not callable with exactly one positional argument.
         """
         if not inspect.iscoroutinefunction(fn):
             raise TypeError(
                 f"ws_handler expects an async function, got {type(fn).__name__}. "
                 "Use 'async def' to define your handler."
             )
+        # Validate signature at registration time (not at first request) so
+        # 0-arg / 2-required-arg coroutine mistakes surface at import.
+        try:
+            sig = inspect.signature(fn)
+            sig.bind(None)  # one positional placeholder for the WebSocket
+        except TypeError as exc:
+            raise TypeError(
+                f"ws_handler must accept exactly one positional argument "
+                f"(the WebSocket); got {fn.__qualname__}{inspect.signature(fn)}"
+            ) from exc
+        if self._ws_fn is not None:
+            # Match the HTTP decorator's last-write-wins semantics, but log
+            # so misconfigured apps that double-register a handler aren't
+            # silently downgraded.
+            logger.warning(
+                "ws_handler overwriting previously registered handler %s with %s",
+                getattr(self._ws_fn, "__qualname__", repr(self._ws_fn)),
+                getattr(fn, "__qualname__", repr(fn)),
+            )
         self._ws_fn = fn
+        # Register the route lazily on first decoration so hosts without a
+        # registered handler return HTTP 404 to a WebSocket upgrade rather than
+        # accepting and immediately closing with code 1011.
+        self._ensure_ws_route_registered()
         return fn
 
-    # ------------------------------------------------------------------
-    # Route factory (for cooperative __init__)
-    # ------------------------------------------------------------------
+    def _ensure_ws_route_registered(self) -> None:
+        """Append the ``/invocations_ws`` WebSocketRoute to the router (idempotent).
 
-    @staticmethod
-    def _build_ws_route(endpoint: Callable[[WebSocket], Awaitable[None]]) -> Any:
-        """Return a :class:`~starlette.routing.WebSocketRoute` for ``/invocations_ws``.
-
-        Imported lazily to avoid hard-coding the route type in the public
-        module body and keep the import surface symmetric with the HTTP
-        ``Route`` import in :mod:`._invocation`.
-
-        :param endpoint: The async endpoint to wire to the route.
-        :type endpoint: Callable[[WebSocket], Awaitable[None]]
-        :return: A configured ``WebSocketRoute`` instance.
-        :rtype: ~starlette.routing.WebSocketRoute
+        Starlette's ``router.routes`` is a plain list and may be mutated
+        between construction and first request, so deferring registration
+        until ``@ws_handler`` is called is safe.
         """
         from starlette.routing import WebSocketRoute  # pylint: disable=import-outside-toplevel
 
-        return WebSocketRoute(
-            InvocationsWSConstants.ROUTE_PATH,
-            endpoint,
-            name="invocations_ws",
+        for route in self.router.routes:
+            if isinstance(route, WebSocketRoute) and getattr(route, "path", None) == InvocationsWSConstants.ROUTE_PATH:
+                return
+        self.router.routes.append(
+            WebSocketRoute(
+                InvocationsWSConstants.ROUTE_PATH,
+                self._ws_endpoint,
+                name="invocations_ws",
+            )
         )
 
     # ------------------------------------------------------------------
@@ -191,10 +249,14 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None:
         :param websocket: The incoming Starlette WebSocket.
         :type websocket: ~starlette.websockets.WebSocket
         """
-        # Per-connection identifiers.  Session ID is generated server-side;
-        # the spec carries it in the close-event metric so an operator can
-        # correlate logs/metrics for a given long-lived connection.
-        session_id = str(uuid.uuid4())
+        # Per-connection identifiers.  Honour the platform-injected
+        # ``FOUNDRY_AGENT_SESSION_ID`` (surfaced via ``self.config.session_id``)
+        # so HTTP and WebSocket transports on the same container report the
+        # same session ID; fall back to a fresh UUID when the platform does
+        # not inject one.  Matches the precedence used by the HTTP
+        # ``POST /invocations`` endpoint (minus the query-param override,
+        # which has no equivalent ergonomic on a long-lived WS connection).
+        session_id = self.config.session_id or str(uuid.uuid4())
         start_ns = time.monotonic_ns()
 
         # Open the OTel span before accepting so any tracecontext header
@@ -212,9 +274,9 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None:
         otel_span = span_ctx.__enter__()
         _safe_set_attrs(otel_span, {InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id})
 
-        if self._ws_fn is None:
-            await self._reject_no_handler(websocket, span_ctx, otel_span, session_id, start_ns)
-            return
+        # NOTE: when no ``@ws_handler`` is registered, the route itself is
+        # not registered (see ``_ensure_ws_route_registered``), so this
+        # endpoint is unreachable in that state — Starlette returns 404.
 
         # Accept the upgrade *before* invoking the user handler — per spec.
         try:
@@ -262,10 +324,14 @@ async def _invoke_user_handler(
             is set only for an *unhandled* exception (so the caller can map
             it to span error events and a 1011 close).
         :rtype: tuple[int, Optional[BaseException]]
+        :raises RuntimeError: If no handler is registered (programmer error;
+            the public ``_ws_endpoint`` checks ``_ws_fn`` before calling).
         """
-        assert self._ws_fn is not None  # checked by caller
+        ws_fn = self._ws_fn
+        if ws_fn is None:
+            raise RuntimeError("_invoke_user_handler called with no registered ws_handler")
         try:
-            await self._ws_fn(websocket)
+            await ws_fn(websocket)
             return InvocationsWSConstants.CLOSE_NORMAL, None
         except WebSocketDisconnect as exc:
             # Client (or proxy) closed first — surface their code, not 1011.
@@ -280,49 +346,6 @@ async def _invoke_user_handler(
             )
             return InvocationsWSConstants.CLOSE_INTERNAL_ERROR, exc
 
-    async def _reject_no_handler(
-        self,
-        websocket: WebSocket,
-        span_ctx: Any,
-        otel_span: Any,
-        session_id: str,
-        start_ns: int,
-    ) -> None:
-        """Refuse a WS upgrade when no ``@ws_handler`` is registered.
-
-        :param websocket: The pending WebSocket awaiting upgrade.
-        :type websocket: ~starlette.websockets.WebSocket
-        :param span_ctx: The active ``request_span`` context manager.
-        :type span_ctx: any
-        :param otel_span: The current OTel span (or ``None``).
-        :type otel_span: any
-        :param session_id: Per-connection session ID.
-        :type session_id: str
-        :param start_ns: ``time.monotonic_ns()`` at connection start.
-        :type start_ns: int
-        """
-        logger.error(
-            "WebSocket connection on %s rejected: no @ws_handler registered",
-            InvocationsWSConstants.ROUTE_PATH,
-        )
-        duration_ms = (time.monotonic_ns() - start_ns) // 1_000_000
-        self._emit_close_event(
-            otel_span,
-            session_id,
-            InvocationsWSConstants.CLOSE_INTERNAL_ERROR,
-            duration_ms,
-            error_code="not_implemented",
-            error_message="No ws_handler registered.",
-        )
-        try:
-            span_ctx.__exit__(None, None, None)
-        finally:
-            end_span(otel_span)
-        await websocket.close(
-            code=InvocationsWSConstants.CLOSE_INTERNAL_ERROR,
-            reason="No ws_handler registered",
-        )
-
     async def _finalize_session(
         self,
         *,
@@ -390,21 +413,16 @@ async def _finalize_session(
             error_message=str(handler_exc) if handler_exc is not None else None,
         )
 
-        if handler_exc is not None:
-            try:
-                record_error(otel_span, handler_exc)
-            finally:
-                try:
-                    span_ctx.__exit__(
-                        type(handler_exc), handler_exc, handler_exc.__traceback__,
-                    )
-                finally:
-                    end_span(otel_span)
-        else:
-            try:
-                span_ctx.__exit__(None, None, None)
-            finally:
-                end_span(otel_span)
+        # Always exit the context manager with ``(None, None, None)``:
+        # OpenTelemetry's ``start_as_current_span`` records an exception
+        # event whenever ``__exit__`` is given exception info, so passing
+        # the user exception there *and* calling ``end_span(span, exc=...)``
+        # would double-record.  We take ownership of error recording and
+        # ending via ``end_span`` instead.
+        try:
+            span_ctx.__exit__(None, None, None)
+        finally:
+            end_span(otel_span, exc=handler_exc)
 
     # ------------------------------------------------------------------
     # Close event
@@ -422,8 +440,9 @@ def _emit_close_event(
     ) -> None:
         """Record close-event span attributes and emit a structured log line.
 
-        The log record carries the ``ws.session_id``, ``ws.close_code``,
-        and ``ws.duration_ms`` fields listed in the spec via the standard
+        The log record carries the ``azure.ai.agentserver.invocations_ws.session_id``,
+        ``azure.ai.agentserver.invocations_ws.close_code``, and
+        ``azure.ai.agentserver.invocations_ws.duration_ms`` fields listed in the spec via the standard
         ``logging`` ``extra`` dict — a structured-logging formatter or an
         OTel logging bridge can pick them up directly without having to
         parse the message.
@@ -454,14 +473,22 @@ def _emit_close_event(
             attrs[InvocationsWSConstants.ATTR_SPAN_ERROR_MESSAGE] = error_message
         _safe_set_attrs(otel_span, attrs)
 
+        # NOTE: ``extra`` keys deliberately use dotted names (``azure.ai.agentserver.invocations_ws.session_id``
+        # etc.) so they line up 1:1 with the OTel span-attribute keys defined
+        # in :class:`InvocationsWSConstants`.  This makes log<->trace
+        # correlation trivial in OTel logging bridges.  The trade-off is that
+        # ``%(azure.ai.agentserver.invocations_ws.session_id)s`` printf-style formatters can't address them
+        # directly — use a structured formatter (JSON, OTel) or access via
+        # ``LogRecord.__dict__["azure.ai.agentserver.invocations_ws.session_id"]`` if you need them in plain
+        # ``logging`` output.
         logger.info(
             "invocations_ws connection closed: session_id=%s code=%s duration_ms=%s",
             session_id,
             close_code,
             duration_ms,
             extra={
-                "ws.session_id": session_id,
-                "ws.close_code": close_code,
-                "ws.duration_ms": duration_ms,
+                InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id,
+                InvocationsWSConstants.ATTR_SPAN_CLOSE_CODE: close_code,
+                InvocationsWSConstants.ATTR_SPAN_DURATION_MS: duration_ms,
             },
         )
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py
index 9a66525d3c03..a228f61e145d 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_bidirectional_streaming_agent/ws_bidirectional_streaming_agent.py
@@ -18,12 +18,14 @@
 
 .. note::
 
-   Connection keep-alive is **not** an application concern: the SDK
-   configures Hypercorn to send WebSocket protocol-level Ping frames
-   (opcode 0x9) every ``ws_ping_interval`` seconds (default 30 s).  That
-   is enough to survive Azure APIM / Azure Load Balancer's ~4-minute idle
-   timeout without your handler having to push any application-level
-   heartbeat messages of its own.
+   Connection keep-alive is **not** an application concern: the SDK can
+   ask Hypercorn to send WebSocket protocol-level Ping frames
+   (opcode 0x9) every ``ws_ping_interval`` seconds (disabled by default;
+   enable by setting ``WS_KEEPALIVE_INTERVAL`` or by passing
+   ``InvocationAgentServerHost(ws_ping_interval=<seconds>)``).  When
+   enabled, that is enough to survive Azure APIM / Azure Load Balancer's
+   ~4-minute idle timeout without your handler having to push any
+   application-level heartbeat messages of its own.
 
 Wire protocol (JSON over text frames)
 -------------------------------------
@@ -49,14 +51,13 @@
 Drive it with the ``websockets`` CLI; the server keeps streaming tokens
 for each prompt while you type the next one::
 
+    pip install websockets
     python -m websockets ws://localhost:8088/invocations_ws
     > {"type": "prompt", "id": "p1", "text": "Tell me a story"}
     > {"type": "prompt", "id": "p2", "text": "And another"}
     > {"type": "cancel", "id": "p1"}
     > {"type": "bye"}
 """
-from __future__ import annotations
-
 import asyncio
 import contextlib
 import json
@@ -70,12 +71,16 @@
 from azure.ai.agentserver.invocations import InvocationAgentServerHost
 
 
-logger = logging.getLogger("ws_bidirectional_streaming_agent")
+logger = logging.getLogger("azure.ai.agentserver")
 
 app = InvocationAgentServerHost()
 
 
 # Simulated tokens — in production these would come from a model.
+# For real-world streaming patterns, see the streaming examples in
+# ``azure-ai-inference`` and ``azure-ai-projects``:
+#   https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/ai/azure-ai-inference/samples
+#   https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/ai/azure-ai-projects/samples
 _SIMULATED_TOKENS = [
     "Once", " upon", " a", " time", ",", " in", " a", " land",
     " of", " full", "-", "duplex", " sockets", ",", " a", " server",
@@ -100,18 +105,18 @@ async def handle_invoke(request: Request) -> Response:
 # WebSocket — true bidirectional streaming.
 # ---------------------------------------------------------------------------
 
-async def _generate_tokens(text: str) -> AsyncGenerator[str, None]:
+async def _generate_tokens(_text: str) -> AsyncGenerator[str, None]:
     """Yield simulated tokens with a small per-token delay.
 
     Replace this with a real streaming model call (e.g. Azure OpenAI) in
     production.
 
-    :param text: The user prompt (unused in this demo).
-    :type text: str
+    :param _text: The user prompt (unused in this demo — leading underscore
+        signals "intentionally ignored" to linters).
+    :type _text: str
     :return: An async generator of token strings.
     :rtype: AsyncGenerator[str, None]
     """
-    del text  # demo: ignore prompt content
     for token in _SIMULATED_TOKENS:
         await asyncio.sleep(_TOKEN_DELAY_S)
         yield token
@@ -138,10 +143,23 @@ async def _stream_tokens(
     except asyncio.CancelledError:
         # Best-effort: tell the client we honoured their cancel.  Suppress
         # any send error (the socket may already be closed) and re-raise
-        # so the caller observes the cancellation.
-        with contextlib.suppress(Exception):
+        # so the caller observes the cancellation.  ``BaseException``
+        # rather than ``Exception`` so we don't accidentally swallow a
+        # nested cancellation while emitting the courtesy frame.
+        try:
             await websocket.send_json({"type": "cancelled", "id": prompt_id})
+        except BaseException:  # pylint: disable=broad-exception-caught
+            pass
         raise
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+        # Surface generation errors to the client as a structured frame so
+        # the connection survives a single-prompt failure (parity with the
+        # ``error`` reply the reader emits on bad input).
+        with contextlib.suppress(Exception):
+            await websocket.send_json(
+                {"type": "error", "id": prompt_id, "message": str(exc)},
+            )
+        logger.exception("_stream_tokens failed for prompt %s", prompt_id)
 
 
 async def _reader(
@@ -188,6 +206,9 @@ async def _reader(
                     name=f"stream-{prompt_id}",
                 )
                 in_flight[prompt_id] = task
+                # ``k=prompt_id`` captures the current value via a default
+                # argument so the callback removes the right entry even if
+                # ``prompt_id`` is rebound by a later iteration of the loop.
                 task.add_done_callback(
                     lambda _t, k=prompt_id: in_flight.pop(k, None),
                 )
@@ -197,6 +218,13 @@ async def _reader(
                 task = in_flight.get(prompt_id)
                 if task is not None and not task.done():
                     task.cancel()
+                    # Give the task a brief window to finish its courtesy
+                    # ``cancelled`` frame before we move on — prevents the
+                    # next prompt from racing against an in-flight close.
+                    with contextlib.suppress(
+                        asyncio.TimeoutError, asyncio.CancelledError, Exception,
+                    ):
+                        await asyncio.wait_for(task, timeout=1.0)
 
             elif msg_type == "bye":
                 return
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py
index 2f36c48a3b20..269d9a1d4c21 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/ws_invoke_agent/ws_invoke_agent.py
@@ -19,6 +19,7 @@
     # -> {"echo": {"name": "Alice"}}
 
     # WebSocket turn (with the `websockets` client library)
+    pip install websockets
     python -m websockets ws://localhost:8088/invocations_ws
     # > hello
     # < hello
@@ -45,10 +46,11 @@ async def handle_ws(websocket: WebSocket) -> None:
     """Echo every text frame back over the WebSocket connection.
 
     The SDK has already accepted the connection by the time this function
-    runs, sends WebSocket Ping frames every 30 s in the background to keep
-    Azure APIM / Azure Load Balancer from idling the socket out, and will
-    close the connection cleanly on return.  An uncaught exception here
-    is mapped to RFC 6455 close code 1011.
+    runs and will close it cleanly on return.  An uncaught exception is
+    mapped to RFC 6455 close code 1011.  WebSocket protocol-level Ping/Pong
+    keep-alive is disabled by default; enable it by setting the
+    ``WS_KEEPALIVE_INTERVAL`` environment variable or by passing
+    ``InvocationAgentServerHost(ws_ping_interval=<seconds>)``.
     """
     async for message in websocket.iter_text():
         await websocket.send_text(message)
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py
index 24c3c1eb9500..7f5947655ef3 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/conftest.py
@@ -16,6 +16,7 @@
 
 def pytest_configure(config):
     config.addinivalue_line("markers", "tracing_e2e: end-to-end tracing tests against live Application Insights")
+    config.addinivalue_line("markers", "slow: tests that send large payloads or otherwise take noticeable time in CI")
 
 
 # ---------------------------------------------------------------------------
@@ -236,7 +237,7 @@ def _records_with_ws_extras(records):
     """Filter log records that carry the close-event ``ws.*`` extras."""
     return [
         r for r in records
-        if hasattr(r, "ws.session_id") and hasattr(r, "ws.close_code")
+        if hasattr(r, "azure.ai.agentserver.invocations_ws.session_id") and hasattr(r, "azure.ai.agentserver.invocations_ws.close_code")
     ]
 
 
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py
index 6c2c9723c301..5fd135ab875d 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_close_event.py
@@ -4,7 +4,7 @@
 """Tests for the structured close-event log line emitted by ``/invocations_ws``.
 
 Parity with :mod:`tests.test_request_id` — verifies the spec's required
-fields (``ws.session_id``, ``ws.close_code``, ``ws.duration_ms``) appear
+fields (``azure.ai.agentserver.invocations_ws.session_id``, ``azure.ai.agentserver.invocations_ws.close_code``, ``azure.ai.agentserver.invocations_ws.duration_ms``) appear
 on every connection close, and that handler exception details are NOT
 leaked into the structured payload.
 """
@@ -14,6 +14,7 @@
 from starlette.testclient import TestClient
 from starlette.websockets import WebSocketDisconnect
 
+from azure.ai.agentserver.invocations import InvocationAgentServerHost
 from azure.ai.agentserver.invocations._constants import InvocationsWSConstants
 
 from conftest import (
@@ -41,9 +42,9 @@ def test_ws_close_event_log_contains_required_fields(caplog):
     assert matches, "expected a structured close-event log record"
     rec = matches[-1]
 
-    session_id = getattr(rec, "ws.session_id")
-    close_code = getattr(rec, "ws.close_code")
-    duration_ms = getattr(rec, "ws.duration_ms")
+    session_id = getattr(rec, "azure.ai.agentserver.invocations_ws.session_id")
+    close_code = getattr(rec, "azure.ai.agentserver.invocations_ws.close_code")
+    duration_ms = getattr(rec, "azure.ai.agentserver.invocations_ws.duration_ms")
 
     assert isinstance(session_id, str) and session_id  # generated UUID
     assert close_code == InvocationsWSConstants.CLOSE_NORMAL
@@ -52,7 +53,7 @@ def test_ws_close_event_log_contains_required_fields(caplog):
 
 
 def test_ws_close_event_duration_is_non_negative(caplog):
-    """``ws.duration_ms`` is a non-negative integer derived from a monotonic clock."""
+    """``azure.ai.agentserver.invocations_ws.duration_ms`` is a non-negative integer derived from a monotonic clock."""
     app = _make_echo_ws_app()
     client = TestClient(app)
 
@@ -63,7 +64,7 @@ def test_ws_close_event_duration_is_non_negative(caplog):
 
     matches = _records_with_ws_extras(caplog.records)
     assert matches
-    duration_ms = getattr(matches[-1], "ws.duration_ms")
+    duration_ms = getattr(matches[-1], "azure.ai.agentserver.invocations_ws.duration_ms")
     assert isinstance(duration_ms, int)
     assert duration_ms >= 0
 
@@ -85,7 +86,7 @@ def test_ws_close_event_on_handler_exception_records_1011(caplog):
 
     matches = _records_with_ws_extras(caplog.records)
     assert matches
-    assert getattr(matches[-1], "ws.close_code") == 1011
+    assert getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code") == 1011
 
 
 # ---------------------------------------------------------------------------
@@ -108,6 +109,54 @@ def test_ws_close_event_log_does_not_leak_exception_message(caplog):
     assert matches
     rec = matches[-1]
     # The structured close-event log line carries only the safe ws.* fields.
-    assert not hasattr(rec, "ws.error_message")
+    assert not hasattr(rec, "azure.ai.agentserver.invocations_ws.error.message")
     # And the message itself does not embed the raw exception text.
     assert "boom" not in rec.getMessage()
+
+
+def test_ws_disconnect_with_code_zero_falls_back_to_normal_close(caplog):
+    """A ``WebSocketDisconnect(code=0)`` is reported as the normal close code (1000).
+
+    Some clients/proxies surface a falsy ``code`` on the disconnect; the SDK
+    treats it as a clean close rather than as 1011.
+    """
+    app = InvocationAgentServerHost()
+
+    @app.ws_handler
+    async def handler(websocket):  # noqa: ARG001
+        raise WebSocketDisconnect(code=0)
+
+    client = TestClient(app)
+    with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"):
+        with pytest.raises(WebSocketDisconnect):
+            with client.websocket_connect("/invocations_ws") as ws:
+                ws.receive_text()
+
+    matches = _records_with_ws_extras(caplog.records)
+    assert matches
+    assert (
+        getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code")
+        == InvocationsWSConstants.CLOSE_NORMAL
+    )
+
+
+def test_ws_disconnect_with_code_none_falls_back_to_normal_close(caplog):
+    """A ``WebSocketDisconnect(code=None)`` is reported as 1000."""
+    app = InvocationAgentServerHost()
+
+    @app.ws_handler
+    async def handler(websocket):  # noqa: ARG001
+        raise WebSocketDisconnect(code=None)  # type: ignore[arg-type]
+
+    client = TestClient(app)
+    with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"):
+        with pytest.raises(WebSocketDisconnect):
+            with client.websocket_connect("/invocations_ws") as ws:
+                ws.receive_text()
+
+    matches = _records_with_ws_extras(caplog.records)
+    assert matches
+    assert (
+        getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code")
+        == InvocationsWSConstants.CLOSE_NORMAL
+    )
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py
index b447ff8d1090..54d25b415601 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_decorator_pattern.py
@@ -11,7 +11,6 @@
 from starlette.websockets import WebSocket, WebSocketDisconnect
 
 from azure.ai.agentserver.invocations import InvocationAgentServerHost
-from azure.ai.agentserver.invocations._constants import InvocationsWSConstants
 
 
 # ---------------------------------------------------------------------------
@@ -61,32 +60,56 @@ def test_ws_handler_default_is_none():
     assert app._ws_fn is None  # noqa: SLF001
 
 
-def test_ws_handler_last_registration_wins():
-    """Re-applying ``@ws_handler`` replaces the previous function."""
+def test_ws_handler_last_registration_wins(caplog):
+    """Re-applying ``@ws_handler`` replaces the previous function and warns."""
+    import logging
+
     app = InvocationAgentServerHost()
 
     @app.ws_handler
     async def first(websocket: WebSocket) -> None:  # noqa: ARG001
         return
 
-    @app.ws_handler
-    async def second(websocket: WebSocket) -> None:  # noqa: ARG001
-        return
+    with caplog.at_level(logging.WARNING, logger="azure.ai.agentserver"):
+        @app.ws_handler
+        async def second(websocket: WebSocket) -> None:  # noqa: ARG001
+            return
 
     assert app._ws_fn is second  # noqa: SLF001
+    assert any("overwriting previously registered handler" in r.message for r in caplog.records)
+
+
+def test_ws_handler_rejects_zero_arg_coroutine():
+    """A 0-arg coroutine cannot accept the WebSocket — fail at registration."""
+    app = InvocationAgentServerHost()
+
+    with pytest.raises(TypeError, match="exactly one positional argument"):
+        @app.ws_handler
+        async def bad() -> None:  # type: ignore[misc]
+            return
+
+
+def test_ws_handler_rejects_two_required_arg_coroutine():
+    """A 2-required-arg coroutine cannot be bound with one positional."""
+    app = InvocationAgentServerHost()
+
+    with pytest.raises(TypeError, match="exactly one positional argument"):
+        @app.ws_handler
+        async def bad(websocket: WebSocket, extra: int) -> None:  # noqa: ARG001
+            return
 
 
 # ---------------------------------------------------------------------------
 # Missing handler behaviour (parity with test_missing_invoke_handler_returns_501)
 # ---------------------------------------------------------------------------
 
-def test_ws_with_no_handler_registered_closes_with_1011():
-    """If no @ws_handler is registered, the SDK closes with 1011."""
+def test_ws_with_no_handler_registered_rejects_upgrade():
+    """If no @ws_handler is registered the route is absent and the upgrade is rejected."""
     app = InvocationAgentServerHost()
     client = TestClient(app)
 
-    with pytest.raises(WebSocketDisconnect) as excinfo:
+    # Starlette has no WebSocketRoute matching /invocations_ws — the upgrade
+    # is rejected before any handler-level close code can be sent.
+    with pytest.raises(WebSocketDisconnect):
         with client.websocket_connect("/invocations_ws"):
             pass
-
-    assert excinfo.value.code == InvocationsWSConstants.CLOSE_INTERNAL_ERROR
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py
index a25f1249ee8b..90fc3183e455 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_edge_cases.py
@@ -64,7 +64,7 @@ async def handler(websocket: WebSocket) -> None:
     matches = _records_with_ws_extras(caplog.records)
     assert matches
     rec = matches[-1]
-    close_code = getattr(rec, "ws.close_code")
+    close_code = getattr(rec, "azure.ai.agentserver.invocations_ws.close_code")
     # Server surfaces the client's code — NOT 1011.
     assert close_code == 4001
     # No ERROR-level records — a client disconnect is normal.
@@ -75,10 +75,19 @@ async def handler(websocket: WebSocket) -> None:
 # Handler-managed close
 # ---------------------------------------------------------------------------
 
-def test_ws_handler_explicit_close_does_not_double_close(caplog):
+def test_ws_handler_explicit_close_does_not_double_close(caplog, monkeypatch):
     """If the handler closes the WS itself, the SDK does NOT attempt a second close."""
     app = InvocationAgentServerHost()
     closes_observed: list[int] = []
+    close_calls: list[tuple[int, str]] = []
+
+    real_close = WebSocket.close
+
+    async def _counting_close(self, code=1000, reason=""):
+        close_calls.append((code, reason))
+        return await real_close(self, code=code, reason=reason)
+
+    monkeypatch.setattr(WebSocket, "close", _counting_close, raising=True)
 
     async def handler(websocket: WebSocket) -> None:
         # Hand-roll close so we can verify the SDK skips re-closing.
@@ -97,10 +106,12 @@ async def handler(websocket: WebSocket) -> None:
 
     # Handler ran to completion.
     assert closes_observed == [1]
+    # Exactly ONE close — the handler's. The SDK must not re-close.
+    assert len(close_calls) == 1, f"expected one close, got {close_calls}"
     # Close-event log line still emitted (with the handler's code).
     matches = _records_with_ws_extras(caplog.records)
     assert matches
-    assert getattr(matches[-1], "ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL
+    assert getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL
 
 
 # ---------------------------------------------------------------------------
@@ -118,4 +129,4 @@ def test_ws_empty_connection_closes_normally(caplog):
 
     matches = _records_with_ws_extras(caplog.records)
     assert matches
-    assert getattr(matches[-1], "ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL
+    assert getattr(matches[-1], "azure.ai.agentserver.invocations_ws.close_code") == InvocationsWSConstants.CLOSE_NORMAL
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py
index 9eb988b73713..7fa992566d76 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_invoke.py
@@ -101,6 +101,10 @@ async def handler(websocket: WebSocket) -> None:
         await websocket.receive_text()
 
     client = TestClient(app)
+    # Starlette TestClient surfaces *any* close (including a clean 1000) via
+    # ``WebSocketDisconnect`` when ``receive_text()`` is called after the
+    # server sends a close frame — this is the documented client API for
+    # observing the close, not an indication of error.
     with pytest.raises(WebSocketDisconnect) as excinfo:
         with client.websocket_connect("/invocations_ws") as ws:
             ws.send_text("done")
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py
index e05b228c13a8..334465244870 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_multimodal_protocol.py
@@ -9,6 +9,8 @@
 from starlette.testclient import TestClient
 from starlette.websockets import WebSocket
 
+import pytest
+
 from azure.ai.agentserver.invocations import InvocationAgentServerHost
 
 from conftest import _make_echo_ws_app
@@ -48,6 +50,7 @@ def test_ws_unicode_text_round_trip():
         assert ws.receive_text() == text
 
 
+@pytest.mark.slow
 def test_ws_large_text_frame_round_trip():
     """A ~1 MB text frame round-trips successfully (parity with test_large_payload)."""
     app = _make_echo_ws_app()
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py
index 82567c6df622..e8ead6148b9e 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_ping_interval.py
@@ -17,11 +17,10 @@
 # Default / accepted values
 # ---------------------------------------------------------------------------
 
-def test_ws_ping_interval_default_is_30_seconds():
-    """Default ping interval matches the spec (30 s)."""
+def test_ws_ping_interval_default_is_disabled():
+    """Default ping interval is 0 (disabled) when no arg / env var is set."""
     app = InvocationAgentServerHost()
-    assert app.ws_ping_interval == InvocationsWSConstants.DEFAULT_PING_INTERVAL_S
-    assert app.ws_ping_interval == 30.0
+    assert app.ws_ping_interval == 0.0
 
 
 def test_ws_ping_interval_custom_value():
@@ -84,12 +83,11 @@ def test_ws_ping_interval_propagates_to_hypercorn_config():
     assert getattr(config, "websocket_ping_interval", None) == 20.0
 
 
-def test_ws_ping_interval_zero_does_not_override_hypercorn_default():
-    """Zero leaves Hypercorn's default (None = disabled) intact."""
+def test_ws_ping_interval_zero_sets_attribute_to_none():
+    """Zero explicitly sets Hypercorn's ``websocket_ping_interval`` to ``None``."""
     app = InvocationAgentServerHost(ws_ping_interval=0)
     config = app._build_hypercorn_config("0.0.0.0", 8088)  # noqa: SLF001
-    # Hypercorn default is None — our wiring leaves it unset for 0.
-    assert getattr(config, "websocket_ping_interval", None) is None
+    assert config.websocket_ping_interval is None  # type: ignore[attr-defined]
 
 
 # ---------------------------------------------------------------------------
@@ -101,3 +99,74 @@ def test_ws_ping_interval_property_is_read_only():
     app = InvocationAgentServerHost(ws_ping_interval=20)
     with pytest.raises(AttributeError):
         app.ws_ping_interval = 10  # type: ignore[misc]
+
+
+# ---------------------------------------------------------------------------
+# WS_KEEPALIVE_INTERVAL env var (platform-injected override)
+# ---------------------------------------------------------------------------
+
+def test_ws_keepalive_interval_env_var_used_when_arg_is_none(monkeypatch):
+    """Env var drives the interval when no explicit arg is passed."""
+    monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "45")
+    app = InvocationAgentServerHost()
+    assert app.ws_ping_interval == 45.0
+
+
+def test_ws_keepalive_interval_env_var_zero_disables(monkeypatch):
+    """``WS_KEEPALIVE_INTERVAL=0`` disables keep-alive (mirrors SSE)."""
+    monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "0")
+    app = InvocationAgentServerHost()
+    assert app.ws_ping_interval == 0.0
+    config = app._build_hypercorn_config("0.0.0.0", 8088)  # noqa: SLF001
+    assert getattr(config, "websocket_ping_interval", None) is None
+
+
+def test_ws_keepalive_interval_env_var_float_accepted(monkeypatch):
+    """Fractional env-var values are coerced to ``float``."""
+    monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "12.5")
+    app = InvocationAgentServerHost()
+    assert app.ws_ping_interval == 12.5
+
+
+def test_ws_keepalive_interval_env_var_empty_uses_default(monkeypatch):
+    """An empty env-var value falls back to the default (disabled)."""
+    monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "")
+    app = InvocationAgentServerHost()
+    assert app.ws_ping_interval == 0.0
+
+
+def test_ws_keepalive_interval_env_var_invalid_rejected(monkeypatch):
+    """Non-numeric env-var values surface as ``ValueError`` at startup."""
+    monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "thirty")
+    with pytest.raises(ValueError, match="WS_KEEPALIVE_INTERVAL"):
+        InvocationAgentServerHost()
+
+
+def test_ws_keepalive_interval_env_var_negative_rejected(monkeypatch):
+    """Negative env-var values are programming errors."""
+    monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "-5")
+    with pytest.raises(ValueError, match="WS_KEEPALIVE_INTERVAL"):
+        InvocationAgentServerHost()
+
+
+def test_ws_ping_interval_arg_overrides_env_var(monkeypatch):
+    """Explicit constructor arg wins over the env var."""
+    monkeypatch.setenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, "45")
+    app = InvocationAgentServerHost(ws_ping_interval=10)
+    assert app.ws_ping_interval == 10.0
+
+
+def test_ws_ping_interval_default_wires_none_into_hypercorn(monkeypatch):
+    """With neither arg nor env var, Hypercorn config has ``websocket_ping_interval = None``."""
+    monkeypatch.delenv(InvocationsWSConstants.ENV_WS_KEEPALIVE_INTERVAL, raising=False)
+    app = InvocationAgentServerHost()
+    config = app._build_hypercorn_config("0.0.0.0", 8088)  # noqa: SLF001
+    assert config.websocket_ping_interval is None  # type: ignore[attr-defined]
+
+
+def test_ws_ping_interval_positive_wires_into_hypercorn():
+    """A positive interval propagates as a float to Hypercorn config."""
+    app = InvocationAgentServerHost(ws_ping_interval=12)
+    config = app._build_hypercorn_config("0.0.0.0", 8088)  # noqa: SLF001
+    assert config.websocket_ping_interval == 12.0  # type: ignore[attr-defined]
+
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py
index 4972b76275c7..8629dc6078ad 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_server_routes.py
@@ -21,15 +21,24 @@
 # Route registration
 # ---------------------------------------------------------------------------
 
-def test_ws_route_is_registered():
-    """The /invocations_ws route exists alongside the HTTP routes."""
-    app = InvocationAgentServerHost()
+def test_ws_route_is_registered_when_handler_is_set():
+    """The /invocations_ws route is registered lazily on @ws_handler."""
+    app = _make_echo_ws_app()
     paths = [getattr(r, "path", None) for r in app.routes]
     assert "/invocations_ws" in paths
     assert "/invocations" in paths
     assert "/readiness" in paths
 
 
+def test_ws_route_is_not_registered_without_handler():
+    """Without @ws_handler the WS route is absent (upgrades return 404)."""
+    app = InvocationAgentServerHost()
+    paths = [getattr(r, "path", None) for r in app.routes]
+    assert "/invocations_ws" not in paths
+    # HTTP routes still registered.
+    assert "/invocations" in paths
+
+
 def test_readiness_still_works_with_ws_registered():
     """Adding the WS route doesn't break /readiness."""
     app = _make_echo_ws_app()
@@ -81,7 +90,7 @@ def test_ws_upgrade_on_http_path_fails():
     client = TestClient(app)
     # /invocations is a Route, not a WebSocketRoute — TestClient surfaces
     # this as an immediate WebSocketDisconnect rather than a connect.
-    with pytest.raises((WebSocketDisconnect, Exception)):  # noqa: PT011
+    with pytest.raises(WebSocketDisconnect):
         with client.websocket_connect("/invocations"):
             pass
 
@@ -90,6 +99,6 @@ def test_ws_unknown_path_fails():
     """An unknown WS path is rejected."""
     app = _make_echo_ws_app()
     client = TestClient(app)
-    with pytest.raises((WebSocketDisconnect, Exception)):  # noqa: PT011
+    with pytest.raises(WebSocketDisconnect):
         with client.websocket_connect("/nonexistent"):
             pass
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py
index 43107ed71b7e..2163e8cd9286 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_session_id.py
@@ -20,8 +20,8 @@
 # ---------------------------------------------------------------------------
 
 def _session_ids_from_records(records):
-    """Pull ``ws.session_id`` from each structured close-event record."""
-    return [getattr(r, "ws.session_id") for r in _records_with_ws_extras(records)]
+    """Pull ``azure.ai.agentserver.invocations_ws.session_id`` from each structured close-event record."""
+    return [getattr(r, "azure.ai.agentserver.invocations_ws.session_id") for r in _records_with_ws_extras(records)]
 
 
 # ---------------------------------------------------------------------------
@@ -56,7 +56,7 @@ def test_ws_session_id_is_unique_per_connection(caplog):
                 ws.receive_text()
 
     session_ids = _session_ids_from_records(caplog.records)
-    assert len(session_ids) == 5
+    # Uniqueness is the meaningful invariant — implies len == 5 too.
     assert len(set(session_ids)) == 5
 
 
@@ -78,3 +78,24 @@ def test_ws_session_id_ignores_query_param(caplog):
     # generates a fresh server-side UUID per connection.
     assert session_ids[-1] != "my-custom-session"
     uuid.UUID(session_ids[-1])
+
+
+def test_ws_session_id_uses_foundry_env_var(caplog, monkeypatch):
+    """When ``FOUNDRY_AGENT_SESSION_ID`` is set, every WS session reports it.
+
+    Parity with HTTP ``POST /invocations``: the platform-injected session
+    ID must propagate to both transports so cross-transport correlation
+    works on the same container.
+    """
+    monkeypatch.setenv("FOUNDRY_AGENT_SESSION_ID", "platform-session-abc")
+    app = _make_echo_ws_app()
+    client = TestClient(app)
+
+    with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"):
+        for _ in range(3):
+            with client.websocket_connect("/invocations_ws") as ws:
+                ws.send_text("ping")
+                ws.receive_text()
+
+    session_ids = _session_ids_from_records(caplog.records)
+    assert session_ids == ["platform-session-abc"] * 3
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py
index f78b0c815c92..a83dcfe18690 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py
@@ -200,7 +200,7 @@ def test_ws_session_id_attribute_set_on_span():
     spans = _ws_session_spans(_get_spans())
     attrs = dict(spans[0].attributes)
     # WS-specific key
-    ws_session = attrs.get("ws.session_id")
+    ws_session = attrs.get("azure.ai.agentserver.invocations_ws.session_id")
     # General session key (shared with HTTP)
     ms_session = attrs.get("microsoft.session.id")
     assert ws_session and isinstance(ws_session, str)
@@ -219,7 +219,7 @@ def test_ws_session_id_is_unique_per_span():
 
     spans = _ws_session_spans(_get_spans())
     assert len(spans) == 2
-    session_ids = {dict(s.attributes).get("ws.session_id") for s in spans}
+    session_ids = {dict(s.attributes).get("azure.ai.agentserver.invocations_ws.session_id") for s in spans}
     assert len(session_ids) == 2
 
 
@@ -237,7 +237,7 @@ def test_ws_span_records_close_code_1000_on_clean_return():
 
     spans = _ws_session_spans(_get_spans())
     attrs = dict(spans[0].attributes)
-    assert attrs.get("ws.close_code") == 1000
+    assert attrs.get("azure.ai.agentserver.invocations_ws.close_code") == 1000
 
 
 def test_ws_span_records_close_code_1011_on_handler_exception():
@@ -251,11 +251,11 @@ def test_ws_span_records_close_code_1011_on_handler_exception():
 
     spans = _ws_session_spans(_get_spans())
     attrs = dict(spans[0].attributes)
-    assert attrs.get("ws.close_code") == 1011
+    assert attrs.get("azure.ai.agentserver.invocations_ws.close_code") == 1011
 
 
 def test_ws_span_records_duration_ms():
-    """``ws.duration_ms`` is set as a non-negative integer."""
+    """``azure.ai.agentserver.invocations_ws.duration_ms`` is set as a non-negative integer."""
     server = _make_tracing_ws_server()
     client = TestClient(server)
     with client.websocket_connect("/invocations_ws") as ws:
@@ -264,7 +264,7 @@ def test_ws_span_records_duration_ms():
 
     spans = _ws_session_spans(_get_spans())
     attrs = dict(spans[0].attributes)
-    duration_ms = attrs.get("ws.duration_ms")
+    duration_ms = attrs.get("azure.ai.agentserver.invocations_ws.duration_ms")
     assert isinstance(duration_ms, int)
     assert duration_ms >= 0
 
@@ -287,7 +287,7 @@ async def handler(websocket: WebSocket) -> None:
 
     spans = _ws_session_spans(_get_spans())
     attrs = dict(spans[0].attributes)
-    assert attrs.get("ws.close_code") == 4010
+    assert attrs.get("azure.ai.agentserver.invocations_ws.close_code") == 4010
 
 
 # ---------------------------------------------------------------------------
@@ -310,10 +310,13 @@ def test_ws_handler_exception_records_exception_on_span():
     # Exception is recorded as an event by record_error.
     event_names = [e.name for e in span.events]
     assert "exception" in event_names
+    # Exactly one exception event — belt-and-suspenders against future
+    # double-recording on the same span.
+    assert event_names.count("exception") == 1
 
 
-def test_ws_no_handler_records_error_on_span():
-    """A connection refused for lack of @ws_handler still produces an error-tagged span."""
+def test_ws_no_handler_produces_no_span():
+    """Without @ws_handler the route is absent, so no WS span is created."""
     app = _make_tracing_ws_no_handler()
     client = TestClient(app)
     with pytest.raises(WebSocketDisconnect):
@@ -321,14 +324,7 @@ def test_ws_no_handler_records_error_on_span():
             pass
 
     spans = _ws_session_spans(_get_spans())
-    assert spans
-    attrs = dict(spans[0].attributes)
-    # Close code is 1011 and the error_code/message attrs are populated.
-    assert attrs.get("ws.close_code") == 1011
-    assert attrs.get("azure.ai.agentserver.invocations_ws.error.code") == "not_implemented"
-    assert "ws_handler" in str(
-        attrs.get("azure.ai.agentserver.invocations_ws.error.message", "")
-    )
+    assert spans == []
 
 
 # ---------------------------------------------------------------------------
@@ -383,14 +379,11 @@ def test_ws_agent_name_and_version_in_span_name():
     assert "3.1" in name
 
 
-def test_ws_agent_name_only_in_span_name():
+def test_ws_agent_name_only_in_span_name(monkeypatch):
     """Agent name without version still appears in span name."""
-    env_override = {"FOUNDRY_AGENT_NAME": "ws-solo"}
-    env_copy = os.environ.copy()
-    env_copy.pop("FOUNDRY_AGENT_VERSION", None)
-    env_copy.update(env_override)
-    with patch.dict(os.environ, env_copy, clear=True):
-        server = _make_tracing_ws_server()
+    monkeypatch.setenv("FOUNDRY_AGENT_NAME", "ws-solo")
+    monkeypatch.delenv("FOUNDRY_AGENT_VERSION", raising=False)
+    server = _make_tracing_ws_server()
 
     client = TestClient(server)
     with client.websocket_connect("/invocations_ws") as ws: