Skip to content

Commit 2e7722d

Browse files
authored
fix(core): restrict "System: Please continue" invalid stream retry to Gemini 2 models (#20897)
1 parent 69e15a5 commit 2e7722d

2 files changed

Lines changed: 63 additions & 9 deletions

File tree

packages/core/src/core/client.test.ts

Lines changed: 58 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1892,11 +1892,16 @@ ${JSON.stringify(
18921892
);
18931893
});
18941894

1895-
it('should recursively call sendMessageStream with "Please continue." when InvalidStream event is received', async () => {
1895+
it('should recursively call sendMessageStream with "Please continue." when InvalidStream event is received for Gemini 2 models', async () => {
18961896
vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(
18971897
true,
18981898
);
1899-
// Arrange
1899+
// Arrange - router must return a Gemini 2 model for retry to trigger
1900+
mockRouterService.route.mockResolvedValue({
1901+
model: 'gemini-2.0-flash',
1902+
reason: 'test',
1903+
});
1904+
19001905
const mockStream1 = (async function* () {
19011906
yield { type: GeminiEventType.InvalidStream };
19021907
})();
@@ -1926,7 +1931,7 @@ ${JSON.stringify(
19261931

19271932
// Assert
19281933
expect(events).toEqual([
1929-
{ type: GeminiEventType.ModelInfo, value: 'default-routed-model' },
1934+
{ type: GeminiEventType.ModelInfo, value: 'gemini-2.0-flash' },
19301935
{ type: GeminiEventType.InvalidStream },
19311936
{ type: GeminiEventType.Content, value: 'Continued content' },
19321937
]);
@@ -1937,7 +1942,7 @@ ${JSON.stringify(
19371942
// First call with original request
19381943
expect(mockTurnRunFn).toHaveBeenNthCalledWith(
19391944
1,
1940-
{ model: 'default-routed-model', isChatModel: true },
1945+
{ model: 'gemini-2.0-flash', isChatModel: true },
19411946
initialRequest,
19421947
expect.any(AbortSignal),
19431948
undefined,
@@ -1946,7 +1951,7 @@ ${JSON.stringify(
19461951
// Second call with "Please continue."
19471952
expect(mockTurnRunFn).toHaveBeenNthCalledWith(
19481953
2,
1949-
{ model: 'default-routed-model', isChatModel: true },
1954+
{ model: 'gemini-2.0-flash', isChatModel: true },
19501955
[{ text: 'System: Please continue.' }],
19511956
expect.any(AbortSignal),
19521957
undefined,
@@ -1990,11 +1995,57 @@ ${JSON.stringify(
19901995
expect(mockTurnRunFn).toHaveBeenCalledTimes(1);
19911996
});
19921997

1998+
it('should not retry with "Please continue." when InvalidStream event is received for non-Gemini-2 models', async () => {
1999+
vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(
2000+
true,
2001+
);
2002+
// Arrange - router returns a non-Gemini-2 model
2003+
mockRouterService.route.mockResolvedValue({
2004+
model: 'gemini-3.0-pro',
2005+
reason: 'test',
2006+
});
2007+
2008+
const mockStream1 = (async function* () {
2009+
yield { type: GeminiEventType.InvalidStream };
2010+
})();
2011+
2012+
mockTurnRunFn.mockReturnValueOnce(mockStream1);
2013+
2014+
const mockChat: Partial<GeminiChat> = {
2015+
addHistory: vi.fn(),
2016+
setTools: vi.fn(),
2017+
getHistory: vi.fn().mockReturnValue([]),
2018+
getLastPromptTokenCount: vi.fn(),
2019+
};
2020+
client['chat'] = mockChat as GeminiChat;
2021+
2022+
const initialRequest = [{ text: 'Hi' }];
2023+
const promptId = 'prompt-id-invalid-stream-non-g2';
2024+
const signal = new AbortController().signal;
2025+
2026+
// Act
2027+
const stream = client.sendMessageStream(initialRequest, signal, promptId);
2028+
const events = await fromAsync(stream);
2029+
2030+
// Assert
2031+
expect(events).toEqual([
2032+
{ type: GeminiEventType.ModelInfo, value: 'gemini-3.0-pro' },
2033+
{ type: GeminiEventType.InvalidStream },
2034+
]);
2035+
2036+
// Verify that turn.run was called only once (no retry)
2037+
expect(mockTurnRunFn).toHaveBeenCalledTimes(1);
2038+
});
2039+
19932040
it('should stop recursing after one retry when InvalidStream events are repeatedly received', async () => {
19942041
vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(
19952042
true,
19962043
);
1997-
// Arrange
2044+
// Arrange - router must return a Gemini 2 model for retry to trigger
2045+
mockRouterService.route.mockResolvedValue({
2046+
model: 'gemini-2.0-flash',
2047+
reason: 'test',
2048+
});
19982049
// Always return a new invalid stream
19992050
mockTurnRunFn.mockImplementation(() =>
20002051
(async function* () {
@@ -2025,7 +2076,7 @@ ${JSON.stringify(
20252076
events
20262077
.filter((e) => e.type === GeminiEventType.ModelInfo)
20272078
.map((e) => e.value),
2028-
).toEqual(['default-routed-model']);
2079+
).toEqual(['gemini-2.0-flash']);
20292080

20302081
// Verify that turn.run was called twice
20312082
expect(mockTurnRunFn).toHaveBeenCalledTimes(2);

packages/core/src/core/client.ts

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ import {
6060
applyModelSelection,
6161
createAvailabilityContextProvider,
6262
} from '../availability/policyHelpers.js';
63-
import { resolveModel } from '../config/models.js';
63+
import { resolveModel, isGemini2Model } from '../config/models.js';
6464
import type { RetryAvailabilityContext } from '../utils/retry.js';
6565
import { partToString } from '../utils/partUtils.js';
6666
import { coreEvents, CoreEvent } from '../utils/events.js';
@@ -725,7 +725,10 @@ export class GeminiClient {
725725
}
726726

727727
if (isInvalidStream) {
728-
if (this.config.getContinueOnFailedApiCall()) {
728+
if (
729+
this.config.getContinueOnFailedApiCall() &&
730+
isGemini2Model(modelToUse)
731+
) {
729732
if (isInvalidStreamRetry) {
730733
logContentRetryFailure(
731734
this.config,

0 commit comments

Comments
 (0)