mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-12 12:54:07 -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>(
|
const actualServer = await vi.importActual<typeof ServerConfig>(
|
||||||
'@google/gemini-cli-core',
|
'@google/gemini-cli-core',
|
||||||
);
|
);
|
||||||
|
|
||||||
return {
|
return {
|
||||||
...actualServer,
|
...actualServer,
|
||||||
IdeClient: {
|
IdeClient: {
|
||||||
|
|||||||
@@ -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);
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -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 (
|
||||||
|
|||||||
Reference in New Issue
Block a user