fix(runtime): bound the tracer flush wait on streaming responses

src/bedrock_agentcore/runtime/app.py

@@ -542,6 +542,8 @@ async def _handle_invocation(self, request):
         request_context = self._build_request_context(request)
 
         start_time = time.time()
+        # Streaming responses flush when the stream ends; flushing here as well would only delay the headers.
+        streaming = False
 
         try:
             payload = await request.json()
@@ -578,9 +580,11 @@ async def _handle_invocation(self, request):
             duration = time.time() - start_time
             if inspect.isgenerator(result):
                 self.logger.info("Returning streaming response (generator) (%.3fs)", duration)
+                streaming = True
                 return StreamingResponse(self._sync_stream_with_error_handling(result), media_type="text/event-stream")
             elif inspect.isasyncgen(result):
                 self.logger.info("Returning streaming response (async generator) (%.3fs)", duration)
+                streaming = True
                 return StreamingResponse(self._stream_with_error_handling(result), media_type="text/event-stream")
 
             # If handler returned a Starlette Response directly, pass it through.
@@ -614,7 +618,8 @@ async def _handle_invocation(self, request):
             self.logger.exception("Invocation failed (%.3fs)", duration)
             return JSONResponse({"error": str(e)}, status_code=500)
         finally:
-            await _flush_tracer_provider_async()
+            if not streaming:
+                await _flush_tracer_provider_async()
 
     def _handle_ping(self, request):
         try:

src/bedrock_agentcore/runtime/tracing.py

@@ -35,6 +35,11 @@
 _registration_lock = threading.Lock()
 _registered_on: Optional[object] = None  # the TracerProvider instance we last registered on
 _flush_lock = threading.Lock()
+# How long a response waits for buffered spans to export. A healthy OTLP export takes tens of milliseconds;
+# 5 s leaves wide headroom, matches the export timeout AgentCore Runtime intends to set
+# (OTEL_EXPORTER_OTLP_TIMEOUT=5000 ms), and keeps a stuck exporter from delaying the end of a response
+# by more than a few seconds. Spans the flush misses stay queued for the background exporter.
+_FLUSH_TIMEOUT_MILLIS = 5000
 
 
 def _ensure_baggage_processor_registered() -> None:
@@ -65,12 +70,36 @@ 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."""
+def _flush_tracer_provider(timeout_millis: int = _FLUSH_TIMEOUT_MILLIS) -> None:
+    """Flush buffered spans before AgentCore Runtime freezes the microVM, waiting at most ``timeout_millis``.
+
+    ``force_flush`` cannot be trusted to honor its timeout (OTel SDK 1.45's ``BatchProcessor`` ignores it and
+    exports until the queue drains, which never ends if the OTLP endpoint is unreachable), so it runs on a
+    daemon thread that the caller stops waiting for. A flush that overruns keeps ``_flush_lock`` until it
+    returns, so later calls skip instead of piling up threads behind it.
+    """
     if not _flush_lock.acquire(blocking=False):
         logger.debug("Tracer provider flush already in progress; skipping duplicate flush")
         return
 
+    try:
+        flush_thread = threading.Thread(
+            target=_force_flush_and_release, args=(timeout_millis,), name="agentcore-trace-flush", daemon=True
+        )
+        flush_thread.start()
+    except Exception:
+        _flush_lock.release()
+        logger.debug("Could not start tracer provider flush", exc_info=True)
+        return
+
+    flush_thread.join(timeout_millis / 1000)
+    if flush_thread.is_alive():
+        logger.warning(
+            "Tracer provider flush did not finish within %d ms; continuing without waiting for it", timeout_millis
+        )
+
+
+def _force_flush_and_release(timeout_millis: int) -> None:
     try:
         from opentelemetry import trace
 
@@ -83,7 +112,7 @@ def _flush_tracer_provider(timeout_millis: int = 30000) -> None:
         _flush_lock.release()
 
 
-async def _flush_tracer_provider_async(timeout_millis: int = 30000) -> None:
+async def _flush_tracer_provider_async(timeout_millis: int = _FLUSH_TIMEOUT_MILLIS) -> None:
     """Flush buffered spans without blocking the event loop or response indefinitely."""
     try:
         await asyncio.wait_for(

tests/bedrock_agentcore/runtime/test_app.py

@@ -2,6 +2,7 @@
 import concurrent.futures
 import contextlib
 import contextvars
+import functools
 import json
 import os
 import threading
@@ -14,13 +15,33 @@
 from starlette.testclient import TestClient
 
 import bedrock_agentcore.runtime.app as _app_module
+import bedrock_agentcore.runtime.tracing as _tracing_module
 from bedrock_agentcore.config_bundle.bundle import ConfigBundleRef
 from bedrock_agentcore.runtime import BedrockAgentCoreApp
 from bedrock_agentcore.runtime.app import _parse_runtime_arn
 from bedrock_agentcore.runtime.context import BedrockAgentCoreContext
 from bedrock_agentcore.runtime.models import is_forwardable_header
 
 
+class _TimeoutIgnoringTracerProvider:
+    """Tracer provider whose ``force_flush`` ignores ``timeout_millis``, like OTel SDK 1.45's BatchProcessor.
+
+    With nothing queued, a flush returns at once. Once the handler has produced spans, a flush blocks the way
+    an export to an unreachable OTLP endpoint does, until ``release`` is set.
+    """
+
+    def __init__(self):
+        self.spans_pending = False
+        self.blocked_flushes = 0
+        self.release = threading.Event()
+
+    def force_flush(self, timeout_millis=30000):
+        if self.spans_pending:
+            self.blocked_flushes += 1
+            self.release.wait(timeout=10)
+        return True
+
+
 class TestBedrockAgentCoreApp:
     def test_bedrock_agentcore_initialization(self):
         """Test BedrockAgentCoreApp initializes with correct name and routes."""
@@ -166,7 +187,7 @@ def handler(payload):
         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."""
