From 6d31e9cfe21d2a8c4490b7a6e60e8cef7c8a6410 Mon Sep 17 00:00:00 2001 From: Theo Browne Date: Sat, 5 Sep 2026 06:20:03 -0700 Subject: [PATCH 1/3] perf(web): skip checkpoint map rebuilds while streaming --- .../web/src/components/ChatView.logic.test.ts | 125 ----- apps/web/src/components/ChatView.logic.ts | 48 -- apps/web/src/components/ChatView.tsx | 73 +-- .../chat/MessagesTimeline.logic.test.ts | 431 +++++++++++++----- .../components/chat/MessagesTimeline.logic.ts | 89 +++- .../components/chat/MessagesTimeline.test.tsx | 188 +++++++- .../src/components/chat/MessagesTimeline.tsx | 36 +- 7 files changed, 612 insertions(+), 378 deletions(-) diff --git a/apps/web/src/components/ChatView.logic.test.ts b/apps/web/src/components/ChatView.logic.test.ts index 47173520087a..820d431db4f0 100644 --- a/apps/web/src/components/ChatView.logic.test.ts +++ b/apps/web/src/components/ChatView.logic.test.ts @@ -1,6 +1,5 @@ import { ANTIGRAVITY_DEFAULT_MODEL, - CheckpointRef, EnvironmentId, MessageId, ProjectId, @@ -13,7 +12,6 @@ import { import { afterEach, beforeEach, describe, expect, it, vi } from "vite-plus/test"; import type { Thread, ThreadShell, TurnDiffSummary } from "../types"; -import type { TimelineEntry } from "../session-logic"; import { deriveProviderInstanceEntries, NO_PROVIDER_MODEL_SELECTION } from "../providerInstances"; import type { CodexArtifactTemplate } from "@t3tools/client-runtime/codex-artifact-templates"; import type { RightPanelSurface } from "../rightPanelStore"; @@ -24,7 +22,6 @@ import { branchMismatchKey, buildExpiredTerminalContextToastCopy, buildLoadingThreadFromShell, - buildRevertTurnCountByUserMessageId, buildThreadTurnInterruptInput, createLocalDispatchSnapshot, deriveComposerSendState, @@ -1127,128 +1124,6 @@ describe("resolveComposerInteractionMode", () => { }); }); -describe("buildRevertTurnCountByUserMessageId", () => { - const userMessageId = MessageId.make("rewind-user-message"); - const assistantMessageId = MessageId.make("rewind-assistant-message"); - const turnId = TurnId.make("rewind-turn"); - const timelineEntries = [ - { - id: userMessageId, - kind: "message", - createdAt: now, - message: { - id: userMessageId, - role: "user", - text: "Update the file", - turnId, - createdAt: now, - updatedAt: now, - streaming: false, - }, - }, - { - id: assistantMessageId, - kind: "message", - createdAt: now, - message: { - id: assistantMessageId, - role: "assistant", - text: "Updated the file", - turnId, - createdAt: now, - updatedAt: now, - streaming: false, - }, - }, - ] satisfies ReadonlyArray; - const turnDiffSummaryByAssistantMessageId = new Map([ - [ - assistantMessageId, - { - turnId, - checkpointTurnCount: 1, - checkpointRef: CheckpointRef.make("refs/t3/checkpoints/rewind-turn"), - status: "ready", - files: [], - assistantMessageId, - completedAt: now, - }, - ], - ]); - - it("offers the checkpoint before the user message when conversation rollback is supported", () => { - expect( - buildRevertTurnCountByUserMessageId({ - supportsConversationRollback: true, - timelineEntries, - turnDiffSummaryByAssistantMessageId, - inferredCheckpointTurnCountByTurnId: {}, - }), - ).toEqual(new Map([[userMessageId, 0]])); - }); - - it("offers no rewind action when file checkpoints exist but conversation rollback is unsupported", () => { - expect( - buildRevertTurnCountByUserMessageId({ - supportsConversationRollback: false, - timelineEntries, - turnDiffSummaryByAssistantMessageId, - inferredCheckpointTurnCountByTurnId: {}, - }).size, - ).toBe(0); - }); - - it.each([true, false])( - "returns the previous map when contents are unchanged (rollback supported: %s)", - (supportsConversationRollback) => { - const input = { - supportsConversationRollback, - timelineEntries, - turnDiffSummaryByAssistantMessageId, - inferredCheckpointTurnCountByTurnId: {}, - }; - const previous = buildRevertTurnCountByUserMessageId(input); - const streamed = timelineEntries.map((entry) => - entry.message.role === "assistant" - ? { ...entry, message: { ...entry.message, text: "Updated the file again" } } - : entry, - ); - - expect( - buildRevertTurnCountByUserMessageId({ ...input, timelineEntries: streamed }, previous), - ).toBe(previous); - }, - ); - - it("returns a new map when a revert target changes", () => { - const input = { - supportsConversationRollback: true, - timelineEntries, - turnDiffSummaryByAssistantMessageId, - inferredCheckpointTurnCountByTurnId: {}, - }; - const previous = buildRevertTurnCountByUserMessageId(input); - const next = buildRevertTurnCountByUserMessageId( - { - ...input, - turnDiffSummaryByAssistantMessageId: new Map([ - [ - assistantMessageId, - { - ...turnDiffSummaryByAssistantMessageId.get(assistantMessageId)!, - checkpointTurnCount: 3, - }, - ], - ]), - }, - previous, - ); - - expect(next).not.toBe(previous); - expect(next).toEqual(new Map([[userMessageId, 2]])); - }); -}); - describe("deriveComposerSendState", () => { it("treats expired terminal pills as non-sendable content", () => { const state = deriveComposerSendState({ diff --git a/apps/web/src/components/ChatView.logic.ts b/apps/web/src/components/ChatView.logic.ts index ff0b15071955..1faeb74c9863 100644 --- a/apps/web/src/components/ChatView.logic.ts +++ b/apps/web/src/components/ChatView.logic.ts @@ -38,7 +38,6 @@ import { } from "../types"; import { type ComposerImageAttachment, type DraftThreadState } from "../composerDraftStore"; import * as Schema from "effect/Schema"; -import { shallow } from "zustand/vanilla/shallow"; import { appAtomRegistry } from "../rpc/atomRegistry"; import { environmentThreadDetails } from "../state/threads"; import { @@ -465,53 +464,6 @@ export function getAntigravitySendBlockReason( return null; } -/** - * Maps each user message to the checkpoint turn count a revert should target. - * Returns `previous` when the result is unchanged: streaming text deltas - * rebuild `timelineEntries` per token, and the timeline row projection only - * reuses rows while this Map keeps its identity. - */ -export function buildRevertTurnCountByUserMessageId( - input: { - supportsConversationRollback: boolean; - timelineEntries: ReadonlyArray; - turnDiffSummaryByAssistantMessageId: ReadonlyMap; - inferredCheckpointTurnCountByTurnId: Readonly>; - }, - previous: Map | null = null, -): Map { - const byUserMessageId = new Map(); - const entryCount = input.supportsConversationRollback ? input.timelineEntries.length : 0; - for (let index = 0; index < entryCount; index += 1) { - const entry = input.timelineEntries[index]; - if (!entry || entry.kind !== "message" || entry.message.role !== "user") { - continue; - } - - for (let nextIndex = index + 1; nextIndex < input.timelineEntries.length; nextIndex += 1) { - const nextEntry = input.timelineEntries[nextIndex]; - if (!nextEntry || nextEntry.kind !== "message") { - continue; - } - if (nextEntry.message.role === "user") { - break; - } - const summary = input.turnDiffSummaryByAssistantMessageId.get(nextEntry.message.id); - if (!summary) { - continue; - } - const turnCount = - summary.checkpointTurnCount ?? input.inferredCheckpointTurnCountByTurnId[summary.turnId]; - if (typeof turnCount !== "number") { - break; - } - byUserMessageId.set(entry.message.id, Math.max(0, turnCount - 1)); - break; - } - } - return previous !== null && shallow(previous, byUserMessageId) ? previous : byUserMessageId; -} - export function reconcileMountedTerminalThreadIds(input: { currentThreadIds: ReadonlyArray; openThreadIds: ReadonlyArray; diff --git a/apps/web/src/components/ChatView.tsx b/apps/web/src/components/ChatView.tsx index c0c79c93bec1..db841e53ebfe 100644 --- a/apps/web/src/components/ChatView.tsx +++ b/apps/web/src/components/ChatView.tsx @@ -153,11 +153,9 @@ import { isImageAttachment, type SessionPhase, type Thread, - type TurnDiffSummary, } from "../types"; import { useTheme } from "../hooks/useTheme"; import { writeTextToClipboard } from "../hooks/useCopyToClipboard"; -import { useTurnDiffSummaries } from "../hooks/useTurnDiffSummaries"; import { isCommandPaletteOpen } from "../commandPaletteBus"; import { buildTemporaryWorktreeBranchName } from "@t3tools/shared/git"; import { useMediaQuery } from "../hooks/useMediaQuery"; @@ -367,7 +365,6 @@ import { buildExpiredTerminalContextToastCopy, buildLocalDraftThread, buildLoadingThreadFromShell, - buildRevertTurnCountByUserMessageId, buildThreadTurnInterruptInput, collectUserMessageBlobPreviewUrls, createLocalDispatchSnapshot, @@ -3101,35 +3098,6 @@ export default function ChatView(props: ChatViewProps) { attachDraftHeroComposerAnchorRef, captureDraftHeroComposerRect, ] = useDraftHeroLayoutTransition(isDraftHeroState); - const { turnDiffSummaries, inferredCheckpointTurnCountByTurnId } = - useTurnDiffSummaries(activeThread); - const turnDiffSummaryByAssistantMessageId = useMemo(() => { - const byMessageId = new Map(); - for (const summary of turnDiffSummaries) { - if (!summary.assistantMessageId) continue; - byMessageId.set(summary.assistantMessageId, summary); - } - return byMessageId; - }, [turnDiffSummaries]); - const lastRevertTurnCountRef = useRef | null>(null); - const revertTurnCountByUserMessageId = useMemo(() => { - const next = buildRevertTurnCountByUserMessageId( - { - supportsConversationRollback, - timelineEntries, - turnDiffSummaryByAssistantMessageId, - inferredCheckpointTurnCountByTurnId, - }, - lastRevertTurnCountRef.current, - ); - lastRevertTurnCountRef.current = next; - return next; - }, [ - supportsConversationRollback, - inferredCheckpointTurnCountByTurnId, - timelineEntries, - turnDiffSummaryByAssistantMessageId, - ]); const gitCwd = activeProject ? projectScriptCwd({ @@ -7778,19 +7746,28 @@ export default function ChatView(props: ChatViewProps) { }, [activeThreadRef, isServerThread, onDiffPanelOpen], ); - // Both the Map and the revert handler are read from refs at call-time so - // the callback reference is fully stable and never busts context identity. - const revertTurnCountRef = useRef(revertTurnCountByUserMessageId); - revertTurnCountRef.current = revertTurnCountByUserMessageId; - const onRevertToTurnCountRef = useRef(onRevertToTurnCount); - onRevertToTurnCountRef.current = onRevertToTurnCount; - const onRevertUserMessage = useCallback((messageId: MessageId) => { - const targetTurnCount = revertTurnCountRef.current.get(messageId); - if (typeof targetTurnCount !== "number") { - return; - } - void onRevertToTurnCountRef.current(targetTurnCount); - }, []); + const revertCheckpoints = activeThread?.checkpoints; + const revertHandlerRef = useRef({ + threadKey: activeThreadKey, + checkpoints: revertCheckpoints, + handler: onRevertToTurnCount, + }); + revertHandlerRef.current = { + threadKey: activeThreadKey, + checkpoints: revertCheckpoints, + handler: onRevertToTurnCount, + }; + // A new thread or checkpoint snapshot must not retarget an old row's click. + const onRevertTimelineTurn = useCallback( + (targetTurnCount: number) => { + const current = revertHandlerRef.current; + if (current.threadKey !== activeThreadKey || current.checkpoints !== revertCheckpoints) { + return; + } + void current.handler(targetTurnCount); + }, + [activeThreadKey, revertCheckpoints], + ); // Empty state: no active thread if (!activeThread) { @@ -8104,12 +8081,12 @@ export default function ChatView(props: ChatViewProps) { timelineEntries={timelineEntries} latestTurn={activeLatestTurn} runningTurnId={activeRunningTurnId} - turnDiffSummaryByAssistantMessageId={turnDiffSummaryByAssistantMessageId} + turnDiffSummaries={activeThread.checkpoints} activeThreadEnvironmentId={activeThread.environmentId} routeThreadKey={routeThreadKey} onOpenTurnDiff={onOpenTurnDiff} - revertTurnCountByUserMessageId={revertTurnCountByUserMessageId} - onRevertUserMessage={onRevertUserMessage} + supportsConversationRollback={supportsConversationRollback} + onRevertToTurnCount={onRevertTimelineTurn} onUseArtifactTemplate={useArtifactTemplate} isRevertingCheckpoint={isRevertingCheckpoint} onImageExpand={onExpandTimelineImage} diff --git a/apps/web/src/components/chat/MessagesTimeline.logic.test.ts b/apps/web/src/components/chat/MessagesTimeline.logic.test.ts index 5d310c5fe345..b761348bb0a4 100644 --- a/apps/web/src/components/chat/MessagesTimeline.logic.test.ts +++ b/apps/web/src/components/chat/MessagesTimeline.logic.test.ts @@ -1,5 +1,22 @@ import { describe, expect, it } from "vite-plus/test"; -import { CheckpointRef, MessageId, TurnId } from "@t3tools/contracts"; +import { + CheckpointRef, + EnvironmentId, + EventId, + MessageId, + ProjectId, + ProviderInstanceId, + ThreadId, + TurnId, + type OrchestrationThread, +} from "@t3tools/contracts"; +import { + applyThreadDetailEvent, + createEnvironmentThreadDetailAtoms, + EMPTY_ENVIRONMENT_THREAD_STATE, +} from "@t3tools/client-runtime/state/threads"; +import * as Option from "effect/Option"; +import { AsyncResult, Atom, AtomRegistry } from "effect/unstable/reactivity"; import { computeStableMessagesTimelineRows, computeMessageDurationStart, @@ -12,6 +29,7 @@ import { shouldFollowWorkGroupAppend, shouldPreserveAssistantLineBreaks, type MessagesTimelineRow, + type MessagesTimelineRowsProjection, workEntryDisplayLabel, } from "./MessagesTimeline.logic"; import { @@ -19,8 +37,8 @@ import { deriveTimelineEntries, deriveTimelineEntriesWithState, type WorkLogEntry, + type TimelineEntriesProjection, } from "../../session-logic"; -import { buildRevertTurnCountByUserMessageId } from "../ChatView.logic"; import { isImageAttachment, type ChatMessage, type TurnDiffSummary } from "../../types"; describe("streaming row projection", () => { @@ -97,8 +115,8 @@ describe("streaming row projection", () => { runningTurnId: turnId, isWorking: true, activeTurnStartedAt: time(5), - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, } satisfies Parameters[0]; return { messages, work, timeline, input, time, turnId, historyTurnId }; } @@ -243,52 +261,48 @@ describe("streaming row projection", () => { }, ); - it("reuses rows when the revert map is rebuilt from the streamed entries", () => { + it("owns checkpoint lookups across streaming and equal source snapshots", () => { const initial = fixture("Partial"); - const inferredCheckpointTurnCountByTurnId = { [initial.historyTurnId]: 1 }; - const turnDiffSummaryByAssistantMessageId = new Map([ - [ - MessageId.make("history-assistant"), - { - turnId: initial.historyTurnId, - checkpointTurnCount: 1, - checkpointRef: CheckpointRef.make("refs/t3/checkpoints/history-turn"), - status: "ready", - files: [], - assistantMessageId: MessageId.make("history-assistant"), - completedAt: initial.time(4), - }, - ], - ]); - let revertMap: Map | null = null; - // Mirrors ChatView: the map is derived from each delta's entries. - const build = (timelineEntries: typeof initial.timeline.entries) => { - revertMap = buildRevertTurnCountByUserMessageId( - { - supportsConversationRollback: true, - timelineEntries, - turnDiffSummaryByAssistantMessageId, - inferredCheckpointTurnCountByTurnId, - }, - revertMap, - ); - return { - ...initial.input, - timelineEntries, - turnDiffSummaryByAssistantMessageId, - revertTurnCountByUserMessageId: revertMap, - }; + let checkpointReads = 0; + const summary: TurnDiffSummary = { + turnId: initial.historyTurnId, + checkpointTurnCount: 1, + checkpointRef: CheckpointRef.make("refs/t3/checkpoints/history-turn"), + status: "ready", + files: [], + get assistantMessageId() { + checkpointReads += 1; + return MessageId.make("history-assistant"); + }, + completedAt: initial.time(4), + }; + const input = { + ...initial.input, + turnDiffSummaries: [summary], + supportsConversationRollback: true, + expandedTurnIds: new Set([initial.historyTurnId]), + expandedWorkGroupIds: new Set(), }; - const previous = deriveMessagesTimelineRowsWithState(build(initial.timeline.entries)); + const previous = deriveMessagesTimelineRowsWithState(input); + expect(checkpointReads).toBeGreaterThan(0); expect(previous.rows.some((row) => row.kind === "message" && row.revertTurnCount === 0)).toBe( true, ); const last = initial.messages.at(-1)!; const messages = [...initial.messages.slice(0, -1), { ...last, text: "Partial token" }]; const timeline = deriveTimelineEntriesWithState(messages, [], initial.work, initial.timeline); - const next = deriveMessagesTimelineRowsWithState(build(timeline.entries), previous); - - expect(next.rows).toEqual(deriveMessagesTimelineRows(build(timeline.entries))); + const nextInput = { + ...input, + timelineEntries: timeline.entries, + turnDiffSummaries: [...input.turnDiffSummaries], + latestTurn: { ...input.latestTurn }, + expandedTurnIds: new Set(input.expandedTurnIds), + expandedWorkGroupIds: new Set(input.expandedWorkGroupIds), + }; + checkpointReads = 0; + const next = deriveMessagesTimelineRowsWithState(nextInput, previous); + expect(checkpointReads).toBe(0); + expect(next.rows).toEqual(deriveMessagesTimelineRows(nextInput)); for (const [index, row] of previous.rows.entries()) { if ((row.kind === "message" || row.kind === "assistant-meta") && row.message === last) { expect(next.rows[index]).toMatchObject({ message: { text: "Partial token" } }); @@ -296,6 +310,199 @@ describe("streaming row projection", () => { expect(next.rows[index]).toBe(row); } } + + const changed = deriveMessagesTimelineRowsWithState( + { ...nextInput, turnDiffSummaries: [{ ...summary, checkpointTurnCount: 3 }] }, + next, + ); + expect( + changed.rows.find((row) => row.kind === "message" && row.message.id === messages[0]?.id), + ).toMatchObject({ revertTurnCount: 2 }); + const unsupported = deriveMessagesTimelineRowsWithState( + { ...changed.input, supportsConversationRollback: false }, + changed, + ); + expect( + unsupported.rows.find((row) => row.kind === "message" && row.message.id === messages[0]?.id), + ).toMatchObject({ revertTurnCount: undefined }); + expect( + previous.rows.find((row) => row.kind === "message" && row.message.id === messages[0]?.id), + ).toMatchObject({ revertTurnCount: 0 }); + }); + + it("reuses long-thread rows through detail events, selectors, and attachment previews", () => { + const initial = fixture("Partial"); + let checkpointReads = 0; + const history = Array.from({ length: 250 }, (_, index) => { + const turnId = TurnId.make(`older-turn-${index}`); + const user = { + ...initial.messages[0]!, + id: MessageId.make(`older-user-${index}`), + createdAt: new Date(Date.UTC(2026, 8, 3, 0, 0, index * 5)).toISOString(), + attachments: [ + { + type: "image" as const, + id: "image", + name: "image.png", + mimeType: "image/png", + sizeBytes: 42, + }, + ], + }; + const assistant = { + ...initial.messages[1]!, + id: MessageId.make(`older-assistant-${index}`), + turnId, + createdAt: new Date(Date.UTC(2026, 8, 3, 0, 0, index * 5 + 3)).toISOString(), + }; + const checkpoint: TurnDiffSummary = { + turnId, + checkpointTurnCount: index + 1, + checkpointRef: CheckpointRef.make(`refs/t3/checkpoints/older-${index}`), + status: "ready", + files: [], + get assistantMessageId() { + checkpointReads += 1; + return assistant.id; + }, + completedAt: assistant.createdAt, + }; + return { user, assistant, checkpoint }; + }); + const liveMessage = initial.messages.at(-1)!; + let thread: OrchestrationThread = { + id: ThreadId.make("streaming-thread"), + projectId: ProjectId.make("project"), + title: "Long thread", + modelSelection: { instanceId: ProviderInstanceId.make("codex"), model: "gpt-5.4" }, + runtimeMode: "full-access", + interactionMode: "default", + branch: null, + worktreePath: null, + latestTurn: { + ...initial.input.latestTurn, + requestedAt: initial.time(5), + assistantMessageId: liveMessage.id, + }, + createdAt: initial.time(0), + updatedAt: initial.time(7), + archivedAt: null, + settledOverride: null, + settledAt: null, + deletedAt: null, + messages: [ + ...history.flatMap(({ user, assistant }) => [user, assistant]), + ...initial.messages, + ], + proposedPlans: [], + activities: [], + checkpoints: history.map(({ checkpoint }) => checkpoint), + session: null, + }; + const state = Atom.make( + AsyncResult.success({ ...EMPTY_ENVIRONMENT_THREAD_STATE, data: Option.some(thread) }), + ); + const details = createEnvironmentThreadDetailAtoms(() => state); + const ref = { environmentId: EnvironmentId.make("local"), threadId: thread.id }; + const registry = AtomRegistry.make(); + const unmount = registry.mount(details.detailAtom(ref)); + const preview = createMessageAttachmentPreviewProjector(); + let imageUrl = "https://first.test/image"; + let timeline: TimelineEntriesProjection | null = null; + let projection: MessagesTimelineRowsProjection | null = null; + const project = () => { + const selected = registry.get(details.detailAtom(ref)); + if (selected === null) throw new Error("Missing thread detail"); + const messages = selected.messages.map((message) => preview(message, () => imageUrl)); + timeline = deriveTimelineEntriesWithState( + messages, + selected.proposedPlans, + initial.work, + timeline, + ); + projection = deriveMessagesTimelineRowsWithState( + { + timelineEntries: timeline.entries, + latestTurn: selected.latestTurn, + runningTurnId: + selected.latestTurn?.state === "running" ? selected.latestTurn.turnId : null, + isWorking: selected.latestTurn?.state === "running", + activeTurnStartedAt: selected.latestTurn?.startedAt ?? null, + turnDiffSummaries: selected.checkpoints, + supportsConversationRollback: true, + }, + projection, + ); + return projection; + }; + const send = (text: string, sequence: number, streaming = true) => { + const result = applyThreadDetailEvent(thread, { + eventId: EventId.make(`delta-${sequence}`), + sequence, + commandId: null, + causationEventId: null, + correlationId: null, + metadata: {}, + occurredAt: initial.time(8 + sequence), + aggregateKind: "thread", + aggregateId: thread.id, + type: "thread.message-sent", + payload: { + threadId: thread.id, + messageId: liveMessage.id, + role: "assistant", + text, + turnId: initial.turnId, + streaming, + createdAt: liveMessage.createdAt, + updatedAt: initial.time(8 + sequence), + }, + }); + if (result.kind !== "updated") throw new Error("Message event did not update the thread"); + thread = result.thread; + registry.set( + state, + AsyncResult.success({ ...EMPTY_ENVIRONMENT_THREAD_STATE, data: Option.some(thread) }), + ); + return project(); + }; + + try { + const first = project(); + const saved = structuredClone(first.rows); + expect(checkpointReads).toBeGreaterThan(0); + checkpointReads = 0; + for (let index = 0; index < 10; index += 1) { + const next = send(` ${index}`, index + 1); + for (const [rowIndex, row] of first.rows.entries()) { + if ( + (row.kind === "message" || row.kind === "assistant-meta") && + row.message.id === liveMessage.id + ) + continue; + expect(next.rows[rowIndex]).toBe(row); + } + } + expect(checkpointReads).toBe(0); + const streamed = project(); + expect(streamed.rows).toEqual(deriveMessagesTimelineRows(streamed.input)); + + imageUrl = "https://renewed.test/image"; + const renewed = project(); + expect(renewed.rows[0]).not.toBe(first.rows[0]); + expect(renewed.rows[0]).toMatchObject({ + message: { attachments: [{ previewUrl: imageUrl }] }, + }); + const completed = send("Complete", 11, false); + expect(completed.rows).toEqual(deriveMessagesTimelineRows(completed.input)); + expect( + completed.rows.find((row) => row.kind === "message" && row.message.id === liveMessage.id), + ).toMatchObject({ message: { text: "Complete" }, assistantCopyStreaming: false }); + expect(first.rows).toEqual(saved); + } finally { + unmount(); + registry.dispose(); + } }); it.each(["completion", "turn", "role", "ordering"] as const)( @@ -407,7 +614,7 @@ describe("streaming row projection", () => { ? { ...message, role: "user", turnId: null, createdAt: initial.time(0) } : message, ); - check({ revertTurnCountByUserMessageId: new Map([[MessageId.make("live-user"), 3]]) }); + check({ supportsConversationRollback: true }); }); }); @@ -574,8 +781,8 @@ describe("work entry labels", () => { ], isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const directRow = rows.find((row) => row.kind === "work"); expect(directRow).toMatchObject({ @@ -887,8 +1094,8 @@ describe("deriveMessagesTimelineRows", () => { ], isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows).toEqual([ @@ -950,8 +1157,8 @@ describe("deriveMessagesTimelineRows", () => { expandedTurnIds: new Set(["turn-1" as never]), isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const assistantRows = rows.filter( @@ -1004,8 +1211,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const assistantRows = rows.filter( @@ -1061,10 +1268,8 @@ describe("deriveMessagesTimelineRows", () => { ], isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map([ - ["assistant-1" as never, assistantTurnDiffSummary], - ]), - revertTurnCountByUserMessageId: new Map([["user-1" as never, 1]]), + turnDiffSummaries: [assistantTurnDiffSummary], + supportsConversationRollback: true, }); const userRow = rows.find( @@ -1142,8 +1347,8 @@ describe("deriveMessagesTimelineRows", () => { timelineEntries, isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const foldRow = collapsedRows.find( @@ -1165,8 +1370,8 @@ describe("deriveMessagesTimelineRows", () => { expandedTurnIds: new Set(["turn-1" as never]), isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(expandedRows.map((row) => row.id)).toEqual([ @@ -1235,8 +1440,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }; const rows = deriveMessagesTimelineRows({ ...input, timelineEntries }); @@ -1318,8 +1523,8 @@ describe("deriveMessagesTimelineRows", () => { timelineEntries, isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.map((row) => row.id)).toEqual(["turn-fold:turn-1", "assistant-final-entry"]); @@ -1421,8 +1626,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:00:14Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const foldRow = rows.find( @@ -1458,8 +1663,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows).toEqual([ @@ -1526,8 +1731,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:01:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.map((row) => row.id)).toEqual([ @@ -1580,8 +1785,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:00:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.some((row) => row.kind === "turn-fold")).toBe(false); @@ -1674,8 +1879,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:01:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.some((row) => row.kind === "turn-fold")).toBe(false); @@ -1754,8 +1959,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:00:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.map((row) => row.kind)).toEqual(["working", "work-live"]); @@ -1827,8 +2032,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:00:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.map((row) => row.kind)).toEqual(["working", "work", "message", "work-live"]); @@ -1879,8 +2084,8 @@ describe("deriveMessagesTimelineRows", () => { expandedTurnIds: new Set([turnId]), isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.find((row) => row.kind === "work")).toMatchObject({ @@ -1947,8 +2152,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:00:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.map((row) => row.kind)).toEqual(["working", "work-live", "message", "work-live"]); @@ -1994,8 +2199,8 @@ describe("deriveMessagesTimelineRows", () => { latestTurn: null, isWorking: true, activeTurnStartedAt: "2026-01-01T00:01:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.some((row) => row.kind === "work-live")).toBe(false); @@ -2055,8 +2260,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:00:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.filter((row) => row.kind === "work-live").map((row) => row.entry.id)).toEqual([ @@ -2100,8 +2305,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:00:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const workLiveRow = rows.find((row) => row.kind === "work-live"); @@ -2148,8 +2353,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:00:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const initialRows = deriveRows(null); @@ -2226,8 +2431,8 @@ describe("deriveMessagesTimelineRows", () => { runningTurnId: "turn-2" as never, isWorking: true, activeTurnStartedAt: "2026-01-01T00:01:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.filter((row) => row.kind === "turn-fold").map((row) => row.turnId)).toEqual([ @@ -2271,8 +2476,8 @@ describe("deriveMessagesTimelineRows", () => { expandedTurnIds: new Set(["turn-1" as never]), isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const assistantRows = rows.filter( @@ -2309,8 +2514,8 @@ describe("deriveMessagesTimelineRows", () => { }, isWorking: true, activeTurnStartedAt: "2026-01-01T00:00:00Z", - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const assistantRow = rows.find( @@ -2373,8 +2578,8 @@ describe("deriveMessagesTimelineRows", () => { timelineEntries, isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }; const collapsedRows = deriveMessagesTimelineRows(baseInput); const expandedRows = deriveMessagesTimelineRows({ @@ -2467,8 +2672,8 @@ describe("deriveMessagesTimelineRows", () => { timelineEntries, isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(row).toMatchObject({ @@ -2509,8 +2714,8 @@ describe("deriveMessagesTimelineRows", () => { expandedTurnIds: new Set([turnId]), runningTurnId: isWorking ? turnId : null, activeTurnStartedAt: isWorking ? createdAt : null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }; const expandedRows = deriveMessagesTimelineRows({ ...input, @@ -2549,8 +2754,8 @@ describe("deriveMessagesTimelineRows", () => { timelineEntries, isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.find((row) => row.kind === "work-toggle")).toMatchObject({ @@ -2596,8 +2801,8 @@ describe("deriveMessagesTimelineRows", () => { timelineEntries, isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); expect(rows.find((row) => row.kind === "work-toggle")).toMatchObject({ @@ -2653,8 +2858,8 @@ describe("computeStableMessagesTimelineRows", () => { runningTurnId: turnId, isWorking: true, activeTurnStartedAt: startedAt, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }; const assistantEntry = { id: "assistant-entry", @@ -2731,8 +2936,8 @@ describe("computeStableMessagesTimelineRows", () => { ], isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const initial = computeStableMessagesTimelineRows(rows, { @@ -2780,8 +2985,8 @@ describe("computeStableMessagesTimelineRows", () => { ], isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const firstRows = createRows(); @@ -2836,8 +3041,8 @@ describe("computeStableMessagesTimelineRows", () => { ], isWorking: false, activeTurnStartedAt: null, - turnDiffSummaryByAssistantMessageId: new Map(), - revertTurnCountByUserMessageId: new Map(), + turnDiffSummaries: [], + supportsConversationRollback: false, }); const initial = computeStableMessagesTimelineRows(firstRows, { diff --git a/apps/web/src/components/chat/MessagesTimeline.logic.ts b/apps/web/src/components/chat/MessagesTimeline.logic.ts index 1288cd6fad8b..08e4f8d26037 100644 --- a/apps/web/src/components/chat/MessagesTimeline.logic.ts +++ b/apps/web/src/components/chat/MessagesTimeline.logic.ts @@ -18,6 +18,7 @@ export { } from "@t3tools/client-runtime/work-log/presentation"; import { formatDuration, + inferCheckpointTurnCountByTurnId, isStreamingMessageTextUpdate, workEntryDisplayIndicatesToolFailure, workEntryIndicatesToolSuccess, @@ -776,6 +777,45 @@ function attachTrailingToolGroupsToAssistant( return result; } +/** Match each user message to the next assistant checkpoint. */ +function buildRevertTurnCountByUserMessageId(input: { + supportsConversationRollback: boolean; + timelineEntries: ReadonlyArray; + turnDiffSummaryByAssistantMessageId: ReadonlyMap; + inferredCheckpointTurnCountByTurnId: Readonly>; +}): Map { + const byUserMessageId = new Map(); + const entryCount = input.supportsConversationRollback ? input.timelineEntries.length : 0; + for (let index = 0; index < entryCount; index += 1) { + const entry = input.timelineEntries[index]; + if (!entry || entry.kind !== "message" || entry.message.role !== "user") { + continue; + } + + for (let nextIndex = index + 1; nextIndex < input.timelineEntries.length; nextIndex += 1) { + const nextEntry = input.timelineEntries[nextIndex]; + if (!nextEntry || nextEntry.kind !== "message") { + continue; + } + if (nextEntry.message.role === "user") { + break; + } + const summary = input.turnDiffSummaryByAssistantMessageId.get(nextEntry.message.id); + if (!summary) { + continue; + } + const turnCount = + summary.checkpointTurnCount ?? input.inferredCheckpointTurnCountByTurnId[summary.turnId]; + if (typeof turnCount !== "number") { + break; + } + byUserMessageId.set(entry.message.id, Math.max(0, turnCount - 1)); + break; + } + } + return byUserMessageId; +} + export function deriveMessagesTimelineRows(input: { timelineEntries: ReadonlyArray; latestTurn?: TimelineLatestTurn | null; @@ -784,9 +824,23 @@ export function deriveMessagesTimelineRows(input: { expandedWorkGroupIds?: ReadonlySet; isWorking: boolean; activeTurnStartedAt: string | null; - turnDiffSummaryByAssistantMessageId: ReadonlyMap; - revertTurnCountByUserMessageId: ReadonlyMap; + turnDiffSummaries: ReadonlyArray; + supportsConversationRollback: boolean; }): MessagesTimelineRow[] { + const turnDiffSummaryByAssistantMessageId = new Map(); + for (const summary of input.turnDiffSummaries) { + if (summary.assistantMessageId) { + turnDiffSummaryByAssistantMessageId.set(summary.assistantMessageId, summary); + } + } + const revertTurnCountByUserMessageId = buildRevertTurnCountByUserMessageId({ + supportsConversationRollback: input.supportsConversationRollback, + timelineEntries: input.timelineEntries, + turnDiffSummaryByAssistantMessageId, + inferredCheckpointTurnCountByTurnId: input.supportsConversationRollback + ? inferCheckpointTurnCountByTurnId(input.turnDiffSummaries) + : {}, + }); const nextRows: MessagesTimelineRow[] = []; const durationStartByMessageId = computeMessageDurationStart( input.timelineEntries.flatMap((entry) => (entry.kind === "message" ? [entry.message] : [])), @@ -1134,11 +1188,11 @@ export function deriveMessagesTimelineRows(input: { assistantCopyStreaming: timelineEntry.message.streaming || assistantResponseStillInProgress, assistantTurnDiffSummary: timelineEntry.message.role === "assistant" - ? input.turnDiffSummaryByAssistantMessageId.get(timelineEntry.message.id) + ? turnDiffSummaryByAssistantMessageId.get(timelineEntry.message.id) : undefined, revertTurnCount: timelineEntry.message.role === "user" - ? input.revertTurnCountByUserMessageId.get(timelineEntry.message.id) + ? revertTurnCountByUserMessageId.get(timelineEntry.message.id) : undefined, }); } @@ -1168,9 +1222,30 @@ function replaceStreamingMessageRows( input: MessagesTimelineRowsInput, previous: MessagesTimelineRowsProjection, ): MessagesTimelineRow[] | null { - const { timelineEntries: previousEntries, ...previousContext } = previous.input; - const { timelineEntries, ...context } = input; - if (timelineEntries.length !== previousEntries.length || !shallow(previousContext, context)) { + const { + timelineEntries: previousEntries, + turnDiffSummaries: previousSummaries, + latestTurn: previousLatestTurn, + expandedTurnIds: previousExpandedTurns, + expandedWorkGroupIds: previousExpandedGroups, + ...previousContext + } = previous.input; + const { + timelineEntries, + turnDiffSummaries, + latestTurn, + expandedTurnIds, + expandedWorkGroupIds, + ...context + } = input; + if ( + timelineEntries.length !== previousEntries.length || + !shallow(previousContext, context) || + !shallow(previousSummaries, turnDiffSummaries) || + !shallow(previousLatestTurn, latestTurn) || + !shallow(previousExpandedTurns, expandedTurnIds) || + !shallow(previousExpandedGroups, expandedWorkGroupIds) + ) { return null; } const replacements = new Map(); diff --git a/apps/web/src/components/chat/MessagesTimeline.test.tsx b/apps/web/src/components/chat/MessagesTimeline.test.tsx index 21f30fa625d6..fa6d03eb4c78 100644 --- a/apps/web/src/components/chat/MessagesTimeline.test.tsx +++ b/apps/web/src/components/chat/MessagesTimeline.test.tsx @@ -1,12 +1,39 @@ import { CheckpointRef, EnvironmentId, MessageId, TurnId } from "@t3tools/contracts"; import { codexFeedbackMessage } from "@t3tools/client-runtime/state/threads"; -import { act, createRef, useLayoutEffect, type ReactNode, type Ref } from "react"; +import { + act, + createRef, + useLayoutEffect, + type ReactNode, + type Ref, + type ComponentProps, +} from "react"; import { renderToStaticMarkup } from "react-dom/server"; import { create, type ReactTestRenderer } from "react-test-renderer"; import { beforeAll, describe, expect, it, vi } from "vite-plus/test"; import type { LegendListRef, MaintainScrollAtEndOptions } from "@legendapp/list/react"; import { shouldUseRestingComposerLayout } from "../composerFooterLayout"; import { useComposerFocusState } from "./useComposerFocusState"; +import { + deriveTimelineEntriesWithState, + type TimelineEntriesProjection, +} from "../../session-logic"; +import type { ChatMessage, TurnDiffSummary } from "../../types"; + +vi.mock("../ui/tooltip", async () => { + const { cloneElement, isValidElement } = await import("react"); + return { + Tooltip: ({ children }: { children: ReactNode }) => <>{children}, + TooltipTrigger({ + render, + children, + }: ComponentProps) { + if (!isValidElement(render)) return <>{children}; + return children === undefined ? render : cloneElement(render, undefined, children); + }, + TooltipPopup: () => null, + }; +}); vi.mock("@legendapp/list/react", async () => { const legendListTestId = "legend-list"; @@ -154,6 +181,7 @@ beforeAll(async () => { clear: () => {}, }); vi.stubGlobal("window", { + localStorage, matchMedia, addEventListener: () => {}, removeEventListener: () => {}, @@ -184,11 +212,11 @@ function buildProps() { listRef: createRef(), latestTurn: null, runningTurnId: null, - turnDiffSummaryByAssistantMessageId: new Map(), + turnDiffSummaries: [], routeThreadKey: "environment-local:thread-1", onOpenTurnDiff: () => {}, - revertTurnCountByUserMessageId: new Map(), - onRevertUserMessage: () => {}, + supportsConversationRollback: false, + onRevertToTurnCount: () => {}, isRevertingCheckpoint: false, onImageExpand: () => {}, activeThreadEnvironmentId: ACTIVE_THREAD_ENVIRONMENT_ID, @@ -328,6 +356,131 @@ describe("MessagesTimeline", () => { }, ); + it("keeps code controls through streaming and recovery, then uses the current revert target", async () => { + vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true); + const { getSyntaxHighlighterPromise } = await import("../../lib/syntaxHighlighting"); + const highlighter = await getSyntaxHighlighterPromise("text"); + const codeToHtml = highlighter.codeToHtml.bind(highlighter); + let fail = true; + const highlight = vi.spyOn(highlighter, "codeToHtml").mockImplementation((...args) => { + if (fail) throw new Error("Temporary highlighter failure"); + return codeToHtml(...args); + }); + const error = vi.spyOn(console, "error").mockImplementation(() => {}); + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + const onRevertToTurnCount = vi.fn(); + const props = buildProps(); + const turnId = TurnId.make("streamed-checkpoint"); + const user: ChatMessage = buildUserTimelineEntry("Update the file").message; + const assistant: ChatMessage = { + ...buildAssistantTimelineEntry("").message, + id: MessageId.make("streamed-assistant"), + turnId, + createdAt: "2026-03-17T19:12:29.000Z", + }; + let timeline: TimelineEntriesProjection | null = null; + const render = ( + text: string, + streaming = true, + summaries: ReadonlyArray = [], + ) => { + timeline = deriveTimelineEntriesWithState( + [user, { ...assistant, text, streaming }], + [], + [], + timeline, + ); + return ( + + ); + }; + let renderer: ReactTestRenderer | undefined; + try { + await act(async () => { + renderer = create(render("```text\nInitial code\n```\n\nResponse")); + }); + const mounted = renderer!; + const codeBlock = mounted.root.findByProps({ "data-language": "text" }); + const initiallyWrapped = codeBlock.props["data-wrap"] === "true"; + const wrap = mounted.root + .findAllByType("button") + .find( + (button) => + button.props["aria-label"] === (initiallyWrapped ? "Disable line wrap" : "Wrap lines"), + ); + if (!wrap) throw new Error("Missing code wrap control"); + await act(() => wrap.props.onClick()); + const initialHighlightCount = highlight.mock.calls.length; + expect(initialHighlightCount).toBeGreaterThan(0); + for (let index = 0; index < 10; index += 1) { + await act(async () => { + mounted.update(render(`\`\`\`text\nInitial code\n\`\`\`\n\nResponse ${index}`)); + }); + } + expect(highlight).toHaveBeenCalledTimes(initialHighlightCount); + expect(mounted.root.findByProps({ "data-language": "text" })).toBe(codeBlock); + expect(codeBlock.props["data-wrap"]).toBe(String(!initiallyWrapped)); + + fail = false; + const finalText = "```text\nRecovered code\n```\n\nComplete"; + await act(async () => { + mounted.update(render(finalText)); + }); + const recoveredHighlightCount = highlight.mock.calls.length; + expect(recoveredHighlightCount).toBeGreaterThan(initialHighlightCount); + expect(mounted.root.findByProps({ "data-language": "text" })).toBe(codeBlock); + expect(codeBlock.props["data-wrap"]).toBe(String(!initiallyWrapped)); + const summary: TurnDiffSummary = { + turnId, + assistantMessageId: assistant.id, + checkpointTurnCount: 2, + checkpointRef: CheckpointRef.make("refs/t3/checkpoints/streamed"), + status: "ready", + files: [], + completedAt: assistant.createdAt, + }; + await act(async () => { + mounted.update(render(finalText, false, [summary])); + }); + const revert = () => { + const button = mounted.root + .findAllByType("button") + .find((entry) => entry.props["aria-label"] === "Revert to this message"); + if (!button) throw new Error("Missing revert action"); + button.props.onClick(); + }; + await act(revert); + expect(onRevertToTurnCount).toHaveBeenLastCalledWith(1); + await act(async () => { + mounted.update(render(finalText, false, [{ ...summary, checkpointTurnCount: 4 }])); + }); + await act(revert); + expect(onRevertToTurnCount).toHaveBeenLastCalledWith(3); + expect(highlight).toHaveBeenCalledTimes(recoveredHighlightCount); + expect(mounted.root.findByProps({ "data-language": "text" })).toBe(codeBlock); + expect(codeBlock.props["data-wrap"]).toBe(String(!initiallyWrapped)); + } finally { + await act(() => renderer?.unmount()); + highlight.mockRestore(); + error.mockRestore(); + warn.mockRestore(); + } + }); + it("renders a feedback command and its pending response as normal thread messages", () => { const submission = { id: MessageId.make("feedback-command"), @@ -448,22 +601,17 @@ describe("MessagesTimeline", () => { }, }, ]} - turnDiffSummaryByAssistantMessageId={ - new Map([ - [ - assistantMessageId, - { - turnId, - checkpointTurnCount: 1, - checkpointRef: CheckpointRef.make("checkpoint-with-files"), - status: "ready", - files: [{ path: "README.md", kind: "modified", additions: 2, deletions: 1 }], - assistantMessageId, - completedAt: MESSAGE_CREATED_AT, - }, - ], - ]) - } + turnDiffSummaries={[ + { + turnId, + checkpointTurnCount: 1, + checkpointRef: CheckpointRef.make("checkpoint-with-files"), + status: "ready", + files: [{ path: "README.md", kind: "modified", additions: 2, deletions: 1 }], + assistantMessageId, + completedAt: MESSAGE_CREATED_AT, + }, + ]} />, ); diff --git a/apps/web/src/components/chat/MessagesTimeline.tsx b/apps/web/src/components/chat/MessagesTimeline.tsx index c1b09899b8de..a113157aca33 100644 --- a/apps/web/src/components/chat/MessagesTimeline.tsx +++ b/apps/web/src/components/chat/MessagesTimeline.tsx @@ -203,7 +203,7 @@ interface TimelineRowSharedState { workspaceRoot: string | undefined; skills: ReadonlyArray>; activeThreadEnvironmentId: EnvironmentId; - onRevertUserMessage: (messageId: MessageId) => void; + onRevertToTurnCount: (targetTurnCount: number) => void; onUseArtifactTemplate: (template: CodexArtifactTemplate) => void; onImageExpand: (preview: ExpandedImagePreview) => void; onFileOpen: (attachment: ChatFileAttachment) => void; @@ -311,11 +311,11 @@ interface MessagesTimelineProps { timelineEntries: ReturnType; latestTurn: TimelineLatestTurn | null; runningTurnId: TurnId | null; - turnDiffSummaryByAssistantMessageId: Map; + turnDiffSummaries: ReadonlyArray; routeThreadKey: string; onOpenTurnDiff: (turnId: TurnId, filePath?: string) => void; - revertTurnCountByUserMessageId: Map; - onRevertUserMessage: (messageId: MessageId) => void; + supportsConversationRollback: boolean; + onRevertToTurnCount: (targetTurnCount: number) => void; onUseArtifactTemplate?: (template: CodexArtifactTemplate) => void; isRevertingCheckpoint: boolean; onImageExpand: (preview: ExpandedImagePreview) => void; @@ -369,11 +369,11 @@ export const MessagesTimeline = memo(function MessagesTimeline({ timelineEntries, latestTurn, runningTurnId, - turnDiffSummaryByAssistantMessageId, + turnDiffSummaries, routeThreadKey, onOpenTurnDiff, - revertTurnCountByUserMessageId, - onRevertUserMessage, + supportsConversationRollback, + onRevertToTurnCount, onUseArtifactTemplate = NOOP_USE_ARTIFACT_TEMPLATE, isRevertingCheckpoint, onImageExpand, @@ -543,8 +543,8 @@ export const MessagesTimeline = memo(function MessagesTimeline({ expandedWorkGroupIds, isWorking, activeTurnStartedAt, - turnDiffSummaryByAssistantMessageId, - revertTurnCountByUserMessageId, + turnDiffSummaries, + supportsConversationRollback, }, previous?.threadKey === routeThreadKey && previous.workspaceRoot === workspaceRoot ? previous.projection @@ -563,8 +563,8 @@ export const MessagesTimeline = memo(function MessagesTimeline({ expandedWorkGroupIds, isWorking, activeTurnStartedAt, - turnDiffSummaryByAssistantMessageId, - revertTurnCountByUserMessageId, + turnDiffSummaries, + supportsConversationRollback, ]); const rows = useStableRows(rawRows); const minimapItems = useMemo(() => deriveTimelineMinimapItems(rows), [rows]); @@ -729,7 +729,7 @@ export const MessagesTimeline = memo(function MessagesTimeline({ workspaceRoot, skills, activeThreadEnvironmentId, - onRevertUserMessage, + onRevertToTurnCount, onUseArtifactTemplate, onImageExpand, onFileOpen, @@ -753,7 +753,7 @@ export const MessagesTimeline = memo(function MessagesTimeline({ workspaceRoot, skills, activeThreadEnvironmentId, - onRevertUserMessage, + onRevertToTurnCount, onUseArtifactTemplate, onImageExpand, onFileOpen, @@ -1336,7 +1336,7 @@ function UserTimelineRow({ row }: { row: Extract image.name.startsWith("preview-annotation-")); const regularImages = userImages.filter((image) => !image.name.startsWith("preview-annotation-")); - const canRevertAgentWork = typeof row.revertTurnCount === "number"; + const revertTurnCount = row.revertTurnCount; return (
@@ -1498,7 +1498,9 @@ function UserTimelineRow({ row }: { row: Extract
- {canRevertAgentWork && } + {typeof revertTurnCount === "number" && ( + + )} {displayedUserMessage.copyText && ( )} @@ -1509,7 +1511,7 @@ function UserTimelineRow({ row }: { row: Extract ctx.onRevertUserMessage(messageId)} + onClick={() => ctx.onRevertToTurnCount(turnCount)} aria-label="Revert to this message" /> } From f48054c73f8839244168f53af9b03cdd3fefd6f8 Mon Sep 17 00:00:00 2001 From: Theo Browne Date: Sat, 5 Sep 2026 21:31:24 -0700 Subject: [PATCH 2/3] test(web): count all timeline checkpoint lookup reads --- .../chat/MessagesTimeline.logic.test.ts | 40 ++++++++++++------- 1 file changed, 26 insertions(+), 14 deletions(-) diff --git a/apps/web/src/components/chat/MessagesTimeline.logic.test.ts b/apps/web/src/components/chat/MessagesTimeline.logic.test.ts index b761348bb0a4..e1cf7f0dc6e7 100644 --- a/apps/web/src/components/chat/MessagesTimeline.logic.test.ts +++ b/apps/web/src/components/chat/MessagesTimeline.logic.test.ts @@ -263,18 +263,24 @@ describe("streaming row projection", () => { it("owns checkpoint lookups across streaming and equal source snapshots", () => { const initial = fixture("Partial"); - let checkpointReads = 0; + let checkpointLookupReads = 0; const summary: TurnDiffSummary = { turnId: initial.historyTurnId, - checkpointTurnCount: 1, + get checkpointTurnCount() { + checkpointLookupReads += 1; + return 1; + }, checkpointRef: CheckpointRef.make("refs/t3/checkpoints/history-turn"), status: "ready", files: [], get assistantMessageId() { - checkpointReads += 1; + checkpointLookupReads += 1; return MessageId.make("history-assistant"); }, - completedAt: initial.time(4), + get completedAt() { + checkpointLookupReads += 1; + return initial.time(4); + }, }; const input = { ...initial.input, @@ -284,7 +290,7 @@ describe("streaming row projection", () => { expandedWorkGroupIds: new Set(), }; const previous = deriveMessagesTimelineRowsWithState(input); - expect(checkpointReads).toBeGreaterThan(0); + expect(checkpointLookupReads).toBeGreaterThan(0); expect(previous.rows.some((row) => row.kind === "message" && row.revertTurnCount === 0)).toBe( true, ); @@ -299,9 +305,9 @@ describe("streaming row projection", () => { expandedTurnIds: new Set(input.expandedTurnIds), expandedWorkGroupIds: new Set(input.expandedWorkGroupIds), }; - checkpointReads = 0; + checkpointLookupReads = 0; const next = deriveMessagesTimelineRowsWithState(nextInput, previous); - expect(checkpointReads).toBe(0); + expect(checkpointLookupReads).toBe(0); expect(next.rows).toEqual(deriveMessagesTimelineRows(nextInput)); for (const [index, row] of previous.rows.entries()) { if ((row.kind === "message" || row.kind === "assistant-meta") && row.message === last) { @@ -332,7 +338,7 @@ describe("streaming row projection", () => { it("reuses long-thread rows through detail events, selectors, and attachment previews", () => { const initial = fixture("Partial"); - let checkpointReads = 0; + let checkpointLookupReads = 0; const history = Array.from({ length: 250 }, (_, index) => { const turnId = TurnId.make(`older-turn-${index}`); const user = { @@ -357,15 +363,21 @@ describe("streaming row projection", () => { }; const checkpoint: TurnDiffSummary = { turnId, - checkpointTurnCount: index + 1, + get checkpointTurnCount() { + checkpointLookupReads += 1; + return index + 1; + }, checkpointRef: CheckpointRef.make(`refs/t3/checkpoints/older-${index}`), status: "ready", files: [], get assistantMessageId() { - checkpointReads += 1; + checkpointLookupReads += 1; return assistant.id; }, - completedAt: assistant.createdAt, + get completedAt() { + checkpointLookupReads += 1; + return assistant.createdAt; + }, }; return { user, assistant, checkpoint }; }); @@ -470,8 +482,8 @@ describe("streaming row projection", () => { try { const first = project(); const saved = structuredClone(first.rows); - expect(checkpointReads).toBeGreaterThan(0); - checkpointReads = 0; + expect(checkpointLookupReads).toBeGreaterThan(0); + checkpointLookupReads = 0; for (let index = 0; index < 10; index += 1) { const next = send(` ${index}`, index + 1); for (const [rowIndex, row] of first.rows.entries()) { @@ -483,7 +495,7 @@ describe("streaming row projection", () => { expect(next.rows[rowIndex]).toBe(row); } } - expect(checkpointReads).toBe(0); + expect(checkpointLookupReads).toBe(0); const streamed = project(); expect(streamed.rows).toEqual(deriveMessagesTimelineRows(streamed.input)); From 3b4335984056a521e1693f158eaab2881f6b11a9 Mon Sep 17 00:00:00 2001 From: Theo Browne Date: Sun, 6 Sep 2026 02:13:39 -0700 Subject: [PATCH 3/3] fix(web): keep the timeline revert callback stable The revert callback in ChatView depended on the thread key and checkpoint snapshot, so every checkpoint update recreated it and rerendered every TimelineRowCtx consumer. Read the handler through a ref with an empty dependency list instead, and drop the stale-click guard that could only compare a callback against itself. Also remove the react-test-renderer test that asserted data-wrap props and highlighter call counts. The revert target update it covered is already checked in MessagesTimeline.logic.test.ts. Co-Authored-By: Claude Fable 5.1 --- apps/web/src/components/ChatView.tsx | 29 +--- .../components/chat/MessagesTimeline.test.tsx | 155 +----------------- 2 files changed, 8 insertions(+), 176 deletions(-) diff --git a/apps/web/src/components/ChatView.tsx b/apps/web/src/components/ChatView.tsx index db841e53ebfe..3974837f5f21 100644 --- a/apps/web/src/components/ChatView.tsx +++ b/apps/web/src/components/ChatView.tsx @@ -7746,28 +7746,13 @@ export default function ChatView(props: ChatViewProps) { }, [activeThreadRef, isServerThread, onDiffPanelOpen], ); - const revertCheckpoints = activeThread?.checkpoints; - const revertHandlerRef = useRef({ - threadKey: activeThreadKey, - checkpoints: revertCheckpoints, - handler: onRevertToTurnCount, - }); - revertHandlerRef.current = { - threadKey: activeThreadKey, - checkpoints: revertCheckpoints, - handler: onRevertToTurnCount, - }; - // A new thread or checkpoint snapshot must not retarget an old row's click. - const onRevertTimelineTurn = useCallback( - (targetTurnCount: number) => { - const current = revertHandlerRef.current; - if (current.threadKey !== activeThreadKey || current.checkpoints !== revertCheckpoints) { - return; - } - void current.handler(targetTurnCount); - }, - [activeThreadKey, revertCheckpoints], - ); + // The revert handler is read from a ref at call-time so the callback + // reference is fully stable and never busts TimelineRowCtx identity. + const onRevertToTurnCountRef = useRef(onRevertToTurnCount); + onRevertToTurnCountRef.current = onRevertToTurnCount; + const onRevertTimelineTurn = useCallback((targetTurnCount: number) => { + void onRevertToTurnCountRef.current(targetTurnCount); + }, []); // Empty state: no active thread if (!activeThread) { diff --git a/apps/web/src/components/chat/MessagesTimeline.test.tsx b/apps/web/src/components/chat/MessagesTimeline.test.tsx index fa6d03eb4c78..fe3e9ab4417c 100644 --- a/apps/web/src/components/chat/MessagesTimeline.test.tsx +++ b/apps/web/src/components/chat/MessagesTimeline.test.tsx @@ -1,39 +1,12 @@ import { CheckpointRef, EnvironmentId, MessageId, TurnId } from "@t3tools/contracts"; import { codexFeedbackMessage } from "@t3tools/client-runtime/state/threads"; -import { - act, - createRef, - useLayoutEffect, - type ReactNode, - type Ref, - type ComponentProps, -} from "react"; +import { act, createRef, useLayoutEffect, type ReactNode, type Ref } from "react"; import { renderToStaticMarkup } from "react-dom/server"; import { create, type ReactTestRenderer } from "react-test-renderer"; import { beforeAll, describe, expect, it, vi } from "vite-plus/test"; import type { LegendListRef, MaintainScrollAtEndOptions } from "@legendapp/list/react"; import { shouldUseRestingComposerLayout } from "../composerFooterLayout"; import { useComposerFocusState } from "./useComposerFocusState"; -import { - deriveTimelineEntriesWithState, - type TimelineEntriesProjection, -} from "../../session-logic"; -import type { ChatMessage, TurnDiffSummary } from "../../types"; - -vi.mock("../ui/tooltip", async () => { - const { cloneElement, isValidElement } = await import("react"); - return { - Tooltip: ({ children }: { children: ReactNode }) => <>{children}, - TooltipTrigger({ - render, - children, - }: ComponentProps) { - if (!isValidElement(render)) return <>{children}; - return children === undefined ? render : cloneElement(render, undefined, children); - }, - TooltipPopup: () => null, - }; -}); vi.mock("@legendapp/list/react", async () => { const legendListTestId = "legend-list"; @@ -181,7 +154,6 @@ beforeAll(async () => { clear: () => {}, }); vi.stubGlobal("window", { - localStorage, matchMedia, addEventListener: () => {}, removeEventListener: () => {}, @@ -356,131 +328,6 @@ describe("MessagesTimeline", () => { }, ); - it("keeps code controls through streaming and recovery, then uses the current revert target", async () => { - vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true); - const { getSyntaxHighlighterPromise } = await import("../../lib/syntaxHighlighting"); - const highlighter = await getSyntaxHighlighterPromise("text"); - const codeToHtml = highlighter.codeToHtml.bind(highlighter); - let fail = true; - const highlight = vi.spyOn(highlighter, "codeToHtml").mockImplementation((...args) => { - if (fail) throw new Error("Temporary highlighter failure"); - return codeToHtml(...args); - }); - const error = vi.spyOn(console, "error").mockImplementation(() => {}); - const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); - const onRevertToTurnCount = vi.fn(); - const props = buildProps(); - const turnId = TurnId.make("streamed-checkpoint"); - const user: ChatMessage = buildUserTimelineEntry("Update the file").message; - const assistant: ChatMessage = { - ...buildAssistantTimelineEntry("").message, - id: MessageId.make("streamed-assistant"), - turnId, - createdAt: "2026-03-17T19:12:29.000Z", - }; - let timeline: TimelineEntriesProjection | null = null; - const render = ( - text: string, - streaming = true, - summaries: ReadonlyArray = [], - ) => { - timeline = deriveTimelineEntriesWithState( - [user, { ...assistant, text, streaming }], - [], - [], - timeline, - ); - return ( - - ); - }; - let renderer: ReactTestRenderer | undefined; - try { - await act(async () => { - renderer = create(render("```text\nInitial code\n```\n\nResponse")); - }); - const mounted = renderer!; - const codeBlock = mounted.root.findByProps({ "data-language": "text" }); - const initiallyWrapped = codeBlock.props["data-wrap"] === "true"; - const wrap = mounted.root - .findAllByType("button") - .find( - (button) => - button.props["aria-label"] === (initiallyWrapped ? "Disable line wrap" : "Wrap lines"), - ); - if (!wrap) throw new Error("Missing code wrap control"); - await act(() => wrap.props.onClick()); - const initialHighlightCount = highlight.mock.calls.length; - expect(initialHighlightCount).toBeGreaterThan(0); - for (let index = 0; index < 10; index += 1) { - await act(async () => { - mounted.update(render(`\`\`\`text\nInitial code\n\`\`\`\n\nResponse ${index}`)); - }); - } - expect(highlight).toHaveBeenCalledTimes(initialHighlightCount); - expect(mounted.root.findByProps({ "data-language": "text" })).toBe(codeBlock); - expect(codeBlock.props["data-wrap"]).toBe(String(!initiallyWrapped)); - - fail = false; - const finalText = "```text\nRecovered code\n```\n\nComplete"; - await act(async () => { - mounted.update(render(finalText)); - }); - const recoveredHighlightCount = highlight.mock.calls.length; - expect(recoveredHighlightCount).toBeGreaterThan(initialHighlightCount); - expect(mounted.root.findByProps({ "data-language": "text" })).toBe(codeBlock); - expect(codeBlock.props["data-wrap"]).toBe(String(!initiallyWrapped)); - const summary: TurnDiffSummary = { - turnId, - assistantMessageId: assistant.id, - checkpointTurnCount: 2, - checkpointRef: CheckpointRef.make("refs/t3/checkpoints/streamed"), - status: "ready", - files: [], - completedAt: assistant.createdAt, - }; - await act(async () => { - mounted.update(render(finalText, false, [summary])); - }); - const revert = () => { - const button = mounted.root - .findAllByType("button") - .find((entry) => entry.props["aria-label"] === "Revert to this message"); - if (!button) throw new Error("Missing revert action"); - button.props.onClick(); - }; - await act(revert); - expect(onRevertToTurnCount).toHaveBeenLastCalledWith(1); - await act(async () => { - mounted.update(render(finalText, false, [{ ...summary, checkpointTurnCount: 4 }])); - }); - await act(revert); - expect(onRevertToTurnCount).toHaveBeenLastCalledWith(3); - expect(highlight).toHaveBeenCalledTimes(recoveredHighlightCount); - expect(mounted.root.findByProps({ "data-language": "text" })).toBe(codeBlock); - expect(codeBlock.props["data-wrap"]).toBe(String(!initiallyWrapped)); - } finally { - await act(() => renderer?.unmount()); - highlight.mockRestore(); - error.mockRestore(); - warn.mockRestore(); - } - }); - it("renders a feedback command and its pending response as normal thread messages", () => { const submission = { id: MessageId.make("feedback-command"),