diff --git a/packages/app/e2e/regression/session-timeline-collapse-state.spec.ts b/packages/app/e2e/regression/session-timeline-collapse-state.spec.ts index 271c2e2cab..18e63e2a1f 100644 --- a/packages/app/e2e/regression/session-timeline-collapse-state.spec.ts +++ b/packages/app/e2e/regression/session-timeline-collapse-state.spec.ts @@ -15,6 +15,15 @@ type EventPayload = { payload: Record } +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('[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('[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") diff --git a/packages/app/src/pages/session/message-timeline.data.ts b/packages/app/src/pages/session/message-timeline.data.ts index c6294d3e38..0643a424ea 100644 --- a/packages/app/src/pages/session/message-timeline.data.ts +++ b/packages/app/src/pages/session/message-timeline.data.ts @@ -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 diff --git a/packages/app/src/pages/session/message-timeline.tsx b/packages/app/src/pages/session/message-timeline.tsx index 60a161200d..6a37f9b4e7 100644 --- a/packages/app/src/pages/session/message-timeline.tsx +++ b/packages/app/src/pages/session/message-timeline.tsx @@ -257,7 +257,7 @@ function TimelineDiffView(props: { diff: SummaryDiff }) { return (
- +
) } @@ -438,6 +438,14 @@ export function MessageTimeline(props: { }) return result }) + const lastAssistantGroupKey = createMemo(() => { + const result = new Map() + 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 ( ) @@ -1046,6 +1054,7 @@ export function MessageTimeline(props: { toolOpen={toolOpen[part().id] ?? defaultOpen()} onToolOpenChange={(open) => setToolOpen(part().id, open)} deferToolContent={false} + virtualizeDiff={false} /> )} diff --git a/packages/ui/src/components/file.tsx b/packages/ui/src/components/file.tsx index 97d4d69f78..ebf1758b1a 100644 --- a/packages/ui/src/components/file.tsx +++ b/packages/ui/src/components/file.tsx @@ -88,6 +88,7 @@ type DiffBaseProps = FileDiffOptions & mode: "diff" annotations?: DiffLineAnnotation[] preloadedDiff?: DiffPreload + virtualize?: boolean } type DiffPairProps = DiffBaseProps & { @@ -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(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> | 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(props: TextFileProps) { let instance: PierreFile | VirtualizedFile | undefined + let renderMode: Virtualizer | "plain" | undefined let viewer!: Viewer const [local, others] = splitProps(props, textKeys) @@ -861,16 +874,20 @@ function TextViewer(props: TextFileProps) { 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(opts, virtualizer, codeMetrics, workerPool) : new PierreFile(opts, workerPool), + update: (value) => value.setOptions(opts), assign: (value) => { instance = value + renderMode = nextRenderMode }, draw: (value) => { const contents = text() @@ -895,6 +912,7 @@ function TextViewer(props: TextFileProps) { onCleanup(() => { instance?.cleanUp() instance = undefined + renderMode = undefined virtuals.cleanup() }) @@ -907,6 +925,7 @@ function TextViewer(props: TextFileProps) { function DiffViewer(props: DiffFileProps) { let instance: FileDiff | undefined + let renderMode: Virtualizer | "plain" | undefined let dragSide: DiffSelectionSide | undefined let dragEndSide: DiffSelectionSide | undefined let viewer!: Viewer @@ -991,7 +1010,10 @@ function DiffViewer(props: DiffFileProps) { 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(props: DiffFileProps) { 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(props: DiffFileProps) { renderViewer({ viewer, current: instance, + reset: renderMode !== undefined && renderMode !== nextRenderMode, create: () => virtualizer ? new VirtualizedFileDiff(opts, virtualizer, virtualMetrics, workerPool) : new FileDiff(opts, workerPool), + update: (value) => value.setOptions(opts), assign: (value) => { instance = value + renderMode = nextRenderMode }, draw: (value) => { if (local.fileDiff) { @@ -1111,6 +1137,7 @@ function DiffViewer(props: DiffFileProps) { onCleanup(() => { instance?.cleanUp() instance = undefined + renderMode = undefined virtuals.cleanup() dragSide = undefined dragEndSide = undefined diff --git a/packages/ui/src/components/message-part.tsx b/packages/ui/src/components/message-part.tsx index 7e5ae7e51d..89549b5f4d 100644 --- a/packages/ui/src/components/message-part.tsx +++ b/packages/ui/src/components/message-part.tsx @@ -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} /> @@ -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({ } >
- +
@@ -2186,6 +2204,7 @@ ToolRegistry.register({ @@ -2258,7 +2277,7 @@ ToolRegistry.register({ } >
- +
diff --git a/packages/ui/src/components/session-diff.ts b/packages/ui/src/components/session-diff.ts index 52ef0a4bbc..0427de0dc3 100644 --- a/packages/ui/src/components/session-diff.ts +++ b/packages/ui/src/components/session-diff.ts @@ -14,6 +14,13 @@ type LegacyDiff = { type SnapshotDiff = SnapshotFileDiff & { file: string } type ReviewDiff = SnapshotDiff | VcsFileDiff | LegacyDiff +export type DiffSource = Pick +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() +const diffCacheLimit = 16 +const fileDiffCache = new Map() +const patchTextCache = new Map() +// 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(cache: Map, key: K) { + const value = cache.get(key) + if (value === undefined) return + cache.delete(key) + cache.set(key, value) + return value +} + +function setMapCache(cache: Map, 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), } }