From f3f817b1e8b45bd603e215dfe4f83d7429edfc7e Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:19:54 +0100 Subject: [PATCH] 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. --- .../alembic/versions/0037_request_traces.py | 84 +++++++ docsgpt/api/user/tasks.py | 29 +++ docsgpt/storage/db/models.py | 36 +++ .../storage/db/repositories/request_traces.py | 219 ++++++++++++++++++ tests/api/user/test_tasks.py | 84 ++++++- .../db/repositories/test_request_traces.py | 179 ++++++++++++++ tests/storage/db/test_migration_0037.py | 65 ++++++ 7 files changed, 686 insertions(+), 10 deletions(-) create mode 100644 docsgpt/alembic/versions/0037_request_traces.py create mode 100644 docsgpt/storage/db/repositories/request_traces.py create mode 100644 tests/storage/db/repositories/test_request_traces.py create mode 100644 tests/storage/db/test_migration_0037.py diff --git a/docsgpt/alembic/versions/0037_request_traces.py b/docsgpt/alembic/versions/0037_request_traces.py new file mode 100644 index 00000000..4929af62 --- /dev/null +++ b/docsgpt/alembic/versions/0037_request_traces.py @@ -0,0 +1,84 @@ +"""0037 request_traces — one execution trace per request for the Logs UI. + +Each row is the full span tree of one execution (a chat turn, a tool-approval +continuation, a scheduled or webhook run, a search, a graph extraction): +agent runs, LLM calls, tool calls, retrieval and embeddings with their +timings, stored as a ``spans`` JSONB array. The Logs UI loads a trace whole, +so one row per trace keeps the write to a single INSERT and lets retention +and conversation deletion remove a trace in one step. + +``message_id`` cascades: deleting a conversation, or truncating it when a +turn is superseded, removes that turn's traces with it. The link ids +(``request_id``, ``activity_id``, ``workflow_run_id``) carry partial indexes +because only the Logs rows that have them look traces up by them. + +``status`` is the only CHECK; span kinds and sources are validated in code so +new ones need no migration. + +Idempotent both ways. + +Revision ID: 0037_request_traces +Revises: 0036_device_audit_created_idx +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0037_request_traces" +down_revision: Union[str, None] = "0036_device_audit_created_idx" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + """ + CREATE TABLE IF NOT EXISTS request_traces ( + id UUID PRIMARY KEY, + request_id TEXT, + message_id UUID REFERENCES conversation_messages(id) ON DELETE CASCADE, + conversation_id UUID, + activity_id TEXT, + workflow_run_id UUID, + user_id TEXT, + agent_id UUID, + source TEXT NOT NULL, + name TEXT, + status TEXT NOT NULL + CONSTRAINT request_traces_status_chk + CHECK (status IN ('ok', 'error', 'paused', 'cancelled')), + started_at TIMESTAMPTZ NOT NULL, + duration_ms INTEGER NOT NULL DEFAULT 0, + span_count INTEGER NOT NULL DEFAULT 0, + dropped_spans INTEGER NOT NULL DEFAULT 0, + summary JSONB NOT NULL DEFAULT '{}'::jsonb, + spans JSONB NOT NULL DEFAULT '[]'::jsonb, + otel_trace_id TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + """ + ) + op.execute( + """ + CREATE INDEX IF NOT EXISTS request_traces_request_idx + ON request_traces (request_id) WHERE request_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_message_idx + ON request_traces (message_id) WHERE message_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_activity_idx + ON request_traces (activity_id) WHERE activity_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_workflow_run_idx + ON request_traces (workflow_run_id) WHERE workflow_run_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_user_started_idx + ON request_traces (user_id, started_at DESC); + CREATE INDEX IF NOT EXISTS request_traces_agent_started_idx + ON request_traces (agent_id, started_at DESC) WHERE agent_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_created_idx + ON request_traces (created_at); + """ + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS request_traces;") diff --git a/docsgpt/api/user/tasks.py b/docsgpt/api/user/tasks.py index 82d247e7..b34945fe 100644 --- a/docsgpt/api/user/tasks.py +++ b/docsgpt/api/user/tasks.py @@ -608,6 +608,11 @@ def setup_periodic_tasks(sender, **kwargs): cleanup_guardrail_events.s(), name="cleanup-guardrail-events", ) + sender.add_periodic_task( + timedelta(hours=24), + cleanup_traces.s(), + name="cleanup-traces", + ) sender.add_periodic_task( timedelta(hours=24), cleanup_orphan_memories.s(), @@ -827,6 +832,30 @@ def cleanup_guardrail_events(self): return {"deleted": deleted, "ttl_days": ttl_days} +@celery.task(bind=True, acks_late=False) +def cleanup_traces(self): + """Delete ``request_traces`` rows older than ``TRACES_RETENTION_DAYS``. + + Every chat turn, scheduled run and search writes a trace, and each one + carries content previews, so the table is bounded by a retention window + like the other per-request journals. + """ + from docsgpt.core.settings import settings + if not settings.POSTGRES_URI: + return {"deleted": 0, "skipped": "POSTGRES_URI not set"} + + from docsgpt.storage.db.engine import get_engine + from docsgpt.storage.db.repositories.request_traces import ( + RequestTracesRepository, + ) + + ttl_days = settings.TRACES_RETENTION_DAYS + engine = get_engine() + with engine.begin() as conn: + deleted = RequestTracesRepository(conn).purge_older_than(ttl_days) + return {"deleted": deleted, "ttl_days": ttl_days} + + @celery.task(bind=True, acks_late=False) def cleanup_orphan_memories(self): """Sweep orphan memories left by the 0009 FK-to-trigger orphan window. diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index 79773b6e..067dfb4f 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -825,6 +825,42 @@ Index( Index("ix_guardrail_events_message", guardrail_events_table.c.message_id) Index("ix_guardrail_events_created", guardrail_events_table.c.created_at) +# One execution trace per request (chat turn, continuation, scheduled or +# webhook run, search, graph extraction): the span tree as a JSONB array, +# rendered as a waterfall in the Logs UI. ``message_id`` cascades so deleting +# or truncating a conversation removes its traces. Migration 0037. +request_traces_table = Table( + "request_traces", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("request_id", Text), + Column( + "message_id", + UUID(as_uuid=True), + ForeignKey("conversation_messages.id", ondelete="CASCADE"), + ), + Column("conversation_id", UUID(as_uuid=True)), + Column("activity_id", Text), + Column("workflow_run_id", UUID(as_uuid=True)), + Column("user_id", Text), + Column("agent_id", UUID(as_uuid=True)), + Column("source", Text, nullable=False), + Column("name", Text), + # ok | error | paused | cancelled + Column("status", Text, nullable=False), + Column("started_at", DateTime(timezone=True), nullable=False), + Column("duration_ms", Integer, nullable=False, server_default="0"), + Column("span_count", Integer, nullable=False, server_default="0"), + Column("dropped_spans", Integer, nullable=False, server_default="0"), + Column("summary", JSONB, nullable=False, server_default="{}"), + Column("spans", JSONB, nullable=False, server_default="[]"), + Column("otel_trace_id", Text), + Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), +) + +Index("request_traces_user_started_idx", request_traces_table.c.user_id, request_traces_table.c.started_at) +Index("request_traces_created_idx", request_traces_table.c.created_at) + tool_call_attempts_table = Table( "tool_call_attempts", metadata, diff --git a/docsgpt/storage/db/repositories/request_traces.py b/docsgpt/storage/db/repositories/request_traces.py new file mode 100644 index 00000000..1fb5c594 --- /dev/null +++ b/docsgpt/storage/db/repositories/request_traces.py @@ -0,0 +1,219 @@ +"""Repository for ``request_traces``: one stored execution trace per request.""" + +from __future__ import annotations + +import datetime +import json +import logging +from typing import Any, Dict, Iterable, List, Optional + +from sqlalchemy import Connection, text +from sqlalchemy.exc import IntegrityError + +from docsgpt.storage.db.base_repository import looks_like_uuid, row_to_dict +from docsgpt.storage.db.serialization import PGNativeJSONEncoder +from docsgpt.utils import strip_null_bytes + +logger = logging.getLogger(__name__) + +#: Columns a trace can be looked up by, mapped to their SQL type. +REF_FIELDS: Dict[str, str] = { + "id": "uuid", + "request_id": "text", + "message_id": "uuid", + "activity_id": "text", + "workflow_run_id": "uuid", +} + +_SUMMARY_COLUMNS = ( + "id, request_id, message_id, conversation_id, activity_id, workflow_run_id, " + "user_id, agent_id, source, name, status, started_at, duration_ms, " + "span_count, dropped_spans, summary, otel_trace_id" +) + + +def _dump_jsonb(value: Any) -> str: + return json.dumps(strip_null_bytes(value), cls=PGNativeJSONEncoder) + + +def _uuid_or_none(value: Optional[str]) -> Optional[str]: + return str(value) if value and looks_like_uuid(str(value)) else None + + +def _started_at(record: Dict[str, Any]) -> datetime.datetime: + started_ns = record.get("started_at_ns") + if started_ns: + return datetime.datetime.fromtimestamp(started_ns / 1e9, tz=datetime.timezone.utc) + return datetime.datetime.now(datetime.timezone.utc) + + +def _scope_clause(user_id: Optional[str], agent_id: Optional[str]) -> tuple[str, Dict[str, Any]]: + """Owner scoping: an owned agent's traces, else the caller's own traces.""" + if agent_id: + return "agent_id = CAST(:scope_agent AS uuid)", {"scope_agent": str(agent_id)} + return "user_id = :scope_user", {"scope_user": user_id} + + +class RequestTracesRepository: + def __init__(self, conn: Connection) -> None: + self._conn = conn + + def insert(self, record: Dict[str, Any]) -> bool: + """Insert one finished trace (the dict from ``Trace.to_record``). + + A foreign-key failure means the message was deleted before the trace + flushed (e.g. the turn was superseded); the trace describes a turn the + user discarded, so it is dropped rather than raised. The insert runs + in a savepoint so that drop leaves the caller's transaction usable. + + Args: + record: The trace record. + + Returns: + True when the row was written. + """ + params = { + "id": record["id"], + "request_id": record.get("request_id"), + "message_id": _uuid_or_none(record.get("message_id")), + "conversation_id": _uuid_or_none(record.get("conversation_id")), + "activity_id": record.get("activity_id"), + "workflow_run_id": _uuid_or_none(record.get("workflow_run_id")), + "user_id": record.get("user_id"), + "agent_id": _uuid_or_none(record.get("agent_id")), + "source": record.get("source") or "unknown", + "name": strip_null_bytes(record.get("name")), + "status": record.get("status") or "ok", + "started_at": _started_at(record), + "duration_ms": int(record.get("duration_ms") or 0), + "span_count": int(record.get("span_count") or 0), + "dropped_spans": int(record.get("dropped_spans") or 0), + "summary": _dump_jsonb(record.get("summary") or {}), + "spans": _dump_jsonb(record.get("spans") or []), + "otel_trace_id": record.get("otel_trace_id"), + } + statement = text( + """ + INSERT INTO request_traces + (id, request_id, message_id, conversation_id, activity_id, + workflow_run_id, user_id, agent_id, source, name, status, + started_at, duration_ms, span_count, dropped_spans, summary, + spans, otel_trace_id) + VALUES + (CAST(:id AS uuid), :request_id, CAST(:message_id AS uuid), + CAST(:conversation_id AS uuid), :activity_id, + CAST(:workflow_run_id AS uuid), :user_id, CAST(:agent_id AS uuid), + :source, :name, :status, :started_at, :duration_ms, :span_count, + :dropped_spans, CAST(:summary AS jsonb), CAST(:spans AS jsonb), + :otel_trace_id) + ON CONFLICT (id) DO NOTHING + """ + ) + try: + with self._conn.begin_nested(): + result = self._conn.execute(statement, params) + except IntegrityError as exc: + sqlstate = getattr(getattr(exc, "orig", None), "sqlstate", None) + if sqlstate != "23503": + raise + logger.info("Dropped trace %s: its message no longer exists", record["id"]) + return False + return (result.rowcount or 0) > 0 + + def list_by_ref( + self, + field: str, + value: str, + *, + user_id: Optional[str], + agent_id: Optional[str] = None, + limit: int = 20, + ) -> List[dict]: + """Full traces matching ``field = value``, oldest first, owner-scoped. + + Args: + field: One of :data:`REF_FIELDS`. + value: The id to match. + user_id: The caller; used when ``agent_id`` is not given. + agent_id: An agent the caller owns; scopes to that agent's traces. + limit: Maximum traces returned. + + Returns: + Trace dicts including ``spans``; empty for an unknown field or id. + """ + sql_type = REF_FIELDS.get(field) + if sql_type is None or not value: + return [] + if sql_type == "uuid" and not looks_like_uuid(str(value)): + return [] + if not agent_id and not user_id: + return [] + scope, scope_params = _scope_clause(user_id, agent_id) + result = self._conn.execute( + text( + f""" + SELECT {_SUMMARY_COLUMNS}, spans FROM request_traces + WHERE {field} = CAST(:value AS {sql_type}) AND {scope} + ORDER BY started_at + LIMIT :limit + """ + ), + {"value": str(value), "limit": max(1, min(int(limit), 100)), **scope_params}, + ) + return [row_to_dict(row) for row in result.fetchall()] + + def summaries_for_refs( + self, + refs: Dict[str, Iterable[str]], + *, + user_id: Optional[str], + agent_id: Optional[str] = None, + ) -> Dict[str, Dict[str, List[dict]]]: + """Trace summaries (no spans) for a page of Logs rows. + + Args: + refs: ``{field: [ids...]}`` for fields in :data:`REF_FIELDS`. + user_id: The caller; used when ``agent_id`` is not given. + agent_id: An agent the caller owns. + + Returns: + ``{field: {id: [summary, ...]}}`` with summaries oldest first. + """ + out: Dict[str, Dict[str, List[dict]]] = {} + if not agent_id and not user_id: + return out + scope, scope_params = _scope_clause(user_id, agent_id) + for field, values in refs.items(): + sql_type = REF_FIELDS.get(field) + if sql_type is None: + continue + ids = sorted({str(v) for v in values if v}) + if sql_type == "uuid": + ids = [v for v in ids if looks_like_uuid(v)] + if not ids: + continue + result = self._conn.execute( + text( + f""" + SELECT {_SUMMARY_COLUMNS} FROM request_traces + WHERE {field} = ANY(CAST(:ids AS {sql_type}[])) AND {scope} + ORDER BY started_at + """ + ), + {"ids": ids, **scope_params}, + ) + for row in result.fetchall(): + data = row_to_dict(row) + out.setdefault(field, {}).setdefault(str(data[field]), []).append(data) + return out + + def purge_older_than(self, days: int) -> int: + """Delete traces older than ``days`` (the retention window).""" + result = self._conn.execute( + text( + "DELETE FROM request_traces " + "WHERE created_at < NOW() - CAST(:days || ' days' AS interval)" + ), + {"days": str(max(1, days))}, + ) + return result.rowcount or 0 diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index a593bbb5..a3ff2a19 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -299,7 +299,7 @@ class TestSetupPeriodicTasks: setup_periodic_tasks(sender) - assert sender.add_periodic_task.call_count == 14 + assert sender.add_periodic_task.call_count == 15 calls = sender.add_periodic_task.call_args_list @@ -326,20 +326,23 @@ class TestSetupPeriodicTasks: # guardrail_events retention sweep (24h) assert calls[8][0][0] == timedelta(hours=24) assert calls[8][1].get("name") == "cleanup-guardrail-events" - # orphan memories sweep (24h) + # request_traces retention sweep (24h) assert calls[9][0][0] == timedelta(hours=24) - assert calls[9][1].get("name") == "cleanup-orphan-memories" + assert calls[9][1].get("name") == "cleanup-traces" + # orphan memories sweep (24h) + assert calls[10][0][0] == timedelta(hours=24) + assert calls[10][1].get("name") == "cleanup-orphan-memories" # scheduler dispatcher - assert calls[10][1].get("name") == "dispatch-scheduled-runs" + assert calls[11][1].get("name") == "dispatch-scheduled-runs" # schedule runs cleanup (24h) - assert calls[11][0][0] == timedelta(hours=24) - assert calls[11][1].get("name") == "cleanup-schedule-runs" + assert calls[12][0][0] == timedelta(hours=24) + assert calls[12][1].get("name") == "cleanup-schedule-runs" # sandbox session reaper (60s) - assert calls[12][0][0] == timedelta(seconds=60) - assert calls[12][1].get("name") == "reap-sandbox-sessions" + assert calls[13][0][0] == timedelta(seconds=60) + assert calls[13][1].get("name") == "reap-sandbox-sessions" # stale workflow-run reaper (5m) - assert calls[13][0][0] == timedelta(seconds=300) - assert calls[13][1].get("name") == "reap-stale-workflow-runs" + assert calls[14][0][0] == timedelta(seconds=300) + assert calls[14][1].get("name") == "reap-stale-workflow-runs" class TestMcpOauthTask: @@ -679,6 +682,67 @@ class TestCleanupMessageEventsTask: assert [r["sequence_no"] for r in rows] == [1] +class TestCleanupTracesTask: + """Retention janitor for ``request_traces``.""" + + @pytest.mark.unit + def test_skips_when_postgres_uri_missing(self, monkeypatch): + from docsgpt.api.user.tasks import cleanup_traces + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "POSTGRES_URI", None, raising=False) + + assert cleanup_traces.run() == {"deleted": 0, "skipped": "POSTGRES_URI not set"} + + @pytest.mark.unit + def test_deletes_traces_past_retention_window(self, pg_conn, monkeypatch): + import time + import uuid + + from sqlalchemy import text as _text + + from docsgpt.api.user.tasks import cleanup_traces + from docsgpt.core.settings import settings + from docsgpt.storage.db.repositories.request_traces import ( + RequestTracesRepository, + ) + + repo = RequestTracesRepository(pg_conn) + for request_id in ("stale", "fresh"): + repo.insert( + { + "id": str(uuid.uuid4()), + "request_id": request_id, + "user_id": "u1", + "source": "stream", + "status": "ok", + "started_at_ns": time.time_ns(), + } + ) + pg_conn.execute( + _text( + "UPDATE request_traces SET created_at = now() - interval '45 days' " + "WHERE request_id = 'stale'" + ) + ) + monkeypatch.setattr(settings, "POSTGRES_URI", "postgresql://stub", raising=False) + monkeypatch.setattr(settings, "TRACES_RETENTION_DAYS", 30) + + @contextmanager + def _fake_begin(): + yield pg_conn + + fake_engine = MagicMock() + fake_engine.begin = _fake_begin + + with patch("docsgpt.storage.db.engine.get_engine", return_value=fake_engine): + result = cleanup_traces.run() + + assert result == {"deleted": 1, "ttl_days": 30} + remaining = pg_conn.execute(_text("SELECT request_id FROM request_traces")).scalars().all() + assert remaining == ["fresh"] + + class TestCleanupOrphanMemoriesTask: """Sweeps orphan memories from the FK-to-trigger orphan window.""" diff --git a/tests/storage/db/repositories/test_request_traces.py b/tests/storage/db/repositories/test_request_traces.py new file mode 100644 index 00000000..0dc72b05 --- /dev/null +++ b/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 diff --git a/tests/storage/db/test_migration_0037.py b/tests/storage/db/test_migration_0037.py new file mode 100644 index 00000000..bb453df2 --- /dev/null +++ b/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())" + ) + )