From 4a67762e0050c8c9d58ab03201c10054da74afe5 Mon Sep 17 00:00:00 2001 From: Yulin Li Date: Tue, 19 May 2026 02:58:25 +0000 Subject: [PATCH] Drop framework-level OpenTelemetry span for WebSocket connections Aligns the invocations_ws transport with the new tracing model on feature/enable-a365-tracing: trace context propagation is handled by the core TraceContextMiddleware, and user-created spans inside handlers are correctly parented without a framework-generated root span. - _invocation_ws.py: remove the per-connection 'websocket_session' span (request_span call, otel_span, _safe_set_attrs helper, end_span import, span_ctx finalize block). Drop the dead 'handler_exc' kwarg on _finalize_session and the 'error_message' kwarg on _emit_close_event (handler exception text never reaches the close-event log line, by design). - _constants.py: drop unused ATTR_SPAN_ERROR_MESSAGE; relabel the section comment from 'Span attribute keys' to 'Structured-log extra keys'. - CHANGELOG.md: trim the WS telemetry feature bullet to describe the log-only path; no breaking-change entry (the WS span was added and removed within the same unreleased 1.0.0b4). - tests/test_ws_tracing.py: removed (17 span-assertion tests obsolete; close-event log line still covered by tests/test_ws_close_event.py). 186 tests pass, 2 skipped. --- .../CHANGELOG.md | 2 +- .../ai/agentserver/invocations/_constants.py | 7 +- .../agentserver/invocations/_invocation_ws.py | 141 ++---- .../tests/test_ws_tracing.py | 463 ------------------ 4 files changed, 36 insertions(+), 577 deletions(-) delete mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index 88e9e043f6e4..f173752e130d 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -7,7 +7,7 @@ - Error source classification headers: All HTTP error responses now include `x-platform-error-source` with a value of `user`, `platform`, or `upstream` to indicate which component caused the error. Developer handler exceptions and missing handler registrations are classified as `upstream`. Exceptions tagged with the platform error tag are classified as `platform` and additionally include `x-platform-error-detail` with truncated exception details (max 2048 characters) for diagnostics. - WebSocket protocol support — `InvocationAgentServerHost` now hosts `/invocations_ws` alongside `POST /invocations`. Register the handler with the new `@app.ws_handler` decorator. The route is registered lazily on first decoration, so hosts without a registered handler return HTTP 404. - WebSocket Ping/Pong keep-alive — disabled by default; enable by setting the `WS_KEEPALIVE_INTERVAL` env var (auto-injected by AgentService into hosted-agent containers; surfaced on `app.config.ws_ping_interval` in `azure-ai-agentserver-core>=2.0.0b4`). `0` (or unset) disables keep-alive. Wired through to Hypercorn's `websocket_ping_interval` by `AgentServerHost._build_hypercorn_config`. -- WebSocket telemetry — single connection-scoped `websocket_session` OpenTelemetry span per WS connection, plus a structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms`. Session ID honours the `FOUNDRY_AGENT_SESSION_ID` env var for HTTP/WS correlation. +- WebSocket telemetry — structured close-event log line carrying `azure.ai.agentserver.invocations_ws.session_id`, `close_code`, and `duration_ms` (via the standard `logging` `extra` dict). 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/azure/ai/agentserver/invocations/_constants.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_constants.py index a704df3c530d..760fae10f563 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 @@ -26,8 +26,8 @@ class InvocationConstants: class InvocationsWSConstants: """invocations_ws (WebSocket) protocol constants. - Route, span attribute keys, and ping/pong defaults for the - WebSocket endpoint hosted alongside the HTTP invocations protocol. + Route, structured-log field keys, and close codes for the WebSocket + endpoint hosted alongside the HTTP invocations protocol. """ # Route @@ -37,9 +37,8 @@ class InvocationsWSConstants: CLOSE_NORMAL = 1000 # handler returned cleanly CLOSE_INTERNAL_ERROR = 1011 # handler raised an unhandled exception - # Span attribute keys + # 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" - 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_ws.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation_ws.py index 86246eaf540e..ffdaad5ad2e8 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 @@ -14,7 +14,7 @@ upstream proxy / load-balancer idle timeouts; * 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 +* 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``. @@ -31,15 +31,14 @@ from azure.ai.agentserver.core import ( # pylint: disable=no-name-in-module AgentServerHost, - end_span, ) 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 +# mypy / pyright (resolves ``self.config`` and ``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: @@ -53,23 +52,6 @@ WSHandler = Callable[[WebSocket], Awaitable[None]] -def _safe_set_attrs(span: Any, attrs: dict[str, Any]) -> None: - """Best-effort ``span.set_attribute`` for a batch of attributes. - - :param span: The OTel span (or ``None`` when tracing is disabled). - :type span: any - :param attrs: Mapping of attribute keys to values. - :type attrs: dict[str, any] - """ - if span is None: - return - try: - for key, value in attrs.items(): - span.set_attribute(key, value) - except Exception: # pylint: disable=broad-exception-caught - logger.debug("Failed to set WS span attributes: %s", list(attrs.keys()), exc_info=True) - - class _WSHandlerMixin(_MixinBase): """Pure mixin that adds the ``@app.ws_handler`` decorator and ``/invocations_ws`` route. @@ -79,8 +61,8 @@ class _WSHandlerMixin(_MixinBase): 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``). + ``self.router``) are accessed via duck typing and are typed only for + the static checkers (see ``_MixinBase``). """ # Slots populated by __init__. @@ -216,21 +198,6 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None: 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 - # the client sent is honoured for parenting child spans inside the - # user handler. ``end_on_exit=False`` so we can attach the close - # code / duration before ending. - span_ctx = self.request_span( - websocket.headers, - session_id, - "websocket_session", - operation_name="websocket_session", - session_id=session_id, - end_on_exit=False, - ) - otel_span = span_ctx.__enter__() - _safe_set_attrs(otel_span, {InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id}) - # 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. @@ -241,12 +208,9 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None: except Exception as exc: # pylint: disable=broad-exception-caught await self._finalize_session( websocket=None, - span_ctx=span_ctx, - otel_span=otel_span, session_id=session_id, start_ns=start_ns, close_code=InvocationsWSConstants.CLOSE_INTERNAL_ERROR, - handler_exc=exc, error_code="accept_failed", ) logger.error( @@ -262,16 +226,16 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None: 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 on - # the span, then re-raise via ``finally`` so cancellation is - # never swallowed. + # 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 \u2014 emits the close-event log line, ends the - # span, and best-effort closes the socket \u2014 even when the - # handler raised a ``BaseException`` like ``CancelledError``. + # 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 @@ -281,12 +245,9 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None: error_code = "cancelled" await self._finalize_session( websocket=websocket, - span_ctx=span_ctx, - otel_span=otel_span, session_id=session_id, start_ns=start_ns, close_code=close_code, - handler_exc=handler_exc, error_code=error_code, ) @@ -332,35 +293,25 @@ async def _finalize_session( self, *, websocket: Optional[WebSocket], - span_ctx: Any, - otel_span: Any, session_id: str, start_ns: int, close_code: int, - handler_exc: Optional[BaseException], error_code: Optional[str], ) -> None: - """Close the WS (best-effort), emit metrics, and end the span. + """Close the WS (best-effort) and emit the close-event log line. Called from both the success path and the accept-failure path. :keyword websocket: The connected WebSocket, or ``None`` when the ASGI ``accept`` itself failed (no socket to close). :paramtype websocket: Optional[~starlette.websockets.WebSocket] - :keyword span_ctx: The active ``request_span`` context manager. - :paramtype span_ctx: any - :keyword otel_span: The current OTel span (or ``None`` when tracing is off). - :paramtype otel_span: any :keyword session_id: Per-connection session ID. :paramtype session_id: str :keyword start_ns: ``time.monotonic_ns()`` at connection start. :paramtype start_ns: int :keyword close_code: The RFC 6455 code to report to the client. :paramtype close_code: int - :keyword handler_exc: Unhandled exception raised by the user handler, - or ``None`` for a clean close. - :paramtype handler_exc: Optional[BaseException] - :keyword error_code: Short error tag for span / log; ``None`` for success. + :keyword error_code: Short error tag for the log line; ``None`` for success. :paramtype error_code: Optional[str] """ duration_ms = (time.monotonic_ns() - start_ns) // 1_000_000 @@ -387,90 +338,62 @@ async def _finalize_session( ) self._emit_close_event( - otel_span, session_id, close_code, duration_ms, error_code=error_code, - error_message=str(handler_exc) if handler_exc is not None else None, ) - # 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 # ------------------------------------------------------------------ @staticmethod def _emit_close_event( - otel_span: Any, session_id: str, close_code: int, duration_ms: int, *, error_code: Optional[str] = None, - error_message: Optional[str] = None, ) -> None: - """Record close-event span attributes and emit a structured log line. + """Emit the structured close-event log line for one WS connection. - The log record carries the ``azure.ai.agentserver.invocations_ws.session_id``, + The log record carries ``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 + ``azure.ai.agentserver.invocations_ws.duration_ms`` 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. + OTel logging bridge can pick them up directly without parsing the + message. Exception details are deliberately NOT included here; they + flow through ``logger.error(..., exc_info=True)`` in + ``_invoke_user_handler`` instead. - :param otel_span: The connection span (or ``None`` when tracing is off). - :type otel_span: any :param session_id: Per-connection session ID. :type session_id: str :param close_code: The RFC 6455 close code reported to the client. :type close_code: int :param duration_ms: Connection duration in milliseconds (monotonic). :type duration_ms: int - :keyword error_code: Optional short error tag for span attribute. + :keyword error_code: Optional short error tag for the log record. :paramtype error_code: Optional[str] - :keyword error_message: Optional human-readable error message for - span attribute (NOT included in the log line — exception details - are logged separately by ``logger.error(..., exc_info=True)``). - :paramtype error_message: Optional[str] """ - attrs: dict[str, Any] = { + extra: dict[str, Any] = { InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id, InvocationsWSConstants.ATTR_SPAN_CLOSE_CODE: close_code, InvocationsWSConstants.ATTR_SPAN_DURATION_MS: duration_ms, } if error_code: - attrs[InvocationsWSConstants.ATTR_SPAN_ERROR_CODE] = error_code - if error_message: - 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. + 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`. + # 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__[""]`` for plain ``logging``. logger.info( "invocations_ws connection closed: session_id=%s code=%s duration_ms=%s", session_id, close_code, duration_ms, - extra={ - InvocationsWSConstants.ATTR_SPAN_SESSION_ID: session_id, - InvocationsWSConstants.ATTR_SPAN_CLOSE_CODE: close_code, - InvocationsWSConstants.ATTR_SPAN_DURATION_MS: duration_ms, - }, + extra=extra, ) 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 deleted file mode 100644 index a83dcfe18690..000000000000 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_ws_tracing.py +++ /dev/null @@ -1,463 +0,0 @@ -# --------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# --------------------------------------------------------- -"""Tests for OpenTelemetry tracing on the ``invocations_ws`` protocol. - -These mirror :mod:`tests.test_tracing` (HTTP /invocations) but target the -WebSocket transport: span creation, GenAI attributes, traceparent -propagation, close-code recording, and error recording. -""" -import os -import uuid -from unittest.mock import patch - -import pytest -from starlette.testclient import TestClient -from starlette.websockets import WebSocket, WebSocketDisconnect - -from azure.ai.agentserver.invocations import InvocationAgentServerHost - - -# --------------------------------------------------------------------------- -# Lazy OTel setup with in-memory exporter (mirrors tests/test_tracing.py) -# --------------------------------------------------------------------------- -try: - from opentelemetry import trace - from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider - from opentelemetry.sdk.trace.export import SimpleSpanProcessor - from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter - - _HAS_OTEL = True -except ImportError: - _HAS_OTEL = False - -_MODULE_PROVIDER = None -_MODULE_EXPORTER = None -_MODULE_SETUP_DONE = False - -pytestmark = pytest.mark.skipif(not _HAS_OTEL, reason="opentelemetry not installed") - - -@pytest.fixture(autouse=True) -def _clear_spans(): - """Set up the OTel provider on first use, then clear spans before each test.""" - global _MODULE_PROVIDER, _MODULE_EXPORTER, _MODULE_SETUP_DONE - if _HAS_OTEL and not _MODULE_SETUP_DONE: - _existing = trace.get_tracer_provider() - if hasattr(_existing, "add_span_processor"): - _MODULE_PROVIDER = _existing - else: - _MODULE_PROVIDER = SdkTracerProvider() - trace.set_tracer_provider(_MODULE_PROVIDER) - _MODULE_EXPORTER = InMemorySpanExporter() - _MODULE_PROVIDER.add_span_processor(SimpleSpanProcessor(_MODULE_EXPORTER)) - _MODULE_SETUP_DONE = True - if _MODULE_EXPORTER: - _MODULE_EXPORTER.clear() - - -def _get_spans(): - """Return all captured spans.""" - if _MODULE_EXPORTER: - return _MODULE_EXPORTER.get_finished_spans() - return [] - - -def _ws_session_spans(spans): - """Filter spans created by the ``/invocations_ws`` route.""" - return [s for s in spans if "websocket_session" in s.name] - - -# --------------------------------------------------------------------------- -# Helpers: tracing-enabled servers -# --------------------------------------------------------------------------- - -def _make_tracing_ws_server(**kwargs) -> InvocationAgentServerHost: - """Build an echo WS host with tracing enabled.""" - with patch.dict( - os.environ, - {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, - ): - with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): - server = InvocationAgentServerHost(**kwargs) - - @server.ws_handler - async def handler(websocket: WebSocket) -> None: - async for msg in websocket.iter_text(): - await websocket.send_text(msg) - - return server - - -def _make_failing_tracing_ws_server(**kwargs) -> InvocationAgentServerHost: - """Build a WS host whose handler raises after one received frame.""" - with patch.dict( - os.environ, - {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, - ): - with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): - server = InvocationAgentServerHost(**kwargs) - - @server.ws_handler - async def handler(websocket: WebSocket) -> None: - await websocket.receive_text() - raise ValueError("ws tracing error test") - - return server - - -def _make_tracing_ws_no_handler(**kwargs) -> InvocationAgentServerHost: - """Build a tracing-enabled host with NO ws_handler registered.""" - with patch.dict( - os.environ, - {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, - ): - with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): - return InvocationAgentServerHost(**kwargs) - - -# --------------------------------------------------------------------------- -# Span creation -# --------------------------------------------------------------------------- - -def test_ws_tracing_creates_websocket_session_span(): - """A WS connection creates a span named ``websocket_session``.""" - server = _make_tracing_ws_server() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hello") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert len(spans) == 1 - assert spans[0].name.startswith("websocket_session") - - -def test_ws_tracing_disabled_by_default_still_creates_span(): - """Even without an exporter, the global tracer creates the span.""" - if _MODULE_EXPORTER: - _MODULE_EXPORTER.clear() - - app = InvocationAgentServerHost() - - @app.ws_handler - async def handler(websocket: WebSocket) -> None: - await websocket.send_text("ready") - - client = TestClient(app) - with client.websocket_connect("/invocations_ws") as ws: - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert len(spans) >= 1 - - -# --------------------------------------------------------------------------- -# GenAI attributes (parity with test_genai_attributes_on_invoke_span) -# --------------------------------------------------------------------------- - -def test_ws_span_has_genai_attributes(): - """WS span carries the standard GenAI / service.name attributes.""" - server = _make_tracing_ws_server() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert spans - attrs = dict(spans[0].attributes) - assert attrs.get("gen_ai.system") == "azure.ai.agentserver" - assert attrs.get("gen_ai.provider.name") == "AzureAI Hosted Agents" - assert attrs.get("service.name") == "azure.ai.agentserver" - - -def test_ws_span_operation_name_attribute(): - """``gen_ai.operation.name`` on the WS span is ``websocket_session``.""" - server = _make_tracing_ws_server() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - attrs = dict(spans[0].attributes) - assert attrs.get("gen_ai.operation.name") == "websocket_session" - - -# --------------------------------------------------------------------------- -# Session ID attribute (parity with test_session_id_in_conversation_id) -# --------------------------------------------------------------------------- - -def test_ws_session_id_attribute_set_on_span(): - """The auto-generated session_id is stored on the span under both keys.""" - server = _make_tracing_ws_server() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - attrs = dict(spans[0].attributes) - # WS-specific key - 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) - uuid.UUID(ws_session) - assert ms_session == ws_session - - -def test_ws_session_id_is_unique_per_span(): - """Two WS connections produce two distinct session IDs on their spans.""" - server = _make_tracing_ws_server() - client = TestClient(server) - for _ in range(2): - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("x") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert len(spans) == 2 - session_ids = {dict(s.attributes).get("azure.ai.agentserver.invocations_ws.session_id") for s in spans} - assert len(session_ids) == 2 - - -# --------------------------------------------------------------------------- -# Close-code / duration attributes (WS-specific) -# --------------------------------------------------------------------------- - -def test_ws_span_records_close_code_1000_on_clean_return(): - """A handler that returns normally records ws.close_code = 1000.""" - server = _make_tracing_ws_server() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - attrs = dict(spans[0].attributes) - assert attrs.get("azure.ai.agentserver.invocations_ws.close_code") == 1000 - - -def test_ws_span_records_close_code_1011_on_handler_exception(): - """A failing handler records ws.close_code = 1011.""" - server = _make_failing_tracing_ws_server() - client = TestClient(server) - with pytest.raises(WebSocketDisconnect): - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("trigger") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - attrs = dict(spans[0].attributes) - assert attrs.get("azure.ai.agentserver.invocations_ws.close_code") == 1011 - - -def test_ws_span_records_duration_ms(): - """``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: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - attrs = dict(spans[0].attributes) - duration_ms = attrs.get("azure.ai.agentserver.invocations_ws.duration_ms") - assert isinstance(duration_ms, int) - assert duration_ms >= 0 - - -def test_ws_span_records_client_close_code(): - """A client closing with a custom code lands on the span as ws.close_code.""" - server_with_recv = _make_tracing_ws_server() - - # Override the handler so receive_text raises with the client's code - # (iter_text swallows WebSocketDisconnect — see _invocation_ws docs). - @server_with_recv.ws_handler - async def handler(websocket: WebSocket) -> None: - while True: - await websocket.receive_text() - - client = TestClient(server_with_recv) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("ping") - ws.close(code=4010) - - spans = _ws_session_spans(_get_spans()) - attrs = dict(spans[0].attributes) - assert attrs.get("azure.ai.agentserver.invocations_ws.close_code") == 4010 - - -# --------------------------------------------------------------------------- -# Error recording on span (parity with test_invoke_error_records_exception) -# --------------------------------------------------------------------------- - -def test_ws_handler_exception_records_exception_on_span(): - """When the handler raises, the span has ERROR status and exception event.""" - server = _make_failing_tracing_ws_server() - client = TestClient(server) - with pytest.raises(WebSocketDisconnect): - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("trigger") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert spans - span = spans[0] - assert span.status.status_code.name == "ERROR" - # 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_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): - with client.websocket_connect("/invocations_ws"): - pass - - spans = _ws_session_spans(_get_spans()) - assert spans == [] - - -# --------------------------------------------------------------------------- -# Traceparent propagation (parity with test_traceparent_propagation) -# --------------------------------------------------------------------------- - -def test_ws_traceparent_propagation(): - """A client-supplied traceparent header parents the WS span.""" - server = _make_tracing_ws_server() - trace_id_hex = uuid.uuid4().hex - span_id_hex = uuid.uuid4().hex[:16] - traceparent = f"00-{trace_id_hex}-{span_id_hex}-01" - - client = TestClient(server) - with client.websocket_connect( - "/invocations_ws", - headers={"traceparent": traceparent}, - ) as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert spans - actual_trace_id = format(spans[0].context.trace_id, "032x") - assert actual_trace_id == trace_id_hex - - -# --------------------------------------------------------------------------- -# Agent name / version in span name (parity with test_agent_name_in_span_name) -# --------------------------------------------------------------------------- - -def test_ws_agent_name_and_version_in_span_name(): - """``FOUNDRY_AGENT_NAME``/``_VERSION`` appear in the WS span name.""" - with patch.dict( - os.environ, - { - "FOUNDRY_AGENT_NAME": "my-ws-agent", - "FOUNDRY_AGENT_VERSION": "3.1", - }, - ): - server = _make_tracing_ws_server() - - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert spans - name = spans[0].name - assert "my-ws-agent" in name - assert "3.1" in name - - -def test_ws_agent_name_only_in_span_name(monkeypatch): - """Agent name without version still appears in span name.""" - 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: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert spans - assert "ws-solo" in spans[0].name - - -# --------------------------------------------------------------------------- -# Span kind / shape -# --------------------------------------------------------------------------- - -def test_ws_span_kind_is_server(): - """The WS span is created with ``kind=SERVER``.""" - server = _make_tracing_ws_server() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert spans - # SpanKind.SERVER == 2; compare by name to avoid pinning the enum value. - assert spans[0].kind.name == "SERVER" - - -def test_ws_span_has_no_parent_when_no_traceparent(): - """Without traceparent the WS span is a root span (parent is None).""" - server = _make_tracing_ws_server() - client = TestClient(server) - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _ws_session_spans(_get_spans()) - assert spans - assert spans[0].parent is None - - -# --------------------------------------------------------------------------- -# Coexistence: invoke + ws spans on the same host -# --------------------------------------------------------------------------- - -def test_ws_and_invoke_spans_coexist(): - """A host serving both HTTP /invocations and /invocations_ws produces both spans.""" - from starlette.requests import Request - from starlette.responses import Response - - with patch.dict( - os.environ, - {"APPLICATIONINSIGHTS_CONNECTION_STRING": "InstrumentationKey=00000000-0000-0000-0000-000000000000"}, - ): - with patch("azure.ai.agentserver.core._tracing._setup_distro_export", create=True): - server = InvocationAgentServerHost() - - @server.invoke_handler - async def http_handler(request: Request) -> Response: - return Response(content=b"ok") - - @server.ws_handler - async def ws_handler(websocket: WebSocket) -> None: - async for msg in websocket.iter_text(): - await websocket.send_text(msg) - - client = TestClient(server) - client.post("/invocations", content=b"data") - with client.websocket_connect("/invocations_ws") as ws: - ws.send_text("hi") - ws.receive_text() - - spans = _get_spans() - invoke = [s for s in spans if "invoke_agent" in s.name] - ws_spans = _ws_session_spans(spans) - assert invoke and ws_spans