Fix an issue where the agent stops prematurely (#16269)

This commit is contained in:
Christian Gunderman
2026-01-09 23:05:27 +00:00
committed by GitHub
parent 356f76e545
commit c87d1aed4c
6 changed files with 118 additions and 25 deletions

View File

@@ -70,6 +70,7 @@ describe('useQuotaAndFallback', () => {
setFallbackHandlerSpy = vi.spyOn(mockConfig, 'setFallbackModelHandler');
vi.spyOn(mockConfig, 'setQuotaErrorOccurred');
vi.spyOn(mockConfig, 'setModel');
vi.spyOn(mockConfig, 'setActiveModel');
});
afterEach(() => {
@@ -164,8 +165,8 @@ describe('useQuotaAndFallback', () => {
const intent = await promise!;
expect(intent).toBe('retry_always');
// Verify setModel was called with isFallbackModel=true
expect(mockConfig.setModel).toHaveBeenCalledWith('gemini-flash', true);
// Verify setActiveModel was called
expect(mockConfig.setActiveModel).toHaveBeenCalledWith('gemini-flash');
// The pending request should be cleared from the state
expect(result.current.proQuotaRequest).toBeNull();
@@ -278,8 +279,8 @@ describe('useQuotaAndFallback', () => {
const intent = await promise!;
expect(intent).toBe('retry_always');
// Verify setModel was called with isFallbackModel=true
expect(mockConfig.setModel).toHaveBeenCalledWith('model-B', true);
// Verify setActiveModel was called
expect(mockConfig.setActiveModel).toHaveBeenCalledWith('model-B');
// The pending request should be cleared from the state
expect(result.current.proQuotaRequest).toBeNull();
@@ -336,10 +337,9 @@ To disable gemini-3-pro-preview, disable "Preview features" in /settings.`,
const intent = await promise!;
expect(intent).toBe('retry_always');
// Verify setModel was called with isFallbackModel=true
expect(mockConfig.setModel).toHaveBeenCalledWith(
// Verify setActiveModel was called
expect(mockConfig.setActiveModel).toHaveBeenCalledWith(
'gemini-2.5-pro',
true,
);
expect(result.current.proQuotaRequest).toBeNull();
@@ -425,8 +425,12 @@ To disable gemini-3-pro-preview, disable "Preview features" in /settings.`,
expect(intent).toBe('retry_always');
expect(result.current.proQuotaRequest).toBeNull();
// Verify setModel was called with isFallbackModel=true
expect(mockConfig.setModel).toHaveBeenCalledWith('gemini-flash', true);
// Verify setActiveModel was called
expect(mockConfig.setActiveModel).toHaveBeenCalledWith('gemini-flash');
// Verify quota error flags are reset
expect(mockSetModelSwitchedFromQuotaError).toHaveBeenCalledWith(false);
expect(mockConfig.setQuotaErrorOccurred).toHaveBeenCalledWith(false);
// Check for the "Switched to fallback model" message
expect(mockHistoryManager.addItem).toHaveBeenCalledTimes(1);

View File

@@ -129,21 +129,27 @@ export function useQuotaAndFallback({
setProQuotaRequest(null);
isDialogPending.current = false; // Reset the flag here
if (choice === 'retry_always') {
// Set the model to the fallback model for the current session.
// This ensures the Footer updates and future turns use this model.
// The change is not persisted, so the original model is restored on restart.
config.activateFallbackMode(proQuotaRequest.fallbackModel);
historyManager.addItem(
{
type: MessageType.INFO,
text: `Switched to fallback model ${proQuotaRequest.fallbackModel}`,
},
Date.now(),
);
if (choice === 'retry_always' || choice === 'retry_once') {
// Reset quota error flags to allow the agent loop to continue.
setModelSwitchedFromQuotaError(false);
config.setQuotaErrorOccurred(false);
if (choice === 'retry_always') {
// Set the model to the fallback model for the current session.
// This ensures the Footer updates and future turns use this model.
// The change is not persisted, so the original model is restored on restart.
config.activateFallbackMode(proQuotaRequest.fallbackModel);
historyManager.addItem(
{
type: MessageType.INFO,
text: `Switched to fallback model ${proQuotaRequest.fallbackModel}`,
},
Date.now(),
);
}
}
},
[proQuotaRequest, historyManager, config],
[proQuotaRequest, historyManager, config, setModelSwitchedFromQuotaError],
);
return {