Merge branch 'main' into feat/flush-traces-docs

Author: seratch

src/agents/__init__.py

@@ -203,6 +203,7 @@
     add_trace_processor,
     agent_span,
     custom_span,
+    flush_traces,
     function_span,
     gen_span_id,
     gen_trace_id,
@@ -451,6 +452,7 @@ def enable_verbose_stdout_logging():
     "add_trace_processor",
     "agent_span",
     "custom_span",
+    "flush_traces",
     "function_span",
     "generation_span",
     "get_current_span",

src/agents/extensions/memory/advanced_sqlite_session.py

@@ -4,7 +4,10 @@
 import json
 import sqlite3
 import threading
+from collections.abc import Iterator
+from contextlib import contextmanager
 from pathlib import Path
+from typing import ClassVar
 
 from ..items import TResponseInputItem
 from .session import SessionABC
@@ -20,6 +23,9 @@ class SQLiteSession(SessionABC):
     """
 
     session_settings: SessionSettings | None = None
+    _file_locks: ClassVar[dict[Path, threading.RLock]] = {}
+    _file_lock_counts: ClassVar[dict[Path, int]] = {}
+    _file_locks_guard: ClassVar[threading.Lock] = threading.Lock()
 
     def __init__(
         self,
@@ -46,21 +52,66 @@ def __init__(
         self.sessions_table = sessions_table
         self.messages_table = messages_table
         self._local = threading.local()
-        self._lock = threading.Lock()
 
         # For in-memory databases, we need a shared connection to avoid thread isolation
         # For file databases, we use thread-local connections for better concurrency
         self._is_memory_db = str(db_path) == ":memory:"
+        self._lock_path: Path | None = None
+        self._lock_released = False
         if self._is_memory_db:
-            self._shared_connection = sqlite3.connect(":memory:", check_same_thread=False)
-            self._shared_connection.execute("PRAGMA journal_mode=WAL")
-            self._init_db_for_connection(self._shared_connection)
+            self._lock = threading.RLock()
         else:
-            # For file databases, initialize the schema once since it persists
-            init_conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
-            init_conn.execute("PRAGMA journal_mode=WAL")
-            self._init_db_for_connection(init_conn)
-            init_conn.close()
+            self._lock_path, self._lock = self._acquire_file_lock(Path(self.db_path))
+
+        try:
+            if self._is_memory_db:
+                self._shared_connection = sqlite3.connect(":memory:", check_same_thread=False)
+                self._shared_connection.execute("PRAGMA journal_mode=WAL")
+                self._init_db_for_connection(self._shared_connection)
+            else:
+                # For file databases, initialize the schema once since it persists
+                with self._lock:
+                    init_conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
+                    init_conn.execute("PRAGMA journal_mode=WAL")
+                    self._init_db_for_connection(init_conn)
+                    init_conn.close()
+        except Exception:
+            if self._lock_path is not None and not self._lock_released:
+                self._release_file_lock(self._lock_path)
+                self._lock_released = True
+            raise
+
+    @classmethod
+    def _acquire_file_lock(cls, db_path: Path) -> tuple[Path, threading.RLock]:
+        """Return the path key and process-local lock for sessions sharing one SQLite file."""
+        lock_path = db_path.expanduser().resolve()
+        with cls._file_locks_guard:
+            lock = cls._file_locks.get(lock_path)
+            if lock is None:
+                lock = threading.RLock()
+                cls._file_locks[lock_path] = lock
+                cls._file_lock_counts[lock_path] = 0
+            cls._file_lock_counts[lock_path] += 1
+            return lock_path, lock
+
+    @classmethod
+    def _release_file_lock(cls, lock_path: Path) -> None:
+        """Drop the shared lock for a file-backed DB once the last session closes."""
+        with cls._file_locks_guard:
+            ref_count = cls._file_lock_counts.get(lock_path)
+            if ref_count is None:
+                return
+            if ref_count <= 1:
+                cls._file_lock_counts.pop(lock_path, None)
+                cls._file_locks.pop(lock_path, None)
+            else:
+                cls._file_lock_counts[lock_path] = ref_count - 1
+
+    @contextmanager
+    def _locked_connection(self) -> Iterator[sqlite3.Connection]:
+        """Serialize sqlite3 access while each operation runs in a worker thread."""
+        with self._lock:
+            yield self._get_connection()
 
     def _get_connection(self) -> sqlite3.Connection:
         """Get a database connection."""
@@ -114,6 +165,31 @@ def _init_db_for_connection(self, conn: sqlite3.Connection) -> None:
 
         conn.commit()
 
+    def _insert_items(self, conn: sqlite3.Connection, items: list[TResponseInputItem]) -> None:
+        conn.execute(
+            f"""
+            INSERT OR IGNORE INTO {self.sessions_table} (session_id) VALUES (?)
+        """,
+            (self.session_id,),
+        )
+
+        message_data = [(self.session_id, json.dumps(item)) for item in items]
+        conn.executemany(
+            f"""
+            INSERT INTO {self.messages_table} (session_id, message_data) VALUES (?, ?)
+        """,
+            message_data,
+        )
+
+        conn.execute(
+            f"""
+            UPDATE {self.sessions_table}
+            SET updated_at = CURRENT_TIMESTAMP
+            WHERE session_id = ?
+        """,
+            (self.session_id,),
+        )
+
     async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
         """Retrieve the conversation history for this session.
 
