From bf0cbb2c7bb9b548a344e26ae99531890fb8d1ef Mon Sep 17 00:00:00 2001
From: Sehoon Shon <sshon@google.com>
Date: Wed, 25 Feb 2026 13:28:10 -0500
Subject: [PATCH] fix(core): use customtools model for gemini-3.1-pro-preview
 for api key users

---
 packages/core/src/config/models.ts         |  5 +++++
 packages/core/src/core/client.ts           |  5 +++++
 packages/core/src/core/contentGenerator.ts |  1 +
 packages/core/src/core/geminiChat.ts       | 16 ++++++++++++++--
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/config/models.ts b/packages/core/src/config/models.ts
index 54ea063569..af04718fbf 100644
--- a/packages/core/src/config/models.ts
+++ b/packages/core/src/config/models.ts
@@ -62,6 +62,11 @@ export function resolveModel(
       }
       return PREVIEW_GEMINI_MODEL;
     }
+    case PREVIEW_GEMINI_3_1_MODEL: {
+      return useCustomToolModel
+        ? PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL
+        : PREVIEW_GEMINI_3_1_MODEL;
+    }
     case DEFAULT_GEMINI_MODEL_AUTO: {
       return DEFAULT_GEMINI_MODEL;
     }
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index c94dd5c04d..12eeba7564 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -64,6 +64,7 @@ import type { RetryAvailabilityContext } from '../utils/retry.js';
 import { partToString } from '../utils/partUtils.js';
 import { coreEvents, CoreEvent } from '../utils/events.js';
 import type { LlmRole } from '../telemetry/types.js';
+import { AuthType } from './contentGenerator.js';
 
 const MAX_TURNS = 100;
 
@@ -539,11 +540,15 @@ export class GeminiClient {
       return this.currentSequenceModel;
     }
 
+    const useCustomToolModel =
+      this.config.getContentGeneratorConfig()?.authType === AuthType.USE_GEMINI;
+
     // Availability logic: The configured model is the source of truth,
     // including any permanent fallbacks (config.setModel) or manual overrides.
     return resolveModel(
       this.config.getActiveModel(),
       this.config.getGemini31LaunchedSync?.() ?? false,
+      useCustomToolModel,
     );
   }
 
diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts
index 7adae874aa..258e228d25 100644
--- a/packages/core/src/core/contentGenerator.ts
+++ b/packages/core/src/core/contentGenerator.ts
@@ -151,6 +151,7 @@ export async function createContentGenerator(
       config.authType === AuthType.USE_GEMINI ||
         config.authType === AuthType.USE_VERTEX_AI ||
         ((await gcConfig.getGemini31Launched?.()) ?? false),
+      config.authType === AuthType.USE_GEMINI,
     );
     const customHeadersEnv =
       process.env['GEMINI_CLI_CUSTOM_HEADERS'] || undefined;
diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts
index b7319c8afd..c4a6b8b6d9 100644
--- a/packages/core/src/core/geminiChat.ts
+++ b/packages/core/src/core/geminiChat.ts
@@ -52,6 +52,7 @@ import {
 } from '../availability/policyHelpers.js';
 import { coreEvents } from '../utils/events.js';
 import type { LlmRole } from '../telemetry/types.js';
+import { AuthType } from './contentGenerator.js';
 
 export enum StreamEventType {
   /** A regular content chunk from the API. */
@@ -494,13 +495,24 @@ export class GeminiChat {
 
     const apiCall = async () => {
       const useGemini3_1 = (await this.config.getGemini31Launched?.()) ?? false;
+      const useCustomToolModel =
+        this.config.getContentGeneratorConfig()?.authType ===
+        AuthType.USE_GEMINI;
       // Default to the last used model (which respects arguments/availability selection)
-      let modelToUse = resolveModel(lastModelToUse, useGemini3_1);
+      let modelToUse = resolveModel(
+        lastModelToUse,
+        useGemini3_1,
+        useCustomToolModel,
+      );
 
       // If the active model has changed (e.g. due to a fallback updating the config),
       // we switch to the new active model.
       if (this.config.getActiveModel() !== initialActiveModel) {
-        modelToUse = resolveModel(this.config.getActiveModel(), useGemini3_1);
+        modelToUse = resolveModel(
+          this.config.getActiveModel(),
+          useGemini3_1,
+          useCustomToolModel,
+        );
       }
 
       if (modelToUse !== lastModelToUse) {