diff --git a/docsgpt/alembic/versions/0037_request_traces.py b/docsgpt/alembic/versions/0037_request_traces.py index 4929af62..0f57f343 100644 --- a/docsgpt/alembic/versions/0037_request_traces.py +++ b/docsgpt/alembic/versions/0037_request_traces.py @@ -7,10 +7,15 @@ 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. +``message_id`` cascades: truncating a conversation when a turn is +superseded removes that turn's traces with it. Deleting a conversation also +deletes, by ``conversation_id``, the traces that have no message (scheduled +runs, stateless ``/v1`` rounds, turns whose message was never reserved); the +conversation index serves that. 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. Traces with no log row of +their own (searches, graph builds) are listed per user or agent and source, +hence the source-leading indexes. ``status`` is the only CHECK; span kinds and sources are validated in code so new ones need no migration. @@ -70,10 +75,12 @@ def upgrade() -> None: 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_conversation_idx + ON request_traces (conversation_id) WHERE conversation_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_user_source_started_idx + ON request_traces (user_id, source, started_at DESC); + CREATE INDEX IF NOT EXISTS request_traces_agent_source_started_idx + ON request_traces (agent_id, source, started_at DESC) WHERE agent_id IS NOT NULL; CREATE INDEX IF NOT EXISTS request_traces_created_idx ON request_traces (created_at); """ diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index 067dfb4f..ebd25567 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -827,8 +827,8 @@ 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. +# rendered as a waterfall in the Logs UI. ``message_id`` cascades; deleting a +# conversation also deletes its traces by ``conversation_id``. Migration 0037. request_traces_table = Table( "request_traces", metadata, @@ -858,7 +858,12 @@ request_traces_table = Table( 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_user_source_started_idx", + request_traces_table.c.user_id, + request_traces_table.c.source, + request_traces_table.c.started_at, +) Index("request_traces_created_idx", request_traces_table.c.created_at) tool_call_attempts_table = Table( diff --git a/docsgpt/storage/db/repositories/conversations.py b/docsgpt/storage/db/repositories/conversations.py index 95205350..7d27a3f5 100644 --- a/docsgpt/storage/db/repositories/conversations.py +++ b/docsgpt/storage/db/repositories/conversations.py @@ -576,6 +576,18 @@ class ConversationsRepository: artifacts = ArtifactsRepository(self._conn) paths = artifacts.storage_paths_for_conversation(conversation_id) + # Traces linked to a message go with it (ON DELETE CASCADE). This + # takes the rest -- scheduled runs, stateless /v1 rounds, turns whose + # message was never reserved -- whose tool results and chunk previews + # the user expects gone with the conversation. + self._conn.execute( + text( + "DELETE FROM request_traces WHERE conversation_id = CAST(:id AS uuid) " + "AND EXISTS (SELECT 1 FROM conversations " + "WHERE id = CAST(:id AS uuid) AND user_id = :user_id)" + ), + {"id": conversation_id, "user_id": user_id}, + ) result = self._conn.execute( text( "DELETE FROM conversations " @@ -597,6 +609,14 @@ class ConversationsRepository: artifacts = ArtifactsRepository(self._conn) paths = artifacts.storage_paths_for_user_conversations(user_id) artifacts.delete_for_user_conversations(user_id) + # See ``delete``: traces not linked to a message are removed explicitly. + self._conn.execute( + text( + "DELETE FROM request_traces WHERE conversation_id IN " + "(SELECT id FROM conversations WHERE user_id = :user_id)" + ), + {"user_id": user_id}, + ) result = self._conn.execute( text("DELETE FROM conversations WHERE user_id = :user_id"), {"user_id": user_id}, diff --git a/tests/storage/db/repositories/test_request_traces.py b/tests/storage/db/repositories/test_request_traces.py index 0dc72b05..7d380188 100644 --- a/tests/storage/db/repositories/test_request_traces.py +++ b/tests/storage/db/repositories/test_request_traces.py @@ -141,6 +141,27 @@ class TestMessageLink: assert count == 0 +class TestConversationDeletion: + """Every trace of a conversation goes with it, not only message-linked ones.""" + + def test_trace_without_message_is_deleted_with_its_conversation(self, pg_conn): + conv_id, _msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + # e.g. a scheduled run in this conversation, or a stateless /v1 round. + repo.insert(_record(conversation_id=conv_id, source="schedule")) + repo.insert(_record(request_id="other")) # another conversation's trace + ConversationsRepository(pg_conn).delete(conv_id, "u1") + remaining = pg_conn.execute(text("SELECT request_id FROM request_traces")).scalars().all() + assert remaining == ["other"] + + def test_delete_all_for_user_removes_their_conversation_traces(self, pg_conn): + conv_id, _msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + repo.insert(_record(conversation_id=conv_id)) + ConversationsRepository(pg_conn).delete_all_for_user("u1") + assert pg_conn.execute(text("SELECT count(*) FROM request_traces")).scalar() == 0 + + class TestSummaries: def test_summaries_grouped_by_field_and_id(self, pg_conn): repo = RequestTracesRepository(pg_conn) diff --git a/tests/storage/db/test_migration_0037.py b/tests/storage/db/test_migration_0037.py index bb453df2..e8838a74 100644 --- a/tests/storage/db/test_migration_0037.py +++ b/tests/storage/db/test_migration_0037.py @@ -63,3 +63,16 @@ class TestMigration0037RoundTrip: "VALUES (gen_random_uuid(), 'stream', 'weird', now())" ) ) + + def test_listing_and_deletion_indexes_exist(self, pg_engine): + with pg_engine.connect() as conn: + names = set( + conn.execute( + text("SELECT indexname FROM pg_indexes WHERE tablename = 'request_traces'") + ).scalars() + ) + assert { + "request_traces_conversation_idx", + "request_traces_user_source_started_idx", + "request_traces_agent_source_started_idx", + } <= names