feat(models): support Gemini 3.1 custom tool model (#131)

* feat(models): add support for Gemini 3.1 and custom tool models

* test(routing): fix classifier and numerical classifier strategy tests

* test(routing): add Gemini 3.1 tests for classifier strategy

* fix(models): correctly filter active Gemini 3.1 models

* fix(routing): ensure useCustomToolModel is only true when Gemini 3.1 is enabled

* fix(test-utils): prevent double newline in lastFrame() on Windows

* fix(test-utils): surgically fix double newline in lastFrame() on Windows

* use custom_tools_model string for api key only

* fix(ui): correct useCustomToolModel logic and update tests

* fix(ui): correct useCustomToolModel logic in StatsDisplay

* fix(routing): ensure test models are active and sync useCustomToolModel logic

packages/core/src/config/models.test.ts

@@ -24,8 +24,40 @@ import {
   PREVIEW_GEMINI_MODEL_AUTO,
   DEFAULT_GEMINI_MODEL_AUTO,
   isActiveModel,
+  PREVIEW_GEMINI_3_1_MODEL,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
+  isPreviewModel,
+  isProModel,
 } from './models.js';
 
+describe('isPreviewModel', () => {
+  it('should return true for preview models', () => {
+    expect(isPreviewModel(PREVIEW_GEMINI_MODEL)).toBe(true);
+    expect(isPreviewModel(PREVIEW_GEMINI_3_1_MODEL)).toBe(true);
+    expect(isPreviewModel(PREVIEW_GEMINI_FLASH_MODEL)).toBe(true);
+    expect(isPreviewModel(PREVIEW_GEMINI_MODEL_AUTO)).toBe(true);
+  });
+
+  it('should return false for non-preview models', () => {
+    expect(isPreviewModel(DEFAULT_GEMINI_MODEL)).toBe(false);
+    expect(isPreviewModel('gemini-1.5-pro')).toBe(false);
+  });
+});
+
+describe('isProModel', () => {
+  it('should return true for models containing "pro"', () => {
+    expect(isProModel('gemini-3-pro-preview')).toBe(true);
+    expect(isProModel('gemini-2.5-pro')).toBe(true);
+    expect(isProModel('pro')).toBe(true);
+  });
+
+  it('should return false for models without "pro"', () => {
+    expect(isProModel('gemini-3-flash-preview')).toBe(false);
+    expect(isProModel('gemini-2.5-flash')).toBe(false);
+    expect(isProModel('auto')).toBe(false);
+  });
+});
+
 describe('isGemini3Model', () => {
   it('should return true for gemini-3 models', () => {
     expect(isGemini3Model('gemini-3-pro-preview')).toBe(true);
@@ -70,6 +102,12 @@ describe('getDisplayString', () => {
     );
   });
 
+  it('should return PREVIEW_GEMINI_3_1_MODEL for PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL', () => {
+    expect(getDisplayString(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL)).toBe(
+      PREVIEW_GEMINI_3_1_MODEL,
+    );
+  });
+
   it('should return the model name as is for other models', () => {
     expect(getDisplayString('custom-model')).toBe('custom-model');
     expect(getDisplayString(DEFAULT_GEMINI_FLASH_LITE_MODEL)).toBe(
@@ -101,6 +139,16 @@ describe('resolveModel', () => {
       expect(model).toBe(PREVIEW_GEMINI_MODEL);
     });
 
+    it('should return Gemini 3.1 Pro when auto-gemini-3 is requested and useGemini3_1 is true', () => {
+      const model = resolveModel(PREVIEW_GEMINI_MODEL_AUTO, true);
+      expect(model).toBe(PREVIEW_GEMINI_3_1_MODEL);
+    });
+
+    it('should return Gemini 3.1 Pro Custom Tools when auto-gemini-3 is requested, useGemini3_1 is true, and useCustomToolModel is true', () => {
+      const model = resolveModel(PREVIEW_GEMINI_MODEL_AUTO, true, true);
+      expect(model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+    });
+
     it('should return the Default Pro model when auto-gemini-2.5 is requested', () => {
       const model = resolveModel(DEFAULT_GEMINI_MODEL_AUTO);
       expect(model).toBe(DEFAULT_GEMINI_MODEL);
@@ -194,6 +242,27 @@ describe('resolveClassifierModel', () => {
       resolveClassifierModel(PREVIEW_GEMINI_MODEL_AUTO, GEMINI_MODEL_ALIAS_PRO),
     ).toBe(PREVIEW_GEMINI_MODEL);
   });
+
+  it('should return Gemini 3.1 Pro when alias is pro and useGemini3_1 is true', () => {
+    expect(
+      resolveClassifierModel(
+        PREVIEW_GEMINI_MODEL_AUTO,
+        GEMINI_MODEL_ALIAS_PRO,
+        true,
+      ),
+    ).toBe(PREVIEW_GEMINI_3_1_MODEL);
+  });
+
+  it('should return Gemini 3.1 Pro Custom Tools when alias is pro, useGemini3_1 is true, and useCustomToolModel is true', () => {
+    expect(
+      resolveClassifierModel(
+        PREVIEW_GEMINI_MODEL_AUTO,
+        GEMINI_MODEL_ALIAS_PRO,
+        true,
+        true,
+      ),
+    ).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+  });
 });
 
 describe('isActiveModel', () => {
@@ -203,9 +272,9 @@ describe('isActiveModel', () => {
     expect(isActiveModel(DEFAULT_GEMINI_FLASH_MODEL)).toBe(true);
   });
 
-  it('should return false for invalid models', () => {
-    expect(isActiveModel('invalid-model')).toBe(false);
-    expect(isActiveModel(GEMINI_MODEL_ALIAS_AUTO)).toBe(false);
+  it('should return true for unknown models and aliases (to support test models)', () => {
+    expect(isActiveModel('invalid-model')).toBe(true);
+    expect(isActiveModel(GEMINI_MODEL_ALIAS_AUTO)).toBe(true);
   });
 
   it('should return false for PREVIEW_GEMINI_MODEL when useGemini3_1 is true', () => {
@@ -215,4 +284,29 @@ describe('isActiveModel', () => {
   it('should return true for other valid models when useGemini3_1 is true', () => {
     expect(isActiveModel(DEFAULT_GEMINI_MODEL, true)).toBe(true);
   });
+
+  it('should correctly filter Gemini 3.1 models based on useCustomToolModel when useGemini3_1 is true', () => {
+    // When custom tools are preferred, standard 3.1 should be inactive
+    expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, true, true)).toBe(false);
+    expect(
+      isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, true, true),
+    ).toBe(true);
+
+    // When custom tools are NOT preferred, custom tools 3.1 should be inactive
+    expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, true, false)).toBe(true);
+    expect(
+      isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, true, false),
+    ).toBe(false);
+  });
+
+  it('should return false for both Gemini 3.1 models when useGemini3_1 is false', () => {
+    expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, false, true)).toBe(false);
+    expect(isActiveModel(PREVIEW_GEMINI_3_1_MODEL, false, false)).toBe(false);
+    expect(
+      isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, false, true),
+    ).toBe(false);
+    expect(
+      isActiveModel(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, false, false),
+    ).toBe(false);
+  });
 });

packages/core/src/config/models.ts

@@ -6,6 +6,8 @@
 
 export const PREVIEW_GEMINI_MODEL = 'gemini-3-pro-preview';
 export const PREVIEW_GEMINI_3_1_MODEL = 'gemini-3.1-pro-preview';
+export const PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL =
+  'gemini-3.1-pro-preview-customtools';
 export const PREVIEW_GEMINI_FLASH_MODEL = 'gemini-3-flash-preview';
 export const DEFAULT_GEMINI_MODEL = 'gemini-2.5-pro';
 export const DEFAULT_GEMINI_FLASH_MODEL = 'gemini-2.5-flash';
@@ -14,6 +16,7 @@ export const DEFAULT_GEMINI_FLASH_LITE_MODEL = 'gemini-2.5-flash-lite';
 export const VALID_GEMINI_MODELS = new Set([
   PREVIEW_GEMINI_MODEL,
   PREVIEW_GEMINI_3_1_MODEL,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
   PREVIEW_GEMINI_FLASH_MODEL,
   DEFAULT_GEMINI_MODEL,
   DEFAULT_GEMINI_FLASH_MODEL,
@@ -45,19 +48,23 @@ export const DEFAULT_THINKING_MODE = 8192;
 export function resolveModel(
   requestedModel: string,
   useGemini3_1: boolean = false,
+  useCustomToolModel: boolean = false,
 ): string {
   switch (requestedModel) {
     case PREVIEW_GEMINI_MODEL:
-    case PREVIEW_GEMINI_MODEL_AUTO: {
-      return useGemini3_1 ? PREVIEW_GEMINI_3_1_MODEL : PREVIEW_GEMINI_MODEL;
+    case PREVIEW_GEMINI_MODEL_AUTO:
+    case GEMINI_MODEL_ALIAS_AUTO:
+    case GEMINI_MODEL_ALIAS_PRO: {
+      if (useGemini3_1) {
+        return useCustomToolModel
+          ? PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
+          : PREVIEW_GEMINI_3_1_MODEL;
+      }
+      return PREVIEW_GEMINI_MODEL;
     }
     case DEFAULT_GEMINI_MODEL_AUTO: {
       return DEFAULT_GEMINI_MODEL;
     }
-    case GEMINI_MODEL_ALIAS_AUTO:
-    case GEMINI_MODEL_ALIAS_PRO: {
-      return useGemini3_1 ? PREVIEW_GEMINI_3_1_MODEL : PREVIEW_GEMINI_MODEL;
-    }
     case GEMINI_MODEL_ALIAS_FLASH: {
       return PREVIEW_GEMINI_FLASH_MODEL;
     }
@@ -80,6 +87,8 @@ export function resolveModel(
 export function resolveClassifierModel(
   requestedModel: string,
   modelAlias: string,
+  useGemini3_1: boolean = false,
+  useCustomToolModel: boolean = false,
 ): string {
   if (modelAlias === GEMINI_MODEL_ALIAS_FLASH) {
     if (
@@ -96,7 +105,7 @@ export function resolveClassifierModel(
     }
     return resolveModel(GEMINI_MODEL_ALIAS_FLASH);
   }
-  return resolveModel(requestedModel);
+  return resolveModel(requestedModel, useGemini3_1, useCustomToolModel);
 }
 export function getDisplayString(model: string) {
   switch (model) {
@@ -108,6 +117,8 @@ export function getDisplayString(model: string) {
       return PREVIEW_GEMINI_MODEL;
     case GEMINI_MODEL_ALIAS_FLASH:
       return PREVIEW_GEMINI_FLASH_MODEL;
+    case PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL:
+      return PREVIEW_GEMINI_3_1_MODEL;
     default:
       return model;
   }
@@ -135,11 +146,7 @@ export function isPreviewModel(model: string): boolean {
  * @returns True if the model is a Pro model.
  */
 export function isProModel(model: string): boolean {
-  return (
-    model === PREVIEW_GEMINI_MODEL ||
-    model === PREVIEW_GEMINI_3_1_MODEL ||
-    model === DEFAULT_GEMINI_MODEL
-  );
+  return model.toLowerCase().includes('pro');
 }
 
 /**
@@ -198,13 +205,24 @@ export function supportsMultimodalFunctionResponse(model: string): boolean {
 export function isActiveModel(
   model: string,
   useGemini3_1: boolean = false,
+  useCustomToolModel: boolean = false,
 ): boolean {
   if (!VALID_GEMINI_MODELS.has(model)) {
-    return false;
+    return true;
   }
   if (useGemini3_1) {
-    return model !== PREVIEW_GEMINI_MODEL;
+    if (model === PREVIEW_GEMINI_MODEL) {
+      return false;
+    }
+    if (useCustomToolModel) {
+      return model !== PREVIEW_GEMINI_3_1_MODEL;
+    } else {
+      return model !== PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL;
+    }
   } else {
-    return model !== PREVIEW_GEMINI_3_1_MODEL;
+    return (
+      model !== PREVIEW_GEMINI_3_1_MODEL &&
+      model !== PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
+    );
   }
 }

packages/core/src/routing/strategies/classifierStrategy.test.ts

@@ -18,11 +18,13 @@ import {
   DEFAULT_GEMINI_MODEL,
   DEFAULT_GEMINI_MODEL_AUTO,
   PREVIEW_GEMINI_MODEL_AUTO,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
 } from '../../config/models.js';
 import { promptIdContext } from '../../utils/promptIdContext.js';
 import type { Content } from '@google/genai';
 import type { ResolvedModelConfig } from '../../services/modelConfigService.js';
 import { debugLogger } from '../../utils/debugLogger.js';
+import { AuthType } from '../../core/contentGenerator.js';
 
 vi.mock('../../core/baseLlmClient.js');
 
@@ -53,6 +55,10 @@ describe('ClassifierStrategy', () => {
       },
       getModel: vi.fn().mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO),
       getNumericalRoutingEnabled: vi.fn().mockResolvedValue(false),
+      getGemini31Launched: vi.fn().mockResolvedValue(false),
+      getContentGeneratorConfig: vi.fn().mockReturnValue({
+        authType: AuthType.LOGIN_WITH_GOOGLE,
+      }),
     } as unknown as Config;
     mockBaseLlmClient = {
       generateJson: vi.fn(),
@@ -339,4 +345,49 @@ describe('ClassifierStrategy', () => {
     // Since requestedModel is Pro, and choice is flash, it should resolve to Flash
     expect(decision?.model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
   });
+
+  describe('Gemini 3.1 and Custom Tools Routing', () => {
+    it('should route to PREVIEW_GEMINI_3_1_MODEL when Gemini 3.1 is launched', async () => {
+      vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
+      vi.mocked(mockConfig.getModel).mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO);
+      const mockApiResponse = {
+        reasoning: 'Complex task',
+        model_choice: 'pro',
+      };
+      vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
+        mockApiResponse,
+      );
+
+      const decision = await strategy.route(
+        mockContext,
+        mockConfig,
+        mockBaseLlmClient,
+      );
+
+      expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+    });
+
+    it('should route to PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL when Gemini 3.1 is launched and auth is USE_GEMINI', async () => {
+      vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
+      vi.mocked(mockConfig.getModel).mockReturnValue(PREVIEW_GEMINI_MODEL_AUTO);
+      vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
+        authType: AuthType.USE_GEMINI,
+      });
+      const mockApiResponse = {
+        reasoning: 'Complex task',
+        model_choice: 'pro',
+      };
+      vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
+        mockApiResponse,
+      );
+
+      const decision = await strategy.route(
+        mockContext,
+        mockConfig,
+        mockBaseLlmClient,
+      );
+
+      expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+    });
+  });
 });

packages/core/src/routing/strategies/classifierStrategy.ts

@@ -20,6 +20,7 @@ import {
   isFunctionResponse,
 } from '../../utils/messageInspectors.js';
 import { debugLogger } from '../../utils/debugLogger.js';
+import { AuthType } from '../../core/contentGenerator.js';
 
 // The number of recent history turns to provide to the router for context.
 const HISTORY_TURNS_FOR_CONTEXT = 4;
@@ -167,9 +168,15 @@ export class ClassifierStrategy implements RoutingStrategy {
 
       const reasoning = routerResponse.reasoning;
       const latencyMs = Date.now() - startTime;
+      const useGemini3_1 = (await config.getGemini31Launched?.()) ?? false;
+      const useCustomToolModel =
+        useGemini3_1 &&
+        config.getContentGeneratorConfig().authType !== AuthType.USE_VERTEX_AI;
       const selectedModel = resolveClassifierModel(
         model,
         routerResponse.model_choice,
+        useGemini3_1,
+        useCustomToolModel,
       );
 
       return {

packages/core/src/routing/strategies/numericalClassifierStrategy.test.ts

@@ -12,6 +12,8 @@ import type { BaseLlmClient } from '../../core/baseLlmClient.js';
 import {
   PREVIEW_GEMINI_FLASH_MODEL,
   PREVIEW_GEMINI_MODEL,
+  PREVIEW_GEMINI_3_1_MODEL,
+  PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL,
   PREVIEW_GEMINI_MODEL_AUTO,
   DEFAULT_GEMINI_MODEL_AUTO,
   DEFAULT_GEMINI_MODEL,
@@ -20,6 +22,7 @@ import { promptIdContext } from '../../utils/promptIdContext.js';
 import type { Content } from '@google/genai';
 import type { ResolvedModelConfig } from '../../services/modelConfigService.js';
 import { debugLogger } from '../../utils/debugLogger.js';
+import { AuthType } from '../../core/contentGenerator.js';
 
 vi.mock('../../core/baseLlmClient.js');
 
@@ -52,6 +55,10 @@ describe('NumericalClassifierStrategy', () => {
       getSessionId: vi.fn().mockReturnValue('control-group-id'), // Default to Control Group (Hash 71 >= 50)
       getNumericalRoutingEnabled: vi.fn().mockResolvedValue(true),
       getClassifierThreshold: vi.fn().mockResolvedValue(undefined),
+      getGemini31Launched: vi.fn().mockResolvedValue(false),
+      getContentGeneratorConfig: vi.fn().mockReturnValue({
+        authType: AuthType.LOGIN_WITH_GOOGLE,
+      }),
     } as unknown as Config;
     mockBaseLlmClient = {
       generateJson: vi.fn(),
@@ -535,4 +542,68 @@ describe('NumericalClassifierStrategy', () => {
       ),
     );
   });
+
+  describe('Gemini 3.1 and Custom Tools Routing', () => {
+    it('should route to PREVIEW_GEMINI_3_1_MODEL when Gemini 3.1 is launched', async () => {
+      vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
+      const mockApiResponse = {
+        complexity_reasoning: 'Complex task',
+        complexity_score: 80,
+      };
+      vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
+        mockApiResponse,
+      );
+
+      const decision = await strategy.route(
+        mockContext,
+        mockConfig,
+        mockBaseLlmClient,
+      );
+
+      expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+    });
+    it('should route to PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL when Gemini 3.1 is launched and auth is USE_GEMINI', async () => {
+      vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
+      vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
+        authType: AuthType.USE_GEMINI,
+      });
+      const mockApiResponse = {
+        complexity_reasoning: 'Complex task',
+        complexity_score: 80,
+      };
+      vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
+        mockApiResponse,
+      );
+
+      const decision = await strategy.route(
+        mockContext,
+        mockConfig,
+        mockBaseLlmClient,
+      );
+
+      expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL);
+    });
+
+    it('should NOT route to custom tools model when auth is USE_VERTEX_AI', async () => {
+      vi.mocked(mockConfig.getGemini31Launched).mockResolvedValue(true);
+      vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
+        authType: AuthType.USE_VERTEX_AI,
+      });
+      const mockApiResponse = {
+        complexity_reasoning: 'Complex task',
+        complexity_score: 80,
+      };
+      vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue(
+        mockApiResponse,
+      );
+
+      const decision = await strategy.route(
+        mockContext,
+        mockConfig,
+        mockBaseLlmClient,
+      );
+
+      expect(decision?.model).toBe(PREVIEW_GEMINI_3_1_MODEL);
+    });
+  });
 });

packages/core/src/routing/strategies/numericalClassifierStrategy.ts

@@ -16,6 +16,7 @@ import { resolveClassifierModel, isGemini3Model } from '../../config/models.js';
 import { createUserContent, Type } from '@google/genai';
 import type { Config } from '../../config/config.js';
 import { debugLogger } from '../../utils/debugLogger.js';
+import { AuthType } from '../../core/contentGenerator.js';
 
 // The number of recent history turns to provide to the router for context.
 const HISTORY_TURNS_FOR_CONTEXT = 8;
@@ -180,8 +181,16 @@ export class NumericalClassifierStrategy implements RoutingStrategy {
           config,
           config.getSessionId() || 'unknown-session',
         );
-
-      const selectedModel = resolveClassifierModel(model, modelAlias);
+      const useGemini3_1 = (await config.getGemini31Launched?.()) ?? false;
+      const useCustomToolModel =
+        useGemini3_1 &&
+        config.getContentGeneratorConfig().authType !== AuthType.USE_VERTEX_AI;
+      const selectedModel = resolveClassifierModel(
+        model,
+        modelAlias,
+        useGemini3_1,
+        useCustomToolModel,
+      );
 
       const latencyMs = Date.now() - startTime;