fix(ui): Resolve race condition in prompt restoration on context overflow #13460
Conversation
Summary of ChangesHello @SandyTao520, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a critical race condition within the UI's prompt restoration feature. When a user's input exceeds the context window, the system attempts to restore the prompt to the input buffer. Previously, this process could fail or restore an incorrect, outdated message due to the asynchronous nature of state updates. The implemented solution introduces a robust state-based synchronization mechanism, ensuring that the prompt is only restored once the underlying message history is consistent, thereby guaranteeing the correct and most recent prompt is always presented to the user. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request aims to fix a race condition in prompt restoration after a context window overflow. The proposed solution, while addressing the original issue, introduces a new potential race condition. My review includes a critical comment on the implementation in AppContainer.tsx, detailing this new issue and proposing a simpler, more robust alternative. The changes to the test file AppContainer.test.tsx are appropriate for testing the asynchronous nature of the fix.
| useEffect(() => { | ||
| if (pendingPromptRestore) { | ||
| // Check if userMessages is up to date with historyManager | ||
| const lastHistoryItem = historyManager.history.findLast( | ||
| (item) => item.type === 'user' && item.text, | ||
| ); | ||
| const lastUserMessage = userMessages.at(-1); | ||
|
|
||
| if ( | ||
| lastHistoryItem && | ||
| lastHistoryItem.text === lastUserMessage && | ||
| lastUserMessage | ||
| ) { | ||
| buffer.setText(lastUserMessage); | ||
| setPendingPromptRestore(false); | ||
| } | ||
| } | ||
| }, [ | ||
| pendingPromptRestore, | ||
| userMessages, | ||
| historyManager.history, | ||
| buffer, | ||
| setPendingPromptRestore, | ||
| ]); |
There was a problem hiding this comment.
While this useEffect approach does fix the original race condition, it introduces a new, more subtle one. If a user submits a second prompt very quickly after the one that caused the context overflow, this logic could restore the second prompt instead of the one that overflowed. This happens because the effect always looks for the latest user message in historyManager.history and userMessages to be in sync, but it doesn't track which message was supposed to be restored.
A simpler and more robust solution would be to avoid the useEffect and the pendingPromptRestore state altogether. You can directly access the correct history in the onCancelSubmit handler.
I suggest the following changes:
- Remove the
pendingPromptRestorestate. - Remove this
useEffectblock. - Update
cancelHandlerRef.currentto get the last message directly fromhistoryManager.historyand addhistoryManager.historyto itsuseCallbackdependency array.
Here's how cancelHandlerRef.current could look:
cancelHandlerRef.current = useCallback(
(shouldRestorePrompt: boolean = true) => {
const pendingHistoryItems = [
...pendingSlashCommandHistoryItems,
...pendingGeminiHistoryItems,
];
if (isToolExecuting(pendingHistoryItems)) {
buffer.setText(''); // Just clear the prompt
return;
}
let textToSet = '';
if (shouldRestorePrompt) {
const lastHistoryItem = historyManager.history.findLast(
(item) => item.type === 'user' && item.text,
);
textToSet = lastHistoryItem?.text || '';
}
const queuedText = getQueuedMessagesText();
if (queuedText) {
textToSet = textToSet ? `${textToSet}\n\n${queuedText}` : queuedText;
clearQueue();
}
if (textToSet || !shouldRestorePrompt) {
buffer.setText(textToSet);
}
},
[
buffer,
getQueuedMessagesText,
clearQueue,
pendingSlashCommandHistoryItems,
pendingGeminiHistoryItems,
historyManager.history, // Add this dependency
],
);This approach is more direct, less complex, and avoids the potential for new race conditions.
|
Size Change: +671 B (0%) Total Size: 21.1 MB ℹ️ View Unchanged
|
1 participant
TLDR
This PR fixes a race condition where restoring the user's prompt after a 'Context Window Overflow' event would sometimes fail or restore an outdated message.
Dive Deeper
When a user submits a prompt that exceeds the context window, the system triggers a cancellation and attempts to restore the prompt to the input buffer. Previously, this logic relied on
userMessagesstate, which is updated asynchronously via auseEffecthook.The
ContextWindowWillOverflowevent fires almost immediately after submission, often before theuserMessagesstate has had a chance to synchronize with the new history item. As a result, the restoration logic would read stale state, effectively 'restoring' the previous prompt instead of the one just submitted.This fix introduces a state-based synchronization mechanism:
pendingPromptRestoreflags when a restoration is requested.useEffecthook watches this flag and compares the latestuserMessagesentry with the actual source of truth (historyManager.history).This approach eliminates the race condition by deferring the UI update until the data is consistent.
Reviewer Test Plan
To verify this fix:
Testing Matrix
Linked issues / bugs
Resolves a race condition in prompt restoration logic.