From e15e250fddeabd0c46ac0dab2c962017c278e1b7 Mon Sep 17 00:00:00 2001 From: jariy17 Date: Tue, 27 Oct 2026 16:34:00 +0000 Subject: [PATCH] fix(runtime): flush tracing spans after invocations --- src/bedrock_agentcore/runtime/app.py | 8 +- src/bedrock_agentcore/runtime/tracing.py | 31 +++++++ tests/bedrock_agentcore/runtime/test_app.py | 56 +++++++++++- .../bedrock_agentcore/runtime/test_tracing.py | 88 ++++++++++++++++++- 4 files changed, 180 insertions(+), 3 deletions(-) diff --git a/src/bedrock_agentcore/runtime/app.py b/src/bedrock_agentcore/runtime/app.py index c291b255..efae2faa 100644 --- a/src/bedrock_agentcore/runtime/app.py +++ b/src/bedrock_agentcore/runtime/app.py @@ -49,7 +49,7 @@ PingStatus, is_forwardable_header, ) -from .tracing import _ensure_baggage_processor_registered +from .tracing import _ensure_baggage_processor_registered, _flush_tracer_provider, _flush_tracer_provider_async from .utils import convert_complex_objects # Sentinel so we only parse OTEL_RESOURCE_ATTRIBUTES once per process. @@ -613,6 +613,8 @@ async def _handle_invocation(self, request): duration = time.time() - start_time self.logger.exception("Invocation failed (%.3fs)", duration) return JSONResponse({"error": str(e)}, status_code=500) + finally: + await _flush_tracer_provider_async() def _handle_ping(self, request): try: @@ -894,6 +896,8 @@ async def _stream_with_error_handling(self, generator): "message": "An error occurred during streaming", } yield self._convert_to_sse(error_event) + finally: + await _flush_tracer_provider_async() def _safe_serialize_to_json_string(self, obj): """Safely serialize object directly to JSON string with progressive fallback handling. @@ -949,3 +953,5 @@ def _sync_stream_with_error_handling(self, generator): "message": "An error occurred during streaming", } yield self._convert_to_sse(error_event) + finally: + _flush_tracer_provider() diff --git a/src/bedrock_agentcore/runtime/tracing.py b/src/bedrock_agentcore/runtime/tracing.py index 7c0cf9cb..449c13ee 100644 --- a/src/bedrock_agentcore/runtime/tracing.py +++ b/src/bedrock_agentcore/runtime/tracing.py @@ -22,6 +22,7 @@ provider after the app is constructed. """ +import asyncio import logging import threading from typing import Optional @@ -33,6 +34,7 @@ # Module-level state for provider-tracking auto-registration. _registration_lock = threading.Lock() _registered_on: Optional[object] = None # the TracerProvider instance we last registered on +_flush_lock = threading.Lock() def _ensure_baggage_processor_registered() -> None: @@ -63,6 +65,35 @@ def _ensure_baggage_processor_registered() -> None: logger.debug("Could not register BaggageSpanProcessor", exc_info=True) +def _flush_tracer_provider(timeout_millis: int = 30000) -> None: + """Flush buffered spans before AgentCore Runtime freezes the microVM.""" + if not _flush_lock.acquire(blocking=False): + logger.debug("Tracer provider flush already in progress; skipping duplicate flush") + return + + try: + from opentelemetry import trace + + trace.get_tracer_provider().force_flush(timeout_millis=timeout_millis) + except ImportError: + logger.debug("opentelemetry-api not installed; tracer provider flush skipped") + except Exception: + logger.debug("Could not flush tracer provider", exc_info=True) + finally: + _flush_lock.release() + + +async def _flush_tracer_provider_async(timeout_millis: int = 30000) -> None: + """Flush buffered spans without blocking the event loop or response indefinitely.""" + try: + await asyncio.wait_for( + asyncio.to_thread(_flush_tracer_provider, timeout_millis), + timeout=timeout_millis / 1000, + ) + except asyncio.TimeoutError: + logger.debug("Timed out waiting for tracer provider flush after %d ms", timeout_millis) + + def _get_base_class() -> type: """Return the OTel SDK SpanProcessor base if available, otherwise object. diff --git a/tests/bedrock_agentcore/runtime/test_app.py b/tests/bedrock_agentcore/runtime/test_app.py index 60169a5a..7ce404c6 100644 --- a/tests/bedrock_agentcore/runtime/test_app.py +++ b/tests/bedrock_agentcore/runtime/test_app.py @@ -8,7 +8,7 @@ import time from datetime import datetime from decimal import Decimal -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from starlette.testclient import TestClient @@ -151,6 +151,60 @@ def handler(payload): assert response.status_code == 500 assert response.json() == {"error": "Test error"} + def test_invocation_flushes_tracer_provider(self): + """Buffered spans are flushed before a non-streaming response returns.""" + bedrock_agentcore = BedrockAgentCoreApp() + + @bedrock_agentcore.entrypoint + def handler(payload): + return {"ok": True} + + with patch("bedrock_agentcore.runtime.app._flush_tracer_provider_async", new_callable=AsyncMock) as mock_flush: + response = TestClient(bedrock_agentcore).post("/invocations", json={}) + + assert response.status_code == 200 + mock_flush.assert_called_once_with() + + def test_sync_stream_flushes_after_generator_finishes(self): + """Sync streams flush once when prepared and again after their final chunk.""" + bedrock_agentcore = BedrockAgentCoreApp() + + @bedrock_agentcore.entrypoint + def handler(payload): + yield {"ok": True} + + with ( + patch( + "bedrock_agentcore.runtime.app._flush_tracer_provider_async", new_callable=AsyncMock + ) as mock_async_flush, + patch("bedrock_agentcore.runtime.app._flush_tracer_provider") as mock_sync_flush, + ): + response = TestClient(bedrock_agentcore).post("/invocations", json={}) + + assert response.status_code == 200 + mock_async_flush.assert_awaited_once_with() + mock_sync_flush.assert_called_once_with() + + def test_async_stream_flushes_after_generator_finishes(self): + """Async streams flush once when prepared and again after their final chunk.""" + bedrock_agentcore = BedrockAgentCoreApp() + + @bedrock_agentcore.entrypoint + async def handler(payload): + yield {"ok": True} + + with ( + patch( + "bedrock_agentcore.runtime.app._flush_tracer_provider_async", new_callable=AsyncMock + ) as mock_async_flush, + patch("bedrock_agentcore.runtime.app._flush_tracer_provider") as mock_sync_flush, + ): + response = TestClient(bedrock_agentcore).post("/invocations", json={}) + + assert response.status_code == 200 + mock_async_flush.assert_awaited_once_with() + mock_sync_flush.assert_called_once_with() + def test_async_handler_without_context(self): """Test async handler without context parameter.""" bedrock_agentcore = BedrockAgentCoreApp() diff --git a/tests/bedrock_agentcore/runtime/test_tracing.py b/tests/bedrock_agentcore/runtime/test_tracing.py index 020b4595..21f79867 100644 --- a/tests/bedrock_agentcore/runtime/test_tracing.py +++ b/tests/bedrock_agentcore/runtime/test_tracing.py @@ -5,7 +5,12 @@ import pytest -from bedrock_agentcore.runtime.tracing import BaggageSpanProcessor, _ensure_baggage_processor_registered +from bedrock_agentcore.runtime.tracing import ( + BaggageSpanProcessor, + _ensure_baggage_processor_registered, + _flush_tracer_provider, + _flush_tracer_provider_async, +) # --------------------------------------------------------------------------- # BaggageSpanProcessor.on_start @@ -254,6 +259,87 @@ def _call(): assert mock_provider.add_span_processor.call_count == 1 +# --------------------------------------------------------------------------- +# _flush_tracer_provider +# --------------------------------------------------------------------------- + + +class TestFlushTracerProvider: + def test_flushes_current_provider(self): + mock_provider = MagicMock() + + with patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider): + _flush_tracer_provider(timeout_millis=123) + + mock_provider.force_flush.assert_called_once_with(timeout_millis=123) + + def test_provider_error_is_swallowed(self): + mock_provider = MagicMock() + mock_provider.force_flush.side_effect = RuntimeError("exporter unavailable") + + with patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider): + _flush_tracer_provider() + + def test_duplicate_flush_is_skipped_while_one_is_in_progress(self): + started = threading.Event() + release = threading.Event() + mock_provider = MagicMock() + + def blocking_flush(**_): + started.set() + release.wait(timeout=1) + + mock_provider.force_flush.side_effect = blocking_flush + + with patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider): + flush_thread = threading.Thread(target=_flush_tracer_provider) + flush_thread.start() + try: + assert started.wait(timeout=1) + _flush_tracer_provider() + finally: + release.set() + flush_thread.join(timeout=1) + + mock_provider.force_flush.assert_called_once() + + async def test_async_flush_runs_on_worker_thread(self): + event_loop_thread = threading.get_ident() + worker_threads = [] + mock_provider = MagicMock() + mock_provider.force_flush.side_effect = lambda **_: worker_threads.append(threading.get_ident()) + + with patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider): + await _flush_tracer_provider_async(timeout_millis=1000) + + assert worker_threads + assert worker_threads[0] != event_loop_thread + + async def test_async_flush_bounds_response_wait(self): + started = threading.Event() + release = threading.Event() + finished = threading.Event() + mock_provider = MagicMock() + + def blocking_flush(**_): + started.set() + try: + release.wait(timeout=1) + finally: + finished.set() + + mock_provider.force_flush.side_effect = blocking_flush + + try: + with patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider): + await _flush_tracer_provider_async(timeout_millis=20) + assert started.is_set() + assert not finished.is_set() + finally: + release.set() + finished.wait(timeout=1) + + # --------------------------------------------------------------------------- # End-to-end: baggage header → ContextVar → span attribute (via app.py) # ---------------------------------------------------------------------------