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
1089 }).size,1089 }).size,
1090 ).toBe(0);1090 ).toBe(0);
1091 });1091 });
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 });
1092});1142});
1093 1143
1094describe("deriveComposerSendState", () => {1144describe("deriveComposerSendState", () => {
apps/web/src/components/ChatView.logic.ts+19−11
@@ -38,6 +38,7 @@ import { · 37 unmodified lines
38} from "../types";38} from "../types";
39import { type ComposerImageAttachment, type DraftThreadState } from "../composerDraftStore";39import { type ComposerImageAttachment, type DraftThreadState } from "../composerDraftStore";
40import * as Schema from "effect/Schema";40import * as Schema from "effect/Schema";
41import { shallow } from "zustand/vanilla/shallow";
41import { appAtomRegistry } from "../rpc/atomRegistry";42import { appAtomRegistry } from "../rpc/atomRegistry";
42import { environmentThreadDetails } from "../state/threads";43import { environmentThreadDetails } from "../state/threads";
43import {44import {
@@ -464,17 +465,24 @@ export function getAntigravitySendBlockReason( · 420 unmodified lines
464 return null;465 return null;
465}466}
466 467
467export function buildRevertTurnCountByUserMessageId(input: {468/**
468 supportsConversationRollback: boolean;469 * Maps each user message to the checkpoint turn count a revert should target.
469 timelineEntries: ReadonlyArray<TimelineEntry>;470 * Returns `previous` when the result is unchanged: streaming text deltas
470 turnDiffSummaryByAssistantMessageId: ReadonlyMap<MessageId, TurnDiffSummary>;471 * rebuild `timelineEntries` per token, and the timeline row projection only
471 inferredCheckpointTurnCountByTurnId: Readonly<Record<string, number | undefined>>;472 * reuses rows while this Map keeps its identity.
472}) {473 */
474export 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> {
473 const byUserMessageId = new Map<MessageId, number>();483 const byUserMessageId = new Map<MessageId, number>();
474 if (!input.supportsConversationRollback) {484 const entryCount = input.supportsConversationRollback ? input.timelineEntries.length : 0;
475 return byUserMessageId;485 for (let index = 0; index < entryCount; index += 1) {
476 }
477 for (let index = 0; index < input.timelineEntries.length; index += 1) {
478 const entry = input.timelineEntries[index];486 const entry = input.timelineEntries[index];
479 if (!entry || entry.kind !== "message" || entry.message.role !== "user") {487 if (!entry || entry.kind !== "message" || entry.message.role !== "user") {
480 continue;488 continue;
@@ -501,7 +509,7 @@ export function buildRevertTurnCountByUserMessageId(input: { · 20 unmodified lines
501 break;509 break;
502 }510 }
503 }511 }
504 return byUserMessageId;512 return previous !== null && shallow(previous, byUserMessageId) ? previous : byUserMessageId;
505}513}
506 514
507export function reconcileMountedTerminalThreadIds(input: {515export function reconcileMountedTerminalThreadIds(input: {
apps/web/src/components/ChatView.tsx+15−11
@@ -2971,21 +2971,25 @@ export default function ChatView(props: ChatViewProps) { · 2970 unmodified lines
2971 }2971 }
2972 return byMessageId;2972 return byMessageId;
2973 }, [turnDiffSummaries]);2973 }, [turnDiffSummaries]);
2974 const revertTurnCountByUserMessageId = useMemo(2974 const lastRevertTurnCountRef = useRef<Map<MessageId, number> | null>(null);
2975 () =>2975 const revertTurnCountByUserMessageId = useMemo(() => {
2976 buildRevertTurnCountByUserMessageId({2976 const next = buildRevertTurnCountByUserMessageId(
2977 {
2977 supportsConversationRollback,2978 supportsConversationRollback,
2978 timelineEntries,2979 timelineEntries,
2979 turnDiffSummaryByAssistantMessageId,2980 turnDiffSummaryByAssistantMessageId,
2980 inferredCheckpointTurnCountByTurnId,2981 inferredCheckpointTurnCountByTurnId,
2981 }),2982 },
2982 [2983 lastRevertTurnCountRef.current,
2983 supportsConversationRollback,2984 );
2984 inferredCheckpointTurnCountByTurnId,2985 lastRevertTurnCountRef.current = next;
2985 timelineEntries,2986 return next;
2986 turnDiffSummaryByAssistantMessageId,2987 }, [
2987 ],2988 supportsConversationRollback,
2988 );2989 inferredCheckpointTurnCountByTurnId,
2990 timelineEntries,
2991 turnDiffSummaryByAssistantMessageId,
2992 ]);
2989 2993
2990 const gitCwd = activeProject2994 const gitCwd = activeProject
2991 ? projectScriptCwd({2995 ? projectScriptCwd({
apps/web/src/components/chat/MessagesTimeline.logic.test.ts+57−1
@@ -1,5 +1,5 @@
1import { describe, expect, it } from "vite-plus/test";1import { describe, expect, it } from "vite-plus/test";
2import { MessageId, TurnId } from "@t3tools/contracts";2import { CheckpointRef, MessageId, TurnId } from "@t3tools/contracts";
3import {3import {
4 computeStableMessagesTimelineRows,4 computeStableMessagesTimelineRows,
5 computeMessageDurationStart,5 computeMessageDurationStart,
@@ -20,6 +20,7 @@ import { · 14 unmodified lines
20 deriveTimelineEntriesWithState,20 deriveTimelineEntriesWithState,
21 type WorkLogEntry,21 type WorkLogEntry,
22} from "../../session-logic";22} from "../../session-logic";
23import { buildRevertTurnCountByUserMessageId } from "../ChatView.logic";
23import { isImageAttachment, type ChatMessage, type TurnDiffSummary } from "../../types";24import { isImageAttachment, type ChatMessage, type TurnDiffSummary } from "../../types";
24 25
25describe("streaming row projection", () => {26describe("streaming row projection", () => {
@@ -242,6 +243,61 @@ describe("streaming row projection", () => { · 216 unmodified lines
242 },243 },
243 );244 );
244 245
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
245 it.each(["completion", "turn", "role", "ordering"] as const)(301 it.each(["completion", "turn", "role", "ordering"] as const)(
246 "rebuilds row structure for a %s change with otherwise unchanged controls",302 "rebuilds row structure for a %s change with otherwise unchanged controls",
247 (change) => {303 (change) => {