Skip to content

Commit cbdce3d

Browse files
refactor: roll out stream helpers to kenari minimax vercel mistral specs (#1105)
Co-authored-by: Roomote <roomote@roomote.dev>
1 parent a9cf1d3 commit cbdce3d

4 files changed

Lines changed: 255 additions & 349 deletions

File tree

src/api/providers/__tests__/kenari.spec.ts

Lines changed: 61 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import { kenariDefaultModelId } from "@roo-code/types"
1717
import { KenariHandler } from "../kenari"
1818
import { getModels } from "../fetchers/modelCache"
1919
import { ApiHandlerOptions } from "../../../shared/api"
20+
import { asyncStreamFrom, collectStream } from "../../../test-utils/stream"
2021

2122
vitest.mock("openai")
2223
vitest.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

Comments
 (0)