core: improve conversation loading performance by batching database queries

Reduces memory usage and speeds up conversation loading by using pagination
and inArray queries instead of loading all messages at once
This commit is contained in:
Dax Raad
2026-01-26 12:33:18 -05:00
parent 2b05833c32
commit 5e1639de2b

View File

@@ -6,7 +6,7 @@ import { Identifier } from "../id/id"
import { LSP } from "../lsp" import { LSP } from "../lsp"
import { Snapshot } from "@/snapshot" import { Snapshot } from "@/snapshot"
import { fn } from "@/util/fn" import { fn } from "@/util/fn"
import { Database, eq, desc } from "@/storage/db" import { Database, eq, desc, inArray } from "@/storage/db"
import { MessageTable, PartTable } from "./session.sql" import { MessageTable, PartTable } from "./session.sql"
import { ProviderTransform } from "@/provider/transform" import { ProviderTransform } from "@/provider/transform"
import { STATUS_CODES } from "http" import { STATUS_CODES } from "http"
@@ -608,27 +608,56 @@ export namespace MessageV2 {
} }
export const stream = fn(Identifier.schema("session"), async function* (sessionID) { export const stream = fn(Identifier.schema("session"), async function* (sessionID) {
const size = 50
let offset = 0
while (true) {
const rows = Database.use((db) => const rows = Database.use((db) =>
db db
.select() .select()
.from(MessageTable) .from(MessageTable)
.where(eq(MessageTable.session_id, sessionID)) .where(eq(MessageTable.session_id, sessionID))
.orderBy(desc(MessageTable.created_at)) .orderBy(desc(MessageTable.created_at))
.limit(size)
.offset(offset)
.all(), .all(),
) )
if (rows.length === 0) break
const ids = rows.map((row) => row.id)
const partsByMessage = new Map<string, MessageV2.Part[]>()
if (ids.length > 0) {
const partRows = Database.use((db) =>
db
.select()
.from(PartTable)
.where(inArray(PartTable.message_id, ids))
.orderBy(PartTable.message_id, PartTable.id)
.all(),
)
for (const row of partRows) {
const list = partsByMessage.get(row.message_id)
if (list) list.push(row.data)
else partsByMessage.set(row.message_id, [row.data])
}
}
for (const row of rows) { for (const row of rows) {
yield { yield {
info: row.data, info: row.data,
parts: await parts(row.id), parts: partsByMessage.get(row.id) ?? [],
} }
} }
offset += rows.length
if (rows.length < size) break
}
}) })
export const parts = fn(Identifier.schema("message"), async (message_id) => { export const parts = fn(Identifier.schema("message"), async (message_id) => {
const rows = Database.use((db) => db.select().from(PartTable).where(eq(PartTable.message_id, message_id)).all()) const rows = Database.use((db) =>
const result = rows.map((row) => row.data) db.select().from(PartTable).where(eq(PartTable.message_id, message_id)).orderBy(PartTable.id).all(),
result.sort((a, b) => (a.id > b.id ? 1 : -1)) )
return result return rows.map((row) => row.data)
}) })
export const get = fn( export const get = fn(