@google/gemini-cli 0.12.0-nightly.20251023.c4c0c0d1 → 0.12.0-preview.0

package/dist/src/ui/components/InputPrompt.test.js

@@ -5,7 +5,7 @@ import { jsx as _jsx } from "react/jsx-runtime";
  * SPDX-License-Identifier: Apache-2.0
  */
 import { renderWithProviders } from '../../test-utils/render.js';
-import { waitFor, act } from '@testing-library/react';
+import { act } from 'react';
 import { InputPrompt } from './InputPrompt.js';
 import { ApprovalMode } from '@google/gemini-cli-core';
 import * as path from 'node:path';
@@ -207,63 +207,68 @@ describe('InputPrompt', () => {
             streamingState: StreamingState.Idle,
         };
     });
-    const wait = (ms = 50) => new Promise((resolve) => setTimeout(resolve, ms));
     it('should call shellHistory.getPreviousCommand on up arrow in shell mode', async () => {
         props.shellModeActive = true;
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\u001B[A');
-        await wait();
-        expect(mockShellHistory.getPreviousCommand).toHaveBeenCalled();
+        await act(async () => {
+            stdin.write('\u001B[A');
+        });
+        await vi.waitFor(() => expect(mockShellHistory.getPreviousCommand).toHaveBeenCalled());
         unmount();
     });
     it('should call shellHistory.getNextCommand on down arrow in shell mode', async () => {
         props.shellModeActive = true;
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\u001B[B');
-        await wait();
-        expect(mockShellHistory.getNextCommand).toHaveBeenCalled();
+        await act(async () => {
+            stdin.write('\u001B[B');
+        });
+        await vi.waitFor(() => expect(mockShellHistory.getNextCommand).toHaveBeenCalled());
         unmount();
     });
     it('should set the buffer text when a shell history command is retrieved', async () => {
         props.shellModeActive = true;
         vi.mocked(mockShellHistory.getPreviousCommand).mockReturnValue('previous command');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\u001B[A');
-        await wait();
-        expect(mockShellHistory.getPreviousCommand).toHaveBeenCalled();
-        expect(props.buffer.setText).toHaveBeenCalledWith('previous command');
+        await act(async () => {
+            stdin.write('\u001B[A');
+        });
+        await vi.waitFor(() => {
+            expect(mockShellHistory.getPreviousCommand).toHaveBeenCalled();
+            expect(props.buffer.setText).toHaveBeenCalledWith('previous command');
+        });
         unmount();
     });
     it('should call shellHistory.addCommandToHistory on submit in shell mode', async () => {
         props.shellModeActive = true;
         props.buffer.setText('ls -l');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\r');
-        await wait();
-        expect(mockShellHistory.addCommandToHistory).toHaveBeenCalledWith('ls -l');
-        expect(props.onSubmit).toHaveBeenCalledWith('ls -l');
+        await act(async () => {
+            stdin.write('\r');
+        });
+        await vi.waitFor(() => {
+            expect(mockShellHistory.addCommandToHistory).toHaveBeenCalledWith('ls -l');
+            expect(props.onSubmit).toHaveBeenCalledWith('ls -l');
+        });
         unmount();
     });
     it('should NOT call shell history methods when not in shell mode', async () => {
         props.buffer.setText('some text');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\u001B[A'); // Up arrow
-        await wait();
-        stdin.write('\u001B[B'); // Down arrow
-        await wait();
-        stdin.write('\r'); // Enter
-        await wait();
+        await act(async () => {
+            stdin.write('\u001B[A'); // Up arrow
+        });
+        await vi.waitFor(() => expect(mockInputHistory.navigateUp).toHaveBeenCalled());
+        await act(async () => {
+            stdin.write('\u001B[B'); // Down arrow
+        });
+        await vi.waitFor(() => expect(mockInputHistory.navigateDown).toHaveBeenCalled());
+        await act(async () => {
+            stdin.write('\r'); // Enter
+        });
+        await vi.waitFor(() => expect(props.onSubmit).toHaveBeenCalledWith('some text'));
         expect(mockShellHistory.getPreviousCommand).not.toHaveBeenCalled();
         expect(mockShellHistory.getNextCommand).not.toHaveBeenCalled();
         expect(mockShellHistory.addCommandToHistory).not.toHaveBeenCalled();
-        expect(mockInputHistory.navigateUp).toHaveBeenCalled();
-        expect(mockInputHistory.navigateDown).toHaveBeenCalled();
-        expect(props.onSubmit).toHaveBeenCalledWith('some text');
         unmount();
     });
     it('should call completion.navigateUp for both up arrow and Ctrl+P when suggestions are showing', async () => {
@@ -277,13 +282,15 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('/mem');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
         // Test up arrow
-        stdin.write('\u001B[A'); // Up arrow
-        await wait();
-        stdin.write('\u0010'); // Ctrl+P
-        await wait();
-        expect(mockCommandCompletion.navigateUp).toHaveBeenCalledTimes(2);
+        await act(async () => {
+            stdin.write('\u001B[A'); // Up arrow
+        });
+        await vi.waitFor(() => expect(mockCommandCompletion.navigateUp).toHaveBeenCalledTimes(1));
+        await act(async () => {
+            stdin.write('\u0010'); // Ctrl+P
+        });
+        await vi.waitFor(() => expect(mockCommandCompletion.navigateUp).toHaveBeenCalledTimes(2));
         expect(mockCommandCompletion.navigateDown).not.toHaveBeenCalled();
         unmount();
     });
@@ -298,13 +305,15 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('/mem');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
         // Test down arrow
-        stdin.write('\u001B[B'); // Down arrow
-        await wait();
-        stdin.write('\u000E'); // Ctrl+N
-        await wait();
-        expect(mockCommandCompletion.navigateDown).toHaveBeenCalledTimes(2);
+        await act(async () => {
+            stdin.write('\u001B[B'); // Down arrow
+        });
+        await vi.waitFor(() => expect(mockCommandCompletion.navigateDown).toHaveBeenCalledTimes(1));
+        await act(async () => {
+            stdin.write('\u000E'); // Ctrl+N
+        });
+        await vi.waitFor(() => expect(mockCommandCompletion.navigateDown).toHaveBeenCalledTimes(2));
         expect(mockCommandCompletion.navigateUp).not.toHaveBeenCalled();
         unmount();
     });
@@ -315,15 +324,22 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('some text');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\u001B[A'); // Up arrow
-        await wait();
-        stdin.write('\u001B[B'); // Down arrow
-        await wait();
-        stdin.write('\u0010'); // Ctrl+P
-        await wait();
-        stdin.write('\u000E'); // Ctrl+N
-        await wait();
+        await act(async () => {
+            stdin.write('\u001B[A'); // Up arrow
+        });
+        await vi.waitFor(() => expect(mockInputHistory.navigateUp).toHaveBeenCalled());
+        await act(async () => {
+            stdin.write('\u001B[B'); // Down arrow
+        });
+        await vi.waitFor(() => expect(mockInputHistory.navigateDown).toHaveBeenCalled());
+        await act(async () => {
+            stdin.write('\u0010'); // Ctrl+P
+        });
+        await vi.waitFor(() => { });
+        await act(async () => {
+            stdin.write('\u000E'); // Ctrl+N
+        });
+        await vi.waitFor(() => { });
         expect(mockCommandCompletion.navigateUp).not.toHaveBeenCalled();
         expect(mockCommandCompletion.navigateDown).not.toHaveBeenCalled();
         unmount();
@@ -338,23 +354,27 @@ describe('InputPrompt', () => {
             vi.mocked(clipboardUtils.clipboardHasImage).mockResolvedValue(true);
             vi.mocked(clipboardUtils.saveClipboardImage).mockResolvedValue('/test/.gemini-clipboard/clipboard-123.png');
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
             // Send Ctrl+V
-            stdin.write('\x16'); // Ctrl+V
-            await wait();
-            expect(clipboardUtils.clipboardHasImage).toHaveBeenCalled();
-            expect(clipboardUtils.saveClipboardImage).toHaveBeenCalledWith(props.config.getTargetDir());
-            expect(clipboardUtils.cleanupOldClipboardImages).toHaveBeenCalledWith(props.config.getTargetDir());
-            expect(mockBuffer.replaceRangeByOffset).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\x16'); // Ctrl+V
+            });
+            await vi.waitFor(() => {
+                expect(clipboardUtils.clipboardHasImage).toHaveBeenCalled();
+                expect(clipboardUtils.saveClipboardImage).toHaveBeenCalledWith(props.config.getTargetDir());
+                expect(clipboardUtils.cleanupOldClipboardImages).toHaveBeenCalledWith(props.config.getTargetDir());
+                expect(mockBuffer.replaceRangeByOffset).toHaveBeenCalled();
+            });
             unmount();
         });
         it('should not insert anything when clipboard has no image', async () => {
             vi.mocked(clipboardUtils.clipboardHasImage).mockResolvedValue(false);
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x16'); // Ctrl+V
-            await wait();
-            expect(clipboardUtils.clipboardHasImage).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\x16'); // Ctrl+V
+            });
+            await vi.waitFor(() => {
+                expect(clipboardUtils.clipboardHasImage).toHaveBeenCalled();
+            });
             expect(clipboardUtils.saveClipboardImage).not.toHaveBeenCalled();
             expect(mockBuffer.setText).not.toHaveBeenCalled();
             unmount();
@@ -363,10 +383,12 @@ describe('InputPrompt', () => {
             vi.mocked(clipboardUtils.clipboardHasImage).mockResolvedValue(true);
             vi.mocked(clipboardUtils.saveClipboardImage).mockResolvedValue(null);
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x16'); // Ctrl+V
-            await wait();
-            expect(clipboardUtils.saveClipboardImage).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\x16'); // Ctrl+V
+            });
+            await vi.waitFor(() => {
+                expect(clipboardUtils.saveClipboardImage).toHaveBeenCalled();
+            });
             expect(mockBuffer.setText).not.toHaveBeenCalled();
             unmount();
         });
