fix(session): restore message pagination hydration

Author: vpetrigo

packages/opencode/src/server/routes/instance/httpapi/handlers/session.ts

@@ -1,4 +1,5 @@
 import { PermissionV1 } from "@opencode-ai/core/v1/permission"
+import { Database } from "@opencode-ai/core/database/database"
 import { Agent } from "@/agent/agent"
 import { SessionV1 } from "@opencode-ai/core/v1/session"
 import { EventV2Bridge } from "@/event-v2-bridge"
@@ -48,6 +49,7 @@ const tryParseJson = (text: string) =>
 export const sessionHandlers = HttpApiBuilder.group(InstanceHttpApi, "session", (handlers) =>
   Effect.gen(function* () {
     const session = yield* Session.Service
+    const database = yield* Database.Service
     const shareSvc = yield* SessionShare.Service
     const promptSvc = yield* SessionPrompt.Service
     const revertSvc = yield* SessionRevert.Service
@@ -128,7 +130,7 @@ export const sessionHandlers = HttpApiBuilder.group(InstanceHttpApi, "session",
           limit: ctx.query.limit,
           before: ctx.query.before,
           after: ctx.query.after,
-        }),
+        }).pipe(Effect.provideService(Database.Service, database)),
       )
       if (!page.cursor) return page.items
 

packages/opencode/src/session/message-v2.ts

@@ -105,6 +105,7 @@ function hydrate(rows: (typeof MessageTable.$inferSelect)[]) {
   const partByMessage = new Map<string, Part[]>()
   return Effect.gen(function* () {
     if (ids.length > 0) {
+      const { db } = yield* Database.Service
       const partRows = yield* db
         .select()
         .from(PartTable)
@@ -442,22 +443,25 @@ export const page = Effect.fn("MessageV2.page")(function* (input: {
     : after
       ? and(eq(MessageTable.session_id, input.sessionID), newer(after))
       : eq(MessageTable.session_id, input.sessionID)
-  const rows = yield* Database.use((db) =>
-    db
-      .select()
-      .from(MessageTable)
-      .where(where)
-      .orderBy(
-        after ? asc(MessageTable.time_created) : desc(MessageTable.time_created),
-        after ? asc(MessageTable.id) : desc(MessageTable.id),
-      )
-      .limit(input.limit + 1)
-      .all(),
-  )
-  if (rows.length === 0) {
-    const row = yield* Database.use((db) =>
-      db.select({ id: SessionTable.id }).from(SessionTable).where(eq(SessionTable.id, input.sessionID)).get(),
+  const { db } = yield* Database.Service
+  const rows = yield* db
+    .select()
+    .from(MessageTable)
+    .where(where)
+    .orderBy(
+      after ? asc(MessageTable.time_created) : desc(MessageTable.time_created),
+      after ? asc(MessageTable.id) : desc(MessageTable.id),
     )
+    .limit(input.limit + 1)
+    .all()
+    .pipe(Effect.orDie)
+  if (rows.length === 0) {
+    const row = yield* db
+      .select({ id: SessionTable.id })
+      .from(SessionTable)
+      .where(eq(SessionTable.id, input.sessionID))
+      .get()
+      .pipe(Effect.orDie)
     if (!row) return yield* new NotFoundError({ message: `Session not found: ${input.sessionID}` })
     return {
       items: [] as WithParts[],

packages/tui/src/context/sync.tsx

@@ -604,31 +604,56 @@ export const {
         },
         async sync(sessionID: string) {
           if (fullSyncedSessions.has(sessionID)) return
-          const [session, messages, todo, diff] = await Promise.all([
-            sdk.client.session.get({ sessionID }, { throwOnError: true }),
-            sdk.client.session.messages({ sessionID, limit: INITIAL_PAGE_SIZE }),
-            sdk.client.session.todo({ sessionID }),
-            sdk.client.session.diff({ sessionID }),
-          ])
-          const olderCursor = (messages.response?.headers.get("x-next-cursor") as string | null | undefined) ?? null
-          setStore(
-            produce((draft) => {
-              const match = search(draft.session, sessionID, (s) => s.id)
-              if (match.found) draft.session[match.index] = session.data!
-              if (!match.found) draft.session.splice(match.index, 0, session.data!)
-              draft.todo[sessionID] = todo.data ?? []
-              const infos: (typeof draft.message)[string] = []
-              for (const message of messages.data ?? []) {
-                infos.push(message.info)
-                draft.part[message.info.id] = message.parts
-              }
-              draft.message[sessionID] = infos
-              draft.session_diff[sessionID] = diff.data ?? []
-              draft.messageOlderCursor[sessionID] = olderCursor
-              draft.messageNewerCursor[sessionID] = null
-            }),
-          )
-          if (!olderCursor) fullSyncedSessions.add(sessionID)
+          const hydration = { messages: new Set<string>(), parts: new Set<string>() }
+          hydratingSessions.set(sessionID, hydration)
+          try {
+            const [session, messages, todo, diff] = await Promise.all([
+              sdk.client.session.get({ sessionID }, { throwOnError: true }),
+              sdk.client.session.messages({ sessionID, limit: INITIAL_PAGE_SIZE }),
+              sdk.client.session.todo({ sessionID }),
+              sdk.client.session.diff({ sessionID }),
+            ])
+            const olderCursor = (messages.response?.headers.get("x-next-cursor") as string | null | undefined) ?? null
+            setStore(
+              produce((draft) => {
+                const match = search(draft.session, sessionID, (s) => s.id)
+                if (match.found) draft.session[match.index] = session.data!
+                if (!match.found) draft.session.splice(match.index, 0, session.data!)
+                draft.todo[sessionID] = todo.data ?? []
+                if (messages.data) {
+                  const existing = draft.message[sessionID] ?? []
+                  const infos: (typeof draft.message)[string] = []
+                  const hydratedIDs = new Set<string>()
+                  for (const message of messages.data) {
+                    hydratedIDs.add(message.info.id)
+                    const current = existing.find((item) => item.id === message.info.id)
+                    if (current) {
+                      infos.push(current)
+                      continue
+                    }
+                    if (hydration.messages.has(message.info.id)) continue
+                    infos.push(message.info)
+                    draft.part[message.info.id] = message.parts
+                  }
+                  for (const message of existing) {
+                    if (hydratedIDs.has(message.id)) continue
+                    infos.push(message)
+                  }
+                  while (infos.length > INITIAL_PAGE_SIZE) {
+                    const evicted = infos.shift()
+                    if (evicted) delete draft.part[evicted.id]
+                  }
+                  draft.message[sessionID] = infos
+                  draft.messageOlderCursor[sessionID] = olderCursor
+                  draft.messageNewerCursor[sessionID] = null
+                }
+                draft.session_diff[sessionID] = diff.data ?? []
+              }),
+            )
+            if (!olderCursor) fullSyncedSessions.add(sessionID)
+          } finally {
+            hydratingSessions.delete(sessionID)
+          }
         },
         async loadOlderMessages(sessionID: string) {
           const cursor = store.messageOlderCursor[sessionID]