Azure · yulin-li · Jul 11, 2026 · Jul 11, 2026 · Jul 11, 2026 · Jul 11, 2026
@@ -2,13 +2,9 @@
 
 ## 2.0.0b9 (Unreleased)
 
-### Features Added
-
-### Breaking Changes
-
 ### Bugs Fixed
 
-### Other Changes
+- Extended W3C trace context and baggage propagation to WebSocket connections so spans created by `invocations_ws` handlers inherit caller context and A365 correlation data.
 
 ## 2.0.0b8 (2026-07-22)
 

@@ -289,8 +289,9 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]:  # noqa: RUF
         )
 
         # Extract W3C trace context (traceparent/tracestate) and baggage
-        # from incoming HTTP requests so that any spans created downstream
-        # (e.g. by MAF / agent-framework) are children of the caller's trace.
+        # from incoming HTTP requests and WebSocket connections so that any
+        # spans created downstream (e.g. by MAF / agent-framework) are
+        # children of the caller's trace.
         # We do NOT create a SERVER span ourselves — we only propagate context.
         from azure.ai.agentserver.core._tracing import TraceContextMiddleware  # pylint: disable=import-outside-toplevel
         self.add_middleware(TraceContextMiddleware)

@@ -25,7 +25,7 @@
 **Span operations:**
 
 - :class:`TraceContextMiddleware` — ASGI middleware that extracts W3C trace
-  context and baggage from incoming headers
+  context and baggage from incoming request headers
 - :func:`end_span` / :func:`record_error` — span lifecycle helpers
 - :func:`trace_stream` — wrap streaming responses with span lifecycle
 - :func:`set_current_span` / :func:`detach_context` — explicit context management
@@ -268,10 +268,11 @@ class TraceContextMiddleware:
     """Pure-ASGI middleware that propagates W3C trace context and baggage.
 
     Extracts ``traceparent``, ``tracestate``, and ``baggage`` headers from
-    incoming HTTP requests using the standard W3C propagators and attaches
-    the resulting context for the duration of the request.  This ensures
-    that any spans created downstream (e.g. by agent-framework / MAF) are
-    automatically children of the caller's trace.
+    incoming HTTP requests and WebSocket connections using the standard W3C
+    propagators and attaches the resulting context for the duration of the
+    request or connection.  This ensures that any spans created downstream
+    (e.g. by agent-framework / MAF) are automatically children of the
+    caller's trace.
 
     This middleware does **not** create its own span — it only propagates
     the incoming context so that downstream instrumentation inherits it.
@@ -284,7 +285,7 @@ def __init__(self, app: Any) -> None:
         self.app = app
 
     async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
-        if scope["type"] != "http":
+        if scope["type"] not in ("http", "websocket"):
             await self.app(scope, receive, send)
             return
 

@@ -1,5 +1,11 @@
 # Release History
 
+## 1.0.0b8 (Unreleased)
+
+### Bugs Fixed
+
+- Added A365 telemetry correlation to `invocations_ws` by propagating the WebSocket session ID as OpenTelemetry baggage for handler spans and logs, using the cross-protocol `azure.ai.agentserver.session_id` attribute.
+
 ## 1.0.0b7 (2026-07-22)
 
 ### Features Added

@@ -296,12 +296,12 @@ app.run()
 - 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.
+- Emits a structured close-event log line carrying `azure.ai.agentserver.session_id`, `azure.ai.agentserver.invocations_ws.close_code`, and `azure.ai.agentserver.invocations_ws.duration_ms`.
 - Inherits `/readiness`, OpenTelemetry export, graceful shutdown, and the `x-platform-server` identity header from `azure-ai-agentserver-core`.
 
-### Per-connection tracing
+### Per-connection telemetry correlation
 
-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.
+The SDK does **not** create its own connection ("websocket_session") span. Instead, for each connection it extracts the caller's W3C trace context (`traceparent`/`tracestate`) and baggage from the WebSocket upgrade request and attaches the per-connection session ID as the `azure.ai.agentserver.session_id` baggage entry for the lifetime of the connection. Any child spans your handler opens — e.g. via `opentelemetry.trace.get_tracer(...).start_as_current_span(...)` — are therefore parented to the caller's trace and carry the A365 session correlation, which the core enrichment processor stamps onto them. The connection close is reported on the structured close-event log line described above (`azure.ai.agentserver.session_id`, `...invocations_ws.close_code`, `...invocations_ws.duration_ms`), not as a span.
 
 ### Handler signature
 

@@ -38,7 +38,6 @@ class InvocationsWSConstants:
     CLOSE_INTERNAL_ERROR = 1011  # handler raised an unhandled exception
 
     # Structured-log ``extra`` 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"
@@ -35,6 +35,7 @@
     PLATFORM_ERROR_TAG,
     USER_ID,
 )
+from azure.ai.agentserver.core._tracing import _BAGGAGE_SESSION_ID
 
 from ._constants import InvocationConstants
 from ._invocation_ws import _WSHandlerMixin
@@ -512,7 +513,7 @@ async def _create_invocation_endpoint(self, request: Request) -> Response:
             "azure.ai.agentserver.invocation_id", invocation_id, context=ctx,
         )
         ctx = _otel_baggage.set_baggage(
-            "azure.ai.agentserver.session_id", session_id, context=ctx,
+            _BAGGAGE_SESSION_ID, session_id, context=ctx,
         )
         baggage_token = _otel_context.attach(ctx)
 

@@ -15,7 +15,7 @@
 * a clean close on handler return (code 1000) or a 1011 close on uncaught
   handler exceptions;
 * a structured close-event log line carrying
-  ``azure.ai.agentserver.invocations_ws.session_id``,
+  ``azure.ai.agentserver.session_id``,
   ``azure.ai.agentserver.invocations_ws.close_code``, and
   ``azure.ai.agentserver.invocations_ws.duration_ms``.
 """
