History

perf(web): keep timeline row reuse engaged while text streams (#9909)

Julius Marminge committed 2 months ago

Commit aefef951 parent ea0487c4 files · +141 −23Browse files

Unified · Split

4 files
  1. apps/web/src/components/ChatView.logic.test.ts +50 −0
  2. apps/web/src/components/ChatView.logic.ts +19 −11
  3. apps/web/src/components/ChatView.tsx +15 −11
  4. apps/web/src/components/chat/MessagesTimeline.logic.test.ts +57 −1
apps/web/src/components/ChatView.logic.test.ts+50−0
@@ -1089,6 +1089,56 @@ describe("buildRevertTurnCountByUserMessageId", () => { · 1088 unmodified lines
10891089 }).size,
10901090 ).toBe(0);
10911091 });
1092+
1093+ it.each([true, false])(
1094+ "returns the previous map when contents are unchanged (rollback supported: %s)",
1095+ (supportsConversationRollback) => {
1096+ const input = {
1097+ supportsConversationRollback,
1098+ timelineEntries,
1099+ turnDiffSummaryByAssistantMessageId,
1100+ inferredCheckpointTurnCountByTurnId: {},
1101+ };
1102+ const previous = buildRevertTurnCountByUserMessageId(input);
1103+ const streamed = timelineEntries.map((entry) =>
1104+ entry.message.role === "assistant"
1105+ ? { ...entry, message: { ...entry.message, text: "Updated the file again" } }
1106+ : entry,
1107+ );
1108+
1109+ expect(
1110+ buildRevertTurnCountByUserMessageId({ ...input, timelineEntries: streamed }, previous),
1111+ ).toBe(previous);
1112+ },
1113+ );
1114+
1115+ it("returns a new map when a revert target changes", () => {
1116+ const input = {
1117+ supportsConversationRollback: true,
1118+ timelineEntries,
1119+ turnDiffSummaryByAssistantMessageId,
1120+ inferredCheckpointTurnCountByTurnId: {},
1121+ };
1122+ const previous = buildRevertTurnCountByUserMessageId(input);
1123+ const next = buildRevertTurnCountByUserMessageId(
1124+ {
1125+ ...input,
1126+ turnDiffSummaryByAssistantMessageId: new Map([
1127+ [
1128+ assistantMessageId,
1129+ {
1130+ ...turnDiffSummaryByAssistantMessageId.get(assistantMessageId)!,
1131+ checkpointTurnCount: 3,
1132+ },
1133+ ],
1134+ ]),
1135+ },
1136+ previous,
1137+ );
1138+
1139+ expect(next).not.toBe(previous);
1140+ expect(next).toEqual(new Map([[userMessageId, 2]]));
1141+ });
10921142 });
10931143
10941144 describe("deriveComposerSendState", () => {
apps/web/src/components/ChatView.logic.ts+19−11
@@ -38,6 +38,7 @@ import { · 37 unmodified lines
3838 } from "../types";
3939 import { type ComposerImageAttachment, type DraftThreadState } from "../composerDraftStore";
4040 import * as Schema from "effect/Schema";
41+import { shallow } from "zustand/vanilla/shallow";
4142 import { appAtomRegistry } from "../rpc/atomRegistry";
4243 import { environmentThreadDetails } from "../state/threads";
4344 import {
@@ -464,17 +465,24 @@ export function getAntigravitySendBlockReason( · 420 unmodified lines
464465 return null;
465466 }
466467
467−export function buildRevertTurnCountByUserMessageId(input: {
468− supportsConversationRollback: boolean;
469− timelineEntries: ReadonlyArray<TimelineEntry>;
470− turnDiffSummaryByAssistantMessageId: ReadonlyMap<MessageId, TurnDiffSummary>;
471− inferredCheckpointTurnCountByTurnId: Readonly<Record<string, number | undefined>>;
472−}) {
468+/**
469+ * Maps each user message to the checkpoint turn count a revert should target.
470+ * Returns `previous` when the result is unchanged: streaming text deltas
471+ * rebuild `timelineEntries` per token, and the timeline row projection only
472+ * reuses rows while this Map keeps its identity.
473+ */
474+export function buildRevertTurnCountByUserMessageId(
475+ input: {
476+ supportsConversationRollback: boolean;
477+ timelineEntries: ReadonlyArray<TimelineEntry>;
478+ turnDiffSummaryByAssistantMessageId: ReadonlyMap<MessageId, TurnDiffSummary>;
479+ inferredCheckpointTurnCountByTurnId: Readonly<Record<string, number | undefined>>;
480+ },
481+ previous: Map<MessageId, number> | null = null,
482+): Map<MessageId, number> {
473483 const byUserMessageId = new Map<MessageId, number>();
474− if (!input.supportsConversationRollback) {
475− return byUserMessageId;
476− }
477− for (let index = 0; index < input.timelineEntries.length; index += 1) {
484+ const entryCount = input.supportsConversationRollback ? input.timelineEntries.length : 0;
485+ for (let index = 0; index < entryCount; index += 1) {
478486 const entry = input.timelineEntries[index];
479487 if (!entry || entry.kind !== "message" || entry.message.role !== "user") {
480488 continue;
@@ -501,7 +509,7 @@ export function buildRevertTurnCountByUserMessageId(input: { · 20 unmodified lines
501509 break;
502510 }
503511 }
504− return byUserMessageId;
512+ return previous !== null && shallow(previous, byUserMessageId) ? previous : byUserMessageId;
505513 }
506514
507515 export function reconcileMountedTerminalThreadIds(input: {
apps/web/src/components/ChatView.tsx+15−11
@@ -2971,21 +2971,25 @@ export default function ChatView(props: ChatViewProps) { · 2970 unmodified lines
29712971 }
29722972 return byMessageId;
29732973 }, [turnDiffSummaries]);
2974− const revertTurnCountByUserMessageId = useMemo(
2975− () =>
2976− buildRevertTurnCountByUserMessageId({
2974+ const lastRevertTurnCountRef = useRef<Map<MessageId, number> | null>(null);
2975+ const revertTurnCountByUserMessageId = useMemo(() => {
2976+ const next = buildRevertTurnCountByUserMessageId(
2977+ {
29772978 supportsConversationRollback,
29782979 timelineEntries,
29792980 turnDiffSummaryByAssistantMessageId,
29802981 inferredCheckpointTurnCountByTurnId,
2981− }),
2982− [
2983− supportsConversationRollback,
2984− inferredCheckpointTurnCountByTurnId,
2985− timelineEntries,
2986− turnDiffSummaryByAssistantMessageId,
2987− ],
2988− );
2982+ },
2983+ lastRevertTurnCountRef.current,
2984+ );
2985+ lastRevertTurnCountRef.current = next;
2986+ return next;
2987+ }, [
2988+ supportsConversationRollback,
2989+ inferredCheckpointTurnCountByTurnId,
2990+ timelineEntries,
2991+ turnDiffSummaryByAssistantMessageId,
2992+ ]);
29892993
29902994 const gitCwd = activeProject
29912995 ? projectScriptCwd({
apps/web/src/components/chat/MessagesTimeline.logic.test.ts+57−1
@@ -1,5 +1,5 @@
11 import { describe, expect, it } from "vite-plus/test";
2−import { MessageId, TurnId } from "@t3tools/contracts";
2+import { CheckpointRef, MessageId, TurnId } from "@t3tools/contracts";
33 import {
44 computeStableMessagesTimelineRows,
55 computeMessageDurationStart,
@@ -20,6 +20,7 @@ import { · 14 unmodified lines
2020 deriveTimelineEntriesWithState,
2121 type WorkLogEntry,
2222 } from "../../session-logic";
23+import { buildRevertTurnCountByUserMessageId } from "../ChatView.logic";
2324 import { isImageAttachment, type ChatMessage, type TurnDiffSummary } from "../../types";
2425
2526 describe("streaming row projection", () => {
@@ -242,6 +243,61 @@ describe("streaming row projection", () => { · 216 unmodified lines
242243 },
243244 );
244245
246+ it("reuses rows when the revert map is rebuilt from the streamed entries", () => {
247+ const initial = fixture("Partial");
248+ const inferredCheckpointTurnCountByTurnId = { [initial.historyTurnId]: 1 };
249+ const turnDiffSummaryByAssistantMessageId = new Map<MessageId, TurnDiffSummary>([
250+ [
251+ MessageId.make("history-assistant"),
252+ {
253+ turnId: initial.historyTurnId,
254+ checkpointTurnCount: 1,
255+ checkpointRef: CheckpointRef.make("refs/t3/checkpoints/history-turn"),
256+ status: "ready",
257+ files: [],
258+ assistantMessageId: MessageId.make("history-assistant"),
259+ completedAt: initial.time(4),
260+ },
261+ ],
262+ ]);
263+ let revertMap: Map<MessageId, number> | null = null;
264+ // Mirrors ChatView: the map is derived from each delta's entries.
265+ const build = (timelineEntries: typeof initial.timeline.entries) => {
266+ revertMap = buildRevertTurnCountByUserMessageId(
267+ {
268+ supportsConversationRollback: true,
269+ timelineEntries,
270+ turnDiffSummaryByAssistantMessageId,
271+ inferredCheckpointTurnCountByTurnId,
272+ },
273+ revertMap,
274+ );
275+ return {
276+ ...initial.input,
277+ timelineEntries,
278+ turnDiffSummaryByAssistantMessageId,
279+ revertTurnCountByUserMessageId: revertMap,
280+ };
281+ };
282+ const previous = deriveMessagesTimelineRowsWithState(build(initial.timeline.entries));
283+ expect(previous.rows.some((row) => row.kind === "message" && row.revertTurnCount === 0)).toBe(
284+ true,
285+ );
286+ const last = initial.messages.at(-1)!;
287+ const messages = [...initial.messages.slice(0, -1), { ...last, text: "Partial token" }];
288+ const timeline = deriveTimelineEntriesWithState(messages, [], initial.work, initial.timeline);
289+ const next = deriveMessagesTimelineRowsWithState(build(timeline.entries), previous);
290+
291+ expect(next.rows).toEqual(deriveMessagesTimelineRows(build(timeline.entries)));
292+ for (const [index, row] of previous.rows.entries()) {
293+ if ((row.kind === "message" || row.kind === "assistant-meta") && row.message === last) {
294+ expect(next.rows[index]).toMatchObject({ message: { text: "Partial token" } });
295+ } else {
296+ expect(next.rows[index]).toBe(row);
297+ }
298+ }
299+ });
300+
245301 it.each(["completion", "turn", "role", "ordering"] as const)(
246302 "rebuilds row structure for a %s change with otherwise unchanged controls",
247303 (change) => {