fix(app): preload more timeline messages (#36172)

Author: Hona

packages/app/e2e/regression/session-timeline-history-root.spec.ts

@@ -17,12 +17,14 @@ import { mockOpenCodeServer } from "../utils/mock-server"
 import { installSseTransport } from "../utils/sse-transport"
 import { expectSessionTitle } from "../utils/waits"
 
-const assistants = Array.from({ length: 14 }, (_, index) =>
+const initialPageSize = 20
+const historyPageSize = 200
+const assistants = Array.from({ length: initialPageSize + 1 }, (_, index) =>
   assistantMessage([textPart(`prt_history_root_${index}`, `Assistant response ${index}`)], {
     id: `msg_${String(index + 1001).padStart(4, "0")}_history_root_assistant`,
     parentID: userID,
     created: 1700000001000 + index * 1_000,
-    completed: index < 13,
+    completed: index < initialPageSize,
   }),
 )
 const messages = [userMessage(), ...assistants]
@@ -46,7 +48,7 @@ const scenarios = [
 test.use({ viewport: { width: 646, height: 1385 } })
 
 for (const scenario of scenarios) {
-  test(`keeps the latest user turn visible through ${scenario.name}`, async ({ page }) => {
+  test(`keeps visible timeline content visible through ${scenario.name}`, async ({ page }) => {
     const requests: { before?: string; phase: "start" | "end" }[] = []
     const pages: { before?: string; limit: number }[] = []
     const roots: { sessionID: string; messageID: string }[] = []
@@ -101,36 +103,60 @@ for (const scenario of scenarios) {
         }
       },
     })
-    await page.addInitScript(
-      ({ userPartID, lastPartID }) => {
-        const state = { armed: false, hidden: false, samples: 0, stop: false }
-        ;(window as Window & { __historyRootProbe?: typeof state }).__historyRootProbe = state
-        const sample = () => {
-          if (state.armed) {
-            const virtual = document.querySelector<HTMLElement>("[data-timeline-virtual-content]")
-            const viewport = virtual?.closest<HTMLElement>(".scroll-view__viewport")
-            const view = viewport?.getBoundingClientRect()
-            const visible = (partID: string) => {
-              const part = viewport?.querySelector<HTMLElement>(`[data-timeline-part-id="${partID}"]`)
-              const rect = part?.getBoundingClientRect()
-              return (
-                !!rect &&
-                !!view &&
-                rect.width > 0 &&
-                rect.height > 0 &&
-                rect.bottom > view.top &&
-                rect.top < view.bottom
-              )
-            }
-            if (!virtual || !visible(userPartID) || !visible(lastPartID)) state.hidden = true
-            state.samples++
+    await page.addInitScript(() => {
+      const visibleParts = () => {
+        const virtual = document.querySelector<HTMLElement>("[data-timeline-virtual-content]")
+        const viewport = virtual?.closest<HTMLElement>(".scroll-view__viewport")
+        const view = viewport?.getBoundingClientRect()
+        if (!viewport || !view) return []
+        return [...viewport.querySelectorAll<HTMLElement>("[data-timeline-part-id]")]
+          .filter((part) => {
+            const rect = part.getBoundingClientRect()
+            return rect.width > 0 && rect.height > 0 && rect.bottom > view.top && rect.top < view.bottom
+          })
+          .flatMap((part) => (part.dataset.timelinePartId ? [part.dataset.timelinePartId] : []))
+      }
+      const state = {
+        armed: false,
+        hidden: false,
+        visibleParts: [] as string[],
+        samples: 0,
+        stop: false,
+        arm() {
+          state.visibleParts = visibleParts()
+          state.armed = true
+        },
+      }
+      ;(window as Window & { __historyRootProbe?: typeof state }).__historyRootProbe = state
+      const sample = () => {
+        if (state.armed) {
+          const virtual = document.querySelector<HTMLElement>("[data-timeline-virtual-content]")
+          const viewport = virtual?.closest<HTMLElement>(".scroll-view__viewport")
+          const view = viewport?.getBoundingClientRect()
+          const visible = (partID: string) => {
+            const part = viewport?.querySelector<HTMLElement>(`[data-timeline-part-id="${CSS.escape(partID)}"]`)
+            const rect = part?.getBoundingClientRect()
+            return (
+              !!rect &&
+              !!view &&
+              rect.width > 0 &&
+              rect.height > 0 &&
+              rect.bottom > view.top &&
+              rect.top < view.bottom
+            )
           }
-          if (!state.stop) requestAnimationFrame(() => setTimeout(sample, 0))
+          if (
+            !virtual ||
+            state.visibleParts.length === 0 ||
+            state.visibleParts.some((partID) => !visible(partID))
+          )
+            state.hidden = true
+          state.samples++
         }
-        requestAnimationFrame(() => setTimeout(sample, 0))
-      },
-      { userPartID, lastPartID },
-    )
+        if (!state.stop) requestAnimationFrame(() => setTimeout(sample, 0))
+      }
+      requestAnimationFrame(() => setTimeout(sample, 0))
+    })
 
     await page.goto(`/${base64Encode(directory)}/session/${sessionID}`)
     await transport.waitForConnection()
@@ -143,23 +169,28 @@ for (const scenario of scenarios) {
       "messages:start:latest",
       "messages:end:latest",
       `message:${userID}`,
-      `messages:start:${messages.at(-2)!.info.id}`,
+      `messages:start:${messages.at(-initialPageSize)!.info.id}`,
     ])
+    await expect(page.locator('[data-timeline-part-id^="prt_history_root_"]')).toHaveCount(initialPageSize)
     await page.evaluate(() => {
       ;(
         window as Window & {
-          __historyRootProbe?: { armed: boolean }
+          __historyRootProbe?: { arm(): void }
         }
-      ).__historyRootProbe!.armed = true
+      ).__historyRootProbe!.arm()
     })
     await waitForProbeSamples(page, 0)
-    expect(await historyRootHidden(page)).toBe(false)
+    expect(await visibleContentHidden(page)).toBe(false)
     const beforeHistory = await probeSamples(page)
     history.resolve()
-    await expect(page.locator('[data-timeline-part-id^="prt_history_root_"]')).toHaveCount(14)
+    await expect(page.locator('[data-timeline-part-id^="prt_history_root_"]')).toHaveCount(assistants.length)
+    await expect.poll(() => requests.filter((request) => request.phase === "end").length).toBe(2)
     await expect(page.getByRole("button", { name: "Stop" })).toBeVisible()
     await waitForProbeSamples(page, beforeHistory)
-    expect(pages[0]).toEqual({ before: undefined, limit: 2 })
+    expect(pages).toEqual([
+      { before: undefined, limit: initialPageSize },
+      { before: messages.at(-initialPageSize)!.info.id, limit: historyPageSize },
+    ])
     expect(roots).toEqual([{ sessionID, messageID: userID }])
 
     const message = messageUpdated(scenario.info)
@@ -213,7 +244,7 @@ async function waitForProbeSamples(page: Page, after: number) {
   )
 }
 
-function historyRootHidden(page: Page) {
+function visibleContentHidden(page: Page) {
   return page.evaluate(
     () => (window as Window & { __historyRootProbe?: { hidden: boolean } }).__historyRootProbe!.hidden,
   )

packages/app/src/context/server-session.test.ts

@@ -174,7 +174,7 @@ describe("server session", () => {
     await ctx.store.sync("root")
 
     expect(ctx.get).toEqual([{ sessionID: "root" }])
-    expect(ctx.messages).toEqual([{ sessionID: "root", limit: 2, before: undefined }])
+    expect(ctx.messages).toEqual([{ sessionID: "root", limit: 20, before: undefined }])
     expect(ctx.store.data.message.root).toEqual([])
   })
 
@@ -194,7 +194,7 @@ describe("server session", () => {
 
     await store.sync("child")
 
-    expect(client.requests).toEqual([{ sessionID: "child", limit: 2, before: undefined }])
+    expect(client.requests).toEqual([{ sessionID: "child", limit: 20, before: undefined }])
     expect(client.rootRequests).toEqual([{ sessionID: "child", messageID: user.id }])
     expect(store.data.message.child).toEqual([user, ...assistants])
     expect(store.history.more("child")).toBe(true)

packages/app/src/context/server-session.ts

@@ -21,7 +21,7 @@ import { dropSessionCaches, pickSessionCacheEvictions, SESSION_CACHE_LIMIT } fro
 const cmp = (a: string, b: string) => (a < b ? -1 : a > b ? 1 : 0)
 const cmpMessage = (a: Message, b: Message) => a.time.created - b.time.created || cmp(a.id, b.id)
 const SKIP_PARTS = new Set(["patch", "step-start", "step-finish"])
-const initialMessagePageSize = 2
+const initialMessagePageSize = 20
 const historyMessagePageSize = 200
 const sessionInfoLimit = 2_048
 const emptyIDs: ReadonlySet<string> = new Set()