diff --git a/studio/backend/storage/profile_stats_db.py b/studio/backend/storage/profile_stats_db.py index 6f0d05de96..f5c6795ea3 100644 --- a/studio/backend/storage/profile_stats_db.py +++ b/studio/backend/storage/profile_stats_db.py @@ -48,13 +48,19 @@ _cache: dict[str, Any] = {"fingerprint": None, "expires_at": 0.0, "payload": Non def _as_float(value: Any) -> Optional[float]: - """Coerce JSON numbers defensively; metadata is written by the client.""" + """Coerce JSON numbers defensively; metadata is written by the client. + + json accepts integers of any width, and float() raises OverflowError past + ~1e308, so one oversized counter would 500 the whole panel. + """ if isinstance(value, bool) or value is None: return None if isinstance(value, (int, float)): - return ( - float(value) if value == value and value not in (float("inf"), float("-inf")) else None - ) + try: + number = float(value) + except (OverflowError, ValueError): + return None + return number if number == number and number not in (float("inf"), float("-inf")) else None return None @@ -187,33 +193,39 @@ class _MessageFold: bucket["threads"].add(thread_id) -def _fork_keepers(conn) -> set[str]: - """One fork per source, elected to stand in for originals that are gone. +def _fork_keepers(conn) -> dict[tuple[str, int, str], str]: + """For each original message, the one clone elected to stand in for it. - A clone is normally ignored because the original is counted instead. When - the original is not there any more, whether its thread was deleted or just - that message was pruned, the clones become the only record. Letting every - sibling count them would multiply the usage, so exactly one fork per source - is allowed to: the one holding the most copied rows, which is the fork that - branched latest and therefore carries the longest ancestry. + A clone is normally ignored because the original is counted instead. Once + the original is gone, whether its thread was deleted or just that row was + pruned, the clones become the only record. Letting every sibling count them + would multiply the usage, so exactly one may. + + Electing per message rather than per fork matters because fork_chat_thread + copies one parent_id branch, not the whole thread: sibling forks taken from + a retry and a regeneration hold different rows, and a per-fork winner would + silently drop whatever only the loser carries. """ 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 + SELECT m.thread_id, m.created_at, m.role, + t.forked_from_thread_id AS source_id + FROM chat_messages m + JOIN chat_threads t ON t.id = m.thread_id WHERE t.forked_from_thread_id IS NOT NULL + AND m.created_at < t.created_at """ - ).fetchall() + ) - best: dict[str, tuple[int, str]] = {} + best: dict[tuple[str, int, str], 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()} + key = (row["source_id"], _as_int(row["created_at"]), row["role"]) + thread_id = row["thread_id"] + current = best.get(key) + # Any stable winner works; lowest id keeps the choice reproducible. + if current is None or thread_id < current: + best[key] = thread_id + return best def _surviving_original_keys(conn) -> set[tuple[str, int, str]]: @@ -296,7 +308,7 @@ def _fold_messages(conn, zone) -> _MessageFold: source_id = row["forked_from_thread_id"] if source_id and created_at < _as_int(row["thread_created_at"]): original = (source_id, created_at, row["role"]) - if original in surviving or thread_id not in keepers: + if original in surviving or keepers.get(original) != thread_id: continue fold.messages += 1 @@ -423,6 +435,11 @@ def _superseded(prefix: str = "r.") -> str: ``delete_run`` never clears the flag, so the continuation has to still be there. Otherwise deleting it would strand the source at zero while its row and metrics stay visible in history. + + The continuation also has to have reached the source's step. ``create_run`` + claims the source the moment a resume starts, but ``final_step`` is only + written on the first metric flush, so a continuation that fails before then + would take the source's completed work down with it. """ return f""" {prefix}resume_blocked = 1 @@ -431,6 +448,8 @@ def _superseded(prefix: str = "r.") -> str: SELECT 1 FROM training_runs continuation WHERE continuation.output_dir = {prefix}output_dir AND continuation.started_at > {prefix}started_at + AND COALESCE(continuation.final_step, 0) + >= COALESCE({prefix}final_step, 0) ) """ diff --git a/studio/backend/tests/test_profile_stats.py b/studio/backend/tests/test_profile_stats.py index e36c5440d1..489f53625e 100644 --- a/studio/backend/tests/test_profile_stats.py +++ b/studio/backend/tests/test_profile_stats.py @@ -887,3 +887,120 @@ def test_route_does_not_block_the_event_loop(stats_db, monkeypatch): # ~50 ticks fit in 0.5s; a blocking call on the loop would yield 0. assert ticks > 10, f"event loop stalled during stats computation ({ticks} ticks)" + + +def test_an_oversized_token_counter_does_not_break_the_panel(stats_db): + """json parses ints of any width; float() gives up long before that.""" + now = datetime.now().replace(hour = 12, minute = 0, second = 0, microsecond = 0) + conn = studio_db.get_connection() + try: + _seed_thread(conn, "sane", "m", [(now - timedelta(hours = 1), _metadata(100, 50))]) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, model_id, created_at, updated_at) " + "VALUES ('bad', 'Bad', 'base', 'm', ?, ?)", + (_ms(now), _ms(now)), + ) + huge = _metadata(100, 50) + # Wider than float can hold, so every counter reads as unusable. + oversized = int("9" * 309) + for field in ("promptTokens", "completionTokens", "totalTokens"): + huge["contextUsage"][field] = oversized + huge["timing"]["tokenCount"] = oversized + conn.execute( + "INSERT INTO chat_messages (id, thread_id, role, content_json, metadata_json, " + "created_at) VALUES ('bad-a0', 'bad', 'assistant', '[]', ?, ?)", + (json.dumps(huge), _ms(now)), + ) + conn.commit() + finally: + conn.close() + + # The row degrades to zero instead of raising, and the absurd counter + # never reaches the totals; the healthy thread still reports. + stats = compute_profile_stats(days = 7) + assert stats["totals"]["totalTokens"] == 150 + assert stats["totals"]["messages"] == 3 + + +def test_divergent_sibling_forks_keep_their_own_branch_messages(stats_db): + """fork_chat_thread copies one parent_id branch, so siblings differ.""" + 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))]) + # A regeneration of the same turn, a second later. + regenerated = older + REPLY_DELAY + timedelta(seconds = 1) + conn.execute( + "INSERT INTO chat_messages (id, thread_id, role, content_json, metadata_json, " + "created_at) VALUES ('root-a1', 'root', 'assistant', '[]', ?, ?)", + (json.dumps(_metadata(200, 100)), _ms(regenerated)), + ) + # One fork per branch: each carries only its own reply. + for fork_id, stamp, meta in ( + ("forkA", older + REPLY_DELAY, _metadata(100, 50)), + ("forkB", regenerated, _metadata(200, 100)), + ): + 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(meta), _ms(stamp)), + ) + conn.commit() + finally: + conn.close() + + # Both originals counted, both clones suppressed. + assert compute_profile_stats(days = 7)["totals"]["totalTokens"] == 450 + + 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) + + # Each branch survives in exactly one fork, so nothing is lost or doubled. + assert stats["totals"]["totalTokens"] == 450 + assert stats["totals"]["messages"] == 2 + + +def test_a_resume_that_never_logged_a_step_keeps_the_source_counters(stats_db): + """create_run claims the source before the continuation flushes a metric.""" + conn = studio_db.get_connection() + try: + conn.execute( + "INSERT INTO training_runs (id, status, model_name, dataset_name, config_json, " + "started_at, total_steps, final_step, duration_seconds, output_dir, resume_blocked) " + "VALUES ('src', 'stopped', 'm', 'd', '{}', '2026-01-01T10:00:00', 20, 10, 600, " + "'/runs/out', 1)", + ) + # Errored before its first training step: no final_step, no metrics. + conn.execute( + "INSERT INTO training_runs (id, status, model_name, dataset_name, config_json, " + "started_at, total_steps, final_step, duration_seconds, output_dir, resume_blocked) " + "VALUES ('cont', 'error', 'm', 'd', '{}', '2026-01-02T10:00:00', 20, NULL, 5, " + "'/runs/out', 0)", + ) + conn.executemany( + "INSERT INTO training_metrics (run_id, step, num_tokens) VALUES (?, ?, ?)", + [("src", step, step * 100) for step in range(1, 11)], + ) + conn.commit() + finally: + conn.close() + + training = compute_profile_stats(days = 7)["training"] + + # The source's completed work is still the only work there is. + assert training["steps"] == 10 + assert training["tokens"] == 1000 diff --git a/studio/frontend/src/features/profile/components/stats/token-activity-card.tsx b/studio/frontend/src/features/profile/components/stats/token-activity-card.tsx index 4e6b1230bd..7a11ec7769 100644 --- a/studio/frontend/src/features/profile/components/stats/token-activity-card.tsx +++ b/studio/frontend/src/features/profile/components/stats/token-activity-card.tsx @@ -12,6 +12,7 @@ import { heatLevel, parseDayKey, seriesForMode, + windowBaseline, } from "../../utils/stats-format"; import { StatsCard } from "./stat-primitives"; @@ -39,6 +40,7 @@ function buildColumns( daily: ProfileStatsDay[], values: number[], columns: number, + mode: ActivityMode, ): Cell[][] { if (daily.length === 0 || columns <= 0) return []; @@ -52,6 +54,9 @@ function buildColumns( const capacity = columns * DAYS_PER_WEEK - trailing; const start = Math.max(0, daily.length - capacity); const visible = daily.slice(start); + // Cumulative is a running total over what the grid shows, so a narrow card + // that drops older weeks has to rebase off the last hidden day. + const baseline = windowBaseline(values, start, mode); const cells: Cell[] = []; // Pad so every column is a Monday-started week. @@ -62,7 +67,11 @@ function buildColumns( cells.push({ key: `pad-${index}`, day: null, value: 0 }); } for (const [index, day] of visible.entries()) { - cells.push({ key: day.date, day, value: values[start + index] ?? 0 }); + cells.push({ + key: day.date, + day, + value: (values[start + index] ?? 0) - baseline, + }); } const grid: Cell[][] = []; @@ -245,8 +254,8 @@ export function TokenActivityCard({ daily }: { daily: ProfileStatsDay[] }) { const shaded = mode === "daily"; const values = useMemo(() => seriesForMode(daily, mode), [daily, mode]); const grid = useMemo( - () => buildColumns(daily, values, columns), - [daily, values, columns], + () => buildColumns(daily, values, columns, mode), + [daily, values, columns, mode], ); // The app language, not the browser's: those differ whenever the user picks // a language in Settings, and it has to be a dependency so switching while diff --git a/studio/frontend/src/features/profile/utils/stats-format.ts b/studio/frontend/src/features/profile/utils/stats-format.ts index aee026f4c1..4f1b1ea16f 100644 --- a/studio/frontend/src/features/profile/utils/stats-format.ts +++ b/studio/frontend/src/features/profile/utils/stats-format.ts @@ -103,6 +103,20 @@ export type ActivityMode = "daily" | "weekly" | "cumulative"; * since the backend caps the series and seeding it with everything older would * flatten every bar against a baseline the grid cannot show. */ +/** + * What to subtract from a cumulative series once the grid drops older days. + * Without it the first visible bar opens at the hidden total and the whole + * window flattens against a baseline the user cannot see. + */ +export function windowBaseline( + values: number[], + start: number, + mode: ActivityMode, +): number { + if (mode !== "cumulative" || start <= 0) return 0; + return values[start - 1] ?? 0; +} + export function seriesForMode( daily: Array<{ date: string; tokens: number }>, mode: ActivityMode, diff --git a/studio/frontend/tests/profile-stats-format.test.ts b/studio/frontend/tests/profile-stats-format.test.ts index 7bf464ce97..16a5ed13e7 100644 --- a/studio/frontend/tests/profile-stats-format.test.ts +++ b/studio/frontend/tests/profile-stats-format.test.ts @@ -11,6 +11,7 @@ import { heatLevel, parseDayKey, seriesForMode, + windowBaseline, } from "../src/features/profile/utils/stats-format.ts"; test("compact numbers match the tile format", () => { @@ -79,3 +80,34 @@ test("series modes reshape the same daily data", () => { assert.deepEqual(seriesForMode(daily, "weekly"), [35, 35, 35, 100]); assert.deepEqual(seriesForMode([], "weekly"), []); }); + +test("a trimmed cumulative window rebases off the last hidden day", () => { + const daily = [ + { date: "2026-01-01", tokens: 1000 }, + { date: "2026-01-02", tokens: 2000 }, + { date: "2026-01-03", tokens: 5 }, + { date: "2026-01-04", tokens: 10 }, + ]; + const values = seriesForMode(daily, "cumulative"); + assert.deepEqual(values, [1000, 3000, 3005, 3015]); + + // Showing only the last two days: without rebasing, both bars sit at ~3000 + // and the 5 vs 10 difference is invisible. + const baseline = windowBaseline(values, 2, "cumulative"); + assert.equal(baseline, 3000); + assert.deepEqual( + values.slice(2).map((value) => value - baseline), + [5, 15], + ); +}); + +test("nothing is rebased when the window shows everything", () => { + const values = [1, 3, 6]; + assert.equal(windowBaseline(values, 0, "cumulative"), 0); +}); + +test("only cumulative rebases: daily and weekly are already per-window", () => { + const values = [10, 20, 5]; + assert.equal(windowBaseline(values, 2, "daily"), 0); + assert.equal(windowBaseline(values, 2, "weekly"), 0); +});