diff --git a/src/bedrock_agentcore/runtime/app.py b/src/bedrock_agentcore/runtime/app.py index efae2faa..c291b255 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, _flush_tracer_provider, _flush_tracer_provider_async +from .tracing import _ensure_baggage_processor_registered from .utils import convert_complex_objects # Sentinel so we only parse OTEL_RESOURCE_ATTRIBUTES once per process. @@ -613,8 +613,6 @@ 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: @@ -896,8 +894,6 @@ 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. @@ -953,5 +949,3 @@ 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 449c13ee..7c0cf9cb 100644 --- a/src/bedrock_agentcore/runtime/tracing.py +++ b/src/bedrock_agentcore/runtime/tracing.py @@ -22,7 +22,6 @@ provider after the app is constructed. """ -import asyncio import logging import threading from typing import Optional @@ -34,7 +33,6 @@ # 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: @@ -65,35 +63,6 @@ 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 7ce404c6..60169a5a 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 AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest from starlette.testclient import TestClient @@ -151,60 +151,6 @@ 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 21f79867..020b4595 100644 --- a/tests/bedrock_agentcore/runtime/test_tracing.py +++ b/tests/bedrock_agentcore/runtime/test_tracing.py @@ -5,12 +5,7 @@ import pytest -from bedrock_agentcore.runtime.tracing import ( - BaggageSpanProcessor, - _ensure_baggage_processor_registered, - _flush_tracer_provider, - _flush_tracer_provider_async, -) +from bedrock_agentcore.runtime.tracing import BaggageSpanProcessor, _ensure_baggage_processor_registered # --------------------------------------------------------------------------- # BaggageSpanProcessor.on_start @@ -259,87 +254,6 @@ 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) # ---------------------------------------------------------------------------