diff --git a/src/agents/tracing/processors.py b/src/agents/tracing/processors.py index b61f3e7976..ed050a7096 100644 --- a/src/agents/tracing/processors.py +++ b/src/agents/tracing/processors.py @@ -85,6 +85,8 @@ def __init__( self.base_delay = base_delay self.max_delay = max_delay self._shutdown_event = threading.Event() + self._shutdown_generation = 0 + self._shutdown_lock = threading.Lock() # Keep a client open for connection pooling across multiple export calls self._client = httpx2.Client(timeout=httpx2.Timeout(timeout=60, connect=5.0)) @@ -534,8 +536,17 @@ def close(self): """Close the underlying HTTP client.""" self._client.close() - def _request_shutdown(self) -> None: - self._shutdown_event.set() + def _request_shutdown(self) -> int: + with self._shutdown_lock: + self._shutdown_generation += 1 + generation = self._shutdown_generation + self._shutdown_event.set() + return generation + + def _reset_shutdown(self, generation: int) -> None: + with self._shutdown_lock: + if generation == self._shutdown_generation: + self._shutdown_event.clear() class BatchTraceProcessor(TracingProcessor): @@ -580,6 +591,7 @@ def __init__( self._thread_start_lock = threading.Lock() self._export_lock = threading.Lock() self._shutdown_deadline: float | None = None + self._exporter_shutdown_generation: int | None = None def _ensure_thread_started(self) -> None: # Fast path without holding the lock @@ -628,7 +640,7 @@ def shutdown(self, timeout: float | None = None): if timeout is not None: request_exporter_shutdown = getattr(self._exporter, "_request_shutdown", None) if callable(request_exporter_shutdown): - request_exporter_shutdown() + self._exporter_shutdown_generation = request_exporter_shutdown() deadline = None if timeout is None else time.monotonic() + timeout self._shutdown_deadline = deadline @@ -640,9 +652,14 @@ def shutdown(self, timeout: float | None = None): logger.warning( "[non-fatal] Tracing: shutdown timeout reached; dropping queued traces." ) + else: + self._reset_exporter_shutdown() else: # No background thread: process any remaining items synchronously. - self._export_batches(deadline=deadline) + try: + self._export_batches(deadline=deadline) + finally: + self._reset_exporter_shutdown() def force_flush(self): """ @@ -650,22 +667,35 @@ def force_flush(self): """ self._export_batches() + def _reset_exporter_shutdown(self) -> None: + if self._exporter_shutdown_generation is None: + return + reset_exporter_shutdown = getattr(self._exporter, "_reset_shutdown", None) + if callable(reset_exporter_shutdown): + reset_exporter_shutdown(self._exporter_shutdown_generation) + def _run(self): - while not self._shutdown_event.is_set(): - current_time = time.monotonic() - queue_size = self._queue.qsize() - - # If it's time for a scheduled flush or queue is above the trigger threshold - if current_time >= self._next_export_time or queue_size >= self._export_trigger_size: - self._export_batches() - # Reset the next scheduled flush time - self._next_export_time = time.monotonic() + self._schedule_delay - else: - # Sleep a short interval so we don't busy-wait. - time.sleep(0.2) + try: + while not self._shutdown_event.is_set(): + current_time = time.monotonic() + queue_size = self._queue.qsize() + + # If it's time for a scheduled flush or queue is above the trigger threshold + if ( + current_time >= self._next_export_time + or queue_size >= self._export_trigger_size + ): + self._export_batches() + # Reset the next scheduled flush time + self._next_export_time = time.monotonic() + self._schedule_delay + else: + # Sleep a short interval so we don't busy-wait. + time.sleep(0.2) - # Final drain after shutdown - self._export_batches(deadline=self._shutdown_deadline) + # Final drain after shutdown + self._export_batches(deadline=self._shutdown_deadline) + finally: + self._reset_exporter_shutdown() def _export_batches(self, deadline: float | None = None): """Drains the queue and exports in batches of up to `max_batch_size` until the queue diff --git a/tests/test_trace_processor.py b/tests/test_trace_processor.py index 07e975ccb9..d16bc0873c 100644 --- a/tests/test_trace_processor.py +++ b/tests/test_trace_processor.py @@ -612,6 +612,85 @@ def test_batch_trace_processor_shutdown_without_timeout_preserves_export_retries exporter.close() +@patch("httpx2.Client") +def test_reusing_exporter_after_shutdown_preserves_retries(mock_client): + mock_response = MagicMock() + mock_response.status_code = 504 + mock_client.return_value.post.return_value = mock_response + + exporter = BackendSpanExporter( + api_key="test_key", + max_retries=3, + base_delay=0.1, + max_delay=0.2, + ) + first_processor = BatchTraceProcessor(exporter=exporter) + first_processor.shutdown(timeout=1.0) + + second_processor = BatchTraceProcessor(exporter=exporter) + with patch.object(exporter._shutdown_event, "wait", return_value=False) as wait_for_retry: + second_processor._queue.put_nowait(get_span(second_processor)) + second_processor.force_flush() + + assert mock_client.return_value.post.call_count == 3 + assert wait_for_retry.call_count == 2 + + exporter.close() + + +@patch("httpx2.Client") +def test_reusing_exporter_does_not_reset_shutdown_while_previous_processor_runs(mock_client): + exporter = BackendSpanExporter(api_key="test_key") + exporter._request_shutdown() + + BatchTraceProcessor(exporter=exporter) + + assert exporter._shutdown_event.is_set() + + +@patch("httpx2.Client") +def test_exporter_shutdown_reset_only_applies_to_current_generation(mock_client): + exporter = BackendSpanExporter(api_key="test_key") + + old_generation = exporter._request_shutdown() + new_generation = exporter._request_shutdown() + exporter._reset_shutdown(old_generation) + assert exporter._shutdown_event.is_set() + + exporter._reset_shutdown(new_generation) + assert not exporter._shutdown_event.is_set() + + exporter.close() + + +@patch("httpx2.Client") +def test_reusing_exporter_does_not_cancel_old_worker(mock_client): + export_started = threading.Event() + release_export = threading.Event() + + class BlockingExporter(BackendSpanExporter): + def export(self, items): + export_started.set() + assert release_export.wait(timeout=2.0) + + exporter = BlockingExporter(api_key="test_key") + first_processor = BatchTraceProcessor(exporter=exporter, schedule_delay=0.01) + first_processor.on_span_end(get_span(first_processor)) + assert export_started.wait(timeout=2.0) + + first_processor.shutdown(timeout=0.01) + second_processor = BatchTraceProcessor(exporter=exporter) + assert exporter._shutdown_event.is_set() + + release_export.set() + first_processor._worker_thread.join(timeout=2.0) + assert not first_processor._worker_thread.is_alive() + assert not exporter._shutdown_event.is_set() + + second_processor.shutdown() + exporter.close() + + @pytest.mark.serial @pytest.mark.review_optional def test_tracing_atexit_cleanup_timeout_preserves_process_exit_code_on_504() -> None: