refactor: Centralize and improve model fallback handling

packages/a2a-server/src/agent/task.ts

@@ -88,12 +88,11 @@ export class Task {
     this.eventBus = eventBus;
     this.completedToolCalls = [];
     this._resetToolCompletionPromise();
-    this.config.setFlashFallbackHandler(
-      async (currentModel: string, fallbackModel: string): Promise<boolean> => {
-        config.setModel(fallbackModel); // gemini-cli-core sets to DEFAULT_GEMINI_FLASH_MODEL
-        // Switch model for future use but return false to stop current retry
-        return false;
-      },
+    this.config.setFallbackModelHandler(
+      // For a2a-server, we want to automatically switch to the fallback model
+      // for future requests without retrying the current one. The 'stop'
+      // intent achieves this.
+      async () => 'stop',
     );
   }
 

packages/a2a-server/src/utils/testing_utils.ts

@@ -43,7 +43,7 @@ export function createMockConfig(
     getContentGeneratorConfig: vi.fn().mockReturnValue({ model: 'gemini-pro' }),
     getModel: vi.fn().mockReturnValue('gemini-pro'),
     getUsageStatisticsEnabled: vi.fn().mockReturnValue(false),
-    setFlashFallbackHandler: vi.fn(),
+    setFallbackModelHandler: vi.fn(),
     initialize: vi.fn().mockResolvedValue(undefined),
     getProxy: vi.fn().mockReturnValue(undefined),
     getHistory: vi.fn().mockReturnValue([]),

packages/cli/src/ui/AppContainer.test.tsx

@@ -4,31 +4,55 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import {
+  describe,
+  it,
+  expect,
+  vi,
+  beforeEach,
+  afterEach,
+  type Mock,
+} from 'vitest';
 import { render, cleanup } from 'ink-testing-library';
 import { AppContainer } from './AppContainer.js';
 import { type Config, makeFakeConfig } from '@google/gemini-cli-core';
 import type { LoadedSettings } from '../config/settings.js';
 import type { InitializationResult } from '../core/initializer.js';
+import { useQuotaAndFallback } from './hooks/useQuotaAndFallback.js';
+import { UIStateContext, type UIState } from './contexts/UIStateContext.js';
+import {
+  UIActionsContext,
+  type UIActions,
+} from './contexts/UIActionsContext.js';
+import { useContext } from 'react';
+
+// Helper component will read the context values provided by AppContainer
+// so we can assert against them in our tests.
+let capturedUIState: UIState;
+let capturedUIActions: UIActions;
+function TestContextConsumer() {
+  capturedUIState = useContext(UIStateContext)!;
+  capturedUIActions = useContext(UIActionsContext)!;
+  return null;
+}
 
-// Mock App component to isolate AppContainer testing
 vi.mock('./App.js', () => ({
-  App: () => 'App Component',
+  App: TestContextConsumer,
 }));
 
-// Mock all the hooks and utilities
-vi.mock('./hooks/useHistory.js');
+vi.mock('./hooks/useQuotaAndFallback.js');
+vi.mock('./hooks/useHistoryManager.js');
 vi.mock('./hooks/useThemeCommand.js');
-vi.mock('./hooks/useAuthCommand.js');
+vi.mock('./auth/useAuth.js');
 vi.mock('./hooks/useEditorSettings.js');
 vi.mock('./hooks/useSettingsCommand.js');
-vi.mock('./hooks/useSlashCommandProcessor.js');
+vi.mock('./hooks/slashCommandProcessor.js');
 vi.mock('./hooks/useConsoleMessages.js');
 vi.mock('./hooks/useTerminalSize.js', () => ({
   useTerminalSize: vi.fn(() => ({ columns: 80, rows: 24 })),
 }));
 vi.mock('./hooks/useGeminiStream.js');
-vi.mock('./hooks/useVim.js');
+vi.mock('./hooks/vim.js');
 vi.mock('./hooks/useFocus.js');
 vi.mock('./hooks/useBracketedPaste.js');
 vi.mock('./hooks/useKeypress.js');
@@ -40,7 +64,7 @@ vi.mock('./hooks/useWorkspaceMigration.js');
 vi.mock('./hooks/useGitBranchName.js');
 vi.mock('./contexts/VimModeContext.js');
 vi.mock('./contexts/SessionContext.js');
-vi.mock('./hooks/useTextBuffer.js');
+vi.mock('./components/shared/text-buffer.js');
 vi.mock('./hooks/useLogger.js');
 
 // Mock external utilities
@@ -49,14 +73,153 @@ vi.mock('../utils/handleAutoUpdate.js');
 vi.mock('./utils/ConsolePatcher.js');
 vi.mock('../utils/cleanup.js');
 
+import { useHistory } from './hooks/useHistoryManager.js';
+import { useThemeCommand } from './hooks/useThemeCommand.js';
+import { useAuthCommand } from './auth/useAuth.js';
+import { useEditorSettings } from './hooks/useEditorSettings.js';
+import { useSettingsCommand } from './hooks/useSettingsCommand.js';
+import { useSlashCommandProcessor } from './hooks/slashCommandProcessor.js';
+import { useConsoleMessages } from './hooks/useConsoleMessages.js';
+import { useGeminiStream } from './hooks/useGeminiStream.js';
+import { useVim } from './hooks/vim.js';
+import { useFolderTrust } from './hooks/useFolderTrust.js';
+import { useMessageQueue } from './hooks/useMessageQueue.js';
+import { useAutoAcceptIndicator } from './hooks/useAutoAcceptIndicator.js';
+import { useWorkspaceMigration } from './hooks/useWorkspaceMigration.js';
+import { useGitBranchName } from './hooks/useGitBranchName.js';
+import { useVimMode } from './contexts/VimModeContext.js';
+import { useSessionStats } from './contexts/SessionContext.js';
+import { useTextBuffer } from './components/shared/text-buffer.js';
+import { useLogger } from './hooks/useLogger.js';
+import { useLoadingIndicator } from './hooks/useLoadingIndicator.js';
+
 describe('AppContainer State Management', () => {
   let mockConfig: Config;
   let mockSettings: LoadedSettings;
   let mockInitResult: InitializationResult;
 
+  // Create typed mocks for all hooks
+  const mockedUseQuotaAndFallback = useQuotaAndFallback as Mock;
+  const mockedUseHistory = useHistory as Mock;
+  const mockedUseThemeCommand = useThemeCommand as Mock;
+  const mockedUseAuthCommand = useAuthCommand as Mock;
+  const mockedUseEditorSettings = useEditorSettings as Mock;
+  const mockedUseSettingsCommand = useSettingsCommand as Mock;
+  const mockedUseSlashCommandProcessor = useSlashCommandProcessor as Mock;
+  const mockedUseConsoleMessages = useConsoleMessages as Mock;
+  const mockedUseGeminiStream = useGeminiStream as Mock;
+  const mockedUseVim = useVim as Mock;
+  const mockedUseFolderTrust = useFolderTrust as Mock;
+  const mockedUseMessageQueue = useMessageQueue as Mock;
+  const mockedUseAutoAcceptIndicator = useAutoAcceptIndicator as Mock;
+  const mockedUseWorkspaceMigration = useWorkspaceMigration as Mock;
+  const mockedUseGitBranchName = useGitBranchName as Mock;
+  const mockedUseVimMode = useVimMode as Mock;
+  const mockedUseSessionStats = useSessionStats as Mock;
+  const mockedUseTextBuffer = useTextBuffer as Mock;
+  const mockedUseLogger = useLogger as Mock;
+  const mockedUseLoadingIndicator = useLoadingIndicator as Mock;
+
   beforeEach(() => {
     vi.clearAllMocks();
 
+    capturedUIState = null!;
+    capturedUIActions = null!;
+
+    // **Provide a default return value for EVERY mocked hook.**
+    mockedUseQuotaAndFallback.mockReturnValue({
+      proQuotaRequest: null,
+      handleProQuotaChoice: vi.fn(),
+    });
+    mockedUseHistory.mockReturnValue({
+      history: [],
+      addItem: vi.fn(),
+      updateItem: vi.fn(),
+      clearItems: vi.fn(),
+      loadHistory: vi.fn(),
+    });
+    mockedUseThemeCommand.mockReturnValue({
+      isThemeDialogOpen: false,
+      openThemeDialog: vi.fn(),
+      handleThemeSelect: vi.fn(),
+      handleThemeHighlight: vi.fn(),
+    });
+    mockedUseAuthCommand.mockReturnValue({
+      authState: 'authenticated',
+      setAuthState: vi.fn(),
+      authError: null,
+      onAuthError: vi.fn(),
+    });
+    mockedUseEditorSettings.mockReturnValue({
+      isEditorDialogOpen: false,
+      openEditorDialog: vi.fn(),
+      handleEditorSelect: vi.fn(),
+      exitEditorDialog: vi.fn(),
+    });
+    mockedUseSettingsCommand.mockReturnValue({
+      isSettingsDialogOpen: false,
+      openSettingsDialog: vi.fn(),
+      closeSettingsDialog: vi.fn(),
+    });
+    mockedUseSlashCommandProcessor.mockReturnValue({
+      handleSlashCommand: vi.fn(),
+      slashCommands: [],
+      pendingHistoryItems: [],
+      commandContext: {},
+      shellConfirmationRequest: null,
+      confirmationRequest: null,
+    });
+    mockedUseConsoleMessages.mockReturnValue({
+      consoleMessages: [],
+      handleNewMessage: vi.fn(),
+      clearConsoleMessages: vi.fn(),
+    });
+    mockedUseGeminiStream.mockReturnValue({
+      streamingState: 'idle',
+      submitQuery: vi.fn(),
+      initError: null,
+      pendingHistoryItems: [],
+      thought: null,
+      cancelOngoingRequest: vi.fn(),
+    });
+    mockedUseVim.mockReturnValue({ handleInput: vi.fn() });
+    mockedUseFolderTrust.mockReturnValue({
+      isFolderTrustDialogOpen: false,
+      handleFolderTrustSelect: vi.fn(),
+      isRestarting: false,
+    });
+    mockedUseMessageQueue.mockReturnValue({
+      messageQueue: [],
+      addMessage: vi.fn(),
+      clearQueue: vi.fn(),
+      getQueuedMessagesText: vi.fn().mockReturnValue(''),
+    });
+    mockedUseAutoAcceptIndicator.mockReturnValue(false);
+    mockedUseWorkspaceMigration.mockReturnValue({
+      showWorkspaceMigrationDialog: false,
+      workspaceExtensions: [],
+      onWorkspaceMigrationDialogOpen: vi.fn(),
+      onWorkspaceMigrationDialogClose: vi.fn(),
+    });
+    mockedUseGitBranchName.mockReturnValue('main');
+    mockedUseVimMode.mockReturnValue({
+      isVimEnabled: false,
+      toggleVimEnabled: vi.fn(),
+    });
+    mockedUseSessionStats.mockReturnValue({ stats: {} });
+    mockedUseTextBuffer.mockReturnValue({
+      text: '',
+      setText: vi.fn(),
+      // Add other properties if AppContainer uses them
+    });
+    mockedUseLogger.mockReturnValue({
+      getPreviousUserMessages: vi.fn().mockResolvedValue([]),
+    });
+    mockedUseLoadingIndicator.mockReturnValue({
+      elapsedTime: '0.0s',
+      currentLoadingPhrase: '',
+    });
+
     // Mock Config
     mockConfig = makeFakeConfig();
 
@@ -325,7 +488,73 @@ describe('AppContainer State Management', () => {
       expect(() => unmount()).not.toThrow();
     });
   });
-});
 
-// TODO: Add comprehensive integration test once all hook mocks are complete
-// For now, the 14 passing unit tests provide good coverage of AppContainer functionality
+  describe('Quota and Fallback Integration', () => {
+    it('passes a null proQuotaRequest to UIStateContext by default', () => {
+      // The default mock from beforeEach already sets proQuotaRequest to null
+      render(
+        <AppContainer
+          config={mockConfig}
+          settings={mockSettings}
+          version="1.0.0"
+          initializationResult={mockInitResult}
+        />,
+      );
+
+      // Assert that the context value is as expected
+      expect(capturedUIState.proQuotaRequest).toBeNull();
+    });
+
+    it('passes a valid proQuotaRequest to UIStateContext when provided by the hook', () => {
+      // Arrange: Create a mock request object that a UI dialog would receive
+      const mockRequest = {
+        failedModel: 'gemini-pro',
+        fallbackModel: 'gemini-flash',
+        resolve: vi.fn(),
+      };
+      mockedUseQuotaAndFallback.mockReturnValue({
+        proQuotaRequest: mockRequest,
+        handleProQuotaChoice: vi.fn(),
+      });
+
+      // Act: Render the container
+      render(
+        <AppContainer
+          config={mockConfig}
+          settings={mockSettings}
+          version="1.0.0"
+          initializationResult={mockInitResult}
+        />,
+      );
+
+      // Assert: The mock request is correctly passed through the context
+      expect(capturedUIState.proQuotaRequest).toEqual(mockRequest);
+    });
+
+    it('passes the handleProQuotaChoice function to UIActionsContext', () => {
+      // Arrange: Create a mock handler function
+      const mockHandler = vi.fn();
+      mockedUseQuotaAndFallback.mockReturnValue({
+        proQuotaRequest: null,
+        handleProQuotaChoice: mockHandler,
+      });
+
+      // Act: Render the container
+      render(
+        <AppContainer
+          config={mockConfig}
+          settings={mockSettings}
+          version="1.0.0"
+          initializationResult={mockInitResult}
+        />,
+      );
+
+      // 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
+      capturedUIActions.handleProQuotaChoice('auth');
+      expect(mockHandler).toHaveBeenCalledWith('auth');
+    });
+  });
+});