diff --git a/frontend/src/core/threads/hooks.ts b/frontend/src/core/threads/hooks.ts index 1201e1619..760b54d9e 100644 --- a/frontend/src/core/threads/hooks.ts +++ b/frontend/src/core/threads/hooks.ts @@ -365,6 +365,107 @@ export function mergeMessages( }); } +/** + * Derive the live turns that context summarization is about to drop and that + * therefore must be re-archived into history. + * + * Summarization emits `RemoveMessage(ALL)` + a hidden summary + the retained + * tail. Everything in the current live thread before the first retained visible + * message is being removed; we keep those (minus the summary control messages + * already tracked) so the UI can still show the full conversation (#3825). + */ +export function computeSummarizationMovedMessages( + currentMessages: Message[], + summarizationMessages: Message[], + summarizedMessageIds: ReadonlySet, +): Message[] { + const firstRetainedVisibleIdentity = summarizationMessages + .filter((message) => message.type !== "remove") + .filter((message) => !isHiddenFromUIMessage(message)) + .map(messageIdentity) + .find(isNonEmptyString); + + const moved: Message[] = []; + for (const message of currentMessages) { + if ( + firstRetainedVisibleIdentity && + messageIdentity(message) === firstRetainedVisibleIdentity + ) { + break; + } + if (!summarizedMessageIds.has(message.id ?? "")) { + moved.push(message); + } + } + return moved; +} + +/** + * Overlay the messages rescued from context summarization on top of the + * (possibly stale) visible history so the merged view never drops them. + * + * Background (#3825): after summarization the backend removes every live + * message (`RemoveMessage(ALL)`) and `onUpdateEvent` re-archives the removed + * messages into history through an async `setState`. The live thread messages + * are owned by the LangGraph SDK external store while the archived history is + * React state, so a render can observe the post-summary (shrunk) thread before + * the archive `setState` commits — leaving the rescued messages in neither + * merge input. Reading them from a synchronous buffer here keeps the merge + * correct at every render regardless of how the two state channels interleave. + * + * The rescued messages are the oldest live turns, so they follow whatever the + * already-loaded history holds. Only messages still missing from history are + * appended: once history absorbs a rescued message, its live copy stays + * authoritative (the buffered copy is an older snapshot and must never overwrite + * it), and ordering is preserved. + */ +export function resolvePreservedHistory( + visibleHistory: Message[], + pendingArchivedMessages: Message[], +): Message[] { + if (pendingArchivedMessages.length === 0) { + return visibleHistory; + } + const presentIdentities = new Set( + visibleHistory.map(messageIdentity).filter(isNonEmptyString), + ); + const missing = pendingArchivedMessages.filter((message) => { + const identity = messageIdentity(message); + // Identity-less messages are intentionally skipped: without a stable + // identity they cannot be matched against history to drain or dedupe, so + // overlaying them would risk a permanent duplicate. They are still archived + // through appendMessages and surface via the normal history path instead. + return identity !== undefined && !presentIdentities.has(identity); + }); + if (missing.length === 0) { + return visibleHistory; + } + return [...visibleHistory, ...missing]; +} + +/** + * Drop the archive-buffer entries that the canonical history state has already + * absorbed. This keeps the buffer a transient bridge across the async gap + * rather than a second long-lived source of truth — otherwise a stale copy + * could resurrect a message that history later filtered out (e.g. a superseded + * or regenerated run). + */ +export function pruneConfirmedArchivedMessages( + pendingArchivedMessages: Message[], + visibleHistory: Message[], +): Message[] { + if (pendingArchivedMessages.length === 0) { + return pendingArchivedMessages; + } + const confirmedIdentities = new Set( + visibleHistory.map(messageIdentity).filter(isNonEmptyString), + ); + return pendingArchivedMessages.filter((message) => { + const identity = messageIdentity(message); + return !identity || !confirmedIdentities.has(identity); + }); +} + function getMessagesAfterBaseline( messages: Message[], baselineMessageIds: ReadonlySet, @@ -734,24 +835,19 @@ export function useThreadStream({ summarizedRef.current?.add(m.id ?? ""); } } - const firstRetainedVisibleIdentity = _messages - .filter((message) => message.type !== "remove") - .filter((message) => !isHiddenFromUIMessage(message)) - .map(messageIdentity) - .find(isNonEmptyString); - const _currentMessages = [...messagesRef.current]; - const _movedMessages: Message[] = []; - for (const m of _currentMessages) { - if ( - firstRetainedVisibleIdentity && - messageIdentity(m) === firstRetainedVisibleIdentity - ) { - break; - } - if (!summarizedRef.current?.has(m.id ?? "")) { - _movedMessages.push(m); - } - } + const _movedMessages = computeSummarizationMovedMessages( + messagesRef.current, + _messages, + summarizedRef.current ?? new Set(), + ); + // Buffer the rescued messages synchronously so the merge can keep + // displaying them immediately, even though appendMessages below only + // updates the archived-history state asynchronously (#3825). + pendingArchivedMessagesRef.current = dedupeMessagesByIdentity([ + ...pendingArchivedMessagesRef.current, + ..._movedMessages, + ]); + pendingArchiveThreadIdRef.current = threadIdRef.current; appendMessages(_movedMessages); messagesRef.current = []; } @@ -896,6 +992,17 @@ export function useThreadStream({ const latestMessageCountsRef = useRef({ humanMessageCount }); const sendInFlightRef = useRef(false); const messagesRef = useRef([]); + // Synchronous bridge for messages rescued from context summarization. The + // archived-history `setState` (via appendMessages) lands on a different + // schedule than the live thread external store, so the merge reads this buffer + // to avoid dropping rescued messages in the render window before history + // catches up (#3825). + const pendingArchivedMessagesRef = useRef([]); + // The thread the rescue buffer belongs to, captured when onUpdateEvent fills + // it. The merge only overlays the buffer when this matches the viewed + // `threadId`, so a previous thread's rescued messages can never flash into + // another thread or the new-chat screen (#3825). + const pendingArchiveThreadIdRef = useRef(null); const summarizedRef = useRef>(null); // Track human message count before sending to prevent clearing optimistic // messages before the server's human message arrives (e.g. when AI messages @@ -912,6 +1019,8 @@ export function useThreadStream({ startedRef.current = false; sendInFlightRef.current = false; messagesRef.current = []; + pendingArchivedMessagesRef.current = []; + pendingArchiveThreadIdRef.current = null; summarizedRef.current = new Set(); pendingUsageBaselineMessageIdsRef.current = new Set(); setPendingSupersededRunIds(new Set()); @@ -920,6 +1029,16 @@ export function useThreadStream({ latestMessageCountsRef.current.humanMessageCount; }, [threadId]); + // Release archive-buffer entries once the canonical history state has absorbed + // them, so the synchronous bridge stays transient and never resurrects a + // message that history later filters out (e.g. a superseded run) (#3825). + useEffect(() => { + pendingArchivedMessagesRef.current = pruneConfirmedArchivedMessages( + pendingArchivedMessagesRef.current, + visibleHistory, + ); + }, [visibleHistory]); + useEffect(() => { if (optimisticThreadId && optimisticThreadId !== currentViewThreadId) { setOptimisticMessages([]); @@ -1295,8 +1414,18 @@ export function useThreadStream({ humanMessageCount, ); + // Overlay the summarization rescue buffer only onto the history of the thread + // it was captured from. visibleHistory is gated on `threadId`, so comparing the + // same prop keeps the buffer from flashing into another thread or the new-chat + // screen, and reading it here (instead of clearing a ref during render) is + // concurrent-mode safe (#3825). + const rescueBuffer = pendingArchivedMessagesRef.current; + const effectiveHistory = + rescueBuffer.length > 0 && pendingArchiveThreadIdRef.current === threadId + ? resolvePreservedHistory(visibleHistory, rescueBuffer) + : visibleHistory; const mergedMessages = mergeMessages( - visibleHistory, + effectiveHistory, persistedMessages, visibleOptimisticMessages, ); diff --git a/frontend/tests/unit/core/threads/message-merge.test.ts b/frontend/tests/unit/core/threads/message-merge.test.ts index 2eb5693dd..264e346d9 100644 --- a/frontend/tests/unit/core/threads/message-merge.test.ts +++ b/frontend/tests/unit/core/threads/message-merge.test.ts @@ -4,6 +4,7 @@ import { expect, test } from "@rstest/core"; import { buildRunMessagesUrl, buildVisibleHistoryMessages, + computeSummarizationMovedMessages, findLatestUnloadedRunIndex, getNextRunMessagesBeforeSeq, getOldestRunMessageSeq, @@ -12,7 +13,9 @@ import { getVisibleOptimisticMessages, MAX_CONSECUTIVE_EMPTY_RUN_LOADS, mergeMessages, + pruneConfirmedArchivedMessages, removeSetItems, + resolvePreservedHistory, runMessagesPageHasMore, shouldAutoContinueOnEmptyRun, } from "@/core/threads/hooks"; @@ -587,3 +590,250 @@ test("shouldAutoContinueOnEmptyRun input must use the post-filter visible count, expect(shouldAutoContinueOnEmptyRun(filteredVisibleCount, 0)).toBe(true); expect(shouldAutoContinueOnEmptyRun(rawPageSize, 0)).toBe(false); }); + +// Regression coverage for #3825: after context summarization the backend emits +// RemoveMessage(ALL) + summary + retained, and onUpdateEvent rescues the removed +// messages into history via an async setState. The live thread.messages (an +// external store) and the archived history (React state) update through two +// independent scheduling channels, so a render can observe the post-summary +// (shrunk) thread while the rescued messages have NOT yet landed in +// visibleHistory. resolvePreservedHistory overlays a synchronous archive buffer +// so the merge never loses those messages regardless of the interleaving. + +const summarizationHuman1 = { + id: "human-1", + type: "human", + content: "round 1 question", +} as Message; +const summarizationAi1 = { + id: "ai-1", + type: "ai", + content: "round 1 answer", +} as Message; +const summarizationHuman2 = { + id: "human-2", + type: "human", + content: "round 2 question", +} as Message; +const summarizationAi2 = { + id: "ai-2", + type: "ai", + content: "round 2 answer (retained)", +} as Message; +const summarizationMovedMessages = [ + summarizationHuman1, + summarizationAi1, + summarizationHuman2, +]; + +test("resolvePreservedHistory keeps rescued messages while history state is still stale (regression for #3825)", () => { + // visibleHistory has not yet absorbed the rescued messages (async setState + // from appendMessages is still pending in this render). + const staleHistory: Message[] = []; + + expect( + resolvePreservedHistory(staleHistory, summarizationMovedMessages), + ).toEqual(summarizationMovedMessages); +}); + +test("resolvePreservedHistory appends rescued messages after already-loaded history", () => { + const olderLoadedHuman = { + id: "older-human", + type: "human", + content: "older loaded turn", + } as Message; + + expect( + resolvePreservedHistory([olderLoadedHuman], summarizationMovedMessages), + ).toEqual([olderLoadedHuman, ...summarizationMovedMessages]); +}); + +test("resolvePreservedHistory does not duplicate or reorder once history state catches up", () => { + // visibleHistory now contains the rescued messages (appendMessages committed), + // but the synchronous buffer still holds them this render. + expect( + resolvePreservedHistory( + summarizationMovedMessages, + summarizationMovedMessages, + ), + ).toEqual(summarizationMovedMessages); +}); + +test("resolvePreservedHistory returns history unchanged when nothing is pending archival", () => { + const history = [summarizationHuman1, summarizationAi1]; + expect(resolvePreservedHistory(history, [])).toBe(history); +}); + +test("merge keeps the full conversation across summarization even when visibleHistory lags (regression for #3825)", () => { + // Hidden summary (name === "summary") + the retained latest answer is all the + // live thread carries after RemoveMessage(ALL). + const hiddenSummary = { + id: "summary-1", + type: "human", + name: "summary", + content: "conversation summary", + } as Message; + const postSummaryThread = [hiddenSummary, summarizationAi2]; + + // The bad render: visibleHistory is still empty, so without the buffer the + // rescued round-1/2 messages exist in neither merge input and are lost. + const effectiveHistory = resolvePreservedHistory( + [], + summarizationMovedMessages, + ); + const merged = mergeMessages(effectiveHistory, postSummaryThread, []); + + expect(merged.map((m) => m.id)).toEqual([ + "human-1", + "ai-1", + "human-2", + "summary-1", + "ai-2", + ]); +}); + +test("pruneConfirmedArchivedMessages drops messages history has absorbed but keeps the rest", () => { + // History has caught up on the first two rescued messages only. + expect( + pruneConfirmedArchivedMessages(summarizationMovedMessages, [ + summarizationHuman1, + summarizationAi1, + ]), + ).toEqual([summarizationHuman2]); +}); + +test("pruneConfirmedArchivedMessages keeps every pending message while history has not caught up", () => { + expect( + pruneConfirmedArchivedMessages(summarizationMovedMessages, []), + ).toEqual(summarizationMovedMessages); +}); + +test("resolvePreservedHistory prefers the live history copy over a stale buffered duplicate (#3825 review #3)", () => { + // Same identity, but the buffered copy is an older snapshot. The live history + // copy (e.g. the finalized answer) must win — the buffer only fills gaps, it + // must never overwrite a message history already shows. + const staleBuffered = { + id: "ai-1", + type: "ai", + content: "streaming partial", + } as Message; + const liveFinal = { + id: "ai-1", + type: "ai", + content: "finalized answer", + } as Message; + + expect(resolvePreservedHistory([liveFinal], [staleBuffered])).toEqual([ + liveFinal, + ]); +}); + +test("computeSummarizationMovedMessages returns the live turns dropped before the retained boundary (regression for #3825)", () => { + const removeAll = { + id: "__remove_all__", + type: "remove", + content: "", + } as Message; + const hiddenSummary = { + id: "summary-1", + type: "human", + name: "summary", + content: "conversation summary", + } as Message; + const liveThreadBeforeSummary = [ + summarizationHuman1, + summarizationAi1, + summarizationHuman2, + summarizationAi2, + ]; + // Summarization emits RemoveMessage(ALL) + hidden summary + retained answer. + const summarizationMessages = [removeAll, hiddenSummary, summarizationAi2]; + + expect( + computeSummarizationMovedMessages( + liveThreadBeforeSummary, + summarizationMessages, + new Set([hiddenSummary.id!]), + ), + ).toEqual([summarizationHuman1, summarizationAi1, summarizationHuman2]); +}); + +test("computeSummarizationMovedMessages excludes already-summarized control messages", () => { + const priorSummary = { + id: "summary-0", + type: "human", + name: "summary", + content: "earlier summary", + } as Message; + const liveThreadBeforeSummary = [ + priorSummary, + summarizationHuman1, + summarizationAi1, + summarizationAi2, + ]; + const summarizationMessages = [ + { id: "__remove_all__", type: "remove", content: "" } as Message, + { + id: "summary-1", + type: "human", + name: "summary", + content: "new summary", + } as Message, + summarizationAi2, + ]; + + // priorSummary is in the summarized set, so it must not be re-archived. + expect( + computeSummarizationMovedMessages( + liveThreadBeforeSummary, + summarizationMessages, + new Set([priorSummary.id!, "summary-1"]), + ), + ).toEqual([summarizationHuman1, summarizationAi1]); +}); + +test("full summarization rescue pipeline keeps the conversation when history state lags (regression for #3825)", () => { + // Exercises the whole rescue algorithm the hook runs: derive the moved + // messages, buffer them, then merge against the post-summary thread while the + // archived-history React state is still stale (empty). + const removeAll = { + id: "__remove_all__", + type: "remove", + content: "", + } as Message; + const hiddenSummary = { + id: "summary-1", + type: "human", + name: "summary", + content: "conversation summary", + } as Message; + const liveThreadBeforeSummary = [ + summarizationHuman1, + summarizationAi1, + summarizationHuman2, + summarizationAi2, + ]; + const summarizationMessages = [removeAll, hiddenSummary, summarizationAi2]; + + const moved = computeSummarizationMovedMessages( + liveThreadBeforeSummary, + summarizationMessages, + new Set([hiddenSummary.id!]), + ); + const staleHistory: Message[] = []; + const postSummaryThread = [hiddenSummary, summarizationAi2]; + + const merged = mergeMessages( + resolvePreservedHistory(staleHistory, moved), + postSummaryThread, + [], + ); + + expect(merged.map((m) => m.id)).toEqual([ + "human-1", + "ai-1", + "human-2", + "summary-1", + "ai-2", + ]); +});