@machina.ai/cell-cli 1.16.0-rc2 → 1.18.4-rc1

package/dist/src/ui/AppContainer.test.js

@@ -6,35 +6,53 @@ import { jsx as _jsx } from "react/jsx-runtime";
  */
 import { describe, it, expect, vi, beforeEach, afterEach, } from 'vitest';
 import { render } from '../test-utils/render.js';
+import { waitFor } from '../test-utils/async.js';
 import { cleanup } from 'ink-testing-library';
 import { act, useContext } from 'react';
 import { AppContainer } from './AppContainer.js';
 import { SettingsContext } from './contexts/SettingsContext.js';
-import { makeFakeConfig, CoreEvent, } from '@google/gemini-cli-core';
+import { makeFakeConfig, CoreEvent, AuthType, } from '@google/gemini-cli-core';
 // Mock coreEvents
 const mockCoreEvents = vi.hoisted(() => ({
     on: vi.fn(),
     off: vi.fn(),
-    drainFeedbackBacklog: vi.fn(),
+    drainBacklogs: vi.fn(),
     emit: vi.fn(),
 }));
+// Mock IdeClient
+const mockIdeClient = vi.hoisted(() => ({
+    getInstance: vi.fn().mockReturnValue(new Promise(() => { })),
+}));
+// Mock stdout
+const mocks = vi.hoisted(() => ({
+    mockStdout: { write: vi.fn() },
+}));
 vi.mock('@google/gemini-cli-core', async (importOriginal) => {
     const actual = await importOriginal();
     return {
         ...actual,
         coreEvents: mockCoreEvents,
+        IdeClient: mockIdeClient,
+        writeToStdout: vi.fn((...args) => process.stdout.write(...args)),
+        writeToStderr: vi.fn((...args) => process.stderr.write(...args)),
+        patchStdio: vi.fn(() => () => { }),
+        createInkStdio: vi.fn(() => ({
+            stdout: process.stdout,
+            stderr: process.stderr,
+        })),
+        enableMouseEvents: vi.fn(),
+        disableMouseEvents: vi.fn(),
     };
 });
 import { useQuotaAndFallback } from './hooks/useQuotaAndFallback.js';
 import { UIStateContext } from './contexts/UIStateContext.js';
 import { UIActionsContext, } from './contexts/UIActionsContext.js';
 // Mock useStdout to capture terminal title writes
-let mockStdout;
 vi.mock('ink', async (importOriginal) => {
     const actual = await importOriginal();
     return {
         ...actual,
-        useStdout: () => ({ stdout: mockStdout }),
+        useStdout: () => ({ stdout: mocks.mockStdout }),
         measureElement: vi.fn(),
     };
 });
@@ -82,10 +100,6 @@ vi.mock('../utils/events.js');
 vi.mock('../utils/handleAutoUpdate.js');
 vi.mock('./utils/ConsolePatcher.js');
 vi.mock('../utils/cleanup.js');
-vi.mock('./utils/mouse.js', () => ({
-    enableMouseEvents: vi.fn(),
-    disableMouseEvents: vi.fn(),
-}));
 import { useHistory } from './hooks/useHistoryManager.js';
 import { useThemeCommand } from './hooks/useThemeCommand.js';
 import { useAuthCommand } from './auth/useAuth.js';
@@ -109,9 +123,8 @@ import { useLoadingIndicator } from './hooks/useLoadingIndicator.js';
 import { useKeypress } from './hooks/useKeypress.js';
 import { measureElement } from 'ink';
 import { useTerminalSize } from './hooks/useTerminalSize.js';
-import { ShellExecutionService } from '@google/gemini-cli-core';
+import { ShellExecutionService, writeToStdout, enableMouseEvents, disableMouseEvents, } from '@google/gemini-cli-core';
 import {} from '../config/extension-manager.js';