+        """Sync streams flush only after their final chunk, not before the response starts."""
         bedrock_agentcore = BedrockAgentCoreApp()
 
         @bedrock_agentcore.entrypoint
@@ -182,11 +203,11 @@ def handler(payload):
             response = TestClient(bedrock_agentcore).post("/invocations", json={})
 
         assert response.status_code == 200
-        mock_async_flush.assert_awaited_once_with()
+        mock_async_flush.assert_not_awaited()
         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."""
+    def test_async_generator_entrypoint_flushes_after_stream_finishes(self):
+        """Async-generator entrypoints are bridged to a sync stream, which flushes once at the end."""
         bedrock_agentcore = BedrockAgentCoreApp()
 
         @bedrock_agentcore.entrypoint
@@ -202,9 +223,119 @@ async def handler(payload):
             response = TestClient(bedrock_agentcore).post("/invocations", json={})
 
         assert response.status_code == 200
-        mock_async_flush.assert_awaited_once_with()
+        mock_async_flush.assert_not_awaited()
         mock_sync_flush.assert_called_once_with()
 
+    def test_returned_async_generator_flushes_after_stream_finishes(self):
+        """A handler that returns an async generator streams on the async path, which flushes once at the end."""
+        bedrock_agentcore = BedrockAgentCoreApp()
+
+        async def stream():
+            yield {"ok": True}
+
+        @bedrock_agentcore.entrypoint
+        def handler(payload):
+            return stream()
+
+        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
+        assert 'data: {"ok": true}' in response.text
+        mock_async_flush.assert_awaited_once_with()
+        mock_sync_flush.assert_not_called()
+
+    def test_stream_headers_not_delayed_by_flush_of_earlier_spans(self):
+        """Spans queued before a streaming request must not hold back the response headers."""
+        provider = _TimeoutIgnoringTracerProvider()
+        provider.spans_pending = True
+        bedrock_agentcore = BedrockAgentCoreApp()
+
+        @bedrock_agentcore.entrypoint
+        def handler(payload):
+            yield {"ok": True}
+
+        try:
+            with (
+                patch("opentelemetry.trace.get_tracer_provider", return_value=provider),
+                patch(
+                    "bedrock_agentcore.runtime.app._flush_tracer_provider_async",
+                    functools.partial(_tracing_module._flush_tracer_provider_async, timeout_millis=3000),
+                ),
+                patch(
+                    "bedrock_agentcore.runtime.app._flush_tracer_provider",
+                    functools.partial(_tracing_module._flush_tracer_provider, timeout_millis=100),
+                ),
+            ):
+                begin = time.monotonic()
+                response = TestClient(bedrock_agentcore).post("/invocations", json={})
+                elapsed = time.monotonic() - begin
+        finally:
+            provider.release.set()
+            assert _tracing_module._flush_lock.acquire(timeout=5)
+            _tracing_module._flush_lock.release()
+
+        assert response.status_code == 200
+        assert elapsed < 2, f"response took {elapsed:.2f}s; a pre-stream flush held it back"
+        assert provider.blocked_flushes == 1
+
+    @pytest.mark.parametrize("handler_kind", ["sync_generator", "async_generator"])
+    def test_stream_closes_when_force_flush_ignores_timeout(self, handler_kind):
+        """A stuck exporter must not hold the stream open (OTel SDK 1.45 ignores ``timeout_millis``).
+
+        Async-generator entrypoints are bridged to a sync generator, so both kinds end in
+        ``_sync_stream_with_error_handling``.
+        """
+        provider = _TimeoutIgnoringTracerProvider()
+        bedrock_agentcore = BedrockAgentCoreApp()
+
+        if handler_kind == "sync_generator":
+
+            @bedrock_agentcore.entrypoint
+            def handler(payload):
+                provider.spans_pending = True
+                yield {"ok": True}
+
+        else:
+
+            @bedrock_agentcore.entrypoint
+            async def handler(payload):
+                provider.spans_pending = True
+                yield {"ok": True}
+
+        result = {}
+
+        def post():
+            result["response"] = TestClient(bedrock_agentcore).post("/invocations", json={})
+
+        client_thread = threading.Thread(target=post, daemon=True)
+        try:
+            with (
+                patch("opentelemetry.trace.get_tracer_provider", return_value=provider),
+                patch(
+                    "bedrock_agentcore.runtime.app._flush_tracer_provider",
+                    functools.partial(_tracing_module._flush_tracer_provider, timeout_millis=100),
+                ),
+            ):
+                client_thread.start()
+                client_thread.join(timeout=5)
+                assert not client_thread.is_alive(), "stream stayed open while force_flush was blocked"
+        finally:
+            provider.release.set()
+            client_thread.join(timeout=5)
+            # Let the abandoned flush finish so it does not hold the module-level lock for later tests.
+            assert _tracing_module._flush_lock.acquire(timeout=5)
+            _tracing_module._flush_lock.release()
+
+        assert provider.blocked_flushes == 1
+        assert result["response"].status_code == 200
+        assert 'data: {"ok": true}' in result["response"].text
+
     def test_async_handler_without_context(self):
         """Test async handler without context parameter."""
         bedrock_agentcore = BedrockAgentCoreApp()

tests/bedrock_agentcore/runtime/test_tracing.py

@@ -1,13 +1,15 @@
 """Tests for BaggageSpanProcessor and auto-registration helper."""
 
 import threading
+import time
 from unittest.mock import MagicMock, patch
 
 import pytest
 
 from bedrock_agentcore.runtime.tracing import (
     BaggageSpanProcessor,
     _ensure_baggage_processor_registered,
+    _flush_lock,
     _flush_tracer_provider,
     _flush_tracer_provider_async,
 )
@@ -303,6 +305,51 @@ def blocking_flush(**_):
 
         mock_provider.force_flush.assert_called_once()
 
+    def test_flush_wait_is_bounded_when_provider_ignores_timeout(self):
+        """OTel SDK 1.45 ``force_flush`` ignores ``timeout_millis``; the caller must still get control back."""
+        started = threading.Event()
+        release = threading.Event()
+        mock_provider = MagicMock()
+
+        def blocking_flush(**_):
+            started.set()
+            release.wait(timeout=5)
+
+        mock_provider.force_flush.side_effect = blocking_flush
+
+        try:
+            with patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider):
+                begin = time.monotonic()
+                _flush_tracer_provider(timeout_millis=50)
+                assert time.monotonic() - begin < 2
+                assert started.is_set()
+
+                # The abandoned flush still holds the lock, so later flushes skip instead of stacking threads.
+                _flush_tracer_provider(timeout_millis=50)
+                assert mock_provider.force_flush.call_count == 1
+        finally:
+            release.set()
+
+        # Once the abandoned flush returns it releases the lock, and the next flush runs normally.
+        assert _flush_lock.acquire(timeout=5)
+        _flush_lock.release()
+        mock_provider.force_flush.side_effect = None
+        with patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider):
+            _flush_tracer_provider(timeout_millis=50)
+        assert mock_provider.force_flush.call_count == 2
+
+    def test_lock_is_released_when_flush_thread_cannot_start(self):
+        mock_provider = MagicMock()
+
+        with (
+            patch("opentelemetry.trace.get_tracer_provider", return_value=mock_provider),
+            patch("threading.Thread.start", side_effect=RuntimeError("can't start new thread")),
+        ):
+            _flush_tracer_provider()
+
+        mock_provider.force_flush.assert_not_called()
+        assert not _flush_lock.locked()
+
     async def test_async_flush_runs_on_worker_thread(self):
         event_loop_thread = threading.get_ident()
         worker_threads = []