fix(copilot/frontend): guard mid-turn poll against session switches

Mirror the request-time-sessionId pattern from usePeekOnBoundary into
useMidTurnDrainPromotion: capture sessionId at request time, compare to
a live ref on resolve, bail if the user switched sessions while the GET
was in flight. Without this, a slow peek for session A could promote
chips into session B's message list after a switch.

Cancellation-flag was tried first but is too broad — this effect re-runs
on every chip-append, which would wrongly invalidate an in-flight poll
for the same session. The sessionId comparison only invalidates on
actual session changes, preserving the chip-append-during-poll race
guarantee.

Added a regression test that holds a peek in flight, switches sessions
mid-resolve, and asserts no promotion fires on the old session's
setMessages.
This commit is contained in:
majdyz
2026-04-30 11:47:13 +07:00
parent e59fe5af76
commit 8e13d4cb27
2 changed files with 91 additions and 5 deletions

View File

@@ -357,6 +357,65 @@ describe("useCopilotPendingChips", () => {
expect(result.current.queuedMessages).toEqual([]);
});
it("mid-turn poll: a peek that resolves after the user switches sessions does not promote into the new session", async () => {
vi.useFakeTimers();
const setMessagesA = vi.fn();
const setMessagesB = vi.fn();
type Props = { sessionId: string; setMessages: typeof setMessagesA };
const { result, rerender } = renderHook<
ReturnType<typeof useCopilotPendingChips>,
Props
>(
({ sessionId, setMessages }) =>
useCopilotPendingChips({
sessionId,
status: "streaming",
messages: [user("u1"), assistant("a1")],
setMessages,
}),
{
initialProps: { sessionId: "sess-A", setMessages: setMessagesA },
},
);
act(() => {
result.current.appendChip("chipA");
});
// Poll fires for sess-A but resolves AFTER we've switched to sess-B.
let resolvePeek!: (value: unknown) => void;
peekMock.mockImplementationOnce(
() =>
new Promise((resolve) => {
resolvePeek = resolve;
}),
);
await act(async () => {
vi.advanceTimersByTime(2_000);
});
// User switches sessions.
rerender({ sessionId: "sess-B", setMessages: setMessagesB });
// Old poll resolves now — backend says count=0 — but we've already
// moved on to sess-B. The session-id guard must short-circuit so
// setMessagesA is never called for the promotion path.
setMessagesA.mockClear();
await act(async () => {
resolvePeek({ status: 200, data: { count: 0, messages: [] } });
await Promise.resolve();
await Promise.resolve();
});
const promotedCallA = setMessagesA.mock.calls.find(([arg]) => {
if (typeof arg !== "function") return false;
const updated = (arg as (p: UIMessage[]) => UIMessage[])([]);
return updated.some((m) => m.id.startsWith("promoted-"));
});
expect(promotedCallA).toBeUndefined();
});
it("mid-turn poll: chip appended during in-flight poll survives the drain", async () => {
vi.useFakeTimers();
const setMessages = vi.fn();

View File

@@ -125,10 +125,13 @@ function usePeekOnBoundary({
// auto-continue effect's duplicate turn-start peek.
if (!sessionChanged && !isIdle && !turnStarting) return;
// Capture the sessionId at request time so a response that resolves
// after the user switched sessions can't bleed old-session chips into
// the new session. Without this guard, a slow peek for session A
// could overwrite chips just restored for session B.
// Capture sessionId at request time and compare to the live ref on
// resolve. Without this, a peek that resolves after the user
// switched sessions could bleed old-session chips into the new
// session. We deliberately don't use a per-effect cancelled flag
// here because this effect re-runs on every status change too, and
// we don't want chip-appends in another effect to invalidate this
// peek's result.
const requestSessionId = sessionId;
void getV2GetPendingMessages(sessionId).then((res) => {
if (prevSessionIdRef.current !== requestSessionId) return;
@@ -276,13 +279,32 @@ function useMidTurnDrainPromotion({
setMessages: (updater: (prev: UIMessage[]) => UIMessage[]) => void;
setChips: (updater: ChipUpdater) => void;
}) {
// Live ref tracks the latest sessionId so a poll captured at request
// time can detect a session switch on resolve. A cancellation flag
// would also fire on chip-append (this effect re-runs on every chip
// change), wrongly aborting an in-flight poll for the same session —
// the sessionId comparison only invalidates on actual session changes.
const latestSessionIdRef = useRef<string | null>(sessionId);
useEffect(() => {
latestSessionIdRef.current = sessionId;
}, [sessionId]);
useEffect(() => {
if (!sessionId) return;
const isActive = status === "streaming" || status === "submitted";
if (!isActive || chips.length === 0) return;
const requestSessionId = sessionId;
const isCurrentSession = () =>
latestSessionIdRef.current === requestSessionId;
const interval = setInterval(() => {
void pollBackendAndPromote(sessionId, chips, setMessages, setChips);
void pollBackendAndPromote(
sessionId,
chips,
setMessages,
setChips,
isCurrentSession,
);
}, MID_TURN_POLL_MS);
return () => clearInterval(interval);
}, [sessionId, status, chips, setMessages, setChips]);
@@ -293,6 +315,7 @@ async function pollBackendAndPromote(
snapshotChips: Chip[],
setMessages: (updater: (prev: UIMessage[]) => UIMessage[]) => void,
setChips: (updater: ChipUpdater) => void,
isCurrentSession: () => boolean,
): Promise<void> {
let backendCount: number;
try {
@@ -302,6 +325,10 @@ async function pollBackendAndPromote(
} catch {
return; // harmless; next tick or hydration will reconcile
}
// Bail if the user switched sessions while the GET was in flight —
// promoting these chips to messages for a different session would
// leak old-session bubbles into the new session.
if (!isCurrentSession()) return;
if (snapshotChips.length === 0) return;
if (backendCount >= snapshotChips.length) return;