@@ -380,11 +402,13 @@ describe('InputPrompt', () => {
             mockBuffer.lines = ['Hello world'];
             mockBuffer.replaceRangeByOffset = vi.fn();
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x16'); // Ctrl+V
-            await wait();
-            // Should insert at cursor position with spaces
-            expect(mockBuffer.replaceRangeByOffset).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\x16'); // Ctrl+V
+            });
+            await vi.waitFor(() => {
+                // Should insert at cursor position with spaces
+                expect(mockBuffer.replaceRangeByOffset).toHaveBeenCalled();
+            });
             // Get the actual call to see what path was used
             const actualCall = vi.mocked(mockBuffer.replaceRangeByOffset).mock
                 .calls[0];
@@ -399,10 +423,12 @@ describe('InputPrompt', () => {
                 .mockImplementation(() => { });
             vi.mocked(clipboardUtils.clipboardHasImage).mockRejectedValue(new Error('Clipboard error'));
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x16'); // Ctrl+V
-            await wait();
-            expect(consoleErrorSpy).toHaveBeenCalledWith('Error handling clipboard image:', expect.any(Error));
+            await act(async () => {
+                stdin.write('\x16'); // Ctrl+V
+            });
+            await vi.waitFor(() => {
+                expect(consoleErrorSpy).toHaveBeenCalledWith('Error handling clipboard image:', expect.any(Error));
+            });
             expect(mockBuffer.setText).not.toHaveBeenCalled();
             consoleErrorSpy.mockRestore();
             unmount();
@@ -418,10 +444,10 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('/mem');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\t'); // Press Tab
-        await wait();
-        expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0);
+        await act(async () => {
+            stdin.write('\t'); // Press Tab
+        });
+        await vi.waitFor(() => expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0));
         unmount();
     });
     it('should append a sub-command when the parent command is already complete', async () => {
@@ -437,10 +463,10 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('/memory ');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\t'); // Press Tab
-        await wait();
-        expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(1);
+        await act(async () => {
+            stdin.write('\t'); // Press Tab
+        });
+        await vi.waitFor(() => expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(1));
         unmount();
     });
     it('should handle the "backspace" edge case correctly', async () => {
@@ -457,11 +483,12 @@ describe('InputPrompt', () => {
         // The user has backspaced, so the query is now just '/memory'
         props.buffer.setText('/memory');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\t'); // Press Tab
-        await wait();
+        await act(async () => {
+            stdin.write('\t'); // Press Tab
+        });
+        await vi.waitFor(() => 
         // It should NOT become '/show'. It should correctly become '/memory show'.
-        expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0);
+        expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0));
         unmount();
     });
     it('should complete a partial argument for a command', async () => {
@@ -474,10 +501,10 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('/chat resume fi-');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\t'); // Press Tab
-        await wait();
-        expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0);
+        await act(async () => {
+            stdin.write('\t'); // Press Tab
+        });
+        await vi.waitFor(() => expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0));
         unmount();
     });
     it('should autocomplete on Enter when suggestions are active, without submitting', async () => {
@@ -489,11 +516,13 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('/mem');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\r');
-        await wait();
-        // The app should autocomplete the text, NOT submit.
-        expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0);
+        await act(async () => {
+            stdin.write('\r');
+        });
+        await vi.waitFor(() => {
+            // The app should autocomplete the text, NOT submit.
+            expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0);
+        });
         expect(props.onSubmit).not.toHaveBeenCalled();
         unmount();
     });
