Skip to content

perf(flow-chat): optimize long-round streaming projection and scroll rendering #3778

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Raing5Days wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(flow-chat): optimize long-round streaming projection and scroll rendering #3778

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 14 additions & 7 deletions src/web-ui/src/flow_chat/components/modern/VirtualMessageList.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1303,14 +1303,15 @@ const VirtualMessageListSession = forwardRef<VirtualMessageListRef, VirtualMessa
}, [activeSessionId, reader, viewportAnchor, viewportOwner, virtualItems]);

const publishViewportSnapshot = useCallback((sample?: FlowChatViewportGeometry) => {
if (!onViewportSnapshotRef.current) return;
const snapshot = captureViewportSnapshot(sample);
if (snapshot) onViewportSnapshotRef.current?.(snapshot);
if (snapshot) onViewportSnapshotRef.current(snapshot);
}, [captureViewportSnapshot]);
const publishViewportSnapshotRef = useRef(publishViewportSnapshot);
publishViewportSnapshotRef.current = publishViewportSnapshot;

const scheduleViewportSnapshot = useCallback(() => {
if (viewportSnapshotFrameRef.current !== null) return;
if (!onViewportSnapshotRef.current || viewportSnapshotFrameRef.current !== null) return;
viewportSnapshotFrameRef.current = requestAnimationFrame(() => {
viewportSnapshotFrameRef.current = null;
publishViewportSnapshotRef.current();
Expand Down Expand Up @@ -1605,20 +1606,21 @@ const VirtualMessageListSession = forwardRef<VirtualMessageListRef, VirtualMessa
} = nativeScrollCallbacksRef.current;
// Share this event's read phase. It cannot outlive the event: commits,
// ResizeObserver and explicit navigation must measure their own layout.
const geometry = readFlowChatViewportGeometry(scrollerElement);
const position = viewportOwner.readReaderScrollPosition();
const previous = readerScrollPositionRef.current;
readerScrollPositionRef.current = position;
const owner = viewportOwner.currentOwner();
const delta = previous === null ? 0 : position - previous;
const direction = delta < -0.5 ? 'before' : delta > 0.5 ? 'after' : undefined;
let geometrySample: FlowChatViewportGeometry | undefined;
const getGeometry = () => (geometrySample ??= readFlowChatViewportGeometry(scrollerElement));
// Synchronous corrections are removed by the register. Smooth owned
// navigation/follow scrolls are excluded here; unowned momentum counts.
if (direction && (owner === null || owner === 'user-gesture')) {
captureLeadingExtent(readFlowChatScrollTop(scrollerElement), true, geometry);
captureLeadingExtent(readFlowChatScrollTop(scrollerElement), true, getGeometry());
historyPager.readerIntent(direction);
}
if (owner === 'one-shot-navigation') captureLeadingExtent(readFlowChatScrollTop(scrollerElement), true, geometry);
if (owner === 'one-shot-navigation') captureLeadingExtent(readFlowChatScrollTop(scrollerElement), true, getGeometry());
/*
* A scroll under a scrollbar press is the one case where a plain scroll
* event does carry intent — the press is what qualifies it. Left
Expand All @@ -1638,9 +1640,14 @@ const VirtualMessageListSession = forwardRef<VirtualMessageListRef, VirtualMessa
* Re-anchor to wherever the user has arrived, and only there. The rule
* for "was this scroll the user's" lives with the anchor.
*/
viewportAnchor.captureAnchorForScroll(readFlowChatScrollTop(scrollerElement) === geometry.scrollTop ? geometry : undefined);
viewportAnchor.captureAnchorForScroll(
geometrySample && readFlowChatScrollTop(scrollerElement) === geometrySample.scrollTop ? geometrySample : undefined,
);
scheduleVisibleTurnInfoUpdate();
publishViewportSnapshot(readFlowChatScrollTop(scrollerElement) === geometry.scrollTop ? geometry : undefined);
if (onViewportSnapshotRef.current) {
const sample = getGeometry();
publishViewportSnapshot(readFlowChatScrollTop(scrollerElement) === sample.scrollTop ? sample : undefined);
}
/*
* Paging is a question about where the reader is, so a scroll is its
* primary input. Through a ref: this listener must not be torn down and
Expand Down
102 changes: 66 additions & 36 deletions src/web-ui/src/flow_chat/components/modern/useFlowChatSearch.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

import { useState, useMemo, useCallback } from 'react';
import type { VirtualItem } from '../../store/modernFlowChatStore';
import { buildInlineFlowGroupData, getProjectedModelRoundGroups } from '../../grouping/roundGroups';
import { getCachedInlineFlowGroupData, getProjectedModelRoundGroups } from '../../grouping/roundGroups';

interface SearchableFlowItem {
id?: string;
Expand Down Expand Up @@ -46,6 +46,7 @@ export interface UseFlowChatSearchReturn {

interface SearchableSource {
content: string;
foldedContent: string;
flowItemId?: string;
expandableIds?: readonly string[];
}
Expand All @@ -63,53 +64,62 @@ function flowItemSearchSources(
...(outerExpandableId ? [outerExpandableId] : []),
...(item.type === 'thinking' && item.id ? [item.id] : []),
];
const content = item.content ?? '';

return [{
content: item.content ?? '',
content,
foldedContent: content.toLowerCase(),
flowItemId: item.id,
expandableIds: expandableIds.length > 0 ? expandableIds : undefined,
}];
});
}

const virtualItemSearchSourcesCache = new WeakMap<VirtualItem, SearchableSource[]>();

function getVirtualItemSearchSources(item: VirtualItem): SearchableSource[] {
const cached = virtualItemSearchSourcesCache.get(item);
if (cached) return cached;
let sources: SearchableSource[];
if (item.type === 'background-run') {
return [{ content: item.turn.userMessage.content }, ...flowItemSearchSources(item.turn.modelRounds.flatMap(round => round.items))];
}
if (item.type === 'user-message' || item.type === 'user-steering-message') {
return [{ content: item.data?.content ?? '' }];
}
if (item.type === 'model-round') {
return getProjectedModelRoundGroups(item).flatMap(group => group.type !== 'critical'
? flowItemSearchSources(group.items, buildInlineFlowGroupData(item.data.id, group).groupId)
const userContent = item.turn.userMessage.content;
sources = [
{ content: userContent, foldedContent: userContent.toLowerCase() },
...flowItemSearchSources(item.turn.modelRounds.flatMap(round => round.items)),
];
} else if (item.type === 'user-message' || item.type === 'user-steering-message') {
const content = item.data?.content ?? '';
sources = [{ content, foldedContent: content.toLowerCase() }];
} else if (item.type === 'model-round') {
sources = getProjectedModelRoundGroups(item).flatMap((group, index) => group.type !== 'critical'
? flowItemSearchSources(group.items, getCachedInlineFlowGroupData(item.data, index, group).groupId)
: flowItemSearchSources([group.item]));
} else if (item.type === 'explore-group') {
sources = flowItemSearchSources(item.data.allItems, item.data.groupId);
} else if (item.type === 'turn-completion-notice') {
const content = item.data.reasonCode;
sources = [{ content, foldedContent: content.toLowerCase() }];
} else if (item.type === 'turn-failure-notice') {
const content = [
item.data.error,
item.data.errorDetail?.provider,
item.data.errorDetail?.providerCode,
item.data.errorDetail?.providerMessage,
item.data.errorDetail?.requestId,
].filter(Boolean).join(' ');
sources = [{ content, foldedContent: content.toLowerCase() }];
} else {
sources = [];
}
if (item.type === 'explore-group') {
return flowItemSearchSources(item.data.allItems, item.data.groupId);
}
if (item.type === 'turn-completion-notice') {
return [{ content: item.data.reasonCode }];
}
if (item.type === 'turn-failure-notice') {
return [{
content: [
item.data.error,
item.data.errorDetail?.provider,
item.data.errorDetail?.providerCode,
item.data.errorDetail?.providerMessage,
item.data.errorDetail?.requestId,
].filter(Boolean).join(' '),
}];
}
return [];
virtualItemSearchSourcesCache.set(item, sources);
return sources;
}

function countQueryOccurrences(content: string, foldedQuery: string): number {
const haystack = content.toLowerCase();
function countFoldedQueryOccurrences(foldedContent: string, foldedQuery: string): number {
let count = 0;
let from = 0;
for (;;) {
const index = haystack.indexOf(foldedQuery, from);
const index = foldedContent.indexOf(foldedQuery, from);
if (index < 0) {
return count;
}
Expand All @@ -118,20 +128,36 @@ function countQueryOccurrences(content: string, foldedQuery: string): number {
}
}

const emptySearchMatches: SearchMatch[] = [];
const perItemSearchMatchCache = new WeakMap<VirtualItem, {
query: string;
virtualItemIndex: number;
matches: readonly SearchMatch[];
}>();

export function buildFlowChatSearchMatches(
virtualItems: readonly VirtualItem[],
searchQuery: string,
): SearchMatch[] {
const trimmed = searchQuery.trim();
if (!trimmed) return [];
if (!trimmed) return emptySearchMatches;
const query = trimmed.toLowerCase();
const matches: SearchMatch[] = [];

virtualItems.forEach((item, virtualItemIndex) => {
for (let virtualItemIndex = 0; virtualItemIndex < virtualItems.length; virtualItemIndex += 1) {
const item = virtualItems[virtualItemIndex];
const cached = perItemSearchMatchCache.get(item);
if (cached && cached.query === query && cached.virtualItemIndex === virtualItemIndex) {
for (let i = 0; i < cached.matches.length; i += 1) {
matches.push(cached.matches[i]);
}
continue;
}
const itemMatches: SearchMatch[] = [];
for (const source of getVirtualItemSearchSources(item)) {
const occurrenceCount = countQueryOccurrences(source.content, query);
const occurrenceCount = countFoldedQueryOccurrences(source.foldedContent, query);
for (let occurrenceIndex = 0; occurrenceIndex < occurrenceCount; occurrenceIndex += 1) {
matches.push({
itemMatches.push({
virtualItemIndex,
turnId: item.turnId,
type: item.type,
Expand All @@ -141,7 +167,11 @@ export function buildFlowChatSearchMatches(
});
}
}
});
perItemSearchMatchCache.set(item, { query, virtualItemIndex, matches: itemMatches });
for (let i = 0; i < itemMatches.length; i += 1) {
matches.push(itemMatches[i]);
}
}

return matches;
}
Expand Down
43 changes: 30 additions & 13 deletions src/web-ui/src/flow_chat/components/modern/useFlowChatVirtualizer.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,25 @@ function useMeasuredVirtualizer(options: VirtualizerOptions<HTMLElement, HTMLEle

const requestRender = useCallback((sync: boolean) => {
if (!mounted.current || (!sync && requested.current)) return;
if (commits.current < MAX_SYNCHRONOUS_VIRTUALIZER_COMMITS) {
pending.current = false;
requested.current = true;
commits.current += 1;
if (frame.current === null) {
frame.current = requestAnimationFrame(() => {
frame.current = null;
commits.current = 0;
if (!mounted.current || !pending.current) return;
pending.current = false;
requested.current = true;
commits.current = 1;
rerender();
});
}
if (sync) flushSync(rerender);
else rerender();
return;
}
pending.current = true;
if (frame.current === null) {
frame.current = requestAnimationFrame(() => {
Expand All @@ -56,12 +75,6 @@ function useMeasuredVirtualizer(options: VirtualizerOptions<HTMLElement, HTMLEle
rerender();
});
}
if (commits.current >= MAX_SYNCHRONOUS_VIRTUALIZER_COMMITS) return;
pending.current = false;
requested.current = true;
commits.current += 1;
if (sync) flushSync(rerender);
else rerender();
}, []);

instance.setOptions({
Expand Down Expand Up @@ -756,7 +769,8 @@ export function useFlowChatVirtualizer<T>({
const key = keys[index] ?? getItemKeyRef.current(item);
const knownHeight = getKnownItemHeightPx(item);
if (knownHeight !== undefined) next.set(key, knownHeight);
if (knownHeight === undefined && !previous.has(key)) continue;
const prevHeight = previous.get(key);
if (knownHeight === prevHeight) continue;

// A mounted row has real geometry. An unmounted row must have its old
// key-based measurement replaced now; otherwise virtual padding keeps
Expand Down Expand Up @@ -814,15 +828,18 @@ export function useFlowChatVirtualizer<T>({
measuredOriginRef.current = endOrigin;
// Stable DOM rows keep their chronological index while an origin handoff
// changes the library index. ResizeObserver must see the new mapping too.
const mountedRows = [...(scroller?.querySelectorAll<HTMLElement>('[data-virtual-index]') ?? [])]
.filter(element => element.closest('[data-flowchat-virtual-viewport]') === scroller);
for (const element of mountedRows) {
element.setAttribute('data-flowchat-measure-index', String(toVirtualIndex(Number(element.getAttribute('data-virtual-index')))));
const attached = originChanged
? [...(scroller?.querySelectorAll<HTMLElement>('[data-virtual-index]') ?? [])]
.filter(element => element.closest('[data-flowchat-virtual-viewport]') === scroller)
: [...attachingRowsRef.current].filter(element => element.isConnected);
if (originChanged) {
for (const element of attached) {
element.setAttribute('data-flowchat-measure-index', String(toVirtualIndex(Number(element.getAttribute('data-virtual-index')))));
}
}
const attached = originChanged ? mountedRows : [...attachingRowsRef.current].filter(element => element.isConnected);
attachingRowsRef.current.clear();
virtualizer.measureElement(null);
if (!attached.length || isViewportSuspendedRef.current()) return;
virtualizer.measureElement(null);
// Read before any viewport shift or size-cache notification can write.
measurementBatchRef.current = new Map(attached.map(element => [element, measureRowHeight(element)]));
try {
Expand Down
Loading
Loading