fix(runtime): flush tracing spans after invocations

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()

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.
 

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()

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)
 # ---------------------------------------------------------------------------