@@ -127,8 +203,7 @@ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
         session_limit = resolve_session_limit(limit, self.session_settings)
 
         def _get_items_sync():
-            conn = self._get_connection()
-            with self._lock if self._is_memory_db else threading.Lock():
+            with self._locked_connection() as conn:
                 if session_limit is None:
                     # Fetch all items in chronological order
                     cursor = conn.execute(
@@ -180,36 +255,8 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
             return
 
         def _add_items_sync():
-            conn = self._get_connection()
-
-            with self._lock if self._is_memory_db else threading.Lock():
-                # Ensure session exists
-                conn.execute(
-                    f"""
-                    INSERT OR IGNORE INTO {self.sessions_table} (session_id) VALUES (?)
-                """,
-                    (self.session_id,),
-                )
-
-                # Add items
-                message_data = [(self.session_id, json.dumps(item)) for item in items]
-                conn.executemany(
-                    f"""
-                    INSERT INTO {self.messages_table} (session_id, message_data) VALUES (?, ?)
-                """,
-                    message_data,
-                )
-
-                # Update session timestamp
-                conn.execute(
-                    f"""
-                    UPDATE {self.sessions_table}
-                    SET updated_at = CURRENT_TIMESTAMP
-                    WHERE session_id = ?
-                """,
-                    (self.session_id,),
-                )
-
+            with self._locked_connection() as conn:
+                self._insert_items(conn, items)
                 conn.commit()
 
         await asyncio.to_thread(_add_items_sync)
@@ -222,8 +269,7 @@ async def pop_item(self) -> TResponseInputItem | None:
         """
 
         def _pop_item_sync():
-            conn = self._get_connection()
-            with self._lock if self._is_memory_db else threading.Lock():
+            with self._locked_connection() as conn:
                 # Use DELETE with RETURNING to atomically delete and return the most recent item
                 cursor = conn.execute(
                     f"""
@@ -259,8 +305,7 @@ async def clear_session(self) -> None:
         """Clear all items for this session."""
 
         def _clear_session_sync():
-            conn = self._get_connection()
-            with self._lock if self._is_memory_db else threading.Lock():
+            with self._locked_connection() as conn:
                 conn.execute(
                     f"DELETE FROM {self.messages_table} WHERE session_id = ?",
                     (self.session_id,),
@@ -281,3 +326,6 @@ def close(self) -> None:
         else:
             if hasattr(self._local, "connection"):
                 self._local.connection.close()
+            if self._lock_path is not None and not self._lock_released:
+                self._release_file_lock(self._lock_path)
+                self._lock_released = True

src/agents/memory/sqlite_session.py

@@ -42,6 +42,7 @@
     "add_trace_processor",
     "agent_span",
     "custom_span",
+    "flush_traces",
     "function_span",
     "generation_span",
     "get_current_span",
@@ -108,3 +109,14 @@ def set_tracing_export_api_key(api_key: str) -> None:
     Set the OpenAI API key for the backend exporter.
     """
     default_exporter().set_api_key(api_key)
+
+
+def flush_traces() -> None:
+    """Force immediate export of buffered traces and spans.
+
+    The default ``BatchTraceProcessor`` already exports traces periodically in the
+    background. Call this when a worker, background job, or request handler needs
+    traces to be visible immediately after a unit of work finishes instead of
+    waiting for the next scheduled flush.
+    """
+    get_trace_provider().force_flush()

src/agents/tracing/__init__.py

@@ -491,6 +491,7 @@ def __init__(
         # We lazily start the background worker thread the first time a span/trace is queued.
         self._worker_thread: threading.Thread | None = None
         self._thread_start_lock = threading.Lock()
+        self._export_lock = threading.Lock()
 
     def _ensure_thread_started(self) -> None:
         # Fast path without holding the lock
@@ -571,25 +572,26 @@ def _export_batches(self, force: bool = False):
         """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.
         """
-        while True:
-            items_to_export: list[Span[Any] | Trace] = []
+        with self._export_lock:
+            while True:
+                items_to_export: list[Span[Any] | Trace] = []
+
+                # Gather a batch of spans up to max_batch_size
+                while not self._queue.empty() and (
+                    force or len(items_to_export) < self._max_batch_size
+                ):
+                    try:
+                        items_to_export.append(self._queue.get_nowait())
+                    except queue.Empty:
+                        # Another thread might have emptied the queue between checks
+                        break
 
-            # Gather a batch of spans up to max_batch_size
-            while not self._queue.empty() and (
-                force or len(items_to_export) < self._max_batch_size
-            ):
-                try:
-                    items_to_export.append(self._queue.get_nowait())
-                except queue.Empty:
-                    # Another thread might have emptied the queue between checks
+                # If we collected nothing, we're done
+                if not items_to_export:
                     break
 
-            # If we collected nothing, we're done
-            if not items_to_export:
-                break
-
-            # Export the batch
-            self._exporter.export(items_to_export)
+                # Export the batch
+                self._exporter.export(items_to_export)
 
 
 # Lazily initialized defaults to avoid creating network clients or threading

src/agents/tracing/processors.py

@@ -188,9 +188,21 @@ def create_span(
     ) -> Span[TSpanData]:
         """Create a new span."""
 
-    @abstractmethod
+    def force_flush(self) -> None:
+        """Force all registered processors to flush buffered traces/spans immediately.
+
+        The default implementation is a no-op so existing custom ``TraceProvider``
+        implementations continue to work without adding this method.
+        """
+        return None
+
     def shutdown(self) -> None:
-        """Clean up any resources used by the provider."""
+        """Clean up any resources used by the provider.
+
+        The default implementation is a no-op so existing custom ``TraceProvider``
+        implementations continue to work without adding this method.
+        """
+        return None
 
 
 class DefaultTraceProvider(TraceProvider):
@@ -365,7 +377,19 @@ def create_span(
             trace_metadata=trace_metadata,
         )
 
+    def force_flush(self) -> None:
+        """Force all processors to flush their buffers immediately."""
+        self._refresh_disabled_flag()
+        if self._disabled:
+            return
+
+        try:
+            self._multi_processor.force_flush()
+        except Exception as e:
+            logger.error(f"Error flushing trace provider: {e}")
+
     def shutdown(self) -> None:
+        self._refresh_disabled_flag()
         if self._disabled:
             return