Quellcode durchsuchen

fix: deduplicate metric keys during migration

Problem: Existing databases with duplicate non-null usage_stats metric keys failed initialization when the migration created the unique session/key index.

Risk: The migration deletes later duplicate rows for each session/key while retaining MIN(id); NULL metric keys remain untouched and concurrent initialization stays serialized by BEGIN IMMEDIATE.
zhenyu.hu vor 2 Wochen
Ursprung
Commit
1d5f65f582
2 geänderte Dateien mit 104 neuen und 0 gelöschten Zeilen
  1. 12 0
      src/agent_lab/infrastructure/sqlite_store.py
  2. 92 0
      tests/test_sqlite_store.py

+ 12 - 0
src/agent_lab/infrastructure/sqlite_store.py

@@ -635,6 +635,18 @@ class SQLiteSessionStore:
             self._connection.execute(
                 "DROP INDEX IF EXISTS usage_stats_metric_key_uq"
             )
+            self._connection.execute(
+                """
+                DELETE FROM usage_stats
+                WHERE metric_key IS NOT NULL
+                  AND id NOT IN (
+                      SELECT MIN(id)
+                      FROM usage_stats
+                      WHERE metric_key IS NOT NULL
+                      GROUP BY session_id, metric_key
+                  )
+                """
+            )
             self._connection.execute(
                 """
                 CREATE UNIQUE INDEX IF NOT EXISTS usage_stats_metric_key_uq

+ 92 - 0
tests/test_sqlite_store.py

@@ -316,6 +316,70 @@ def test_sqlite_store_migrates_legacy_usage_and_turn_metrics_additively(tmp_path
     assert "session_id, metric_key" in metric_index_sql
 
 
+def test_sqlite_store_deduplicates_metric_keys_before_creating_unique_index(
+    tmp_path,
+):
+    database_path = tmp_path / "duplicate-metric-keys.sqlite3"
+    connection = sqlite3.connect(database_path)
+    connection.execute(
+        """
+        CREATE TABLE usage_stats (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            session_id TEXT NOT NULL,
+            turn_index INTEGER NOT NULL,
+            round_index INTEGER NOT NULL,
+            prompt_tokens INTEGER NOT NULL,
+            completion_tokens INTEGER NOT NULL,
+            total_tokens INTEGER NOT NULL,
+            cached_tokens INTEGER NOT NULL,
+            ttft_ms INTEGER,
+            elapsed_ms INTEGER NOT NULL,
+            metric_key TEXT,
+            created_at TEXT NOT NULL
+        )
+        """
+    )
+    connection.executemany(
+        """
+        INSERT INTO usage_stats (
+            id, session_id, turn_index, round_index, prompt_tokens,
+            completion_tokens, total_tokens, cached_tokens, ttft_ms,
+            elapsed_ms, metric_key, created_at
+        ) VALUES (?, 'legacy', 1, 1, 0, 0, 0, 0, NULL, 1, ?, 'now')
+        """,
+        [
+            (7, "chat:1:1"),
+            (11, "chat:1:1"),
+            (13, None),
+            (17, None),
+            (19, None),
+        ],
+    )
+    connection.commit()
+    connection.close()
+
+    SQLiteSessionStore(database_path).list_sessions()
+
+    migrated = sqlite3.connect(database_path)
+    rows = migrated.execute(
+        """
+        SELECT id, session_id, metric_key
+        FROM usage_stats
+        ORDER BY id
+        """
+    ).fetchall()
+    indexes = [row[1] for row in migrated.execute("PRAGMA index_list(usage_stats)")]
+    migrated.close()
+
+    assert rows == [
+        (7, "legacy", "chat:1:1"),
+        (13, "legacy", None),
+        (17, "legacy", None),
+        (19, "legacy", None),
+    ]
+    assert indexes.count("usage_stats_metric_key_uq") == 1
+
+
 def test_sqlite_store_session_list_totals_are_not_multiplied_by_turns(tmp_path):
     store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
     session_id = store.create_session(title="two turns", config={})
@@ -556,10 +620,26 @@ def test_sqlite_store_migration_is_safe_across_concurrent_initialization(
             cached_tokens INTEGER NOT NULL,
             ttft_ms INTEGER,
             elapsed_ms INTEGER NOT NULL,
+            metric_key TEXT,
             created_at TEXT NOT NULL
         )
         """
     )
+    connection.executemany(
+        """
+        INSERT INTO usage_stats (
+            id, session_id, turn_index, round_index, prompt_tokens,
+            completion_tokens, total_tokens, cached_tokens, ttft_ms,
+            elapsed_ms, metric_key, created_at
+        ) VALUES (?, 'legacy', 1, 1, 0, 0, 0, 0, NULL, 1, ?, 'now')
+        """,
+        [
+            (7, "chat:1:1"),
+            (11, "chat:1:1"),
+            (13, None),
+            (17, None),
+        ],
+    )
     connection.execute(
         """
         INSERT INTO messages
@@ -628,6 +708,13 @@ def test_sqlite_store_migration_is_safe_across_concurrent_initialization(
     metric_indexes = [
         row[1] for row in migrated.execute("PRAGMA index_list(usage_stats)")
     ]
+    usage_rows = migrated.execute(
+        """
+        SELECT id, session_id, metric_key
+        FROM usage_stats
+        ORDER BY id
+        """
+    ).fetchall()
     migrated.close()
     assert turn_columns.count("wall_time_ms") == 1
     for column_name in (
@@ -641,3 +728,8 @@ def test_sqlite_store_migration_is_safe_across_concurrent_initialization(
     ):
         assert usage_columns.count(column_name) == 1
     assert metric_indexes.count("usage_stats_metric_key_uq") == 1
+    assert usage_rows == [
+        (7, "legacy", "chat:1:1"),
+        (13, "legacy", None),
+        (17, "legacy", None),
+    ]