Fix tracing exporter reuse after shutdown

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,32 +652,50 @@ 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):
         """
         Forces an immediate flush of all queued spans.
         """
         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

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: