fix(core): allow Cloud Shell users to use PRO_MODEL_NO_ACCESS experiment (#25702)

This commit is contained in:
Sehoon Shon 2026-04-20 13:29:58 -04:00 committed by GitHub
parent 8573650253
commit 0e5c5b6f49
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 57 additions and 1 deletions

View file

@ -710,6 +710,59 @@ describe('Server Config (config.ts)', () => {
);
});
describe('getProModelNoAccessSync', () => {
it('should return experiment value for AuthType.LOGIN_WITH_GOOGLE', async () => {
vi.mocked(getExperiments).mockResolvedValue({
experimentIds: [],
flags: {
[ExperimentFlags.PRO_MODEL_NO_ACCESS]: {
boolValue: true,
},
},
});
const config = new Config(baseParams);
vi.mocked(createContentGeneratorConfig).mockResolvedValue({
authType: AuthType.LOGIN_WITH_GOOGLE,
});
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
expect(config.getProModelNoAccessSync()).toBe(true);
});
it('should return experiment value for AuthType.COMPUTE_ADC', async () => {
vi.mocked(getExperiments).mockResolvedValue({
experimentIds: [],
flags: {
[ExperimentFlags.PRO_MODEL_NO_ACCESS]: {
boolValue: true,
},
},
});
const config = new Config(baseParams);
vi.mocked(createContentGeneratorConfig).mockResolvedValue({
authType: AuthType.COMPUTE_ADC,
});
await config.refreshAuth(AuthType.COMPUTE_ADC);
expect(config.getProModelNoAccessSync()).toBe(true);
});
it('should return false for other auth types even if experiment is true', async () => {
vi.mocked(getExperiments).mockResolvedValue({
experimentIds: [],
flags: {
[ExperimentFlags.PRO_MODEL_NO_ACCESS]: {
boolValue: true,
},
},
});
const config = new Config(baseParams);
vi.mocked(createContentGeneratorConfig).mockResolvedValue({
authType: AuthType.USE_GEMINI,
});
await config.refreshAuth(AuthType.USE_GEMINI);
expect(config.getProModelNoAccessSync()).toBe(false);
});
});
describe('getRequestTimeoutMs', () => {
it('should return undefined if the flag is not set', () => {
const config = new Config(baseParams);

View file

@ -3160,7 +3160,10 @@ export class Config implements McpContext, AgentLoopContext {
* Note: This method should only be called after startup, once experiments have been loaded.
*/
getProModelNoAccessSync(): boolean {
if (this.contentGeneratorConfig?.authType !== AuthType.LOGIN_WITH_GOOGLE) {
if (
this.contentGeneratorConfig?.authType !== AuthType.LOGIN_WITH_GOOGLE &&
this.contentGeneratorConfig?.authType !== AuthType.COMPUTE_ADC
) {
return false;
}
return (