fix: make tracing shutdown best-effort on process exit (#3343)

Author: seratch

src/agents/tracing/processors.py

@@ -7,8 +7,9 @@
 import random
 import threading
 import time
+from collections.abc import Callable
 from functools import cached_property
-from typing import Any
+from typing import Any, cast
 
 import httpx
 
@@ -103,6 +104,9 @@ def project(self):
         return self._project or os.environ.get("OPENAI_PROJECT_ID")
 
     def export(self, items: list[Trace | Span[Any]]) -> None:
+        self._export_with_deadline(items, deadline=None)
+
+    def _export_with_deadline(self, items: list[Trace | Span[Any]], deadline: float | None) -> None:
         if not items:
             return
 
@@ -143,9 +147,23 @@ def export(self, items: list[Trace | Span[Any]]) -> None:
             attempt = 0
             delay = self.base_delay
             while True:
+                request_timeout = self._timeout_for_deadline(deadline)
+                if deadline is not None and request_timeout is None:
+                    logger.warning(
+                        "[non-fatal] Tracing: export deadline reached, giving up on this batch."
+                    )
+                    break
+
                 attempt += 1
                 try:
-                    response = self._client.post(url=self.endpoint, headers=headers, json=payload)
+                    request_kwargs: dict[str, Any] = {
+                        "url": self.endpoint,
+                        "headers": headers,
+                        "json": payload,
+                    }
+                    if request_timeout is not None:
+                        request_kwargs["timeout"] = request_timeout
+                    response = self._client.post(**request_kwargs)
 
                     # If the response is successful, break out of the loop
                     if response.status_code < 300:
@@ -178,9 +196,41 @@ def export(self, items: list[Trace | Span[Any]]) -> None:
 
                 # Exponential backoff + jitter
                 sleep_time = delay + random.uniform(0, 0.1 * delay)  # 10% jitter
-                time.sleep(sleep_time)
+                if not self._sleep_before_retry(sleep_time, deadline):
+                    break
                 delay = min(delay * 2, self.max_delay)
 
+    def _timeout_for_deadline(self, deadline: float | None) -> httpx.Timeout | None:
+        if deadline is None:
+            return None
+
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            return None
+
+        connect_timeout = min(5.0, remaining)
+        return httpx.Timeout(remaining, connect=connect_timeout)
+
+    def _sleep_before_retry(self, sleep_time: float, deadline: float | None) -> bool:
+        if deadline is None:
+            time.sleep(sleep_time)
+            return True
+
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            logger.warning("[non-fatal] Tracing: export deadline reached before retry, giving up.")
+            return False
+
+        if sleep_time >= remaining:
+            time.sleep(remaining)
+            logger.warning(
+                "[non-fatal] Tracing: export deadline reached during retry backoff, giving up."
+            )
+            return False
+
+        time.sleep(sleep_time)
+        return True
+
     def _should_sanitize_for_openai_tracing_api(self) -> bool:
         return self.endpoint.rstrip("/") == self._OPENAI_TRACING_INGEST_ENDPOINT.rstrip("/")
 
@@ -502,6 +552,7 @@ def __init__(
         self._worker_thread: threading.Thread | None = None
         self._thread_start_lock = threading.Lock()
         self._export_lock = threading.Lock()
+        self._shutdown_deadline: float | None = None
 
     def _ensure_thread_started(self) -> None:
         # Fast path without holding the lock
@@ -547,13 +598,19 @@ def shutdown(self, timeout: float | None = None):
         Called when the application stops. We signal our thread to stop, then join it.
         """
         self._shutdown_event.set()
+        deadline = None if timeout is None else time.monotonic() + timeout
+        self._shutdown_deadline = deadline
 
         # Only join if we ever started the background thread; otherwise flush synchronously.
         if self._worker_thread and self._worker_thread.is_alive():
             self._worker_thread.join(timeout=timeout)
+            if self._worker_thread.is_alive():
+                logger.warning(
+                    "[non-fatal] Tracing: shutdown timeout reached; dropping queued traces."
+                )
         else:
             # No background thread: process any remaining items synchronously.
-            self._export_batches(force=True)
+            self._export_batches(force=True, deadline=deadline)
 
     def force_flush(self):
         """
@@ -576,14 +633,20 @@ def _run(self):
                 time.sleep(0.2)
 
         # Final drain after shutdown
-        self._export_batches(force=True)
+        self._export_batches(force=True, deadline=self._shutdown_deadline)
 
-    def _export_batches(self, force: bool = False):
+    def _export_batches(self, force: bool = False, deadline: float | None = None):
         """Drains the queue and exports in batches. If force=True, export everything.
         Otherwise, export up to `max_batch_size` repeatedly until the queue is completely empty.
         """
         with self._export_lock:
             while True:
+                if deadline is not None and time.monotonic() >= deadline:
+                    logger.warning(
+                        "[non-fatal] Tracing: export deadline reached; dropping queued traces."
+                    )
+                    break
+
                 items_to_export: list[Span[Any] | Trace] = []
 
                 # Gather a batch of spans up to max_batch_size
@@ -604,7 +667,15 @@ def _export_batches(self, force: bool = False):
                 # cannot kill the background worker thread and silently strand all
                 # subsequent spans in the queue.
                 try:
-                    self._exporter.export(items_to_export)
+                    export_with_deadline = getattr(self._exporter, "_export_with_deadline", None)
+                    if deadline is not None and callable(export_with_deadline):
+                        export_fn = cast(
+                            Callable[[list[Trace | Span[Any]], float | None], None],
+                            export_with_deadline,
+                        )
+                        export_fn(items_to_export, deadline)
+                    else:
+                        self._exporter.export(items_to_export)
                 except Exception as exc:
                     logger.error(
                         "[non-fatal] Tracing: exporter raised %s; dropping batch of %d items",

src/agents/tracing/provider.py

@@ -3,10 +3,12 @@
 import logging
 import os
 import threading
+import time
 import uuid
 from abc import ABC, abstractmethod
 from datetime import datetime, timezone
-from typing import Any
+from inspect import Parameter, signature
+from typing import Any, cast
 
 from ..logger import logger
 from .config import TracingConfig
@@ -41,6 +43,24 @@ def _has_closed_stream_handler(log: logging.Logger) -> bool:
         return
 
 
+def _remaining_timeout(deadline: float | None) -> float | None:
+    if deadline is None:
+        return None
+    return max(0.0, deadline - time.monotonic())
+
+
+def _supports_shutdown_timeout(processor: TracingProcessor) -> bool:
+    try:
+        parameters = signature(processor.shutdown).parameters
+    except (TypeError, ValueError):
+        return False
+
+    for parameter in parameters.values():
+        if parameter.kind == Parameter.VAR_KEYWORD:
+            return True
+    return "timeout" in parameters
+
+
 def _is_noop_id(value: str | None) -> bool:
     return value == "no-op"
 
@@ -119,14 +139,24 @@ def on_span_end(self, span: Span[Any]) -> None:
             except Exception as e:
                 logger.error(f"Error in trace processor {processor} during on_span_end: {e}")
 
-    def shutdown(self) -> None:
+    def shutdown(self, timeout: float | None = None) -> None:
         """
         Called when the application stops.
         """
+        deadline = None if timeout is None else time.monotonic() + timeout
         for processor in self._processors:
             _safe_debug(f"Shutting down trace processor {processor}")
             try:
-                processor.shutdown()
+                processor_timeout = _remaining_timeout(deadline)
+                if processor_timeout is not None and processor_timeout <= 0:
+                    logger.warning(
+                        "[non-fatal] Tracing: shutdown timeout reached before processor cleanup."
+                    )
+                    return
+                if processor_timeout is not None and _supports_shutdown_timeout(processor):
+                    cast(Any, processor.shutdown)(timeout=processor_timeout)
+                else:
+                    processor.shutdown()
             except Exception as e:
                 logger.error(f"Error shutting down trace processor {processor}: {e}")
 
@@ -405,13 +435,13 @@ def force_flush(self) -> None:
         except Exception as e:
             logger.error(f"Error flushing trace provider: {e}")
 
-    def shutdown(self) -> None:
+    def shutdown(self, timeout: float | None = None) -> None:
         self._refresh_disabled_flag()
         if self._disabled:
             return
 
         try:
             _safe_debug("Shutting down trace provider")
-            self._multi_processor.shutdown()
+            self._multi_processor.shutdown(timeout=timeout)
         except Exception as e:
             logger.error(f"Error shutting down trace provider: {e}")

src/agents/tracing/setup.py

@@ -7,6 +7,7 @@
 if TYPE_CHECKING:
     from .provider import TraceProvider
 
+_DEFAULT_SHUTDOWN_TIMEOUT = 5.0
 GLOBAL_TRACE_PROVIDER: TraceProvider | None = None
 _GLOBAL_TRACE_PROVIDER_LOCK = threading.Lock()
 _SHUTDOWN_HANDLER_REGISTERED = False
@@ -15,6 +16,11 @@
 def _shutdown_global_trace_provider() -> None:
     provider = GLOBAL_TRACE_PROVIDER
     if provider is not None:
+        from .provider import DefaultTraceProvider
+
+        if isinstance(provider, DefaultTraceProvider):
+            provider.shutdown(timeout=_DEFAULT_SHUTDOWN_TIMEOUT)
+            return
         provider.shutdown()