fix(core): restrict "System: Please continue" invalid stream retry to Gemini 2 models (#20897)

This commit is contained in:
Sandy Tao
2026-03-02 15:21:13 -08:00
committed by GitHub
parent 69e15a50d1
commit 2e7722d6a3
2 changed files with 63 additions and 9 deletions

View File

@@ -1892,11 +1892,16 @@ ${JSON.stringify(
);
});
it('should recursively call sendMessageStream with "Please continue." when InvalidStream event is received', async () => {
it('should recursively call sendMessageStream with "Please continue." when InvalidStream event is received for Gemini 2 models', async () => {
vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(
true,
);
// Arrange
// Arrange - router must return a Gemini 2 model for retry to trigger
mockRouterService.route.mockResolvedValue({
model: 'gemini-2.0-flash',
reason: 'test',
});
const mockStream1 = (async function* () {
yield { type: GeminiEventType.InvalidStream };
})();
@@ -1926,7 +1931,7 @@ ${JSON.stringify(
// Assert
expect(events).toEqual([
{ type: GeminiEventType.ModelInfo, value: 'default-routed-model' },
{ type: GeminiEventType.ModelInfo, value: 'gemini-2.0-flash' },
{ type: GeminiEventType.InvalidStream },
{ type: GeminiEventType.Content, value: 'Continued content' },
]);
@@ -1937,7 +1942,7 @@ ${JSON.stringify(
// First call with original request
expect(mockTurnRunFn).toHaveBeenNthCalledWith(
1,
{ model: 'default-routed-model', isChatModel: true },
{ model: 'gemini-2.0-flash', isChatModel: true },
initialRequest,
expect.any(AbortSignal),
undefined,
@@ -1946,7 +1951,7 @@ ${JSON.stringify(
// Second call with "Please continue."
expect(mockTurnRunFn).toHaveBeenNthCalledWith(
2,
{ model: 'default-routed-model', isChatModel: true },
{ model: 'gemini-2.0-flash', isChatModel: true },
[{ text: 'System: Please continue.' }],
expect.any(AbortSignal),
undefined,
@@ -1990,11 +1995,57 @@ ${JSON.stringify(
expect(mockTurnRunFn).toHaveBeenCalledTimes(1);
});
it('should not retry with "Please continue." when InvalidStream event is received for non-Gemini-2 models', async () => {
vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(
true,
);
// Arrange - router returns a non-Gemini-2 model
mockRouterService.route.mockResolvedValue({
model: 'gemini-3.0-pro',
reason: 'test',
});
const mockStream1 = (async function* () {
yield { type: GeminiEventType.InvalidStream };
})();
mockTurnRunFn.mockReturnValueOnce(mockStream1);
const mockChat: Partial<GeminiChat> = {
addHistory: vi.fn(),
setTools: vi.fn(),
getHistory: vi.fn().mockReturnValue([]),
getLastPromptTokenCount: vi.fn(),
};
client['chat'] = mockChat as GeminiChat;
const initialRequest = [{ text: 'Hi' }];
const promptId = 'prompt-id-invalid-stream-non-g2';
const signal = new AbortController().signal;
// Act
const stream = client.sendMessageStream(initialRequest, signal, promptId);
const events = await fromAsync(stream);
// Assert
expect(events).toEqual([
{ type: GeminiEventType.ModelInfo, value: 'gemini-3.0-pro' },
{ type: GeminiEventType.InvalidStream },
]);
// Verify that turn.run was called only once (no retry)
expect(mockTurnRunFn).toHaveBeenCalledTimes(1);
});
it('should stop recursing after one retry when InvalidStream events are repeatedly received', async () => {
vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(
true,
);
// Arrange
// Arrange - router must return a Gemini 2 model for retry to trigger
mockRouterService.route.mockResolvedValue({
model: 'gemini-2.0-flash',
reason: 'test',
});
// Always return a new invalid stream
mockTurnRunFn.mockImplementation(() =>
(async function* () {
@@ -2025,7 +2076,7 @@ ${JSON.stringify(
events
.filter((e) => e.type === GeminiEventType.ModelInfo)
.map((e) => e.value),
).toEqual(['default-routed-model']);
).toEqual(['gemini-2.0-flash']);
// Verify that turn.run was called twice
expect(mockTurnRunFn).toHaveBeenCalledTimes(2);

View File

@@ -60,7 +60,7 @@ import {
applyModelSelection,
createAvailabilityContextProvider,
} from '../availability/policyHelpers.js';
import { resolveModel } from '../config/models.js';
import { resolveModel, isGemini2Model } from '../config/models.js';
import type { RetryAvailabilityContext } from '../utils/retry.js';
import { partToString } from '../utils/partUtils.js';
import { coreEvents, CoreEvent } from '../utils/events.js';
@@ -725,7 +725,10 @@ export class GeminiClient {
}
if (isInvalidStream) {
if (this.config.getContinueOnFailedApiCall()) {
if (
this.config.getContinueOnFailedApiCall() &&
isGemini2Model(modelToUse)
) {
if (isInvalidStreamRetry) {
logContentRetryFailure(
this.config,