@@ -17,6 +17,7 @@ import { kenariDefaultModelId } from "@roo-code/types"
1717import { KenariHandler } from "../kenari"
1818import { getModels } from "../fetchers/modelCache"
1919import { ApiHandlerOptions } from "../../../shared/api"
20+ import { asyncStreamFrom , collectStream } from "../../../test-utils/stream"
2021
2122vitest . mock ( "openai" )
2223vitest . mock ( "delay" , ( ) => ( { default : vitest . fn ( ( ) => Promise . resolve ( ) ) } ) )
@@ -84,9 +85,9 @@ describe("KenariHandler", () => {
8485
8586 describe ( "createMessage" , ( ) => {
8687 beforeEach ( ( ) => {
87- mockCreate . mockImplementation ( async ( ) => ( {
88- [ Symbol . asyncIterator ] : async function * ( ) {
89- yield {
88+ mockCreate . mockImplementation ( async ( ) =>
89+ asyncStreamFrom ( [
90+ {
9091 choices : [
9192 {
9293 delta : {
@@ -104,28 +105,25 @@ describe("KenariHandler", () => {
104105 } ,
105106 ] ,
106107 usage : null ,
107- }
108- yield {
108+ } ,
109+ {
109110 choices : [ { delta : { } , index : 0 } ] ,
110111 usage : {
111112 prompt_tokens : 12 ,
112113 completion_tokens : 7 ,
113114 total_tokens : 19 ,
114115 prompt_tokens_details : { cached_tokens : 4 } ,
115116 } ,
116- }
117- } ,
118- } ) )
117+ } ,
118+ ] ) ,
119+ )
119120 } )
120121
121122 it ( "streams text, reasoning, tool-call and usage chunks" , async ( ) => {
122123 const handler = new KenariHandler ( mockOptions )
123124 const messages : Anthropic . Messages . MessageParam [ ] = [ { role : "user" , content : "Hi" } ]
124125
125- const chunks = [ ]
126- for await ( const chunk of handler . createMessage ( "You are helpful." , messages ) ) {
127- chunks . push ( chunk )
128- }
126+ const chunks = await collectStream ( handler . createMessage ( "You are helpful." , messages ) )
129127
130128 expect ( chunks ) . toContainEqual ( { type : "text" , text : "Hello" } )
131129 expect ( chunks ) . toContainEqual ( { type : "reasoning" , text : "thinking…" } )
@@ -145,37 +143,31 @@ describe("KenariHandler", () => {
145143 } )
146144
147145 it ( "yields nothing for a chunk whose delta has no content, reasoning or tool calls" , async ( ) => {
148- mockCreate . mockImplementation ( async ( ) => ( {
149- [ Symbol . asyncIterator ] : async function * ( ) {
150- yield { choices : [ { delta : { } , index : 0 } ] , usage : null }
151- yield { choices : [ ] , usage : null }
152- } ,
153- } ) )
146+ mockCreate . mockImplementation ( async ( ) =>
147+ asyncStreamFrom ( [
148+ { choices : [ { delta : { } , index : 0 } ] , usage : null } ,
149+ { choices : [ ] , usage : null } ,
150+ ] ) ,
151+ )
154152
155153 const handler = new KenariHandler ( mockOptions )
156- const chunks = [ ]
157- for await ( const chunk of handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) ) {
158- chunks . push ( chunk )
159- }
154+ const chunks = await collectStream ( handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) )
160155
161156 expect ( chunks ) . toEqual ( [ ] )
162157 } )
163158
164159 it ( "streams tool call chunks even when the function name and arguments are missing" , async ( ) => {
165- mockCreate . mockImplementation ( async ( ) => ( {
166- [ Symbol . asyncIterator ] : async function * ( ) {
167- yield {
160+ mockCreate . mockImplementation ( async ( ) =>
161+ asyncStreamFrom ( [
162+ {
168163 choices : [ { delta : { tool_calls : [ { index : 1 } ] } , index : 0 } ] ,
169164 usage : null ,
170- }
171- } ,
172- } ) )
165+ } ,
166+ ] ) ,
167+ )
173168
174169 const handler = new KenariHandler ( mockOptions )
175- const chunks = [ ]
176- for await ( const chunk of handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) ) {
177- chunks . push ( chunk )
178- }
170+ const chunks = await collectStream ( handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) )
179171
180172 expect ( chunks ) . toEqual ( [
181173 {
@@ -189,20 +181,17 @@ describe("KenariHandler", () => {
189181 } )
190182
191183 it ( "reports undefined cache reads when usage has no prompt_tokens_details" , async ( ) => {
192- mockCreate . mockImplementation ( async ( ) => ( {
193- [ Symbol . asyncIterator ] : async function * ( ) {
194- yield {
184+ mockCreate . mockImplementation ( async ( ) =>
185+ asyncStreamFrom ( [
186+ {
195187 choices : [ { delta : { } , index : 0 } ] ,
196188 usage : { prompt_tokens : 3 , completion_tokens : 2 , total_tokens : 5 } ,
197- }
198- } ,
199- } ) )
189+ } ,
190+ ] ) ,
191+ )
200192
201193 const handler = new KenariHandler ( mockOptions )
202- const chunks = [ ]
203- for await ( const chunk of handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) ) {
204- chunks . push ( chunk )
205- }
194+ const chunks = await collectStream ( handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) )
206195
207196 expect ( chunks ) . toEqual ( [
208197 {
@@ -215,39 +204,33 @@ describe("KenariHandler", () => {
215204 } )
216205
217206 it ( "skips the reasoning chunk when reasoning_content is an empty string" , async ( ) => {
218- mockCreate . mockImplementation ( async ( ) => ( {
219- [ Symbol . asyncIterator ] : async function * ( ) {
220- yield {
207+ mockCreate . mockImplementation ( async ( ) =>
208+ asyncStreamFrom ( [
209+ {
221210 choices : [ { delta : { content : "Hi" , reasoning_content : "" } , index : 0 } ] ,
222211 usage : null ,
223- }
224- } ,
225- } ) )
212+ } ,
213+ ] ) ,
214+ )
226215
227216 const handler = new KenariHandler ( mockOptions )
228- const chunks = [ ]
229- for await ( const chunk of handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) ) {
230- chunks . push ( chunk )
231- }
217+ const chunks = await collectStream ( handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) )
232218
233219 expect ( chunks ) . toEqual ( [ { type : "text" , text : "Hi" } ] )
234220 } )
235221
236222 it ( "emits reasoning from the OpenRouter-style `reasoning` field when reasoning_content is absent" , async ( ) => {
237- mockCreate . mockImplementation ( async ( ) => ( {
238- [ Symbol . asyncIterator ] : async function * ( ) {
239- yield {
223+ mockCreate . mockImplementation ( async ( ) =>
224+ asyncStreamFrom ( [
225+ {
240226 choices : [ { delta : { content : "Hi" , reasoning : "thinking…" } , index : 0 } ] ,
241227 usage : null ,
242- }
243- } ,
244- } ) )
228+ } ,
229+ ] ) ,
230+ )
245231
246232 const handler = new KenariHandler ( mockOptions )
247- const chunks = [ ]
248- for await ( const chunk of handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) ) {
249- chunks . push ( chunk )
250- }
233+ const chunks = await collectStream ( handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) )
251234
252235 expect ( chunks ) . toContainEqual ( { type : "reasoning" , text : "thinking…" } )
253236 } )
@@ -264,9 +247,7 @@ describe("KenariHandler", () => {
264247 } )
265248
266249 const handler = new KenariHandler ( { kenariApiKey : "test-key" , kenariModelId : "openai/o3-mini" } )
267- for await ( const _chunk of handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) ) {
268- void _chunk // drain
269- }
250+ await collectStream ( handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) )
270251
271252 expect ( mockCreate ) . toHaveBeenCalledWith (
272253 expect . objectContaining ( {
@@ -278,40 +259,35 @@ describe("KenariHandler", () => {
278259
279260 it ( "sends an explicitly configured model temperature" , async ( ) => {
280261 const handler = new KenariHandler ( { ...mockOptions , modelTemperature : 0.7 } )
281- for await ( const _chunk of handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) ) {
282- void _chunk // drain
283- }
262+ await collectStream ( handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) )
284263
285264 expect ( mockCreate ) . toHaveBeenCalledWith ( expect . objectContaining ( { temperature : 0.7 } ) )
286265 } )
287266
288267 it ( "honors metadata.parallelToolCalls false" , async ( ) => {
289268 const handler = new KenariHandler ( mockOptions )
290- for await ( const _chunk of handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] , {
291- taskId : "task-1" ,
292- parallelToolCalls : false ,
293- } ) ) {
294- void _chunk // drain
295- }
269+ await collectStream (
270+ handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] , {
271+ taskId : "task-1" ,
272+ parallelToolCalls : false ,
273+ } ) ,
274+ )
296275
297276 expect ( mockCreate ) . toHaveBeenCalledWith ( expect . objectContaining ( { parallel_tool_calls : false } ) )
298277 } )
299278
300279 it ( "reports zero usage when the upstream counts are zero" , async ( ) => {
301- mockCreate . mockImplementation ( async ( ) => ( {
302- [ Symbol . asyncIterator ] : async function * ( ) {
303- yield {
280+ mockCreate . mockImplementation ( async ( ) =>
281+ asyncStreamFrom ( [
282+ {
304283 choices : [ { delta : { content : "x" } , index : 0 } ] ,
305284 usage : { prompt_tokens : 0 , completion_tokens : 0 } ,
306- }
307- } ,
308- } ) )
285+ } ,
286+ ] ) ,
287+ )
309288
310289 const handler = new KenariHandler ( mockOptions )
311- const chunks = [ ]
312- for await ( const chunk of handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) ) {
313- chunks . push ( chunk )
314- }
290+ const chunks = await collectStream ( handler . createMessage ( "sys" , [ { role : "user" , content : "Hi" } ] ) )
315291
316292 expect ( chunks ) . toContainEqual ( {
317293 type : "usage" ,
@@ -324,9 +300,7 @@ describe("KenariHandler", () => {
324300 it ( "requests a streaming completion with usage included" , async ( ) => {
325301 const handler = new KenariHandler ( mockOptions )
326302 const messages : Anthropic . Messages . MessageParam [ ] = [ { role : "user" , content : "Hi" } ]
327- for await ( const _chunk of handler . createMessage ( "sys" , messages ) ) {
328- void _chunk // drain
329- }
303+ await collectStream ( handler . createMessage ( "sys" , messages ) )
330304
331305 expect ( mockCreate ) . toHaveBeenCalledWith (
332306 expect . objectContaining ( {
0 commit comments