@@ -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 ) ;
0 commit comments