@@ -27,11 +27,13 @@
 from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any, Optional
 
+from opentelemetry import baggage as _otel_baggage, context as _otel_context
 from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
 
 from azure.ai.agentserver.core import (  # pylint: disable=no-name-in-module
     AgentServerHost,
 )
+from azure.ai.agentserver.core._tracing import _BAGGAGE_SESSION_ID
 
 from ._constants import InvocationsWSConstants
 
@@ -198,58 +200,74 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None:
         session_id = self.config.session_id or str(uuid.uuid4())
         start_ns = time.monotonic_ns()
 
-        # 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.
+        # Preserve caller baggage extracted by TraceContextMiddleware and add
+        # the session correlation key consumed by the A365 enrichment
+        # processor for child spans and logs.
+        ctx = _otel_context.get_current()
+        ctx = _otel_baggage.set_baggage(
+            _BAGGAGE_SESSION_ID, session_id, context=ctx,
+        )
+        baggage_token = _otel_context.attach(ctx)
         try:
-            await websocket.accept()
-        except Exception as exc:  # pylint: disable=broad-exception-caught
-            await self._finalize_session(
-                websocket=None,
-                session_id=session_id,
-                start_ns=start_ns,
-                close_code=InvocationsWSConstants.CLOSE_INTERNAL_ERROR,
-                error_code="accept_failed",
-            )
-            logger.error(
-                "WebSocket accept failed for session %s: %s",
-                session_id, exc, exc_info=True,
-            )
-            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.
 
-        close_code: int = InvocationsWSConstants.CLOSE_NORMAL
-        handler_exc: Optional[BaseException] = None
-        try:
-            close_code, handler_exc = await self._invoke_user_handler(websocket, session_id)
-        except BaseException as exc:  # pylint: disable=broad-exception-caught
-            # ``_invoke_user_handler`` catches ``Exception`` but not
-            # ``BaseException`` (notably ``asyncio.CancelledError``).  Capture
-            # the exception so the ``finally`` block below can record it,
-            # then re-raise via ``finally`` so cancellation is never
-            # swallowed.
-            close_code = InvocationsWSConstants.CLOSE_INTERNAL_ERROR
-            handler_exc = exc
-            raise
+            # Accept the upgrade *before* invoking the user handler — per spec.
+            try:
+                await websocket.accept()
+            except Exception as exc:  # pylint: disable=broad-exception-caught
+                await self._finalize_session(
+                    websocket=None,
+                    session_id=session_id,
+                    start_ns=start_ns,
+                    close_code=InvocationsWSConstants.CLOSE_INTERNAL_ERROR,
+                    error_code="accept_failed",
+                )
+                logger.error(
+                    "WebSocket accept failed for session %s: %s",
+                    session_id, exc, exc_info=True,
+                )
+                return
+
+            close_code: int = InvocationsWSConstants.CLOSE_NORMAL
+            handler_exc: Optional[BaseException] = None
+            try:
+                close_code, handler_exc = await self._invoke_user_handler(
+                    websocket, session_id
+                )
+            except BaseException as exc:  # pylint: disable=broad-exception-caught
+                # ``_invoke_user_handler`` catches ``Exception`` but not
+                # ``BaseException`` (notably ``asyncio.CancelledError``).
+                # Capture the exception so the ``finally`` block below can
+                # record it, then re-raise via ``finally`` so cancellation is
+                # never swallowed.
+                close_code = InvocationsWSConstants.CLOSE_INTERNAL_ERROR
+                handler_exc = exc
+                raise
+            finally:
+                # Always finalize — emits the close-event log line and
+                # best-effort closes the socket — even when the handler
+                # raised a ``BaseException`` like ``CancelledError``.
+                error_code: Optional[str]
+                if handler_exc is None:
+                    error_code = None
+                elif isinstance(handler_exc, Exception):
+                    error_code = "internal_error"
+                else:
+                    error_code = "cancelled"
+                await self._finalize_session(
+                    websocket=websocket,
+                    session_id=session_id,
+                    start_ns=start_ns,
+                    close_code=close_code,
+                    error_code=error_code,
+                )
         finally:
