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