diff --git a/packages/app/src/agent-stream/chat-find/index.web.tsx b/packages/app/src/agent-stream/chat-find/index.web.tsx index eadcec4e22..df61d4e859 100644 --- a/packages/app/src/agent-stream/chat-find/index.web.tsx +++ b/packages/app/src/agent-stream/chat-find/index.web.tsx @@ -56,6 +56,7 @@ export function ChatFind({ agentId, serverId, epoch, + rewindSeq, items, viewportRef, revealLoadedMessage, @@ -95,8 +96,8 @@ export function ChatFind({ ); const state = useSyncExternalStore(model.subscribe, model.getSnapshot, model.getSnapshot); useEffect(() => { - model.updateHistory(epoch, items); - }, [model, epoch, items]); + model.updateHistory(epoch, items, rewindSeq); + }, [model, epoch, items, rewindSeq]); useEffect(() => () => model.close(), [model]); useEffect(() => { if (!active) model.close(); diff --git a/packages/app/src/agent-stream/chat-find/model.test.ts b/packages/app/src/agent-stream/chat-find/model.test.ts index 9dab7b337d..16b544624f 100644 --- a/packages/app/src/agent-stream/chat-find/model.test.ts +++ b/packages/app/src/agent-stream/chat-find/model.test.ts @@ -55,7 +55,7 @@ it("loads an older location and searches a response merged across suppressed tas }, async load(_epoch, seq) { loaded.push(seq); - model.updateHistory("epoch", history); + model.updateHistory("epoch", history, null); }, async reveal(_id, _query, occurrence) { revealed.push(occurrence); @@ -64,7 +64,7 @@ it("loads an older location and searches a response merged across suppressed tas clear() {}, }); try { - model.updateHistory("epoch", []); + model.updateHistory("epoch", [], null); model.open(); model.setQuery("target"); await expect.poll(() => model.getSnapshot().phase).toBe("ready"); @@ -99,7 +99,7 @@ it("skips a source-only candidate and highlights a displayed match", async () => clear() {}, }); try { - model.updateHistory("epoch", [message(1), message(2)]); + model.updateHistory("epoch", [message(1), message(2)], null); model.open(); model.setQuery("target"); await expect.poll(() => model.getSnapshot().selectedItemId).toBe("row-2"); @@ -129,7 +129,7 @@ it("ignores replies after closing and does not navigate", async () => { }, clear() {}, }); - model.updateHistory("epoch", [message(1)]); + model.updateHistory("epoch", [message(1)], null); model.open(); model.setQuery("target"); await expect.poll(() => calls).toEqual(["search"]); @@ -155,7 +155,7 @@ it("retains a failed search for retry and distinguishes failure from zero matche }; const model = new ChatFindModel(operations); try { - model.updateHistory("epoch", []); + model.updateHistory("epoch", [], null); model.open(); model.setQuery("target"); await expect.poll(() => model.getSnapshot().phase).toBe("error"); @@ -185,7 +185,7 @@ it("reports a failed timeline load as a connection failure", async () => { clear() {}, }); try { - model.updateHistory("epoch", []); + model.updateHistory("epoch", [], null); model.open(); model.setQuery("target"); await expect.poll(() => model.getSnapshot().phase).toBe("error"); @@ -207,24 +207,87 @@ it("reports a search result that outlived its epoch as a changed history", async clear() {}, }); try { - model.updateHistory("epoch", [message(1)]); + model.updateHistory("epoch", [message(1)], null); model.open(); model.setQuery("target"); await expect.poll(() => model.getSnapshot().phase).toBe("ready"); - model.updateHistory("epoch-2", [message(1)]); + model.updateHistory("epoch-2", [message(1)], null); expect(model.getSnapshot()).toMatchObject({ phase: "error", failure: "historyChanged" }); } finally { model.close(); } }); +it("reports results that a rewind took messages from as a changed history", async () => { + const model = new ChatFindModel({ + async search() { + return response([{ seq: 2, role: "assistant" }]); + }, + async load() {}, + async reveal() { + return { count: 1, occurrence: 0 }; + }, + clear() {}, + }); + try { + model.updateHistory("epoch", [message(1), message(2)], null); + model.open(); + model.setQuery("target"); + await expect.poll(() => model.getSnapshot().phase).toBe("ready"); + expect(model.getSnapshot()).toMatchObject({ selectedItemId: "row-2", count: 1 }); + // Another client rewinds the turn holding the match; the epoch stays. + model.updateHistory("epoch", [message(1)], 3); + expect(model.getSnapshot()).toMatchObject({ + phase: "error", + failure: "historyChanged", + selectedItemId: null, + }); + } finally { + model.close(); + } +}); + +it("stops a search still receiving pages when a rewind lands between them", async () => { + let pages = 0; + let releaseSecondPage!: () => void; + const secondPage = new Promise((resolve) => { + releaseSecondPage = resolve; + }); + const model = new ChatFindModel({ + async search(_query, cursor) { + pages += 1; + if (cursor === undefined) + return { ...response([{ seq: 2, role: "assistant" }]), nextCursor: 2 }; + await secondPage; + return response([{ seq: 1, role: "assistant" }]); + }, + async load() {}, + async reveal() { + return { count: 1, occurrence: 0 }; + }, + clear() {}, + }); + try { + model.updateHistory("epoch", [message(1), message(2)], null); + model.open(); + model.setQuery("target"); + await expect.poll(() => pages).toBe(2); + model.updateHistory("epoch", [message(1)], 3); + releaseSecondPage(); + await expect.poll(() => model.getSnapshot().failure).toBe("historyChanged"); + expect(model.getSnapshot()).toMatchObject({ phase: "error", selectedItemId: null }); + } finally { + model.close(); + } +}); + it("can load a search location before the initial chat history has arrived", async () => { const model = new ChatFindModel({ async search() { return response([{ seq: 1, role: "assistant" }]); }, async load() { - model.updateHistory("epoch", [message(1)]); + model.updateHistory("epoch", [message(1)], null); }, async reveal() { return { count: 1, occurrence: 0 }; @@ -232,7 +295,7 @@ it("can load a search location before the initial chat history has arrived", asy clear() {}, }); try { - model.updateHistory(null, []); + model.updateHistory(null, [], null); model.open(); model.setQuery("target"); await expect.poll(() => model.getSnapshot().phase).toBe("ready"); @@ -260,7 +323,7 @@ it("reports the position across the whole chat and corrects it with rendered cou clear() {}, }); try { - model.updateHistory("epoch", [message(1), message(2)]); + model.updateHistory("epoch", [message(1), message(2)], null); model.open(); model.setQuery("target"); await expect.poll(() => model.getSnapshot().phase).toBe("ready"); @@ -297,7 +360,7 @@ it("counts one match per message from a host that reports no counts", async () = clear() {}, }); try { - model.updateHistory("epoch", [message(1), message(2)]); + model.updateHistory("epoch", [message(1), message(2)], null); model.open(); model.setQuery("target"); await expect.poll(() => model.getSnapshot().phase).toBe("ready"); diff --git a/packages/app/src/agent-stream/chat-find/model.ts b/packages/app/src/agent-stream/chat-find/model.ts index 98acc9840f..d916a6ba2b 100644 --- a/packages/app/src/agent-stream/chat-find/model.ts +++ b/packages/app/src/agent-stream/chat-find/model.ts @@ -76,6 +76,10 @@ export class ChatFindModel { private items: StreamItem[] = []; private epoch: string | null = null; private resultEpoch: string | null = null; + /** The latest rewind applied to the history. A rewind keeps the epoch and retires messages. */ + private rewindSeq: number | null = null; + /** The latest rewind when the current search started; undefined while no search runs or holds results. */ + private searchedRewindSeq: number | null | undefined = undefined; private locations: Location[] = []; private resolved = new Map(); private current = -1; @@ -95,6 +99,7 @@ export class ChatFindModel { this.listeners.forEach((listener) => listener()); } private cancel() { + this.searchedRewindSeq = undefined; this.abort.abort(); this.abort = new AbortController(); clearTimeout(this.timer); @@ -144,6 +149,7 @@ export class ChatFindModel { }); } private async search(signal: AbortSignal) { + this.searchedRewindSeq = this.rewindSeq; try { let cursor: number | undefined; do { @@ -168,15 +174,14 @@ export class ChatFindModel { this.fail(error, signal); } } - updateHistory(epoch: string | null, items: StreamItem[]) { - if ( - this.resultEpoch !== null && - epoch !== null && - epoch !== this.resultEpoch && - this.state.open - ) { + /** `rewindSeq` is the sequence number of the latest rewind applied to this history. */ + updateHistory(epoch: string | null, items: StreamItem[], rewindSeq: number | null) { + const replaced = this.resultEpoch !== null && epoch !== null && epoch !== this.resultEpoch; + const rewound = this.searchedRewindSeq !== undefined && rewindSeq !== this.searchedRewindSeq; + if ((replaced || rewound) && this.state.open) { this.fail(new ChatFindFailureError("historyChanged"), this.abort.signal); } + this.rewindSeq = rewindSeq; this.epoch = epoch; this.items = items; this.historyListeners.forEach((listener) => listener()); diff --git a/packages/app/src/agent-stream/chat-find/types.ts b/packages/app/src/agent-stream/chat-find/types.ts index 4fdf543a5e..2338a09540 100644 --- a/packages/app/src/agent-stream/chat-find/types.ts +++ b/packages/app/src/agent-stream/chat-find/types.ts @@ -9,6 +9,8 @@ export interface ChatFindProps { agentId: string; serverId: string; epoch: string | null; + /** The sequence number of the latest rewind applied to this history, if any. */ + rewindSeq: number | null; items: StreamItem[]; viewportRef: RefObject; revealLoadedMessage(messageId: string): boolean; diff --git a/packages/app/src/agent-stream/view.tsx b/packages/app/src/agent-stream/view.tsx index 761e73bac6..b5e0c84b01 100644 --- a/packages/app/src/agent-stream/view.tsx +++ b/packages/app/src/agent-stream/view.tsx @@ -406,6 +406,10 @@ const AgentStreamViewComponent = forwardRef state.sessions[resolvedServerId]?.agentTimelineCursor.get(agentId)?.epoch ?? null, ); + const timelineRewindSeq = useSessionStore((state) => { + const rewind = state.sessions[resolvedServerId]?.agentTimelineLatestRewind.get(agentId); + return rewind && rewind.epoch === timelineEpoch ? rewind.seq : null; + }); const isTimelineDetached = useSessionStore( (state) => state.sessions[resolvedServerId]?.agentTimelineHasNewer.get(agentId) === true, ); @@ -1160,6 +1164,7 @@ const AgentStreamViewComponent = forwardRef; } +/** The latest conversation rewind applied to an agent's history, by its sequence number. */ +export interface AgentTimelineRewind { + epoch: string; + seq: number; +} + export type AgentTimelineState = | { status: "cold" } | { status: "painted"; items: StreamItem[] } @@ -490,6 +496,7 @@ export interface SessionState { agentTimelineUnsettledTurn: Map; agentTimelineHasOlder: Map; agentTimelineHasNewer: Map; + agentTimelineLatestRewind: Map; agentTimelineOlderFetchInFlight: Map; historySyncGeneration: number; agentHistorySyncGeneration: Map; @@ -609,6 +616,12 @@ interface SessionStoreActions { serverId: string, state: Map | ((prev: Map) => Map), ) => void; + /** Records a rewind applied to the agent's history; an older one of the same history is ignored. */ + recordAgentTimelineRewind: ( + serverId: string, + agentId: string, + rewind: AgentTimelineRewind, + ) => void; setAgentTimelineOlderFetchInFlight: ( serverId: string, state: Map | ((prev: Map) => Map), @@ -745,6 +758,7 @@ function createInitialSessionState( agentTimelineUnsettledTurn: new Map(), agentTimelineHasOlder: new Map(), agentTimelineHasNewer: new Map(), + agentTimelineLatestRewind: new Map(), agentTimelineOlderFetchInFlight: new Map(), historySyncGeneration: 0, agentHistorySyncGeneration: new Map(), @@ -1386,6 +1400,24 @@ export const useSessionStore = create()( }); }, + recordAgentTimelineRewind: (serverId, agentId, rewind) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + const current = session.agentTimelineLatestRewind.get(agentId); + if (current && current.epoch === rewind.epoch && current.seq >= rewind.seq) return prev; + const next = new Map(session.agentTimelineLatestRewind); + next.set(agentId, rewind); + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, agentTimelineLatestRewind: next }, + }, + }; + }); + }, + setAgentTimelineOlderFetchInFlight: (serverId, state) => { set((prev) => { const session = prev.sessions[serverId]; diff --git a/packages/app/src/timeline/replica.test.ts b/packages/app/src/timeline/replica.test.ts index d0b4b78777..ec482ef811 100644 --- a/packages/app/src/timeline/replica.test.ts +++ b/packages/app/src/timeline/replica.test.ts @@ -617,6 +617,10 @@ describe("conversation rewind through the owner", () => { return texts([...tail, ...(session?.agentStreamHead.get(AGENT_ID) ?? [])]); } + function latestRewind() { + return useSessionStore.getState().sessions[SERVER_ID]?.agentTimelineLatestRewind.get(AGENT_ID); + } + it("retires B from the viewed and persisted timeline without refetching", () => { useSessionStore.getState().initializeSession(SERVER_ID, null); const commits: CachedTimeline[] = []; @@ -658,6 +662,7 @@ describe("conversation rewind through the owner", () => { expect(viewed()).toEqual(["A", "reply A", "C"]); expect(texts(commits.at(-1)?.items ?? [])).toEqual(["A", "reply A", "C"]); expect(commits.at(-1)?.range).toMatchObject({ epoch: "epoch-1", endSeq: 6 }); + expect(latestRewind()).toEqual({ epoch: "epoch-1", seq: 5 }); owner.dispose(); }); @@ -679,6 +684,7 @@ describe("conversation rewind through the owner", () => { }); expect(viewed()).toEqual(["A", "reply A", "C"]); + expect(latestRewind()).toEqual({ epoch: "epoch-1", seq: 5 }); owner.dispose(); }); diff --git a/packages/app/src/timeline/session-stream-reducers.ts b/packages/app/src/timeline/session-stream-reducers.ts index 45eacf274c..6da93d4f39 100644 --- a/packages/app/src/timeline/session-stream-reducers.ts +++ b/packages/app/src/timeline/session-stream-reducers.ts @@ -2038,6 +2038,10 @@ export function createSessionAgentStreamReducerQueue( }; }, commit: (agentId, result, events) => { + for (const { event, seq, epoch } of events) { + if (event.type !== "rewind" || seq === undefined || epoch === undefined) continue; + useSessionStore.getState().recordAgentTimelineRewind(serverId, agentId, { epoch, seq }); + } if ( result.changedTail || result.changedHead || diff --git a/packages/app/src/timeline/viewed-timeline-sync.ts b/packages/app/src/timeline/viewed-timeline-sync.ts index ab8fe896ca..7011c5efd3 100644 --- a/packages/app/src/timeline/viewed-timeline-sync.ts +++ b/packages/app/src/timeline/viewed-timeline-sync.ts @@ -236,6 +236,13 @@ function commitProcessedTimeline(input: { const agentId = payload.agentId; const store = useSessionStore.getState(); if (result.commit !== "discard") { + const latestRewindSeq = Math.max(...(payload.retirements ?? []).map(({ seq }) => seq)); + if (Number.isFinite(latestRewindSeq)) { + store.recordAgentTimelineRewind(serverId, agentId, { + epoch: payload.epoch, + seq: latestRewindSeq, + }); + } store.applyAgentTimelineResponseState(serverId, agentId, { items: result.tail, head: result.head,