Per-Auth Method Feature Flag for Model Routing (#11333)

Co-authored-by: Abhi <43648792+abhipatel12@users.noreply.github.com>
This commit is contained in:
Victor May
2025-10-21 09:25:51 -04:00
committed by GitHub
parent 81772c4252
commit 0e7b3951c2
3 changed files with 57 additions and 13 deletions
+1
View File
@@ -85,6 +85,7 @@ vi.mock('@google/gemini-cli-core', async () => {
const actualServer = await vi.importActual<typeof ServerConfig>( const actualServer = await vi.importActual<typeof ServerConfig>(
'@google/gemini-cli-core', '@google/gemini-cli-core',
); );
return { return {
...actualServer, ...actualServer,
IdeClient: { IdeClient: {
+36 -11
View File
@@ -574,27 +574,52 @@ describe('Server Config (config.ts)', () => {
}); });
}); });
describe('UseModelRouter Configuration', () => { describe('Model Router with Auth', () => {
it('should default useModelRouter to false when not provided', () => { it('should disable model router by default for oauth-personal', async () => {
const config = new Config(baseParams); const config = new Config({
...baseParams,
useModelRouter: true,
});
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
expect(config.getUseModelRouter()).toBe(false); expect(config.getUseModelRouter()).toBe(false);
}); });
it('should set useModelRouter to true when provided as true', () => { it('should enable model router by default for other auth types', async () => {
const paramsWithModelRouter: ConfigParameters = { const config = new Config({
...baseParams, ...baseParams,
useModelRouter: true, useModelRouter: true,
}; });
const config = new Config(paramsWithModelRouter); await config.refreshAuth(AuthType.USE_GEMINI);
expect(config.getUseModelRouter()).toBe(true); expect(config.getUseModelRouter()).toBe(true);
}); });
it('should set useModelRouter to false when explicitly provided as false', () => { it('should disable model router for specified auth type', async () => {
const paramsWithModelRouter: ConfigParameters = { const config = new Config({
...baseParams,
useModelRouter: true,
disableModelRouterForAuth: [AuthType.USE_GEMINI],
});
await config.refreshAuth(AuthType.USE_GEMINI);
expect(config.getUseModelRouter()).toBe(false);
});
it('should enable model router for other auth type', async () => {
const config = new Config({
...baseParams,
useModelRouter: true,
disableModelRouterForAuth: [],
});
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
expect(config.getUseModelRouter()).toBe(true);
});
it('should keep model router disabled when useModelRouter is false', async () => {
const config = new Config({
...baseParams, ...baseParams,
useModelRouter: false, useModelRouter: false,
}; disableModelRouterForAuth: [AuthType.USE_GEMINI],
const config = new Config(paramsWithModelRouter); });
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
expect(config.getUseModelRouter()).toBe(false); expect(config.getUseModelRouter()).toBe(false);
}); });
}); });
+20 -2
View File
@@ -47,6 +47,7 @@ import {
DEFAULT_GEMINI_EMBEDDING_MODEL, DEFAULT_GEMINI_EMBEDDING_MODEL,
DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_MODEL_AUTO,
DEFAULT_THINKING_MODE, DEFAULT_THINKING_MODE,
} from './models.js'; } from './models.js';
import { shouldAttemptBrowserLaunch } from '../utils/browser.js'; import { shouldAttemptBrowserLaunch } from '../utils/browser.js';
@@ -278,6 +279,7 @@ export interface ConfigParameters {
policyEngineConfig?: PolicyEngineConfig; policyEngineConfig?: PolicyEngineConfig;
output?: OutputSettings; output?: OutputSettings;
useModelRouter?: boolean; useModelRouter?: boolean;
disableModelRouterForAuth?: AuthType[];
enableMessageBusIntegration?: boolean; enableMessageBusIntegration?: boolean;
codebaseInvestigatorSettings?: CodebaseInvestigatorSettings; codebaseInvestigatorSettings?: CodebaseInvestigatorSettings;
continueOnFailedApiCall?: boolean; continueOnFailedApiCall?: boolean;
@@ -374,7 +376,9 @@ export class Config {
private readonly messageBus: MessageBus; private readonly messageBus: MessageBus;
private readonly policyEngine: PolicyEngine; private readonly policyEngine: PolicyEngine;
private readonly outputSettings: OutputSettings; private readonly outputSettings: OutputSettings;
private readonly useModelRouter: boolean; private useModelRouter: boolean;
private readonly initialUseModelRouter: boolean;
private readonly disableModelRouterForAuth?: AuthType[];
private readonly enableMessageBusIntegration: boolean; private readonly enableMessageBusIntegration: boolean;
private readonly codebaseInvestigatorSettings: CodebaseInvestigatorSettings; private readonly codebaseInvestigatorSettings: CodebaseInvestigatorSettings;
private readonly continueOnFailedApiCall: boolean; private readonly continueOnFailedApiCall: boolean;
@@ -470,7 +474,11 @@ export class Config {
this.enableToolOutputTruncation = params.enableToolOutputTruncation ?? true; this.enableToolOutputTruncation = params.enableToolOutputTruncation ?? true;
this.useSmartEdit = params.useSmartEdit ?? true; this.useSmartEdit = params.useSmartEdit ?? true;
this.useWriteTodos = params.useWriteTodos ?? false; this.useWriteTodos = params.useWriteTodos ?? false;
this.useModelRouter = params.useModelRouter ?? false; this.initialUseModelRouter = params.useModelRouter ?? false;
this.useModelRouter = this.initialUseModelRouter;
this.disableModelRouterForAuth = params.disableModelRouterForAuth ?? [
AuthType.LOGIN_WITH_GOOGLE,
];
this.enableMessageBusIntegration = this.enableMessageBusIntegration =
params.enableMessageBusIntegration ?? false; params.enableMessageBusIntegration ?? false;
this.codebaseInvestigatorSettings = { this.codebaseInvestigatorSettings = {
@@ -541,6 +549,16 @@ export class Config {
} }
async refreshAuth(authMethod: AuthType) { async refreshAuth(authMethod: AuthType) {
this.useModelRouter = this.initialUseModelRouter;
if (this.disableModelRouterForAuth?.includes(authMethod)) {
this.useModelRouter = false;
if (this.model === DEFAULT_GEMINI_MODEL_AUTO) {
this.model = DEFAULT_GEMINI_MODEL;
}
} else {
this.model = DEFAULT_GEMINI_MODEL_AUTO;
}
// Vertex and Genai have incompatible encryption and sending history with // Vertex and Genai have incompatible encryption and sending history with
// thoughtSignature from Genai to Vertex will fail, we need to strip them // thoughtSignature from Genai to Vertex will fail, we need to strip them
if ( if (