@@ -89,6 +89,119 @@ describe("GeminiHandler", () => {
8989 )
9090 } )
9191
92+ it ( "should handle reasoning chunks correctly when intermediate reasoning is enabled" , async ( ) => {
93+ // Setup the mock implementation to return an async generator with reasoning chunks
94+ ; ( handler [ "client" ] . models . generateContentStream as any ) . mockResolvedValue ( {
95+ [ Symbol . asyncIterator ] : async function * ( ) {
96+ yield {
97+ candidates : [
98+ {
99+ content : {
100+ parts : [ { thought : true , text : "Let me think about this..." } , { text : "Hello" } ] ,
101+ } ,
102+ } ,
103+ ] ,
104+ }
105+ yield {
106+ candidates : [
107+ {
108+ content : {
109+ parts : [ { thought : true , text : "I need to consider..." } , { text : " world!" } ] ,
110+ } ,
111+ } ,
112+ ] ,
113+ }
114+ yield { usageMetadata : { promptTokenCount : 10 , candidatesTokenCount : 5 , thoughtsTokenCount : 20 } }
115+ } ,
116+ } )
117+
118+ const stream = handler . createMessage ( systemPrompt , mockMessages )
119+ const chunks = [ ]
120+
121+ for await ( const chunk of stream ) {
122+ chunks . push ( chunk )
123+ }
124+
125+ // Should have 6 chunks: 2 reasoning + 2 text + 2 reasoning + 2 text + usage
126+ expect ( chunks . length ) . toBe ( 5 )
127+ expect ( chunks [ 0 ] ) . toEqual ( { type : "reasoning" , text : "Let me think about this..." } )
128+ expect ( chunks [ 1 ] ) . toEqual ( { type : "text" , text : "Hello" } )
129+ expect ( chunks [ 2 ] ) . toEqual ( { type : "reasoning" , text : "I need to consider..." } )
130+ expect ( chunks [ 3 ] ) . toEqual ( { type : "text" , text : " world!" } )
131+ expect ( chunks [ 4 ] ) . toEqual ( {
132+ type : "usage" ,
133+ inputTokens : 10 ,
134+ outputTokens : 5 ,
135+ reasoningTokens : 20 ,
136+ } )
137+ } )
138+
139+ it ( "should suppress reasoning chunks when geminiDisableIntermediateReasoning is enabled" , async ( ) => {
140+ // Create a new handler with the setting enabled
141+ const handlerWithDisabledReasoning = new GeminiHandler ( {
142+ apiKey : "test-key" ,
143+ apiModelId : GEMINI_20_FLASH_THINKING_NAME ,
144+ geminiApiKey : "test-key" ,
145+ geminiDisableIntermediateReasoning : true ,
146+ } )
147+
148+ // Replace the client with our mock
149+ handlerWithDisabledReasoning [ "client" ] = {
150+ models : {
151+ generateContentStream : vitest . fn ( ) ,
152+ generateContent : vitest . fn ( ) ,
153+ getGenerativeModel : vitest . fn ( ) ,
154+ } ,
155+ } as any
156+
157+ // Setup the mock implementation to return an async generator with reasoning chunks
158+ ; ( handlerWithDisabledReasoning [ "client" ] . models . generateContentStream as any ) . mockResolvedValue ( {
159+ [ Symbol . asyncIterator ] : async function * ( ) {
160+ yield {
161+ candidates : [
162+ {
163+ content : {
164+ parts : [ { thought : true , text : "Let me think about this..." } , { text : "Hello" } ] ,
165+ } ,
166+ } ,
167+ ] ,
168+ }
169+ yield {
170+ candidates : [
171+ {
172+ content : {
173+ parts : [ { thought : true , text : "I need to consider..." } , { text : " world!" } ] ,
174+ } ,
175+ } ,
176+ ] ,
177+ }
178+ yield { usageMetadata : { promptTokenCount : 10 , candidatesTokenCount : 5 , thoughtsTokenCount : 20 } }
179+ } ,
180+ } )
181+
182+ const stream = handlerWithDisabledReasoning . createMessage ( systemPrompt , mockMessages )
183+ const chunks = [ ]
184+
185+ for await ( const chunk of stream ) {
186+ chunks . push ( chunk )
187+ }
188+
189+ // Should have only 3 chunks: 2 text + usage (reasoning chunks should be suppressed)
190+ expect ( chunks . length ) . toBe ( 3 )
191+ expect ( chunks [ 0 ] ) . toEqual ( { type : "text" , text : "Hello" } )
192+ expect ( chunks [ 1 ] ) . toEqual ( { type : "text" , text : " world!" } )
193+ expect ( chunks [ 2 ] ) . toEqual ( {
194+ type : "usage" ,
195+ inputTokens : 10 ,
196+ outputTokens : 5 ,
197+ reasoningTokens : 20 ,
198+ } )
199+
200+ // Verify no reasoning chunks were yielded
201+ const reasoningChunks = chunks . filter ( ( chunk ) => chunk . type === "reasoning" )
202+ expect ( reasoningChunks . length ) . toBe ( 0 )
203+ } )
204+
92205 it ( "should handle API errors" , async ( ) => {
93206 const mockError = new Error ( "Gemini API error" )
94207 ; ( handler [ "client" ] . models . generateContentStream as any ) . mockRejectedValue ( mockError )
0 commit comments