From 73a641281b3b74cbc9f4d0008ee756d7f28953e9 Mon Sep 17 00:00:00 2001 From: Harsheet Shah Date: Thu, 10 Sep 2026 11:20:31 +0530 Subject: [PATCH 1/2] [agentserver] Make per-request span flush non-blocking The Responses endpoint flushed spans synchronously in the request `finally` block via `flush_spans()`, which runs `TracerProvider.force_flush` inline. On the async handler this blocks the asyncio event loop until the exporter drains, serialising concurrent requests behind a single export (head-of-line blocking) and adding the export time to every response. Add two helpers to azure-ai-agentserver-core: - `flush_spans_async`: offloads the blocking force_flush to a worker thread so it never stalls the event loop (same durability guarantee). - `schedule_flush_spans`: fire-and-forget flush that returns immediately so the response is not delayed (requires a platform drain window before freeze). The Responses hot-path flush now dispatches on `AGENTSERVER_FLUSH_MODE`: `async` (default, off the event loop), `background` (respond first), or `sync` (legacy). Default `async` removes event-loop head-of-line blocking with no change to durability. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 93f793d2-1b86-4677-908c-3722d4fef290 --- .../azure-ai-agentserver-core/CHANGELOG.md | 6 ++ .../azure/ai/agentserver/core/__init__.py | 4 ++ .../azure/ai/agentserver/core/_tracing.py | 61 +++++++++++++++++++ .../CHANGELOG.md | 10 +++ .../responses/hosting/_endpoint_handler.py | 28 +++++++-- 5 files changed, 104 insertions(+), 5 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index 2c4a1c06b069..16962094b335 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -4,6 +4,12 @@ ### Features Added +- Added `flush_spans_async`, a non-blocking variant of `flush_spans` that + offloads the blocking `TracerProvider.force_flush` to a worker thread so it + does not stall the asyncio event loop, and `schedule_flush_spans`, a + fire-and-forget helper that flushes in the background without delaying the + caller. + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/__init__.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/__init__.py index 10cb683493c7..f5e625daf929 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/__init__.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/__init__.py @@ -30,7 +30,9 @@ detach_context, end_span, flush_spans, + flush_spans_async, record_error, + schedule_flush_spans, set_current_span, trace_stream, ) @@ -52,6 +54,8 @@ "end_span", "experimental", "flush_spans", + "flush_spans_async", + "schedule_flush_spans", "get_request_context", "record_error", "read_request_id", diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py index 40753926997d..22b791e4645b 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py @@ -33,6 +33,7 @@ OpenTelemetry is a required dependency — these functions always create real spans. Azure Monitor export is optional (auto-configured by the distro). """ +import asyncio # pylint: disable=do-not-import-asyncio from collections.abc import AsyncIterable, AsyncIterator # pylint: disable=import-error from contextlib import contextmanager, nullcontext import logging @@ -583,6 +584,66 @@ def flush_spans(timeout_millis: int = 5000) -> None: logger.debug("TracerProvider.force_flush() failed", exc_info=True) +# Strong references to in-flight background flush tasks so they are not garbage +# collected before completing (asyncio only holds weak references to tasks). +_BG_FLUSH_TASKS: set = set() + + +async def flush_spans_async(timeout_millis: int = 5000) -> None: + """Non-blocking variant of :func:`flush_spans`. + + ``TracerProvider.force_flush`` blocks the calling thread until the exporter + drains its queue. On the request hot path -- which runs inside an ``async`` + handler -- that blocks the asyncio event loop, serialising every concurrent + request behind a single export (head-of-line blocking). Offload the + blocking call to the default thread pool so the event loop stays free to + send the response and service other requests concurrently. + + No-op when the OTel SDK is not installed or the provider does not support + ``force_flush``. + + :param timeout_millis: Maximum time to wait for the flush, in milliseconds. + Defaults to 5000 (5 seconds). + :type timeout_millis: int + """ + provider = trace.get_tracer_provider() + flush = getattr(provider, "force_flush", None) + if flush is None: + return + try: + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, flush, timeout_millis) + except Exception: # pylint: disable=broad-exception-caught + logger.debug("TracerProvider.force_flush() (async) failed", exc_info=True) + + +def schedule_flush_spans(timeout_millis: int = 5000) -> None: + """Schedule a span flush as a background task and return immediately. + + Unlike :func:`flush_spans` / :func:`flush_spans_async`, this does not delay + the caller (i.e. the HTTP response) by the export duration. A strong + reference to the task is retained until it completes so it is not garbage + collected. Falls back to a synchronous flush when no event loop is running. + + .. note:: + Only safe when the hosting platform guarantees a brief drain window + before it suspends/freezes the process after sending a response; + otherwise the final request's spans may be lost. + + :param timeout_millis: Maximum time to wait for the flush, in milliseconds. + Defaults to 5000 (5 seconds). + :type timeout_millis: int + """ + try: + loop = asyncio.get_running_loop() + except RuntimeError: + flush_spans(timeout_millis) + return + task = loop.create_task(flush_spans_async(timeout_millis)) + _BG_FLUSH_TASKS.add(task) + task.add_done_callback(_BG_FLUSH_TASKS.discard) + + def record_error(span: Any, exc: BaseException) -> None: """Record an exception and ERROR status on a span. diff --git a/sdk/agentserver/azure-ai-agentserver-responses/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-responses/CHANGELOG.md index 8f9837aa7d44..0bd24288f42a 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-responses/CHANGELOG.md @@ -4,6 +4,16 @@ ### Bugs Fixed +- The per-request span flush in the Responses endpoint no longer blocks the + asyncio event loop. The synchronous `flush_spans()` call in the request + `finally` block ran `TracerProvider.force_flush` inline, which blocks the + event loop until the exporter drains and serialises concurrent requests + behind a single export (head-of-line blocking). It now awaits the + non-blocking `flush_spans_async()` by default. A new `AGENTSERVER_FLUSH_MODE` + environment variable selects the strategy: `async` (default, off the event + loop), `background` (return the response first, flush in the background -- + requires a platform drain window), or `sync` (legacy blocking behaviour). + - Scoped durable multi-turn task IDs with `FOUNDRY_AGENT_SESSION_GUID` when available, preventing recreated same-name sessions from colliding with task tombstones. Existing pre-rollout active chains remain resumable through a diff --git a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py index d3b6ea436155..2b891a00ad7d 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py @@ -13,6 +13,7 @@ import asyncio # pylint: disable=do-not-import-asyncio import contextvars import logging +import os import threading from typing import TYPE_CHECKING, Any, cast @@ -24,7 +25,9 @@ from azure.ai.agentserver.core import ( # pylint: disable=import-error,no-name-in-module FoundryAgentRequestContext, flush_spans, + flush_spans_async, reset_request_context, + schedule_flush_spans, set_request_context, ) from azure.ai.agentserver.core.tasks import ( @@ -949,11 +952,26 @@ async def _iter_with_context(): # type: ignore[return] _conversation_id_var.reset(cid_token) _streaming_var.reset(str_token) reset_request_context(platform_ctx_token) - # Flush pending spans before the response is sent. - # BatchSpanProcessor exports on a timer; in hosted sandboxes - # the platform may freeze the process after the HTTP response, - # losing any buffered spans (e.g. LangGraph per-node spans). - flush_spans() + # Flush pending spans before the process may be frozen. + # ``force_flush`` blocks the calling thread until the exporter + # drains; doing that inline on this ``async`` handler blocks the + # event loop and serialises concurrent requests behind one export. + # AGENTSERVER_FLUSH_MODE selects the strategy: + # "async" (default) -> await flush_spans_async(): off the event + # loop; same durability, no head-of-line + # blocking under concurrency. + # "background" -> schedule_flush_spans(): return the response + # first, flush in the background. Lowest + # latency, but requires the platform to grant + # a brief drain window before freezing. + # "sync" -> flush_spans(): legacy blocking behaviour. + _flush_mode = os.environ.get("AGENTSERVER_FLUSH_MODE", "async").lower() + if _flush_mode == "sync": + flush_spans() + elif _flush_mode == "background": + schedule_flush_spans() + else: + await flush_spans_async() try: _otel_context.detach(baggage_token) except ValueError: From 8fe884abde46ea1c52c12160380f1f80d25c12c8 Mon Sep 17 00:00:00 2001 From: Harsheet Shah Date: Thu, 10 Sep 2026 12:03:22 +0530 Subject: [PATCH 2/2] Address review: bound background flush, raise dep floor, add tests - Coalesce background flushing: at most one flush task runs at a time; concurrent requests collapse into a single follow-up pass instead of spawning a retained task per request (bounded under load). - Raise azure-ai-agentserver-core floor to >=2.2.0b2 in responses, since the handler now imports flush_spans_async/schedule_flush_spans (added in b2). - Record flush_spans_async/schedule_flush_spans in core api.md. - Extract _flush_spans_for_mode dispatch helper (case/whitespace-insensitive, unknown values fall back to the async default). - Add tests: flush_spans_async non-blocking/timeout/exception/no-op; schedule_flush_spans coalescing + sync fallback; handler flush-mode dispatch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 93f793d2-1b86-4677-908c-3722d4fef290 --- .../azure-ai-agentserver-core/CHANGELOG.md | 4 +- .../azure-ai-agentserver-core/api.md | 6 + .../azure/ai/agentserver/core/_tracing.py | 50 ++++++-- .../tests/test_tracing.py | 120 ++++++++++++++++++ .../responses/hosting/_endpoint_handler.py | 65 +++++++--- .../pyproject.toml | 4 +- .../tests/test_flush_dispatch.py | 68 ++++++++++ 7 files changed, 285 insertions(+), 32 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-responses/tests/test_flush_dispatch.py diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index 16962094b335..9138c3f9037d 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -8,7 +8,9 @@ offloads the blocking `TracerProvider.force_flush` to a worker thread so it does not stall the asyncio event loop, and `schedule_flush_spans`, a fire-and-forget helper that flushes in the background without delaying the - caller. + caller. Background flushes are coalesced: at most one runs at a time and + concurrent requests collapse into a single follow-up flush, so the work does + not grow with the request rate. ### Breaking Changes diff --git a/sdk/agentserver/azure-ai-agentserver-core/api.md b/sdk/agentserver/azure-ai-agentserver-core/api.md index 20445d98573e..fce3fa0df5d1 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/api.md +++ b/sdk/agentserver/azure-ai-agentserver-core/api.md @@ -56,12 +56,18 @@ namespace azure.ai.agentserver.core def azure.ai.agentserver.core.resolve_state_subdir(name: str) -> Path: ... + def azure.ai.agentserver.core.schedule_flush_spans(timeout_millis: int = 5000) -> None: ... + + def azure.ai.agentserver.core.set_current_span(span: Any) -> Any: ... def azure.ai.agentserver.core.set_request_context(context: FoundryAgentRequestContext) -> Token[FoundryAgentRequestContext]: ... + async def azure.ai.agentserver.core.flush_spans_async:async(timeout_millis: int = 5000) -> None: ... + + async def azure.ai.agentserver.core.trace_stream:async(iterator: AsyncIterable[StreamContent], span: Any) -> AsyncIterator[StreamContent]: ... diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py index 22b791e4645b..53bc1043c204 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py @@ -584,9 +584,18 @@ def flush_spans(timeout_millis: int = 5000) -> None: logger.debug("TracerProvider.force_flush() failed", exc_info=True) -# Strong references to in-flight background flush tasks so they are not garbage -# collected before completing (asyncio only holds weak references to tasks). -_BG_FLUSH_TASKS: set = set() +# A single coalesced background flush runs at a time. ``force_flush`` drains +# the provider *globally*, so concurrent per-request flushes would be redundant +# work. Instead of spawning one task per request (which lets ``_bg_flush_task`` +# / the executor queue grow without bound under load), requests that arrive +# while a flush is in flight set ``_bg_flush_pending``; the running task then +# performs exactly one follow-up flush afterwards to capture spans produced +# during the active flush. This bounds in-flight background work to a single +# task regardless of request rate. The module-level reference also keeps the +# task alive (asyncio only holds weak references to tasks). Access is confined +# to the event-loop thread, so no lock is required. +_bg_flush_task: "Optional[asyncio.Task[None]]" = None +_bg_flush_pending: bool = False async def flush_spans_async(timeout_millis: int = 5000) -> None: @@ -617,13 +626,30 @@ async def flush_spans_async(timeout_millis: int = 5000) -> None: logger.debug("TracerProvider.force_flush() (async) failed", exc_info=True) +async def _coalesced_flush(timeout_millis: int) -> None: + """Run a background flush, then one more pass per pending coalesced request. + + Because ``force_flush`` drains the provider globally, a single trailing + flush captures the spans of every request that arrived while a flush was + already running -- no need for a task (or export) per request. + """ + global _bg_flush_pending # pylint: disable=global-statement + await flush_spans_async(timeout_millis) + while _bg_flush_pending: + _bg_flush_pending = False + await flush_spans_async(timeout_millis) + + def schedule_flush_spans(timeout_millis: int = 5000) -> None: - """Schedule a span flush as a background task and return immediately. + """Schedule a coalesced background span flush and return immediately. Unlike :func:`flush_spans` / :func:`flush_spans_async`, this does not delay - the caller (i.e. the HTTP response) by the export duration. A strong - reference to the task is retained until it completes so it is not garbage - collected. Falls back to a synchronous flush when no event loop is running. + the caller (i.e. the HTTP response) by the export duration. At most one + background flush task runs at a time: calls made while a flush is in flight + are coalesced into a single follow-up flush rather than spawning a task per + request, so neither the retained task reference nor the executor queue grows + with the request rate. Falls back to a synchronous flush when no event loop + is running. .. note:: Only safe when the hosting platform guarantees a brief drain window @@ -634,14 +660,18 @@ def schedule_flush_spans(timeout_millis: int = 5000) -> None: Defaults to 5000 (5 seconds). :type timeout_millis: int """ + global _bg_flush_task, _bg_flush_pending # pylint: disable=global-statement try: loop = asyncio.get_running_loop() except RuntimeError: flush_spans(timeout_millis) return - task = loop.create_task(flush_spans_async(timeout_millis)) - _BG_FLUSH_TASKS.add(task) - task.add_done_callback(_BG_FLUSH_TASKS.discard) + if _bg_flush_task is not None and not _bg_flush_task.done(): + # A flush is already draining the provider globally; record that more + # spans arrived so the running task performs one more pass afterwards. + _bg_flush_pending = True + return + _bg_flush_task = loop.create_task(_coalesced_flush(timeout_millis)) def record_error(span: Any, exc: BaseException) -> None: diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py index c6dad95025ab..59a753439250 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py @@ -3,7 +3,9 @@ # --------------------------------------------------------- """Tests for tracing configuration — not invocation spans (those live in the invocations package).""" +import asyncio import os +import pytest from functools import partial from threading import Event, Thread from typing import Any, Optional @@ -25,6 +27,7 @@ _BaggageLogRecordProcessor, _FoundryEnrichmentSpanProcessor, ) +from azure.ai.agentserver.core import _tracing class _CollectorExporter(SpanExporter): @@ -938,3 +941,120 @@ def test_does_not_overwrite_existing_log_attributes(self) -> None: assert attrs["gen_ai.agent.name"] == "existing-name" assert attrs["gen_ai.agent.version"] == "0.0.1" assert attrs["microsoft.session.id"] == "existing-session" + + +# ------------------------------------------------------------------ # +# flush_spans_async / schedule_flush_spans (non-blocking flush) +# ------------------------------------------------------------------ # + + +class _BlockingFlushProvider: + """Fake TracerProvider whose force_flush blocks until released. + + Records each timeout it is called with so tests can assert pass-through + and coalescing behaviour. + """ + + def __init__(self, block_first: bool = False) -> None: + self.calls: list = [] + self._block_first = block_first + self.started = Event() + self.release = Event() + + def force_flush(self, timeout_millis: int = 30000) -> bool: + first = len(self.calls) == 0 + self.calls.append(timeout_millis) + if first: + self.started.set() + if self._block_first: + self.release.wait(5) + return True + + +async def _wait_for(flag: Event, *, timeout_s: float = 5.0) -> None: + """Yield to the event loop until *flag* is set (proves the loop is free).""" + deadline = asyncio.get_running_loop().time() + timeout_s + while not flag.is_set(): + assert asyncio.get_running_loop().time() < deadline, "flag never set" + await asyncio.sleep(0.001) + + +class TestFlushSpansAsync: + """`flush_spans_async` offloads the blocking export off the event loop.""" + + @pytest.mark.asyncio + async def test_does_not_block_event_loop_and_passes_timeout(self) -> None: + provider = _BlockingFlushProvider(block_first=True) + with mock.patch.object(_tracing.trace, "get_tracer_provider", return_value=provider): + flush_task = asyncio.create_task(_tracing.flush_spans_async(1234)) + # If force_flush ran inline it would block here; instead this + # coroutine keeps running while the export blocks in a worker thread. + await _wait_for(provider.started) + assert not flush_task.done(), "flush should still be draining off-loop" + provider.release.set() + await flush_task + assert provider.calls == [1234] # timeout argument forwarded + + @pytest.mark.asyncio + async def test_swallows_exceptions(self) -> None: + class _BoomProvider: + def force_flush(self, timeout_millis: int = 30000) -> bool: + raise RuntimeError("boom") + + with mock.patch.object(_tracing.trace, "get_tracer_provider", return_value=_BoomProvider()): + await _tracing.flush_spans_async() # must not raise + + @pytest.mark.asyncio + async def test_noop_when_provider_has_no_force_flush(self) -> None: + class _NoFlushProvider: + pass + + with mock.patch.object(_tracing.trace, "get_tracer_provider", return_value=_NoFlushProvider()): + await _tracing.flush_spans_async() # must not raise + + +class TestScheduleFlushSpans: + """`schedule_flush_spans` runs a single coalesced background flush.""" + + def setup_method(self) -> None: + _tracing._bg_flush_task = None + _tracing._bg_flush_pending = False + + def teardown_method(self) -> None: + _tracing._bg_flush_task = None + _tracing._bg_flush_pending = False + + def test_falls_back_to_sync_without_running_loop(self) -> None: + with mock.patch.object(_tracing, "flush_spans") as m_sync: + _tracing.schedule_flush_spans(777) + m_sync.assert_called_once_with(777) + assert _tracing._bg_flush_task is None + + @pytest.mark.asyncio + async def test_coalesces_concurrent_requests_into_one_followup(self) -> None: + provider = _BlockingFlushProvider(block_first=True) + with mock.patch.object(_tracing.trace, "get_tracer_provider", return_value=provider): + # First call spawns exactly one background flush task. + _tracing.schedule_flush_spans() + task = _tracing._bg_flush_task + assert task is not None + await _wait_for(provider.started) + # 50 requests arrive while the flush is in flight. They must NOT + # each spawn a task (bounded); they coalesce into one pending pass. + for _ in range(50): + _tracing.schedule_flush_spans() + assert _tracing._bg_flush_task is task + assert _tracing._bg_flush_pending is True + provider.release.set() + await task + # Initial flush + one coalesced follow-up == 2 exports, not 51. + assert len(provider.calls) == 2 + assert _tracing._bg_flush_pending is False + + @pytest.mark.asyncio + async def test_single_request_flushes_once(self) -> None: + provider = _BlockingFlushProvider(block_first=False) + with mock.patch.object(_tracing.trace, "get_tracer_provider", return_value=provider): + _tracing.schedule_flush_spans() + await _tracing._bg_flush_task + assert len(provider.calls) == 1 diff --git a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py index 2b891a00ad7d..aaa44ddc2ff8 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py @@ -190,6 +190,46 @@ def _get_scope_request_id(request: Request) -> str | None: _streaming_var: contextvars.ContextVar[str] = contextvars.ContextVar("Streaming", default="") +_FLUSH_MODE_ENV = "AGENTSERVER_FLUSH_MODE" +_DEFAULT_FLUSH_MODE = "async" + + +async def _flush_spans_for_mode(mode: str) -> None: + """Dispatch span flushing according to *mode* (see ``AGENTSERVER_FLUSH_MODE``). + + ``force_flush`` blocks the calling thread until the exporter drains; doing + that inline on this ``async`` handler blocks the event loop and serialises + concurrent requests behind one export. The mode selects the strategy: + + * ``"async"`` (default) -> :func:`flush_spans_async`: off the event loop; + same durability, no head-of-line blocking under concurrency. + * ``"background"`` -> :func:`schedule_flush_spans`: return the response + first and flush in the background (lowest latency, but needs the platform + to grant a brief drain window before freezing). + * ``"sync"`` -> :func:`flush_spans`: legacy blocking behaviour. + + Any unrecognised value falls back to the ``"async"`` default (fail safe: + never silently drop telemetry). + + :param mode: The flush mode; matched case-insensitively. + :type mode: str + """ + normalized = (mode or "").strip().lower() + if normalized == "sync": + flush_spans() + elif normalized == "background": + schedule_flush_spans() + else: + if normalized and normalized != _DEFAULT_FLUSH_MODE: + logger.warning( + "Unrecognised %s=%r; falling back to %r flush mode.", + _FLUSH_MODE_ENV, + mode, + _DEFAULT_FLUSH_MODE, + ) + await flush_spans_async() + + class _ResponseLogFilter(logging.Filter): """Attach response-scope IDs to every log record from context vars. @@ -953,25 +993,12 @@ async def _iter_with_context(): # type: ignore[return] _streaming_var.reset(str_token) reset_request_context(platform_ctx_token) # Flush pending spans before the process may be frozen. - # ``force_flush`` blocks the calling thread until the exporter - # drains; doing that inline on this ``async`` handler blocks the - # event loop and serialises concurrent requests behind one export. - # AGENTSERVER_FLUSH_MODE selects the strategy: - # "async" (default) -> await flush_spans_async(): off the event - # loop; same durability, no head-of-line - # blocking under concurrency. - # "background" -> schedule_flush_spans(): return the response - # first, flush in the background. Lowest - # latency, but requires the platform to grant - # a brief drain window before freezing. - # "sync" -> flush_spans(): legacy blocking behaviour. - _flush_mode = os.environ.get("AGENTSERVER_FLUSH_MODE", "async").lower() - if _flush_mode == "sync": - flush_spans() - elif _flush_mode == "background": - schedule_flush_spans() - else: - await flush_spans_async() + # ``AGENTSERVER_FLUSH_MODE`` selects the strategy (see + # ``_flush_spans_for_mode``); the default keeps the flush off the + # event loop without dropping telemetry. + await _flush_spans_for_mode( + os.environ.get(_FLUSH_MODE_ENV, _DEFAULT_FLUSH_MODE) + ) try: _otel_context.detach(baggage_token) except ValueError: diff --git a/sdk/agentserver/azure-ai-agentserver-responses/pyproject.toml b/sdk/agentserver/azure-ai-agentserver-responses/pyproject.toml index d5eafa3ae3cd..de331ffee57c 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/pyproject.toml +++ b/sdk/agentserver/azure-ai-agentserver-responses/pyproject.toml @@ -18,7 +18,7 @@ classifiers = [ "Programming Language :: Python :: 3.14", ] dependencies = [ - "azure-ai-agentserver-core>=2.2.0b1,<2.3.0", + "azure-ai-agentserver-core>=2.2.0b2,<2.3.0", "azure-core>=1.37.0,<2.0.0", "isodate>=0.6.1,<1.0.0", "aiohttp>=3.10.0,<4.0.0a0", @@ -76,5 +76,5 @@ in_bundle = false [tool.azure-sdk-build] verifytypes = false latestdependency = false -# azure-ai-agentserver-core>=2.2.0b1 is not yet on PyPI +# azure-ai-agentserver-core>=2.2.0b2 is not yet on PyPI mindependency = false diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/test_flush_dispatch.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/test_flush_dispatch.py new file mode 100644 index 000000000000..e879b0307437 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-responses/tests/test_flush_dispatch.py @@ -0,0 +1,68 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Tests for AGENTSERVER_FLUSH_MODE dispatch on the response hot path.""" + +from unittest import mock + +import pytest + +from azure.ai.agentserver.responses.hosting import _endpoint_handler as eh + + +_MODE_TO_HELPER = { + "flush_spans": "sync flush", + "schedule_flush_spans": "background flush", + "flush_spans_async": "async flush", +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode, expected_helper", + [ + # sync + ("sync", "flush_spans"), + ("SYNC", "flush_spans"), + (" sync ", "flush_spans"), + # background + ("background", "schedule_flush_spans"), + ("Background", "schedule_flush_spans"), + # async (explicit default) + ("async", "flush_spans_async"), + ("ASYNC", "flush_spans_async"), + # fail-safe fallback to async for empty / unknown values + ("", "flush_spans_async"), + ("bogus", "flush_spans_async"), + ("backgroundx", "flush_spans_async"), + ], +) +async def test_flush_mode_dispatch(mode: str, expected_helper: str) -> None: + """Each mode routes to exactly one helper; case/whitespace-insensitive.""" + with mock.patch.object(eh, "flush_spans") as m_sync, mock.patch.object( + eh, "schedule_flush_spans" + ) as m_bg, mock.patch.object( + eh, "flush_spans_async", new_callable=mock.AsyncMock + ) as m_async: + await eh._flush_spans_for_mode(mode) + + called = { + "flush_spans": m_sync.called, + "schedule_flush_spans": m_bg.called, + "flush_spans_async": m_async.called, + } + assert called[expected_helper] is True, f"{expected_helper} not called for {mode!r}" + for helper, was_called in called.items(): + if helper != expected_helper: + assert was_called is False, f"{helper} unexpectedly called for {mode!r}" + + +@pytest.mark.asyncio +async def test_flush_mode_dispatch_awaits_async_helper() -> None: + """The default async path is actually awaited (not fire-and-forget).""" + with mock.patch.object(eh, "flush_spans"), mock.patch.object( + eh, "schedule_flush_spans" + ), mock.patch.object( + eh, "flush_spans_async", new_callable=mock.AsyncMock + ) as m_async: + await eh._flush_spans_for_mode("async") + m_async.assert_awaited_once()