Release YB (v0.51.672): pre-aggregate cron sidebar ordering for missing-index scans (#4962)

Release YB (v0.51.672): pre-aggregate cron sidebar ordering for missing-index scans (#4962)

CHANGELOG.md

@@ -3,6 +3,12 @@
 
 ## [Unreleased]
 
+## [v0.51.672] — 2026-06-26 — Release YB (faster cron sidebar rebuild on stores missing the message index)
+
+### Fixed
+
+- **The cron/CLI sidebar rebuild is much faster on a `state.db` that lacks the message-session index.** When `idx_messages_session` is missing (an unmigrated/older store), the cron-only capped rescue scan ran a correlated per-row latest-message subquery that could take multiple seconds cold. That pass now uses a single pre-aggregated latest-message ordering (one grouped scan) for the missing-index case only; the common indexed path is unchanged (the correlated ordering is still the faster plan on a healthy migrated store), and the session set/order is identical to before. Follow-up to the #4842 chain, independent of #4952. Thanks @rodboev. (#4962)
+
 ## [v0.51.671] — 2026-06-26 — Release YA (sidebar polls no longer stall behind a slow session-list rebuild)
 
 ### Fixed

api/agent_sessions.py

@@ -492,25 +492,26 @@ def read_importable_agent_session_rows(
         use_messages_join = messages_has_session_id
         count_col = 'id' if 'id' in message_cols else 'session_id'
 
-        # Defensive index prime (#3887). The candidate-ordering query below sorts
-        # sessions by a correlated ``MAX(mx.timestamp)`` subquery over ``messages``.
-        # That is fast only when the agent's standard
-        # ``idx_messages_session ON messages(session_id, timestamp)`` index exists.
-        # A normally-migrated hermes-agent state.db has it, but a db that lost its
-        # migrations (older hermes-agent, or a hand-rebuilt/reimported db) does
-        # not — and the subquery then degrades to a full ``messages`` scan per
-        # candidate session, stalling ``/api/sessions`` for seconds on every
-        # refresh (the 5s-TTL cache never settles). Priming the index is a no-op
-        # (~free) when it already exists, and self-heals an affected db in
-        # milliseconds. Best-effort: degrade silently on a read-only db or any
-        # error so the listing never fails because of the prime.
+        # Defensive index prime (#3887). The normal candidate-ordering shape uses
+        # the agent's standard ``idx_messages_session ON messages(session_id,
+        # timestamp)`` index; without it, large cron-only scans degrade badly.
+        # Writable dbs self-heal by recreating the index. Read-only or locked dbs
+        # fall back to the pre-aggregated cron-only path below instead of failing.
+        messages_index_present = False
         if messages_has_session_id and messages_has_timestamp:
             try:
-                cur.execute(
-                    "CREATE INDEX IF NOT EXISTS idx_messages_session "
-                    "ON messages(session_id, timestamp)"
-                )
-                conn.commit()
+                cur.execute("PRAGMA index_list(messages)")
+                messages_index_present = any(str(row[1]) == "idx_messages_session" for row in cur.fetchall())
+            except sqlite3.Error:
+                messages_index_present = False
+            try:
+                if not messages_index_present:
+                    cur.execute(
+                        "CREATE INDEX IF NOT EXISTS idx_messages_session "
+                        "ON messages(session_id, timestamp)"
+                    )
+                    conn.commit()
+                    messages_index_present = True
             except sqlite3.Error:
                 pass  # read-only db / locked / older schema — degrade gracefully
 
@@ -532,21 +533,13 @@ def read_importable_agent_session_rows(
             join_clause = ""
             group_by_clause = ""
 
-        if use_messages_join and messages_has_timestamp:
-            order_by_clause = "ORDER BY COALESCE(MAX(m.timestamp), s.started_at) DESC"
-            candidate_order_clause = (
-                "ORDER BY COALESCE(\n"
-                "                        (SELECT MAX(mx.timestamp) FROM messages mx WHERE mx.session_id = s.id),\n"
-                "                        s.started_at\n"
-                "                    ) DESC,\n"
-                "                    s.started_at DESC"
-            )
-        else:
-            order_by_clause = "ORDER BY s.started_at DESC"
-            candidate_order_clause = "ORDER BY s.started_at DESC"
+        order_by_clause = "ORDER BY s.started_at DESC"
+        latest_messages_cte = None
+        candidate_order_clause = "ORDER BY s.started_at DESC"
 
         where_clauses = ["s.source IS NOT NULL"]
         params: list[object] = []
+        included = ()
         if include_sources:
             included = tuple(str(source) for source in include_sources if source)
             if included:
@@ -560,6 +553,32 @@ def read_importable_agent_session_rows(
                 where_clauses.append(f"s.source NOT IN ({placeholders})")
                 params.extend(excluded)
 
+        use_preaggregated_candidate_order = (
+            use_messages_join
+            and messages_has_timestamp
+            and included == ("cron",)
+            and not messages_index_present
+        )
+        if use_preaggregated_candidate_order:
+            order_by_clause = "ORDER BY COALESCE(MAX(m.timestamp), s.started_at) DESC"
+            latest_messages_cte = (
+                "latest_messages AS (\n"
+                "                    SELECT mx.session_id AS session_id, MAX(mx.timestamp) AS last_message_at\n"
+                "                    FROM messages mx\n"
+                "                    GROUP BY mx.session_id\n"
+                "                )"
+            )
+            candidate_order_clause = "ORDER BY COALESCE(lm.last_message_at, s.started_at) DESC, s.started_at DESC"
+        elif use_messages_join and messages_has_timestamp:
+            order_by_clause = "ORDER BY COALESCE(MAX(m.timestamp), s.started_at) DESC"
+            candidate_order_clause = (
+                "ORDER BY COALESCE(\n"
+                "                        (SELECT MAX(mx.timestamp) FROM messages mx WHERE mx.session_id = s.id),\n"
+                "                        s.started_at\n"
+                "                    ) DESC,\n"
+                "                    s.started_at DESC"
+            )
+
         select_sql = f"""
             SELECT s.id, s.title, s.model, s.message_count,
                    s.started_at, s.source,
@@ -592,15 +611,38 @@ def read_importable_agent_session_rows(
             # Oversampling preserves room for hidden compression segments or
             # other rows filtered after projection.
             candidate_limit = max(result_limit * 8, result_limit)
+            if latest_messages_cte:
+                candidate_cte = (
+                    "WITH {latest_messages_cte}, candidates AS (\n"
+                    "                    SELECT s.id\n"
+                    "                    FROM sessions s\n"
+                    "                    LEFT JOIN latest_messages lm ON lm.session_id = s.id\n"
+                    "                    WHERE {where_clause}\n"
+                    "                    {candidate_order_clause}\n"
+                    "                    LIMIT ?\n"
+                    "                )"
+                ).format(
+                    latest_messages_cte=latest_messages_cte,
+                    where_clause=" AND ".join(where_clauses),
+                    candidate_order_clause=candidate_order_clause,
+                )
+            else:
+                candidate_cte = (
+                    "WITH candidates AS (\n"
+                    "                    SELECT s.id\n"
+                    "                    FROM sessions s\n"
+                    "                    WHERE {where_clause}\n"
+                    "                    {candidate_order_clause}\n"
+                    "                    LIMIT ?\n"
+                    "                )"
+                ).format(
+                    where_clause=" AND ".join(where_clauses),
+                    candidate_order_clause=candidate_order_clause,
+                )
+
             cur.execute(
                 f"""
-                WITH candidates AS (
-                    SELECT s.id
-                    FROM sessions s
-                    WHERE {' AND '.join(where_clauses)}
-                    {candidate_order_clause}
-                    LIMIT ?
-                )
+                {candidate_cte}
                 {select_sql}
                 FROM sessions s
                 JOIN candidates c ON c.id = s.id

tests/test_issue2628_cli_sessions_perf.py

@@ -4,12 +4,15 @@ import pathlib
 import sqlite3
 import time
 
+import pytest
+
 import api.agent_sessions as agent_sessions
 
+_REAL_SQLITE_CONNECT = sqlite3.connect
 REPO_ROOT = pathlib.Path(__file__).resolve().parents[1]
 
 
-def _make_state_db(path, *, sessions=80, messages_per_session=3):
+def _make_state_db(path, *, sessions=80, messages_per_session=3, create_messages_index=True, source="cli", session_source="cli"):
     conn = sqlite3.connect(str(path))
     conn.executescript(
         """
@@ -33,7 +36,6 @@ def _make_state_db(path, *, sessions=80, messages_per_session=3):
             content TEXT,
             timestamp REAL
         );
-        CREATE INDEX idx_messages_session ON messages(session_id, timestamp);
         """
     )
     base = time.time() - sessions
@@ -44,19 +46,153 @@ def _make_state_db(path, *, sessions=80, messages_per_session=3):
             """
             INSERT INTO sessions
             (id, source, session_source, title, model, started_at, message_count, parent_session_id, ended_at, end_reason)
-            VALUES (?, 'cli', 'cli', ?, 'openai/gpt-5', ?, ?, NULL, NULL, NULL)
+            VALUES (?, ?, ?, ?, 'openai/gpt-5', ?, ?, NULL, NULL, NULL)
             """,
-            (sid, sid, started, messages_per_session),
+            (sid, source, session_source, sid, started, messages_per_session),
         )
         for j in range(messages_per_session):
             conn.execute(
                 "INSERT INTO messages (id, session_id, role, content, timestamp) VALUES (?, ?, ?, 'hello', ?)",
                 (f"msg_{i:04d}_{j:02d}", sid, "user" if j == 0 else "assistant", started + j / 10),
             )
+    if create_messages_index:
+        conn.execute("CREATE INDEX idx_messages_session ON messages(session_id, timestamp)")
     conn.commit()
     conn.close()
 
 
+def _newest_first_reference_ids(db_path, *, include_sources=None, exclude_sources=("webui",)):
+    where_clauses = ["s.source IS NOT NULL"]
+    params = []
+    if include_sources:
+        placeholders = ", ".join("?" for _ in include_sources)
+        where_clauses.append(f"s.source IN ({placeholders})")
+        params.extend(include_sources)
+    if exclude_sources:
+        placeholders = ", ".join("?" for _ in exclude_sources)
+        where_clauses.append(f"s.source NOT IN ({placeholders})")
+        params.extend(exclude_sources)
+    conn = sqlite3.connect(f"file:{db_path}?mode=ro&immutable=1", uri=True)
+    conn.row_factory = sqlite3.Row
+    try:
+        rows = conn.execute(
+            f"""
+            SELECT s.id
+            FROM sessions s
+            LEFT JOIN messages m ON m.session_id = s.id
+            WHERE {' AND '.join(where_clauses)}
+            GROUP BY s.id
+            ORDER BY COALESCE(MAX(m.timestamp), s.started_at) DESC
+            """,
+            params,
+        ).fetchall()
+        return [row["id"] for row in rows]
+    finally:
+        conn.close()
+
+
+def _execute_candidate_ordering_baseline_sql(
+    db_path,
+    candidate_limit,
+    *,
+    budget_ops,
+    interval=1,
+    include_sources=None,
+    exclude_sources=("webui",),
+):
+    where_clauses = ["s.source IS NOT NULL"]
+    params = []
+    if include_sources:
+        placeholders = ", ".join("?" for _ in include_sources)
+        where_clauses.append(f"s.source IN ({placeholders})")
+        params.extend(include_sources)
+    if exclude_sources:
+        placeholders = ", ".join("?" for _ in exclude_sources)
+        where_clauses.append(f"s.source NOT IN ({placeholders})")
+        params.extend(exclude_sources)
+
+    def _on_progress():
+        nonlocal steps
+        steps += 1
+        return 1 if steps > budget_ops else 0
+
+    steps = 0
+    conn = sqlite3.connect(f"file:{db_path}?mode=ro&immutable=1", uri=True)
+    conn.set_progress_handler(_on_progress, interval)
+    try:
+        return conn.execute(
+            f"""
+            WITH candidates AS (
+                SELECT s.id
+                FROM sessions s
+                WHERE {' AND '.join(where_clauses)}
+                ORDER BY COALESCE(
+                    (SELECT MAX(mx.timestamp) FROM messages mx WHERE mx.session_id = s.id),
+                    s.started_at
+                ) DESC,
+                s.started_at DESC
+                LIMIT ?
+            )
+            SELECT s.id
+            FROM sessions s
+            JOIN candidates c ON c.id = s.id
+            LEFT JOIN messages m ON m.session_id = s.id
+            GROUP BY s.id
+            ORDER BY COALESCE(MAX(m.timestamp), s.started_at) DESC
+            """,
+            [*params, candidate_limit],
+        ).fetchall()
+    finally:
+        conn.close()
+
+
+def _make_connect_with_progress_budget(*, budget_ops, interval=1):
+    def _connect(database, *_, **__):
+        steps = {"count": 0}
+        database_uri = str(database)
+        if database_uri.startswith("file:"):
+            target_uri = database_uri
+        else:
+            target_uri = f"file:{database_uri}?mode=ro&immutable=1"
+
+        def _on_progress():
+            steps["count"] += 1
+            return 1 if steps["count"] > budget_ops else 0
+
+        conn = _REAL_SQLITE_CONNECT(
+            target_uri,
+            uri=True,
+        )
+        conn.set_progress_handler(_on_progress, interval)
+        return conn
+
+    return _connect
+
+
+def _make_connect_with_progress_counter(*, interval=1):
+    steps = {"count": 0}
+
+    def _connect(database, *_, **__):
+        database_uri = str(database)
+        if database_uri.startswith("file:"):
+            target_uri = database_uri
+        else:
+            target_uri = f"file:{database_uri}?mode=ro&immutable=1"
+
+        def _on_progress():
+            steps["count"] += 1
+            return 0
+
+        conn = _REAL_SQLITE_CONNECT(
+            target_uri,
+            uri=True,
+        )
+        conn.set_progress_handler(_on_progress, interval)
+        return conn
+
+    return _connect, steps
+
+
 def test_importable_agent_rows_push_sidebar_limit_into_sql(tmp_path):
     """A capped sidebar scan should not aggregate the entire state.db first."""
     db = tmp_path / "state.db"
@@ -71,10 +207,82 @@ def test_importable_agent_rows_push_sidebar_limit_into_sql(tmp_path):
     src = (REPO_ROOT / "api" / "agent_sessions.py").read_text()
     assert "WITH candidates AS" in src
     assert "JOIN candidates c ON c.id = s.id" in src
-    assert "SELECT MAX(mx.timestamp) FROM messages mx WHERE mx.session_id = s.id" in src
+    assert "latest_messages AS" in src
+    assert "LEFT JOIN latest_messages lm ON lm.session_id = s.id" in src
+    assert 'included == ("cron",)' in src
+    assert "not messages_index_present" in src
+    assert "PRAGMA index_list(messages)" in src
+    assert "CREATE INDEX IF NOT EXISTS idx_messages_session" in src
+    assert "_CRON_PREAGGREGATE_CANDIDATE_ORDER_MIN_MESSAGES" not in src
+    assert "MAX(mx.timestamp) FROM messages mx WHERE mx.session_id = s.id" in src
     assert "candidate_limit = max(result_limit * 8, result_limit)" in src
 
 
+def test_importable_agent_rows_candidate_ordering_stays_under_progress_budget(tmp_path, monkeypatch):
+    """Cron-only missing-index scans should fail under the old shape budget, then pass after pre-aggregation."""
+    db = tmp_path / "state.db"
+    _make_state_db(
+        db,
+        sessions=120,
+        messages_per_session=900,
+        create_messages_index=False,
+        source="cron",
+        session_source="cron",
+    )
+    reference_ids = _newest_first_reference_ids(db, include_sources=("cron",), exclude_sources=None)
+    candidate_limit = max(20 * 8, 20)
+    progress_interval = 100
+
+    original_connect = agent_sessions.sqlite3.connect
+    connect_with_progress_counter, progress_counter = _make_connect_with_progress_counter(
+        interval=progress_interval
+    )
+    monkeypatch.setattr(agent_sessions.sqlite3, "connect", connect_with_progress_counter)
+    try:
+        measured_rows = agent_sessions.read_importable_agent_session_rows(
+            db,
+            limit=20,
+            exclude_sources=None,
+            include_sources=("cron",),
+        )
+        assert [row["id"] for row in measured_rows] == reference_ids[:20]
+    finally:
+        # Keep this helper isolated; the baseline must still run without the
+        # counting handler to validate raw cost differences.
+        monkeypatch.setattr(agent_sessions.sqlite3, "connect", original_connect)
+
+    # Give the head path a small deterministic margin, then require the old
+    # correlated query to exceed the same budget on the missing-index branch.
+    progress_budget_ops = max(progress_counter["count"] + 200, 1)
+
+    with pytest.raises(sqlite3.OperationalError, match="interrupted"):
+        _execute_candidate_ordering_baseline_sql(
+            db,
+            candidate_limit,
+            budget_ops=progress_budget_ops,
+            interval=progress_interval,
+            include_sources=("cron",),
+            exclude_sources=None,
+        )
+
+    monkeypatch.setattr(
+        agent_sessions.sqlite3,
+        "connect",
+        _make_connect_with_progress_budget(
+            budget_ops=progress_budget_ops,
+            interval=progress_interval,
+        ),
+    )
+
+    rows = agent_sessions.read_importable_agent_session_rows(
+        db,
+        limit=20,
+        exclude_sources=None,
+        include_sources=("cron",),
+    )
+    assert [row["id"] for row in rows] == reference_ids[:20]
+
+
 def test_importable_agent_rows_limit_includes_resumed_old_session(tmp_path):
     """The capped candidate window must not hide old sessions resumed recently."""
     db = tmp_path / "state.db"