perf(code): make threads list fast on large session databases (#4005)

Makes `dcode threads list` fast on large session databases.

On a real profile (~12 GB `sessions.db`, 196k checkpoints, 1,505
threads) `threads list` took **~60 s**. Profiling showed the cost was
**not** where it looked — it was the top-level listing query, not the
per-thread message counts.

## Root cause

LangGraph's `SqliteSaver` stores each checkpoint's full state blob
**inline in the same `checkpoints` row** as the small `metadata` field.
The listing query only reads `metadata` (latest `updated_at`,
`agent_name`, `git_branch`, `cwd` per thread), but with no covering
index SQLite full-scans the table — dragging **11.8 GB of state blobs**
through I/O just to read **234 MB of metadata**. That scan alone
measured **59.6 s**.

## Fix 1 — covering index (primary win)

Add an idempotent `idx_dcode_threads_list` covering index carrying
exactly the expressions the `list_threads` `GROUP BY` reads. The planner
now answers from the index alone (`SCAN … USING INDEX
idx_dcode_threads_list`) and never touches the blob-bearing rows.

Measured on the real DB above:

| | Query time | Plan |
|---|---|---|
| Before | **59.6 s** | full scan via PK autoindex (drags 12 GB) |
| After | **~0.04 s** | index-only scan (~50 MB index) |

- One-time build (~40–60 s, a single table scan) on the first `threads
list` for an existing DB; near-zero for new/small DBs.
- Adds ~50 MB. Index creation is non-fatal on failure (read-only DB /
lock contention) — the query falls back to the previous scan and logs.

## Fix 2 — linear message-count reconstruction (secondary)

For delta-channel threads whose latest checkpoint doesn't inline
`messages`, the count is rebuilt from the `writes` table. The old path
deserialized each write in its own `run_in_executor` hop and folded
deltas one at a time through `add_messages` — O(n²) per thread. Now each
chunk is decoded and reduced in a single worker-thread hop, and the
common append/clear history folds in one linear `add_messages` pass.
Specific delete-by-ID deltas and any reducer error fall back to the
exact sequential fold, so displayed counts are unchanged. A 4,000-write
thread now counts in ~30 ms instead of seconds.

## Robustness & docs

- Per-thread fault isolation: a single malformed thread yields a missing
count instead of failing the whole list load.
- Tightened docstrings: the delete-by-ID fallback is documented as
conservative (batch and sequential folds agree on the *count* in
practice; the sequential fold just guarantees it), not a demonstrated
divergence.

---------

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>

libs/code/deepagents_code/sessions.py

@@ -311,6 +311,56 @@ async def _table_exists(conn: aiosqlite.Connection, table: str) -> bool:
         return await cursor.fetchone() is not None
 
 
+_THREADS_LIST_INDEX = "idx_dcode_threads_list"
+"""Covering index that makes the `list_threads` GROUP BY an index-only scan.
+
+LangGraph's `SqliteSaver` stores each checkpoint's full state blob inline in the
+`checkpoints` row alongside the small `metadata` field. The thread-list query
+only needs `metadata` (per-thread latest `updated_at`, `agent_name`, etc.), but
+without a covering index SQLite scans the whole table — dragging every state
+blob through I/O. On a large profile (e.g. ~12 GB of blobs) that scan takes
+tens of seconds. This index carries exactly the expressions the query reads, so
+the planner satisfies the GROUP BY from the index alone and never touches the
+blob-bearing rows, turning a ~60 s scan into a sub-second lookup.
+
+The column order (leading `thread_id`) also lets the GROUP BY consume the index
+in order. Keep the indexed expressions in sync with the `list_threads` query.
+"""
+
+
+async def _ensure_threads_list_index(conn: aiosqlite.Connection) -> None:
+    """Create the `list_threads` covering index if it does not already exist.
+
+    Idempotent: `CREATE INDEX IF NOT EXISTS` is a near-instant catalog check once
+    the index exists. The one-time build on a pre-existing large database costs a
+    single full table scan (seconds to tens of seconds), after which every
+    `list_threads` call is a sub-second index-only scan. Runs in the aiosqlite
+    worker thread, so it does not block the event loop.
+
+    A failure here is non-fatal: the list query still returns correct results via
+    the slower table scan, so we log and continue rather than break `threads
+    list` (e.g. on a read-only database or under write-lock contention).
+    """
+    try:
+        await conn.execute(
+            f"CREATE INDEX IF NOT EXISTS {_THREADS_LIST_INDEX} ON checkpoints("
+            "thread_id, "
+            "json_extract(metadata, '$.updated_at'), "
+            "checkpoint_id, "
+            "json_extract(metadata, '$.agent_name'), "
+            "json_extract(metadata, '$.git_branch'), "
+            "json_extract(metadata, '$.cwd'))"
+        )
+        await conn.commit()
+    except Exception:
+        logger.warning(
+            "Failed to create the %s index; `threads list` will fall back to a "
+            "full table scan and may be slow on large databases",
+            _THREADS_LIST_INDEX,
+            exc_info=True,
+        )
+
+
 async def list_threads(
     agent_name: str | None = None,
     limit: int = 20,
@@ -345,6 +395,11 @@ async def list_threads(
         if not await _table_exists(conn, "checkpoints"):
             return []
 
+        # Ensure the covering index exists before the GROUP BY below, so the
+        # query is an index-only scan instead of a full scan over the (large,
+        # blob-bearing) checkpoints table.
+        await _ensure_threads_list_index(conn)
+
         if sort_by not in {"updated", "created"}:
             msg = f"Invalid sort_by {sort_by!r}; expected 'updated' or 'created'"
             raise ValueError(msg)
@@ -915,11 +970,16 @@ async def _load_message_counts_from_writes_batch(
     `add_messages` as a count-equivalent stand-in for the channel's actual
     reducer (`_messages_delta_reducer`): both dedup by ID and honor
     `RemoveMessage` / `REMOVE_ALL_MESSAGES`, so they produce the same final
-    message set. The reducer is batching-invariant, so folding the write history
-    from empty yields the same value LangGraph reconstructs from the latest
-    snapshot plus subsequent deltas. An `Overwrite` write resets the accumulator
-    to its value, matching the net effect of `DeltaChannel.replay_writes` (where
-    the last `Overwrite` is the reset point).
+    message set. An `Overwrite` write resets the accumulator to its value,
+    matching the net effect of `DeltaChannel.replay_writes` (where the last
+    `Overwrite` is the reset point).
+
+    Reduction runs in a single worker-thread hop per chunk (decode is CPU-bound
+    and a long thread can have thousands of writes; dispatching per row both
+    serialized the work and added an executor round-trip each time). The common
+    append-and-clear history folds in one `add_messages` pass (linear), which is
+    why a busy thread no longer takes seconds to count. See
+    `_count_messages_from_deltas` for the fold and its exact-fold fallback.
 
     Only the root namespace (`checkpoint_ns = ''`) is counted, matching both the
     inline path and the conversation the `/threads` selector cares about;
@@ -949,14 +1009,12 @@ async def _load_message_counts_from_writes_batch(
     if not thread_ids:
         return {}
 
-    from langgraph.graph.message import add_messages
-    from langgraph.types import Overwrite
-
     loop = asyncio.get_running_loop()
-    # Accumulate one reduced message list per thread across chunks. Ordering by
-    # (checkpoint_id, task_id, idx) replays deltas oldest-to-newest, matching
-    # how LangGraph applies them on load.
-    reduced: dict[str, list[Any]] = {}
+    results: dict[str, int] = {}
+    # Chunks partition by thread, so every write for a given thread lands in the
+    # same query; each thread is counted exactly once. Ordering by
+    # (checkpoint_id, task_id, idx) replays deltas oldest-to-newest, matching how
+    # LangGraph applies them on load.
     for start in range(0, len(thread_ids), _SQLITE_MAX_VARIABLE_NUMBER):
         chunk = thread_ids[start : start + _SQLITE_MAX_VARIABLE_NUMBER]
         placeholders = ",".join("?" * len(chunk))
@@ -971,32 +1029,144 @@ async def _load_message_counts_from_writes_batch(
         async with conn.execute(query, chunk) as cursor:
             rows = await cursor.fetchall()
 
-        for tid, type_str, value_blob in rows:
-            if not type_str or not value_blob:
-                continue
-            try:
-                delta = await loop.run_in_executor(
-                    None, serde.loads_typed, (type_str, value_blob)
-                )
-                if isinstance(delta, Overwrite):
-                    # Overwrite resets the channel; its value becomes the base.
-                    value = delta.value
-                    reduced[tid] = list(value) if isinstance(value, list) else []
-                else:
-                    # `add_messages` with a list left arg returns a list.
-                    reduced[tid] = cast(
-                        "list[Any]", add_messages(reduced.get(tid, []), delta)
-                    )
-            except Exception:
-                logger.warning(
-                    "Failed to replay messages write for thread %s; "
-                    "message count may be inaccurate",
-                    tid,
-                    exc_info=True,
-                )
-                continue
+        chunk_counts = await loop.run_in_executor(
+            None, _reduce_message_write_rows, list(rows), serde
+        )
+        results.update(chunk_counts)
 
-    return {tid: len(messages) for tid, messages in reduced.items()}
+    return results
+
+
+def _reduce_message_write_rows(
+    rows: list[tuple[str, str | None, bytes | None]],
+    serde: JsonPlusSerializer,
+) -> dict[str, int]:
+    """Decode `messages`-channel write rows and count messages per thread.
+
+    Runs synchronously in a worker thread. Rows must be ordered so each thread's
+    deltas are oldest-to-newest. Undecodable rows are skipped (logged), matching
+    the per-row error handling of the previous implementation.
+
+    Returns:
+        Mapping of thread ID to reconstructed message count.
+    """
+    deltas_by_thread: dict[str, list[Any]] = {}
+    for tid, type_str, value_blob in rows:
+        if not type_str or not value_blob:
+            continue
+        try:
+            delta = serde.loads_typed((type_str, value_blob))
+        except Exception:
+            logger.warning(
+                "Failed to replay messages write for thread %s; "
+                "message count may be inaccurate",
+                tid,
+                exc_info=True,
+            )
+            continue
+        deltas_by_thread.setdefault(tid, []).append(delta)
+
+    counts: dict[str, int] = {}
+    for tid, deltas in deltas_by_thread.items():
+        try:
+            counts[tid] = _count_messages_from_deltas(deltas)
+        except Exception:
+            # Keep one malformed thread from failing the whole `threads list`
+            # load: skip it (its count is simply absent) rather than propagating.
+            logger.warning(
+                "Failed to count messages for thread %s; omitting its count",
+                tid,
+                exc_info=True,
+            )
+    return counts
+
+
+def _count_messages_from_deltas(deltas: list[Any]) -> int:
+    """Count messages from an ordered list of `messages`-channel write deltas.
+
+    Fast path: appends and full-clears (`REMOVE_ALL_MESSAGES`, `Overwrite`) fold
+    into one `add_messages` pass — O(n) instead of the O(n^2) incremental fold,
+    so threads with thousands of writes count in milliseconds. For these ops the
+    single-pass result is count-equivalent to the sequential fold (both dedup by
+    ID, and clears collapse to the post-clear tail).
+
+    Slow path: a specific `RemoveMessage` (delete-by-ID) or any reducer error
+    falls back to the exact sequential fold as a conservative measure. A
+    delete-by-ID concatenated into the single `buffer` can make batch
+    `add_messages` raise (the target ID may be absent at that buffer position),
+    and we do not rely on unproven count-equivalence of batched removal. In
+    practice the two folds still agree on the count for these histories; the
+    sequential fold simply guarantees it. Such deletes are rare in linear dcode
+    histories, so the common case stays on the fast path.
+
+    Returns:
+        Number of messages after reducing the deltas.
+    """
+    from langchain_core.messages import RemoveMessage
+    from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages
+    from langgraph.types import Overwrite
+
+    buffer: list[Any] = []
+    needs_exact_fold = False
+    for delta in deltas:
+        if isinstance(delta, Overwrite):
+            value = delta.value
+            buffer = list(value) if isinstance(value, list) else []
+            continue
+        items = delta if isinstance(delta, list) else [delta]
+        for item in items:
+            if isinstance(item, RemoveMessage):
+                if item.id == REMOVE_ALL_MESSAGES:
+                    buffer = []
+                else:
+                    needs_exact_fold = True
+                    break
+            else:
+                buffer.append(item)
+        if needs_exact_fold:
+            break
+
+    if not needs_exact_fold:
+        try:
+            return len(cast("list[Any]", add_messages([], buffer)))
+        except Exception:
+            logger.debug(
+                "Batched message-count fold failed; using sequential fold",
+                exc_info=True,
+            )
+
+    return _incremental_message_count(deltas)
+
+
+def _incremental_message_count(deltas: list[Any]) -> int:
+    """Count messages by folding deltas sequentially through `add_messages`.
+
+    Exact reference reduction: applies one delta at a time, resetting on
+    `Overwrite` and skipping any delta the reducer rejects (e.g. a delete for an
+    absent ID). Used as the fallback when the batched fast path cannot guarantee
+    a matching count.
+
+    Returns:
+        Number of messages after the sequential fold.
+    """
+    from langgraph.graph.message import add_messages
+    from langgraph.types import Overwrite
+
+    reduced: list[Any] = []
+    for delta in deltas:
+        if isinstance(delta, Overwrite):
+            value = delta.value
+            reduced = list(value) if isinstance(value, list) else []
+            continue
+        try:
+            reduced = cast("list[Any]", add_messages(reduced, delta))
+        except Exception:
+            logger.warning(
+                "Failed to replay messages write; message count may be inaccurate",
+                exc_info=True,
+            )
+            continue
+    return len(reduced)
 
 
 def _summarize_checkpoint(data: object) -> _CheckpointSummary:

libs/code/tests/unit_tests/test_sessions.py

@@ -174,6 +174,53 @@ class TestThreadFunctions:
             assert by_id["thread2"]["cwd"] == "/tmp/workspace"
             assert by_id["thread3"]["cwd"] is None
 
+    def test_list_threads_creates_covering_index(self, temp_db):
+        """`list_threads` creates the covering index and the plan uses it.
+
+        Regression guard for the `threads list` slowdown on large profiles: the
+        GROUP BY must be an index-only scan of `idx_dcode_threads_list`, not a
+        full scan of the blob-bearing checkpoints table.
+        """
+        with patch.object(sessions, "get_db_path", return_value=temp_db):
+            asyncio.run(sessions.list_threads())
+
+        conn = sqlite3.connect(str(temp_db))
+        try:
+            index_names = {
+                r[0]
+                for r in conn.execute(
+                    "SELECT name FROM sqlite_master WHERE type='index'"
+                )
+            }
+            assert "idx_dcode_threads_list" in index_names
+
+            plan = " ".join(
+                str(row[3])
+                for row in conn.execute(
+                    "EXPLAIN QUERY PLAN "
+                    "SELECT thread_id, "
+                    "MAX(json_extract(metadata, '$.updated_at')) u, "
+                    "MAX(checkpoint_id), "
+                    "MAX(json_extract(metadata, '$.agent_name')), "
+                    "MAX(json_extract(metadata, '$.git_branch')), "
+                    "MAX(json_extract(metadata, '$.cwd')) "
+                    "FROM checkpoints GROUP BY thread_id ORDER BY u DESC LIMIT 20"
+                )
+            )
+            # Index-only scan of the covering index, not the PK autoindex (which
+            # would drag the checkpoint state blobs through I/O).
+            assert "idx_dcode_threads_list" in plan
+            assert "sqlite_autoindex_checkpoints_1" not in plan
+        finally:
+            conn.close()
+
+    def test_list_threads_index_creation_is_idempotent(self, temp_db):
+        """Repeated `list_threads` calls succeed once the index already exists."""
+        with patch.object(sessions, "get_db_path", return_value=temp_db):
+            first = asyncio.run(sessions.list_threads())
+            second = asyncio.run(sessions.list_threads())
+        assert len(first) == len(second) == 3
+
     def test_list_threads_filter_by_agent(self, temp_db):
         """List filters by agent name."""
         with patch.object(sessions, "get_db_path", return_value=temp_db):
@@ -2711,6 +2758,163 @@ class TestLoadMessageCountsFromWritesBatch:
         # The good write still counts; the corrupt one is skipped.
         assert results == {"t1": 1}
 
+    async def test_large_append_history_counts_correctly(self) -> None:
+        """A long append-only history counts correctly via the one-pass fold.
+
+        Correctness-at-scale check for the `threads list` speedup. Note this
+        does not guard against a perf regression on its own: the old O(n^2)
+        fold returns the same count (just slowly), so a revert to the quadratic
+        path would still pass. Wall-clock assertions are too flaky for CI; a
+        codspeed benchmark would be the real regression guard.
+        """
+        serde = JsonPlusSerializer()
+
+        import aiosqlite
+
+        n = 4000
+        async with aiosqlite.connect(":memory:") as conn:
+            await conn.execute(
+                "CREATE TABLE writes "
+                "(thread_id TEXT, checkpoint_ns TEXT, checkpoint_id TEXT, "
+                "task_id TEXT, idx INTEGER, channel TEXT, type TEXT, value BLOB)"
+            )
+            rows = []
+            for i in range(n):
+                type_str, blob = serde.dumps_typed(
+                    [{"type": "human", "content": "x", "id": f"m{i}"}]
+                )
+                rows.append(("t1", f"cp_{i:06d}", "task1", type_str, blob))
+            await conn.executemany(
+                "INSERT INTO writes VALUES (?, '', ?, ?, 0, 'messages', ?, ?)",
+                rows,
+            )
+            await conn.commit()
+
+            results = await sessions._load_message_counts_from_writes_batch(  # pyright: ignore[reportPrivateUsage]
+                conn, ["t1"], serde
+            )
+
+        assert results == {"t1": n}
+
+
+class TestCountMessagesFromDeltas:
+    """Tests for the delta-folding message counter and its exact fallback."""
+
+    def test_fast_path_dedups_repeated_ids(self) -> None:
+        """Repeated IDs across deltas collapse to one message (fast path)."""
+        from langchain_core.messages import AIMessage
+
+        deltas = [
+            [AIMessage(content="draft", id="a1")],
+            [AIMessage(content="final", id="a1")],
+        ]
+        assert sessions._count_messages_from_deltas(deltas) == 1  # pyright: ignore[reportPrivateUsage]
+
+    def test_remove_all_then_append_resets(self) -> None:
+        """`REMOVE_ALL_MESSAGES` clears the buffer before later appends."""
+        from langchain_core.messages import HumanMessage, RemoveMessage
+        from langgraph.graph.message import REMOVE_ALL_MESSAGES
+
+        deltas = [
+            [HumanMessage(content="a", id="h1")],
+            [HumanMessage(content="b", id="h2")],
+            [
+                RemoveMessage(id=REMOVE_ALL_MESSAGES),
+                HumanMessage(content="fresh", id="h9"),
+            ],
+        ]
+        assert sessions._count_messages_from_deltas(deltas) == 1  # pyright: ignore[reportPrivateUsage]
+
+    def test_overwrite_resets_buffer(self) -> None:
+        """An `Overwrite` delta replaces the accumulated buffer."""
+        from langchain_core.messages import HumanMessage
+        from langgraph.types import Overwrite
+
+        deltas = [
+            [HumanMessage(content="a", id="h1")],
+            [HumanMessage(content="b", id="h2")],
+            Overwrite(value=[HumanMessage(content="fresh", id="h9")]),
+        ]
+        assert sessions._count_messages_from_deltas(deltas) == 1  # pyright: ignore[reportPrivateUsage]
+
+    def test_duplicate_id_within_one_delta_counts_once(self) -> None:
+        """Two messages sharing an ID in one delta collapse to a single count.
+
+        Pins the case excluded from the property test: batch and sequential
+        `add_messages` may order/identify differently, but the *count* the
+        counter reports is the same (both dedup by ID), so the fast path is
+        safe here and needs no exact-fold routing.
+        """
+        from langchain_core.messages import AIMessage
+
+        deltas = [
+            [AIMessage(content="draft", id="d1"), AIMessage(content="final", id="d1")]
+        ]
+        assert sessions._count_messages_from_deltas(deltas) == 1  # pyright: ignore[reportPrivateUsage]
+        assert sessions._count_messages_from_deltas(  # pyright: ignore[reportPrivateUsage]
+            deltas
+        ) == sessions._incremental_message_count(deltas)  # pyright: ignore[reportPrivateUsage]
+
+    def test_specific_remove_matches_incremental_fold(self) -> None:
+        """A delete-by-ID routes through the exact fold and matches it."""
+        from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
+
+        deltas = [
+            [HumanMessage(content="a", id="h1"), AIMessage(content="b", id="a1")],
+            [RemoveMessage(id="a1")],
+        ]
+        assert sessions._count_messages_from_deltas(  # pyright: ignore[reportPrivateUsage]
+            deltas
+        ) == sessions._incremental_message_count(deltas)  # pyright: ignore[reportPrivateUsage]
+        assert sessions._count_messages_from_deltas(deltas) == 1  # pyright: ignore[reportPrivateUsage]
+
+    def test_fast_and_exact_agree_on_realistic_histories(self) -> None:
+        """Fast path and exact fold agree on append/clear/overwrite histories.
+
+        Covers the realistic shapes the counter sees (unique-ID appends,
+        streaming updates that reuse an ID, compaction clears, and snapshot
+        overwrites). Excludes duplicate-ID-within-one-delta, whose batch-vs-
+        sequential *ordering/identity* equivalence we haven't characterized;
+        the counts still match there (see
+        `test_duplicate_id_within_one_delta_counts_once`).
+        """
+        import random
+
+        from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
+        from langgraph.graph.message import REMOVE_ALL_MESSAGES
+        from langgraph.types import Overwrite
+
+        rng = random.Random(20240117)
+        for _ in range(500):
+            deltas: list[object] = []
+            next_id = 0
+            for _ in range(rng.randint(0, 25)):
+                roll = rng.random()
+                if roll < 0.6:
+                    next_id += 1
+                    cls = rng.choice([HumanMessage, AIMessage])
+                    deltas.append([cls(content="x", id=f"m{next_id}")])
+                elif roll < 0.75 and next_id:
+                    # Streaming update: reuse a recent ID.
+                    deltas.append(
+                        [AIMessage(content="y", id=f"m{rng.randint(1, next_id)}")]
+                    )
+                elif roll < 0.85:
+                    deltas.append([RemoveMessage(id=REMOVE_ALL_MESSAGES)])
+                else:
+                    count = rng.randint(0, 3)
+                    deltas.append(
+                        Overwrite(
+                            value=[
+                                HumanMessage(content="o", id=f"ov{j}")
+                                for j in range(count)
+                            ]
+                        )
+                    )
+            fast = sessions._count_messages_from_deltas(deltas)  # pyright: ignore[reportPrivateUsage]
+            exact = sessions._incremental_message_count(deltas)  # pyright: ignore[reportPrivateUsage]
+            assert fast == exact, deltas
+
 
 class TestInitialPromptFromMessages:
     """Tests for the message-list parser used by the writes-table reader."""