(bug): fix to queue drain

Author: levilentz

src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py

@@ -145,6 +145,9 @@ async def on_span_start(self, span: Span) -> None:
             items=[sgp_span.to_request_params()]
         )
 
+        # Input has been serialized and sent; clear it on the retained span to
+        # release memory.  on_span_end only needs output/metadata/end_time.
+        sgp_span.input = None  # type: ignore[assignment]
         self._spans[span.id] = sgp_span
 
     @override

src/agentex/lib/core/tracing/span_queue.py

@@ -12,6 +12,8 @@
 
 logger = make_logger(__name__)
 
+_DEFAULT_BATCH_SIZE = 50
+
 
 class SpanEventType(str, Enum):
     START = "start"
@@ -28,15 +30,18 @@ class _SpanQueueItem:
 class AsyncSpanQueue:
     """Background FIFO queue for async span processing.
 
-    Span events are enqueued synchronously (non-blocking) and processed
-    sequentially by a background drain task. This keeps tracing HTTP calls
-    off the critical request path while preserving start-before-end ordering.
+    Span events are enqueued synchronously (non-blocking) and drained by a
+    background task.  Items are processed in batches: all START events in a
+    batch are flushed concurrently, then all END events, so that per-span
+    start-before-end ordering is preserved while HTTP calls for independent
+    spans execute in parallel.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, batch_size: int = _DEFAULT_BATCH_SIZE) -> None:
         self._queue: asyncio.Queue[_SpanQueueItem] = asyncio.Queue()
         self._drain_task: asyncio.Task[None] | None = None
         self._stopping = False
+        self._batch_size = batch_size
 
     def enqueue(
         self,
@@ -54,9 +59,45 @@ def _ensure_drain_running(self) -> None:
         if self._drain_task is None or self._drain_task.done():
             self._drain_task = asyncio.create_task(self._drain_loop())
 
+    # ------------------------------------------------------------------
+    # Drain loop
+    # ------------------------------------------------------------------
+
     async def _drain_loop(self) -> None:
         while True:
-            item = await self._queue.get()
+            # Block until at least one item is available.
+            first = await self._queue.get()
+            batch: list[_SpanQueueItem] = [first]
+
+            # Opportunistically grab more ready items (non-blocking).
+            while len(batch) < self._batch_size:
+                try:
+                    batch.append(self._queue.get_nowait())
+                except asyncio.QueueEmpty:
+                    break
+
+            try:
+                # Separate START and END events.  Processing all STARTs before
+                # ENDs ensures that on_span_start completes before on_span_end
+                # for any span whose both events land in the same batch.
+                starts = [i for i in batch if i.event_type == SpanEventType.START]
+                ends = [i for i in batch if i.event_type == SpanEventType.END]
+
+                if starts:
+                    await self._process_items(starts)
+                if ends:
+                    await self._process_items(ends)
+            finally:
+                for _ in batch:
+                    self._queue.task_done()
+                # Release span data for GC.
+                batch.clear()
+
+    @staticmethod
+    async def _process_items(items: list[_SpanQueueItem]) -> None:
+        """Process a list of span events concurrently."""
+
+        async def _handle(item: _SpanQueueItem) -> None:
             try:
                 if item.event_type == SpanEventType.START:
                     coros = [p.on_span_start(item.span) for p in item.processors]
@@ -72,9 +113,15 @@ async def _drain_loop(self) -> None:
                             exc_info=result,
                         )
             except Exception:
-                logger.exception("Unexpected error in span queue drain loop for span %s", item.span.id)
-            finally:
-                self._queue.task_done()
+                logger.exception(
+                    "Unexpected error in span queue for span %s", item.span.id
+                )
+
+        await asyncio.gather(*[_handle(item) for item in items])
+
+    # ------------------------------------------------------------------
+    # Shutdown
+    # ------------------------------------------------------------------
 
     async def shutdown(self, timeout: float = 30.0) -> None:
         self._stopping = True

src/agentex/lib/core/tracing/trace.py

@@ -109,7 +109,7 @@ def end_span(
         if span.end_time is None:
             span.end_time = datetime.now(UTC)
 
-        span.input = recursive_model_dump(span.input) if span.input else None
+        # input was already serialized at start_span; skip redundant re-serialization
         span.output = recursive_model_dump(span.output) if span.output else None
         span.data = recursive_model_dump(span.data) if span.data else None
 
@@ -252,12 +252,17 @@ async def end_span(
         if span.end_time is None:
             span.end_time = datetime.now(UTC)
 
-        span.input = recursive_model_dump(span.input) if span.input else None
+        # input was already serialized at start_span; skip redundant re-serialization
         span.output = recursive_model_dump(span.output) if span.output else None
         span.data = recursive_model_dump(span.data) if span.data else None
 
         if self.processors:
-            self._span_queue.enqueue(SpanEventType.END, span.model_copy(deep=True), self.processors)
+            end_copy = span.model_copy(deep=True)
+            # input was already sent with the START event; drop it from the END
+            # copy to avoid retaining large payloads (system prompts, full
+            # conversation histories) in the async queue.
+            end_copy.input = None
+            self._span_queue.enqueue(SpanEventType.END, end_copy, self.processors)
 
         return span
 

tests/lib/core/tracing/processors/test_sgp_tracing_processor.py

@@ -162,3 +162,18 @@ async def test_span_end_for_unknown_span_is_noop(self):
         await processor.on_span_end(span)
 
         assert len(processor._spans) == 0
+
+    async def test_sgp_span_input_cleared_after_start(self):
+        """After on_span_start sends the data, sgp_span.input should be None to release memory."""
+        processor, _ = self._make_processor()
+
+        with patch(f"{MODULE}.create_span", side_effect=lambda **kw: _make_mock_sgp_span()):
+            span = _make_span()
+            span.input = {"system_prompt": "x" * 10_000}
+            await processor.on_span_start(span)
+
+        assert len(processor._spans) == 1
+        sgp_span = next(iter(processor._spans.values()))
+        assert sgp_span.input is None, (
+            "SGP span input should be cleared after upsert to release memory"
+        )