@@ -514,18 +543,19 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('/?');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\t'); // Press Tab for autocomplete
-        await wait();
-        expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0);
+        await act(async () => {
+            stdin.write('\t'); // Press Tab for autocomplete
+        });
+        await vi.waitFor(() => expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0));
         unmount();
     });
     it('should not submit on Enter when the buffer is empty or only contains whitespace', async () => {
         props.buffer.setText('   '); // Set buffer to whitespace
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\r'); // Press Enter
-        await wait();
+        await act(async () => {
+            stdin.write('\r'); // Press Enter
+        });
+        await vi.waitFor(() => { });
         expect(props.onSubmit).not.toHaveBeenCalled();
         unmount();
     });
@@ -537,10 +567,10 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('/clear');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\r');
-        await wait();
-        expect(props.onSubmit).toHaveBeenCalledWith('/clear');
+        await act(async () => {
+            stdin.write('\r');
+        });
+        await vi.waitFor(() => expect(props.onSubmit).toHaveBeenCalledWith('/clear'));
         unmount();
     });
     it('should submit directly on Enter when a complete leaf command is typed', async () => {
@@ -551,10 +581,10 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('/clear');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\r');
-        await wait();
-        expect(props.onSubmit).toHaveBeenCalledWith('/clear');
+        await act(async () => {
+            stdin.write('\r');
+        });
+        await vi.waitFor(() => expect(props.onSubmit).toHaveBeenCalledWith('/clear'));
         unmount();
     });
     it('should autocomplete an @-path on Enter without submitting', async () => {
@@ -566,10 +596,10 @@ describe('InputPrompt', () => {
         });
         props.buffer.setText('@src/components/');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\r');
-        await wait();
-        expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0);
+        await act(async () => {
+            stdin.write('\r');
+        });
+        await vi.waitFor(() => expect(mockCommandCompletion.handleAutocomplete).toHaveBeenCalledWith(0));
         expect(props.onSubmit).not.toHaveBeenCalled();
         unmount();
     });
@@ -579,31 +609,36 @@ describe('InputPrompt', () => {
         mockBuffer.cursor = [0, 11];
         mockBuffer.lines = ['first line\\'];
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\r');
-        await wait();
+        await act(async () => {
+            stdin.write('\r');
+        });
+        await vi.waitFor(() => {
+            expect(props.buffer.backspace).toHaveBeenCalled();
+            expect(props.buffer.newline).toHaveBeenCalled();
+        });
         expect(props.onSubmit).not.toHaveBeenCalled();
-        expect(props.buffer.backspace).toHaveBeenCalled();
-        expect(props.buffer.newline).toHaveBeenCalled();
         unmount();
     });
     it('should clear the buffer on Ctrl+C if it has text', async () => {
         props.buffer.setText('some text to clear');
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\x03'); // Ctrl+C character
-        await wait();
-        expect(props.buffer.setText).toHaveBeenCalledWith('');
-        expect(mockCommandCompletion.resetCompletionState).toHaveBeenCalled();
+        await act(async () => {
+            stdin.write('\x03'); // Ctrl+C character
+        });
+        await vi.waitFor(() => {
+            expect(props.buffer.setText).toHaveBeenCalledWith('');
+            expect(mockCommandCompletion.resetCompletionState).toHaveBeenCalled();
+        });
         expect(props.onSubmit).not.toHaveBeenCalled();
         unmount();
     });
     it('should NOT clear the buffer on Ctrl+C if it is empty', async () => {
         props.buffer.text = '';
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-        await wait();
-        stdin.write('\x03'); // Ctrl+C character
-        await wait();
+        await act(async () => {
+            stdin.write('\x03'); // Ctrl+C character
+        });
+        await vi.waitFor(() => { });
         expect(props.buffer.setText).not.toHaveBeenCalled();
         unmount();
     });
@@ -699,8 +734,9 @@ describe('InputPrompt', () => {
                     : [],
             });
             const { unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            expect(mockedUseCommandCompletion).toHaveBeenCalledWith(mockBuffer, ['/test/project/src'], path.join('test', 'project', 'src'), mockSlashCommands, mockCommandContext, false, false, expect.any(Object));
+            await vi.waitFor(() => {
+                expect(mockedUseCommandCompletion).toHaveBeenCalledWith(mockBuffer, ['/test/project/src'], path.join('test', 'project', 'src'), mockSlashCommands, mockCommandContext, false, false, expect.any(Object));
+            });
             unmount();
         });
     });
