fix(app): keep timeline diffs mounted

This commit is contained in:
LukeParkerDev 2026-05-20 16:23:51 +10:00
parent aedbe006a0
commit ebd2a2964a
6 changed files with 294 additions and 77 deletions

View file

@ -15,6 +15,15 @@ type EventPayload = {
payload: Record<string, unknown>
}
declare global {
interface Window {
__timelineDiffProbe: {
reset: () => void
shadowRoots: () => number
}
}
}
const userMessage = {
info: {
id: userMessageID,
@ -124,6 +133,44 @@ test.describe("regression: session timeline local row state", () => {
streamedTextVisible: true,
})
})
test("does not remount an edit diff when sibling parts or diff counts update", async ({ page }) => {
const events: EventPayload[] = []
await installDiffProbe(page)
await mockServer(page, events)
await configurePage(page)
await page.goto(`/${base64Encode(directory)}/session/${sessionID}`)
await expect(page.getByRole("heading", { name: title })).toBeVisible()
const wrapper = page.locator(`[data-timeline-part-id="${editPartID}"]`).first()
await expect(wrapper).toBeVisible()
await expect(wrapper.locator('[data-component="file"][data-mode="diff"]').first()).toBeVisible()
await markDiffProbe(page)
events.push({
directory,
payload: {
type: "message.part.updated",
properties: { part: streamedTextPart },
},
})
await expect(page.locator(`[data-timeline-part-id="${textPartID}"]`).first()).toBeVisible({ timeout: 10_000 })
expect(await readDiffProbe(page)).toEqual({ fileMarker: "before", shadowRoots: 0, toolMarker: "before" })
await markDiffProbe(page)
events.push({
directory,
payload: {
type: "message.part.updated",
properties: { part: editPartWithAdditions(2) },
},
})
await expect(wrapper.locator('[data-slot="diff-changes-additions"]').filter({ hasText: "+2" }).first()).toBeVisible({ timeout: 10_000 })
expect(await readDiffProbe(page)).toEqual({ fileMarker: "before", shadowRoots: 0, toolMarker: "before" })
})
})
async function configurePage(page: Page) {
@ -166,6 +213,63 @@ async function readToolState(page: Page) {
}), textPartID)
}
async function installDiffProbe(page: Page) {
await page.addInitScript(() => {
let shadowRootCount = 0
const attachShadow = Element.prototype.attachShadow
Element.prototype.attachShadow = function (init) {
shadowRootCount += 1
return attachShadow.call(this, init)
}
window.__timelineDiffProbe = {
reset: () => {
shadowRootCount = 0
},
shadowRoots: () => shadowRootCount,
}
})
}
async function markDiffProbe(page: Page) {
await page.locator(`[data-timeline-part-id="${editPartID}"]`).first().evaluate((element) => {
const tool = element as HTMLElement
const file = tool.querySelector<HTMLElement>('[data-component="file"][data-mode="diff"]')
if (!file) throw new Error("missing edit diff file")
tool.dataset.timelineProbe = "before"
file.dataset.timelineProbe = "before"
window.__timelineDiffProbe.reset()
})
}
async function readDiffProbe(page: Page) {
return page.locator(`[data-timeline-part-id="${editPartID}"]`).first().evaluate((element) => {
const tool = element as HTMLElement
const file = tool.querySelector<HTMLElement>('[data-component="file"][data-mode="diff"]')
return {
fileMarker: file?.dataset.timelineProbe,
shadowRoots: window.__timelineDiffProbe.shadowRoots(),
toolMarker: tool.dataset.timelineProbe,
}
})
}
function editPartWithAdditions(additions: number) {
return {
...editPart,
state: {
...editPart.state,
metadata: {
...editPart.state.metadata,
filediff: {
...editPart.state.metadata.filediff,
additions,
},
},
},
}
}
function readExpanded(element: Element) {
const trigger = element.querySelector('[data-slot="collapsible-trigger"]')
const aria = trigger?.getAttribute("aria-expanded")

View file

@ -23,7 +23,6 @@ export type TimelineRowMap = {
userMessageID: string
group: PartGroup
previousAssistantPart: boolean
lastAssistantPart: boolean
}
Thinking: { userMessageID: string; reasoningHeading?: string }
Retry: { userMessageID: string }
@ -50,7 +49,6 @@ export namespace TimelineRow {
userMessageID: string
group: PartGroup
previousAssistantPart: boolean
lastAssistantPart: boolean
}> {}
export class Thinking extends Data.TaggedClass("Thinking")<{
userMessageID: string
@ -151,8 +149,6 @@ export namespace Timeline {
),
]
: groupParts(assistantPartRefs).map((group) => ({ type: "part" as const, group }))
const assistantGroupCount = assistantItems.filter((item) => item.type === "part").length
if (comments.length > 0)
rows.push(
new TimelineRow.CommentStrip({
@ -195,7 +191,6 @@ export namespace Timeline {
userMessageID: userMessage.id,
group: item.group,
previousAssistantPart: assistantGroupIndex > 0,
lastAssistantPart: assistantGroupIndex === assistantGroupCount - 1,
}),
)
assistantGroupIndex += 1

View file

@ -257,7 +257,7 @@ function TimelineDiffView(props: { diff: SummaryDiff }) {
return (
<div data-slot="session-turn-diff-view" data-scrollable>
<Dynamic component={fileComponent} mode="diff" fileDiff={view.fileDiff} />
<Dynamic component={fileComponent} mode="diff" virtualize={false} fileDiff={view.fileDiff} />
</div>
)
}
@ -438,6 +438,14 @@ export function MessageTimeline(props: {
})
return result
})
const lastAssistantGroupKey = createMemo(() => {
const result = new Map<string, string>()
timelineRows().forEach((row) => {
if (row._tag !== "AssistantPart") return
result.set(row.userMessageID, row.group.key)
})
return result
})
const keepMounted = createMemo(() => {
const id = activeMessageID()
if (!id) return
@ -1010,7 +1018,7 @@ export function MessageTimeline(props: {
return (
<ContextToolGroup
parts={parts()}
busy={workingTurn(row().userMessageID) && row().lastAssistantPart}
busy={workingTurn(row().userMessageID) && lastAssistantGroupKey().get(row().userMessageID) === row().group.key}
onSizeChange={measureTimeline}
/>
)
@ -1046,6 +1054,7 @@ export function MessageTimeline(props: {
toolOpen={toolOpen[part().id] ?? defaultOpen()}
onToolOpenChange={(open) => setToolOpen(part().id, open)}
deferToolContent={false}
virtualizeDiff={false}
/>
)}
</Show>

View file

@ -88,6 +88,7 @@ type DiffBaseProps<T> = FileDiffOptions<T> &
mode: "diff"
annotations?: DiffLineAnnotation<T>[]
preloadedDiff?: DiffPreload<T>
virtualize?: boolean
}
type DiffPairProps<T> = DiffBaseProps<T> & {
@ -123,7 +124,7 @@ const sharedKeys = [
] as const
const textKeys = ["file", ...sharedKeys] as const
const diffKeys = ["fileDiff", "before", "after", ...sharedKeys] as const
const diffKeys = ["fileDiff", "before", "after", "virtualize", ...sharedKeys] as const
// ---------------------------------------------------------------------------
// Shared viewer hook
@ -482,17 +483,24 @@ function notifyRendered(opts: {
function renderViewer<I extends RenderTarget>(opts: {
viewer: Viewer
current: I | undefined
reset?: boolean
create: () => I
update?: (value: I) => void
assign: (value: I) => void
draw: (value: I) => void
onReady: () => void
}) {
clearReadyWatcher(opts.viewer.ready)
opts.current?.cleanUp()
const next = opts.create()
opts.assign(next)
const reset = opts.reset === true && opts.current !== undefined
if (reset) opts.current?.cleanUp()
const next = reset || !opts.current ? opts.create() : opts.current
if (reset || !opts.current) {
opts.viewer.container.innerHTML = ""
opts.assign(next)
} else {
opts.update?.(next)
}
opts.viewer.container.innerHTML = ""
opts.draw(next)
applyViewerScheme(opts.viewer.getHost())
@ -566,7 +574,7 @@ function createLocalVirtualStrategy(host: () => HTMLDivElement | undefined, enab
}
}
function createSharedVirtualStrategy(host: () => HTMLDivElement | undefined): VirtualStrategy {
function createSharedVirtualStrategy(host: () => HTMLDivElement | undefined, enabled: () => boolean): VirtualStrategy {
let shared: NonNullable<ReturnType<typeof acquireVirtualizer>> | undefined
const release = () => {
@ -576,6 +584,10 @@ function createSharedVirtualStrategy(host: () => HTMLDivElement | undefined): Vi
return {
get: () => {
if (!enabled()) {
release()
return
}
if (shared) return shared.virtualizer
const container = host()
@ -689,6 +701,7 @@ function ViewerShell(props: {
function TextViewer<T>(props: TextFileProps<T>) {
let instance: PierreFile<T> | VirtualizedFile<T> | undefined
let renderMode: Virtualizer | "plain" | undefined
let viewer!: Viewer
const [local, others] = splitProps(props, textKeys)
@ -861,16 +874,20 @@ function TextViewer<T>(props: TextFileProps<T>) {
const isVirtual = virtual()
const virtualizer = virtuals.get()
const nextRenderMode = isVirtual && virtualizer ? virtualizer : "plain"
renderViewer({
viewer,
current: instance,
reset: renderMode !== undefined && renderMode !== nextRenderMode,
create: () =>
isVirtual && virtualizer
? new VirtualizedFile<T>(opts, virtualizer, codeMetrics, workerPool)
: new PierreFile<T>(opts, workerPool),
update: (value) => value.setOptions(opts),
assign: (value) => {
instance = value
renderMode = nextRenderMode
},
draw: (value) => {
const contents = text()
@ -895,6 +912,7 @@ function TextViewer<T>(props: TextFileProps<T>) {
onCleanup(() => {
instance?.cleanUp()
instance = undefined
renderMode = undefined
virtuals.cleanup()
})
@ -907,6 +925,7 @@ function TextViewer<T>(props: TextFileProps<T>) {
function DiffViewer<T>(props: DiffFileProps<T>) {
let instance: FileDiff<T> | undefined
let renderMode: Virtualizer | "plain" | undefined
let dragSide: DiffSelectionSide | undefined
let dragEndSide: DiffSelectionSide | undefined
let viewer!: Viewer
@ -991,7 +1010,10 @@ function DiffViewer<T>(props: DiffFileProps<T>) {
adapter,
)
const virtuals = createSharedVirtualStrategy(() => viewer.container)
const virtuals = createSharedVirtualStrategy(
() => viewer.container,
() => local.virtualize !== false,
)
const large = createMemo(() => {
if (local.fileDiff) {
@ -1056,6 +1078,7 @@ function DiffViewer<T>(props: DiffFileProps<T>) {
const opts = options()
const workerPool = large() ? getWorkerPool("unified") : getWorkerPool(props.diffStyle)
const virtualizer = virtuals.get()
const nextRenderMode = virtualizer ?? "plain"
const beforeContents = typeof local.before?.contents === "string" ? local.before.contents : ""
const afterContents = typeof local.after?.contents === "string" ? local.after.contents : ""
const done = preserve(viewer)
@ -1070,12 +1093,15 @@ function DiffViewer<T>(props: DiffFileProps<T>) {
renderViewer({
viewer,
current: instance,
reset: renderMode !== undefined && renderMode !== nextRenderMode,
create: () =>
virtualizer
? new VirtualizedFileDiff<T>(opts, virtualizer, virtualMetrics, workerPool)
: new FileDiff<T>(opts, workerPool),
update: (value) => value.setOptions(opts),
assign: (value) => {
instance = value
renderMode = nextRenderMode
},
draw: (value) => {
if (local.fileDiff) {
@ -1111,6 +1137,7 @@ function DiffViewer<T>(props: DiffFileProps<T>) {
onCleanup(() => {
instance?.cleanUp()
instance = undefined
renderMode = undefined
virtuals.cleanup()
dragSide = undefined
dragEndSide = undefined

View file

@ -177,6 +177,7 @@ export interface MessagePartProps {
toolOpen?: boolean
onToolOpenChange?: (open: boolean) => void
deferToolContent?: boolean
virtualizeDiff?: boolean
showAssistantCopyPartID?: string | null
turnDurationMs?: number
}
@ -291,7 +292,7 @@ function getDirectory(path: string | undefined) {
}
import type { IconProps } from "./icon"
import { normalize } from "./session-diff"
import { normalize, resolveFileDiff } from "./session-diff"
export type ToolInfo = {
icon: IconProps["name"]
@ -1269,6 +1270,7 @@ export function Part(props: MessagePartProps) {
toolOpen={props.toolOpen}
onToolOpenChange={props.onToolOpenChange}
deferToolContent={props.deferToolContent}
virtualizeDiff={props.virtualizeDiff}
showAssistantCopyPartID={props.showAssistantCopyPartID}
turnDurationMs={props.turnDurationMs}
/>
@ -1288,6 +1290,7 @@ export interface ToolProps {
open?: boolean
onOpenChange?: (open: boolean) => void
deferContent?: boolean
virtualizeDiff?: boolean
forceOpen?: boolean
locked?: boolean
}
@ -1433,6 +1436,7 @@ PART_MAPPING["tool"] = function ToolPartDisplay(props) {
open={controlledOpen()}
onOpenChange={props.onToolOpenChange ? handleToolOpenChange : undefined}
deferContent={props.deferToolContent}
virtualizeDiff={props.virtualizeDiff}
/>
</Match>
</Switch>
@ -1936,15 +1940,29 @@ ToolRegistry.register({
const path = createMemo(() => props.metadata?.filediff?.file || props.input.filePath || "")
const filename = () => getFilename(props.input.filePath ?? "")
const pending = () => props.status === "pending" || props.status === "running"
const diffSource = createMemo(
() => {
const filediff = props.metadata?.filediff
if (!filediff) return
return {
file: filediff.file || props.input.filePath || "",
patch: typeof filediff.patch === "string" ? filediff.patch : undefined,
before: typeof filediff.before === "string" ? filediff.before : undefined,
after: typeof filediff.after === "string" ? filediff.after : undefined,
}
},
undefined,
{
equals: (a, b) =>
a?.file === b?.file && a?.patch === b?.patch && a?.before === b?.before && a?.after === b?.after,
},
)
const fileCompProps = createMemo(() => {
try {
if (props.metadata?.filediff) {
const diff = normalize({
...props.metadata?.filediff,
status: "modified",
})
const fileDiff = diff.fileDiff
const source = diffSource()
if (source) {
const fileDiff = resolveFileDiff(source)
if (fileDiff) return { fileDiff, hunkSeparators: fileDiff.isPartial ? "simple" : "line-info-basic" }
}
} catch {}
@ -2002,7 +2020,7 @@ ToolRegistry.register({
}
>
<div data-component="edit-content">
<Dynamic component={fileComponent} mode="diff" {...fileCompProps()} />
<Dynamic component={fileComponent} mode="diff" virtualize={props.virtualizeDiff} {...fileCompProps()} />
</div>
</ToolFileAccordion>
</Show>
@ -2186,6 +2204,7 @@ ToolRegistry.register({
<Dynamic
component={fileComponent}
mode="diff"
virtualize={props.virtualizeDiff}
fileDiff={file.view.fileDiff}
hunkSeparators={file.view.fileDiff.isPartial ? "simple" : "line-info-basic"}
/>
@ -2258,7 +2277,7 @@ ToolRegistry.register({
}
>
<div data-component="apply-patch-file-diff">
<Dynamic component={fileComponent} mode="diff" fileDiff={single()!.view.fileDiff} />
<Dynamic component={fileComponent} mode="diff" virtualize={props.virtualizeDiff} fileDiff={single()!.view.fileDiff} />
</div>
</ToolFileAccordion>
</BasicTool>

View file

@ -14,6 +14,13 @@ type LegacyDiff = {
type SnapshotDiff = SnapshotFileDiff & { file: string }
type ReviewDiff = SnapshotDiff | VcsFileDiff | LegacyDiff
export type DiffSource = Pick<LegacyDiff, "file" | "patch" | "before" | "after">
type PatchData = {
before: string
after: string
patch: string
patchIsPartial: boolean
}
export type ViewDiff = {
file: string
@ -24,66 +31,115 @@ export type ViewDiff = {
fileDiff: FileDiffMetadata
}
const cache = new Map<string, FileDiffMetadata>()
const diffCacheLimit = 16
const fileDiffCache = new Map<string, FileDiffMetadata>()
const patchTextCache = new Map<string, PatchData>()
// Keep this before structuredPatch/formatPatch; those dominate huge diff metadata updates.
const contentPatchCache: { file: string; before: string; after: string; value: PatchData }[] = []
function patch(diff: ReviewDiff) {
function mapCache<K, V>(cache: Map<K, V>, key: K) {
const value = cache.get(key)
if (value === undefined) return
cache.delete(key)
cache.set(key, value)
return value
}
function setMapCache<K, V>(cache: Map<K, V>, key: K, value: V) {
cache.delete(key)
cache.set(key, value)
while (cache.size > diffCacheLimit) cache.delete(cache.keys().next().value!)
return value
}
function patch(diff: DiffSource) {
if (typeof diff.patch === "string") {
try {
const [patch] = parsePatch(diff.patch)
const beforeLines: Array<{ text: string; newline: boolean }> = []
const afterLines: Array<{ text: string; newline: boolean }> = []
let previous: "-" | "+" | " " | undefined
return patchFromText(diff.patch)
}
const patchIsPartial = patch.hunks.every((h) => h.oldStart > 1)
return patchFromContent(
diff.file,
typeof diff.before === "string" ? diff.before : "",
typeof diff.after === "string" ? diff.after : "",
)
}
for (const hunk of patch.hunks) {
for (const line of hunk.lines) {
if (line.startsWith("\\")) {
if (previous === "-" || previous === " ") {
const before = beforeLines.at(-1)
if (before) before.newline = false
}
if (previous === "+" || previous === " ") {
const after = afterLines.at(-1)
if (after) after.newline = false
}
continue
function patchFromText(value: string): PatchData {
const cached = mapCache(patchTextCache, value)
if (cached) return cached
return setMapCache(patchTextCache, value, parsePatchText(value))
}
function parsePatchText(value: string): PatchData {
try {
const [patch] = parsePatch(value)
const beforeLines: Array<{ text: string; newline: boolean }> = []
const afterLines: Array<{ text: string; newline: boolean }> = []
let previous: "-" | "+" | " " | undefined
const patchIsPartial = patch.hunks.every((h) => h.oldStart > 1)
for (const hunk of patch.hunks) {
for (const line of hunk.lines) {
if (line.startsWith("\\")) {
if (previous === "-" || previous === " ") {
const before = beforeLines.at(-1)
if (before) before.newline = false
}
if (line.startsWith("-")) {
beforeLines.push({ text: line.slice(1), newline: true })
previous = "-"
} else if (line.startsWith("+")) {
afterLines.push({ text: line.slice(1), newline: true })
previous = "+"
} else {
// context line (starts with ' ')
beforeLines.push({ text: line.slice(1), newline: true })
afterLines.push({ text: line.slice(1), newline: true })
previous = " "
if (previous === "+" || previous === " ") {
const after = afterLines.at(-1)
if (after) after.newline = false
}
continue
}
if (line.startsWith("-")) {
beforeLines.push({ text: line.slice(1), newline: true })
previous = "-"
} else if (line.startsWith("+")) {
afterLines.push({ text: line.slice(1), newline: true })
previous = "+"
} else {
// context line (starts with ' ')
beforeLines.push({ text: line.slice(1), newline: true })
afterLines.push({ text: line.slice(1), newline: true })
previous = " "
}
}
return {
before: beforeLines.map((line) => line.text + (line.newline ? "\n" : "")).join(""),
after: afterLines.map((line) => line.text + (line.newline ? "\n" : "")).join(""),
patch: diff.patch,
patchIsPartial,
}
} catch {
return { before: "", after: "", patch: diff.patch, patchIsPartial: false }
}
return {
before: beforeLines.map((line) => line.text + (line.newline ? "\n" : "")).join(""),
after: afterLines.map((line) => line.text + (line.newline ? "\n" : "")).join(""),
patch: value,
patchIsPartial,
}
} catch {
return { before: "", after: "", patch: value, patchIsPartial: false }
}
return {
before: "before" in diff && typeof diff.before === "string" ? diff.before : "",
after: "after" in diff && typeof diff.after === "string" ? diff.after : "",
}
function patchFromContent(file: string, before: string, after: string): PatchData {
const index = contentPatchCache.findIndex(
(entry) => entry.file === file && entry.before === before && entry.after === after,
)
if (index !== -1) {
const entry = contentPatchCache[index]!
contentPatchCache.splice(index, 1)
contentPatchCache.push(entry)
return entry.value
}
const value = {
before,
after,
patch: formatPatch(
structuredPatch(
diff.file,
diff.file,
"before" in diff && typeof diff.before === "string" ? diff.before : "",
"after" in diff && typeof diff.after === "string" ? diff.after : "",
file,
file,
before,
after,
"",
"",
{ context: Number.MAX_SAFE_INTEGER },
@ -91,30 +147,37 @@ function patch(diff: ReviewDiff) {
),
patchIsPartial: false,
}
contentPatchCache.push({ file, before, after, value })
while (contentPatchCache.length > diffCacheLimit) contentPatchCache.shift()
return value
}
function file(file: string, patch: string, before: string, after: string, partial = false) {
const hit = cache.get(patch)
function fileDiff(file: string, patch: string, before: string, after: string, partial = false) {
const hit = mapCache(fileDiffCache, patch)
if (hit) return hit
let value: FileDiffMetadata | undefined
if (partial) value = parsePatchFiles(patch)[0]?.files[0]
if (value === undefined) value = parseDiffFromFile({ name: file, contents: before }, { name: file, contents: after })
cache.set(patch, value)
return value
return setMapCache(fileDiffCache, patch, value)
}
export function resolveFileDiff(diff: DiffSource) {
const next = patch(diff)
return fileDiff(diff.file, next.patch, next.before, next.after, next.patchIsPartial)
}
export function normalize(diff: ReviewDiff): ViewDiff {
const next = patch(diff)
const fileDiff = file(diff.file, next.patch, next.before, next.after, next.patchIsPartial)
return {
file: diff.file,
patch: next.patch,
additions: diff.additions,
deletions: diff.deletions,
status: diff.status,
fileDiff,
fileDiff: fileDiff(diff.file, next.patch, next.before, next.after, next.patchIsPartial),
}
}