-            # Always finalize — emits the close-event log line and
-            # best-effort closes the socket — even when the handler
-            # raised a ``BaseException`` like ``CancelledError``.
-            error_code: Optional[str]
-            if handler_exc is None:
-                error_code = None
-            elif isinstance(handler_exc, Exception):
-                error_code = "internal_error"
-            else:
-                error_code = "cancelled"
-            await self._finalize_session(
-                websocket=websocket,
-                session_id=session_id,
-                start_ns=start_ns,
-                close_code=close_code,
-                error_code=error_code,
-            )
+            try:
+                _otel_context.detach(baggage_token)
+            except ValueError:
+                pass
 
     async def _invoke_user_handler(
         self, websocket: WebSocket, session_id: str,
@@ -358,7 +376,7 @@ def _emit_close_event(
     ) -> None:
         """Emit the structured close-event log line for one WS connection.
 
-        The log record carries ``azure.ai.agentserver.invocations_ws.session_id``,
+        The log record carries ``azure.ai.agentserver.session_id``,
         ``azure.ai.agentserver.invocations_ws.close_code``, and
         ``azure.ai.agentserver.invocations_ws.duration_ms`` via the standard
         ``logging`` ``extra`` dict — a structured-logging formatter or an
@@ -377,16 +395,19 @@ def _emit_close_event(
         :paramtype error_code: Optional[str]
         """
         extra: dict[str, Any] = {
-            InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id,
+            _BAGGAGE_SESSION_ID: session_id,
             InvocationsWSConstants.ATTR_SPAN_CLOSE_CODE: close_code,
             InvocationsWSConstants.ATTR_SPAN_DURATION_MS: duration_ms,
         }
         if error_code:
             extra[InvocationsWSConstants.ATTR_SPAN_ERROR_CODE] = error_code
 
-        # NOTE: ``extra`` keys deliberately use dotted names
-        # (``azure.ai.agentserver.invocations_ws.session_id`` etc.) so they
-        # line up 1:1 with the keys defined in :class:`InvocationsWSConstants`.
+        # NOTE: ``extra`` keys deliberately use dotted names so they line up
+        # 1:1 with their source constants — the session-ID key comes from the
+        # shared ``_BAGGAGE_SESSION_ID`` (``azure.ai.agentserver.session_id``)
+        # in ``azure-ai-agentserver-core`` so HTTP and WebSocket logs correlate,
+        # while the close-code / duration / error keys come from
+        # :class:`InvocationsWSConstants`.
         # The trade-off is that printf-style log formatters can't address
         # them directly — use a structured (JSON / OTel) formatter, or
         # access via ``LogRecord.__dict__["<key>"]`` for plain ``logging``.

@@ -2,4 +2,4 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-VERSION = "1.0.0b7"
+VERSION = "1.0.0b8"
@@ -21,7 +21,7 @@ classifiers = [
 keywords = ["azure", "azure sdk", "agent", "agentserver", "invocations"]
 
 dependencies = [
-    "azure-ai-agentserver-core>=2.0.0b8",
+    "azure-ai-agentserver-core>=2.0.0b9",
     # Constraint on the transitive aiohttp: the `--pre` CI install otherwise
     # resolves the unbuildable aiohttp 4.0.0a1 alpha. Cap must be <4.0.0a0
     # (<4.0.0 still admits 4.0.0a1 under PEP 440).

@@ -261,7 +261,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, "azure.ai.agentserver.invocations_ws.session_id") and hasattr(r, "azure.ai.agentserver.invocations_ws.close_code")
+        if hasattr(r, "azure.ai.agentserver.session_id") and hasattr(r, "azure.ai.agentserver.invocations_ws.close_code")
     ]
 
 

@@ -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 (``azure.ai.agentserver.invocations_ws.session_id``, ``azure.ai.agentserver.invocations_ws.close_code``, ``azure.ai.agentserver.invocations_ws.duration_ms``) appear
+fields (``azure.ai.agentserver.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.
 """
@@ -42,7 +42,7 @@ 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, "azure.ai.agentserver.invocations_ws.session_id")
+    session_id = getattr(rec, "azure.ai.agentserver.session_id")
     close_code = getattr(rec, "azure.ai.agentserver.invocations_ws.close_code")
     duration_ms = getattr(rec, "azure.ai.agentserver.invocations_ws.duration_ms")
 

@@ -20,8 +20,8 @@
 # ---------------------------------------------------------------------------
 
 def _session_ids_from_records(records):
-    """Pull ``azure.ai.agentserver.invocations_ws.session_id`` from each structured close-event record."""
-    return [getattr(r, "azure.ai.agentserver.invocations_ws.session_id") for r in _records_with_ws_extras(records)]
+    """Pull ``azure.ai.agentserver.session_id`` from each structured close-event record."""
+    return [getattr(r, "azure.ai.agentserver.session_id") for r in _records_with_ws_extras(records)]
 
 
 # ---------------------------------------------------------------------------