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.
This commit is contained in:
arc53-machine committed 2026-09-23 22:42:05 +01:00
1 parent 0591580014
commit 778830208b
7 files changed
+91 -19

No files matched your search

+1 -1
View File
@@ -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
@@ -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');
+1 -1
View File
@@ -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
+28
View File
@@ -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
+1 -1
View File
@@ -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)