-import { enableMouseEvents, disableMouseEvents } from './utils/mouse.js';
 describe('AppContainer State Management', () => {
     let mockConfig;
     let mockSettings;
@@ -147,7 +160,7 @@ describe('AppContainer State Management', () => {
     beforeEach(() => {
         vi.clearAllMocks();
         // Initialize mock stdout for terminal title tests
-        mockStdout = { write: vi.fn() };
+        mocks.mockStdout.write.mockClear();
         // Mock computeWindowTitle function to centralize title logic testing
         vi.mock('../utils/windowTitle.js', async () => ({
             computeWindowTitle: vi.fn((folderName) => 
@@ -254,7 +267,7 @@ describe('AppContainer State Management', () => {
             // Add other properties if AppContainer uses them
         });
         mockedUseLogger.mockReturnValue({
-            getPreviousUserMessages: vi.fn().mockResolvedValue([]),
+            getPreviousUserMessages: vi.fn().mockReturnValue(new Promise(() => { })),
         });
         mockedUseLoadingIndicator.mockReturnValue({
             elapsedTime: '0.0s',
@@ -300,17 +313,13 @@ describe('AppContainer State Management', () => {
     describe('Basic Rendering', () => {
         it('renders without crashing with minimal props', async () => {
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
         it('renders with startup warnings', async () => {
             const startupWarnings = ['Warning 1', 'Warning 2'];
             const { unmount } = renderAppContainer({ startupWarnings });
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
     });
@@ -323,9 +332,7 @@ describe('AppContainer State Management', () => {
             const { unmount } = renderAppContainer({
                 initResult: initResultWithError,
             });
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
         it('handles debug mode state', () => {
@@ -339,31 +346,23 @@ describe('AppContainer State Management', () => {
     describe('Context Providers', () => {
         it('provides AppContext with correct values', async () => {
             const { unmount } = renderAppContainer({ version: '2.0.0' });
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             // Should render and unmount cleanly
             expect(() => unmount()).not.toThrow();
         });
         it('provides UIStateContext with state management', async () => {
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
         it('provides UIActionsContext with action handlers', async () => {
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
         it('provides ConfigContext with config object', async () => {
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
     });
@@ -378,9 +377,7 @@ describe('AppContainer State Management', () => {
                 },
             };
             const { unmount } = renderAppContainer({ settings: settingsAllHidden });
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
         it('handles settings with memory usage enabled', async () => {
@@ -393,18 +390,14 @@ describe('AppContainer State Management', () => {
                 },
             };
             const { unmount } = renderAppContainer({ settings: settingsWithMemory });
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
     });
     describe('Version Handling', () => {
         it.each(['1.0.0', '2.1.3-beta', '3.0.0-nightly'])('handles version format: %s', async (version) => {
             const { unmount } = renderAppContainer({ version });
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
     });
@@ -416,9 +409,6 @@ describe('AppContainer State Management', () => {
             });
             // Should still render without crashing - errors should be handled internally
             const { unmount } = renderAppContainer({ config: errorConfig });
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
             unmount();
         });
         it('handles undefined settings gracefully', async () => {
@@ -426,9 +416,7 @@ describe('AppContainer State Management', () => {
                 merged: {},
             };
             const { unmount } = renderAppContainer({ settings: undefinedSettings });
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
         });
     });
@@ -724,11 +712,10 @@ describe('AppContainer State Management', () => {
         it('passes a null proQuotaRequest to UIStateContext by default', async () => {
             // The default mock from beforeEach already sets proQuotaRequest to null
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
+            await waitFor(() => {
+                // Assert that the context value is as expected
+                expect(capturedUIState.proQuotaRequest).toBeNull();
             });
-            // Assert that the context value is as expected
-            expect(capturedUIState.proQuotaRequest).toBeNull();
             unmount();
         });
         it('passes a valid proQuotaRequest to UIStateContext when provided by the hook', async () => {
@@ -744,11 +731,10 @@ describe('AppContainer State Management', () => {
             });
             // Act: Render the container
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
+            await waitFor(() => {
+                // Assert: The mock request is correctly passed through the context
+                expect(capturedUIState.proQuotaRequest).toEqual(mockRequest);
             });
-            // Assert: The mock request is correctly passed through the context
-            expect(capturedUIState.proQuotaRequest).toEqual(mockRequest);
             unmount();
         });
         it('passes the handleProQuotaChoice function to UIActionsContext', async () => {
@@ -760,11 +746,10 @@ describe('AppContainer State Management', () => {
             });
             // Act: Render the container
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
+            await waitFor(() => {
+                // Assert: The action in the context is the mock handler we provided
+                expect(capturedUIActions.handleProQuotaChoice).toBe(mockHandler);
             });
-            // Assert: The action in the context is the mock handler we provided
-            expect(capturedUIActions.handleProQuotaChoice).toBe(mockHandler);
             // You can even verify that the plumbed function is callable
             act(() => {
                 capturedUIActions.handleProQuotaChoice('retry_later');
@@ -776,7 +761,12 @@ describe('AppContainer State Management', () => {
     describe('Terminal Title Update Feature', () => {
         beforeEach(() => {
             // Reset mock stdout for each test
-            mockStdout = { write: vi.fn() };
+            mocks.mockStdout.write.mockClear();
+        });
+        it('verifies useStdout is mocked', async () => {
+            const { useStdout } = await import('ink');
+            const { stdout } = useStdout();
+            expect(stdout).toBe(mocks.mockStdout);
         });
         it('should not update terminal title when showStatusInTitle is false', () => {
             // Arrange: Set up mock settings with showStatusInTitle disabled
@@ -796,7 +786,7 @@ describe('AppContainer State Management', () => {
                 settings: mockSettingsWithShowStatusFalse,
             });
             // Assert: Check that no title-related writes occurred
-            const titleWrites = mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
+            const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
             expect(titleWrites).toHaveLength(0);
             unmount();
         });
@@ -818,7 +808,7 @@ describe('AppContainer State Management', () => {
                 settings: mockSettingsWithHideTitleTrue,
             });
             // Assert: Check that no title-related writes occurred
-            const titleWrites = mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
+            const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
             expect(titleWrites).toHaveLength(0);
             unmount();
         });
@@ -850,7 +840,7 @@ describe('AppContainer State Management', () => {
                 settings: mockSettingsWithTitleEnabled,
             });
             // Assert: Check that title was updated with thought subject
-            const titleWrites = mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
+            const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
             expect(titleWrites).toHaveLength(1);
             expect(titleWrites[0][0]).toBe(`\x1b]2;${thoughtSubject.padEnd(80, ' ')}\x07`);
             unmount();
@@ -882,7 +872,7 @@ describe('AppContainer State Management', () => {
                 settings: mockSettingsWithTitleEnabled,
             });
             // Assert: Check that title was updated with default Idle text
-            const titleWrites = mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
+            const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
             expect(titleWrites).toHaveLength(1);
             expect(titleWrites[0][0]).toBe(`\x1b]2;${'Gemini - workspace'.padEnd(80, ' ')}\x07`);
             unmount();
@@ -915,7 +905,7 @@ describe('AppContainer State Management', () => {
                 settings: mockSettingsWithTitleEnabled,
             });
             // Assert: Check that title was updated with confirmation text
-            const titleWrites = mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
+            const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
             expect(titleWrites).toHaveLength(1);
             expect(titleWrites[0][0]).toBe(`\x1b]2;${thoughtSubject.padEnd(80, ' ')}\x07`);
             unmount();
@@ -948,7 +938,7 @@ describe('AppContainer State Management', () => {
                 settings: mockSettingsWithTitleEnabled,
             });
             // Assert: Check that title is padded to exactly 80 characters
-            const titleWrites = mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
+            const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
             expect(titleWrites).toHaveLength(1);
             const calledWith = titleWrites[0][0];
             const expectedTitle = shortTitle.padEnd(80, ' ');
@@ -986,7 +976,7 @@ describe('AppContainer State Management', () => {
                 settings: mockSettingsWithTitleEnabled,
             });
             // Assert: Check that the correct ANSI escape sequence is used
-            const titleWrites = mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
+            const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
             expect(titleWrites).toHaveLength(1);
             const expectedEscapeSequence = `\x1b]2;${title.padEnd(80, ' ')}\x07`;
             expect(titleWrites[0][0]).toBe(expectedEscapeSequence);
@@ -1021,7 +1011,7 @@ describe('AppContainer State Management', () => {
                 settings: mockSettingsWithTitleEnabled,
             });
             // Assert: Check that title was updated with CLI_TITLE value
-            const titleWrites = mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
+            const titleWrites = mocks.mockStdout.write.mock.calls.filter((call) => call[0].includes('\x1b]2;'));
             expect(titleWrites).toHaveLength(1);
             expect(titleWrites[0][0]).toBe(`\x1b]2;${'Custom Gemini Title'.padEnd(80, ' ')}\x07`);
             unmount();
@@ -1102,12 +1092,7 @@ describe('AppContainer State Management', () => {
                 activePtyId: 'some-id',
             });
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
-            // Assert: The shell should be resized to a minimum height of 1, not a negative number.
-            // The old code would have tried to set a negative height.
-            expect(resizePtySpy).toHaveBeenCalled();
+            await waitFor(() => expect(resizePtySpy).toHaveBeenCalled());
             const lastCall = resizePtySpy.mock.calls[resizePtySpy.mock.calls.length - 1];
             // Check the height argument specifically
             expect(lastCall[2]).toBe(1);
@@ -1198,7 +1183,7 @@ describe('AppContainer State Management', () => {
                 await setupKeypressTest();
                 pressKey({ name: 'c', ctrl: true }, 2);
                 expect(mockCancelOngoingRequest).toHaveBeenCalledTimes(2);
-                expect(mockHandleSlashCommand).toHaveBeenCalledWith('/quit');
+                expect(mockHandleSlashCommand).toHaveBeenCalledWith('/quit', undefined, undefined, false);
                 unmount();
             });
             it('should reset press count after a timeout', async () => {
@@ -1228,7 +1213,7 @@ describe('AppContainer State Management', () => {
             it('should quit on second press if buffer is empty', async () => {
                 await setupKeypressTest();
                 pressKey({ name: 'd', ctrl: true }, 2);
-                expect(mockHandleSlashCommand).toHaveBeenCalledWith('/quit');
+                expect(mockHandleSlashCommand).toHaveBeenCalledWith('/quit', undefined, undefined, false);
                 unmount();
             });
             it('should reset press count after a timeout', async () => {
@@ -1269,7 +1254,7 @@ describe('AppContainer State Management', () => {
             unmount = renderResult.unmount;
         };
         beforeEach(() => {
-            mockStdout.write.mockClear();
+            mocks.mockStdout.write.mockClear();
             mockedUseKeypress.mockImplementation((callback) => {
                 handleGlobalKeypress = callback;
             });
@@ -1292,7 +1277,7 @@ describe('AppContainer State Management', () => {
         ])('$modeName', ({ isAlternateMode, shouldEnable }) => {
             it(`should ${shouldEnable ? 'toggle' : 'NOT toggle'} mouse off when Ctrl+S is pressed`, async () => {
                 await setupCopyModeTest(isAlternateMode);
-                mockStdout.write.mockClear(); // Clear initial enable call
+                mocks.mockStdout.write.mockClear(); // Clear initial enable call
                 act(() => {
                     handleGlobalKeypress({
                         name: 's',
@@ -1316,7 +1301,7 @@ describe('AppContainer State Management', () => {
             if (shouldEnable) {
                 it('should toggle mouse back on when Ctrl+S is pressed again', async () => {
                     await setupCopyModeTest(isAlternateMode);
-                    mockStdout.write.mockClear();
+                    writeToStdout.mockClear();
                     // Turn it on (disable mouse)
                     act(() => {
                         handleGlobalKeypress({
@@ -1362,7 +1347,7 @@ describe('AppContainer State Management', () => {
                         });
                     });
                     rerender();
-                    mockStdout.write.mockClear();
+                    writeToStdout.mockClear();
                     // Press any other key
                     act(() => {
                         handleGlobalKeypress({
@@ -1391,9 +1376,7 @@ describe('AppContainer State Management', () => {
                 closeModelDialog: vi.fn(),
             });
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             expect(capturedUIState.isModelDialogOpen).toBe(true);
             unmount();
         });
@@ -1405,9 +1388,7 @@ describe('AppContainer State Management', () => {
                 closeModelDialog: mockCloseModelDialog,
             });
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             // Verify that the actions are correctly passed through context
             act(() => {
                 capturedUIActions.closeModelDialog();
@@ -1419,26 +1400,20 @@ describe('AppContainer State Management', () => {
     describe('CoreEvents Integration', () => {
         it('subscribes to UserFeedback and drains backlog on mount', async () => {
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             expect(mockCoreEvents.on).toHaveBeenCalledWith(CoreEvent.UserFeedback, expect.any(Function));
-            expect(mockCoreEvents.drainFeedbackBacklog).toHaveBeenCalledTimes(1);
+            expect(mockCoreEvents.drainBacklogs).toHaveBeenCalledTimes(1);
             unmount();
         });
         it('unsubscribes from UserFeedback on unmount', async () => {
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             unmount();
             expect(mockCoreEvents.off).toHaveBeenCalledWith(CoreEvent.UserFeedback, expect.any(Function));
         });
         it('adds history item when UserFeedback event is received', async () => {
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
-            });
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
             // Get the registered handler
             const handler = mockCoreEvents.on.mock.calls.find((call) => call[0] === CoreEvent.UserFeedback)?.[1];
             expect(handler).toBeDefined();
@@ -1460,10 +1435,8 @@ describe('AppContainer State Management', () => {
             // Arrange: Mock initial model
             vi.spyOn(mockConfig, 'getModel').mockReturnValue('initial-model');
             const { unmount } = renderAppContainer();
-            await act(async () => {
-                await vi.waitFor(() => {
-                    expect(capturedUIState?.currentModel).toBe('initial-model');
-                });
+            await waitFor(() => {
+                expect(capturedUIState?.currentModel).toBe('initial-model');
             });
             // Get the registered handler for ModelChanged
             const handler = mockCoreEvents.on.mock.calls.find((call) => call[0] === CoreEvent.ModelChanged)?.[1];
@@ -1495,10 +1468,115 @@ describe('AppContainer State Management', () => {
             });
             // The main assertion is that the render does not throw.
             const { unmount } = render(_jsx(SettingsContext.Provider, { value: mockSettings, children: _jsx(AppContainer, { config: mockConfig, version: "1.0.0", initializationResult: mockInitResult }) }));
+            await waitFor(() => expect(resizePtySpy).toHaveBeenCalled());
+            unmount();
+        });
+    });
+    describe('Banner Text', () => {
+        it('should render placeholder banner text for USE_GEMINI auth type', async () => {
+            const config = makeFakeConfig();
+            vi.spyOn(config, 'getContentGeneratorConfig').mockReturnValue({
+                authType: AuthType.USE_GEMINI,
+                apiKey: 'fake-key',
+            });
+            const { unmount } = renderAppContainer();
+            await waitFor(() => {
+                expect(capturedUIState.bannerData.defaultText).toBeDefined();
+                unmount();
+            });
+        });
+    });
+    describe('onCancelSubmit Behavior', () => {
+        let mockSetText;
+        // Helper to extract arguments from the useGeminiStream hook call
+        // This isolates the positional argument dependency to a single location
+        const extractUseGeminiStreamArgs = (args) => ({
+            onCancelSubmit: args[13],
+        });
+        beforeEach(() => {
+            mockSetText = vi.fn();
+            mockedUseTextBuffer.mockReturnValue({
+                text: '',
+                setText: mockSetText,
+            });
+        });
+        it('clears the prompt when onCancelSubmit is called with shouldRestorePrompt=false', async () => {
+            const { unmount } = renderAppContainer();
+            await waitFor(() => expect(capturedUIState).toBeTruthy());
+            const { onCancelSubmit } = extractUseGeminiStreamArgs(mockedUseGeminiStream.mock.lastCall);
+            act(() => {
+                onCancelSubmit(false);
+            });
+            expect(mockSetText).toHaveBeenCalledWith('');
+            unmount();
+        });
+        it('restores the prompt when onCancelSubmit is called with shouldRestorePrompt=true (or undefined)', async () => {
+            mockedUseLogger.mockReturnValue({
+                getPreviousUserMessages: vi
+                    .fn()
+                    .mockResolvedValue(['previous message']),
+            });
+            const { unmount } = renderAppContainer();
+            await waitFor(() => expect(capturedUIState.userMessages).toContain('previous message'));
+            const { onCancelSubmit } = extractUseGeminiStreamArgs(mockedUseGeminiStream.mock.lastCall);
             await act(async () => {
-                await new Promise((resolve) => setTimeout(resolve, 0));
+                onCancelSubmit(true);
+            });
+            expect(mockSetText).toHaveBeenCalledWith('previous message');
+            unmount();
+        });
+        it('correctly restores prompt even if userMessages is stale (race condition fix)', async () => {
+            // Setup initial history with one message
+            const initialHistory = [{ type: 'user', text: 'Previous Prompt' }];
+            mockedUseHistory.mockReturnValue({
+                history: initialHistory,
+                addItem: vi.fn(),
+                updateItem: vi.fn(),
+                clearItems: vi.fn(),
+                loadHistory: vi.fn(),
+            });
+            let resolveLoggerPromise;
+            const loggerPromise = new Promise((resolve) => {
+                resolveLoggerPromise = resolve;
+            });
+            // Mock logger to control when userMessages updates
+            const getPreviousUserMessagesMock = vi
+                .fn()
+                .mockResolvedValueOnce([]) // Initial mount
+                .mockReturnValueOnce(loggerPromise); // Second render (simulated update)
+            mockedUseLogger.mockReturnValue({
+                getPreviousUserMessages: getPreviousUserMessagesMock,
+            });
+            const { unmount, rerender } = renderAppContainer();
+            // Wait for userMessages to be populated with 'Previous Prompt'
+            await waitFor(() => expect(capturedUIState.userMessages).toContain('Previous Prompt'));
+            // Simulate a new prompt being added (e.g., user sent it, but it overflowed)
+            const newPrompt = 'Current Prompt that Overflowed';
+            const newHistory = [...initialHistory, { type: 'user', text: newPrompt }];
+            mockedUseHistory.mockReturnValue({
+                history: newHistory,
+                addItem: vi.fn(),
+                updateItem: vi.fn(),
+                clearItems: vi.fn(),
+                loadHistory: vi.fn(),
+            });
+            // Rerender to reflect the history change.
+            // This triggers the effect to update userMessages, but it hangs on loggerPromise.
+            rerender(getAppContainer());
+            const { onCancelSubmit } = extractUseGeminiStreamArgs(mockedUseGeminiStream.mock.lastCall);
+            // Call onCancelSubmit immediately. userMessages is still stale (has only 'Previous Prompt')
+            // because the effect is waiting on loggerPromise.
+            act(() => {
+                onCancelSubmit(true);
+            });
+            // Now resolve the promise to let the effect complete and update userMessages
+            await act(async () => {
+                resolveLoggerPromise([]);
+            });
+            // With the fix, it should have waited for userMessages to update and then set the new prompt
+            await waitFor(() => {
+                expect(mockSetText).toHaveBeenCalledWith(newPrompt);
             });
-            expect(resizePtySpy).toHaveBeenCalled();
             unmount();
         });
     });