@@ -708,32 +744,38 @@ describe('InputPrompt', () => {
         it('should not call buffer.handleInput when vim mode is enabled and vim handles the input', async () => {
             props.vimHandleInput = vi.fn().mockReturnValue(true); // Mock that vim handled it.
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('i');
-            await wait();
-            expect(props.vimHandleInput).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('i');
+            });
+            await vi.waitFor(() => {
+                expect(props.vimHandleInput).toHaveBeenCalled();
+            });
             expect(mockBuffer.handleInput).not.toHaveBeenCalled();
             unmount();
         });
         it('should call buffer.handleInput when vim mode is enabled but vim does not handle the input', async () => {
             props.vimHandleInput = vi.fn().mockReturnValue(false); // Mock that vim did NOT handle it.
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('i');
-            await wait();
-            expect(props.vimHandleInput).toHaveBeenCalled();
-            expect(mockBuffer.handleInput).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('i');
+            });
+            await vi.waitFor(() => {
+                expect(props.vimHandleInput).toHaveBeenCalled();
+                expect(mockBuffer.handleInput).toHaveBeenCalled();
+            });
             unmount();
         });
         it('should call handleInput when vim mode is disabled', async () => {
             // Mock vimHandleInput to return false (vim didn't handle the input)
             props.vimHandleInput = vi.fn().mockReturnValue(false);
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('i');
-            await wait();
-            expect(props.vimHandleInput).toHaveBeenCalled();
-            expect(mockBuffer.handleInput).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('i');
+            });
+            await vi.waitFor(() => {
+                expect(props.vimHandleInput).toHaveBeenCalled();
+                expect(mockBuffer.handleInput).toHaveBeenCalled();
+            });
             unmount();
         });
     });
@@ -741,21 +783,24 @@ describe('InputPrompt', () => {
         it('should handle bracketed paste when not focused', async () => {
             props.focus = false;
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x1B[200~pasted text\x1B[201~');
-            await wait();
-            expect(mockBuffer.handleInput).toHaveBeenCalledWith(expect.objectContaining({
-                paste: true,
-                sequence: 'pasted text',
-            }));
+            await act(async () => {
+                stdin.write('\x1B[200~pasted text\x1B[201~');
+            });
+            await vi.waitFor(() => {
+                expect(mockBuffer.handleInput).toHaveBeenCalledWith(expect.objectContaining({
+                    paste: true,
+                    sequence: 'pasted text',
+                }));
+            });
             unmount();
         });
         it('should ignore regular keypresses when not focused', async () => {
             props.focus = false;
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('a');
-            await wait();
+            await act(async () => {
+                stdin.write('a');
+            });
+            await vi.waitFor(() => { });
             expect(mockBuffer.handleInput).not.toHaveBeenCalled();
             unmount();
         });
@@ -817,9 +862,10 @@ describe('InputPrompt', () => {
                 mockBuffer.viewportVisualLines = [text];
                 mockBuffer.visualCursor = visualCursor;
                 const { stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-                await wait();
-                const frame = stdout.lastFrame();
-                expect(frame).toContain(expected);
+                await vi.waitFor(() => {
+                    const frame = stdout.lastFrame();
+                    expect(frame).toContain(expected);
+                });
                 unmount();
             });
         });
@@ -863,9 +909,10 @@ describe('InputPrompt', () => {
                 mockBuffer.visualCursor = visualCursor;
                 mockBuffer.visualToLogicalMap = visualToLogicalMap;
                 const { stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-                await wait();
-                const frame = stdout.lastFrame();
-                expect(frame).toContain(expected);
+                await vi.waitFor(() => {
+                    const frame = stdout.lastFrame();
+                    expect(frame).toContain(expected);
+                });
                 unmount();
             });
             it('should display cursor on a blank line in a multiline block', async () => {
@@ -880,11 +927,12 @@ describe('InputPrompt', () => {
                     [2, 0],
                 ];
                 const { stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-                await wait();
-                const frame = stdout.lastFrame();
-                const lines = frame.split('\n');
-                // The line with the cursor should just be an inverted space inside the box border
-                expect(lines.find((l) => l.includes(chalk.inverse(' ')))).not.toBeUndefined();
+                await vi.waitFor(() => {
+                    const frame = stdout.lastFrame();
+                    const lines = frame.split('\n');
+                    // The line with the cursor should just be an inverted space inside the box border
+                    expect(lines.find((l) => l.includes(chalk.inverse(' ')))).not.toBeUndefined();
+                });
                 unmount();
             });
         });
@@ -904,15 +952,16 @@ describe('InputPrompt', () => {
                 [2, 0], // 'world' is logical line 2, col 0
             ];
             const { stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            const frame = stdout.lastFrame();
-            // Check that all lines, including the empty one, are rendered.
-            // This implicitly tests that the Box wrapper provides height for the empty line.
-            expect(frame).toContain('hello');
-            expect(frame).toContain(`world${chalk.inverse(' ')}`);
-            const outputLines = frame.split('\n');
-            // The number of lines should be 2 for the border plus 3 for the content.
-            expect(outputLines.length).toBe(5);
+            await vi.waitFor(() => {
+                const frame = stdout.lastFrame();
+                // Check that all lines, including the empty one, are rendered.
+                // This implicitly tests that the Box wrapper provides height for the empty line.
+                expect(frame).toContain('hello');
+                expect(frame).toContain(`world${chalk.inverse(' ')}`);
+                const outputLines = frame.split('\n');
+                // The number of lines should be 2 for the border plus 3 for the content.
+                expect(outputLines.length).toBe(5);
+            });
             unmount();
         });
     });
@@ -932,16 +981,18 @@ describe('InputPrompt', () => {
             },
         ])('should handle multiline paste $description', async ({ pastedText }) => {
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
             // Simulate a bracketed paste event from the terminal
-            stdin.write(`\x1b[200~${pastedText}\x1b[201~`);
-            await wait();
-            // Verify that the buffer's handleInput was called once with the full text
-            expect(props.buffer.handleInput).toHaveBeenCalledTimes(1);
-            expect(props.buffer.handleInput).toHaveBeenCalledWith(expect.objectContaining({
-                paste: true,
-                sequence: pastedText,
-            }));
+            await act(async () => {
+                stdin.write(`\x1b[200~${pastedText}\x1b[201~`);
+            });
+            await vi.waitFor(() => {
+                // Verify that the buffer's handleInput was called once with the full text
+                expect(props.buffer.handleInput).toHaveBeenCalledTimes(1);
+                expect(props.buffer.handleInput).toHaveBeenCalledWith(expect.objectContaining({
+                    paste: true,
+                    sequence: pastedText,
+                }));
+            });
             unmount();
         });
     });
