diff --git a/studio/backend/storage/profile_stats_db.py b/studio/backend/storage/profile_stats_db.py index 041d1e2d6e..8c1332c46e 100644 --- a/studio/backend/storage/profile_stats_db.py +++ b/studio/backend/storage/profile_stats_db.py @@ -186,13 +186,44 @@ class _MessageFold: bucket["threads"].add(thread_id) +def _orphan_fork_keepers(conn) -> set[str]: + """Forks whose source is gone and that should still count their copies. + + ``forked_from_thread_id`` is not a foreign key, so deleting a source leaves + its forks holding the only copies of that ancestry. Enabling all of them + would multiply the usage by the number of siblings, so exactly one per dead + source is kept: the one carrying the most copied rows, which is the fork + that branched latest and therefore holds the longest ancestry. + """ + rows = conn.execute( + """ + SELECT t.id, t.forked_from_thread_id AS source_id, + (SELECT COUNT(*) FROM chat_messages m + WHERE m.thread_id = t.id AND m.created_at < t.created_at) AS copied + FROM chat_threads t + WHERE t.forked_from_thread_id IS NOT NULL + AND NOT EXISTS (SELECT 1 FROM chat_threads s WHERE s.id = t.forked_from_thread_id) + """ + ).fetchall() + + best: dict[str, tuple[int, str]] = {} + for row in rows: + rank = (_as_int(row["copied"]), row["id"]) + current = best.get(row["source_id"]) + if current is None or rank > current: + best[row["source_id"]] = rank + return {thread_id for _, thread_id in best.values()} + + def _fold_messages(conn, zone) -> _MessageFold: fold = _MessageFold() + keepers = _orphan_fork_keepers(conn) rows = conn.execute( """ SELECT m.thread_id, m.role, m.metadata_json, m.attachments_json, m.created_at, - t.title, t.model_id, t.model_type, - t.created_at AS thread_created_at, src.id AS fork_source_id + t.title, t.model_id, t.model_type, t.pair_id, + t.created_at AS thread_created_at, t.forked_from_thread_id, + src.id AS fork_source_id FROM chat_messages m LEFT JOIN chat_threads t ON t.id = m.thread_id -- forked_from_thread_id is not a foreign key, so it outlives the source. @@ -232,16 +263,26 @@ def _fold_messages(conn, zone) -> _MessageFold: thread_messages = 0 previous_created = None + # Compare mode stores one thread per pane under a shared pair_id, and + # the sidebar shows them as a single conversation. Count them that way + # too, or one comparison inflates the chat total and drags the average + # tokens per chat down. + conversation_id = row["pair_id"] or thread_id + # A fork is its own visible conversation, so it counts towards the chat # total from the moment it exists, before any new turn is added. - fold.threads.add(thread_id) + fold.threads.add(conversation_id) # Forking clones the whole ancestry into the new thread, keeping each # copy's original timestamp. Counting those again would double every # metric for the branched-from conversation, so skip anything older - # than the fork itself, but only while the source survives to be - # counted in its place. - if row["fork_source_id"] and created_at < _as_int(row["thread_created_at"]): + # than the fork itself. Once the source is gone one fork is elected to + # keep its copies, since they are then the only record left. + if ( + row["forked_from_thread_id"] + and created_at < _as_int(row["thread_created_at"]) + and (row["fork_source_id"] or thread_id not in keepers) + ): continue fold.messages += 1 @@ -333,7 +374,7 @@ def _fold_messages(conn, zone) -> _MessageFold: if stamp is not None: fold.by_hour[stamp.hour] += 1 fold.by_weekday[stamp.weekday()] += 1 - fold.note_day(stamp.date(), message_tokens, thread_id) + fold.note_day(stamp.date(), message_tokens, conversation_id) close_thread() return fold diff --git a/studio/backend/tests/test_profile_stats.py b/studio/backend/tests/test_profile_stats.py index ed5af2aff4..f5a00c3ec9 100644 --- a/studio/backend/tests/test_profile_stats.py +++ b/studio/backend/tests/test_profile_stats.py @@ -560,6 +560,74 @@ def test_deleting_the_source_thread_keeps_the_forks_copies(stats_db): assert stats["totals"]["messages"] == 1 +def test_sibling_forks_do_not_multiply_a_deleted_source(stats_db): + """Two forks of one thread must not both re-count the shared ancestry.""" + now = datetime.now().replace(hour = 12, minute = 0, second = 0, microsecond = 0) + older = now - timedelta(hours = 3) + conn = studio_db.get_connection() + try: + _seed_thread(conn, "root", "m", [(older, _metadata(100, 50))]) + for fork_id in ("forkA", "forkB"): + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, model_id, created_at, " + "updated_at, forked_from_thread_id, forked_from_message_id) " + "VALUES (?, 'fork', 'base', 'm', ?, ?, 'root', 'root-a0')", + (fork_id, _ms(now), _ms(now)), + ) + conn.execute( + "INSERT INTO chat_messages (id, thread_id, role, content_json, " + "metadata_json, created_at) VALUES (?, ?, 'assistant', '[]', ?, ?)", + (f"{fork_id}-a0", fork_id, json.dumps(_metadata(100, 50)), _ms(older)), + ) + conn.commit() + finally: + conn.close() + + assert compute_profile_stats(days = 7)["totals"]["totalTokens"] == 150 + + conn = studio_db.get_connection() + try: + conn.execute("DELETE FROM chat_threads WHERE id = 'root'") + conn.commit() + finally: + conn.close() + + invalidate_profile_stats_cache() + stats = compute_profile_stats(days = 7) + + # Exactly one surviving copy is counted, not one per sibling fork. + assert stats["totals"]["totalTokens"] == 150 + assert stats["totals"]["messages"] == 1 + + +def test_comparison_panes_count_as_one_chat(stats_db): + """Compare mode stores a thread per pane; the sidebar shows one chat.""" + now = datetime.now() + conn = studio_db.get_connection() + try: + for pane in ("left", "right"): + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, model_id, pair_id, " + "created_at, updated_at) VALUES (?, 'compare', 'base', 'm', 'pair-1', ?, ?)", + (pane, _ms(now), _ms(now)), + ) + conn.execute( + "INSERT INTO chat_messages (id, thread_id, role, content_json, " + "metadata_json, created_at) VALUES (?, ?, 'assistant', '[]', ?, ?)", + (f"{pane}-a0", pane, json.dumps(_metadata(100, 50)), _ms(now)), + ) + conn.commit() + finally: + conn.close() + + stats = compute_profile_stats(days = 7) + + assert stats["totals"]["threads"] == 1 + # Both panes still contribute their own messages and tokens. + assert stats["totals"]["messages"] == 2 + assert stats["totals"]["totalTokens"] == 300 + + def test_deleting_a_continuation_restores_its_source(stats_db): """delete_run leaves resume_blocked set, so supersession needs a live tail.""" conn = studio_db.get_connection() diff --git a/studio/frontend/src/features/profile/utils/stats-format.ts b/studio/frontend/src/features/profile/utils/stats-format.ts index 518585e213..aee026f4c1 100644 --- a/studio/frontend/src/features/profile/utils/stats-format.ts +++ b/studio/frontend/src/features/profile/utils/stats-format.ts @@ -99,7 +99,9 @@ export type ActivityMode = "daily" | "weekly" | "cumulative"; /** * Recast the dense daily series for the selected mode. Weekly sums each * calendar week onto its days so the grid shows week-level intensity; - * cumulative shows the running lifetime total, which only ever grows. + * cumulative is the running total across the displayed window, not lifetime, + * since the backend caps the series and seeding it with everything older would + * flatten every bar against a baseline the grid cannot show. */ export function seriesForMode( daily: Array<{ date: string; tokens: number }>,