From fdba261d113cb689671c91fa3c22c5aeb43a6ac9 Mon Sep 17 00:00:00 2001 From: Alex Date: Sat, 5 Sep 2026 09:49:22 +0100 Subject: [PATCH] fix(db): make compression-point deduplication atomic The duplicate check for append_compression_point ran as a SELECT before the UPDATE, so two sessions persisting the same point could both pass it. The predicate now lives inside the UPDATE's CASE expression and is evaluated under the row lock; a point without query_index or compressed_summary never matches and is always appended. Adds a two-session regression test that races the same point and expects one. --- .../storage/db/repositories/conversations.py | 76 +++++++++---------- .../db/repositories/test_conversations.py | 40 ++++++++++ 2 files changed, 78 insertions(+), 38 deletions(-) diff --git a/application/storage/db/repositories/conversations.py b/application/storage/db/repositories/conversations.py index b432f988..aa01af17 100644 --- a/application/storage/db/repositories/conversations.py +++ b/application/storage/db/repositories/conversations.py @@ -508,49 +508,47 @@ class ConversationsRepository: return False # Idempotent on the same point: two persist paths (the mid-execution # handler and the route's end-of-stream save) used to both append it. - last_point = self._conn.execute( - text( - "SELECT compression_metadata -> 'compression_points' -> -1 " - "FROM conversations WHERE id = CAST(:id AS uuid)" - ), - {"id": conversation_id}, - ).scalar() - if ( - isinstance(last_point, dict) - and point.get("query_index") is not None - and point.get("compressed_summary") is not None - and last_point.get("query_index") == point.get("query_index") - and last_point.get("compressed_summary") == point.get("compressed_summary") - ): - return True + # The duplicate check lives inside the UPDATE so it runs under the + # row lock — a separate SELECT would let two sessions both pass it. + # A point without ``query_index``/``compressed_summary`` never matches + # (NULL comparisons are false), so it is always appended. + dup_index = point.get("query_index") + dup_summary = point.get("compressed_summary") result = self._conn.execute( text( """ UPDATE conversations SET - compression_metadata = jsonb_set( - COALESCE(compression_metadata, '{}'::jsonb), - '{compression_points}', - COALESCE( - ( - SELECT jsonb_agg(elem ORDER BY rn) - FROM ( - SELECT - elem, - row_number() OVER () AS rn, - count(*) OVER () AS cnt - FROM jsonb_array_elements( - COALESCE( - compression_metadata -> 'compression_points', - '[]'::jsonb - ) || jsonb_build_array(CAST(:point AS jsonb)) - ) AS elem - ) ranked - WHERE rn > cnt - :max_points + compression_metadata = CASE + WHEN (compression_metadata -> 'compression_points' -> -1 + ->> 'query_index') = CAST(:dup_index AS text) + AND (compression_metadata -> 'compression_points' -> -1 + ->> 'compressed_summary') = CAST(:dup_summary AS text) + THEN compression_metadata + ELSE jsonb_set( + COALESCE(compression_metadata, '{}'::jsonb), + '{compression_points}', + COALESCE( + ( + SELECT jsonb_agg(elem ORDER BY rn) + FROM ( + SELECT + elem, + row_number() OVER () AS rn, + count(*) OVER () AS cnt + FROM jsonb_array_elements( + COALESCE( + compression_metadata -> 'compression_points', + '[]'::jsonb + ) || jsonb_build_array(CAST(:point AS jsonb)) + ) AS elem + ) ranked + WHERE rn > cnt - :max_points + ), + '[]'::jsonb ), - '[]'::jsonb - ), - true - ), + true + ) + END, updated_at = now() WHERE id = CAST(:id AS uuid) """ @@ -559,6 +557,8 @@ class ConversationsRepository: "id": conversation_id, "point": json.dumps(point, cls=PGNativeJSONEncoder), "max_points": int(max_points), + "dup_index": None if dup_index is None else str(dup_index), + "dup_summary": None if dup_summary is None else str(dup_summary), }, ) return result.rowcount > 0 diff --git a/tests/storage/db/repositories/test_conversations.py b/tests/storage/db/repositories/test_conversations.py index 0ab638da..5241f464 100644 --- a/tests/storage/db/repositories/test_conversations.py +++ b/tests/storage/db/repositories/test_conversations.py @@ -393,6 +393,46 @@ class TestCompressionMetadata: points = repo.get(conv["id"], "user-1")["compression_metadata"]["compression_points"] assert [p["query_index"] for p in points] == [3, 5] + def test_append_compression_point_dedupes_across_concurrent_sessions(self, pg_engine): + """The duplicate check runs inside the UPDATE, under the row lock, so + two sessions racing to persist the same point store it once.""" + import threading + import time + + with pg_engine.connect() as setup: + with setup.begin(): + conv = _repo(setup).create("user-1", "c") + point = {"query_index": 3, "compressed_summary": "S", "original_token_count": 90} + results, errors = [], [] + + def append(hold_before_commit): + try: + with pg_engine.connect() as conn: + with conn.begin(): + results.append( + _repo(conn).append_compression_point( + conv["id"], dict(point), max_points=3 + ) + ) + time.sleep(hold_before_commit) + except Exception as exc: # pragma: no cover - surfaced by the assertion + errors.append(exc) + + first = threading.Thread(target=append, args=(0.5,)) + second = threading.Thread(target=append, args=(0.0,)) + first.start() + time.sleep(0.1) # second blocks on the row lock until first commits + second.start() + first.join() + second.join() + + assert not errors + assert results == [True, True] + with pg_engine.connect() as conn: + fetched = _repo(conn).get(conv["id"], "user-1") + points = fetched["compression_metadata"]["compression_points"] + assert [p["query_index"] for p in points] == [3] + class TestResolveAgentRef: """The repo must translate Mongo ObjectId-shaped ``agent_id`` values