@@ -963,11 +1014,13 @@ describe('InputPrompt', () => {
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
             await vi.runAllTimersAsync();
             // Simulate a paste operation (this should set the paste protection)
-            act(() => {
+            await act(async () => {
                 stdin.write(`\x1b[200~pasted content\x1b[201~`);
             });
             // Simulate an Enter key press immediately after paste
-            stdin.write('\r');
+            await act(async () => {
+                stdin.write('\r');
+            });
             await vi.runAllTimersAsync();
             // Verify that onSubmit was NOT called due to recent paste protection
             expect(props.onSubmit).not.toHaveBeenCalled();
@@ -981,7 +1034,7 @@ describe('InputPrompt', () => {
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
             await vi.runAllTimersAsync();
             // Simulate a paste operation (this sets the protection)
-            act(() => {
+            await act(async () => {
                 stdin.write('\x1b[200~pasted text\x1b[201~');
             });
             await vi.runAllTimersAsync();
@@ -990,7 +1043,9 @@ describe('InputPrompt', () => {
                 await vi.advanceTimersByTimeAsync(50);
             });
             // Now Enter should work normally
-            stdin.write('\r');
+            await act(async () => {
+                stdin.write('\r');
+            });
             await vi.runAllTimersAsync();
             expect(props.onSubmit).toHaveBeenCalledWith('pasted text');
             expect(props.buffer.newline).not.toHaveBeenCalled();
@@ -1011,10 +1066,14 @@ describe('InputPrompt', () => {
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }), { kittyProtocolEnabled: true });
             await vi.runAllTimersAsync();
             // Simulate a paste operation
-            stdin.write('\x1b[200~some pasted stuff\x1b[201~');
+            await act(async () => {
+                stdin.write('\x1b[200~some pasted stuff\x1b[201~');
+            });
             await vi.runAllTimersAsync();
             // Simulate an Enter key press immediately after paste
-            stdin.write('\r');
+            await act(async () => {
+                stdin.write('\r');
+            });
             await vi.runAllTimersAsync();
             // Verify that onSubmit was called
             expect(props.onSubmit).toHaveBeenCalledWith('pasted command');
@@ -1026,7 +1085,9 @@ describe('InputPrompt', () => {
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
             await vi.runAllTimersAsync();
             // Press Enter without any recent paste
-            stdin.write('\r');
+            await act(async () => {
+                stdin.write('\r');
+            });
             await vi.runAllTimersAsync();
             // Verify that onSubmit was called normally
             expect(props.onSubmit).toHaveBeenCalledWith('normal command');
@@ -1039,13 +1100,19 @@ describe('InputPrompt', () => {
             props.onEscapePromptChange = onEscapePromptChange;
             props.buffer.setText('text to clear');
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }), { kittyProtocolEnabled: false });
-            await wait();
-            stdin.write('\x1B');
-            await wait();
-            stdin.write('\x1B');
-            await wait();
-            expect(props.buffer.setText).toHaveBeenCalledWith('');
-            expect(mockCommandCompletion.resetCompletionState).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\x1B');
+            });
+            await vi.waitFor(() => {
+                expect(onEscapePromptChange).toHaveBeenCalledWith(true);
+            });
+            await act(async () => {
+                stdin.write('\x1B');
+            });
+            await vi.waitFor(() => {
+                expect(props.buffer.setText).toHaveBeenCalledWith('');
+                expect(mockCommandCompletion.resetCompletionState).toHaveBeenCalled();
+            });
             unmount();
         });
         it('should reset escape state on any non-ESC key', async () => {
@@ -1053,12 +1120,16 @@ describe('InputPrompt', () => {
             props.onEscapePromptChange = onEscapePromptChange;
             props.buffer.setText('some text');
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }), { kittyProtocolEnabled: false });
-            stdin.write('\x1B');
-            await waitFor(() => {
+            await act(async () => {
+                stdin.write('\x1B');
+            });
+            await vi.waitFor(() => {
                 expect(onEscapePromptChange).toHaveBeenCalledWith(true);
             });
-            stdin.write('a');
-            await waitFor(() => {
+            await act(async () => {
+                stdin.write('a');
+            });
+            await vi.waitFor(() => {
                 expect(onEscapePromptChange).toHaveBeenCalledWith(false);
             });
             unmount();
@@ -1066,10 +1137,10 @@ describe('InputPrompt', () => {
         it('should handle ESC in shell mode by disabling shell mode', async () => {
             props.shellModeActive = true;
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }), { kittyProtocolEnabled: false });
-            await wait();
-            stdin.write('\x1B');
-            await wait();
-            expect(props.setShellModeActive).toHaveBeenCalledWith(false);
+            await act(async () => {
+                stdin.write('\x1B');
+            });
+            await vi.waitFor(() => expect(props.setShellModeActive).toHaveBeenCalledWith(false));
             unmount();
         });
         it('should handle ESC when completion suggestions are showing', async () => {
@@ -1079,10 +1150,10 @@ describe('InputPrompt', () => {
                 suggestions: [{ label: 'suggestion', value: 'suggestion' }],
             });
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }), { kittyProtocolEnabled: false });
-            await wait();
-            stdin.write('\x1B');
-            await wait();
-            expect(mockCommandCompletion.resetCompletionState).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\x1B');
+            });
+            await vi.waitFor(() => expect(mockCommandCompletion.resetCompletionState).toHaveBeenCalled());
             unmount();
         });
         it('should not call onEscapePromptChange when not provided', async () => {
@@ -1091,20 +1162,23 @@ describe('InputPrompt', () => {
             props.buffer.setText('some text');
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }), { kittyProtocolEnabled: false });
             await vi.runAllTimersAsync();
-            stdin.write('\x1B');
+            await act(async () => {
+                stdin.write('\x1B');
+            });
             await vi.runAllTimersAsync();
             vi.useRealTimers();
             unmount();
         });
         it('should not interfere with existing keyboard shortcuts', async () => {
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }), { kittyProtocolEnabled: false });
-            await wait();
-            stdin.write('\x0C');
-            await wait();
-            expect(props.onClearScreen).toHaveBeenCalled();
-            stdin.write('\x01');
-            await wait();
-            expect(props.buffer.move).toHaveBeenCalledWith('home');
+            await act(async () => {
+                stdin.write('\x0C');
+            });
+            await vi.waitFor(() => expect(props.onClearScreen).toHaveBeenCalled());
+            await act(async () => {
+                stdin.write('\x01');
+            });
+            await vi.waitFor(() => expect(props.buffer.move).toHaveBeenCalledWith('home'));
             unmount();
         });
     });
