split context (#24623)

2026-04-08 04:10:35 -07:00 · 2026-04-06 10:20:38 -07:00
parent c96cb09e09
commit 70f6d6a992
20 changed files with 651 additions and 400 deletions
--- a/packages/cli/src/test-utils/render.tsx
+++ b/packages/cli/src/test-utils/render.tsx
@@ -598,6 +598,9 @@ const mockUIActions: UIActions = {
  clearAccountSuspension: vi.fn(),
 };

+import { type TextBuffer } from '../ui/components/shared/text-buffer.js';
+import { InputContext, type InputState } from '../ui/contexts/InputContext.js';
+
 let capturedOverflowState: OverflowState | undefined;
 let capturedOverflowActions: OverflowActions | undefined;
 const ContextCapture: React.FC<{ children: React.ReactNode }> = ({
@@ -614,6 +617,7 @@ export const renderWithProviders = async (
    shellFocus = true,
    settings = mockSettings,
    uiState: providedUiState,
+    inputState: providedInputState,
    width,
    mouseEventsEnabled = false,
    config,
@@ -625,6 +629,7 @@ export const renderWithProviders = async (
    shellFocus?: boolean;
    settings?: LoadedSettings;
    uiState?: Partial<UIState>;
+    inputState?: Partial<InputState>;
    width?: number;
    mouseEventsEnabled?: boolean;
    config?: Config;
@@ -659,6 +664,18 @@ export const renderWithProviders = async (
    },
  ) as UIState;

+  const inputState = {
+    buffer: { text: '' } as unknown as TextBuffer,
+    userMessages: [],
+    shellModeActive: false,
+    showEscapePrompt: false,
+    copyModeEnabled: false,
+    inputWidth: 80,
+    suggestionsWidth: 40,
+    ...(providedUiState as unknown as Partial<InputState>),
+    ...providedInputState,
+  };
+
  if (persistentState?.get) {
    persistentStateMock.get.mockImplementation(persistentState.get);
  }
@@ -708,63 +725,65 @@ export const renderWithProviders = async (
    <AppContext.Provider value={appState}>
      <ConfigContext.Provider value={config}>
        <SettingsContext.Provider value={settings}>
-          <UIStateContext.Provider value={finalUiState}>
-            <VimModeProvider>
-              <ShellFocusContext.Provider value={shellFocus}>
-                <SessionStatsProvider>
-                  <StreamingContext.Provider
-                    value={finalUiState.streamingState}
-                  >
-                    <UIActionsContext.Provider value={finalUIActions}>
-                      <OverflowProvider>
-                        <ToolActionsProvider
-                          config={config}
-                          toolCalls={allToolCalls}
-                          isExpanded={
-                            toolActions?.isExpanded ??
-                            vi.fn().mockReturnValue(false)
-                          }
-                          toggleExpansion={
-                            toolActions?.toggleExpansion ?? vi.fn()
-                          }
-                          toggleAllExpansion={
-                            toolActions?.toggleAllExpansion ?? vi.fn()
-                          }
-                        >
-                          <AskUserActionsProvider
-                            request={null}
-                            onSubmit={vi.fn()}
-                            onCancel={vi.fn()}
+          <InputContext.Provider value={inputState}>
+            <UIStateContext.Provider value={finalUiState}>
+              <VimModeProvider>
+                <ShellFocusContext.Provider value={shellFocus}>
+                  <SessionStatsProvider>
+                    <StreamingContext.Provider
+                      value={finalUiState.streamingState}
+                    >
+                      <UIActionsContext.Provider value={finalUIActions}>
+                        <OverflowProvider>
+                          <ToolActionsProvider
+                            config={config}
+                            toolCalls={allToolCalls}
+                            isExpanded={
+                              toolActions?.isExpanded ??
+                              vi.fn().mockReturnValue(false)
+                            }
+                            toggleExpansion={
+                              toolActions?.toggleExpansion ?? vi.fn()
+                            }
+                            toggleAllExpansion={
+                              toolActions?.toggleAllExpansion ?? vi.fn()
+                            }
                          >
-                            <KeypressProvider>
-                              <MouseProvider
-                                mouseEventsEnabled={mouseEventsEnabled}
-                              >
-                                <TerminalProvider>
-                                  <ScrollProvider>
-                                    <ContextCapture>
-                                      <Box
-                                        width={terminalWidth}
-                                        flexShrink={0}
-                                        flexGrow={0}
-                                        flexDirection="column"
-                                      >
-                                        {comp}
-                                      </Box>
-                                    </ContextCapture>
-                                  </ScrollProvider>
-                                </TerminalProvider>
-                              </MouseProvider>
-                            </KeypressProvider>
-                          </AskUserActionsProvider>
-                        </ToolActionsProvider>
-                      </OverflowProvider>
-                    </UIActionsContext.Provider>
-                  </StreamingContext.Provider>
-                </SessionStatsProvider>
-              </ShellFocusContext.Provider>
-            </VimModeProvider>
-          </UIStateContext.Provider>
+                            <AskUserActionsProvider
+                              request={null}
+                              onSubmit={vi.fn()}
+                              onCancel={vi.fn()}
+                            >
+                              <KeypressProvider>
+                                <MouseProvider
+                                  mouseEventsEnabled={mouseEventsEnabled}
+                                >
+                                  <TerminalProvider>
+                                    <ScrollProvider>
+                                      <ContextCapture>
+                                        <Box
+                                          width={terminalWidth}
+                                          flexShrink={0}
+                                          flexGrow={0}
+                                          flexDirection="column"
+                                        >
+                                          {comp}
+                                        </Box>
+                                      </ContextCapture>
+                                    </ScrollProvider>
+                                  </TerminalProvider>
+                                </MouseProvider>
+                              </KeypressProvider>
+                            </AskUserActionsProvider>
+                          </ToolActionsProvider>
+                        </OverflowProvider>
+                      </UIActionsContext.Provider>
+                    </StreamingContext.Provider>
+                  </SessionStatsProvider>
+                </ShellFocusContext.Provider>
+              </VimModeProvider>
+            </UIStateContext.Provider>
+          </InputContext.Provider>
        </SettingsContext.Provider>
      </ConfigContext.Provider>
    </AppContext.Provider>
--- a/packages/cli/src/ui/AppContainer.test.tsx
+++ b/packages/cli/src/ui/AppContainer.test.tsx
@@ -122,13 +122,17 @@ vi.mock('ink', async (importOriginal) => {
  };
 });

+import { InputContext, type InputState } from './contexts/InputContext.js';
+
 // Helper component will read the context values provided by AppContainer
 // so we can assert against them in our tests.
 let capturedUIState: UIState;
+let capturedInputState: InputState;
 let capturedUIActions: UIActions;
 let capturedOverflowActions: OverflowActions;
 function TestContextConsumer() {
  capturedUIState = useContext(UIStateContext)!;
+  capturedInputState = useContext(InputContext)!;
  capturedUIActions = useContext(UIActionsContext)!;
  capturedOverflowActions = useOverflowActions()!;
  return null;
@@ -3036,7 +3040,7 @@ describe('AppContainer State Management', () => {
      });

      const { unmount } = await act(async () => renderAppContainer());
-      expect(capturedUIState.userMessages).toContain('previous message');
+      expect(capturedInputState.userMessages).toContain('previous message');

      const { onCancelSubmit } = extractUseGeminiStreamArgs(
        mockedUseGeminiStream.mock.lastCall!,
@@ -3064,8 +3068,8 @@ describe('AppContainer State Management', () => {
      const { rerender, unmount } = await act(async () => renderAppContainer());

      // Verify userMessages is populated from inputHistory
-      expect(capturedUIState.userMessages).toContain('first prompt');
-      expect(capturedUIState.userMessages).toContain('second prompt');
+      expect(capturedInputState.userMessages).toContain('first prompt');
+      expect(capturedInputState.userMessages).toContain('second prompt');

      // Clear the conversation history (simulating /clear command)
      const mockClearItems = vi.fn();
@@ -3084,8 +3088,8 @@ describe('AppContainer State Management', () => {

      // Verify that userMessages still contains the input history
      // (it should not be affected by clearing conversation history)
-      expect(capturedUIState.userMessages).toContain('first prompt');
-      expect(capturedUIState.userMessages).toContain('second prompt');
+      expect(capturedInputState.userMessages).toContain('first prompt');
+      expect(capturedInputState.userMessages).toContain('second prompt');

      unmount();
    });
--- a/packages/cli/src/ui/AppContainer.tsx
+++ b/packages/cli/src/ui/AppContainer.tsx
@@ -194,6 +194,8 @@ import {
 } from './hooks/useVisibilityToggle.js';
 import { useKeyMatchers } from './hooks/useKeyMatchers.js';

+import { InputContext } from './contexts/InputContext.js';
+
 /**
 * The fraction of the terminal width to allocate to the shell.
 * This provides horizontal padding.
@@ -2328,6 +2330,27 @@ Logging in with Google... Restarting Gemini CLI to continue.
    };
  }, [config, refreshStatic]);

+  const inputState = useMemo(
+    () => ({
+      buffer,
+      userMessages: inputHistory,
+      shellModeActive,
+      showEscapePrompt,
+      copyModeEnabled,
+      inputWidth,
+      suggestionsWidth,
+    }),
+    [
+      buffer,
+      inputHistory,
+      shellModeActive,
+      showEscapePrompt,
+      copyModeEnabled,
+      inputWidth,
+      suggestionsWidth,
+    ],
+  );
+
  const uiState: UIState = useMemo(
    () => ({
      history: historyManager.history,
@@ -2371,11 +2394,6 @@ Logging in with Google... Restarting Gemini CLI to continue.
      initError,
      pendingGeminiHistoryItems,
      thought,
-      shellModeActive,
-      userMessages: inputHistory,
-      buffer,
-      inputWidth,
-      suggestionsWidth,
      isInputActive,
      isResuming,
      shouldShowIdePrompt,
@@ -2391,7 +2409,6 @@ Logging in with Google... Restarting Gemini CLI to continue.
      renderMarkdown,
      ctrlCPressedOnce: ctrlCPressCount >= 1,
      ctrlDPressedOnce: ctrlDPressCount >= 1,
-      showEscapePrompt,
      shortcutsHelpVisible,
      cleanUiDetailsVisible,
      isFocused,
@@ -2443,7 +2460,6 @@ Logging in with Google... Restarting Gemini CLI to continue.
      embeddedShellFocused,
      showDebugProfiler,
      customDialog,
-      copyModeEnabled,
      transientMessage,
      bannerData,
      bannerVisible,
@@ -2498,11 +2514,6 @@ Logging in with Google... Restarting Gemini CLI to continue.
      initError,
      pendingGeminiHistoryItems,
      thought,
-      shellModeActive,
-      inputHistory,
-      buffer,
-      inputWidth,
-      suggestionsWidth,
      isInputActive,
      isResuming,
      shouldShowIdePrompt,
@@ -2518,7 +2529,6 @@ Logging in with Google... Restarting Gemini CLI to continue.
      renderMarkdown,
      ctrlCPressCount,
      ctrlDPressCount,
-      showEscapePrompt,
      shortcutsHelpVisible,
      cleanUiDetailsVisible,
      isFocused,
@@ -2570,7 +2580,6 @@ Logging in with Google... Restarting Gemini CLI to continue.
      customDialog,
      apiKeyDefaultValue,
      authState,
-      copyModeEnabled,
      transientMessage,
      bannerData,
      bannerVisible,
@@ -2757,32 +2766,34 @@ Logging in with Google... Restarting Gemini CLI to continue.

  return (
    <UIStateContext.Provider value={uiState}>
-      <UIActionsContext.Provider value={uiActions}>
-        <ConfigContext.Provider value={config}>
-          <AppContext.Provider
-            value={{
-              version: props.version,
-              startupWarnings: props.startupWarnings || [],
-            }}
-          >
-            <ToolActionsProvider
-              config={config}
-              toolCalls={allToolCalls}
-              isExpanded={isExpanded}
-              toggleExpansion={toggleExpansion}
-              toggleAllExpansion={toggleAllExpansion}
+      <InputContext.Provider value={inputState}>
+        <UIActionsContext.Provider value={uiActions}>
+          <ConfigContext.Provider value={config}>
+            <AppContext.Provider
+              value={{
+                version: props.version,
+                startupWarnings: props.startupWarnings || [],
+              }}
            >
-              <ShellFocusContext.Provider value={isFocused}>
-                <MouseProvider mouseEventsEnabled={mouseMode}>
-                  <ScrollProvider>
-                    <App key={`app-${forceRerenderKey}`} />
-                  </ScrollProvider>
-                </MouseProvider>
-              </ShellFocusContext.Provider>
-            </ToolActionsProvider>
-          </AppContext.Provider>
-        </ConfigContext.Provider>
-      </UIActionsContext.Provider>
+              <ToolActionsProvider
+                config={config}
+                toolCalls={allToolCalls}
+                isExpanded={isExpanded}
+                toggleExpansion={toggleExpansion}
+                toggleAllExpansion={toggleAllExpansion}
+              >
+                <ShellFocusContext.Provider value={isFocused}>
+                  <MouseProvider mouseEventsEnabled={mouseMode}>
+                    <ScrollProvider>
+                      <App key={`app-${forceRerenderKey}`} />
+                    </ScrollProvider>
+                  </MouseProvider>
+                </ShellFocusContext.Provider>
+              </ToolActionsProvider>
+            </AppContext.Provider>
+          </ConfigContext.Provider>
+        </UIActionsContext.Provider>
+      </InputContext.Provider>
    </UIStateContext.Provider>
  );
 };
--- a/packages/cli/src/ui/components/Composer.test.tsx
+++ b/packages/cli/src/ui/components/Composer.test.tsx
@@ -245,20 +245,37 @@ const createMockConfig = (overrides = {}): Config =>
    ...overrides,
  }) as unknown as Config;

+import { InputContext, type InputState } from '../contexts/InputContext.js';
+
 const renderComposer = async (
  uiState: UIState,
  settings = createMockSettings({ ui: {} }),
  config = createMockConfig(),
  uiActions = createMockUIActions(),
+  inputStateOverrides: Partial<InputState> = {},
 ) => {
+  const inputState = {
+    buffer: { text: '' } as unknown as TextBuffer,
+    userMessages: [],
+    shellModeActive: false,
+    showEscapePrompt: false,
+    copyModeEnabled: false,
+    inputWidth: 80,
+    suggestionsWidth: 40,
+    ...(uiState as unknown as Partial<InputState>),
+    ...inputStateOverrides,
+  };
+
  const result = await render(
    <ConfigContext.Provider value={config as unknown as Config}>
      <SettingsContext.Provider value={settings as unknown as LoadedSettings}>
-        <UIStateContext.Provider value={uiState}>
-          <UIActionsContext.Provider value={uiActions}>
-            <Composer isFocused={true} />
-          </UIActionsContext.Provider>
-        </UIStateContext.Provider>
+        <InputContext.Provider value={inputState}>
+          <UIStateContext.Provider value={uiState}>
+            <UIActionsContext.Provider value={uiActions}>
+              <Composer isFocused={true} />
+            </UIActionsContext.Provider>
+          </UIStateContext.Provider>
+        </InputContext.Provider>
      </SettingsContext.Provider>
    </ConfigContext.Provider>,
  );
@@ -541,7 +558,6 @@ describe('Composer', () => {
      const uiState = createMockUIState({
        ctrlCPressedOnce: false,
        ctrlDPressedOnce: false,
-        showEscapePrompt: false,
      });

      const { lastFrame } = await renderComposer(uiState);
@@ -631,7 +647,6 @@ describe('Composer', () => {
      async (mode) => {
        const uiState = createMockUIState({
          showApprovalModeIndicator: mode,
-          shellModeActive: false,
        });

        const { lastFrame } = await renderComposer(uiState);
@@ -641,11 +656,15 @@ describe('Composer', () => {
    );

    it('shows ShellModeIndicator when shell mode is active', async () => {
-      const uiState = createMockUIState({
-        shellModeActive: true,
-      });
+      const uiState = createMockUIState();

-      const { lastFrame } = await renderComposer(uiState);
+      const { lastFrame } = await renderComposer(
+        uiState,
+        undefined,
+        undefined,
+        undefined,
+        { shellModeActive: true },
+      );

      expect(lastFrame()).toMatch(/ShellModeIndic[\s\S]*tor/);
    });
@@ -724,11 +743,16 @@ describe('Composer', () => {
    it('shows Esc rewind prompt in minimal mode without showing full UI', async () => {
      const uiState = createMockUIState({
        cleanUiDetailsVisible: false,
-        showEscapePrompt: true,
        history: [{ id: 1, type: 'user', text: 'msg' }],
      });

-      const { lastFrame } = await renderComposer(uiState);
+      const { lastFrame } = await renderComposer(
+        uiState,
+        undefined,
+        undefined,
+        undefined,
+        { showEscapePrompt: true },
+      );
      const output = lastFrame();
      expect(output).toContain('Press Esc again to rewind.');
      expect(output).not.toContain('ContextSummaryDisplay');
@@ -828,11 +852,16 @@ describe('Composer', () => {
  describe('Shortcuts Hint', () => {
    it('restores shortcuts hint after 200ms debounce when buffer is empty', async () => {
      const uiState = createMockUIState({
-        buffer: { text: '' } as unknown as TextBuffer,
        cleanUiDetailsVisible: false,
      });

-      const { lastFrame } = await renderComposer(uiState);
+      const { lastFrame } = await renderComposer(
+        uiState,
+        undefined,
+        undefined,
+        undefined,
+        { buffer: { text: '' } as unknown as TextBuffer },
+      );

      await act(async () => {
        await vi.advanceTimersByTimeAsync(250);
@@ -845,11 +874,16 @@ describe('Composer', () => {

    it('hides shortcuts hint when text is typed in buffer', async () => {
      const uiState = createMockUIState({
-        buffer: { text: 'hello' } as unknown as TextBuffer,
        cleanUiDetailsVisible: false,
      });

-      const { lastFrame } = await renderComposer(uiState);
+      const { lastFrame } = await renderComposer(
+        uiState,
+        undefined,
+        undefined,
+        undefined,
+        { buffer: { text: 'hello' } as unknown as TextBuffer },
+      );

      expect(lastFrame()).not.toContain('press tab twice for more');
      expect(lastFrame()).not.toContain('? for shortcuts');
--- a/packages/cli/src/ui/components/Composer.tsx
+++ b/packages/cli/src/ui/components/Composer.tsx
@@ -9,6 +9,7 @@ import { useState, useEffect } from 'react';
 import { useConfig } from '../contexts/ConfigContext.js';
 import { useSettings } from '../contexts/SettingsContext.js';
 import { useUIState } from '../contexts/UIStateContext.js';
+import { useInputState } from '../contexts/InputContext.js';
 import { useUIActions } from '../contexts/UIActionsContext.js';
 import { useVimMode } from '../contexts/VimModeContext.js';
 import { useAlternateBuffer } from '../hooks/useAlternateBuffer.js';
@@ -30,6 +31,7 @@ import { appEvents, AppEvent } from '../../utils/events.js';

 export const Composer = ({ isFocused = true }: { isFocused?: boolean }) => {
  const uiState = useUIState();
+  const inputState = useInputState();
  const uiActions = useUIActions();
  const settings = useSettings();
  const config = useConfig();
@@ -81,12 +83,12 @@ export const Composer = ({ isFocused = true }: { isFocused?: boolean }) => {
    return null;
  }

-  const hasToast = shouldShowToast(uiState);
+  const showToast = shouldShowToast(uiState, inputState);
  const hideUiDetailsForSuggestions =
    suggestionsVisible && suggestionsPosition === 'above';

  // Mini Mode VIP Flags (Pure Content Triggers)
-  const showMinimalToast = hasToast;
+  const showMinimalToast = showToast;

  return (
    <Box
@@ -141,17 +143,12 @@ export const Composer = ({ isFocused = true }: { isFocused?: boolean }) => {

      {uiState.isInputActive && (
        <InputPrompt
-          buffer={uiState.buffer}
-          inputWidth={uiState.inputWidth}
-          suggestionsWidth={uiState.suggestionsWidth}
          onSubmit={uiActions.handleFinalSubmit}
-          userMessages={uiState.userMessages}
          setBannerVisible={uiActions.setBannerVisible}
          onClearScreen={uiActions.handleClearScreen}
          config={config}
          slashCommands={uiState.slashCommands || []}
          commandContext={uiState.commandContext}
-          shellModeActive={uiState.shellModeActive}
          setShellModeActive={uiActions.setShellModeActive}
          approvalMode={uiState.showApprovalModeIndicator}
          onEscapePromptChange={uiActions.onEscapePromptChange}
@@ -165,7 +162,7 @@ export const Composer = ({ isFocused = true }: { isFocused?: boolean }) => {
              ? vimMode === 'INSERT'
                ? "  Press 'Esc' for NORMAL mode."
                : "  Press 'i' for INSERT mode."
-              : uiState.shellModeActive
+              : inputState.shellModeActive
                ? '  Type your shell command'
                : '  Type your message or @path/to/file'
          }
@@ -173,7 +170,6 @@ export const Composer = ({ isFocused = true }: { isFocused?: boolean }) => {
          streamingState={uiState.streamingState}
          suggestionsPosition={suggestionsPosition}
          onSuggestionsVisibilityChange={setSuggestionsVisible}
-          copyModeEnabled={uiState.copyModeEnabled}
        />
      )}

--- a/packages/cli/src/ui/components/CopyModeWarning.test.tsx
+++ b/packages/cli/src/ui/components/CopyModeWarning.test.tsx
@@ -4,34 +4,36 @@
 * SPDX-License-Identifier: Apache-2.0
 */

-import { render } from '../../test-utils/render.js';
 import { CopyModeWarning } from './CopyModeWarning.js';
 import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { useUIState, type UIState } from '../contexts/UIStateContext.js';
+import { renderWithProviders } from '../../test-utils/render.js';
+import { useInputState } from '../contexts/InputContext.js';

-vi.mock('../contexts/UIStateContext.js');
+vi.mock('../contexts/InputContext.js');

 describe('CopyModeWarning', () => {
-  const mockUseUIState = vi.mocked(useUIState);
-
  beforeEach(() => {
    vi.clearAllMocks();
  });

  it('renders nothing when copy mode is disabled', async () => {
-    mockUseUIState.mockReturnValue({
+    vi.mocked(useInputState).mockReturnValue({
      copyModeEnabled: false,
-    } as unknown as UIState);
-    const { lastFrame, unmount } = await render(<CopyModeWarning />);
+    } as unknown as ReturnType<typeof useInputState>);
+    const { lastFrame, unmount } = await renderWithProviders(
+      <CopyModeWarning />,
+    );
    expect(lastFrame({ allowEmpty: true })).toBe('');
    unmount();
  });

  it('renders warning when copy mode is enabled', async () => {
-    mockUseUIState.mockReturnValue({
+    vi.mocked(useInputState).mockReturnValue({
      copyModeEnabled: true,
-    } as unknown as UIState);
-    const { lastFrame, unmount } = await render(<CopyModeWarning />);
+    } as unknown as ReturnType<typeof useInputState>);
+    const { lastFrame, unmount } = await renderWithProviders(
+      <CopyModeWarning />,
+    );
    expect(lastFrame()).toContain('In Copy Mode');
    expect(lastFrame()).toContain('Use Page Up/Down to scroll');
    expect(lastFrame()).toContain('Press Ctrl+S or any other key to exit');
--- a/packages/cli/src/ui/components/CopyModeWarning.tsx
+++ b/packages/cli/src/ui/components/CopyModeWarning.tsx
@@ -6,11 +6,11 @@

 import type React from 'react';
 import { Box, Text } from 'ink';
-import { useUIState } from '../contexts/UIStateContext.js';
+import { useInputState } from '../contexts/InputContext.js';
 import { theme } from '../semantic-colors.js';

 export const CopyModeWarning: React.FC = () => {
-  const { copyModeEnabled } = useUIState();
+  const { copyModeEnabled } = useInputState();

  return (
    <Box height={1}>
--- a/packages/cli/src/ui/components/ExitPlanModeDialog.test.tsx
+++ b/packages/cli/src/ui/components/ExitPlanModeDialog.test.tsx
@@ -169,6 +169,11 @@ Implement a comprehensive authentication system with multiple providers.
          getUseTerminalBuffer: () => false,
        } as unknown as import('@google/gemini-cli-core').Config,
        settings: createMockSettings({ ui: { useAlternateBuffer } }),
+        inputState: {
+          buffer: { text: '' } as never,
+          showEscapePrompt: false,
+          shellModeActive: false,
+        },
      },
    );
  };
@@ -472,6 +477,11 @@ Implement a comprehensive authentication system with multiple providers.
              settings: createMockSettings({
                ui: { useAlternateBuffer: useAlternateBuffer ?? true },
              }),
+              inputState: {
+                buffer: { text: '' } as never,
+                showEscapePrompt: false,
+                shellModeActive: false,
+              },
            },
          ),
        );
--- a/packages/cli/src/ui/components/Footer.tsx
+++ b/packages/cli/src/ui/components/Footer.tsx
@@ -26,6 +26,7 @@ import { useUIState } from '../contexts/UIStateContext.js';
 import { useConfig } from '../contexts/ConfigContext.js';
 import { useSettings } from '../contexts/SettingsContext.js';
 import { useVimMode } from '../contexts/VimModeContext.js';
+import { useInputState } from '../contexts/InputContext.js';
 import {
  ALL_ITEMS,
  type FooterItemId,
@@ -173,6 +174,7 @@ interface FooterColumn {

 export const Footer: React.FC = () => {
  const uiState = useUIState();
+  const { copyModeEnabled } = useInputState();
  const config = useConfig();
  const settings = useSettings();
  const { vimEnabled, vimMode } = useVimMode();
@@ -365,10 +367,7 @@ export const Footer: React.FC = () => {
          id,
          header,
          () => (
-            <MemoryUsageDisplay
-              color={itemColor}
-              isActive={!uiState.copyModeEnabled}
-            />
+            <MemoryUsageDisplay color={itemColor} isActive={!copyModeEnabled} />
          ),
          10,
        );
--- a/packages/cli/src/ui/components/InputPrompt.test.tsx
+++ b/packages/cli/src/ui/components/InputPrompt.test.tsx
--- a/packages/cli/src/ui/components/InputPrompt.tsx
+++ b/packages/cli/src/ui/components/InputPrompt.tsx
@@ -70,6 +70,7 @@ import { getSafeLowColorBackground } from '../themes/color-utils.js';
 import { isLowColorDepth } from '../utils/terminalUtils.js';
 import { useShellFocusState } from '../contexts/ShellFocusContext.js';
 import { useUIState } from '../contexts/UIStateContext.js';
+import { useInputState } from '../contexts/InputContext.js';
 import {
  appEvents,
  AppEvent,
@@ -104,18 +105,13 @@ export type ScrollableItem =
  | { type: 'ghostLine'; ghostLine: string; index: number };

 export interface InputPromptProps {
-  buffer: TextBuffer;
  onSubmit: (value: string) => void;
-  userMessages: readonly string[];
  onClearScreen: () => void;
  config: Config;
  slashCommands: readonly SlashCommand[];
  commandContext: CommandContext;
  placeholder?: string;
  focus?: boolean;
-  inputWidth: number;
-  suggestionsWidth: number;
-  shellModeActive: boolean;
  setShellModeActive: (value: boolean) => void;
  approvalMode: ApprovalMode;
  onEscapePromptChange?: (showPrompt: boolean) => void;
@@ -128,7 +124,6 @@ export interface InputPromptProps {
  onQueueMessage?: (message: string) => void;
  suggestionsPosition?: 'above' | 'below';
  setBannerVisible: (visible: boolean) => void;
-  copyModeEnabled?: boolean;
 }

 // The input content, input container, and input suggestions list may have different widths
@@ -199,18 +194,13 @@ export function tryTogglePasteExpansion(buffer: TextBuffer): boolean {
 }

 export const InputPrompt: React.FC<InputPromptProps> = ({
-  buffer,
  onSubmit,
-  userMessages,
  onClearScreen,
  config,
  slashCommands,
  commandContext,
  placeholder = '  Type your message or @path/to/file',
  focus = true,
-  inputWidth,
-  suggestionsWidth,
-  shellModeActive,
  setShellModeActive,
  approvalMode,
  onEscapePromptChange,
@@ -223,8 +213,16 @@ export const InputPrompt: React.FC<InputPromptProps> = ({
  onQueueMessage,
  suggestionsPosition = 'below',
  setBannerVisible,
-  copyModeEnabled = false,
 }) => {
+  const inputState = useInputState();
+  const {
+    buffer,
+    userMessages,
+    shellModeActive,
+    copyModeEnabled,
+    inputWidth,
+    suggestionsWidth,
+  } = inputState;
  const isHelpDismissKey = useIsHelpDismissKey();
  const keyMatchers = useKeyMatchers();
  const { stdout } = useStdout();
--- a/packages/cli/src/ui/components/StatusRow.test.tsx
+++ b/packages/cli/src/ui/components/StatusRow.test.tsx
@@ -9,7 +9,7 @@ import { StatusRow } from './StatusRow.js';
 import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
 import { useComposerStatus } from '../hooks/useComposerStatus.js';
 import { type UIState } from '../contexts/UIStateContext.js';
-import { type TextBuffer } from '../components/shared/text-buffer.js';
+
 import { type SessionStatsState } from '../contexts/SessionContext.js';
 import { type ThoughtSummary } from '../types.js';
 import { ApprovalMode } from '@google/gemini-cli-core';
@@ -29,13 +29,11 @@ describe('<StatusRow />', () => {
    elapsedTime: 0,
    currentWittyPhrase: undefined,
    activeHooks: [],
-    buffer: { text: '' } as unknown as TextBuffer,
    sessionStats: { lastPromptTokenCount: 0 } as unknown as SessionStatsState,
    shortcutsHelpVisible: false,
    contextFileNames: [],
    showApprovalModeIndicator: ApprovalMode.DEFAULT,
    allowPlanMode: false,
-    shellModeActive: false,
    renderMarkdown: true,
    currentModel: 'gemini-3',
  };
--- a/packages/cli/src/ui/components/StatusRow.tsx
+++ b/packages/cli/src/ui/components/StatusRow.tsx
@@ -153,6 +153,8 @@ export const StatusNode: React.FC<{
  );
 };

+import { useInputState } from '../contexts/InputContext.js';
+
 export const StatusRow: React.FC<StatusRowProps> = ({
  showUiDetails,
  isNarrow,
@@ -162,6 +164,7 @@ export const StatusRow: React.FC<StatusRowProps> = ({
  hasPendingActionRequired,
 }) => {
  const uiState = useUIState();
+  const inputState = useInputState();
  const settings = useSettings();
  const {
    isInteractiveShellWaiting,
@@ -225,7 +228,7 @@ export const StatusRow: React.FC<StatusRowProps> = ({
      settings.merged.ui.showShortcutsHint &&
      !hideUiDetailsForSuggestions &&
      !hasPendingActionRequired &&
-      uiState.buffer.text.length === 0
+      inputState.buffer.text.length === 0
    ) {
      return showUiDetails ? '? for shortcuts' : 'press tab twice for more';
    }
@@ -391,13 +394,14 @@ export const StatusRow: React.FC<StatusRowProps> = ({
          >
            {showUiDetails ? (
              <>
-                {!hideUiDetailsForSuggestions && !uiState.shellModeActive && (
-                  <ApprovalModeIndicator
-                    approvalMode={uiState.showApprovalModeIndicator}
-                    allowPlanMode={uiState.allowPlanMode}
-                  />
-                )}
-                {uiState.shellModeActive && (
+                {!hideUiDetailsForSuggestions &&
+                  !inputState.shellModeActive && (
+                    <ApprovalModeIndicator
+                      approvalMode={uiState.showApprovalModeIndicator}
+                      allowPlanMode={uiState.allowPlanMode}
+                    />
+                  )}
+                {inputState.shellModeActive && (
                  <Box marginLeft={LAYOUT.INDICATOR_LEFT_MARGIN}>
                    <ShellModeIndicator />
                  </Box>
--- a/packages/cli/src/ui/components/ToastDisplay.test.tsx
+++ b/packages/cli/src/ui/components/ToastDisplay.test.tsx
@@ -9,16 +9,24 @@ import { renderWithProviders } from '../../test-utils/render.js';
 import { ToastDisplay, shouldShowToast } from './ToastDisplay.js';
 import { TransientMessageType } from '../../utils/events.js';
 import { type UIState } from '../contexts/UIStateContext.js';
+import { type InputState } from '../contexts/InputContext.js';
 import { type TextBuffer } from './shared/text-buffer.js';
 import { type HistoryItem } from '../types.js';

-const renderToastDisplay = async (uiState: Partial<UIState> = {}) =>
+const renderToastDisplay = async (
+  uiState: Partial<UIState> = {},
+  inputState: Partial<InputState> = {},
+) =>
  renderWithProviders(<ToastDisplay />, {
    uiState: {
-      buffer: { text: '' } as TextBuffer,
      history: [] as HistoryItem[],
      ...uiState,
    },
+    inputState: {
+      buffer: { text: '' } as TextBuffer,
+      showEscapePrompt: false,
+      ...inputState,
+    },
  });

 describe('ToastDisplay', () => {
@@ -27,86 +35,121 @@ describe('ToastDisplay', () => {
  });

  describe('shouldShowToast', () => {
-    const baseState: Partial<UIState> = {
+    const baseUIState: Partial<UIState> = {
      ctrlCPressedOnce: false,
      transientMessage: null,
      ctrlDPressedOnce: false,
-      showEscapePrompt: false,
-      buffer: { text: '' } as TextBuffer,
      history: [] as HistoryItem[],
      queueErrorMessage: null,
      showIsExpandableHint: false,
    };

+    const baseInputState: Partial<InputState> = {
+      showEscapePrompt: false,
+      buffer: { text: '' } as TextBuffer,
+    };
+
    it('returns false for default state', () => {
-      expect(shouldShowToast(baseState as UIState)).toBe(false);
+      expect(
+        shouldShowToast(baseUIState as UIState, baseInputState as InputState),
+      ).toBe(false);
    });

    it('returns true when showIsExpandableHint is true', () => {
      expect(
-        shouldShowToast({
-          ...baseState,
-          showIsExpandableHint: true,
-        } as UIState),
+        shouldShowToast(
+          {
+            ...baseUIState,
+            showIsExpandableHint: true,
+          } as UIState,
+          baseInputState as InputState,
+        ),
      ).toBe(true);
    });

    it('returns true when ctrlCPressedOnce is true', () => {
      expect(
-        shouldShowToast({ ...baseState, ctrlCPressedOnce: true } as UIState),
+        shouldShowToast(
+          { ...baseUIState, ctrlCPressedOnce: true } as UIState,
+          baseInputState as InputState,
+        ),
      ).toBe(true);
    });

    it('returns true when transientMessage is present', () => {
      expect(
-        shouldShowToast({
-          ...baseState,
-          transientMessage: { text: 'test', type: TransientMessageType.Hint },
-        } as UIState),
+        shouldShowToast(
+          {
+            ...baseUIState,
+            transientMessage: { text: 'test', type: TransientMessageType.Hint },
+          } as UIState,
+          baseInputState as InputState,
+        ),
      ).toBe(true);
    });

    it('returns true when ctrlDPressedOnce is true', () => {
      expect(
-        shouldShowToast({ ...baseState, ctrlDPressedOnce: true } as UIState),
+        shouldShowToast(
+          { ...baseUIState, ctrlDPressedOnce: true } as UIState,
+          baseInputState as InputState,
+        ),
      ).toBe(true);
    });

    it('returns true when showEscapePrompt is true and buffer is NOT empty', () => {
      expect(
-        shouldShowToast({
-          ...baseState,
-          showEscapePrompt: true,
-          buffer: { text: 'some text' } as TextBuffer,
-        } as UIState),
+        shouldShowToast(
+          {
+            ...baseUIState,
+          } as UIState,
+          {
+            ...baseInputState,
+            showEscapePrompt: true,
+            buffer: { text: 'some text' } as TextBuffer,
+          } as InputState,
+        ),
      ).toBe(true);
    });

    it('returns true when showEscapePrompt is true and history is NOT empty', () => {
      expect(
-        shouldShowToast({
-          ...baseState,
-          showEscapePrompt: true,
-          history: [{ id: '1' } as unknown as HistoryItem],
-        } as UIState),
+        shouldShowToast(
+          {
+            ...baseUIState,
+            history: [{ id: '1' } as unknown as HistoryItem],
+          } as UIState,
+          {
+            ...baseInputState,
+            showEscapePrompt: true,
+          } as InputState,
+        ),
      ).toBe(true);
    });

    it('returns false when showEscapePrompt is true but buffer and history are empty', () => {
      expect(
-        shouldShowToast({
-          ...baseState,
-          showEscapePrompt: true,
-        } as UIState),
+        shouldShowToast(
+          {
+            ...baseUIState,
+          } as UIState,
+          {
+            ...baseInputState,
+            showEscapePrompt: true,
+          } as InputState,
+        ),
      ).toBe(false);
    });

    it('returns true when queueErrorMessage is present', () => {
      expect(
-        shouldShowToast({
-          ...baseState,
-          queueErrorMessage: 'error',
-        } as UIState),
+        shouldShowToast(
+          {
+            ...baseUIState,
+            queueErrorMessage: 'error',
+          } as UIState,
+          baseInputState as InputState,
+        ),
      ).toBe(true);
    });
  });
@@ -151,18 +194,25 @@ describe('ToastDisplay', () => {
  });

  it('renders Escape prompt when buffer is empty', async () => {
-    const { lastFrame } = await renderToastDisplay({
-      showEscapePrompt: true,
-      history: [{ id: 1, type: 'user', text: 'test' }] as HistoryItem[],
-    });
+    const { lastFrame } = await renderToastDisplay(
+      {
+        history: [{ id: 1, type: 'user', text: 'test' }] as HistoryItem[],
+      },
+      {
+        showEscapePrompt: true,
+      },
+    );
    expect(lastFrame()).toMatchSnapshot();
  });

  it('renders Escape prompt when buffer is NOT empty', async () => {
-    const { lastFrame } = await renderToastDisplay({
-      showEscapePrompt: true,
-      buffer: { text: 'some text' } as TextBuffer,
-    });
+    const { lastFrame } = await renderToastDisplay(
+      {},
+      {
+        showEscapePrompt: true,
+        buffer: { text: 'some text' } as TextBuffer,
+      },
+    );
    expect(lastFrame()).toMatchSnapshot();
  });

--- a/packages/cli/src/ui/components/ToastDisplay.tsx
+++ b/packages/cli/src/ui/components/ToastDisplay.tsx
@@ -8,15 +8,19 @@ import type React from 'react';
 import { Text } from 'ink';
 import { theme } from '../semantic-colors.js';
 import { useUIState, type UIState } from '../contexts/UIStateContext.js';
+import { useInputState, type InputState } from '../contexts/InputContext.js';
 import { TransientMessageType } from '../../utils/events.js';

-export function shouldShowToast(uiState: UIState): boolean {
+export function shouldShowToast(
+  uiState: UIState,
+  inputState: InputState,
+): boolean {
  return (
    uiState.ctrlCPressedOnce ||
    Boolean(uiState.transientMessage) ||
    uiState.ctrlDPressedOnce ||
-    (uiState.showEscapePrompt &&
-      (uiState.buffer.text.length > 0 || uiState.history.length > 0)) ||
+    (inputState.showEscapePrompt &&
+      (inputState.buffer.text.length > 0 || uiState.history.length > 0)) ||
    Boolean(uiState.queueErrorMessage) ||
    uiState.showIsExpandableHint
  );
@@ -24,6 +28,7 @@ export function shouldShowToast(uiState: UIState): boolean {

 export const ToastDisplay: React.FC = () => {
  const uiState = useUIState();
+  const inputState = useInputState();

  if (uiState.ctrlCPressedOnce) {
    return (
@@ -46,8 +51,8 @@ export const ToastDisplay: React.FC = () => {
    );
  }

-  if (uiState.showEscapePrompt) {
-    const isPromptEmpty = uiState.buffer.text.length === 0;
+  if (inputState.showEscapePrompt) {
+    const isPromptEmpty = inputState.buffer.text.length === 0;
    const hasHistory = uiState.history.length > 0;

    if (isPromptEmpty && !hasHistory) {
--- a/packages/cli/src/ui/components/shared/VirtualizedList.tsx
+++ b/packages/cli/src/ui/components/shared/VirtualizedList.tsx
@@ -85,32 +85,43 @@ const VirtualizedListItem = memo(
    width,
    containerWidth,
    itemKey,
-    itemRef,
+    index,
+    onSetRef,
  }: {
    content: React.ReactElement;
    shouldBeStatic: boolean;
    width: number | string | undefined;
    containerWidth: number;
    itemKey: string;
-    itemRef: (el: DOMElement | null) => void;
-  }) => (
-    <Box width="100%" flexDirection="column" flexShrink={0} ref={itemRef}>
-      {shouldBeStatic ? (
-        <StaticRender
-          width={typeof width === 'number' ? width : containerWidth}
-          key={
-            itemKey +
-            '-static-' +
-            (typeof width === 'number' ? width : containerWidth)
-          }
-        >
-          {content}
-        </StaticRender>
-      ) : (
-        content
-      )}
-    </Box>
-  ),
+    index: number;
+    onSetRef: (index: number, el: DOMElement | null) => void;
+  }) => {
+    const itemRef = useCallback(
+      (el: DOMElement | null) => {
+        onSetRef(index, el);
+      },
+      [index, onSetRef],
+    );
+
+    return (
+      <Box width="100%" flexDirection="column" flexShrink={0} ref={itemRef}>
+        {shouldBeStatic ? (
+          <StaticRender
+            width={typeof width === 'number' ? width : containerWidth}
+            key={
+              itemKey +
+              '-static-' +
+              (typeof width === 'number' ? width : containerWidth)
+            }
+          >
+            {content}
+          </StaticRender>
+        ) : (
+          content
+        )}
+      </Box>
+    );
+  },
 );

 VirtualizedListItem.displayName = 'VirtualizedListItem';
@@ -195,6 +206,10 @@ function VirtualizedList<T>(
  const containerObserverRef = useRef<ResizeObserver | null>(null);
  const nodeToKeyRef = useRef(new WeakMap<DOMElement, string>());

+  const onSetRef = useCallback((index: number, el: DOMElement | null) => {
+    itemRefs.current[index] = el;
+  }, []);
+
  const containerRefCallback = useCallback((node: DOMElement | null) => {
    containerObserverRef.current?.disconnect();
    containerRef.current = node;
@@ -517,7 +532,6 @@ function VirtualizedList<T>(
    observedNodes.current = currentNodes;
  });

-  const renderedItems = [];
  const renderRangeStart =
    renderStatic || overflowToBackbuffer ? 0 : startIndex;
  const renderRangeEnd = renderStatic ? data.length - 1 : endIndex;
@@ -533,7 +547,12 @@ function VirtualizedList<T>(
    process.env['NODE_ENV'] === 'test' ||
    (width !== undefined && typeof width === 'number');

-  if (isReady) {
+  const renderedItems = useMemo(() => {
+    if (!isReady) {
+      return [];
+    }
+
+    const items = [];
    for (let i = renderRangeStart; i <= renderRangeEnd; i++) {
      const item = data[i];
      if (item) {
@@ -545,7 +564,7 @@ function VirtualizedList<T>(
        const content = renderItem({ item, index: i });
        const key = keyExtractor(item, i);

-        renderedItems.push(
+        items.push(
          <VirtualizedListItem
            key={key}
            itemKey={key}
@@ -553,16 +572,28 @@ function VirtualizedList<T>(
            shouldBeStatic={shouldBeStatic}
            width={width}
            containerWidth={containerWidth}
-            itemRef={(el: DOMElement | null) => {
-              if (i >= renderRangeStart && i <= renderRangeEnd) {
-                itemRefs.current[i] = el;
-              }
-            }}
+            index={i}
+            onSetRef={onSetRef}
          />,
        );
      }
    }
-  }
+    return items;
+  }, [
+    isReady,
+    renderRangeStart,
+    renderRangeEnd,
+    data,
+    startIndex,
+    endIndex,
+    renderStatic,
+    isStaticItem,
+    renderItem,
+    keyExtractor,
+    width,
+    containerWidth,
+    onSetRef,
+  ]);

  const { getScrollTop, setPendingScrollTop } = useBatchedScroll(scrollTop);

--- a/packages/cli/src/ui/contexts/InputContext.tsx
+++ b/packages/cli/src/ui/contexts/InputContext.tsx
@@ -0,0 +1,28 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { createContext, useContext } from 'react';
+import type { TextBuffer } from '../components/shared/text-buffer.js';
+
+export interface InputState {
+  buffer: TextBuffer;
+  userMessages: string[];
+  shellModeActive: boolean;
+  showEscapePrompt: boolean;
+  copyModeEnabled: boolean | undefined;
+  inputWidth: number;
+  suggestionsWidth: number;
+}
+
+export const InputContext = createContext<InputState | null>(null);
+
+export const useInputState = () => {
+  const context = useContext(InputContext);
+  if (!context) {
+    throw new Error('useInputState must be used within an InputProvider');
+  }
+  return context;
+};
--- a/packages/cli/src/ui/contexts/UIStateContext.tsx
+++ b/packages/cli/src/ui/contexts/UIStateContext.tsx
@@ -17,7 +17,7 @@ import type {
  PermissionConfirmationRequest,
 } from '../types.js';
 import type { CommandContext, SlashCommand } from '../commands/types.js';
-import type { TextBuffer } from '../components/shared/text-buffer.js';
+
 import type {
  IdeContext,
  ApprovalMode,
@@ -143,11 +143,6 @@ export interface UIState {
  initError: string | null;
  pendingGeminiHistoryItems: HistoryItemWithoutId[];
  thought: ThoughtSummary | null;
-  shellModeActive: boolean;
-  userMessages: string[];
-  buffer: TextBuffer;
-  inputWidth: number;
-  suggestionsWidth: number;
  isInputActive: boolean;
  isResuming: boolean;
  shouldShowIdePrompt: boolean;
@@ -162,7 +157,6 @@ export interface UIState {
  renderMarkdown: boolean;
  ctrlCPressedOnce: boolean;
  ctrlDPressedOnce: boolean;
-  showEscapePrompt: boolean;
  shortcutsHelpVisible: boolean;
  cleanUiDetailsVisible: boolean;
  elapsedTime: number;
@@ -207,7 +201,6 @@ export interface UIState {
  embeddedShellFocused: boolean;
  showDebugProfiler: boolean;
  showFullTodos: boolean;
-  copyModeEnabled: boolean;
  bannerData: {
    defaultText: string;
    warningText: string;
--- a/packages/cli/src/ui/layouts/DefaultAppLayout.test.tsx
+++ b/packages/cli/src/ui/layouts/DefaultAppLayout.test.tsx
@@ -5,8 +5,11 @@
 */

 import { render } from '../../test-utils/render.js';
-import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { DefaultAppLayout } from './DefaultAppLayout.js';
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { useInputState } from '../contexts/InputContext.js';
+
+vi.mock('../contexts/InputContext.js');
 import { StreamingState } from '../types.js';
 import { Text } from 'ink';
 import type { UIState } from '../contexts/UIStateContext.js';
@@ -95,6 +98,9 @@ const createMockShell = (pid: number): BackgroundTask => ({
 describe('<DefaultAppLayout />', () => {
  beforeEach(() => {
    vi.clearAllMocks();
+    vi.mocked(useInputState).mockReturnValue({
+      copyModeEnabled: false,
+    } as unknown as ReturnType<typeof useInputState>);
    // Reset mock state defaults
    mockUIState.backgroundTasks = new Map();
    mockUIState.activeBackgroundTaskPid = null;
--- a/packages/cli/src/ui/layouts/DefaultAppLayout.tsx
+++ b/packages/cli/src/ui/layouts/DefaultAppLayout.tsx
@@ -17,9 +17,11 @@ import { useAlternateBuffer } from '../hooks/useAlternateBuffer.js';
 import { CopyModeWarning } from '../components/CopyModeWarning.js';
 import { BackgroundTaskDisplay } from '../components/BackgroundTaskDisplay.js';
 import { StreamingState } from '../types.js';
+import { useInputState } from '../contexts/InputContext.js';

 export const DefaultAppLayout: React.FC = () => {
  const uiState = useUIState();
+  const { copyModeEnabled } = useInputState();
  const isAlternateBuffer = useAlternateBuffer();

  const { rootUiRef, terminalHeight } = uiState;
@@ -62,9 +64,7 @@ export const DefaultAppLayout: React.FC = () => {
        flexShrink={0}
        flexGrow={0}
        width={uiState.terminalWidth}
-        height={
-          uiState.copyModeEnabled ? uiState.stableControlsHeight : undefined
-        }
+        height={copyModeEnabled ? uiState.stableControlsHeight : undefined}
      >
        <Notifications />
        <CopyModeWarning />