mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-16 17:11:04 -07:00
Per-Auth Method Feature Flag for Model Routing (#11333)
Co-authored-by: Abhi <43648792+abhipatel12@users.noreply.github.com>
This commit is contained in:
@@ -85,6 +85,7 @@ vi.mock('@google/gemini-cli-core', async () => {
|
||||
const actualServer = await vi.importActual<typeof ServerConfig>(
|
||||
'@google/gemini-cli-core',
|
||||
);
|
||||
|
||||
return {
|
||||
...actualServer,
|
||||
IdeClient: {
|
||||
|
||||
@@ -574,27 +574,52 @@ describe('Server Config (config.ts)', () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe('UseModelRouter Configuration', () => {
|
||||
it('should default useModelRouter to false when not provided', () => {
|
||||
const config = new Config(baseParams);
|
||||
describe('Model Router with Auth', () => {
|
||||
it('should disable model router by default for oauth-personal', async () => {
|
||||
const config = new Config({
|
||||
...baseParams,
|
||||
useModelRouter: true,
|
||||
});
|
||||
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
|
||||
expect(config.getUseModelRouter()).toBe(false);
|
||||
});
|
||||
|
||||
it('should set useModelRouter to true when provided as true', () => {
|
||||
const paramsWithModelRouter: ConfigParameters = {
|
||||
it('should enable model router by default for other auth types', async () => {
|
||||
const config = new Config({
|
||||
...baseParams,
|
||||
useModelRouter: true,
|
||||
};
|
||||
const config = new Config(paramsWithModelRouter);
|
||||
});
|
||||
await config.refreshAuth(AuthType.USE_GEMINI);
|
||||
expect(config.getUseModelRouter()).toBe(true);
|
||||
});
|
||||
|
||||
it('should set useModelRouter to false when explicitly provided as false', () => {
|
||||
const paramsWithModelRouter: ConfigParameters = {
|
||||
it('should disable model router for specified auth type', async () => {
|
||||
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,
|
||||
useModelRouter: false,
|
||||
};
|
||||
const config = new Config(paramsWithModelRouter);
|
||||
disableModelRouterForAuth: [AuthType.USE_GEMINI],
|
||||
});
|
||||
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
|
||||
expect(config.getUseModelRouter()).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -47,6 +47,7 @@ import {
|
||||
DEFAULT_GEMINI_EMBEDDING_MODEL,
|
||||
DEFAULT_GEMINI_FLASH_MODEL,
|
||||
DEFAULT_GEMINI_MODEL,
|
||||
DEFAULT_GEMINI_MODEL_AUTO,
|
||||
DEFAULT_THINKING_MODE,
|
||||
} from './models.js';
|
||||
import { shouldAttemptBrowserLaunch } from '../utils/browser.js';
|
||||
@@ -278,6 +279,7 @@ export interface ConfigParameters {
|
||||
policyEngineConfig?: PolicyEngineConfig;
|
||||
output?: OutputSettings;
|
||||
useModelRouter?: boolean;
|
||||
disableModelRouterForAuth?: AuthType[];
|
||||
enableMessageBusIntegration?: boolean;
|
||||
codebaseInvestigatorSettings?: CodebaseInvestigatorSettings;
|
||||
continueOnFailedApiCall?: boolean;
|
||||
@@ -374,7 +376,9 @@ export class Config {
|
||||
private readonly messageBus: MessageBus;
|
||||
private readonly policyEngine: PolicyEngine;
|
||||
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 codebaseInvestigatorSettings: CodebaseInvestigatorSettings;
|
||||
private readonly continueOnFailedApiCall: boolean;
|
||||
@@ -470,7 +474,11 @@ export class Config {
|
||||
this.enableToolOutputTruncation = params.enableToolOutputTruncation ?? true;
|
||||
this.useSmartEdit = params.useSmartEdit ?? true;
|
||||
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 =
|
||||
params.enableMessageBusIntegration ?? false;
|
||||
this.codebaseInvestigatorSettings = {
|
||||
@@ -541,6 +549,16 @@ export class Config {
|
||||
}
|
||||
|
||||
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
|
||||
// thoughtSignature from Genai to Vertex will fail, we need to strip them
|
||||
if (
|
||||
|
||||
Reference in New Issue
Block a user