diff --git a/.chronus/changes/wujin-voice-tracing-v2-2026-8-17.md b/.chronus/changes/wujin-voice-tracing-v2-2026-8-17.md new file mode 100644 index 000000000000..8d7ddd005e04 --- /dev/null +++ b/.chronus/changes/wujin-voice-tracing-v2-2026-8-17.md @@ -0,0 +1,7 @@ +--- +changeKind: feature +packages: + - azure-ai-agentserver-invocations +--- + +Added content-free Voice connection and callback tracing, application-declared target-turn spans, W3C propagation filtering, source-aware connection termination facts, and unsampled aggregate metrics. \ No newline at end of file diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md index f3c22da547dd..d04a26178df8 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md @@ -1,5 +1,27 @@ # Release History +## 1.1.0b2 (2026-08-17) + +### Features Added + +- Added W3C Voice connection tracing with one `agentserver.connection` span, + content-free `voice.callback` dispatch spans, and aggregate duration and + propagation-failure metrics. +- Added application-declared target-turn tracing through + `Session.start_target_turn`, `TargetTurn.activate`, and explicit + `TargetTurn.complete` outcomes. The SDK does not infer response lifecycle or + own application tasks. +- Added source-aware `Session.termination` for classifying unfinished + application work during connection cleanup. + +### Samples + +- Updated `basic_voice_agent` to declare target turns around real background + generation work and report truthful response, timeout, cancellation, + end-call, and transport outcomes, with bounded per-connection concurrency and + retained model output. Unfinished turns now distinguish clean abandonment, + application/server errors, and protocol or transport loss. + ## 1.1.0b1 (2026-08-11) ### Samples diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/README.md index c7518d70fdb4..abbfc233a1e4 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/README.md @@ -297,13 +297,9 @@ 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.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 tracing - -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. - ### 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. @@ -317,6 +313,8 @@ The handler receives a Starlette [`WebSocket`][starlette-ws] and returns `None`. send-only `Session`: ```python +import asyncio + from azure.ai.agentserver.invocations.voice import ( ResponseCreated, ResponseDone, @@ -325,6 +323,8 @@ from azure.ai.agentserver.invocations.voice import ( SessionReady, SessionRejected, SessionStart, + TargetTurnOrigin, + TargetTurnOutcome, UserMessage, VoiceAgentServerHost, new_item_id, @@ -349,17 +349,44 @@ async def on_session_start(session: Session, event: SessionStart) -> None: async def on_user_message(session: Session, event: UserMessage) -> None: response_id = new_response_id() item_id = new_item_id() - await session.send( - ResponseCreated(response_id=response_id, in_reply_to=(event.item_id,)) - ) - await session.send( - ResponseOutputTextDone( + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + response_started = False + output_item_count = 0 + try: + with turn.activate(): + await session.send( + ResponseCreated(response_id=response_id, in_reply_to=(event.item_id,)) + ) + response_started = True + await session.send( + ResponseOutputTextDone( + response_id=response_id, + item_id=item_id, + text="Hello from the hosted text agent.", + ) + ) + output_item_count = 1 + await session.send(ResponseDone(response_id=response_id)) + turn.complete( + outcome=TargetTurnOutcome.RESPONSE, response_id=response_id, - item_id=item_id, - text="Hello from the hosted text agent.", + output_item_count=1, ) - ) - await session.send(ResponseDone(response_id=response_id)) + except asyncio.CancelledError: + turn.complete( + outcome=TargetTurnOutcome.CANCELLED, + response_id=response_id if response_started else None, + output_item_count=output_item_count, + ) + raise + except Exception: + if not turn.is_completed: + turn.complete( + outcome=TargetTurnOutcome.ERROR, + response_id=response_id if response_started else None, + output_item_count=output_item_count, + ) + raise ``` The submodule is deliberately a thin typed event relay. It decodes one inbound frame, @@ -368,6 +395,37 @@ serializes concurrent WebSocket writes. It does **not** own pending responses, terminal arbitration, timeout/cancel operations, generation tasks, history, or reconnect state. +### Voice tracing + +The typed Voice endpoint extracts W3C context from each WebSocket upgrade and +creates one `agentserver.connection` span for the physical connection. Each +registered event dispatch creates a sibling `voice.callback` span. Application +code opts into a target-decision `invoke_agent` span by calling +`Session.start_target_turn`, activating the returned handle around all model, +tool, retrieval, and custom descendant work, and completing it once with the +application-known outcome. + +```text +Hosted Agents invoke_agent +└── agentserver.connection + ├── voice.callback + └── invoke_agent # only after start_target_turn(...) + └── customer model/tool spans +``` + +The SDK never discovers, retains, or awaits application tasks and never infers +response facts from `Session.send`. Applications that do not call +`start_target_turn` do not receive an automatic target `invoke_agent` span. +`Session.termination` exposes the first source-aware physical terminal fact to +`on_connection_terminating` so the application can classify unfinished work. + +Voice propagation accepts only W3C trace context, the five correlation baggage +keys produced by Hosted Agents, and a bounded `x-request-id`. Transcript, +output, prompt, tool argument/result, and arbitrary baggage content are not +added to SDK spans, metrics, or diagnostics. Valid upstream sampling and +`tracestate` are preserved; unsampled operations still contribute aggregate +connection and target duration metrics. + When the peer or proxy closes the WebSocket, `@app.on_disconnect` receives a local `SessionDisconnected` event. This callback represents only the observed peer disconnect. diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/api.md b/sdk/agentserver/azure-ai-agentserver-invocations/api.md index 65952b6f5b13..983c360ddca3 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/api.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/api.md @@ -467,11 +467,20 @@ namespace azure.ai.agentserver.invocations.voice @experimental class azure.ai.agentserver.invocations.voice.Session: + property termination: SessionTermination | None # Read-only def __init__(self) -> None: ... async def send(self, message: OutboundVoiceMessage) -> None: ... + def start_target_turn( + self, + *, + input_count: int, + origin: TargetTurnOrigin | str, + trigger_context: SpanContext | None = ... + ) -> TargetTurn: ... + @experimental @dataclass(frozen=True, kw_only=True, repr=False) @@ -586,14 +595,16 @@ namespace azure.ai.agentserver.invocations.voice def __hash__() -> None: ... def __init__( + self, + *, + caller: Mapping[str, Any] | None = ..., + greeting: str | None = ..., id: str, - ts: str, + no_input_timeout_ms: int | None = ..., protocol_version: str, reconnect: bool, response_timeouts: ResponseTimeouts, - greeting: str | None = None, - no_input_timeout_ms: int | None = None, - caller: Mapping = None + ts: str ) -> None: ... def __setattr__() -> None: ... @@ -601,6 +612,56 @@ namespace azure.ai.agentserver.invocations.voice def _voice_model_repr(self: Any) -> str: ... + @experimental + class azure.ai.agentserver.invocations.voice.SessionTermination(str, Enum, metaclass=CaseInsensitiveEnumMeta): + ACCEPT_ERROR = "accept_error" + CALLBACK_ERROR = "callback_error" + CANCELLED = "cancelled" + COMPLETED = "completed" + INTERNAL_ERROR = "internal_error" + PROTOCOL_ERROR = "protocol_error" + TRANSPORT_ERROR = "transport_error" + + + @experimental + class azure.ai.agentserver.invocations.voice.TargetTurn: + property is_completed: bool # Read-only + + def __init__(self) -> None: ... + + def activate(self) -> ContextManager[None]: ... + + def complete( + self, + *, + outcome: TargetTurnOutcome | str, + output_item_count: int | None = ..., + response_id: str | None = ... + ) -> None: ... + + + @experimental + class azure.ai.agentserver.invocations.voice.TargetTurnOrigin(str, Enum, metaclass=CaseInsensitiveEnumMeta): + NO_INPUT = "no_input" + OTHER = "other" + PROACTIVE = "proactive" + RECOVERY = "recovery" + USER = "user" + + + @experimental + class azure.ai.agentserver.invocations.voice.TargetTurnOutcome(str, Enum, metaclass=CaseInsensitiveEnumMeta): + ABANDONED = "abandoned" + CANCELLED = "cancelled" + END_CALL = "end_call" + ERROR = "error" + NONE = "none" + OTHER = "other" + RESPONSE = "response" + TIMEOUT = "timeout" + TRANSPORT_ERROR = "transport_error" + + @experimental @dataclass(frozen=True, kw_only=True, repr=False) class azure.ai.agentserver.invocations.voice.UserMessage(_InboundMessage): @@ -680,6 +741,13 @@ namespace azure.ai.agentserver.invocations.voice property routes: list[BaseRoute] # Read-only property ws_ping_interval: float # Read-only + async def __call__( + self, + scope: Scope, + receive: Receive, + send: Send + ) -> None: ... + def __init__( self, *, diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/api.metadata.yml b/sdk/agentserver/azure-ai-agentserver-invocations/api.metadata.yml index 1ec41090eb4d..952d7fdba9da 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/api.metadata.yml +++ b/sdk/agentserver/azure-ai-agentserver-invocations/api.metadata.yml @@ -1,3 +1,3 @@ -apiMdSha256: 8c81198c94f9b6d95391cc58b11d3769ec0de8850d93a2c5e736046e82d25f63 +apiMdSha256: e4e11fea9355277129dc5bea863e6c2252d8df0028b9e30dd2a29a53811389b1 parserVersion: 0.3.31 pythonVersion: 3.11.15 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 760fae10f563..06a1225ba2a1 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 @@ -1,7 +1,9 @@ # --------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -from azure.ai.agentserver.core._platform_headers import SESSION_ID as _SESSION_ID # pylint: disable=import-error,no-name-in-module +from azure.ai.agentserver.core._platform_headers import ( + SESSION_ID as _SESSION_ID, +) # pylint: disable=import-error,no-name-in-module class InvocationConstants: @@ -35,10 +37,21 @@ class InvocationsWSConstants: # Close codes (RFC 6455) CLOSE_NORMAL = 1000 # handler returned cleanly + CLOSE_GOING_AWAY = 1001 # peer is leaving or restarting CLOSE_INTERNAL_ERROR = 1011 # handler raised an unhandled exception + NORMAL_CLOSE_CODES = frozenset({CLOSE_NORMAL, CLOSE_GOING_AWAY}) + PROTOCOL_CLOSE_CODES = frozenset({1002, 1003, 1007, 1008, 1009, 1010}) # 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" + + +def _classify_websocket_close_code(close_code: int) -> str: + if close_code in InvocationsWSConstants.NORMAL_CLOSE_CODES: + return "completed" + if close_code in InvocationsWSConstants.PROTOCOL_CLOSE_CODES: + return "protocol_error" + return "transport_error" diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py index b7e3203b2ccd..868fc81d5d1f 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py @@ -2,4 +2,4 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -VERSION = "1.1.0b1" +VERSION = "1.1.0b2" diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/__init__.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/__init__.py index ec9cfc7169b3..b1ec4b3184e3 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/__init__.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/__init__.py @@ -34,7 +34,8 @@ new_message_id, new_response_id, ) -from ._session import Session +from ._session import Session, SessionTermination +from ._turn import TargetTurn, TargetTurnOrigin, TargetTurnOutcome from ._voice_host import ( BargeInCallback, ConnectionTerminatingCallback, @@ -85,6 +86,10 @@ "SessionRejected", "SessionStart", "SessionStartCallback", + "SessionTermination", + "TargetTurn", + "TargetTurnOrigin", + "TargetTurnOutcome", "UserMessage", "UserMessageCallback", "UserNoInput", diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_session.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_session.py index e6f9483948dc..457691a57aad 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_session.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_session.py @@ -9,16 +9,20 @@ import contextvars import sys from collections.abc import Coroutine, MutableMapping +from enum import Enum from threading import Lock from typing import Any, TypeVar, cast +from opentelemetry.trace import SpanContext from starlette.types import Send from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState from azure.ai.agentserver.core import experimental +from azure.core import CaseInsensitiveEnumMeta from ._codec import encode_outbound_message from ._models import OutboundVoiceMessage, SessionDisconnected +from ._turn import TargetTurn, TargetTurnOrigin _VOICE_SESSION_SCOPE_KEY = "azure.ai.agentserver.invocations.voice.session" _VOICE_CLOSE_CODE_SCOPE_KEY = "azure.ai.agentserver.invocations.voice.close_code" @@ -33,6 +37,19 @@ _TransportResultT = TypeVar("_TransportResultT") +@experimental +class SessionTermination(str, Enum, metaclass=CaseInsensitiveEnumMeta): + """First source-aware terminal fact for one physical Voice connection.""" + + COMPLETED = "completed" + CANCELLED = "cancelled" + PROTOCOL_ERROR = "protocol_error" + ACCEPT_ERROR = "accept_error" + CALLBACK_ERROR = "callback_error" + INTERNAL_ERROR = "internal_error" + TRANSPORT_ERROR = "transport_error" + + class _OperationCancellationState: __slots__ = ("cancellation", "cancellation_requests", "operation_cancel_requested") @@ -230,10 +247,14 @@ async def _run_transport_operation( async def _run_close_attempt( send_lock: asyncio.Lock, + websocket: WebSocket, send: Send, message: dict[str, Any], ) -> None: async with send_lock: + scope = getattr(websocket, "scope", None) + if isinstance(scope, MutableMapping) and _VOICE_DISCONNECT_EVENT_SCOPE_KEY in scope: + return await send(message) @@ -262,7 +283,7 @@ def _start_close_attempt( raise RuntimeError("Voice WebSocket close attempt limit reached") _CLOSE_ATTEMPT_RESERVATIONS += 1 websocket.application_state = WebSocketState.DISCONNECTED - close_coroutine = _run_close_attempt(send_lock, send, message) + close_coroutine = _run_close_attempt(send_lock, websocket, send, message) try: task = contextvars.Context().run(asyncio.create_task, close_coroutine, name="voice_websocket_close") except BaseException as creation_error: # pylint: disable=broad-exception-caught @@ -294,23 +315,37 @@ class Session: registered callbacks. """ - __slots__ = ("_send_lock", "_terminal", "_websocket") + __slots__ = ( + "_connection_context", + "_send_lock", + "_terminal", + "_termination", + "_websocket", + ) + _connection_context: Any _send_lock: asyncio.Lock _terminal: bool + _termination: SessionTermination | None _websocket: WebSocket def __init__(self) -> None: raise TypeError("Session instances are created by VoiceAgentServerHost") @classmethod - def _create(cls, websocket: WebSocket) -> "Session": + def _create(cls, websocket: WebSocket, *, connection_context: Any = None) -> "Session": current = cls._current(websocket) if current is not None: + # pylint: disable=protected-access + if not current._terminal and current._connection_context is None and connection_context is not None: + current._connection_context = connection_context + # pylint: enable=protected-access return current instance = object.__new__(cls) instance._websocket = websocket + instance._connection_context = connection_context instance._send_lock = asyncio.Lock() instance._terminal = False + instance._termination = None scope = getattr(websocket, "scope", None) if isinstance(scope, MutableMapping): scope[_VOICE_SESSION_SCOPE_KEY] = instance @@ -330,8 +365,16 @@ def _release(cls, websocket: WebSocket, session: "Session") -> None: if isinstance(scope, MutableMapping) and scope.get(_VOICE_SESSION_SCOPE_KEY) is session: del scope[_VOICE_SESSION_SCOPE_KEY] - def _begin_termination(self) -> None: + @property + def termination(self) -> SessionTermination | None: + """Return the first physical connection terminal fact, when committed.""" + return self._termination + + def _begin_termination(self, termination: SessionTermination | None = None) -> None: self._terminal = True + self._connection_context = None + if self._termination is None and termination is not None: + self._termination = termination def _start_close(self, code: int, reason: str) -> asyncio.Task[None] | None: self._begin_termination() @@ -368,6 +411,32 @@ def _ensure_writable(self) -> None: if self._terminal: raise RuntimeError("Voice Session is terminating") + def start_target_turn( + self, + *, + origin: TargetTurnOrigin | str, + input_count: int, + trigger_context: SpanContext | None = None, + ) -> TargetTurn: + """Start one application-owned target-agent decision trace. + + :keyword origin: Application-declared decision origin. + :paramtype origin: TargetTurnOrigin or str + :keyword input_count: Number of inputs consumed by the decision. + :paramtype input_count: int + :keyword trigger_context: Optional ended trigger context to link to this decision. + :paramtype trigger_context: opentelemetry.trace.SpanContext or None + :return: A handle the application activates around model/tool work and completes explicitly. + :rtype: TargetTurn + """ + self._ensure_writable() + return TargetTurn._create( # pylint: disable=protected-access + self._connection_context, + origin=origin, + input_count=input_count, + trigger_context=trigger_context, + ) + async def send(self, message: OutboundVoiceMessage) -> None: """Encode and send one explicit agent-to-Bridge event. @@ -389,7 +458,7 @@ async def send_frame() -> BaseException | None: await self._websocket.send_text(frame) except WebSocketDisconnect as error: if cancellation_state.cancellation is None or _find_cancellation(error) is None: - self._begin_termination() + self._begin_termination(SessionTermination.TRANSPORT_ERROR) scope = getattr(self._websocket, "scope", None) if isinstance(scope, MutableMapping): raw_reason = getattr(error, "reason", None) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_tracing.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_tracing.py new file mode 100644 index 000000000000..dd784ae193fd --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_tracing.py @@ -0,0 +1,340 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Package-local OpenTelemetry helpers for the typed Voice relay.""" + +from __future__ import annotations + +import time +from typing import Any + +from opentelemetry import ( + context as otel_context, + metrics as otel_metrics, + trace as otel_trace, +) + +from .._constants import _classify_websocket_close_code +from .._version import VERSION + + +_SCOPE_NAME = "azure.ai.agentserver.invocations.voice" +_SCHEMA_URL = "https://opentelemetry.io/schemas/gen-ai-dev/1.42.0-dev" + + +def _get_tracer() -> Any: + try: + return otel_trace.get_tracer(_SCOPE_NAME, VERSION, schema_url=_SCHEMA_URL) + except BaseException: # pylint: disable=broad-exception-caught + return None + + +def _get_meter() -> Any: + try: + return otel_metrics.get_meter(_SCOPE_NAME, VERSION, schema_url=_SCHEMA_URL) + except BaseException: # pylint: disable=broad-exception-caught + return None + + +def _create_histogram(meter: Any, name: str, boundaries: tuple[float, ...]) -> Any: + if meter is None: + return None + try: + return meter.create_histogram( + name, + unit="s", + explicit_bucket_boundaries_advisory=boundaries, + ) + except BaseException: # pylint: disable=broad-exception-caught + return None + + +def _create_counter(meter: Any, name: str) -> Any: + if meter is None: + return None + try: + return meter.create_counter(name) + except BaseException: # pylint: disable=broad-exception-caught + return None + + +_TRACER = _get_tracer() +_METER = _get_meter() +_CONNECTION_DURATION = _create_histogram( + _METER, + "azure.ai.agentserver.voice.connection.duration", + ( + 1, + 5, + 10, + 30, + 60, + 120, + 300, + 600, + 1800, + 3600, + 7200, + ), +) +_TARGET_DURATION = _create_histogram( + _METER, + "gen_ai.invoke_agent.duration", + ( + 0.1, + 0.2, + 0.4, + 0.8, + 1.6, + 3.2, + 6.4, + 12.8, + 25.6, + 51.2, + 102.4, + 204.8, + 409.6, + ), +) +_PROPAGATION_FAILURES = _create_counter( + _METER, + "azure.ai.agentserver.trace_context.propagation_failures", +) + + +def _current_context() -> Any: + try: + return otel_context.get_current() + except BaseException: # pylint: disable=broad-exception-caught + return None + + +def _attach_context(context: Any) -> tuple[Any, Any] | None: + if context is None: + return None + previous = _current_context() + if previous is None: + return None + try: + return otel_context.attach(context), previous + except BaseException: # pylint: disable=broad-exception-caught + return None + + +def _reset_context(attachment: tuple[Any, Any] | None) -> None: + if attachment is None: + return + token, previous = attachment + try: + token.var.reset(token) + return + except BaseException: # pylint: disable=broad-exception-caught + pass + try: + token.var.set(previous) + except BaseException: # pylint: disable=broad-exception-caught + pass + + +class _SpanScope: + """One returned span and its lexical context attachment.""" + + __slots__ = ( + "_attachment", + "_completed", + "_context", + "_end_time_ns", + "_span", + "_start_ns", + ) + + def __init__( + self, + span: Any = None, + attachment: tuple[Any, Any] | None = None, + *, + context: Any = None, + start_ns: int | None = None, + ) -> None: + self._span = span + self._attachment = attachment + self._completed = False + self._context = context + self._end_time_ns: int | None = None + self._start_ns = time.monotonic_ns() if start_ns is None else start_ns + + @classmethod + def start( + cls, + name: str, + *, + kind: otel_trace.SpanKind, + parent_context: Any, + attributes: dict[str, Any] | None = None, + ) -> "_SpanScope": + if parent_context is None or _TRACER is None: + return cls(start_ns=time.monotonic_ns()) + start_ns = time.monotonic_ns() + try: + span = _TRACER.start_span( + name, + context=parent_context, + kind=kind, + attributes=attributes, + ) + except BaseException: # pylint: disable=broad-exception-caught + return cls(start_ns=start_ns) + span_context = None + try: + span_context = otel_trace.set_span_in_context(span, parent_context) + attachment = _attach_context(span_context) + except BaseException: # pylint: disable=broad-exception-caught + attachment = None + if span_context is None or attachment is None: + try: + span.end() + except BaseException: # pylint: disable=broad-exception-caught + pass + return cls(start_ns=start_ns) + return cls(span, attachment, context=span_context, start_ns=start_ns) + + @property + def context(self) -> Any: + """Return the explicit child context only after lexical attachment succeeded.""" + return self._context + + @property + def is_active(self) -> bool: + """Whether span creation and lexical attachment both succeeded. + + :return: Whether this scope can parent semantic descendants. + :rtype: bool + """ + return self._span is not None and self._attachment is not None and self._context is not None + + @property + def is_completed(self) -> bool: + """Whether a terminal connection outcome already committed. + + :return: Whether the connection outcome is immutable. + :rtype: bool + """ + return self._completed + + def complete_connection(self, outcome: str, close_code: int) -> None: + if self._completed: + return + self._completed = True + self._set_attributes( + { + "azure.ai.agentserver.invocations_ws.close_code": close_code, + "bridge.outcome": outcome, + } + ) + if outcome not in {"completed", "cancelled"}: + self._set_error(outcome) + end_ns = time.monotonic_ns() + self._end_time_ns = time.time_ns() + metric_attributes = {"bridge.outcome": outcome} + if outcome not in {"completed", "cancelled"}: + metric_attributes["error.type"] = outcome + _record_metric( + _CONNECTION_DURATION, + max(0.0, (end_ns - self._start_ns) / 1_000_000_000), + metric_attributes, + ) + + def record_callback_error(self, error_type: str) -> None: + self._set_error(error_type) + + def set_attributes(self, attributes: dict[str, Any]) -> None: + """Best-effort enrichment before the span terminal boundary. + + :param attributes: Content-free span attributes to add. + :type attributes: dict[str, Any] + """ + self._set_attributes(attributes) + + def _set_attributes(self, attributes: dict[str, Any]) -> None: + if self._span is None: + return + for name, value in attributes.items(): + try: + self._span.set_attribute(name, value) + except BaseException: # pylint: disable=broad-exception-caught + pass + + def _set_error(self, error_type: str) -> None: + if self._span is None: + return + try: + self._span.set_attribute("error.type", error_type) + except BaseException: # pylint: disable=broad-exception-caught + pass + try: + self._span.set_status(otel_trace.StatusCode.ERROR) + except BaseException: # pylint: disable=broad-exception-caught + pass + + def close(self) -> None: + attachment = self._attachment + self._attachment = None + _reset_context(attachment) + span = self._span + self._span = None + self._context = None + if span is None: + return + try: + span.end(end_time=self._end_time_ns) + except BaseException: # pylint: disable=broad-exception-caught + pass + + +def _connection_outcome(close_code: int, error_code: str | None) -> str: + error_outcomes = { + "cancelled": "cancelled", + "accept_failed": "accept_error", + "internal_error": "callback_error", + } + if error_code is not None: + return error_outcomes.get(error_code, "internal_error") + return _classify_websocket_close_code(close_code) + + +def _record_target_duration( + start_ns: int, + origin: str, + outcome: str, + error_type: str | None, + *, + end_ns: int, +) -> None: + attributes = {"turn.origin": origin, "bridge.outcome": outcome} + if error_type is not None: + attributes["error.type"] = error_type + _record_metric( + _TARGET_DURATION, + max(0.0, (end_ns - start_ns) / 1_000_000_000), + attributes, + ) + + +def _record_propagation_failure(error_type: str) -> None: + try: + _PROPAGATION_FAILURES.add( + 1, + { + "azure.ai.agentserver.trace_context.propagation.hop": "hosted_agents_to_agentserver", + "error.type": error_type, + }, + ) + except BaseException: # pylint: disable=broad-exception-caught + pass + + +def _record_metric(instrument: Any, value: float, attributes: dict[str, str]) -> None: + try: + instrument.record(value, attributes) + except BaseException: # pylint: disable=broad-exception-caught + pass diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_turn.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_turn.py new file mode 100644 index 000000000000..e3e337cdaa00 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_turn.py @@ -0,0 +1,375 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Application-owned tracing handle for one target-agent decision.""" + +from __future__ import annotations + +import asyncio # pylint: disable=do-not-import-asyncio +import re +import threading +import time +import weakref +from enum import Enum +from typing import Any, ContextManager, cast + +from opentelemetry import trace as otel_trace +from opentelemetry.trace import Link, SpanContext, TraceState + +from azure.ai.agentserver.core import experimental +from azure.core import CaseInsensitiveEnumMeta + +from ._codec import _validate_prefixed_identifier_value +from ._tracing import _TRACER, _attach_context, _record_target_duration, _reset_context + + +_SPAN_CONTEXT_TYPE = cast(type[Any], SpanContext) +_SAFE_RESPONSE_ID = re.compile(r"^r_[0-9a-f]{32}$") + + +@experimental +class TargetTurnOrigin(str, Enum, metaclass=CaseInsensitiveEnumMeta): + """Application-declared origin of a target-agent decision.""" + + USER = "user" + NO_INPUT = "no_input" + PROACTIVE = "proactive" + RECOVERY = "recovery" + OTHER = "other" + + +@experimental +class TargetTurnOutcome(str, Enum, metaclass=CaseInsensitiveEnumMeta): + """Immutable terminal outcome of a target-agent decision.""" + + RESPONSE = "response" + NONE = "none" + TIMEOUT = "timeout" + ERROR = "error" + CANCELLED = "cancelled" + END_CALL = "end_call" + TRANSPORT_ERROR = "transport_error" + ABANDONED = "abandoned" + OTHER = "other" + + +def _normalize_enum(enum_type: type[Enum], value: Any, name: str) -> Any: + if isinstance(value, enum_type): + return value + if not isinstance(value, str): + raise TypeError(f"{name} must be a string or {enum_type.__name__}") + if not value.strip(): + raise ValueError(f"{name} must not be empty") + try: + return enum_type(value) + except ValueError: + return enum_type("other") + + +def _validate_count(value: Any, name: str, *, minimum: int) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if value < minimum or value > 2**63 - 1: + raise ValueError(f"{name} must be between {minimum} and {2**63 - 1}") + return value + + +def _validate_completion( + outcome: TargetTurnOutcome | str, + response_id: str | None, + output_item_count: int | None, +) -> tuple[TargetTurnOutcome, str | None, int | None]: + normalized_outcome = _normalize_enum(TargetTurnOutcome, outcome, "outcome") + validated_response_id = None + if response_id is not None: + validated_response_id = _validate_prefixed_identifier_value(response_id, "response_id", "r_") + validated_output_count = None + if output_item_count is not None: + validated_output_count = _validate_count(output_item_count, "output_item_count", minimum=0) + + if normalized_outcome is TargetTurnOutcome.RESPONSE: + if validated_response_id is None or validated_output_count is None or validated_output_count < 1: + raise ValueError("response outcome requires a response_id and at least one output item") + elif normalized_outcome is TargetTurnOutcome.NONE: + if validated_response_id is not None or validated_output_count != 0: + raise ValueError("none outcome requires no response_id and zero output items") + elif validated_output_count is not None and validated_output_count > 0 and validated_response_id is None: + raise ValueError("positive output_item_count requires a response_id") + + return normalized_outcome, validated_response_id, validated_output_count + + +def _is_safe_telemetry_response_id(value: str) -> bool: + return _SAFE_RESPONSE_ID.fullmatch(value) is not None + + +class _TargetTurnActivation: + __slots__ = ("_owner", "_entered") + + def __init__(self, owner: "TargetTurn") -> None: + self._owner = owner + self._entered = False + + def __enter__(self) -> None: + if self._entered: + raise RuntimeError("Target turn activation has already been entered") + self._owner._enter_activation() # pylint: disable=protected-access + self._entered = True + + def __exit__(self, _exc_type: Any, _exc_value: Any, _traceback: Any) -> None: + if self._entered: + self._owner._exit_activation() # pylint: disable=protected-access + self._entered = False + + +@experimental +class TargetTurn: + """Application-owned trace lifetime for one target-agent decision.""" + + __slots__ = ( + "_activated", + "_activation_attachment", + "_activation_task", + "_active", + "_affinity", + "_completed", + "_connection_context", + "_origin", + "_span", + "_start_ns", + ) + _activated: bool + _activation_attachment: ( + tuple[ + tuple[Any, Any] | None, + tuple[Any, Any] | None, + ] + | None + ) + _activation_task: asyncio.Task[Any] | None + _active: bool + _affinity: tuple[ + int, + weakref.ReferenceType[asyncio.AbstractEventLoop] | None, + ] + _completed: bool + _connection_context: Any + _origin: str + _span: Any + _start_ns: int + + def __init__(self) -> None: + raise TypeError("TargetTurn instances are created by Session.start_target_turn") + + @classmethod + def _create( + cls, + connection_context: Any, + *, + origin: TargetTurnOrigin | str, + input_count: int, + trigger_context: SpanContext | None = None, + ) -> "TargetTurn": + normalized_origin = _normalize_enum(TargetTurnOrigin, origin, "origin") + minimum = 0 if normalized_origin in {TargetTurnOrigin.PROACTIVE, TargetTurnOrigin.OTHER} else 1 + validated_input_count = _validate_count(input_count, "input_count", minimum=minimum) + if normalized_origin is TargetTurnOrigin.PROACTIVE and validated_input_count != 0: + raise ValueError("input_count must be zero for proactive target turns") + links: tuple[Link, ...] = () + if trigger_context is not None: + if not isinstance(trigger_context, _SPAN_CONTEXT_TYPE): + raise TypeError("trigger_context must be a SpanContext or None") + if not trigger_context.is_valid: + raise ValueError("trigger_context must contain a valid trace ID and span ID") + try: + safe_trigger = SpanContext( + trace_id=trigger_context.trace_id, + span_id=trigger_context.span_id, + is_remote=trigger_context.is_remote, + trace_flags=trigger_context.trace_flags, + trace_state=TraceState(), + ) + links = (Link(safe_trigger),) + except BaseException: # pylint: disable=broad-exception-caught + links = () + + instance = object.__new__(cls) + instance._activated = False + instance._activation_attachment = None + instance._activation_task = None + instance._active = False + instance._completed = False + instance._connection_context = connection_context + try: + loop_reference = weakref.ref(asyncio.get_running_loop()) + except RuntimeError: + loop_reference = None + instance._affinity = (threading.get_ident(), loop_reference) + instance._origin = normalized_origin.value + instance._span = None + instance._start_ns = time.monotonic_ns() + if connection_context is not None and _TRACER is not None: + try: + instance._span = _TRACER.start_span( + "invoke_agent", + context=connection_context, + kind=otel_trace.SpanKind.INTERNAL, + attributes={ + "gen_ai.operation.name": "invoke_agent", + "turn.origin": normalized_origin.value, + "bridge.input.count": validated_input_count, + }, + links=links, + ) + except BaseException: # pylint: disable=broad-exception-caught + instance._span = None + return instance + + @property + def is_completed(self) -> bool: + """Whether the first valid completion has committed. + + :return: Whether this target turn is complete. + :rtype: bool + """ + return self._completed + + def activate(self) -> ContextManager[None]: + """Return the single lexical activation scope for this target turn. + + The application must await every task that creates target descendants + before leaving this scope. The scope must exit in the task that entered it. + + :return: A synchronous context manager that makes this target turn current. + :rtype: contextlib.AbstractContextManager[None] + """ + return _TargetTurnActivation(self) + + def complete( + self, + *, + outcome: TargetTurnOutcome | str, + response_id: str | None = None, + output_item_count: int | None = None, + ) -> None: + """Complete the target turn with application-owned terminal facts. + + Validation depends on ``outcome``: + + * ``TargetTurnOutcome.RESPONSE`` requires ``response_id`` and an + ``output_item_count`` of at least 1. + * ``TargetTurnOutcome.NONE`` requires no ``response_id`` and an + ``output_item_count`` equal to 0. + * Other outcomes allow both values to be omitted. If + ``output_item_count`` is positive, ``response_id`` is required. + + :keyword outcome: First immutable terminal outcome. + :paramtype outcome: TargetTurnOutcome or str + :keyword response_id: Real response identifier subject to the outcome rules above. + :paramtype response_id: str or None + :keyword output_item_count: Completed output item count subject to the outcome rules above. + :paramtype output_item_count: int or None + """ + self._check_affinity() + if self._completed: + return + if self._active: + raise RuntimeError("Target turn cannot complete while its activation is active") + + normalized_outcome, validated_response_id, validated_output_count = _validate_completion( + outcome, response_id, output_item_count + ) + + self._completed = True + span = self._span + error_type = ( + normalized_outcome.value + if normalized_outcome + in { + TargetTurnOutcome.TIMEOUT, + TargetTurnOutcome.ERROR, + TargetTurnOutcome.TRANSPORT_ERROR, + TargetTurnOutcome.ABANDONED, + } + else None + ) + end_monotonic_ns = time.monotonic_ns() + end_time_ns = time.time_ns() + try: + if span is not None: + attributes: dict[str, Any] = {"bridge.outcome": normalized_outcome.value} + if validated_output_count is not None: + attributes["bridge.output.item_count"] = validated_output_count + if validated_response_id is not None and _is_safe_telemetry_response_id(validated_response_id): + attributes["gen_ai.response.id"] = validated_response_id + if error_type is not None: + attributes["error.type"] = normalized_outcome.value + for name, value in attributes.items(): + try: + span.set_attribute(name, value) + except BaseException: # pylint: disable=broad-exception-caught + pass + if error_type is not None: + try: + span.set_status(otel_trace.StatusCode.ERROR) + except BaseException: # pylint: disable=broad-exception-caught + pass + try: + span.end(end_time=end_time_ns) + except BaseException: # pylint: disable=broad-exception-caught + pass + finally: + _record_target_duration( + self._start_ns, + self._origin, + normalized_outcome.value, + error_type, + end_ns=end_monotonic_ns, + ) + self._span = None + self._connection_context = None + self._activation_task = None + self._activation_attachment = None + + def _check_affinity(self) -> None: + thread_id, loop_reference = self._affinity + if threading.get_ident() != thread_id: + raise RuntimeError("Target turn must be used on its creation thread") + if loop_reference is not None: + try: + current_loop = asyncio.get_running_loop() + except RuntimeError as exc: + raise RuntimeError("Target turn must be used on its creation event loop") from exc + if current_loop is not loop_reference(): + raise RuntimeError("Target turn must be used on its creation event loop") + + def _enter_activation(self) -> None: + self._check_affinity() + if self._completed: + raise RuntimeError("Target turn has already completed") + if self._activated: + raise RuntimeError("Target turn has already been activated") + self._activated = True + self._active = True + self._activation_task = asyncio.current_task() if self._affinity[1] is not None else None + connection_attachment = _attach_context(self._connection_context) + target_attachment = None + if connection_attachment is not None and self._span is not None: + try: + target_context = otel_trace.set_span_in_context(self._span, self._connection_context) + except BaseException: # pylint: disable=broad-exception-caught + target_context = None + if target_context is not None: + target_attachment = _attach_context(target_context) + self._activation_attachment = (connection_attachment, target_attachment) + + def _exit_activation(self) -> None: + if self._affinity[1] is not None and asyncio.current_task() is not self._activation_task: + raise RuntimeError("Target turn activation must exit in the task that entered it") + if self._activation_attachment is not None: + connection_attachment, target_attachment = self._activation_attachment + _reset_context(target_attachment) + _reset_context(connection_attachment) + self._activation_attachment = None + self._activation_task = None + self._active = False diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_voice_host.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_voice_host.py index da487efeff97..643d550ebf75 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_voice_host.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/voice/_voice_host.py @@ -8,19 +8,28 @@ import asyncio # pylint: disable=do-not-import-asyncio import inspect import logging +import os +import re import time import uuid from collections.abc import Awaitable, Callable, MutableMapping from typing import Any, NoReturn, TypeVar, cast +from urllib.parse import unquote_to_bytes -from opentelemetry import baggage as _otel_baggage, context as _otel_context +from opentelemetry import ( + baggage as _otel_baggage, + context as _otel_context, + trace as _otel_trace, +) from opentelemetry.propagators.textmap import Getter +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator from starlette.routing import Host, Match, Mount, Router, WebSocketRoute from starlette.types import Receive, Scope, Send from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState from azure.ai.agentserver.core import ( FoundryAgentRequestContext, + configure_observability as _CORE_CONFIGURE_OBSERVABILITY, experimental, set_request_context, ) @@ -31,7 +40,7 @@ USER_ID, ) -from .._constants import InvocationsWSConstants +from .._constants import InvocationsWSConstants, _classify_websocket_close_code from .._invocation import InvocationAgentServerHost from . import _session as _session_transport from ._codec import MAX_FRAME_BYTES, VoiceProtocolError, decode_inbound_message @@ -49,7 +58,14 @@ UserNoInput, UserSpeechStarted, ) -from ._session import Session +from ._session import Session, SessionTermination +from ._tracing import ( + _SpanScope, + _attach_context, + _connection_outcome, + _record_propagation_failure, + _reset_context, +) SessionStartCallback = Callable[[Session, SessionStart], Awaitable[None]] UserMessageCallback = Callable[[Session, UserMessage], Awaitable[None]] @@ -67,12 +83,92 @@ _CallbackT = TypeVar("_CallbackT", bound=Callable[..., Awaitable[None]]) _AwaitedT = TypeVar("_AwaitedT") _VoiceCallback = Callable[[Session, Any], Awaitable[None]] -logger = logging.getLogger("azure.ai.agentserver") + + +class _NoOpVoiceLogger: + @staticmethod + def debug(*_args: Any, **_kwargs: Any) -> None: + return None + + @staticmethod + def info(*_args: Any, **_kwargs: Any) -> None: + return None + + @staticmethod + def error(*_args: Any, **_kwargs: Any) -> None: + return None + + +def _get_voice_logger() -> Any: + try: + return logging.getLogger("azure.ai.agentserver") + except BaseException: # pylint: disable=broad-exception-caught + return _NoOpVoiceLogger() + + +def _get_voice_trace_propagator() -> Any: + try: + return TraceContextTextMapPropagator() + except BaseException: # pylint: disable=broad-exception-caught + return None + + +def _new_voice_context() -> Any: + try: + return _otel_context.Context() + except BaseException: # pylint: disable=broad-exception-caught + return None + + +logger = _get_voice_logger() +# Voice host and Session intentionally cooperate through package-private transport hooks. +# pylint: disable=protected-access _VOICE_AUTHORITY_ROUTE = object() _VOICE_CLOSE_CODE = _session_transport._VOICE_CLOSE_CODE_SCOPE_KEY # pylint: disable=protected-access _VOICE_DISCONNECT_EVENT = _session_transport._VOICE_DISCONNECT_EVENT_SCOPE_KEY # pylint: disable=protected-access _VOICE_TERMINATION_DEADLINE = "azure.ai.agentserver.invocations.voice.termination_deadline" +_VOICE_CONNECTION_TRACE = "azure.ai.agentserver.invocations.voice.connection_trace" +_VOICE_CONNECTION_CONTEXT = "azure.ai.agentserver.invocations.voice.connection_context" _VOICE_ROUTE_CONFLICT = "VoiceAgentServerHost cannot own /invocations_ws because the route is already registered" +_VOICE_TRACE_PROPAGATOR = _get_voice_trace_propagator() +_VOICE_BAGGAGE_KEYS = frozenset( + { + "azure.ai.agentserver.session_id", + "microsoft.a365.agent.blueprint.id", + "user.id", + "gen_ai.agent.id", + "microsoft.tenant.id", + } +) +_VOICE_SESSION_ID = re.compile(r"^[A-Za-z0-9_-]{8,128}$") +_VOICE_OPAQUE_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:@/-]*$") +_VOICE_PROTOCOL_VERSION = re.compile(r"^[0-9]{1,3}(?:\.[0-9]{1,3}){1,2}$") +_TRACESTATE_KEY = r"[a-z][_0-9a-z\-*\/]{0,255}|" + r"[a-z0-9][_0-9a-z\-*\/]{0,240}@[a-z][_0-9a-z\-*\/]{0,13}" +_TRACESTATE_VALUE = r"[\x20-\x2b\x2d-\x3c\x3e-\x7e]{0,255}[\x21-\x2b\x2d-\x3c\x3e-\x7e]" +_TRACESTATE_MEMBER = re.compile(rf"({_TRACESTATE_KEY})=({_TRACESTATE_VALUE})[ \t]*") +_VALID_PERCENT_ESCAPE = re.compile(r"%(?:[0-9A-Fa-f]{2})") + + +def _configure_voice_observability( + *, + connection_string: str | None = None, + log_level: str | None = None, + enable_sensitive_data: bool = False, +) -> None: + # AgentServerHost derives this callback argument from the same environment + # variable, defaulting it to True when unset. Voice requires an explicit + # environment opt-in for sensitive Agent Framework instrumentation, so + # ignore the inherited value and resolve the variable with default False. + del enable_sensitive_data + configured = os.environ.get("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "false") + try: + _CORE_CONFIGURE_OBSERVABILITY( + connection_string=connection_string, + log_level=log_level, + enable_sensitive_data=configured.strip().lower() in {"1", "true"}, + ) + except BaseException: # pylint: disable=broad-exception-caught + pass def _is_async_callable(callback: Callable[..., Any]) -> bool: @@ -135,19 +231,132 @@ def keys(self, carrier: list[tuple[bytes, bytes]]) -> list[str]: _VOICE_HEADER_GETTER = _VoiceHeaderGetter() +def _is_valid_voice_correlation_id(key: str, value: str) -> bool: + try: + encoded = value.encode("ascii") + except UnicodeEncodeError: + return False + if key == "azure.ai.agentserver.session_id": + return _VOICE_SESSION_ID.fullmatch(value) is not None + return len(encoded) <= 256 and _VOICE_OPAQUE_ID.fullmatch(value) is not None + + +def _decode_voice_baggage_value(value: str) -> str | None: + without_escapes = _VALID_PERCENT_ESCAPE.sub("", value) + if "%" in without_escapes: + return None + try: + return unquote_to_bytes(value).decode("utf-8", errors="strict") + except (UnicodeDecodeError, ValueError): + return None + + +def _copy_voice_baggage(raw_headers: list[tuple[bytes, bytes]], context: Any) -> tuple[Any, bool]: + raw_values = [value.decode("latin-1") for name, value in raw_headers if name.lower() == b"baggage"] + if not raw_values: + return context, False + header = ",".join(raw_values) + if len(header.encode("latin-1")) > 8192: + return context, True + members = header.split(",") + if len(members) > 180: + return context, True + approved_members: dict[str, list[str]] = {} + invalid = False + for member in members: + candidate = member.strip() + if "=" not in candidate: + if candidate in _VOICE_BAGGAGE_KEYS: + invalid = True + continue + key, encoded_value = candidate.split("=", 1) + if key not in _VOICE_BAGGAGE_KEYS: + continue + approved_members.setdefault(key, []).append(encoded_value) + + for key, encoded_values in approved_members.items(): + if len(encoded_values) != 1: + invalid = True + continue + encoded_value = encoded_values[0] + if ";" in encoded_value: + invalid = True + continue + value = _decode_voice_baggage_value(encoded_value) + if value is not None and _is_valid_voice_correlation_id(key, value): + context = _otel_baggage.set_baggage(key, value, context=context) + else: + invalid = True + return context, invalid + + +def _has_valid_tracestate(raw_headers: list[tuple[bytes, bytes]]) -> bool: + values = [value.decode("latin-1") for name, value in raw_headers if name.lower() == b"tracestate"] + if not values: + return True + header = ",".join(values) + if len(header.encode("latin-1")) > 512: + return False + members = re.split(r"[ \t]*,[ \t]*", header) + if len(members) > 32: + return False + keys: set[str] = set() + for member in members: + if not member: + return False + match = _TRACESTATE_MEMBER.fullmatch(member) + if match is None or match.group(1) in keys: + return False + keys.add(match.group(1)) + return True + + def _extract_voice_websocket_context(websocket: WebSocket) -> Any: + failure: str | None = None try: + base_context = _new_voice_context() + if base_context is None or _VOICE_TRACE_PROPAGATOR is None: + _record_propagation_failure("extraction_error") + return None raw_headers: list[tuple[bytes, bytes]] = websocket.scope.get("headers", []) - from opentelemetry.propagate import extract # pylint: disable=import-outside-toplevel - - context = extract(carrier=raw_headers, getter=_VOICE_HEADER_GETTER) + traceparents = [value for name, value in raw_headers if name.lower() == b"traceparent"] + if not traceparents: + failure = "missing" + elif len(traceparents) != 1: + failure = "invalid" + trace_headers = raw_headers + valid_tracestate = _has_valid_tracestate(raw_headers) + if not valid_tracestate: + failure = "invalid" + trace_headers = [(name, value) for name, value in raw_headers if name.lower() != b"tracestate"] + context = _VOICE_TRACE_PROPAGATOR.extract( + carrier=trace_headers, + context=base_context, + getter=_VOICE_HEADER_GETTER, + ) + if context is None: + _record_propagation_failure("extraction_error") + return None + if traceparents and not _otel_trace.get_current_span(context).get_span_context().is_valid: + failure = "invalid" + context, invalid_baggage = _copy_voice_baggage(raw_headers, context) + if invalid_baggage: + failure = "invalid" request_ids = _VOICE_HEADER_GETTER.get(raw_headers, REQUEST_ID) or [] - request_id = next((value for value in request_ids if value), None) + request_id = request_ids[0] if len(request_ids) == 1 else None + if request_ids and ( + request_id is None or not request_id or not _is_valid_voice_correlation_id("x_request_id", request_id) + ): + failure = "invalid" + request_id = None if request_id: context = _otel_baggage.set_baggage("x_request_id", request_id, context=context) + if failure is not None: + _record_propagation_failure(failure) return context except BaseException: # pylint: disable=broad-exception-caught - return _otel_context.Context() + _record_propagation_failure("extraction_error") + return None def _selected_voice_close_code(websocket: WebSocket, default_code: int) -> int: @@ -166,6 +375,9 @@ def _select_voice_close_code(websocket: WebSocket, code: int) -> None: def _raise_voice_disconnect(websocket: WebSocket, code: int, reason: str) -> NoReturn: _select_voice_close_code(websocket, code) + session = Session._current(websocket) # pylint: disable=protected-access + if session is not None: + session._begin_termination(SessionTermination.PROTOCOL_ERROR) # pylint: disable=protected-access raise WebSocketDisconnect(code=code, reason=reason) @@ -203,6 +415,29 @@ def _selected_voice_termination_deadline(websocket: WebSocket) -> float: return asyncio.get_running_loop().time() + _session_transport.CLOSE_TIMEOUT_SECONDS +def _commit_voice_session_termination( + session: Session, + *, + handler_error: BaseException | None, + disconnect_event: SessionDisconnected | None, + close_code: int, + accept_failed: bool, +) -> None: + if session.termination is not None: + return + if accept_failed: + termination = SessionTermination.ACCEPT_ERROR + elif isinstance(handler_error, asyncio.CancelledError): + termination = SessionTermination.CANCELLED + elif handler_error is not None: + termination = SessionTermination.CALLBACK_ERROR + elif disconnect_event is not None: + termination = SessionTermination(_classify_websocket_close_code(int(disconnect_event.code))) + else: + termination = SessionTermination(_classify_websocket_close_code(close_code)) + session._begin_termination(termination) + + def _peek_voice_disconnect_event(websocket: WebSocket) -> SessionDisconnected | None: scope = getattr(websocket, "scope", None) if not isinstance(scope, MutableMapping): @@ -223,7 +458,9 @@ async def _raise_pending_cancellation() -> None: await asyncio.sleep(0) -async def _raise_pending_or_consumed_cancellation(cancellation_requests: int | None) -> None: +async def _raise_pending_or_consumed_cancellation( + cancellation_requests: int | None, +) -> None: await _raise_pending_cancellation() current_requests = _task_cancellation_requests() if cancellation_requests is not None and current_requests is not None and current_requests > cancellation_requests: @@ -243,12 +480,59 @@ async def _await_with_cancellation_guard( return result -async def _receive_voice_transport_message(websocket: WebSocket) -> MutableMapping[str, Any]: +async def _receive_voice_transport_message( + websocket: WebSocket, +) -> MutableMapping[str, Any]: message = await _session_transport._run_transport_operation(websocket.receive()) # pylint: disable=protected-access await _raise_pending_cancellation() return message +async def _receive_voice_event( + websocket: WebSocket, + session: Session, +) -> InboundVoiceMessage | None: + try: + raw_message = await _receive_voice_transport_message(websocket) + except OSError as error: + code = 1006 + _select_voice_close_code(websocket, code) + session._begin_termination(SessionTermination.TRANSPORT_ERROR) # pylint: disable=protected-access + _begin_voice_termination(websocket, session) + websocket.scope.setdefault( + _VOICE_DISCONNECT_EVENT, + SessionDisconnected(code=code), + ) + raise WebSocketDisconnect(code=code) from error + raw_type = raw_message.get("type") + if raw_type == "websocket.disconnect": + code = int(raw_message.get("code") or 1000) + raw_reason = raw_message.get("reason") + reason = raw_reason if isinstance(raw_reason, str) else None + _select_voice_close_code(websocket, code) + session._begin_termination( # pylint: disable=protected-access + SessionTermination(_classify_websocket_close_code(code)) + ) + _begin_voice_termination(websocket, session) + websocket.scope.setdefault( + _VOICE_DISCONNECT_EVENT, + SessionDisconnected(code=code, reason=reason), + ) + raise WebSocketDisconnect(code=code, reason=reason) + if raw_type != "websocket.receive": + _raise_voice_disconnect(websocket, 1002, "Invalid Voice WebSocket event") + frame = raw_message.get("text") + if frame is None: + _raise_voice_disconnect(websocket, 1003, "Voice messages must be text frames") + try: + return decode_inbound_message(frame) + except VoiceProtocolError as exc: + try: + _raise_voice_disconnect(websocket, exc.close_code, "Invalid Voice message") + except WebSocketDisconnect as disconnect: + raise disconnect from exc + + @experimental class VoiceAgentServerHost(InvocationAgentServerHost): """Invocations host with typed Voice event decorators. @@ -257,6 +541,12 @@ class VoiceAgentServerHost(InvocationAgentServerHost): Agent code owns IDs, application tasks, response lifecycle, terminal-event correlation, cancellation, history, and reconnect restoration. + Default Voice observability enables sensitive Agent Framework + instrumentation only when + ``OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT`` is explicitly set to + ``true`` or ``1``. Pass a custom ``configure_observability`` callable in + ``kwargs`` to use a different policy. + :param openapi_spec: Optional OpenAPI document inherited from Invocations. :param asyncapi_spec_json: Optional AsyncAPI JSON document. :param asyncapi_spec_yaml: Optional AsyncAPI YAML document. @@ -274,6 +564,8 @@ def __init__( self._voice_callbacks: dict[str, _VoiceCallback] = {} self._connection_terminating_callback: ConnectionTerminatingCallback | None = None self._voice_route: _VoiceWebSocketRoute | None = None + if "configure_observability" not in kwargs: + kwargs["configure_observability"] = _configure_voice_observability super().__init__( openapi_spec=openapi_spec, asyncapi_spec_json=asyncapi_spec_json, @@ -308,12 +600,25 @@ def _ensure_ws_route_registered(self) -> None: async def _ws_endpoint(self, websocket: WebSocket) -> None: session_id = self.config.session_id or str(uuid.uuid4()) start_ns = time.monotonic_ns() - trace_token = None + parent_attachment = None + connection_trace = _SpanScope() + scope = getattr(websocket, "scope", None) try: - try: - trace_token = _otel_context.attach(_extract_voice_websocket_context(websocket)) - except BaseException: # pylint: disable=broad-exception-caught - trace_token = None + extracted_context = _extract_voice_websocket_context(websocket) + parent_attachment = _attach_context(extracted_context) + if parent_attachment is not None: + connection_attributes: dict[str, Any] = {"network.protocol.name": "websocket"} + if _is_valid_voice_correlation_id("azure.ai.agentserver.session_id", session_id): + connection_attributes[InvocationsWSConstants.ATTR_SPAN_SESSION_ID] = session_id + connection_trace = _SpanScope.start( + "agentserver.connection", + kind=_otel_trace.SpanKind.SERVER, + parent_context=extracted_context, + attributes=connection_attributes, + ) + if isinstance(scope, MutableMapping): + scope[_VOICE_CONNECTION_TRACE] = connection_trace + scope[_VOICE_CONNECTION_CONTEXT] = connection_trace.context if connection_trace.is_active else None platform_token = set_request_context( FoundryAgentRequestContext( call_id=websocket.headers.get(FOUNDRY_CALL_ID) or None, @@ -322,14 +627,25 @@ async def _ws_endpoint(self, websocket: WebSocket) -> None: ) ) try: - await self._run_voice_endpoint(websocket, session_id, start_ns) + await self._run_voice_endpoint(websocket, session_id, start_ns, connection_trace) finally: platform_token.var.reset(platform_token) finally: - if trace_token is not None: - trace_token.var.reset(trace_token) - - async def _run_voice_endpoint(self, websocket: WebSocket, session_id: str, start_ns: int) -> None: + if isinstance(scope, MutableMapping): + scope.pop(_VOICE_CONNECTION_TRACE, None) + scope.pop(_VOICE_CONNECTION_CONTEXT, None) + if not connection_trace.is_completed: + connection_trace.complete_connection("internal_error", InvocationsWSConstants.CLOSE_INTERNAL_ERROR) + connection_trace.close() + _reset_context(parent_attachment) + + async def _run_voice_endpoint( + self, + websocket: WebSocket, + session_id: str, + start_ns: int, + connection_trace: _SpanScope, + ) -> None: try: accept_error, voice_session, close_code, handler_exc, pending_error = ( await self._run_voice_connection_context( @@ -343,6 +659,8 @@ async def _run_voice_endpoint(self, websocket: WebSocket, session_id: str, start start_ns=start_ns, close_code=InvocationsWSConstants.CLOSE_INTERNAL_ERROR, error_code="cancelled", + connection_trace=connection_trace, + outcome="cancelled", ) raise @@ -357,12 +675,13 @@ async def _run_voice_endpoint(self, websocket: WebSocket, session_id: str, start handler_exc=None, pending_error=None, error_code_override="accept_failed", + connection_trace=connection_trace, ) self._report_voice_accept_failure( session_id, start_ns, - accept_error, emit_event=voice_session is None, + connection_trace=connection_trace, ) return @@ -376,13 +695,20 @@ async def _run_voice_endpoint(self, websocket: WebSocket, session_id: str, start close_code=close_code, handler_exc=handler_exc, pending_error=pending_error, + connection_trace=connection_trace, ) async def _run_voice_connection_context( self, websocket: WebSocket, session_id: str, - ) -> tuple[Exception | None, Session | None, int, BaseException | None, BaseException | None]: + ) -> tuple[ + Exception | None, + Session | None, + int, + BaseException | None, + BaseException | None, + ]: accept_error: Exception | None = None voice_session: Session | None = None close_code = InvocationsWSConstants.CLOSE_NORMAL @@ -396,7 +722,11 @@ async def _run_voice_connection_context( session_id, ) elif websocket.application_state == WebSocketState.CONNECTED: - voice_session = Session._create(websocket) # pylint: disable=protected-access + voice_session = Session._create( # pylint: disable=protected-access + websocket, + connection_context=websocket.scope.get(_VOICE_CONNECTION_CONTEXT), + ) + voice_session._begin_termination(SessionTermination.ACCEPT_ERROR) # pylint: disable=protected-access _begin_voice_termination(websocket, voice_session) close_code = InvocationsWSConstants.CLOSE_INTERNAL_ERROR finally: @@ -429,7 +759,10 @@ async def _run_accepted_voice_handler( websocket: WebSocket, session_id: str, ) -> tuple[Session, int, BaseException | None, BaseException | None]: - voice_session = Session._create(websocket) # pylint: disable=protected-access + voice_session = Session._create( # pylint: disable=protected-access + websocket, + connection_context=websocket.scope.get(_VOICE_CONNECTION_CONTEXT), + ) close_code = InvocationsWSConstants.CLOSE_NORMAL handler_exc: BaseException | None = None pending_error: BaseException | None = None @@ -459,6 +792,7 @@ async def _complete_voice_endpoint( close_code: int, handler_exc: BaseException | None, pending_error: BaseException | None, + connection_trace: _SpanScope, error_code_override: str | None = None, ) -> None: deadline = _selected_voice_termination_deadline(websocket) @@ -474,6 +808,13 @@ async def _complete_voice_endpoint( error_code = "cancelled" else: error_code = "internal_error" + _commit_voice_session_termination( + voice_session, + handler_error=handler_exc, + disconnect_event=disconnect_event, + close_code=close_code, + accept_failed=error_code_override == "accept_failed", + ) if cancellation is None and disconnect_event is None and close_code not in {1005, 1006, 1015}: reason = "Internal server error" if close_code == InvocationsWSConstants.CLOSE_INTERNAL_ERROR else "" try: @@ -492,6 +833,7 @@ async def _complete_voice_endpoint( close_error = exc disconnect_event = _take_voice_disconnect_event(websocket) disconnect_error = await self._notify_peer_disconnect( + websocket, voice_session, disconnect_event, ) @@ -507,12 +849,13 @@ async def _complete_voice_endpoint( error_code=( "internal_error" if termination_error is not None or disconnect_error is not None else error_code ), + connection_trace=connection_trace, + outcome=(voice_session.termination.value if voice_session.termination is not None else None), ) if pending_error is not None: raise pending_error self._report_voice_endpoint_errors( - session_id=session_id, handler_error=handler_exc if isinstance(handler_exc, Exception) else None, termination_error=termination_error, disconnect_error=disconnect_error, @@ -522,7 +865,6 @@ async def _complete_voice_endpoint( @staticmethod def _report_voice_endpoint_errors( *, - session_id: str, handler_error: Exception | None, termination_error: BaseException | None, disconnect_error: BaseException | None, @@ -530,22 +872,22 @@ def _report_voice_endpoint_errors( ) -> None: if handler_error is not None: try: - logger.error("Voice WebSocket handler raised for session %s", session_id, exc_info=handler_error) + logger.error("Voice WebSocket handler failed") except BaseException: # pylint: disable=broad-exception-caught pass if termination_error is not None: try: - logger.error("Voice connection termination callback failed", exc_info=termination_error) + logger.error("Voice connection termination callback failed") except BaseException: # pylint: disable=broad-exception-caught pass if disconnect_error is not None: try: - logger.error("Voice disconnect callback failed", exc_info=disconnect_error) + logger.error("Voice disconnect callback failed") except BaseException: # pylint: disable=broad-exception-caught pass if close_error is not None: try: - logger.debug("Error closing Voice WebSocket session %s", session_id, exc_info=close_error) + logger.debug("Voice WebSocket close failed") except BaseException: # pylint: disable=broad-exception-caught pass @@ -553,9 +895,9 @@ def _report_voice_accept_failure( self, session_id: str, start_ns: int, - accept_error: Exception, *, emit_event: bool, + connection_trace: _SpanScope, ) -> None: if emit_event: self._emit_voice_close_event( @@ -563,12 +905,53 @@ def _report_voice_accept_failure( start_ns=start_ns, close_code=InvocationsWSConstants.CLOSE_INTERNAL_ERROR, error_code="accept_failed", + connection_trace=connection_trace, + outcome="accept_error", ) try: - logger.error("Voice WebSocket accept failed for session %s", session_id, exc_info=accept_error) + logger.error("Voice WebSocket accept failed") except BaseException: # pylint: disable=broad-exception-caught pass + async def _invoke_traced_voice_callback( + self, + websocket: WebSocket, + session: Session, + event: InboundVoiceMessage, + callback: _VoiceCallback, + ) -> None: + callback_trace = self._start_voice_callback_trace(websocket, event.type) + try: + await _await_with_cancellation_guard( + callback(session, event), + on_success=( + (lambda: _begin_voice_termination(websocket, session)) if isinstance(event, SessionEnd) else None + ), + ) + except asyncio.CancelledError: + callback_trace.record_callback_error("cancelled") + raise + except BaseException: # pylint: disable=broad-exception-caught + error_type = ( + "transport_error" if session.termination is SessionTermination.TRANSPORT_ERROR else "callback_error" + ) + callback_trace.record_callback_error(error_type) + raise + finally: + callback_trace.close() + + @staticmethod + def _start_voice_callback_trace(websocket: WebSocket, event_type: str) -> _SpanScope: + connection_context = websocket.scope.get(_VOICE_CONNECTION_CONTEXT) + if connection_context is None: + return _SpanScope() + return _SpanScope.start( + "voice.callback", + kind=_otel_trace.SpanKind.INTERNAL, + parent_context=connection_context, + attributes={"voice.event.type": event_type}, + ) + async def _invoke_user_handler( self, websocket: WebSocket, @@ -583,7 +966,10 @@ async def _invoke_user_handler( return InvocationsWSConstants.CLOSE_NORMAL, None except WebSocketDisconnect as exc: _raise_wrapped_cancellation(exc, cancellation_requests) - return int(exc.code) if exc.code else InvocationsWSConstants.CLOSE_NORMAL, None + session = Session._current(websocket) # pylint: disable=protected-access + if session is None or (session.termination is None and _peek_voice_disconnect_event(websocket) is None): + return InvocationsWSConstants.CLOSE_INTERNAL_ERROR, exc + return (int(exc.code) if exc.code else InvocationsWSConstants.CLOSE_NORMAL), None except Exception as exc: # pylint: disable=broad-exception-caught _raise_wrapped_cancellation(exc, cancellation_requests) return InvocationsWSConstants.CLOSE_INTERNAL_ERROR, exc @@ -595,13 +981,37 @@ def _emit_voice_close_event( start_ns: int, close_code: int, error_code: str | None, + connection_trace: _SpanScope, + outcome: str | None = None, ) -> None: duration_ms = (time.monotonic_ns() - start_ns) // 1_000_000 + connection_trace.complete_connection( + outcome or _connection_outcome(close_code, error_code), + close_code, + ) try: self._emit_close_event(session_id, close_code, duration_ms, error_code=error_code) except BaseException: # pylint: disable=broad-exception-caught pass + @staticmethod + def _emit_close_event( + session_id: str, + close_code: int, + duration_ms: int, + *, + error_code: str | None = None, + ) -> None: + extra: dict[str, Any] = { + InvocationsWSConstants.ATTR_SPAN_CLOSE_CODE: close_code, + InvocationsWSConstants.ATTR_SPAN_DURATION_MS: duration_ms, + } + if _is_valid_voice_correlation_id("azure.ai.agentserver.session_id", session_id): + extra[InvocationsWSConstants.ATTR_SPAN_SESSION_ID] = session_id + if error_code is not None: + extra[InvocationsWSConstants.ATTR_SPAN_ERROR_CODE] = error_code + logger.info("Voice connection closed", extra=extra) + def ws_handler(self, fn: Any) -> NoReturn: """Reject raw-handler registration on the typed Voice host. @@ -779,81 +1189,73 @@ def _notify_connection_terminating(self, session: Session) -> BaseException | No async def _notify_peer_disconnect( self, + websocket: WebSocket, session: Session, event: SessionDisconnected | None, ) -> BaseException | None: callback = self._voice_callbacks.get("disconnect") if callback is None or event is None: return None + callback_trace = self._start_voice_callback_trace(websocket, "disconnect") cancellation_requests = _task_cancellation_requests() try: - await callback(session, event) - except asyncio.CancelledError as exc: - current_requests = _task_cancellation_requests() - if ( - cancellation_requests is not None - and current_requests is not None - and current_requests > cancellation_requests - ): + try: + await callback(session, event) + except asyncio.CancelledError as exc: + callback_trace.record_callback_error("cancelled") + current_requests = _task_cancellation_requests() + if ( + cancellation_requests is not None + and current_requests is not None + and current_requests > cancellation_requests + ): + raise + return exc + except Exception as exc: # pylint: disable=broad-exception-caught + callback_trace.record_callback_error("callback_error") + _raise_wrapped_cancellation(exc, cancellation_requests) + return exc + except BaseException as exc: # pylint: disable=broad-exception-caught + callback_trace.record_callback_error("callback_error") + _raise_wrapped_cancellation(exc, cancellation_requests) + return exc + try: + await _raise_pending_or_consumed_cancellation(cancellation_requests) + except asyncio.CancelledError: + callback_trace.record_callback_error("cancelled") raise - return exc - except Exception as exc: # pylint: disable=broad-exception-caught - _raise_wrapped_cancellation(exc, cancellation_requests) - return exc - except BaseException as exc: # pylint: disable=broad-exception-caught - _raise_wrapped_cancellation(exc, cancellation_requests) - return exc - await _raise_pending_or_consumed_cancellation(cancellation_requests) - return None + return None + finally: + callback_trace.close() async def _handle_voice_connection(self, websocket: WebSocket) -> None: bound_session = Session._current(websocket) # pylint: disable=protected-access session = bound_session or Session._create(websocket) # pylint: disable=protected-access try: while True: - raw_message = await _receive_voice_transport_message(websocket) - raw_type = raw_message.get("type") - if raw_type == "websocket.disconnect": - code = int(raw_message.get("code") or 1000) - raw_reason = raw_message.get("reason") - reason = raw_reason if isinstance(raw_reason, str) else None - _select_voice_close_code(websocket, code) - _begin_voice_termination(websocket, session) - websocket.scope.setdefault( - _VOICE_DISCONNECT_EVENT, - SessionDisconnected(code=code, reason=reason), - ) - raise WebSocketDisconnect( - code=code, - reason=reason, - ) - if raw_type != "websocket.receive": - reason = "Invalid Voice WebSocket event" - _raise_voice_disconnect(websocket, 1002, reason) - frame = raw_message.get("text") - if frame is None: - reason = "Voice messages must be text frames" - _raise_voice_disconnect(websocket, 1003, reason) - try: - event = decode_inbound_message(frame) - except VoiceProtocolError as exc: - reason = "Invalid Voice message" - try: - _raise_voice_disconnect(websocket, exc.close_code, reason) - except WebSocketDisconnect as disconnect: - raise disconnect from exc + event = await _receive_voice_event(websocket, session) if event is None: continue + if isinstance(event, SessionStart): + scope = getattr(websocket, "scope", None) + connection_trace = scope.get(_VOICE_CONNECTION_TRACE) if isinstance(scope, MutableMapping) else None + if isinstance(connection_trace, _SpanScope): + attributes: dict[str, Any] = { + "azure.ai.agentserver.invocations_ws.reconnect": event.reconnect, + } + if _VOICE_PROTOCOL_VERSION.fullmatch(event.protocol_version) is not None: + attributes["azure.ai.agentserver.invocations_ws.protocol_version"] = event.protocol_version + connection_trace.set_attributes(attributes) callback = self._voice_callbacks.get(event.type) if callback is not None: - await _await_with_cancellation_guard( - callback(session, cast(InboundVoiceMessage, event)), - on_success=( - (lambda: _begin_voice_termination(websocket, session)) - if isinstance(event, SessionEnd) - else None - ), + await self._invoke_traced_voice_callback( + websocket, + session, + cast(InboundVoiceMessage, event), + callback, ) + if isinstance(event, SessionEnd): + session._begin_termination(SessionTermination.COMPLETED) # pylint: disable=protected-access if isinstance(event, SessionEnd): return finally: @@ -862,17 +1264,18 @@ async def _handle_voice_connection(self, websocket: WebSocket) -> None: Session._release(websocket, session) # pylint: disable=protected-access termination_error = self._notify_connection_terminating(session) disconnect_error = await self._notify_peer_disconnect( + websocket, session, _take_voice_disconnect_event(websocket), ) if termination_error is not None: try: - logger.error("Voice connection termination callback failed", exc_info=termination_error) + logger.error("Voice connection termination callback failed") except BaseException: # pylint: disable=broad-exception-caught pass if disconnect_error is not None: try: - logger.error("Voice disconnect callback failed", exc_info=disconnect_error) + logger.error("Voice disconnect callback failed") except BaseException: # pylint: disable=broad-exception-caught pass diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/dev_requirements.txt b/sdk/agentserver/azure-ai-agentserver-invocations/dev_requirements.txt index 78508f9fa127..48e702960a1c 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/dev_requirements.txt +++ b/sdk/agentserver/azure-ai-agentserver-invocations/dev_requirements.txt @@ -6,5 +6,5 @@ pytest httpx pytest-asyncio -opentelemetry-api>=1.40.0 -opentelemetry-sdk>=1.40.0 \ No newline at end of file +opentelemetry-api>=1.43.0 +opentelemetry-sdk>=1.43.0 \ No newline at end of file diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/pyproject.toml b/sdk/agentserver/azure-ai-agentserver-invocations/pyproject.toml index b581cf85f5df..72e8b3d2032b 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/pyproject.toml +++ b/sdk/agentserver/azure-ai-agentserver-invocations/pyproject.toml @@ -23,6 +23,7 @@ keywords = ["azure", "azure sdk", "agent", "agentserver", "invocations"] dependencies = [ "azure-ai-agentserver-core>=2.1.0b1", "azure-core>=1.37.0", + "opentelemetry-api>=1.43.0", # 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). @@ -36,8 +37,8 @@ dev = [ "azure-monitor-query", "azure-sdk-tools", "httpx", - "opentelemetry-api>=1.40.0", - "opentelemetry-sdk>=1.40.0", + "opentelemetry-api>=1.43.0", + "opentelemetry-sdk>=1.43.0", "pytest-asyncio", "pytest", ] diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/basic_voice_agent/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/samples/basic_voice_agent/README.md index 330e09b76eb9..41a47a44bc59 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/basic_voice_agent/README.md +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/basic_voice_agent/README.md @@ -2,7 +2,7 @@ This sample hosts a text-only agent for the Voice Live Bridge protocol `1.0`. The SDK relays typed events while the application owns response generation, -task cancellation, and correlation. +task cancellation, correlation, and target-turn trace completion. ## Prerequisites @@ -49,6 +49,19 @@ The sample intentionally uses a simulated model stream. Replace `generate_answer` with the application's model call while preserving the application-owned task cleanup shown by the event callbacks. +Each response generation declares one target turn with +`session.start_target_turn(...)`. The generation task activates that handle +around the real model/output work and completes it only after the activation +scope exits. Cancellation callbacks store an application-owned terminal hint +before cancelling the task; the runner or done callback then completes the turn +with the first truthful outcome. The SDK does not inspect or retain the task. +For unfinished work, the sample maps clean peer closure to `abandoned`, +application/server failures to `error`, protocol or transport loss to +`transport_error`, and cancellation to `cancelled`. An explicit application +hint such as `end_call` or `timeout` remains the first winner. +The sample also caps active generations per connection and bounds retained model +output by both UTF-8 bytes and chunk count; excess input receives `response.none`. + `on_connection_terminating` synchronously cancels the sample's generation tasks whenever the connection handler exits. The tasks remain responsible for their own asynchronous resource cleanup, while `on_session_end` provides the graceful diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/basic_voice_agent/basic_voice_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/basic_voice_agent/basic_voice_agent.py index 1e352a982d46..6fd569dadcc3 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/basic_voice_agent/basic_voice_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/basic_voice_agent/basic_voice_agent.py @@ -6,6 +6,7 @@ """ import asyncio +import contextvars import logging from collections.abc import AsyncIterator from dataclasses import dataclass @@ -26,6 +27,10 @@ SessionReady, SessionRejected, SessionStart, + SessionTermination, + TargetTurn, + TargetTurnOrigin, + TargetTurnOutcome, UserMessage, VoiceAgentServerHost, new_item_id, @@ -35,15 +40,24 @@ logger = logging.getLogger("azure.ai.agentserver") app = VoiceAgentServerHost() SUPPORTED_PROTOCOL_VERSION = "1.0" +MAX_ACTIVE_GENERATIONS_PER_SESSION = 8 +MAX_OUTPUT_CHUNKS = 4096 +MAX_OUTPUT_UTF8_BYTES = 512 * 1024 -@dataclass(frozen=True) +@dataclass class Generation: - """Application-owned generation task and its input correlation.""" + """Application-owned generation task, trace, and terminal facts.""" session: Session input_ids: tuple[str, ...] - task: asyncio.Task[None] + response_id: str + turn: TargetTurn + preparation_ready: asyncio.Event + task: asyncio.Task[None] | None = None + response_started: bool = False + output_item_count: int = 0 + outcome_hint: TargetTurnOutcome | None = None GenerationKey = tuple[int, str] @@ -60,33 +74,118 @@ async def generate_answer(text: str) -> AsyncIterator[str]: async def stream_response( - session: Session, + generation: Generation, *, - response_id: str, item_id: str, - in_reply_to: tuple[str, ...], text: str, ) -> None: """Translate one model stream into explicit Bridge output events.""" - await session.send(ResponseCreated(response_id=response_id, in_reply_to=in_reply_to)) + await generation.session.send(ResponseCreated(response_id=generation.response_id, in_reply_to=generation.input_ids)) + generation.response_started = True chunks: list[str] = [] + output_utf8_bytes = 0 async for delta in generate_answer(text): + delta_utf8_bytes = len(delta.encode("utf-8")) + if len(chunks) >= MAX_OUTPUT_CHUNKS or output_utf8_bytes + delta_utf8_bytes > MAX_OUTPUT_UTF8_BYTES: + raise RuntimeError("Voice model output exceeded sample limits") + output_utf8_bytes += delta_utf8_bytes chunks.append(delta) - await session.send( + await generation.session.send( ResponseOutputTextDelta( - response_id=response_id, + response_id=generation.response_id, item_id=item_id, delta=delta, ) ) - await session.send( + await generation.session.send( ResponseOutputTextDone( - response_id=response_id, + response_id=generation.response_id, item_id=item_id, text="".join(chunks), ) ) - await session.send(ResponseDone(response_id=response_id)) + generation.output_item_count = 1 + await generation.session.send(ResponseDone(response_id=generation.response_id)) + + +def completion_facts(generation: Generation) -> tuple[str | None, int]: + """Project only application-committed response facts into tracing.""" + response_id = generation.response_id if generation.response_started else None + return response_id, generation.output_item_count + + +def complete_generation(generation: Generation, outcome: TargetTurnOutcome) -> None: + """Complete one application turn after its activation scope has exited.""" + response_id, output_item_count = completion_facts(generation) + generation.turn.complete( + outcome=outcome, + response_id=response_id, + output_item_count=output_item_count, + ) + + +def termination_outcome(termination: SessionTermination | None) -> TargetTurnOutcome: + """Map a physical connection fact to an unfinished application decision.""" + if termination is None or termination is SessionTermination.CANCELLED: + return TargetTurnOutcome.CANCELLED + if termination is SessionTermination.COMPLETED: + return TargetTurnOutcome.ABANDONED + if termination in {SessionTermination.PROTOCOL_ERROR, SessionTermination.TRANSPORT_ERROR}: + return TargetTurnOutcome.TRANSPORT_ERROR + if termination in { + SessionTermination.ACCEPT_ERROR, + SessionTermination.CALLBACK_ERROR, + SessionTermination.INTERNAL_ERROR, + }: + return TargetTurnOutcome.ERROR + raise AssertionError(f"Unhandled Voice session termination: {termination!r}") + + +def generation_error_outcome(generation: Generation) -> TargetTurnOutcome: + """Prefer already-committed application or connection facts over a generic error.""" + if generation.outcome_hint is not None: + return generation.outcome_hint + if generation.session.termination is not None: + return termination_outcome(generation.session.termination) + return TargetTurnOutcome.ERROR + + +async def send_no_response(session: Session, input_ids: tuple[str, ...], reason: str) -> None: + """Send one bounded no-response decision under an explicit target turn.""" + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=len(input_ids)) + try: + with turn.activate(): + await session.send(ResponseNone(in_reply_to=input_ids, reason=reason)) + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + except asyncio.CancelledError: + if not turn.is_completed: + turn.complete(outcome=termination_outcome(session.termination), output_item_count=0) + raise + except BaseException: + if not turn.is_completed: + outcome = ( + termination_outcome(session.termination) if session.termination is not None else TargetTurnOutcome.ERROR + ) + turn.complete(outcome=outcome, output_item_count=0) + raise + + +async def run_generation(generation: Generation, *, item_id: str, text: str) -> None: + """Run all descendant-producing work under the declared target turn.""" + await generation.preparation_ready.wait() + try: + with generation.turn.activate(): + await stream_response(generation, item_id=item_id, text=text) + complete_generation(generation, TargetTurnOutcome.RESPONSE) + except asyncio.CancelledError: + complete_generation( + generation, + generation.outcome_hint or termination_outcome(generation.session.termination), + ) + raise + except BaseException: + complete_generation(generation, generation_error_outcome(generation)) + raise def generation_finished(key: GenerationKey, completed: asyncio.Task[None]) -> None: @@ -94,36 +193,59 @@ def generation_finished(key: GenerationKey, completed: asyncio.Task[None]) -> No generation = generations.get(key) if generation is None or generation.task is not completed: return - del generations[key] - for input_id in generation.input_ids: - input_key = (id(generation.session), input_id) - if input_generations.get(input_key) == key: - del input_generations[input_key] - if not completed.cancelled() and (exception := completed.exception()) is not None: - logger.error( - "Voice response generation failed", - exc_info=(type(exception), exception, exception.__traceback__), - ) - - -def cancel_generation(session: Session, response_id: str) -> None: + try: + if not generation.turn.is_completed: + if completed.cancelled(): + outcome = generation.outcome_hint or termination_outcome(generation.session.termination) + elif completed.exception() is not None: + outcome = generation_error_outcome(generation) + else: + outcome = TargetTurnOutcome.ABANDONED + complete_generation(generation, outcome) + if not completed.cancelled() and completed.exception() is not None: + logger.error("Voice response generation failed") + finally: + if generations.get(key) is generation: + del generations[key] + for input_id in generation.input_ids: + input_key = (id(generation.session), input_id) + if input_generations.get(input_key) == key: + del input_generations[input_key] + + +def set_outcome_hint(generation: Generation, outcome: TargetTurnOutcome) -> None: + """Commit the first application-known terminal hint.""" + if generation.outcome_hint is None: + generation.outcome_hint = outcome + + +def cancel_generation(session: Session, response_id: str, outcome: TargetTurnOutcome) -> None: """Cancel one application-owned generation task when present.""" generation = generations.get((id(session), response_id)) - if generation is not None: + if generation is not None and generation.task is not None: + set_outcome_hint(generation, outcome) generation.task.cancel() -def cancel_session_generation_tasks(session: Session) -> tuple[asyncio.Task[None], ...]: +def cancel_session_generation_tasks( + session: Session, + outcome: TargetTurnOutcome | None = None, +) -> tuple[asyncio.Task[None], ...]: """Synchronously signal every application-owned task for one connection.""" - tasks = tuple(generation.task for generation in tuple(generations.values()) if generation.session is session) - for task in tasks: - task.cancel() - return tasks + selected = tuple(generation for generation in tuple(generations.values()) if generation.session is session) + tasks = [] + for generation in selected: + if outcome is not None: + set_outcome_hint(generation, outcome) + if generation.task is not None: + tasks.append(generation.task) + generation.task.cancel() + return tuple(tasks) -async def cancel_session_generations(session: Session) -> None: +async def cancel_session_generations(session: Session, outcome: TargetTurnOutcome) -> None: """Cancel and join all application-owned tasks for one connection.""" - tasks = cancel_session_generation_tasks(session) + tasks = cancel_session_generation_tasks(session, outcome) if tasks: await asyncio.gather(*tasks, return_exceptions=True) @@ -135,7 +257,7 @@ async def on_session_start(session: Session, event: SessionStart) -> None: await session.send(SessionRejected(code="protocol_mismatch", retriable=False)) return if event.reconnect: - logger.info("Voice transport reattached; restore durable application state here") + logger.info("Voice transport reattached") await session.send(SessionReady()) @@ -143,76 +265,97 @@ async def on_session_start(session: Session, event: SessionStart) -> None: async def on_user_message(session: Session, event: UserMessage) -> None: """Start generation without blocking later full-duplex control events.""" text = " ".join(part.text for part in event.content if isinstance(part, InputTextPart)) + input_ids = (event.item_id,) if not text: - await session.send(ResponseNone(in_reply_to=(event.item_id,), reason="no_reply_needed")) + await send_no_response(session, input_ids, "no_reply_needed") + return + active_generations = sum(generation.session is session for generation in generations.values()) + if active_generations >= MAX_ACTIVE_GENERATIONS_PER_SESSION: + await send_no_response(session, input_ids, "capacity_exceeded") return response_id = new_response_id() item_id = new_item_id() - input_ids = (event.item_id,) key = (id(session), response_id) - task = asyncio.create_task( - stream_response( - session, - response_id=response_id, - item_id=item_id, - in_reply_to=input_ids, - text=text, - ), - name=f"voice-response-{response_id}", + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=len(input_ids)) + generation = Generation( + session=session, + input_ids=input_ids, + response_id=response_id, + turn=turn, + preparation_ready=asyncio.Event(), ) - generations[key] = Generation(session=session, input_ids=input_ids, task=task) - input_generations[(id(session), event.item_id)] = key - - def on_generation_finished(completed: asyncio.Task[None]) -> None: - generation_finished(key, completed) - - task.add_done_callback(on_generation_finished) + coroutine = run_generation(generation, item_id=item_id, text=text) + task = None + try: + task = asyncio.create_task(coroutine, name=f"voice-response-{response_id}") + generation.task = task + + def on_generation_finished(completed: asyncio.Task[None]) -> None: + generation_finished(key, completed) + + task.add_done_callback(on_generation_finished, context=contextvars.Context()) + generations[key] = generation + input_generations[(id(session), event.item_id)] = key + generation.preparation_ready.set() + except BaseException: + turn.complete(outcome=TargetTurnOutcome.ABANDONED, output_item_count=0) + generations.pop(key, None) + input_generations.pop((id(session), event.item_id), None) + if task is None: + coroutine.close() + else: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + raise @app.on_barge_in async def on_barge_in(session: Session, event: BargeIn) -> None: """Stop generation and reconcile history from the playback snapshot.""" - cancel_generation(session, event.response_id) - logger.info("Caller heard %d characters before barge-in", len(event.heard_text)) + cancel_generation(session, event.response_id, TargetTurnOutcome.CANCELLED) + logger.info("Voice response interrupted") @app.on_response_cancelled async def on_response_cancelled(session: Session, event: ResponseCancelled) -> None: """Handle the terminal outcome of an explicit self-cancel request.""" - cancel_generation(session, event.response_id) + cancel_generation(session, event.response_id, TargetTurnOutcome.CANCELLED) @app.on_response_timeout async def on_response_timeout(session: Session, event: ResponseTimeout) -> None: """Stop the application task targeted by the Bridge timeout.""" if event.response_id is not None: - cancel_generation(session, event.response_id) + cancel_generation(session, event.response_id, TargetTurnOutcome.TIMEOUT) return for input_id in event.item_ids or (): key = input_generations.get((id(session), input_id)) if key is not None and (generation := generations.get(key)) is not None: + set_outcome_hint(generation, TargetTurnOutcome.TIMEOUT) + assert generation.task is not None generation.task.cancel() @app.on_session_end async def on_session_end(session: Session, event: SessionEnd) -> None: """Cancel and join all application tasks for this connection.""" - logger.info("Voice session ended: %s", event.reason) - await cancel_session_generations(session) + del event + logger.info("Voice session ended") + await cancel_session_generations(session, TargetTurnOutcome.END_CALL) @app.on_disconnect async def on_disconnect(session: Session, event: SessionDisconnected) -> None: """Observe a peer transport disconnect.""" - del session - logger.info("Voice transport disconnected with close code %d", event.code) + del session, event + logger.info("Voice transport disconnected") @app.on_connection_terminating def on_connection_terminating(session: Session) -> None: """Synchronously cancel application tasks whenever the handler exits.""" - cancel_session_generation_tasks(session) + cancel_session_generation_tasks(session, termination_outcome(session.termination)) if __name__ == "__main__": diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_basic_voice_agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_basic_voice_agent.py index 296c0c507ffc..ac441e729c0c 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_basic_voice_agent.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_basic_voice_agent.py @@ -14,25 +14,66 @@ import azure.ai.agentserver.invocations.voice as voice from azure.ai.agentserver.invocations.voice import ( InputTextPart, + ResponseNone, ResponseTimeouts, SessionDisconnected, SessionEnd, SessionReady, SessionRejected, SessionStart, + SessionTermination, + TargetTurnOutcome, UserMessage, ) _SAMPLE_ROOT = Path(__file__).parents[2] / "samples" / "basic_voice_agent" +class _Activation: + def __init__(self, turn): + self.turn = turn + + def __enter__(self): + self.turn.activation_count += 1 + + def __exit__(self, _exc_type, _exc_value, _traceback): + return None + + +class _CapturingTurn: + def __init__(self, origin, input_count): + self.origin = origin + self.input_count = input_count + self.activation_count = 0 + self.completions = [] + + def activate(self): + return _Activation(self) + + def complete(self, **kwargs): + if not self.completions: + self.completions.append(kwargs) + + @property + def is_completed(self): + return bool(self.completions) + + class _CapturingSession: def __init__(self): self.messages = [] + self.turns = [] + self.termination = None async def send(self, message): self.messages.append(message) + def start_target_turn(self, *, origin, input_count, trigger_context=None): + del trigger_context + turn = _CapturingTurn(origin, input_count) + self.turns.append(turn) + return turn + @pytest.fixture def sample_module(monkeypatch): @@ -70,6 +111,70 @@ def _session_start(protocol_version): ) +@pytest.mark.parametrize( + ("termination", "expected"), + [ + pytest.param(None, TargetTurnOutcome.CANCELLED, id="local-cancellation"), + pytest.param(SessionTermination.CANCELLED, TargetTurnOutcome.CANCELLED, id="connection-cancelled"), + pytest.param(SessionTermination.COMPLETED, TargetTurnOutcome.ABANDONED, id="clean-peer-close"), + pytest.param(SessionTermination.PROTOCOL_ERROR, TargetTurnOutcome.TRANSPORT_ERROR, id="protocol-error"), + pytest.param(SessionTermination.TRANSPORT_ERROR, TargetTurnOutcome.TRANSPORT_ERROR, id="transport-error"), + pytest.param(SessionTermination.ACCEPT_ERROR, TargetTurnOutcome.ERROR, id="accept-error"), + pytest.param(SessionTermination.CALLBACK_ERROR, TargetTurnOutcome.ERROR, id="callback-error"), + pytest.param(SessionTermination.INTERNAL_ERROR, TargetTurnOutcome.ERROR, id="internal-error"), + ], +) +def test_termination_outcome_preserves_source_semantics(sample_module, termination, expected): + assert sample_module.termination_outcome(termination) is expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("termination", "expected"), + [ + pytest.param(None, TargetTurnOutcome.CANCELLED, id="local-cancellation-negative-control"), + pytest.param( + SessionTermination.TRANSPORT_ERROR, + TargetTurnOutcome.TRANSPORT_ERROR, + id="committed-transport-error", + ), + pytest.param(SessionTermination.CALLBACK_ERROR, TargetTurnOutcome.ERROR, id="committed-callback-error"), + ], +) +async def test_no_response_cancellation_preserves_committed_termination( + sample_module, + termination, + expected, +): + session = _CapturingSession() + send_attempts = [] + + async def cancel_send(message): + send_attempts.append(message) + session.termination = termination + raise asyncio.CancelledError() + + session.send = cancel_send + + with pytest.raises(asyncio.CancelledError): + await sample_module.send_no_response(session, ("in_1",), "no_reply_needed") + + assert len(send_attempts) == 1 + assert len(session.turns) == 1 + turn = session.turns[0] + assert turn.completions == [ + { + "outcome": expected, + "output_item_count": 0, + } + ] + assert not sample_module.generations + assert not sample_module.input_generations + + sample_module.on_connection_terminating(session) + assert len(turn.completions) == 1 + + def test_sample_includes_setup_run_and_bridge_manifest(): readme = (_SAMPLE_ROOT / "README.md").read_text(encoding="utf-8") requirements = (_SAMPLE_ROOT / "requirements.txt").read_text(encoding="utf-8").splitlines() @@ -127,6 +232,7 @@ async def blocked_generation(_text): await sample_module.on_disconnect(session, SessionDisconnected(code=1006)) assert not generation.task.done() + session.termination = SessionTermination.TRANSPORT_ERROR sample_module.on_connection_terminating(session) with pytest.raises(asyncio.CancelledError): @@ -134,6 +240,116 @@ async def blocked_generation(_text): await asyncio.sleep(0) assert generation.task.cancelled() + assert generation.turn.activation_count == 1 + assert generation.turn.completions == [ + { + "outcome": TargetTurnOutcome.TRANSPORT_ERROR, + "response_id": generation.response_id, + "output_item_count": 0, + } + ] + assert not sample_module.generations + assert not sample_module.input_generations + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("termination", "expected"), + [ + pytest.param(SessionTermination.COMPLETED, TargetTurnOutcome.ABANDONED, id="clean-peer-close"), + pytest.param(SessionTermination.CALLBACK_ERROR, TargetTurnOutcome.ERROR, id="callback-error"), + pytest.param(SessionTermination.PROTOCOL_ERROR, TargetTurnOutcome.TRANSPORT_ERROR, id="protocol-error"), + ], +) +async def test_connection_cleanup_projects_source_aware_outcome_once( + sample_module, + monkeypatch, + termination, + expected, +): + generation_started = asyncio.Event() + + async def blocked_generation(_text): + generation_started.set() + await asyncio.Future() + yield "not reached" + + monkeypatch.setattr(sample_module, "generate_answer", blocked_generation) + session = _CapturingSession() + await sample_module.on_user_message( + session, + UserMessage( + id="m_user", + ts="2026-08-12T00:00:00Z", + item_id="in_1", + content=(InputTextPart(text="hello"),), + ), + ) + await asyncio.wait_for(generation_started.wait(), timeout=1) + generation = next(iter(sample_module.generations.values())) + sent_count = len(session.messages) + + session.termination = termination + sample_module.on_connection_terminating(session) + sample_module.on_connection_terminating(session) + + with pytest.raises(asyncio.CancelledError): + await generation.task + await asyncio.sleep(0) + + assert generation.turn.completions == [ + { + "outcome": expected, + "response_id": generation.response_id, + "output_item_count": 0, + } + ] + assert len(session.messages) == sent_count + assert not sample_module.generations + assert not sample_module.input_generations + + sample_module.on_connection_terminating(session) + assert len(generation.turn.completions) == 1 + + +@pytest.mark.asyncio +async def test_explicit_end_call_hint_wins_later_completed_connection(sample_module, monkeypatch): + generation_started = asyncio.Event() + + async def blocked_generation(_text): + generation_started.set() + await asyncio.Future() + yield "not reached" + + monkeypatch.setattr(sample_module, "generate_answer", blocked_generation) + session = _CapturingSession() + await sample_module.on_user_message( + session, + UserMessage( + id="m_user", + ts="2026-08-12T00:00:00Z", + item_id="in_1", + content=(InputTextPart(text="hello"),), + ), + ) + await asyncio.wait_for(generation_started.wait(), timeout=1) + generation = next(iter(sample_module.generations.values())) + + sample_module.cancel_session_generation_tasks(session, TargetTurnOutcome.END_CALL) + session.termination = SessionTermination.COMPLETED + sample_module.on_connection_terminating(session) + + with pytest.raises(asyncio.CancelledError): + await generation.task + await asyncio.sleep(0) + + assert generation.turn.completions == [ + { + "outcome": TargetTurnOutcome.END_CALL, + "response_id": generation.response_id, + "output_item_count": 0, + } + ] assert not sample_module.generations assert not sample_module.input_generations @@ -172,6 +388,243 @@ async def blocked_generation(_text): assert cleanup_finished.is_set() assert generation.task.cancelled() + assert generation.turn.completions == [ + { + "outcome": TargetTurnOutcome.END_CALL, + "response_id": generation.response_id, + "output_item_count": 0, + } + ] + assert not sample_module.generations + assert not sample_module.input_generations + + +@pytest.mark.asyncio +async def test_successful_generation_completes_declared_target_turn(sample_module, monkeypatch): + async def immediate_generation(_text): + yield "hello" + + monkeypatch.setattr(sample_module, "generate_answer", immediate_generation) + session = _CapturingSession() + await sample_module.on_user_message( + session, + UserMessage( + id="m_user", + ts="2026-08-12T00:00:00Z", + item_id="in_1", + content=(InputTextPart(text="hello"),), + ), + ) + generation = next(iter(sample_module.generations.values())) + await generation.task + await asyncio.sleep(0) + + assert generation.turn.activation_count == 1 + assert generation.turn.completions == [ + { + "outcome": TargetTurnOutcome.RESPONSE, + "response_id": generation.response_id, + "output_item_count": 1, + } + ] + assert not sample_module.generations + assert not sample_module.input_generations + + +@pytest.mark.asyncio +async def test_send_side_transport_failure_preserves_physical_outcome(sample_module): + session = _CapturingSession() + + async def fail_send(message): + session.messages.append(message) + session.termination = SessionTermination.TRANSPORT_ERROR + raise OSError("peer transport failed") + + session.send = fail_send + await sample_module.on_user_message( + session, + UserMessage( + id="m_user", + ts="2026-08-12T00:00:00Z", + item_id="in_1", + content=(InputTextPart(text="hello"),), + ), + ) + generation = next(iter(sample_module.generations.values())) + + with pytest.raises(OSError, match="peer transport failed"): + await generation.task + await asyncio.sleep(0) + + assert generation.turn.completions == [ + { + "outcome": TargetTurnOutcome.TRANSPORT_ERROR, + "response_id": None, + "output_item_count": 0, + } + ] + assert not sample_module.generations + assert not sample_module.input_generations + + +@pytest.mark.asyncio +async def test_generation_capacity_replies_none_without_retaining_another_task(sample_module, monkeypatch): + generation_started = asyncio.Event() + + async def blocked_generation(_text): + generation_started.set() + await asyncio.Future() + yield "not reached" + + monkeypatch.setattr(sample_module, "MAX_ACTIVE_GENERATIONS_PER_SESSION", 1) + monkeypatch.setattr(sample_module, "generate_answer", blocked_generation) + session = _CapturingSession() + await sample_module.on_user_message( + session, + UserMessage( + id="m_first", + ts="2026-08-12T00:00:00Z", + item_id="in_1", + content=(InputTextPart(text="first"),), + ), + ) + await asyncio.wait_for(generation_started.wait(), timeout=1) + generation = next(iter(sample_module.generations.values())) + + await sample_module.on_user_message( + session, + UserMessage( + id="m_second", + ts="2026-08-12T00:00:01Z", + item_id="in_2", + content=(InputTextPart(text="second"),), + ), + ) + + assert len(sample_module.generations) == 1 + assert len(sample_module.input_generations) == 1 + refusal = session.messages[-1] + assert isinstance(refusal, ResponseNone) + assert refusal.in_reply_to == ("in_2",) + assert refusal.reason == "capacity_exceeded" + assert session.turns[-1].completions == [ + { + "outcome": TargetTurnOutcome.NONE, + "output_item_count": 0, + } + ] + + generation.task.cancel() + with pytest.raises(asyncio.CancelledError): + await generation.task + await asyncio.sleep(0) + assert not sample_module.generations + assert not sample_module.input_generations + + +@pytest.mark.asyncio +async def test_generation_output_retention_is_bounded(sample_module, monkeypatch): + async def oversized_generation(_text): + yield "abc" + yield "def" + + monkeypatch.setattr(sample_module, "MAX_OUTPUT_CHUNKS", 1) + monkeypatch.setattr(sample_module, "MAX_OUTPUT_UTF8_BYTES", 5) + monkeypatch.setattr(sample_module, "generate_answer", oversized_generation) + session = _CapturingSession() + await sample_module.on_user_message( + session, + UserMessage( + id="m_user", + ts="2026-08-12T00:00:00Z", + item_id="in_1", + content=(InputTextPart(text="hello"),), + ), + ) + generation = next(iter(sample_module.generations.values())) + + with pytest.raises(RuntimeError, match="output exceeded sample limits"): + await generation.task + await asyncio.sleep(0) + + assert generation.turn.completions == [ + { + "outcome": TargetTurnOutcome.ERROR, + "response_id": generation.response_id, + "output_item_count": 0, + } + ] + assert not sample_module.generations + assert not sample_module.input_generations + + +@pytest.mark.asyncio +async def test_ordinary_application_cancellation_is_not_transport_error( + sample_module, + monkeypatch, +): + generation_started = asyncio.Event() + + async def blocked_generation(_text): + generation_started.set() + await asyncio.Future() + yield "not reached" + + monkeypatch.setattr(sample_module, "generate_answer", blocked_generation) + session = _CapturingSession() + await sample_module.on_user_message( + session, + UserMessage( + id="m_user", + ts="2026-08-12T00:00:00Z", + item_id="in_1", + content=(InputTextPart(text="hello"),), + ), + ) + await asyncio.wait_for(generation_started.wait(), timeout=1) + generation = next(iter(sample_module.generations.values())) + + generation.task.cancel() + with pytest.raises(asyncio.CancelledError): + await generation.task + await asyncio.sleep(0) + + assert generation.turn.completions == [ + { + "outcome": TargetTurnOutcome.CANCELLED, + "response_id": generation.response_id, + "output_item_count": 0, + } + ] + + +@pytest.mark.asyncio +async def test_pre_start_cancellation_completes_and_releases_turn(sample_module): + session = _CapturingSession() + await sample_module.on_user_message( + session, + UserMessage( + id="m_user", + ts="2026-08-12T00:00:00Z", + item_id="in_1", + content=(InputTextPart(text="hello"),), + ), + ) + generation = next(iter(sample_module.generations.values())) + + generation.task.cancel() + with pytest.raises(asyncio.CancelledError): + await generation.task + await asyncio.sleep(0) + + assert generation.turn.activation_count == 0 + assert generation.turn.completions == [ + { + "outcome": TargetTurnOutcome.CANCELLED, + "response_id": None, + "output_item_count": 0, + } + ] assert not sample_module.generations assert not sample_module.input_generations diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_session.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_session.py index ad57fab3547f..05d308178f4e 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_session.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_session.py @@ -66,7 +66,13 @@ async def test_session_has_only_transport_and_send_gate_and_serializes_writes(): websocket = _BlockingWebSocket() session = Session._create(websocket) # pylint: disable=protected-access assert not hasattr(session, "__dict__") - assert set(Session.__slots__) == {"_websocket", "_send_lock", "_terminal"} + assert set(Session.__slots__) == { + "_connection_context", + "_send_lock", + "_terminal", + "_termination", + "_websocket", + } first = asyncio.create_task(session.send(SessionReady())) await websocket.entered.wait() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_tracing.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_tracing.py new file mode 100644 index 000000000000..09614319f743 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_tracing.py @@ -0,0 +1,1392 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Tracing contract tests for the typed Voice relay.""" + +import asyncio +import json +import logging +import subprocess +import sys + +import pytest +from opentelemetry import baggage, metrics, trace +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanContext, TraceFlags, TraceState +from starlette.testclient import TestClient +from starlette.websockets import WebSocket, WebSocketDisconnect + +from azure.ai.agentserver.invocations._constants import InvocationsWSConstants +from azure.ai.agentserver.invocations.voice import ( + Session, + SessionDisconnected, + SessionReady, + SessionTermination, + TargetTurnOrigin, + TargetTurnOutcome, + VoiceAgentServerHost, + new_response_id, +) +from azure.ai.agentserver.invocations.voice import _voice_host as voice_host_module +from azure.ai.agentserver.invocations.voice import _tracing as tracing_module +from azure.ai.agentserver.invocations.voice import _turn as turn_module + + +_PROVIDER = None +_EXPORTER = None +_METER_PROVIDER = None +_METRIC_READER = None + + +@pytest.fixture +def spans(): + """Capture spans without replacing a provider another test installed.""" + global _PROVIDER, _EXPORTER + if _PROVIDER is None: + existing = trace.get_tracer_provider() + if hasattr(existing, "add_span_processor"): + _PROVIDER = existing + else: + _PROVIDER = TracerProvider() + trace.set_tracer_provider(_PROVIDER) + _EXPORTER = InMemorySpanExporter() + _PROVIDER.add_span_processor(SimpleSpanProcessor(_EXPORTER)) + _EXPORTER.clear() + return _PROVIDER, _EXPORTER + + +@pytest.fixture +def metric_reader(): + """Install one in-memory reader after module-level proxy instruments exist.""" + global _METER_PROVIDER, _METRIC_READER + if _METER_PROVIDER is None: + _METRIC_READER = InMemoryMetricReader() + _METER_PROVIDER = MeterProvider(metric_readers=[_METRIC_READER]) + metrics.set_meter_provider(_METER_PROVIDER) + return _METRIC_READER + + +def _session_start_frame() -> dict[str, object]: + return { + "type": "session.start", + "id": "m_start", + "ts": "2026-08-17T00:00:00Z", + "protocol_version": "1.0", + "reconnect": False, + "response_timeouts": { + "first_output_ms": 1, + "idle_ms": 2, + "max_duration_ms": 3, + }, + } + + +def _span_by_name(exporter: InMemorySpanExporter, name: str): + matches = [span for span in exporter.get_finished_spans() if span.name == name] + assert len(matches) == 1, [span.name for span in exporter.get_finished_spans()] + return matches[0] + + +def _metric_points(reader: InMemoryMetricReader, name: str): + data = reader.get_metrics_data() + if data is None: + return [] + return [ + point + for resource_metrics in data.resource_metrics + for scope_metrics in resource_metrics.scope_metrics + for metric in scope_metrics.metrics + if metric.name == name + for point in metric.data.data_points + ] + + +def _websocket_with_headers(headers: list[tuple[bytes, bytes]]) -> WebSocket: + async def receive(): + return {"type": "websocket.disconnect", "code": 1000} + + async def send(_message): + return None + + return WebSocket( + { + "type": "websocket", + "asgi": {"version": "3.0", "spec_version": "2.4"}, + "scheme": "ws", + "path": "/invocations_ws", + "raw_path": b"/invocations_ws", + "query_string": b"", + "headers": headers, + "client": ("test", 1), + "server": ("testserver", 80), + "subprotocols": [], + "state": {}, + }, + receive, + send, + ) + + +@pytest.mark.parametrize( + "factory_setup", + [ + "trace.get_tracer = lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError('tracer factory'))", + "metrics.get_meter = lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError('meter factory'))", + """ +class FailingMeter: + def create_histogram(self, *_args, **_kwargs): + raise RuntimeError("histogram factory") + + def create_counter(self, *_args, **_kwargs): + raise RuntimeError("counter factory") + +metrics.get_meter = lambda *_args, **_kwargs: FailingMeter() +""", + """ +import logging +real_get_logger = logging.getLogger + +def fail_voice_logger(name=None): + if name == "azure.ai.agentserver": + raise RuntimeError("logger factory") + return real_get_logger(name) + +logging.getLogger = fail_voice_logger +""", + """ +import opentelemetry.trace.propagation.tracecontext as tracecontext +tracecontext.TraceContextTextMapPropagator = lambda: (_ for _ in ()).throw(RuntimeError("propagator factory")) +""", + ], +) +def test_first_voice_import_survives_throwing_telemetry_factories(factory_setup): + script = f""" +from opentelemetry import metrics, trace +import azure.ai.agentserver.invocations + +{factory_setup} + +from azure.ai.agentserver.invocations.voice import SessionTermination, TargetTurnOutcome + +print(SessionTermination.CANCELLED.value, TargetTurnOutcome.ERROR.value) +""" + completed = subprocess.run( # nosec B603 - fixed interpreter and script + [sys.executable, "-c", script], + check=False, + capture_output=True, + text=True, + ) + + assert completed.returncode == 0, completed.stderr + assert completed.stdout.strip() == "cancelled error" + + +def test_voice_connection_and_callback_have_explicit_semantic_parents(spans): + provider, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + customer_tracer = provider.get_tracer("customer.agent") + + @app.on_session_start + async def on_session_start(session, _event): + with customer_tracer.start_as_current_span("customer.callback"): + pass + await session.send(SessionReady()) + + remote_trace_id = "11111111111111111111111111111111" + remote_span_id = "2222222222222222" + headers = {"traceparent": f"00-{remote_trace_id}-{remote_span_id}-01"} + with TestClient(app).websocket_connect("/invocations_ws", headers=headers) as websocket: + websocket.send_json(_session_start_frame()) + assert websocket.receive_json()["type"] == "session.ready" + websocket.send_json( + { + "type": "session.end", + "id": "m_end", + "ts": "2026-08-17T00:00:01Z", + "reason": "completed", + } + ) + assert websocket.receive()["type"] == "websocket.close" + + connection = _span_by_name(exporter, "agentserver.connection") + callback = _span_by_name(exporter, "voice.callback") + customer = _span_by_name(exporter, "customer.callback") + + assert f"{connection.context.trace_id:032x}" == remote_trace_id + assert connection.parent is not None + assert f"{connection.parent.span_id:016x}" == remote_span_id + assert callback.parent is not None and callback.parent.span_id == connection.context.span_id + assert customer.parent is not None and customer.parent.span_id == callback.context.span_id + assert callback.attributes == {"voice.event.type": "session.start"} + assert connection.attributes["bridge.outcome"] == "completed" + assert not [span for span in exporter.get_finished_spans() if span.name == "invoke_agent"] + + +@pytest.mark.asyncio +async def test_disconnect_callback_is_a_connection_sibling_and_parents_customer_work(spans): + provider, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + customer_tracer = provider.get_tracer("customer.agent") + + @app.on_session_start + async def on_session_start(session, _event): + await session.send(SessionReady()) + + @app.on_disconnect + async def on_disconnect(_session, _event): + with customer_tracer.start_as_current_span("customer.disconnect"): + pass + + inbound_events = [ + {"type": "websocket.connect"}, + {"type": "websocket.receive", "text": json.dumps(_session_start_frame())}, + {"type": "websocket.disconnect", "code": 1000}, + ] + websocket = _websocket_with_headers([(b"traceparent", b"00-11111111111111111111111111111111-2222222222222222-01")]) + websocket._receive = lambda: asyncio.sleep(0, result=inbound_events.pop(0)) # pylint: disable=protected-access + websocket._send = lambda _message: asyncio.sleep(0) # pylint: disable=protected-access + + await asyncio.wait_for(app._ws_endpoint(websocket), timeout=1) # pylint: disable=protected-access + + connection = _span_by_name(exporter, "agentserver.connection") + callbacks = [span for span in exporter.get_finished_spans() if span.name == "voice.callback"] + disconnects = [span for span in callbacks if span.attributes.get("voice.event.type") == "disconnect"] + assert len(disconnects) == 1, [(span.name, dict(span.attributes)) for span in exporter.get_finished_spans()] + disconnect = disconnects[0] + customer = _span_by_name(exporter, "customer.disconnect") + assert disconnect.parent is not None and disconnect.parent.span_id == connection.context.span_id + assert customer.parent is not None and customer.parent.span_id == disconnect.context.span_id + + +@pytest.mark.asyncio +async def test_disconnect_callback_failure_marks_only_content_free_callback_error(spans): + _, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + + @app.on_session_start + async def on_session_start(session, _event): + await session.send(SessionReady()) + + @app.on_disconnect + async def on_disconnect(_session, _event): + raise RuntimeError("private-disconnect-sentinel") + + inbound_events = [ + {"type": "websocket.connect"}, + {"type": "websocket.receive", "text": json.dumps(_session_start_frame())}, + {"type": "websocket.disconnect", "code": 1000}, + ] + websocket = _websocket_with_headers([]) + websocket._receive = lambda: asyncio.sleep(0, result=inbound_events.pop(0)) # pylint: disable=protected-access + websocket._send = lambda _message: asyncio.sleep(0) # pylint: disable=protected-access + + await asyncio.wait_for(app._ws_endpoint(websocket), timeout=1) # pylint: disable=protected-access + + callbacks = [span for span in exporter.get_finished_spans() if span.name == "voice.callback"] + disconnects = [span for span in callbacks if span.attributes.get("voice.event.type") == "disconnect"] + assert len(disconnects) == 1, [(span.name, dict(span.attributes)) for span in exporter.get_finished_spans()] + disconnect = disconnects[0] + assert disconnect.attributes["error.type"] == "callback_error" + assert disconnect.status.status_code is trace.StatusCode.ERROR + assert "private-disconnect-sentinel" not in repr(disconnect.attributes) + + +def test_voice_upgrade_rebuilds_only_hosted_agents_baggage(): + websocket = _websocket_with_headers( + [ + ( + b"traceparent", + b"00-11111111111111111111111111111111-2222222222222222-01", + ), + ( + b"baggage", + b"azure.ai.agentserver.session_id=session-safe," + b"microsoft.a365.agent.blueprint.id=blueprint-safe," + b"user.id=user-safe,gen_ai.agent.id=agent-safe," + b"microsoft.tenant.id=tenant-safe,customer-secret=private-sentinel", + ), + (b"x-request-id", b"request-safe"), + ] + ) + + extracted = voice_host_module._extract_voice_websocket_context(websocket) # pylint: disable=protected-access + + assert baggage.get_baggage("azure.ai.agentserver.session_id", context=extracted) == "session-safe" + assert baggage.get_baggage("microsoft.a365.agent.blueprint.id", context=extracted) == "blueprint-safe" + assert baggage.get_baggage("user.id", context=extracted) == "user-safe" + assert baggage.get_baggage("gen_ai.agent.id", context=extracted) == "agent-safe" + assert baggage.get_baggage("microsoft.tenant.id", context=extracted) == "tenant-safe" + assert baggage.get_baggage("x_request_id", context=extracted) == "request-safe" + assert baggage.get_baggage("customer-secret", context=extracted) is None + + +def test_duplicate_approved_baggage_is_dropped_and_classified(metric_reader): + websocket = _websocket_with_headers( + [ + ( + b"traceparent", + b"00-11111111111111111111111111111111-2222222222222222-01", + ), + ( + b"baggage", + b"azure.ai.agentserver.session_id=session-first," b"azure.ai.agentserver.session_id=session-second", + ), + ] + ) + + extracted = voice_host_module._extract_voice_websocket_context(websocket) # pylint: disable=protected-access + + assert baggage.get_baggage("azure.ai.agentserver.session_id", context=extracted) is None + points = _metric_points( + metric_reader, + "azure.ai.agentserver.trace_context.propagation_failures", + ) + assert any(point.attributes["error.type"] == "invalid" for point in points) + + +@pytest.mark.parametrize( + ("header_name", "header_value"), + [ + (b"baggage", b"customer-secret-private-sentinel"), + (b"tracestate", b"Private-Sentinel=value"), + ], +) +def test_invalid_propagation_never_logs_raw_member(caplog, header_name, header_value): + websocket = _websocket_with_headers( + [ + ( + b"traceparent", + b"00-11111111111111111111111111111111-2222222222222222-01", + ), + (header_name, header_value), + ] + ) + + with caplog.at_level(logging.WARNING): + voice_host_module._extract_voice_websocket_context(websocket) # pylint: disable=protected-access + + assert "private-sentinel" not in caplog.text.lower() + + +def test_unsampled_parent_propagates_without_exporting_semantic_spans(spans): + _, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + observed_contexts = [] + + @app.on_session_start + async def on_session_start(session, _event): + observed_contexts.append(trace.get_current_span().get_span_context()) + await session.send(SessionReady()) + + remote_trace_id = "33333333333333333333333333333333" + remote_span_id = "4444444444444444" + headers = { + "traceparent": f"00-{remote_trace_id}-{remote_span_id}-00", + "tracestate": "vendor=value", + } + with TestClient(app).websocket_connect("/invocations_ws", headers=headers) as websocket: + websocket.send_json(_session_start_frame()) + assert websocket.receive_json()["type"] == "session.ready" + + assert len(observed_contexts) == 1 + observed = observed_contexts[0] + assert f"{observed.trace_id:032x}" == remote_trace_id + assert f"{observed.span_id:016x}" != remote_span_id + assert not observed.trace_flags.sampled + assert observed.trace_state.get("vendor") == "value" + assert exporter.get_finished_spans() == () + + +def test_declared_target_turn_uses_explicit_connection_parent(spans): + provider, exporter = spans + tracer = provider.get_tracer("azure.ai.agentserver.invocations.voice") + customer_tracer = provider.get_tracer("customer.agent") + connection = tracer.start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=2) + with turn.activate(): + with customer_tracer.start_as_current_span("customer.model"): + pass + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + connection.end() + + target = _span_by_name(exporter, "invoke_agent") + customer = _span_by_name(exporter, "customer.model") + assert target.parent is not None and target.parent.span_id == connection.context.span_id + assert customer.parent is not None and customer.parent.span_id == target.context.span_id + assert target.attributes["gen_ai.operation.name"] == "invoke_agent" + assert target.attributes["turn.origin"] == "user" + assert target.attributes["bridge.input.count"] == 2 + assert target.attributes["bridge.output.item_count"] == 0 + assert target.attributes["bridge.outcome"] == "none" + assert target.status.status_code is trace.StatusCode.UNSET + + +@pytest.mark.asyncio +async def test_target_turn_covers_application_owned_background_work(spans): + provider, exporter = spans + tracer = provider.get_tracer("azure.ai.agentserver.invocations.voice") + customer_tracer = provider.get_tracer("customer.agent") + connection = tracer.start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + response_id = new_response_id() + started = asyncio.Event() + release = asyncio.Event() + + async def background_work(): + with turn.activate(): + with customer_tracer.start_as_current_span("customer.background"): + started.set() + await release.wait() + turn.complete( + outcome=TargetTurnOutcome.RESPONSE, + response_id=response_id, + output_item_count=1, + ) + + task = asyncio.create_task(background_work()) + await asyncio.wait_for(started.wait(), timeout=1) + assert not [span for span in exporter.get_finished_spans() if span.name == "invoke_agent"] + + release.set() + await asyncio.wait_for(task, timeout=1) + connection.end() + + target = _span_by_name(exporter, "invoke_agent") + customer = _span_by_name(exporter, "customer.background") + assert customer.parent is not None and customer.parent.span_id == target.context.span_id + assert target.attributes["gen_ai.response.id"] == response_id + assert target.attributes["bridge.output.item_count"] == 1 + assert target.attributes["bridge.outcome"] == "response" + + +def test_target_turn_rejects_completion_while_active_and_second_activation(spans): + provider, _ = spans + connection = provider.get_tracer("test.connection").start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + turn = session.start_target_turn(origin=TargetTurnOrigin.NO_INPUT, input_count=1) + + with turn.activate(): + with pytest.raises(RuntimeError, match="active"): + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + + with pytest.raises(RuntimeError, match="activated"): + with turn.activate(): + pass + + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + assert turn.is_completed + connection.end() + + +@pytest.mark.parametrize( + "kwargs", + [ + {"outcome": TargetTurnOutcome.RESPONSE, "output_item_count": 1}, + { + "outcome": TargetTurnOutcome.RESPONSE, + "response_id": "r_real", + "output_item_count": 0, + }, + {"outcome": TargetTurnOutcome.NONE}, + { + "outcome": TargetTurnOutcome.NONE, + "response_id": "r_real", + "output_item_count": 0, + }, + {"outcome": TargetTurnOutcome.ERROR, "output_item_count": 1}, + ], +) +def test_target_turn_rejects_contradictory_completion_facts(spans, kwargs): + provider, _ = spans + connection = provider.get_tracer("test.connection").start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + + with pytest.raises((TypeError, ValueError)): + turn.complete(**kwargs) + + assert not turn.is_completed + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + connection.end() + + +@pytest.mark.parametrize( + ("scenario", "expected"), + [ + ("completed", SessionTermination.COMPLETED), + ("protocol_error", SessionTermination.PROTOCOL_ERROR), + ("callback_error", SessionTermination.CALLBACK_ERROR), + ], +) +def test_connection_termination_fact_is_visible_before_cleanup(spans, scenario, expected): + _, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + observed = [] + + @app.on_session_start + async def on_session_start(session, _event): + if scenario == "callback_error": + raise RuntimeError("private callback detail") + await session.send(SessionReady()) + + @app.on_connection_terminating + def on_connection_terminating(session): + observed.append(session.termination) + + with TestClient(app).websocket_connect("/invocations_ws") as websocket: + if scenario == "protocol_error": + websocket.send_text("not-json") + assert websocket.receive()["code"] == 1002 + else: + websocket.send_json(_session_start_frame()) + if scenario == "completed": + assert websocket.receive_json()["type"] == "session.ready" + websocket.send_json( + { + "type": "session.end", + "id": "m_end", + "ts": "2026-08-17T00:00:01Z", + "reason": "completed", + } + ) + assert websocket.receive()["type"] == "websocket.close" + + assert observed == [expected] + connection = _span_by_name(exporter, "agentserver.connection") + assert connection.attributes["bridge.outcome"] == expected.value + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("close_code", "expected_termination"), + [ + pytest.param(1000, SessionTermination.COMPLETED, id="normal-baseline"), + pytest.param(1001, SessionTermination.COMPLETED, id="going-away"), + pytest.param(1002, SessionTermination.PROTOCOL_ERROR, id="protocol-error"), + pytest.param(1003, SessionTermination.PROTOCOL_ERROR, id="unsupported-data"), + pytest.param(1007, SessionTermination.PROTOCOL_ERROR, id="invalid-payload"), + pytest.param(1008, SessionTermination.PROTOCOL_ERROR, id="policy-violation"), + pytest.param(1009, SessionTermination.PROTOCOL_ERROR, id="message-too-big"), + pytest.param(1010, SessionTermination.PROTOCOL_ERROR, id="mandatory-extension"), + pytest.param(1011, SessionTermination.TRANSPORT_ERROR, id="internal-error-negative-control"), + pytest.param(1006, SessionTermination.TRANSPORT_ERROR, id="abnormal-negative-control"), + ], +) +async def test_peer_disconnect_classification_matches_cleanup_and_telemetry( + spans, + metric_reader, + caplog, + close_code, + expected_termination, +): + _, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + callback_order = [] + terminations = [] + disconnects = [] + rejected_writes = [] + later_callbacks = [] + sent_messages = [] + inbound_events = [ + {"type": "websocket.connect"}, + {"type": "websocket.disconnect", "code": close_code, "reason": "peer close"}, + {"type": "websocket.receive", "text": json.dumps(_session_start_frame())}, + ] + + @app.on_session_start + async def on_session_start(_session, _event): + later_callbacks.append("session.start") + + @app.on_connection_terminating + def on_connection_terminating(session): + callback_order.append("terminating") + terminations.append(session.termination) + + @app.on_disconnect + async def on_disconnect(session, event): + callback_order.append("disconnect") + disconnects.append((session.termination, event.code, event.reason)) + with pytest.raises(RuntimeError, match="terminating"): + await session.send(SessionReady()) + rejected_writes.append(True) + + async def receive(): + return inbound_events.pop(0) + + async def send(message): + sent_messages.append(message) + + websocket = _websocket_with_headers([]) + websocket._receive = receive # pylint: disable=protected-access + websocket._send = send # pylint: disable=protected-access + + expected_outcome = expected_termination.value + expected_metric_attributes = {"bridge.outcome": expected_outcome} + if expected_termination is not SessionTermination.COMPLETED: + expected_metric_attributes["error.type"] = expected_outcome + metric_name = "azure.ai.agentserver.voice.connection.duration" + before_metric_count = sum( + point.count + for point in _metric_points(metric_reader, metric_name) + if point.attributes == expected_metric_attributes + ) + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + await asyncio.wait_for(app._ws_endpoint(websocket), timeout=1) # pylint: disable=protected-access + + assert callback_order == ["terminating", "disconnect"] + assert terminations == [expected_termination] + assert disconnects == [(expected_termination, close_code, "peer close")] + assert rejected_writes == [True] + assert later_callbacks == [] + assert len(inbound_events) == 1 + assert [message["type"] for message in sent_messages] == ["websocket.accept"] + assert Session._current(websocket) is None # pylint: disable=protected-access + close_records = [record for record in caplog.records if record.getMessage() == "Voice connection closed"] + assert len(close_records) == 1 + assert getattr(close_records[0], InvocationsWSConstants.ATTR_SPAN_CLOSE_CODE) == close_code + + connection = _span_by_name(exporter, "agentserver.connection") + assert connection.attributes["bridge.outcome"] == expected_outcome + fallback_websocket = _websocket_with_headers([]) + fallback_session = Session._create(fallback_websocket) # pylint: disable=protected-access + try: + voice_host_module._commit_voice_session_termination( # pylint: disable=protected-access + fallback_session, + handler_error=None, + disconnect_event=SessionDisconnected(code=close_code, reason="peer close"), + close_code=close_code, + accept_failed=False, + ) + assert fallback_session.termination is expected_termination + finally: + Session._release(fallback_websocket, fallback_session) # pylint: disable=protected-access + assert tracing_module._connection_outcome(close_code, None) == expected_outcome # pylint: disable=protected-access + points = _metric_points(metric_reader, metric_name) + after_metric_count = sum(point.count for point in points if point.attributes == expected_metric_attributes) + assert after_metric_count - before_metric_count == 1 + if expected_termination is SessionTermination.COMPLETED: + assert "error.type" not in connection.attributes + assert connection.status.status_code is trace.StatusCode.UNSET + else: + assert connection.attributes["error.type"] == expected_outcome + assert connection.status.status_code is trace.StatusCode.ERROR + + +def test_missing_and_invalid_context_record_sanitized_failure_metrics(metric_reader): + metric_name = "azure.ai.agentserver.trace_context.propagation_failures" + before = {point.attributes["error.type"]: point.value for point in _metric_points(metric_reader, metric_name)} + voice_host_module._extract_voice_websocket_context(_websocket_with_headers([])) # pylint: disable=protected-access + voice_host_module._extract_voice_websocket_context( # pylint: disable=protected-access + _websocket_with_headers([(b"traceparent", b"private-invalid-traceparent")]) + ) + + points = _metric_points( + metric_reader, + metric_name, + ) + dimensions = {tuple(sorted(point.attributes.items())) for point in points} + assert dimensions == { + ( + ( + "azure.ai.agentserver.trace_context.propagation.hop", + "hosted_agents_to_agentserver", + ), + ("error.type", "invalid"), + ), + ( + ( + "azure.ai.agentserver.trace_context.propagation.hop", + "hosted_agents_to_agentserver", + ), + ("error.type", "missing"), + ), + } + after = {point.attributes["error.type"]: point.value for point in points} + assert after["missing"] - before.get("missing", 0) == 1 + assert after["invalid"] - before.get("invalid", 0) == 1 + + +def test_unsampled_operations_still_record_duration_metrics(spans, metric_reader): + _, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + + @app.on_session_start + async def on_session_start(session, _event): + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + with turn.activate(): + pass + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + await session.send(SessionReady()) + + headers = { + "traceparent": "00-55555555555555555555555555555555-6666666666666666-00", + } + with TestClient(app).websocket_connect("/invocations_ws", headers=headers) as websocket: + websocket.send_json(_session_start_frame()) + assert websocket.receive_json()["type"] == "session.ready" + websocket.send_json( + { + "type": "session.end", + "id": "m_end", + "ts": "2026-08-17T00:00:01Z", + "reason": "completed", + } + ) + assert websocket.receive()["type"] == "websocket.close" + + assert exporter.get_finished_spans() == () + connection_points = _metric_points( + metric_reader, + "azure.ai.agentserver.voice.connection.duration", + ) + target_points = _metric_points(metric_reader, "gen_ai.invoke_agent.duration") + assert any(point.attributes == {"bridge.outcome": "completed"} for point in connection_points) + assert any(point.attributes == {"bridge.outcome": "none", "turn.origin": "user"} for point in target_points) + + +def test_target_turn_projects_one_content_free_trigger_link(spans): + provider, exporter = spans + connection = provider.get_tracer("test.connection").start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + trigger = SpanContext( + trace_id=0xAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA, + span_id=0xBBBBBBBBBBBBBBBB, + is_remote=True, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=TraceState.from_header(["private=value"]), + ) + + turn = session.start_target_turn( + origin=TargetTurnOrigin.PROACTIVE, + input_count=0, + trigger_context=trigger, + ) + with turn.activate(): + pass + turn.complete( + outcome=TargetTurnOutcome.RESPONSE, + response_id="r_proactive", + output_item_count=1, + ) + connection.end() + + target = _span_by_name(exporter, "invoke_agent") + assert len(target.links) == 1 + link = target.links[0] + assert link.context.trace_id == trigger.trace_id + assert link.context.span_id == trigger.span_id + assert link.context.trace_flags == trigger.trace_flags + assert link.context.is_remote == trigger.is_remote + assert len(link.context.trace_state) == 0 + assert not link.attributes + + +@pytest.mark.parametrize("constructor_name", ["SpanContext", "TraceState", "Link"]) +def test_trigger_link_factory_failure_does_not_change_target_turn(monkeypatch, spans, constructor_name): + provider, exporter = spans + connection = provider.get_tracer("test.connection").start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + trigger = SpanContext( + trace_id=0xAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA, + span_id=0xBBBBBBBBBBBBBBBB, + is_remote=True, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=TraceState(), + ) + + def fail_constructor(*_args, **_kwargs): + raise RuntimeError("private link factory failure") + + monkeypatch.setattr(turn_module, constructor_name, fail_constructor) + turn = session.start_target_turn( + origin=TargetTurnOrigin.PROACTIVE, + input_count=0, + trigger_context=trigger, + ) + with turn.activate(): + pass + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + connection.end() + + target = _span_by_name(exporter, "invoke_agent") + assert target.links == () + assert target.attributes["bridge.outcome"] == "none" + + +def test_unsafe_response_id_is_not_projected_into_target_span(spans): + provider, exporter = spans + connection = provider.get_tracer("test.connection").start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + with turn.activate(): + pass + + turn.complete( + outcome=TargetTurnOutcome.RESPONSE, + response_id="r_private-secret-token", + output_item_count=1, + ) + connection.end() + + target = _span_by_name(exporter, "invoke_agent") + assert target.attributes["bridge.outcome"] == "response" + assert target.attributes["bridge.output.item_count"] == 1 + assert "gen_ai.response.id" not in target.attributes + assert "private-secret-token" not in repr(target.attributes) + + +def test_terminal_session_rejects_new_target_turn(spans): + provider, _ = spans + connection = provider.get_tracer("test.connection").start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + session._begin_termination(SessionTermination.TRANSPORT_ERROR) # pylint: disable=protected-access + + with pytest.raises(RuntimeError, match="terminating"): + session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + + assert session._connection_context is None # pylint: disable=protected-access + connection.end() + + +def test_error_target_metric_matches_span_outcome(spans, metric_reader): + provider, exporter = spans + connection = provider.get_tracer("test.connection").start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + with turn.activate(): + pass + turn.complete(outcome=TargetTurnOutcome.TIMEOUT, output_item_count=0) + connection.end() + + target = _span_by_name(exporter, "invoke_agent") + assert target.attributes["bridge.outcome"] == "timeout" + assert target.attributes["error.type"] == "timeout" + points = _metric_points(metric_reader, "gen_ai.invoke_agent.duration") + assert any( + point.attributes == {"bridge.outcome": "timeout", "error.type": "timeout", "turn.origin": "user"} + for point in points + ) + + +@pytest.mark.parametrize(("environment_value", "expected"), [(None, False), ("true", True)]) +def test_voice_default_observability_requires_explicit_sensitive_opt_in( + monkeypatch, + environment_value, + expected, +): + calls = [] + if environment_value is None: + monkeypatch.delenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", raising=False) + else: + monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", environment_value) + monkeypatch.setattr( + voice_host_module, + "_CORE_CONFIGURE_OBSERVABILITY", + lambda **kwargs: calls.append(kwargs), + raising=False, + ) + + VoiceAgentServerHost() + + assert len(calls) == 1 + assert calls[0]["enable_sensitive_data"] is expected + + +@pytest.mark.parametrize( + "failure", + [RuntimeError("private-observability-sentinel"), asyncio.CancelledError("private-observability-sentinel")], +) +def test_voice_default_observability_failure_is_content_free_and_fail_open(monkeypatch, caplog, failure): + def fail_observability(**_kwargs): + raise failure + + monkeypatch.setattr(voice_host_module, "_CORE_CONFIGURE_OBSERVABILITY", fail_observability) + + with caplog.at_level(logging.WARNING): + app = VoiceAgentServerHost() + + assert isinstance(app, VoiceAgentServerHost) + assert "private-observability-sentinel" not in caplog.text + + +def test_voice_diagnostics_are_content_free(monkeypatch, caplog, spans): + _, exporter = spans + monkeypatch.setenv("FOUNDRY_AGENT_SESSION_ID", "private session sentinel") + app = VoiceAgentServerHost(configure_observability=None) + + @app.on_session_start + async def on_session_start(_session, _event): + raise RuntimeError("private-exception-sentinel") + + with caplog.at_level(logging.INFO, logger="azure.ai.agentserver"): + with TestClient(app).websocket_connect("/invocations_ws") as websocket: + websocket.send_json(_session_start_frame()) + assert websocket.receive()["type"] == "websocket.close" + + assert "private session sentinel" not in caplog.text + assert "private-exception-sentinel" not in caplog.text + voice_records = [record for record in caplog.records if str(record.msg).startswith("Voice ")] + assert voice_records + for record in voice_records: + assert record.args == () + assert record.exc_info is None + connection = _span_by_name(exporter, "agentserver.connection") + assert "private session sentinel" not in repr(connection.attributes) + + +def test_connection_span_carries_safe_session_and_protocol_attributes(monkeypatch, spans): + _, exporter = spans + monkeypatch.setenv("FOUNDRY_AGENT_SESSION_ID", "session_safe_123") + app = VoiceAgentServerHost(configure_observability=None) + + @app.on_session_start + async def on_session_start(session, _event): + await session.send(SessionReady()) + + with TestClient(app).websocket_connect("/invocations_ws") as websocket: + websocket.send_json(_session_start_frame()) + assert websocket.receive_json()["type"] == "session.ready" + websocket.send_json( + { + "type": "session.end", + "id": "m_end", + "ts": "2026-08-17T00:00:01Z", + "reason": "completed", + } + ) + assert websocket.receive()["type"] == "websocket.close" + + connection = _span_by_name(exporter, "agentserver.connection") + assert connection.attributes["azure.ai.agentserver.invocations_ws.session_id"] == "session_safe_123" + assert connection.attributes["azure.ai.agentserver.invocations_ws.protocol_version"] == "1.0" + assert connection.attributes["azure.ai.agentserver.invocations_ws.reconnect"] is False + + +def test_connection_span_omits_unsafe_protocol_content_without_changing_dispatch(spans): + _, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + observed_protocols = [] + + @app.on_session_start + async def on_session_start(session, event): + observed_protocols.append(event.protocol_version) + await session.send(SessionReady()) + + private_protocol = "private-protocol-sentinel" + frame = _session_start_frame() + frame["protocol_version"] = private_protocol + with TestClient(app).websocket_connect("/invocations_ws") as websocket: + websocket.send_json(frame) + assert websocket.receive_json()["type"] == "session.ready" + websocket.send_json( + { + "type": "session.end", + "id": "m_end", + "ts": "2026-08-17T00:00:01Z", + "reason": "completed", + } + ) + assert websocket.receive()["type"] == "websocket.close" + + assert observed_protocols == [private_protocol] + connection = _span_by_name(exporter, "agentserver.connection") + assert "azure.ai.agentserver.invocations_ws.protocol_version" not in connection.attributes + assert private_protocol not in repr(connection.attributes) + + +@pytest.mark.parametrize("failure_stage", ["parent_attach", "connection_attach"]) +def test_connection_setup_failure_disables_semantic_descendants( + monkeypatch, + spans, + failure_stage, +): + _, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + callback_count = 0 + + if failure_stage == "parent_attach": + monkeypatch.setattr(voice_host_module, "_attach_context", lambda _context: None) + else: + monkeypatch.setattr(tracing_module, "_attach_context", lambda _context: None) + + @app.on_session_start + async def on_session_start(session, _event): + nonlocal callback_count + callback_count += 1 + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + with turn.activate(): + pass + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + await session.send(SessionReady()) + + headers = {"traceparent": "00-77777777777777777777777777777777-8888888888888888-01"} + with TestClient(app).websocket_connect("/invocations_ws", headers=headers) as websocket: + websocket.send_json(_session_start_frame()) + assert websocket.receive_json()["type"] == "session.ready" + websocket.send_json( + { + "type": "session.end", + "id": "m_end", + "ts": "2026-08-17T00:00:01Z", + "reason": "completed", + } + ) + assert websocket.receive()["type"] == "websocket.close" + + assert callback_count == 1 + semantic = [ + span + for span in exporter.get_finished_spans() + if span.name in {"agentserver.connection", "voice.callback", "invoke_agent"} + ] + expected_names = [] if failure_stage == "parent_attach" else ["agentserver.connection"] + assert [span.name for span in semantic] == expected_names + + +def test_context_factory_failure_disables_telemetry_without_changing_wire(monkeypatch, spans): + _, exporter = spans + + def fail_context(): + raise RuntimeError("private context factory failure") + + monkeypatch.setattr(voice_host_module._otel_context, "Context", fail_context) + app = VoiceAgentServerHost(configure_observability=None) + + @app.on_session_start + async def on_session_start(session, _event): + await session.send(SessionReady()) + + with TestClient(app).websocket_connect("/invocations_ws") as websocket: + websocket.send_json(_session_start_frame()) + assert websocket.receive_json()["type"] == "session.ready" + + assert not [ + span + for span in exporter.get_finished_spans() + if span.name in {"agentserver.connection", "voice.callback", "invoke_agent"} + ] + + +@pytest.mark.asyncio +async def test_send_side_peer_loss_marks_callback_and_connection_transport_error(spans): + _, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + inbound_events = [ + {"type": "websocket.connect"}, + {"type": "websocket.receive", "text": json.dumps(_session_start_frame())}, + ] + + @app.on_session_start + async def on_session_start(session, _event): + await session.send(SessionReady()) + + async def receive(): + return inbound_events.pop(0) + + async def send(message): + if message["type"] == "websocket.send": + raise OSError("private peer loss detail") + + websocket = _websocket_with_headers([]) + websocket._receive = receive # pylint: disable=protected-access + websocket._send = send # pylint: disable=protected-access + + await asyncio.wait_for(app._ws_endpoint(websocket), timeout=1) # pylint: disable=protected-access + + connection = _span_by_name(exporter, "agentserver.connection") + callbacks = [span for span in exporter.get_finished_spans() if span.name == "voice.callback"] + callback = next(span for span in callbacks if span.attributes.get("voice.event.type") == "session.start") + assert callback.attributes["error.type"] == "transport_error" + assert callback.status.status_code is trace.StatusCode.ERROR + assert connection.attributes["bridge.outcome"] == "transport_error" + assert "private peer loss detail" not in repr(callback.attributes) + + +@pytest.mark.asyncio +async def test_application_websocket_disconnect_is_callback_error(spans): + _, exporter = spans + app = VoiceAgentServerHost(configure_observability=None) + observed_terminations = [] + disconnects = [] + sent_messages = [] + inbound_events = [ + {"type": "websocket.connect"}, + {"type": "websocket.receive", "text": json.dumps(_session_start_frame())}, + ] + + @app.on_session_start + async def on_session_start(_session, _event): + raise WebSocketDisconnect(code=1000, reason="private callback detail") + + @app.on_connection_terminating + def on_connection_terminating(session): + observed_terminations.append(session.termination) + + @app.on_disconnect + async def on_disconnect(_session, event): + disconnects.append(event) + + async def receive(): + return inbound_events.pop(0) + + async def send(message): + sent_messages.append(message) + + websocket = _websocket_with_headers([]) + websocket._receive = receive # pylint: disable=protected-access + websocket._send = send # pylint: disable=protected-access + + await asyncio.wait_for(app._ws_endpoint(websocket), timeout=1) # pylint: disable=protected-access + + connection = _span_by_name(exporter, "agentserver.connection") + callback = next( + span + for span in exporter.get_finished_spans() + if span.name == "voice.callback" and span.attributes.get("voice.event.type") == "session.start" + ) + assert observed_terminations == [SessionTermination.CALLBACK_ERROR] + assert disconnects == [] + assert [message["type"] for message in sent_messages] == ["websocket.accept", "websocket.close"] + assert sent_messages[-1]["code"] == 1011 + assert callback.attributes["error.type"] == "callback_error" + assert connection.attributes["bridge.outcome"] == "callback_error" + assert "private callback detail" not in repr(callback.attributes) + + +@pytest.mark.parametrize("failure_stage", ["span_start", "span_lifecycle", "meter", "logger"]) +def test_telemetry_failure_does_not_change_wire_or_next_connection(monkeypatch, failure_stage): + class FailingTracer: + @staticmethod + def start_span(*_args, **_kwargs): + raise RuntimeError("private tracer failure") + + class FailingSpan: + def __init__(self): + self._context = SpanContext( + trace_id=0x99999999999999999999999999999999, + span_id=0xAAAAAAAAAAAAAAAA, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=TraceState(), + ) + + def get_span_context(self): + return self._context + + @staticmethod + def set_attribute(*_args, **_kwargs): + raise RuntimeError("private attribute failure") + + @staticmethod + def set_status(*_args, **_kwargs): + raise RuntimeError("private status failure") + + @staticmethod + def end(*_args, **_kwargs): + raise RuntimeError("private end failure") + + class LifecycleTracer: + @staticmethod + def start_span(*_args, **_kwargs): + return FailingSpan() + + class FailingInstrument: + @staticmethod + def add(*_args, **_kwargs): + raise RuntimeError("private counter failure") + + @staticmethod + def record(*_args, **_kwargs): + raise RuntimeError("private histogram failure") + + class FailingHandler(logging.Handler): + def emit(self, _record): + raise RuntimeError("private logger failure") + + if failure_stage == "span_start": + monkeypatch.setattr(tracing_module, "_TRACER", FailingTracer()) + elif failure_stage == "span_lifecycle": + monkeypatch.setattr(tracing_module, "_TRACER", LifecycleTracer()) + elif failure_stage == "meter": + monkeypatch.setattr(tracing_module, "_CONNECTION_DURATION", FailingInstrument()) + monkeypatch.setattr(tracing_module, "_TARGET_DURATION", FailingInstrument()) + monkeypatch.setattr(tracing_module, "_PROPAGATION_FAILURES", FailingInstrument()) + + app = VoiceAgentServerHost(configure_observability=None) + + @app.on_session_start + async def on_session_start(session, _event): + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + with turn.activate(): + pass + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + await session.send(SessionReady()) + + handler = FailingHandler() + if failure_stage == "logger": + logging.getLogger("azure.ai.agentserver").addHandler(handler) + try: + for _ in range(2): + with TestClient(app).websocket_connect("/invocations_ws") as websocket: + websocket.send_json(_session_start_frame()) + assert websocket.receive_json()["type"] == "session.ready" + websocket.send_json( + { + "type": "session.end", + "id": "m_end", + "ts": "2026-08-17T00:00:01Z", + "reason": "completed", + } + ) + assert websocket.receive()["type"] == "websocket.close" + finally: + if failure_stage == "logger": + logging.getLogger("azure.ai.agentserver").removeHandler(handler) + + +def test_target_span_end_is_attempted_after_attribute_failure(monkeypatch, metric_reader): + class FailingSpan: + def __init__(self): + self.end_calls = 0 + self.status_calls = 0 + + @staticmethod + def set_attribute(*_args, **_kwargs): + raise RuntimeError("private attribute failure") + + def set_status(self, *_args, **_kwargs): + self.status_calls += 1 + raise RuntimeError("private status failure") + + def end(self, *_args, **_kwargs): + self.end_calls += 1 + + span = FailingSpan() + + class FailingTracer: + @staticmethod + def start_span(*_args, **_kwargs): + return span + + monkeypatch.setattr(turn_module, "_TRACER", FailingTracer()) + before = sum(point.count for point in _metric_points(metric_reader, "gen_ai.invoke_agent.duration")) + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context( + trace.get_tracer("test.connection").start_span("agentserver.connection") + ), + ) + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + with turn.activate(): + pass + + turn.complete(outcome=TargetTurnOutcome.ERROR, output_item_count=0) + turn.complete(outcome=TargetTurnOutcome.ERROR, output_item_count=0) + + after = sum(point.count for point in _metric_points(metric_reader, "gen_ai.invoke_agent.duration")) + assert span.end_calls == 1 + assert span.status_calls == 1 + assert after - before == 1 + assert turn.is_completed + + +def test_connection_finalizer_attempts_status_and_end_after_attribute_failure( + metric_reader, +): + class FailingSpan: + def __init__(self): + self.end_calls = 0 + self.status_calls = 0 + + @staticmethod + def set_attribute(*_args, **_kwargs): + raise RuntimeError("private attribute failure") + + def set_status(self, *_args, **_kwargs): + self.status_calls += 1 + + def end(self, *_args, **_kwargs): + self.end_calls += 1 + + span = FailingSpan() + scope = tracing_module._SpanScope(span=span) # pylint: disable=protected-access + before = sum( + point.count for point in _metric_points(metric_reader, "azure.ai.agentserver.voice.connection.duration") + ) + + scope.complete_connection("transport_error", 1006) + scope.complete_connection("transport_error", 1006) + scope.close() + + after = sum( + point.count for point in _metric_points(metric_reader, "azure.ai.agentserver.voice.connection.duration") + ) + assert span.status_calls == 1 + assert span.end_calls == 1 + assert after - before == 1 + + +def test_target_attach_failure_falls_back_to_connection_parent(monkeypatch, spans): + provider, exporter = spans + tracer = provider.get_tracer("test.connection") + customer_tracer = provider.get_tracer("customer.agent") + connection = tracer.start_span("agentserver.connection") + session = Session._create( # pylint: disable=protected-access + _websocket_with_headers([]), + connection_context=trace.set_span_in_context(connection), + ) + turn = session.start_target_turn(origin=TargetTurnOrigin.USER, input_count=1) + original_attach = turn_module._attach_context # pylint: disable=protected-access + + def fail_target_attach(candidate): + candidate_span = trace.get_current_span(candidate) + if candidate_span is turn._span: # pylint: disable=protected-access + return None + return original_attach(candidate) + + monkeypatch.setattr(turn_module, "_attach_context", fail_target_attach) + with turn.activate(): + with customer_tracer.start_as_current_span("customer.fallback"): + pass + turn.complete(outcome=TargetTurnOutcome.NONE, output_item_count=0) + connection.end() + + customer = _span_by_name(exporter, "customer.fallback") + assert customer.parent is not None + assert customer.parent.span_id == connection.context.span_id diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_transport_findings.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_transport_findings.py index 6132d4b8fade..c536219837b4 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_transport_findings.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/voice/test_voice_transport_findings.py @@ -9,7 +9,6 @@ import sys import pytest -import opentelemetry.propagate as otel_propagate from opentelemetry import baggage, trace from starlette.applications import Starlette from starlette.routing import Host, Mount, WebSocketRoute @@ -17,7 +16,7 @@ from starlette.websockets import WebSocket, WebSocketDisconnect from azure.ai.agentserver.core import get_request_context -from azure.ai.agentserver.invocations.voice import Session, SessionReady, VoiceAgentServerHost +from azure.ai.agentserver.invocations.voice import Session, SessionReady, SessionTermination, VoiceAgentServerHost from azure.ai.agentserver.invocations.voice import _session as session_module from azure.ai.agentserver.invocations.voice import _voice_host as voice_host_module from azure.ai.agentserver.invocations.voice._codec import MAX_FRAME_BYTES @@ -385,7 +384,7 @@ def test_voice_upgrade_preserves_repeated_w3c_headers(): (b"traceparent", f"00-{expected_trace_id}-2222222222222222-01".encode()), (b"tracestate", b"vendor1=value1"), (b"tracestate", b"vendor2=value2"), - (b"baggage", b"tenant.id=tenant-1"), + (b"baggage", b"microsoft.tenant.id=tenant-1"), (b"baggage", b"region=west"), (b"x-request-id", b""), (b"x-request-id", b"request-first"), @@ -399,9 +398,22 @@ def test_voice_upgrade_preserves_repeated_w3c_headers(): assert f"{span_context.trace_id:032x}" == expected_trace_id assert span_context.trace_state.get("vendor1") == "value1" assert span_context.trace_state.get("vendor2") == "value2" - assert baggage.get_baggage("tenant.id", context=context) == "tenant-1" - assert baggage.get_baggage("region", context=context) == "west" - assert baggage.get_baggage("x_request_id", context=context) == "request-first" + assert baggage.get_baggage("microsoft.tenant.id", context=context) == "tenant-1" + assert baggage.get_baggage("region", context=context) is None + assert baggage.get_baggage("x_request_id", context=context) is None + + +@pytest.mark.parametrize( + ("value_lengths", "expected"), + [((253, 254), True), ((254, 254), False)], +) +def test_voice_tracestate_enforces_total_512_byte_limit(value_lengths, expected): + first_length, second_length = value_lengths + tracestate = f"a={'a' * first_length},b={'b' * second_length}".encode("ascii") + assert len(tracestate) == (512 if expected else 513) + headers = [(b"tracestate", tracestate)] + + assert voice_host_module._has_valid_tracestate(headers) is expected # pylint: disable=protected-access def test_voice_upgrade_rejects_duplicate_traceparent(): @@ -420,7 +432,7 @@ def test_voice_upgrade_ignores_context_extraction_failure(monkeypatch): def fail_extract(*_args, **_kwargs): raise RuntimeError("context extraction failed") - monkeypatch.setattr(otel_propagate, "extract", fail_extract) + monkeypatch.setattr(voice_host_module._VOICE_TRACE_PROPAGATOR, "extract", fail_extract) app = VoiceAgentServerHost(configure_observability=None) @app.on_session_start @@ -1842,6 +1854,43 @@ async def send(message): assert _live_voice_send_tasks() == [] +@pytest.mark.asyncio +async def test_receive_side_transport_failure_is_committed_before_cleanup(): + app = VoiceAgentServerHost(configure_observability=None) + observed_terminations = [] + disconnects = [] + inbound_events = [{"type": "websocket.connect"}] + sent_messages = [] + + @app.on_connection_terminating + def on_connection_terminating(session): + observed_terminations.append(session.termination) + + @app.on_disconnect + async def on_disconnect(_session, event): + disconnects.append((event.code, event.reason)) + + async def receive(): + if inbound_events: + return inbound_events.pop(0) + raise OSError("peer receive failed") + + async def send(message): + sent_messages.append(message) + + websocket = _websocket_with_headers([]) + websocket._receive = receive # pylint: disable=protected-access + websocket._send = send # pylint: disable=protected-access + + await asyncio.wait_for(app._ws_endpoint(websocket), timeout=1) # pylint: disable=protected-access + + assert observed_terminations == [SessionTermination.TRANSPORT_ERROR] + assert disconnects == [(1006, None)] + assert [message["type"] for message in sent_messages] == ["websocket.accept"] + assert Session._current(websocket) is None # pylint: disable=protected-access + assert _live_voice_transport_tasks() == [] + + @pytest.mark.asyncio async def test_voice_send_side_peer_loss_wins_later_receive_disconnect(): app = VoiceAgentServerHost(configure_observability=None) @@ -2060,12 +2109,14 @@ async def send(message): disconnect_key = session_module._VOICE_DISCONNECT_EVENT_SCOPE_KEY # pylint: disable=protected-access assert disconnect_key not in websocket.scope - assert [message["type"] for message in sent_messages] == [ - "websocket.accept", - "websocket.send", - "websocket.close", - ] - assert sent_messages[-1]["code"] == 1002 + expected_message_types = ( + ["websocket.accept", "websocket.send"] + if peer_loss + else ["websocket.accept", "websocket.send", "websocket.close"] + ) + assert [message["type"] for message in sent_messages] == expected_message_types + if not peer_loss: + assert sent_messages[-1]["code"] == 1002 assert close_events == [(1002, None)] assert Session._current(websocket) is None # pylint: disable=protected-access with pytest.raises(RuntimeError, match="terminating"):