|
|
@@ -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),
|
|
|
+ ]
|