diff --git a/api/models.py b/api/models.py index 725fb91d..fd052923 100644 --- a/api/models.py +++ b/api/models.py @@ -3476,14 +3476,34 @@ def _sidebar_title_is_generic_webui(title: str | None) -> bool: return text.startswith(prefix) and text[len(prefix):].isdigit() -def _read_state_db_sidebar_overrides(db_path: Path, session_ids: set[str]) -> dict[str, dict]: +def _read_state_db_sidebar_overrides( + db_path: Path, + session_ids: set[str], + count_session_ids: set[str] | None = None, +) -> dict[str, dict]: """Return cheap state.db source/title overrides for sidebar rows. This intentionally does not chase lineage parents/children. It is used on the /api/sessions hot path before CLI filtering so state.db can correct stale JSON source flags without paying the full lineage-enrichment cost. + + Two-tier cost split (#5132): the ``sessions``-table lookup (source/title/ + message_count) is an indexed primary-key fetch and is run for ALL + ``session_ids`` — its result feeds the source classification that + ``/api/sessions`` filters on BEFORE the lazy lineage correction, so capping + it would silently drop rows (e.g. a stale ``cli`` JSON row whose state.db + source is ``webui``) from the default sidebar. The expensive part is the + ``messages`` aggregation (``COUNT(*)``/``MAX(timestamp)`` GROUP BY), which is + what blocked /api/sessions for 5-18s on power users; that scan is restricted + to ``count_session_ids`` (the top-N paint-priority rows). When + ``count_session_ids`` is None, both tiers cover the full set (caller opted + out of the cap). """ wanted = {str(sid) for sid in (session_ids or set()) if sid} + if count_session_ids is None: + count_wanted = set(wanted) + else: + count_wanted = {str(sid) for sid in count_session_ids if sid} & wanted if not wanted or not db_path.exists(): return {} try: @@ -3549,15 +3569,19 @@ def _read_state_db_sidebar_overrides(db_path: Path, session_ids: set[str]) -> di if entry: overrides[sid] = entry if has_messages_table and messages_has_session_id: + count_chunk = [sid for sid in chunk if sid in count_wanted] + if not count_chunk: + continue + count_placeholders = ','.join('?' * len(count_chunk)) last_at_expr = "MAX(timestamp) AS last_message_at" if messages_has_timestamp else "NULL AS last_message_at" cur.execute( f""" SELECT session_id, COUNT(*) AS actual_message_count, {last_at_expr} FROM messages - WHERE session_id IN ({placeholders}) + WHERE session_id IN ({count_placeholders}) GROUP BY session_id """, - chunk, + count_chunk, ) for row in cur.fetchall(): sid = str(row['session_id']) @@ -3582,33 +3606,32 @@ def _read_state_db_sidebar_overrides(db_path: Path, session_ids: set[str]) -> di def _apply_sidebar_state_db_overrides(sessions: list[dict]) -> None: """Apply state.db source/title overrides without full lineage enrichment. - Cap the state.db lookup to the top-N most recent sessions to bound - wall-clock on power users with thousands of sessions. This mirrors the - identical cap on _enrich_sidebar_lineage_metadata (#4638): the caller - passes an already pinned-first/newest-first sorted list, the sidebar - paints chronologically newest-first, and these source/title corrections - only matter for rows the user can actually see. The full scan on every - poll is the avoidable cost. Issue #5132 / 2026-06-28 triage: - /api/sessions blocked 5-18s in all_sessions.state_db_overrides reading - state.db for 2400+ rows on every concurrent poll, flapping the UI to - "Connection lost". The cap is env-configurable and fails open. - - Rows beyond the cap keep their JSON source/title (lazily corrected when - the user opens the history panel), exactly as with the lineage cap. + Source classification (source/title) is corrected for ALL rows because it + feeds the CLI/WebUI sidebar filter that runs BEFORE the lazy lineage + correction — capping it would silently drop rows whose stale JSON source + disagrees with state.db (#5132 regression guard). Only the expensive + ``messages`` count/last-message aggregation is capped to the top-N most + recent (paint-priority) rows, which is what actually blocked /api/sessions + for 5-18s on power users reading state.db for 2400+ rows on every + concurrent poll (#5132). The cap is env-configurable and fails open; rows + beyond it keep their JSON message-count/last-message until the history panel + opens (lazily corrected, exactly as with the lineage cap #4638). """ import os as _os try: _cap = int(_os.environ.get("HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N", "300")) except (TypeError, ValueError): _cap = 300 + all_ids = {str(s.get('session_id')) for s in sessions if s.get('session_id')} if _cap > 0 and len(sessions) > _cap: - candidates = sessions[:_cap] + count_ids = {str(s.get('session_id')) for s in sessions[:_cap] if s.get('session_id')} else: - candidates = sessions + count_ids = None # cap disabled / under cap -> count every row too try: metadata = _read_state_db_sidebar_overrides( _active_state_db_path(), - {str(s.get('session_id')) for s in candidates if s.get('session_id')}, + all_ids, + count_session_ids=count_ids, ) except Exception: return diff --git a/tests/test_issue5132_state_db_override_top_n_cap.py b/tests/test_issue5132_state_db_override_top_n_cap.py index 362b2701..23e1908d 100644 --- a/tests/test_issue5132_state_db_override_top_n_cap.py +++ b/tests/test_issue5132_state_db_override_top_n_cap.py @@ -1,19 +1,27 @@ -"""Regression tests for #5132 — cap sidebar state.db source/title overrides to top-N. +"""Regression tests for #5132 — cap sidebar state.db count/title overrides to top-N. On power users with thousands of sessions, GET /api/sessions blocked 5-18s in the all_sessions.state_db_overrides stage: _apply_sidebar_state_db_overrides probed -state.db for EVERY row's source/title (2400+ ids) on every concurrent poll, piling -up sqlite reads against the gateway-watcher + agent writes and flapping the UI to -"Connection lost". The sidebar paints pinned-first then newest-first, and the caller -passes an already-sorted list, so overriding only the top-N (default 300) most-recent -rows covers the visible window while bounding wall-clock — exactly the precedent set -for lineage enrichment in #4638. The cap is env-configurable -(HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N) and fails open. +state.db for EVERY row (2400+ ids) on every concurrent poll, piling up sqlite reads +against the gateway-watcher + agent writes and flapping the UI to "Connection lost". -These tests pin: (1) only the top-N ids are probed when the list exceeds the cap, -(2) the env override is honored, (3) a non-positive / unparseable cap disables the -cap (override all), (4) lists at/under the cap probe everything, and (5) a DB error -never breaks /api/sessions. +The fix is a TWO-TIER split (the cap on the whole probe was the first cut and was a +regression — see test_source_classification_is_never_capped below): + + * SOURCE classification (source/title) is read for ALL rows. It feeds the CLI/WebUI + sidebar filter that runs BEFORE the lazy lineage correction, so capping it would + silently drop rows whose stale JSON source disagrees with state.db. + * The expensive MESSAGES aggregation (COUNT(*)/MAX(timestamp) GROUP BY) — the actual + bottleneck — is capped to the top-N (default 300) paint-priority rows. + +The cap is env-configurable (HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N) and fails open. + +These tests pin: (1) source ids are NEVER capped; (2) count ids ARE capped to top-N; +(3) the env override bounds the COUNT tier; (4) a non-positive / unparseable cap; +(5) lists at/under the cap; (6) a DB error never breaks /api/sessions; (7) a probed +row's webui source is applied; and (8) the end-to-end route regression: a stale-CLI +JSON row whose state.db source is webui must stay visible in a WebUI-filtered sidebar +even when it falls beyond the cap. """ from __future__ import annotations @@ -21,11 +29,12 @@ import api.models as models def _capture_probed_ids(monkeypatch): - """Patch _read_state_db_sidebar_overrides to record the id set it is asked for.""" + """Patch _read_state_db_sidebar_overrides to record both id tiers it is asked for.""" seen = {} - def _fake_read(db_path, id_set): - seen["ids"] = set(id_set) + def _fake_read(db_path, id_set, count_session_ids=None): + seen["source_ids"] = set(id_set) + seen["count_ids"] = None if count_session_ids is None else set(count_session_ids) return {} monkeypatch.setattr(models, "_read_state_db_sidebar_overrides", _fake_read) @@ -39,21 +48,34 @@ def _sessions(n): return [{"session_id": f"s{i}"} for i in range(n)] -def test_caps_overrides_to_top_n_default_300(monkeypatch): +def test_source_classification_is_never_capped(monkeypatch): + """SOURCE tier must cover ALL rows (capping it dropped rows from the sidebar — #5132 regression).""" seen = _capture_probed_ids(monkeypatch) monkeypatch.delenv("HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N", raising=False) models._apply_sidebar_state_db_overrides(_sessions(1000)) - assert seen["ids"] == {f"s{i}" for i in range(300)}, ( - "Default cap must probe exactly the top-300 (paint-priority) sessions" + assert seen["source_ids"] == {f"s{i}" for i in range(1000)}, ( + "Source/title classification must be read for EVERY row, never capped" ) -def test_env_override_changes_cap(monkeypatch): +def test_count_aggregation_capped_to_top_n_default_300(monkeypatch): + seen = _capture_probed_ids(monkeypatch) + monkeypatch.delenv("HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N", raising=False) + models._apply_sidebar_state_db_overrides(_sessions(1000)) + assert seen["count_ids"] == {f"s{i}" for i in range(300)}, ( + "The expensive message-count aggregation must be capped to the top-300" + ) + + +def test_env_override_changes_count_cap(monkeypatch): seen = _capture_probed_ids(monkeypatch) monkeypatch.setenv("HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N", "50") models._apply_sidebar_state_db_overrides(_sessions(1000)) - assert seen["ids"] == {f"s{i}" for i in range(50)}, ( - "HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N must bound the probed set" + assert seen["count_ids"] == {f"s{i}" for i in range(50)}, ( + "HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N must bound the COUNT tier" + ) + assert seen["source_ids"] == {f"s{i}" for i in range(1000)}, ( + "...but the source tier stays uncapped" ) @@ -61,30 +83,30 @@ def test_non_positive_cap_disables_capping(monkeypatch): seen = _capture_probed_ids(monkeypatch) monkeypatch.setenv("HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N", "0") models._apply_sidebar_state_db_overrides(_sessions(500)) - assert len(seen["ids"]) == 500, "cap<=0 must override all sessions (cap disabled)" + assert seen["count_ids"] is None, "cap<=0 must disable the count cap (count every row)" + assert len(seen["source_ids"]) == 500 def test_unparseable_cap_falls_back_to_default(monkeypatch): seen = _capture_probed_ids(monkeypatch) monkeypatch.setenv("HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N", "not-a-number") models._apply_sidebar_state_db_overrides(_sessions(1000)) - assert seen["ids"] == {f"s{i}" for i in range(300)}, ( + assert seen["count_ids"] == {f"s{i}" for i in range(300)}, ( "An unparseable cap must fall back to the default 300, not crash" ) -def test_list_under_cap_probes_everything(monkeypatch): +def test_list_under_cap_counts_everything(monkeypatch): seen = _capture_probed_ids(monkeypatch) monkeypatch.delenv("HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N", raising=False) models._apply_sidebar_state_db_overrides(_sessions(120)) - assert seen["ids"] == {f"s{i}" for i in range(120)}, ( - "A list at/under the cap must override every session" - ) + assert seen["count_ids"] is None, "A list at/under the cap counts every row (no cap needed)" + assert seen["source_ids"] == {f"s{i}" for i in range(120)} def test_override_failure_is_swallowed(monkeypatch): """Override lookup must fail open — a DB error never breaks /api/sessions.""" - def _boom(db_path, id_set): + def _boom(db_path, id_set, count_session_ids=None): raise RuntimeError("db down") monkeypatch.setattr(models, "_read_state_db_sidebar_overrides", _boom) @@ -93,12 +115,12 @@ def test_override_failure_is_swallowed(monkeypatch): models._apply_sidebar_state_db_overrides(_sessions(10)) -def test_capped_rows_still_receive_overrides(monkeypatch): - """The top-N rows that ARE probed must still get their metadata applied.""" - def _fake_read(db_path, id_set): - # Return a webui source override for s0 only. +def test_capped_rows_still_receive_source_overrides(monkeypatch): + """A row beyond the COUNT cap must still get its SOURCE metadata applied.""" + def _fake_read(db_path, id_set, count_session_ids=None): + # Return a webui source override for s400 (beyond the default 300 cap). return { - "s0": { + "s400": { "_state_db_source": "webui", "_state_db_source_tag": "webui", "_state_db_raw_source": "webui", @@ -109,9 +131,43 @@ def test_capped_rows_still_receive_overrides(monkeypatch): monkeypatch.setattr(models, "_read_state_db_sidebar_overrides", _fake_read) monkeypatch.setattr(models, "_active_state_db_path", lambda: ":memory:") - sessions: list[dict] = _sessions(400) - sessions[0]["is_cli_session"] = True + sessions: list[dict] = _sessions(500) + sessions[400]["is_cli_session"] = True models._apply_sidebar_state_db_overrides(sessions) - assert sessions[0]["is_cli_session"] is False, ( - "A probed top-N row with a webui state.db source must be corrected" + assert sessions[400]["is_cli_session"] is False, ( + "A row beyond the count cap with a webui state.db source must STILL be corrected " + "(source classification is uncapped)" + ) + + +def test_stale_cli_json_beyond_cap_stays_webui_via_real_db(monkeypatch, tmp_path): + """End-to-end #5132 regression: a stale-CLI JSON row whose state.db source is webui, + sitting BEYOND the top-N cap, must keep source='webui' so it isn't filtered out of a + WebUI-only sidebar. This is the exact row that vanished under the first (cap-everything) cut. + """ + import sqlite3 + + db = tmp_path / "state.db" + conn = sqlite3.connect(str(db)) + conn.execute("CREATE TABLE sessions (id TEXT PRIMARY KEY, source TEXT, session_source TEXT, title TEXT, message_count INTEGER)") + # The stale row lives at index 400 (beyond the 300 cap). state.db says it's webui. + conn.execute("INSERT INTO sessions (id, source, title, message_count) VALUES (?,?,?,?)", + ("s400", "webui", "Stale CLI row", 3)) + conn.commit() + conn.close() + + monkeypatch.setattr(models, "_active_state_db_path", lambda: db) + monkeypatch.delenv("HERMES_WEBUI_STATE_DB_OVERRIDE_TOP_N", raising=False) + + sessions = _sessions(500) + # JSON metadata wrongly marks it CLI. + sessions[400]["is_cli_session"] = True + sessions[400]["session_source"] = "cli" + models._apply_sidebar_state_db_overrides(sessions) + + assert sessions[400]["is_cli_session"] is False, ( + "state.db source=webui must override the stale CLI JSON flag even beyond the cap" + ) + assert sessions[400].get("session_source") == "webui", ( + "session_source must be reclassified to webui from state.db" )