Yulin Li (yulin-li) · Jul 11, 2026 · Jul 11, 2026 · Jul 11, 2026 · Jul 11, 2026 · Jul 11, 2026
@@ -13,6 +13,7 @@
 
 ### Bugs Fixed
 
+- Extended W3C trace context and baggage propagation to WebSocket connections so spans created by `invocations_ws` handlers inherit caller context and A365 correlation data.
 - Preserve host-owned recovery callbacks and task options when loading
   process-wide task descriptors at startup.
 - Exclude protocol-disabled task functions before startup, periodic or inline

@@ -1,4 +1,4 @@
 apiMdSha256: e882d606d64b71d6b7582933cb41e627cb0d3c50650496bbdce9e9e5f5a4fde0
-packageVersion: 2.3.0b1
+packageVersion: 2.3.0b2
 parserVersion: 0.3.31
 pythonVersion: 3.11.9
@@ -516,8 +516,9 @@ async def _after_restore_route(request: Request) -> Response:
         )
 
         # 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
 

@@ -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
@@ -496,10 +496,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.
@@ -512,29 +513,48 @@ 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
 
-        # Build a simple dict of headers for the propagators
+        # Preserve split W3C list headers while retaining the existing
+        # last-value behavior for single-value headers.
         raw_headers: list[tuple[bytes, bytes]] = scope.get("headers", [])
-        headers = {k.decode("latin-1"): v.decode("latin-1") for k, v in raw_headers}
+        headers: dict[str, str] = {}
+        duplicate_traceparent = False
+        for raw_name, raw_value in raw_headers:
+            name = raw_name.decode("latin-1").lower()
+            value = raw_value.decode("latin-1")
+            if name in headers:
+                if name == "traceparent":
+                    duplicate_traceparent = True
+                elif name in ("baggage", "tracestate"):
+                    headers[name] = f"{headers[name]},{value}"
+                else:
+                    headers[name] = value
+            else:
+                headers[name] = value
+        if duplicate_traceparent:
+            headers.pop("traceparent", None)
 
-        # Use the global propagator to extract trace context + baggage
-        from opentelemetry.propagate import extract  # pylint: disable=import-outside-toplevel
-
-        ctx = extract(carrier=headers)
-
-        # Add x-request-id as baggage for downstream propagation
-        x_request_id = headers.get("x-request-id")
-        if x_request_id:
-            ctx = _otel_baggage.set_baggage(
-                "x_request_id",
-                x_request_id,
-                context=ctx,
-            )
+        try:
+            # Use the global propagator to extract trace context + baggage
+            from opentelemetry.propagate import extract  # pylint: disable=import-outside-toplevel
+            ctx = extract(carrier=headers)
+
+            # Add x-request-id as baggage for downstream propagation
+            x_request_id = headers.get("x-request-id")
+            if x_request_id:
+                ctx = _otel_baggage.set_baggage(
+                    "x_request_id", x_request_id, context=ctx,
+                )
+
+            token = _otel_context.attach(ctx)
+        except Exception:  # pylint: disable=broad-exception-caught
+            logger.debug("Failed to propagate incoming trace context")
+            await self.app(scope, receive, send)
+            return
 
-        token = _otel_context.attach(ctx)
         try:
             await self.app(scope, receive, send)
         finally:
@@ -912,7 +932,8 @@ def on_emit(self, log_data: Any) -> None:  # pylint: disable=unused-argument
             entries = _otel_baggage.get_all(context=ctx)
             if entries:
                 for key, value in entries.items():
-                    attrs[key] = value  # type: ignore[index]
+                    if key not in attrs:
+                        attrs[key] = value  # type: ignore[index]
 
             attrs[_ATTR_MICROSOFT_FOUNDRY] = _MICROSOFT_FOUNDRY_VALUE
             if self.agent_name and _ATTR_GEN_AI_AGENT_NAME not in attrs:

@@ -2,7 +2,6 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 """Tests for tracing configuration — not invocation spans (those live in the invocations package)."""
-
 import asyncio
 import logging
 import os
@@ -13,15 +12,15 @@
 from unittest import mock
 
 import pytest
-from opentelemetry import baggage as _otel_baggage, context as _otel_context
+from opentelemetry import baggage as _otel_baggage, context as _otel_context, trace
 from opentelemetry._logs import LogRecord
 from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
 from opentelemetry.sdk._logs.export import InMemoryLogRecordExporter, SimpleLogRecordProcessor
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, SpanExportResult
 from opentelemetry.sdk.resources import Resource
 
