diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index 66396c36..f2dcc79e 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -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([]); + + 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([]); useEffect(() => { const fetchUserMessages = async () => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index dd2428bb..751b869e 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -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, () => {}, + () => {}, + () => {}, ), ); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 63ba961f..58bec431 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -94,6 +94,7 @@ export const useGeminiStream = ( modelSwitchedFromQuotaError: boolean, setModelSwitchedFromQuotaError: React.Dispatch>, onEditorClose: () => void, + onCancelSubmit: () => void, ) => { const [initError, setInitError] = useState(null); const abortControllerRef = useRef(null); @@ -200,6 +201,7 @@ export const useGeminiStream = ( Date.now(), ); setPendingHistoryItem(null); + onCancelSubmit(); setIsResponding(false); } });