Store execution traces in a new request_traces table

One row per execution holding its span tree as JSONB, linked to Logs rows
by request, message, activity and workflow-run ids. message_id cascades so
deleting or truncating a conversation drops its traces; a daily beat task
enforces TRACES_RETENTION_DAYS.

tests/storage/db/repositories/test_request_traces.py

@@ -0,0 +1,179 @@
+"""Tests for RequestTracesRepository against a real Postgres instance."""
+
+from __future__ import annotations
+
+import time
+import uuid
+
+from sqlalchemy import text
+
+from docsgpt.storage.db.repositories.conversations import ConversationsRepository
+from docsgpt.storage.db.repositories.request_traces import RequestTracesRepository
+
+
+def _record(**overrides):
+    base = {
+        "id": str(uuid.uuid4()),
+        "request_id": "req-1",
+        "message_id": None,
+        "conversation_id": None,
+        "activity_id": None,
+        "workflow_run_id": None,
+        "user_id": "u1",
+        "agent_id": None,
+        "source": "stream",
+        "name": "stream",
+        "status": "ok",
+        "started_at_ns": time.time_ns(),
+        "duration_ms": 1234,
+        "span_count": 1,
+        "dropped_spans": 0,
+        "summary": {"llm_calls": 1},
+        "spans": [
+            {
+                "id": "s1",
+                "parent_id": None,
+                "kind": "llm",
+                "name": "chat m",
+                "status": "ok",
+                "offset_ms": 1.0,
+                "duration_ms": 10.0,
+                "attributes": {"gen_ai.request.model": "m"},
+            }
+        ],
+        "otel_trace_id": None,
+    }
+    base.update(overrides)
+    return base
+
+
+def _message(pg_conn, user_id="u1"):
+    convs = ConversationsRepository(pg_conn)
+    conv = convs.create(user_id, "t")
+    msg = convs.reserve_message(
+        str(conv["id"]), prompt="q", placeholder_response="..."
+    )
+    return str(conv["id"]), str(msg["id"])
+
+
+class TestInsert:
+    def test_round_trip(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        record = _record()
+        assert repo.insert(record) is True
+        rows = repo.list_by_ref("request_id", "req-1", user_id="u1")
+        assert len(rows) == 1
+        row = rows[0]
+        assert row["id"] == record["id"]
+        assert row["duration_ms"] == 1234
+        assert row["summary"] == {"llm_calls": 1}
+        assert row["spans"][0]["attributes"]["gen_ai.request.model"] == "m"
+
+    def test_strips_nul_bytes(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        record = _record(
+            spans=[{"id": "s", "name": "a\x00b", "preview": {"result": "x\x00y"}}]
+        )
+        assert repo.insert(record)
+        row = repo.list_by_ref("id", record["id"], user_id="u1")[0]
+        assert row["spans"][0]["name"] == "ab"
+        assert row["spans"][0]["preview"]["result"] == "xy"
+
+    def test_non_uuid_link_ids_become_null(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        record = _record(agent_id="not-a-uuid", conversation_id="nope")
+        assert repo.insert(record)
+        row = repo.list_by_ref("id", record["id"], user_id="u1")[0]
+        assert row["agent_id"] is None
+        assert row["conversation_id"] is None
+
+    def test_missing_message_drops_row_and_keeps_transaction_usable(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        assert repo.insert(_record(message_id=str(uuid.uuid4()))) is False
+        # The savepoint rolled back; the outer transaction still works.
+        assert repo.insert(_record(request_id="req-after")) is True
+
+    def test_duplicate_id_is_ignored(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        record = _record()
+        assert repo.insert(record) is True
+        assert repo.insert(record) is False
+
+
+class TestScoping:
+    def test_other_users_cannot_read(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        repo.insert(_record())
+        assert repo.list_by_ref("request_id", "req-1", user_id="u2") == []
+
+    def test_agent_scope_returns_agent_traces_from_any_user(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        agent_id = str(uuid.uuid4())
+        repo.insert(_record(user_id="someone-else", agent_id=agent_id))
+        rows = repo.list_by_ref("request_id", "req-1", user_id="owner", agent_id=agent_id)
+        assert len(rows) == 1
+
+    def test_unknown_field_and_bad_uuid_return_empty(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        repo.insert(_record())
+        assert repo.list_by_ref("user_id", "u1", user_id="u1") == []
+        assert repo.list_by_ref("message_id", "not-uuid", user_id="u1") == []
+        assert repo.list_by_ref("request_id", "req-1", user_id=None) == []
+
+
+class TestMessageLink:
+    def test_rounds_for_a_message_are_ordered(self, pg_conn):
+        conv_id, msg_id = _message(pg_conn)
+        repo = RequestTracesRepository(pg_conn)
+        first = _record(message_id=msg_id, status="paused", started_at_ns=time.time_ns())
+        second = _record(message_id=msg_id, started_at_ns=time.time_ns() + 1_000_000)
+        repo.insert(second)
+        repo.insert(first)
+        rows = repo.list_by_ref("message_id", msg_id, user_id="u1")
+        assert [r["id"] for r in rows] == [first["id"], second["id"]]
+
+    def test_deleting_the_conversation_cascades(self, pg_conn):
+        conv_id, msg_id = _message(pg_conn)
+        repo = RequestTracesRepository(pg_conn)
+        repo.insert(_record(message_id=msg_id, conversation_id=conv_id))
+        ConversationsRepository(pg_conn).delete(conv_id, "u1")
+        count = pg_conn.execute(text("SELECT count(*) FROM request_traces")).scalar()
+        assert count == 0
+
+
+class TestSummaries:
+    def test_summaries_grouped_by_field_and_id(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        a = _record(request_id="r-a")
+        b = _record(request_id="r-b", activity_id="act-1")
+        repo.insert(a)
+        repo.insert(b)
+        out = repo.summaries_for_refs(
+            {"request_id": ["r-a", "r-b", "r-missing"], "activity_id": ["act-1"]},
+            user_id="u1",
+        )
+        assert set(out["request_id"]) == {"r-a", "r-b"}
+        assert out["activity_id"]["act-1"][0]["id"] == b["id"]
+        assert "spans" not in out["request_id"]["r-a"][0]
+
+    def test_summaries_are_scoped(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        repo.insert(_record(request_id="r-a"))
+        assert repo.summaries_for_refs({"request_id": ["r-a"]}, user_id="u2") == {}
+
+
+class TestPurge:
+    def test_purge_older_than(self, pg_conn):
+        repo = RequestTracesRepository(pg_conn)
+        old = _record(request_id="old")
+        repo.insert(old)
+        repo.insert(_record(request_id="new"))
+        pg_conn.execute(
+            text(
+                "UPDATE request_traces SET created_at = now() - interval '40 days' "
+                "WHERE request_id = 'old'"
+            )
+        )
+        assert repo.purge_older_than(30) == 1
+        assert repo.list_by_ref("request_id", "old", user_id="u1") == []
+        assert len(repo.list_by_ref("request_id", "new", user_id="u1")) == 1

tests/storage/db/test_migration_0037.py

@@ -0,0 +1,65 @@
+"""Migration round-trip test for 0037_request_traces."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+from sqlalchemy import text
+
+
+pytestmark = pytest.mark.integration
+
+
+def _alembic_ini() -> Path:
+    return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini"
+
+
+def _run_alembic(url: str, *args: str) -> None:
+    subprocess.check_call(
+        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args],
+        timeout=60,
+        env={**os.environ, "POSTGRES_URI": url},
+    )
+
+
+def _alembic_version(conn) -> str:
+    return conn.execute(text("SELECT version_num FROM alembic_version")).scalar()
+
+
+def _table_exists(conn, table: str) -> bool:
+    return conn.execute(text("SELECT to_regclass(:t)"), {"t": f"public.{table}"}).scalar() is not None
+
+
+_0037 = "0037_request_traces"
+_0036 = "0036_device_audit_created_idx"
+
+
+class TestMigration0037RoundTrip:
+    def test_head_has_request_traces(self, pg_engine):
+        with pg_engine.connect() as conn:
+            assert _alembic_version(conn) >= _0037
+            assert _table_exists(conn, "request_traces")
+
+    def test_downgrade_drops_then_upgrade_restores(self, pg_engine):
+        url = pg_engine.url.render_as_string(hide_password=False)
+        _run_alembic(url, "downgrade", _0036)
+        with pg_engine.connect() as conn:
+            assert _alembic_version(conn) == _0036
+            assert not _table_exists(conn, "request_traces")
+        _run_alembic(url, "upgrade", "head")
+        with pg_engine.connect() as conn:
+            assert _table_exists(conn, "request_traces")
+
+    def test_status_check_rejects_unknown_status(self, pg_engine):
+        with pytest.raises(Exception):
+            with pg_engine.begin() as conn:
+                conn.execute(
+                    text(
+                        "INSERT INTO request_traces (id, source, status, started_at) "
+                        "VALUES (gen_random_uuid(), 'stream', 'weird', now())"
+                    )
+                )