-from azure.ai.agentserver.core import AgentServerHost
+from azure.ai.agentserver.core import AgentServerHost, _tracing
 from azure.ai.agentserver.core._config import (
     resolve_agent_name,
     resolve_agent_version,
@@ -31,8 +30,8 @@
 from azure.ai.agentserver.core._tracing import (
     _BaggageLogRecordProcessor,
     _FoundryEnrichmentSpanProcessor,
+    TraceContextMiddleware,
 )
-from azure.ai.agentserver.core import _tracing
 
 
 class _CollectorExporter(SpanExporter):
@@ -52,6 +51,110 @@ def force_flush(self, timeout_millis=30000):
         return True
 
 
+def test_trace_context_middleware_preserves_repeated_baggage_headers() -> None:
+    captured_baggage = {}
+
+    async def app(_scope, _receive, _send):
+        captured_baggage.update(_otel_baggage.get_all())
+
+    scope = {
+        "type": "websocket",
+        "headers": [
+            (b"Baggage", b"caller.first=one"),
+            (b"baggage", b"caller.second=two"),
+        ],
+    }
+
+    asyncio.run(TraceContextMiddleware(app)(scope, None, None))
+
+    assert captured_baggage == {
+        "caller.first": "one",
+        "caller.second": "two",
+    }
+
+
+def test_trace_context_middleware_preserves_repeated_tracestate_headers() -> None:
+    captured = {}
+
+    async def app(_scope, _receive, _send):
+        span_context = trace.get_current_span().get_span_context()
+        captured.update(
+            {
+                "is_remote": span_context.is_remote,
+                "vendor1": span_context.trace_state.get("vendor1"),
+                "vendor2": span_context.trace_state.get("vendor2"),
+            }
+        )
+
+    scope = {
+        "type": "websocket",
+        "headers": [
+            (b"traceparent", b"00-11111111111111111111111111111111-2222222222222222-01"),
+            (b"tracestate", b"vendor1=value1"),
+            (b"TraceState", b"vendor2=value2"),
+        ],
+    }
+
+    asyncio.run(TraceContextMiddleware(app)(scope, None, None))
+
+    assert captured == {
+        "is_remote": True,
+        "vendor1": "value1",
+        "vendor2": "value2",
+    }
+
+
+def test_trace_context_middleware_rejects_duplicate_traceparent_headers() -> None:
+    captured_trace_id = None
+
+    async def app(_scope, _receive, _send):
+        nonlocal captured_trace_id
+        captured_trace_id = trace.get_current_span().get_span_context().trace_id
+
+    scope = {
+        "type": "websocket",
+        "headers": [
+            (b"traceparent", b"00-11111111111111111111111111111111-2222222222222222-01"),
+            (b"TraceParent", b"00-33333333333333333333333333333333-4444444444444444-01"),
+        ],
+    }
+
+    asyncio.run(TraceContextMiddleware(app)(scope, None, None))
+
+    assert captured_trace_id == 0
+
+
+def test_trace_context_middleware_propagates_process_control_exceptions() -> None:
+    async def app(_scope, _receive, _send):
+        raise AssertionError("downstream app should not run")
+
+    scope = {"type": "websocket", "headers": []}
+    middleware = TraceContextMiddleware(app)
+
+    with mock.patch("opentelemetry.propagate.extract", side_effect=SystemExit):
+        with pytest.raises(SystemExit):
+            asyncio.run(middleware(scope, None, None))
+
+
+def test_trace_context_middleware_logs_fail_open_propagation_failure(caplog) -> None:
+    called = False
+
+    async def app(_scope, _receive, _send):
+        nonlocal called
+        called = True
+
+    scope = {"type": "websocket", "headers": []}
+    middleware = TraceContextMiddleware(app)
+
+    with caplog.at_level("DEBUG", logger="azure.ai.agentserver"):
+        with mock.patch("opentelemetry.propagate.extract", side_effect=RuntimeError("private-value")):
+            asyncio.run(middleware(scope, None, None))
+
+    assert called is True
+    assert "Failed to propagate incoming trace context" in caplog.text
+    assert "private-value" not in caplog.text
+
+
 # ------------------------------------------------------------------ #
 # Tracing enabled / disabled
 # ------------------------------------------------------------------ #
@@ -967,6 +1070,18 @@ def test_adds_agent_and_fallback_session_attributes(self) -> None:
         assert attrs["gen_ai.agent.version"] == "1.2.3"
         assert attrs["microsoft.session.id"] == "session-fallback-1"
 
+    def test_baggage_does_not_overwrite_existing_attributes(self) -> None:
+        proc = _BaggageLogRecordProcessor()
+        log_data = _FakeLogData({"server.attribute": "server-value"})
+        ctx = _otel_baggage.set_baggage("server.attribute", "caller-value")
+        token = _otel_context.attach(ctx)
+        try:
+            proc.on_emit(log_data)
+        finally:
+            _otel_context.detach(token)
+
+        assert log_data.log_record.attributes["server.attribute"] == "server-value"
+
     def test_prefers_baggage_session_id_over_fallback(self) -> None:
         proc = _BaggageLogRecordProcessor(
             agent_name="agent-a",

@@ -10,6 +10,7 @@
 
 ### 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 while retaining the legacy structured-log field.
 - Apply the Invocations task-runtime switch to both typed Voice WebSocket
   endpoints, including composed Voice/Responses hosts where Responses
   independently enables the task manager.

@@ -297,9 +297,13 @@ 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`.
+- Emits a structured close-event log line carrying `azure.ai.agentserver.session_id`, the legacy `azure.ai.agentserver.invocations_ws.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 telemetry correlation
+
+For each generic `@app.ws_handler` connection, the SDK 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 are therefore parented to the caller's trace and carry A365 session correlation. Generic handlers do not create a connection span; typed Voice hosts provide their own connection tracing.
+
 ### Handler signature
 
 The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`. The full WebSocket API — `iter_text`, `iter_bytes`, `iter_json`, `send_text`, `send_bytes`, `send_json`, `close`, `headers`, `query_params`, `client`, `state` — is available, so application protocols on top of `invocations_ws` are entirely under your control.

@@ -9,9 +9,7 @@
 import contextvars
 import inspect
 import logging
-import re
 import threading
-import uuid
 from collections.abc import AsyncIterator, Awaitable, Callable  # pylint: disable=import-error
 from typing import Any, Optional
 
@@ -36,10 +34,16 @@
     USER_ID,
 )
 from azure.ai.agentserver.core._middleware import _task_runtime_middleware
+from azure.ai.agentserver.core._tracing import _BAGGAGE_SESSION_ID
 from azure.ai.agentserver.core.tasks import resilient_tasks_enabled
 
 from ._constants import InvocationConstants
-from ._invocation_ws import _WSHandlerMixin
+from ._invocation_ws import (  # pylint: disable=unused-import
+    _MAX_ID_LENGTH,
+    _WSHandlerMixin,
+    _sanitize_id,
+    uuid,
+)
 from ._sse import _with_keep_alive
 
 logger = logging.getLogger("azure.ai.agentserver")
@@ -90,12 +94,6 @@ def _classify_error(exc: BaseException) -> tuple[str, Optional[str]]:
         return _ERROR_SOURCE_PLATFORM, detail
     return _ERROR_SOURCE_UPSTREAM, None
 
-
-# Maximum length and allowed characters for user-provided IDs (defense in depth).
-_MAX_ID_LENGTH = 256
-_VALID_ID_RE = re.compile(r"^[a-zA-Z0-9\-_.:]+$")
-
-
 # Context variables for structured logging — concurrency-safe alternative to logger filters.
 _invocation_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("invocation_id", default="")
 _session_id_var: contextvars.ContextVar[str] = contextvars.ContextVar("session_id", default="")
@@ -131,26 +129,6 @@ def _ensure_log_filter() -> None:
         _log_filter_installed = True
 
 
-def _sanitize_id(value: str, fallback: str | None = None) -> str:
-    """Validate a user-provided ID string.
-
-    Returns *value* unchanged when it passes validation, otherwise uses
-    *fallback* or generates a UUID when it is omitted. This prevents malformed IDs from
-    propagating into headers, span attributes, and log messages.
-
-    :param value: The raw ID from a header or query parameter.
-    :type value: str
-    :param fallback: A safe fallback value, or ``None`` to generate a UUID only
-        when *value* fails validation.
-    :type fallback: str | None
-    :return: The validated ID or the fallback.
-    :rtype: str
-    """
-    if not value or len(value) > _MAX_ID_LENGTH or not _VALID_ID_RE.match(value):
-        return fallback if fallback is not None else str(uuid.uuid4())
-    return value
-
-
 def _platform_context_from_request(
     request: Request,
     session_id: str,
@@ -541,9 +519,7 @@ async def _create_invocation_endpoint(self, request: Request) -> Response:
             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)