@@ -1132,12 +1206,11 @@ describe('InputPrompt', () => {
                 activeSuggestionIndex: 0,
             });
             const { stdin, stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
             // Trigger reverse search with Ctrl+R
-            act(() => {
+            await act(async () => {
                 stdin.write('\x12');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 const frame = stdout.lastFrame();
                 expect(frame).toContain('(r:)');
                 expect(frame).toContain('echo hello');
@@ -1148,12 +1221,17 @@ describe('InputPrompt', () => {
         });
         it('resets reverse search state on Escape', async () => {
             const { stdin, stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x12');
-            await wait();
-            stdin.write('\x1B');
-            stdin.write('\u001b[27u'); // Press kitty escape key
-            await waitFor(() => {
+            await act(async () => {
+                stdin.write('\x12');
+            });
+            await vi.waitFor(() => { });
+            await act(async () => {
+                stdin.write('\x1B');
+            });
+            await act(async () => {
+                stdin.write('\u001b[27u'); // Press kitty escape key
+            });
+            await vi.waitFor(() => {
                 expect(stdout.lastFrame()).not.toContain('(r:)');
             });
             expect(stdout.lastFrame()).not.toContain('echo hello');
@@ -1179,20 +1257,21 @@ describe('InputPrompt', () => {
             }));
             const { stdin, stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
             // Enter reverse search mode with Ctrl+R
-            act(() => {
+            await act(async () => {
                 stdin.write('\x12');
             });
             // Verify reverse search is active
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(stdout.lastFrame()).toContain('(r:)');
             });
             // Press Tab to complete the highlighted entry
-            act(() => {
+            await act(async () => {
                 stdin.write('\t');
             });
-            await wait();
-            expect(mockHandleAutocomplete).toHaveBeenCalledWith(0);
-            expect(props.buffer.setText).toHaveBeenCalledWith('echo hello');
+            await vi.waitFor(() => {
+                expect(mockHandleAutocomplete).toHaveBeenCalledWith(0);
+                expect(props.buffer.setText).toHaveBeenCalledWith('echo hello');
+            });
             unmount();
         }, 15000);
         it('submits the highlighted entry on Enter and exits reverse-search', async () => {
@@ -1208,16 +1287,16 @@ describe('InputPrompt', () => {
                 activeSuggestionIndex: 0,
             });
             const { stdin, stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            act(() => {
+            await act(async () => {
                 stdin.write('\x12');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(stdout.lastFrame()).toContain('(r:)');
             });
-            act(() => {
+            await act(async () => {
                 stdin.write('\r');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(stdout.lastFrame()).not.toContain('(r:)');
             });
             expect(props.onSubmit).toHaveBeenCalledWith('echo hello');
@@ -1237,19 +1316,18 @@ describe('InputPrompt', () => {
                 showSuggestions: reverseSearchActiveFromInputPrompt,
             }));
             const { stdin, stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
             // reverse search with Ctrl+R
-            act(() => {
+            await act(async () => {
                 stdin.write('\x12');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(stdout.lastFrame()).toContain('(r:)');
             });
             // Press kitty escape key
-            act(() => {
+            await act(async () => {
                 stdin.write('\u001b[27u');
             });
-            await waitFor(() => {
+            await vi.waitFor(() => {
                 expect(stdout.lastFrame()).not.toContain('(r:)');
                 expect(props.buffer.text).toBe(initialText);
                 expect(props.buffer.cursor).toEqual(initialCursor);
@@ -1263,10 +1341,12 @@ describe('InputPrompt', () => {
             props.buffer.cursor = [1, 2];
             props.buffer.lines = ['line 1', 'line 2', 'line 3'];
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x05'); // Ctrl+E
-            await wait();
-            expect(props.buffer.move).toHaveBeenCalledWith('end');
+            await act(async () => {
+                stdin.write('\x05'); // Ctrl+E
+            });
+            await vi.waitFor(() => {
+                expect(props.buffer.move).toHaveBeenCalledWith('end');
+            });
             expect(props.buffer.moveToOffset).not.toHaveBeenCalled();
             unmount();
         });
@@ -1275,10 +1355,12 @@ describe('InputPrompt', () => {
             props.buffer.cursor = [0, 5];
             props.buffer.lines = ['single line text'];
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x05'); // Ctrl+E
-            await wait();
-            expect(props.buffer.move).toHaveBeenCalledWith('end');
+            await act(async () => {
+                stdin.write('\x05'); // Ctrl+E
+            });
+            await vi.waitFor(() => {
+                expect(props.buffer.move).toHaveBeenCalledWith('end');
+            });
             expect(props.buffer.moveToOffset).not.toHaveBeenCalled();
             unmount();
         });
@@ -1298,15 +1380,15 @@ describe('InputPrompt', () => {
                 activeSuggestionIndex: isActive ? 0 : -1,
             }));
             const { stdin, stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            act(() => {
+            await act(async () => {
                 stdin.write('\x12'); // Ctrl+R
             });
-            await wait();
-            const frame = stdout.lastFrame() ?? '';
-            expect(frame).toContain('(r:)');
-            expect(frame).toContain('git commit');
-            expect(frame).toContain('git push');
+            await vi.waitFor(() => {
+                const frame = stdout.lastFrame() ?? '';
+                expect(frame).toContain('(r:)');
+                expect(frame).toContain('git commit');
+                expect(frame).toContain('git push');
+            });
             unmount();
         });
         it('expands and collapses long suggestion via Right/Left arrows', async () => {
@@ -1321,18 +1403,26 @@ describe('InputPrompt', () => {
                 isLoadingSuggestions: false,
             });
             const { stdin, stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x12');
-            await wait();
-            expect(clean(stdout.lastFrame())).toContain('→');
-            stdin.write('\u001B[C');
-            await wait(200);
-            expect(clean(stdout.lastFrame())).toContain('←');
-            expect(stdout.lastFrame()).toMatchSnapshot('command-search-expanded-match');
-            stdin.write('\u001B[D');
-            await wait();
-            expect(clean(stdout.lastFrame())).toContain('→');
-            expect(stdout.lastFrame()).toMatchSnapshot('command-search-collapsed-match');
+            await act(async () => {
+                stdin.write('\x12');
+            });
+            await vi.waitFor(() => {
+                expect(clean(stdout.lastFrame())).toContain('→');
+            });
+            await act(async () => {
+                stdin.write('\u001B[C');
+            });
+            await vi.waitFor(() => {
+                expect(clean(stdout.lastFrame())).toContain('←');
+            });
+            expect(stdout.lastFrame()).toMatchSnapshot('command-search-render-expanded-match');
+            await act(async () => {
+                stdin.write('\u001B[D');
+            });
+            await vi.waitFor(() => {
+                expect(clean(stdout.lastFrame())).toContain('→');
+            });
+            expect(stdout.lastFrame()).toMatchSnapshot('command-search-render-collapsed-match');
             unmount();
         });
         it('renders match window and expanded view (snapshots)', async () => {
@@ -1349,13 +1439,18 @@ describe('InputPrompt', () => {
                 isLoadingSuggestions: false,
             });
             const { stdin, stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x12');
-            await wait();
-            expect(stdout.lastFrame()).toMatchSnapshot('command-search-collapsed-match');
-            stdin.write('\u001B[C');
-            await wait();
-            expect(stdout.lastFrame()).toMatchSnapshot('command-search-expanded-match');
+            await act(async () => {
+                stdin.write('\x12');
+            });
+            await vi.waitFor(() => {
+                expect(stdout.lastFrame()).toMatchSnapshot('command-search-render-collapsed-match');
+            });
+            await act(async () => {
+                stdin.write('\u001B[C');
+            });
+            await vi.waitFor(() => {
+                expect(stdout.lastFrame()).toMatchSnapshot('command-search-render-expanded-match');
+            });
             unmount();
         });
         it('does not show expand/collapse indicator for short suggestions', async () => {
@@ -1370,12 +1465,16 @@ describe('InputPrompt', () => {
                 isLoadingSuggestions: false,
             });
             const { stdin, stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\x12');
-            await wait();
-            const frame = clean(stdout.lastFrame());
-            expect(frame).not.toContain('→');
-            expect(frame).not.toContain('←');
+            await act(async () => {
+                stdin.write('\x12');
+            });
+            await vi.waitFor(() => {
+                const frame = clean(stdout.lastFrame());
+                // Ensure it rendered the search mode
+                expect(frame).toContain('(r:)');
+                expect(frame).not.toContain('→');
+                expect(frame).not.toContain('←');
+            });
             unmount();
         });
     });
@@ -1385,12 +1484,12 @@ describe('InputPrompt', () => {
             props.popAllMessages = mockPopAllMessages;
             props.buffer.text = '';
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\u001B[A');
-            await wait();
-            expect(mockPopAllMessages).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\u001B[A');
+            });
+            await vi.waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());
             const callback = mockPopAllMessages.mock.calls[0][0];
-            act(() => {
+            await act(async () => {
                 callback('Message 1\n\nMessage 2\n\nMessage 3');
             });
             expect(props.buffer.setText).toHaveBeenCalledWith('Message 1\n\nMessage 2\n\nMessage 3');
@@ -1401,11 +1500,11 @@ describe('InputPrompt', () => {
             props.popAllMessages = mockPopAllMessages;
             props.buffer.text = 'some text';
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\u001B[A');
-            await wait();
+            await act(async () => {
+                stdin.write('\u001B[A');
+            });
+            await vi.waitFor(() => expect(mockInputHistory.navigateUp).toHaveBeenCalled());
             expect(mockPopAllMessages).not.toHaveBeenCalled();
-            expect(mockInputHistory.navigateUp).toHaveBeenCalled();
             unmount();
         });
         it('should handle undefined messages from popAllMessages', async () => {
@@ -1413,12 +1512,12 @@ describe('InputPrompt', () => {
             props.popAllMessages = mockPopAllMessages;
             props.buffer.text = '';
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\u001B[A');
-            await wait();
-            expect(mockPopAllMessages).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\u001B[A');
+            });
+            await vi.waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());
             const callback = mockPopAllMessages.mock.calls[0][0];
-            act(() => {
+            await act(async () => {
                 callback(undefined);
             });
             expect(props.buffer.setText).not.toHaveBeenCalled();
@@ -1433,10 +1532,10 @@ describe('InputPrompt', () => {
             props.buffer.visualCursor = [0, 0];
             props.buffer.visualScrollRow = 0;
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\u001B[A');
-            await wait();
-            expect(mockPopAllMessages).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\u001B[A');
+            });
+            await vi.waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());
             unmount();
         });
         it('should handle single queued message', async () => {
@@ -1444,11 +1543,12 @@ describe('InputPrompt', () => {
             props.popAllMessages = mockPopAllMessages;
             props.buffer.text = '';
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\u001B[A');
-            await wait();
+            await act(async () => {
+                stdin.write('\u001B[A');
+            });
+            await vi.waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());
             const callback = mockPopAllMessages.mock.calls[0][0];
-            act(() => {
+            await act(async () => {
                 callback('Single message');
             });
             expect(props.buffer.setText).toHaveBeenCalledWith('Single message');
@@ -1459,20 +1559,20 @@ describe('InputPrompt', () => {
             props.popAllMessages = mockPopAllMessages;
             props.buffer.text = '   '; // Whitespace only
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\u001B[A');
-            await wait();
-            expect(mockPopAllMessages).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\u001B[A');
+            });
+            await vi.waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());
             unmount();
         });
         it('should not call popAllMessages if it is not provided', async () => {
             props.popAllMessages = undefined;
             props.buffer.text = '';
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\u001B[A');
-            await wait();
-            expect(mockInputHistory.navigateUp).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\u001B[A');
+            });
+            await vi.waitFor(() => expect(mockInputHistory.navigateUp).toHaveBeenCalled());
             unmount();
         });
         it('should navigate input history on fresh start when no queued messages exist', async () => {
@@ -1480,12 +1580,12 @@ describe('InputPrompt', () => {
             props.popAllMessages = mockPopAllMessages;
             props.buffer.text = '';
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\u001B[A');
-            await wait();
-            expect(mockPopAllMessages).toHaveBeenCalled();
+            await act(async () => {
+                stdin.write('\u001B[A');
+            });
+            await vi.waitFor(() => expect(mockPopAllMessages).toHaveBeenCalled());
             const callback = mockPopAllMessages.mock.calls[0][0];
-            act(() => {
+            await act(async () => {
                 callback(undefined);
             });
             expect(mockInputHistory.navigateUp).toHaveBeenCalled();
@@ -1497,33 +1597,31 @@ describe('InputPrompt', () => {
         it('should render correctly in shell mode', async () => {
             props.shellModeActive = true;
             const { stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            expect(stdout.lastFrame()).toMatchSnapshot();
+            await vi.waitFor(() => expect(stdout.lastFrame()).toMatchSnapshot());
             unmount();
         });
         it('should render correctly when accepting edits', async () => {
             props.approvalMode = ApprovalMode.AUTO_EDIT;
             const { stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            expect(stdout.lastFrame()).toMatchSnapshot();
+            await vi.waitFor(() => expect(stdout.lastFrame()).toMatchSnapshot());
             unmount();
         });
         it('should render correctly in yolo mode', async () => {
             props.approvalMode = ApprovalMode.YOLO;
             const { stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            expect(stdout.lastFrame()).toMatchSnapshot();
+            await vi.waitFor(() => expect(stdout.lastFrame()).toMatchSnapshot());
             unmount();
         });
         it('should not show inverted cursor when shell is focused', async () => {
             props.isEmbeddedShellFocused = true;
             props.focus = false;
             const { stdout, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            expect(stdout.lastFrame()).not.toContain(`{chalk.inverse(' ')}`);
-            // This snapshot is good to make sure there was an input prompt but does
-            // not show the inverted cursor because snapshots do not show colors.
-            expect(stdout.lastFrame()).toMatchSnapshot();
+            await vi.waitFor(() => {
+                expect(stdout.lastFrame()).not.toContain(`{chalk.inverse(' ')}`);
+                // This snapshot is good to make sure there was an input prompt but does
+                // not show the inverted cursor because snapshots do not show colors.
+                expect(stdout.lastFrame()).toMatchSnapshot();
+            });
             unmount();
         });
     });
@@ -1531,10 +1629,10 @@ describe('InputPrompt', () => {
         const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }), {
             shellFocus: false,
         });
-        await wait();
-        stdin.write('a');
-        await wait();
-        expect(mockBuffer.handleInput).toHaveBeenCalled();
+        await act(async () => {
+            stdin.write('a');
+        });
+        await vi.waitFor(() => expect(mockBuffer.handleInput).toHaveBeenCalled());
         unmount();
     });
     describe('command queuing while streaming', () => {
@@ -1569,17 +1667,19 @@ describe('InputPrompt', () => {
             props.buffer.text = bufferText;
             props.shellModeActive = shellMode;
             const { stdin, unmount } = renderWithProviders(_jsx(InputPrompt, { ...props }));
-            await wait();
-            stdin.write('\r');
-            await wait();
-            if (shouldSubmit) {
-                expect(props.onSubmit).toHaveBeenCalledWith(bufferText);
-                expect(props.setQueueErrorMessage).not.toHaveBeenCalled();
-            }
-            else {
-                expect(props.onSubmit).not.toHaveBeenCalled();
-                expect(props.setQueueErrorMessage).toHaveBeenCalledWith(errorMessage);
-            }
+            await act(async () => {
+                stdin.write('\r');
+            });
+            await vi.waitFor(() => {
+                if (shouldSubmit) {
+                    expect(props.onSubmit).toHaveBeenCalledWith(bufferText);
+                    expect(props.setQueueErrorMessage).not.toHaveBeenCalled();
+                }
+                else {
+                    expect(props.onSubmit).not.toHaveBeenCalled();
+                    expect(props.setQueueErrorMessage).toHaveBeenCalledWith(errorMessage);
+                }
+            });
             unmount();
         });
     });