Write chat traces off the stream's thread

The OTel replay and the trace INSERT ran in the stream's finally, so a slow
database held the SSE connection open after the last event. The trace is
still frozen when the stream ends, but written on a small writer pool.

docsgpt/api/answer/routes/base.py

@@ -110,7 +110,9 @@ def _traced_stream(
             try:
                 yield from method(*bound.args, **bound.kwargs)
             finally:
-                tracing.flush(trace)
+                # Written on a writer thread so the stream's connection closes
+                # without waiting on the OTel replay and the INSERT.
+                tracing.flush(trace, background=True)
 
     return wrapper
 

docsgpt/tracing/sink.py

@@ -3,29 +3,36 @@
 from __future__ import annotations
 
 import logging
+import threading
+from concurrent.futures import Future, ThreadPoolExecutor
 from typing import Optional
 
 from docsgpt.tracing.core import Trace
 
 logger = logging.getLogger(__name__)
 
+# Writes handed off by ``flush(background=True)``. Two workers keep up with
+# chat traffic (each write is one INSERT); the pool's threads are joined at
+# interpreter exit, so queued traces are still written on a clean shutdown.
+_writer: Optional[ThreadPoolExecutor] = None
+_writer_lock = threading.Lock()
 
-def flush(trace: Optional[Trace], status: Optional[str] = None) -> None:
-    """Finish ``trace`` and persist it once; later calls are no-ops.
+
+def _executor() -> ThreadPoolExecutor:
+    global _writer
+    with _writer_lock:
+        if _writer is None:
+            _writer = ThreadPoolExecutor(max_workers=2, thread_name_prefix="trace-writer")
+        return _writer
+
+
+def _write(trace: Trace) -> None:
+    """Export ``trace`` to OTel, then store it; each sink swallows its own failure.
 
     OTel replay runs first so the exported trace id can be stored with the
-    row. Both sinks swallow their own failures: a trace is diagnostic data
-    and must never fail the request that produced it.
-
-    Args:
-        trace: The trace to write; ``None`` is accepted and ignored.
-        status: Final status; defaults to ``error`` when a top-level span
-            failed, else ``ok``.
+    row. A trace is diagnostic data and must never fail the request that
+    produced it.
     """
-    if trace is None or trace.flushed:
-        return
-    trace.flushed = True
-    trace.finish(status)
     try:
         from docsgpt.tracing.otel import export_trace
 
@@ -45,6 +52,39 @@ def flush(trace: Optional[Trace], status: Optional[str] = None) -> None:
         logger.warning("Failed to store trace %s", trace.id, exc_info=True)
 
 
+def flush(
+    trace: Optional[Trace], status: Optional[str] = None, *, background: bool = False
+) -> Optional[Future]:
+    """Finish ``trace`` and persist it once; later calls are no-ops.
+
+    The trace is frozen immediately, so nothing recorded afterwards is
+    included, whichever way it is written.
+
+    Args:
+        trace: The trace to write; ``None`` is accepted and ignored.
+        status: Final status; defaults to ``error`` when a top-level span
+            failed, else ``ok``.
+        background: Write on a writer thread instead of the caller's. A chat
+            stream flushes this way so its connection closes without waiting
+            on the OTel replay and the INSERT.
+
+    Returns:
+        The pending write when ``background`` is set, else ``None``.
+    """
+    if trace is None or trace.flushed:
+        return None
+    trace.flushed = True
+    trace.finish(status)
+    if background:
+        try:
+            return _executor().submit(_write, trace)
+        except RuntimeError:
+            # The pool is shut down (interpreter exit): write inline instead.
+            pass
+    _write(trace)
+    return None
+
+
 def discard(trace: Optional[Trace]) -> None:
     """Drop ``trace`` without writing it (the request was rejected or superseded)."""
     if trace is None or trace.flushed:

tests/api/answer/test_trace_wiring.py

@@ -22,7 +22,7 @@ def _captured_flushes():
     """Record every flushed trace instead of writing it."""
     flushed = []
 
-    def _fake_flush(trace, status=None):
+    def _fake_flush(trace, status=None, **_kwargs):
         if trace is None or trace.flushed:
             return
         trace.flushed = True

tests/e2e/specs/tier-a/execution-traces.spec.ts

@@ -75,9 +75,10 @@ test.describe('tier-a · execution traces', () => {
       expect(res.ok()).toBeTruthy();
       expect(await res.text()).toContain('"type": "end"');
 
-      // Stored trace: one row for the turn, linked to its message.
+      // Stored trace: one row for the turn, linked to its message. It is
+      // written just after the stream closes, so wait for it.
+      await expect.poll(async () => (await tracesFor(sub)).length).toBe(1);
       const traces = await tracesFor(sub);
-      expect(traces).toHaveLength(1);
       const [trace] = traces;
       expect(trace.source).toBe('stream');
       expect(trace.request_id).toBeTruthy();
@@ -154,6 +155,7 @@ test.describe('tier-a · execution traces', () => {
       });
       expect(await res.text()).toContain('"type": "end"');
 
+      await expect.poll(async () => (await tracesFor(sub)).length).toBe(1);
       const [trace] = await tracesFor(sub);
       const byKind = (kind: string) => trace.spans.filter((s) => s.kind === kind);
       const retrieval = byKind('retrieval');

tests/tracing/test_entry_points.py

@@ -22,7 +22,7 @@ def flushed():
     """Capture flushed traces instead of writing them."""
     captured = []
 
-    def _fake_flush(trace, status=None):
+    def _fake_flush(trace, status=None, **_kwargs):
         if trace is None or trace.flushed:
             return
         trace.flushed = True

tests/tracing/test_sink.py

@@ -98,3 +98,31 @@ class TestDiscard:
             tracing.discard(trace)
             tracing.flush(trace)
         repo.insert.assert_not_called()
+
+
+class TestBackgroundFlush:
+    def test_background_flush_writes_off_the_calling_thread(self):
+        import threading
+
+        trace = _trace_with_span()
+        seen = {}
+
+        def _insert(record):
+            seen["thread"] = threading.current_thread().name
+            return True
+
+        with _patched_store() as repo:
+            repo.insert.side_effect = _insert
+            future = tracing.flush(trace, background=True)
+            assert trace.finished  # frozen at once, written later
+            future.result(timeout=5)
+        assert seen["thread"].startswith("trace-writer")
+        assert repo.insert.call_count == 1
+
+    def test_background_flush_is_still_once(self):
+        trace = _trace_with_span()
+        with _patched_store() as repo:
+            future = tracing.flush(trace, background=True)
+            assert tracing.flush(trace, background=True) is None
+            future.result(timeout=5)
+        assert repo.insert.call_count == 1

tests/worker/test_extract_graph.py

@@ -188,7 +188,7 @@ class TestExtractGraphTrace:
         monkeypatch.setattr("docsgpt.graphrag.extraction.extract_graph_for_source", extract)
         flushed = []
 
-        def _flush(trace, status=None):
+        def _flush(trace, status=None, **_kwargs):
             trace.flushed = True
             trace.finish(status)
             flushed.append(trace)