fix: Restore user input when the user cancels response (#5601)

Co-authored-by: Shi Shu <shii@google.com>
Co-authored-by: Jacob Richman <jacob314@gmail.com>
This commit is contained in:
shishu314 2025-08-06 15:19:10 -04:00 committed by GitHub
parent 6133bea388
commit 1f0ad86544
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 87 additions and 10 deletions

View File

@ -486,6 +486,24 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
setGeminiMdFileCount,
);
const buffer = useTextBuffer({
initialText: '',
viewport: { height: 10, width: inputWidth },
stdin,
setRawMode,
isValidPath,
shellModeActive,
});
const [userMessages, setUserMessages] = useState<string[]>([]);
const handleUserCancel = useCallback(() => {
const lastUserMessage = userMessages.at(-1);
if (lastUserMessage) {
buffer.setText(lastUserMessage);
}
}, [buffer, userMessages]);
const {
streamingState,
submitQuery,
@ -506,6 +524,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
modelSwitchedFromQuotaError,
setModelSwitchedFromQuotaError,
refreshStatic,
handleUserCancel,
);
// Input handling
@ -519,15 +538,6 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
[submitQuery],
);
const buffer = useTextBuffer({
initialText: '',
viewport: { height: 10, width: inputWidth },
stdin,
setRawMode,
isValidPath,
shellModeActive,
});
const { handleInput: vimHandleInput } = useVim(buffer, handleFinalSubmit);
const pendingHistoryItems = [...pendingSlashCommandHistoryItems];
pendingHistoryItems.push(...pendingGeminiHistoryItems);
@ -607,7 +617,6 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
}, [config, config.getGeminiMdFileCount]);
const logger = useLogger();
const [userMessages, setUserMessages] = useState<string[]>([]);
useEffect(() => {
const fetchUserMessages = async () => {

View File

@ -406,6 +406,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
);
},
{
@ -560,6 +562,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -633,6 +637,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -737,6 +743,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -843,6 +851,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -943,6 +953,44 @@ describe('useGeminiStream', () => {
expect(result.current.streamingState).toBe(StreamingState.Idle);
});
it('should call onCancelSubmit handler when escape is pressed', async () => {
const cancelSubmitSpy = vi.fn();
const mockStream = (async function* () {
yield { type: 'content', value: 'Part 1' };
// Keep the stream open
await new Promise(() => {});
})();
mockSendMessageStream.mockReturnValue(mockStream);
const { result } = renderHook(() =>
useGeminiStream(
mockConfig.getGeminiClient(),
[],
mockAddItem,
mockConfig,
mockOnDebugMessage,
mockHandleSlashCommand,
false,
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
false,
() => {},
() => {},
cancelSubmitSpy,
),
);
// Start a query
await act(async () => {
result.current.submitQuery('test query');
});
simulateEscapeKeyPress();
expect(cancelSubmitSpy).toHaveBeenCalled();
});
it('should not do anything if escape is pressed when not responding', () => {
const { result } = renderTestHook();
@ -1202,6 +1250,8 @@ describe('useGeminiStream', () => {
mockPerformMemoryRefresh,
false,
() => {},
() => {},
() => {},
),
);
@ -1253,6 +1303,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -1301,6 +1353,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -1347,6 +1401,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -1394,6 +1450,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -1481,6 +1539,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -1535,6 +1595,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -1611,6 +1673,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);
@ -1663,6 +1727,8 @@ describe('useGeminiStream', () => {
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
),
);

View File

@ -94,6 +94,7 @@ export const useGeminiStream = (
modelSwitchedFromQuotaError: boolean,
setModelSwitchedFromQuotaError: React.Dispatch<React.SetStateAction<boolean>>,
onEditorClose: () => void,
onCancelSubmit: () => void,
) => {
const [initError, setInitError] = useState<string | null>(null);
const abortControllerRef = useRef<AbortController | null>(null);
@ -200,6 +201,7 @@ export const useGeminiStream = (
Date.now(),
);
setPendingHistoryItem(null);
onCancelSubmit();
setIsResponding(false);
}
});