Skip to content

Commit f032e02

Browse files
authored
Merge pull request #3878 from bobleer/bob/fix-flow-chat-peer-races
fix(flow-chat): guard peer turn, cancel, queue and history state races
2 parents 878235c + 2af94de commit f032e02

17 files changed

Lines changed: 680 additions & 62 deletions

‎src/web-ui/src/flow_chat/components/modern/ModernFlowChatContainer.history-restore.test.tsx‎

Lines changed: 121 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,20 +14,25 @@ const fixture = vi.hoisted(() => ({
1414
restore: vi.fn(),
1515
restoreTail: vi.fn(),
1616
loadTurnWindow: vi.fn(),
17+
synchronize: vi.fn(),
18+
extend: vi.fn(),
1719
present: null as ((sessionId: string, presentation: SessionHistoryPresentation) => void) | null,
1820
}));
1921
vi.mock('../../store/FlowChatStore', () => ({ flowChatStore: {
2022
getState: () => ({ activeSessionId: fixture.session?.sessionId, sessions: new Map(fixture.session ? [[fixture.session.sessionId, fixture.session]] : []) }),
2123
retainSessionHistory: () => () => {},
2224
subscribeSelector: () => () => {},
23-
getSessionHistoryViewState: () => undefined,
25+
getSessionHistoryViewState: () => ({
26+
catalog: fixture.session?.turnCatalog ?? null, loadedRanges: [], activeRange: fixture.cached?.range ?? null,
27+
}),
28+
synchronizeSessionHistory: (...args: unknown[]) => fixture.synchronize(...args),
2429
reactivateSessionHistoryWindow: () => fixture.cached,
2530
restoreSessionHistoryWindow: (...args: unknown[]) => fixture.restore(...args),
2631
restoreSessionTailPresentation: (...args: unknown[]) => fixture.restoreTail(...args),
2732
hasPendingSessionHistoryCompletion: () => false,
2833
hasDeferredSessionHistoryProjection: () => false,
2934
releaseSessionHistoryCompletionAfterInitialPaint: () => false,
30-
extendSessionHistoryWindow: () => null,
35+
extendSessionHistoryWindow: (...args: unknown[]) => fixture.extend(...args),
3136
loadSessionTurnWindow: (...args: unknown[]) => fixture.loadTurnWindow(...args),
3237
loadPreviousSessionTurnWindow: (...args: unknown[]) => fixture.loadTurnWindow(...args),
3338
activateSessionHistoryWindow: () => fixture.cached,
@@ -45,6 +50,7 @@ vi.mock('./VirtualMessageList', () => ({ VirtualMessageList: React.forwardRef((p
4550
fixture.listProps = props;
4651
React.useImperativeHandle(ref, () => ({
4752
scrollToTurnEnd: () => true,
53+
scrollToLatestEndPosition: () => {},
4854
captureViewportSnapshot: () => null,
4955
clearSearchMatch: () => {},
5056
prepareTurnNavigation: () => 'pending',
@@ -103,6 +109,14 @@ function session(totalTurnCount = 8, historyState: Session['historyState'] = 're
103109
}],
104110
} as Session;
105111
}
112+
function sessionWithCatalog(count: number, revision: string): Session {
113+
return { ...session(count), turnCatalog: {
114+
schemaVersion: 1, sessionId: 'session-a', revision, totalTurnCount: count,
115+
complete: true, entries: Array.from({ length: count }, (_, ordinal) => ({
116+
ordinal, storageTurnIndex: ordinal, turnId: `turn-${ordinal}`, previewTruncated: false,
117+
})),
118+
} };
119+
}
106120
function presentation(startOrdinal = 4, endOrdinalExclusive = 8): SessionHistoryPresentation {
107121
return {
108122
range: { startOrdinal, endOrdinalExclusive, targetTurnId: `turn-${startOrdinal}`, mode: 'history-window' },
@@ -151,6 +165,8 @@ describe('FlowChat reader restoration across history rewrites', () => {
151165
fixture.session = session();
152166
fixture.cached = presentation();
153167
fixture.restore.mockResolvedValue(null);
168+
fixture.extend.mockReturnValue(null);
169+
fixture.synchronize.mockResolvedValue(undefined);
154170
fixture.listProps = null;
155171
fixture.viewportSettled = true;
156172
fixture.present = null;
@@ -223,6 +239,109 @@ describe('FlowChat reader restoration across history rewrites', () => {
223239
expect(host.querySelector('[data-turn-id="turn-2"]')).not.toBeNull();
224240
});
225241

242+
it.each([['local', 8], ['local', 9], ['ssh', 8], ['ssh', 9], ['peer-device', 8], ['peer-device', 9]] as const)(
243+
'reacquires the reader window when the %s catalog changes after a turn (total: %d)', async (location, count) => {
244+
if (location === 'peer-device') activateSurface('peer-device');
245+
const withCatalog = (total: number, revision: string): Session => ({
246+
...sessionWithCatalog(total, revision),
247+
...(location === 'ssh' ? { remoteConnectionId: 'ssh-target', remoteSshHost: 'build-host' } : {}),
248+
});
249+
fixture.session = withCatalog(8, 'before-turn');
250+
const window = presentation(3, 7);
251+
fixture.cached = window;
252+
stageBookmark(window);
253+
await render();
254+
expectPresented();
255+
// Authoritative completion changes the catalog and invalidates the
256+
// store's cached windows, without changing Session id or ready status.
257+
fixture.cached = null;
258+
const first = deferred<SessionHistoryPresentation | null>();
259+
fixture.restore.mockReturnValueOnce(first.promise);
260+
fixture.session = withCatalog(count, 'after-turn');
261+
await render();
262+
expect(fixture.restore).toHaveBeenCalledWith('session-a', window.range, expect.any(Function));
263+
const firstOwner = fixture.restore.mock.calls[0][2] as () => boolean;
264+
const second = deferred<SessionHistoryPresentation | null>();
265+
fixture.restore.mockReturnValueOnce(second.promise);
266+
fixture.session = withCatalog(count, 'corrected-turn');
267+
await render();
268+
expect(fixture.restore).toHaveBeenCalledTimes(2);
269+
expect(firstOwner()).toBe(false);
270+
await act(async () => first.reject(new Error('Retired history read failed')));
271+
await act(async () => second.resolve(window));
272+
expectPresented();
273+
expect(messages().dataset.presentationMode).toBe('history-window');
274+
expect(host.querySelector('[data-turn-id="turn-3"]')).not.toBeNull();
275+
},
276+
);
277+
278+
it('drops a retained live-tail history projection when its catalog changes', async () => {
279+
fixture.session = sessionWithCatalog(8, 'before-turn');
280+
fixture.cached = presentation(0, 8);
281+
stageBookmark(fixture.cached);
282+
await render();
283+
await act(async () => { fixture.listProps!.onRequestJumpToLatest!(); });
284+
expect(host.querySelector('[data-turn-id="turn-0"]')).not.toBeNull();
285+
fixture.cached = null;
286+
fixture.session = sessionWithCatalog(8, 'after-turn');
287+
await render();
288+
expectPresented();
289+
expect(messages().dataset.presentationMode).toBe('tail');
290+
expect(host.querySelector('[data-turn-id="turn-0"]')).toBeNull();
291+
expect(fixture.restore).not.toHaveBeenCalled();
292+
});
293+
294+
it.each([false, true])('retries an orphaned reader after a failed read (explicit sync: %s)', async synchronizeHistory => {
295+
const window = presentation(3, 7);
296+
fixture.cached = window;
297+
stageBookmark(window);
298+
await render();
299+
fixture.cached = null;
300+
fixture.restore.mockImplementation(async () => {
301+
fixture.cached = window;
302+
return window;
303+
});
304+
fixture.restore.mockRejectedValueOnce(new Error('History read timed out'));
305+
fixture.loadTurnWindow.mockResolvedValue({ status: 'ready', isCurrent: true });
306+
fixture.extend.mockImplementation(() => fixture.cached ? presentation(0, 7) : null);
307+
let failed: unknown;
308+
await act(async () => {
309+
failed = await fixture.listProps!.onHistoryWindowBoundaryIntent!('before', { synchronizeHistory });
310+
});
311+
expect(failed).toBe('not-ready');
312+
expect(fixture.loadTurnWindow).not.toHaveBeenCalled();
313+
expect(host.querySelector('[data-turn-id="turn-3"]')).not.toBeNull();
314+
let outcome: unknown;
315+
await act(async () => {
316+
outcome = await fixture.listProps!.onHistoryWindowBoundaryIntent!('before', { synchronizeHistory });
317+
});
318+
expect(outcome).toBe('applied');
319+
expect(fixture.restore).toHaveBeenCalledWith('session-a', window.range, expect.any(Function));
320+
expect(fixture.restore).toHaveBeenCalledTimes(2);
321+
expect(fixture.synchronize).toHaveBeenCalledTimes(synchronizeHistory ? 2 : 0);
322+
expect(host.querySelector('[data-turn-id="turn-0"]')).not.toBeNull();
323+
});
324+
325+
it('rejects a late orphan recovery after changing devices', async () => {
326+
const window = presentation(3, 7);
327+
fixture.cached = window;
328+
stageBookmark(window);
329+
await render();
330+
fixture.cached = null;
331+
const read = deferred<SessionHistoryPresentation | null>();
332+
fixture.restore.mockReturnValueOnce(read.promise);
333+
let outcome!: Promise<unknown>;
334+
await act(async () => {
335+
outcome = fixture.listProps!.onHistoryWindowBoundaryIntent!('before', { synchronizeHistory: true });
336+
});
337+
expect(fixture.restore).toHaveBeenCalledTimes(1);
338+
activateSurface('peer-device');
339+
await act(async () => { read.resolve(window); await outcome; });
340+
expect(await outcome).toBe('cancelled');
341+
expect(fixture.loadTurnWindow).not.toHaveBeenCalled();
342+
expect(fixture.extend).not.toHaveBeenCalled();
343+
});
344+
226345
it('cancels a pending cold range when a ready history rewrite removes it', async () => {
227346
fixture.cached = null;
228347
stageBookmark();

‎src/web-ui/src/flow_chat/components/modern/ModernFlowChatContainer.tsx‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -777,7 +777,8 @@ export const ModernFlowChatContainer: React.FC<ModernFlowChatContainerProps> = (
777777
useLayoutEffect(() => {
778778
const sessionId = activeSession?.sessionId;
779779
const previousSessionId = activeSessionIdRef.current;
780-
if (previousSessionId && previousSessionId !== sessionId) {
780+
if (previousSessionId && (previousSessionId !== sessionId
781+
|| historyPresentationRef.current?.sessionId === previousSessionId)) {
781782
rememberSessionViewportState(previousSessionId, {
782783
historyPresentation: historyPresentationRef.current?.sessionId === previousSessionId
783784
? historyPresentationRef.current
@@ -851,6 +852,28 @@ export const ModernFlowChatContainer: React.FC<ModernFlowChatContainerProps> = (
851852
}
852853
}, [activeSession?.sessionId, activeSession?.historyState, isRememberedHistoryRangeInvalid, invalidateHistoryPresentationOwner, rememberSessionViewportState, updateViewportIntent, restoreRevision, historyRestoreRetry, surfaceScope]);
853854

855+
const readerCatalogRef = useRef<{ sessionId: string | undefined; revision: string | undefined } | null>(null);
856+
useEffect(() => {
857+
const sessionId = activeSession?.sessionId;
858+
const revision = activeSession?.turnCatalog?.revision;
859+
const previous = readerCatalogRef.current;
860+
readerCatalogRef.current = { sessionId, revision };
861+
if (!sessionId || previous?.sessionId !== sessionId || previous.revision === revision) return;
862+
const intent = viewportIntentRef.current;
863+
if (intent?.kind !== 'turn' || intent.sessionId !== sessionId || intent.source !== 'history-range') {
864+
if (historyPresentationRef.current?.sessionId === sessionId) {
865+
switchToLiveTailForSession(sessionId, { discardRecentHistory: true });
866+
}
867+
return;
868+
}
869+
// The store invalidates cached ranges when the authoritative catalog
870+
// changes. Retire the view's matching owner too, then reload its bookmark
871+
// through the normal history path. Keeping only the rendered copy leaves
872+
// every later extend/synchronize pointed at a window the store no longer owns.
873+
invalidateHistoryPresentationOwner();
874+
setHistoryRestoreRetry(retry => retry + 1);
875+
}, [activeSession?.sessionId, activeSession?.turnCatalog?.revision, invalidateHistoryPresentationOwner, switchToLiveTailForSession]);
876+
854877
useEffect(() => {
855878
const retainedSessionId = continuousProjectionSessionId;
856879
if (!retainedSessionId) {
@@ -1991,6 +2014,27 @@ export const ModernFlowChatContainer: React.FC<ModernFlowChatContainerProps> = (
19912014
return 'not-ready';
19922015
}
19932016
}
2017+
// A catalog refresh or cache eviction can retire the store's window
2018+
// while the reader still displays its copy. Reacquire that exact range
2019+
// before extending it; repeatedly fetching the adjacent page cannot
2020+
// restore a missing active window, even after an explicit source sync.
2021+
const activeRange = flowChatStore.getSessionHistoryViewState(sessionId)?.activeRange;
2022+
if (presentation && (!activeRange
2023+
|| activeRange.startOrdinal !== presentation.range.startOrdinal
2024+
|| activeRange.endOrdinalExclusive !== presentation.range.endOrdinalExclusive)) {
2025+
setHistoryBoundaryState(previous => ({ ...previous, [direction]: 'loading' }));
2026+
try {
2027+
const restored = await flowChatStore.restoreSessionHistoryWindow(sessionId, presentation.range, isCurrent);
2028+
if (!isCurrent()) return 'cancelled';
2029+
if (!restored) throw new Error('The reader history range is no longer available');
2030+
} catch (error) {
2031+
if (!isCurrent() || isSurfaceChangedError(error)) return 'cancelled';
2032+
recordHistoryPagingEvent(sessionId, 'outcome_not_ready', { direction, reason: 'reader-range-restore-failed' });
2033+
log.warn('Failed to reacquire reader history range', { sessionId, direction, error });
2034+
setHistoryBoundaryState(previous => ({ ...previous, [direction]: 'error' }));
2035+
return 'not-ready';
2036+
}
2037+
}
19942038
const session = flowChatStore.getState().sessions.get(sessionId);
19952039
if (scope.surfaceId !== 'local' && !presentation) {
19962040
if (direction === 'after' || session?.isPartial === false) {
@@ -2197,6 +2241,10 @@ export const ModernFlowChatContainer: React.FC<ModernFlowChatContainerProps> = (
21972241
direction,
21982242
reason: `load-${result.status}`,
21992243
targetOrdinal,
2244+
surfaceId: scope.surfaceId,
2245+
remoteWorkspace: Boolean(session?.remoteConnectionId || session?.remoteSshHost),
2246+
requestedCatalogRevision: session?.turnCatalog?.revision ?? null,
2247+
currentCatalogRevision: flowChatStore.getSessionHistoryViewState(sessionId)?.catalog?.revision ?? null,
22002248
});
22012249
setHistoryBoundaryState(previous => ({ ...previous, [direction]: 'error' }));
22022250
return 'not-ready';
@@ -2247,13 +2295,21 @@ export const ModernFlowChatContainer: React.FC<ModernFlowChatContainerProps> = (
22472295
? flowChatStore.extendSessionHistoryWindow(sessionId, direction)
22482296
: flowChatStore.activateSessionHistoryWindowFromTail(sessionId, targetOrdinal);
22492297
if (!nextPresentation) {
2298+
const currentHistoryView = flowChatStore.getSessionHistoryViewState(sessionId);
22502299
if (viewportPreparationStarted) {
22512300
options?.cancelViewportPresentationCommit?.();
22522301
}
22532302
recordHistoryPagingEvent(sessionId, 'outcome_not_ready', {
22542303
direction,
22552304
reason: 'no-next-presentation',
22562305
targetOrdinal,
2306+
surfaceId: scope.surfaceId,
2307+
remoteWorkspace: Boolean(session?.remoteConnectionId || session?.remoteSshHost),
2308+
currentCatalogRevision: currentHistoryView?.catalog?.revision ?? null,
2309+
renderedStartOrdinal: renderedRange.startOrdinal,
2310+
renderedEndOrdinalExclusive: renderedRange.endOrdinalExclusive,
2311+
activeStartOrdinal: currentHistoryView?.activeRange?.startOrdinal ?? null,
2312+
activeEndOrdinalExclusive: currentHistoryView?.activeRange?.endOrdinalExclusive ?? null,
22572313
});
22582314
setHistoryBoundaryState(previous => ({ ...previous, [direction]: 'error' }));
22592315
return 'not-ready';

‎src/web-ui/src/flow_chat/components/modern/UserMessageItem.test.tsx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@ const flowChatStoreMock = vi.hoisted(() => ({
150150
activeSessionId: null,
151151
})),
152152
loadSessionHistory: vi.fn(async () => undefined),
153+
subscribeSelector: vi.fn(() => () => {}),
153154
}));
154155

155156
vi.mock('../../store/FlowChatStore', () => ({

‎src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.test.ts‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,37 @@ describe('external BTW turn projection', () => {
123123
});
124124
});
125125

126+
describe('late dialog start events', () => {
127+
beforeEach(() => { resetFlowChatStore(); stateMachineManager.clear(); });
128+
afterEach(() => { resetFlowChatStore(); stateMachineManager.clear(); });
129+
130+
it.each(['completed', 'cancelled', 'error'] as const)(
131+
'does not restart a turn already projected as %s', status => {
132+
createSessionWithTurn({ ...createFinishingTurn(), status });
133+
stateMachineManager.getOrCreate('session-1');
134+
__test_only__.handleDialogTurnStarted(createFlowChatContext(), {
135+
sessionId: 'session-1', turnId: 'turn-1', turnIndex: 0, userInput: 'Delayed start',
136+
});
137+
expect(stateMachineManager.getCurrentState('session-1')).toBe(SessionExecutionState.IDLE);
138+
expect(FlowChatStore.getInstance().getState().sessions.get('session-1')?.dialogTurns[0].status).toBe(status);
139+
},
140+
);
141+
142+
it('does not let an older admitted turn take ownership from the current turn', async () => {
143+
createSessionWithTurn({ ...createFinishingTurn(), status: 'processing', backendTurnIndex: 1 });
144+
FlowChatStore.getInstance().addDialogTurn('session-1', {
145+
...createFinishingTurn(), id: 'turn-2', status: 'processing', backendTurnIndex: 2,
146+
});
147+
await stateMachineManager.transition('session-1', SessionExecutionEvent.START, {
148+
taskId: 'session-1', dialogTurnId: 'turn-2',
149+
});
150+
__test_only__.handleDialogTurnStarted(createFlowChatContext(), {
151+
sessionId: 'session-1', turnId: 'turn-1', turnIndex: 1, userInput: 'Delayed start',
152+
});
153+
expect(stateMachineManager.get('session-1')?.getContext().currentDialogTurnId).toBe('turn-2');
154+
});
155+
});
156+
126157
describe('Claw bootstrap cancellation', () => {
127158
beforeEach(() => {
128159
resetFlowChatStore();

‎src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1896,12 +1896,25 @@ function handleDialogTurnStarted(context: FlowChatContext, event: any): void {
18961896
return;
18971897
}
18981898

1899+
// Canonical history and direct control events travel independently. A delayed
1900+
// start must not resurrect a settled turn or supersede a newer admitted turn.
1901+
// Optimistic queued turns have no backend index and do not block admission.
1902+
const store = FlowChatStore.getInstance();
1903+
const knownTurns = store.getState().sessions.get(sessionId)?.dialogTurns ?? [];
1904+
const knownTurn = knownTurns.find(turn => turn.id === turnId);
1905+
const admittedIndex = typeof turnIndex === 'number'
1906+
? turnIndex : knownTurn?.backendTurnIndex ?? knownTurn?.storageTurnIndex;
1907+
if (knownTurn && ['completed', 'cancelled', 'error'].includes(knownTurn.status)
1908+
|| admittedIndex !== undefined && knownTurns.some(turn =>
1909+
(turn.backendTurnIndex ?? turn.storageTurnIndex ?? -1) > admittedIndex)) {
1910+
log.debug('Dropped stale DialogTurnStarted', { sessionId, turnId, admittedIndex });
1911+
return;
1912+
}
1913+
18991914
clearPromptSuggestion(sessionId);
19001915
finalizePendingTurnCompletionNow(context, sessionId);
19011916
clearPendingTurnCompletion(context, sessionId, turnId);
19021917

1903-
const store = FlowChatStore.getInstance();
1904-
19051918
// Clean up temp image analysis turn if one exists for this session
19061919
const tempTurnId = pendingImageAnalysisTurns.get(sessionId);
19071920
const hadTempTurn = !!tempTurnId;

0 commit comments

Comments
 (0)