fix(core): ensure stable fallback for restricted preview models (#26999)

This commit is contained in:
Gal Zahavi
2026-05-13 14:46:41 -07:00
committed by GitHub
parent 1814c7f358
commit 77078b3e8a
9 changed files with 50 additions and 28 deletions
+4 -4
View File
@@ -3251,8 +3251,8 @@ describe('Config Quota & Preview Model Access', () => {
vi.mocked(getCodeAssistServer).mockReturnValue(undefined);
const result = await config.refreshUserQuota();
expect(result).toBeUndefined();
// Never set => stays null (unknown); getter returns true so UI shows preview
expect(config.getHasAccessToPreviewModel()).toBe(true);
// Never set => stays null (unknown); getter returns false by default
expect(config.getHasAccessToPreviewModel()).toBe(false);
});
it('should return undefined if retrieveUserQuota fails', async () => {
@@ -3261,8 +3261,8 @@ describe('Config Quota & Preview Model Access', () => {
);
const result = await config.refreshUserQuota();
expect(result).toBeUndefined();
// Never set => stays null (unknown); getter returns true so UI shows preview
expect(config.getHasAccessToPreviewModel()).toBe(true);
// Never set => stays null (unknown); getter returns false by default
expect(config.getHasAccessToPreviewModel()).toBe(false);
});
it('should derive quota from remainingFraction when remainingAmount is missing', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
+7 -6
View File
@@ -1622,10 +1622,7 @@ export class Config implements McpContext, AgentLoopContext {
this.baseLlmClient = new BaseLlmClient(this.contentGenerator, this);
const authType = this.contentGeneratorConfig.authType;
if (
authType === AuthType.USE_GEMINI ||
authType === AuthType.USE_VERTEX_AI
) {
if (authType === AuthType.USE_GEMINI) {
this.setHasAccessToPreviewModel(true);
}
@@ -2227,7 +2224,7 @@ export class Config implements McpContext, AgentLoopContext {
}
getHasAccessToPreviewModel(): boolean {
return this.hasAccessToPreviewModel !== false;
return this.hasAccessToPreviewModel ?? false;
}
setHasAccessToPreviewModel(hasAccess: boolean | null): void {
@@ -2289,7 +2286,6 @@ export class Config implements McpContext, AgentLoopContext {
});
}
}
this.emitQuotaChangedEvent();
}
const hasAccess =
@@ -2297,6 +2293,11 @@ export class Config implements McpContext, AgentLoopContext {
(b) => b.modelId && isPreviewModel(b.modelId, this),
) ?? false;
this.setHasAccessToPreviewModel(hasAccess);
if (quota.buckets) {
this.emitQuotaChangedEvent();
}
return quota;
} catch (e) {
debugLogger.debug('Failed to retrieve user quota', e);