Skip to content

Commit 3965cd9

Browse files
hannesrudolphhannesrudolph
authored
fix: stabilize token/cache accounting across providers and routed Roo metadata (#11448)
1 parent b7857bc commit 3965cd9

35 files changed

Lines changed: 1272 additions & 162 deletions

‎src/api/providers/__tests__/bedrock.spec.ts‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1279,4 +1279,165 @@ describe("AwsBedrockHandler", () => {
12791279
expect(mockCaptureException).toHaveBeenCalled()
12801280
})
12811281
})
1282+
1283+
describe("AI SDK v6 usage field paths", () => {
1284+
const systemPrompt = "You are a helpful assistant"
1285+
const messages: RooMessage[] = [
1286+
{
1287+
role: "user",
1288+
content: "Hello",
1289+
},
1290+
]
1291+
1292+
function setupStream(usage: Record<string, unknown>, providerMetadata: Record<string, unknown> = {}) {
1293+
async function* mockFullStream() {
1294+
yield { type: "text-delta", text: "reply" }
1295+
}
1296+
1297+
mockStreamText.mockReturnValue({
1298+
fullStream: mockFullStream(),
1299+
usage: Promise.resolve(usage),
1300+
providerMetadata: Promise.resolve(providerMetadata),
1301+
})
1302+
}
1303+
1304+
describe("cache tokens", () => {
1305+
it("should read cache tokens from v6 top-level cachedInputTokens", async () => {
1306+
setupStream({ inputTokens: 100, outputTokens: 50, cachedInputTokens: 30 })
1307+
1308+
const generator = handler.createMessage(systemPrompt, messages)
1309+
const chunks: unknown[] = []
1310+
for await (const chunk of generator) {
1311+
chunks.push(chunk)
1312+
}
1313+
1314+
const usageChunk = chunks.find((c: any) => c.type === "usage") as any
1315+
expect(usageChunk).toBeDefined()
1316+
expect(usageChunk.cacheReadTokens).toBe(30)
1317+
})
1318+
1319+
it("should read cache tokens from v6 inputTokenDetails.cacheReadTokens", async () => {
1320+
setupStream({
1321+
inputTokens: 100,
1322+
outputTokens: 50,
1323+
inputTokenDetails: { cacheReadTokens: 25 },
1324+
})
1325+
1326+
const generator = handler.createMessage(systemPrompt, messages)
1327+
const chunks: unknown[] = []
1328+
for await (const chunk of generator) {
1329+
chunks.push(chunk)
1330+
}
1331+
1332+
const usageChunk = chunks.find((c: any) => c.type === "usage") as any
1333+
expect(usageChunk).toBeDefined()
1334+
expect(usageChunk.cacheReadTokens).toBe(25)
1335+
})
1336+
1337+
it("should prefer v6 top-level cachedInputTokens over providerMetadata.bedrock", async () => {
1338+
setupStream(
1339+
{ inputTokens: 100, outputTokens: 50, cachedInputTokens: 30 },
1340+
{ bedrock: { usage: { cacheReadInputTokens: 20 } } },
1341+
)
1342+
1343+
const generator = handler.createMessage(systemPrompt, messages)
1344+
const chunks: unknown[] = []
1345+
for await (const chunk of generator) {
1346+
chunks.push(chunk)
1347+
}
1348+
1349+
const usageChunk = chunks.find((c: any) => c.type === "usage") as any
1350+
expect(usageChunk).toBeDefined()
1351+
expect(usageChunk.cacheReadTokens).toBe(30)
1352+
})
1353+
1354+
it("should fall back to providerMetadata.bedrock.usage.cacheReadInputTokens", async () => {
1355+
setupStream(
1356+
{ inputTokens: 100, outputTokens: 50 },
1357+
{ bedrock: { usage: { cacheReadInputTokens: 20 } } },
1358+
)
1359+
1360+
const generator = handler.createMessage(systemPrompt, messages)
1361+
const chunks: unknown[] = []
1362+
for await (const chunk of generator) {
1363+
chunks.push(chunk)
1364+
}
1365+
1366+
const usageChunk = chunks.find((c: any) => c.type === "usage") as any
1367+
expect(usageChunk).toBeDefined()
1368+
expect(usageChunk.cacheReadTokens).toBe(20)
1369+
})
1370+
1371+
it("should read cacheWriteTokens from v6 inputTokenDetails.cacheWriteTokens", async () => {
1372+
setupStream({
1373+
inputTokens: 100,
1374+
outputTokens: 50,
1375+
inputTokenDetails: { cacheWriteTokens: 15 },
1376+
})
1377+
1378+
const generator = handler.createMessage(systemPrompt, messages)
1379+
const chunks: unknown[] = []
1380+
for await (const chunk of generator) {
1381+
chunks.push(chunk)
1382+
}
1383+
1384+
const usageChunk = chunks.find((c: any) => c.type === "usage") as any
1385+
expect(usageChunk).toBeDefined()
1386+
expect(usageChunk.cacheWriteTokens).toBe(15)
1387+
})
1388+
})
1389+
1390+
describe("reasoning tokens", () => {
1391+
it("should read reasoning tokens from v6 top-level reasoningTokens", async () => {
1392+
setupStream({ inputTokens: 100, outputTokens: 50, reasoningTokens: 40 })
1393+
1394+
const generator = handler.createMessage(systemPrompt, messages)
1395+
const chunks: unknown[] = []
1396+
for await (const chunk of generator) {
1397+
chunks.push(chunk)
1398+
}
1399+
1400+
const usageChunk = chunks.find((c: any) => c.type === "usage") as any
1401+
expect(usageChunk).toBeDefined()
1402+
expect(usageChunk.reasoningTokens).toBe(40)
1403+
})
1404+
1405+
it("should read reasoning tokens from v6 outputTokenDetails.reasoningTokens", async () => {
1406+
setupStream({
1407+
inputTokens: 100,
1408+
outputTokens: 50,
1409+
outputTokenDetails: { reasoningTokens: 35 },
1410+
})
1411+
1412+
const generator = handler.createMessage(systemPrompt, messages)
1413+
const chunks: unknown[] = []
1414+
for await (const chunk of generator) {
1415+
chunks.push(chunk)
1416+
}
1417+
1418+
const usageChunk = chunks.find((c: any) => c.type === "usage") as any
1419+
expect(usageChunk).toBeDefined()
1420+
expect(usageChunk.reasoningTokens).toBe(35)
1421+
})
1422+
1423+
it("should prefer v6 top-level reasoningTokens over outputTokenDetails", async () => {
1424+
setupStream({
1425+
inputTokens: 100,
1426+
outputTokens: 50,
1427+
reasoningTokens: 40,
1428+
outputTokenDetails: { reasoningTokens: 15 },
1429+
})
1430+
1431+
const generator = handler.createMessage(systemPrompt, messages)
1432+
const chunks: unknown[] = []
1433+
for await (const chunk of generator) {
1434+
chunks.push(chunk)
1435+
}
1436+
1437+
const usageChunk = chunks.find((c: any) => c.type === "usage") as any
1438+
expect(usageChunk).toBeDefined()
1439+
expect(usageChunk.reasoningTokens).toBe(40)
1440+
})
1441+
})
1442+
})
12821443
})

‎src/api/providers/__tests__/gemini.spec.ts‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -472,4 +472,168 @@ describe("GeminiHandler", () => {
472472
expect(mockCaptureException).toHaveBeenCalled()
473473
})
474474
})
475+
476+
describe("AI SDK v6 usage field paths", () => {
477+
const mockMessages: RooMessage[] = [
478+
{
479+
role: "user",
480+
content: "Hello",
481+
},
482+
]
483+
const systemPrompt = "You are a helpful assistant"
484+
485+
function setupStream(usage: Record<string, unknown>) {
486+
const mockFullStream = (async function* () {
487+
yield { type: "text-delta", text: "reply" }
488+
})()
489+
490+
mockStreamText.mockReturnValue({
491+
fullStream: mockFullStream,
492+
usage: Promise.resolve(usage),
493+
providerMetadata: Promise.resolve({}),
494+
})
495+
}
496+
497+
describe("cache tokens", () => {
498+
it("should read cache tokens from v6 top-level cachedInputTokens", async () => {
499+
setupStream({ inputTokens: 100, outputTokens: 50, cachedInputTokens: 30 })
500+
501+
const stream = handler.createMessage(systemPrompt, mockMessages)
502+
const chunks = []
503+
for await (const chunk of stream) {
504+
chunks.push(chunk)
505+
}
506+
507+
const usageChunk = chunks.find((c) => c.type === "usage")
508+
expect(usageChunk).toBeDefined()
509+
expect(usageChunk!.cacheReadTokens).toBe(30)
510+
})
511+
512+
it("should read cache tokens from v6 inputTokenDetails.cacheReadTokens", async () => {
513+
setupStream({
514+
inputTokens: 100,
515+
outputTokens: 50,
516+
inputTokenDetails: { cacheReadTokens: 25 },
517+
})
518+
519+
const stream = handler.createMessage(systemPrompt, mockMessages)
520+
const chunks = []
521+
for await (const chunk of stream) {
522+
chunks.push(chunk)
523+
}
524+
525+
const usageChunk = chunks.find((c) => c.type === "usage")
526+
expect(usageChunk).toBeDefined()
527+
expect(usageChunk!.cacheReadTokens).toBe(25)
528+
})
529+
530+
it("should prefer v6 top-level cachedInputTokens over legacy details", async () => {
531+
setupStream({
532+
inputTokens: 100,
533+
outputTokens: 50,
534+
cachedInputTokens: 30,
535+
details: { cachedInputTokens: 20 },
536+
})
537+
538+
const stream = handler.createMessage(systemPrompt, mockMessages)
539+
const chunks = []
540+
for await (const chunk of stream) {
541+
chunks.push(chunk)
542+
}
543+
544+
const usageChunk = chunks.find((c) => c.type === "usage")
545+
expect(usageChunk).toBeDefined()
546+
expect(usageChunk!.cacheReadTokens).toBe(30)
547+
})
548+
549+
it("should fall back to legacy details.cachedInputTokens", async () => {
550+
setupStream({
551+
inputTokens: 100,
552+
outputTokens: 50,
553+
details: { cachedInputTokens: 20 },
554+
})
555+
556+
const stream = handler.createMessage(systemPrompt, mockMessages)
557+
const chunks = []
558+
for await (const chunk of stream) {
559+
chunks.push(chunk)
560+
}
561+
562+
const usageChunk = chunks.find((c) => c.type === "usage")
563+
expect(usageChunk).toBeDefined()
564+
expect(usageChunk!.cacheReadTokens).toBe(20)
565+
})
566+
})
567+
568+
describe("reasoning tokens", () => {
569+
it("should read reasoning tokens from v6 top-level reasoningTokens", async () => {
570+
setupStream({ inputTokens: 100, outputTokens: 50, reasoningTokens: 40 })
571+
572+
const stream = handler.createMessage(systemPrompt, mockMessages)
573+
const chunks = []
574+
for await (const chunk of stream) {
575+
chunks.push(chunk)
576+
}
577+
578+
const usageChunk = chunks.find((c) => c.type === "usage")
579+
expect(usageChunk).toBeDefined()
580+
expect(usageChunk!.reasoningTokens).toBe(40)
581+
})
582+
583+
it("should read reasoning tokens from v6 outputTokenDetails.reasoningTokens", async () => {
584+
setupStream({
585+
inputTokens: 100,
586+
outputTokens: 50,
587+
outputTokenDetails: { reasoningTokens: 35 },
588+
})
589+
590+
const stream = handler.createMessage(systemPrompt, mockMessages)
591+
const chunks = []
592+
for await (const chunk of stream) {
593+
chunks.push(chunk)
594+
}
595+
596+
const usageChunk = chunks.find((c) => c.type === "usage")
597+
expect(usageChunk).toBeDefined()
598+
expect(usageChunk!.reasoningTokens).toBe(35)
599+
})
600+
601+
it("should prefer v6 top-level reasoningTokens over legacy details", async () => {
602+
setupStream({
603+
inputTokens: 100,
604+
outputTokens: 50,
605+
reasoningTokens: 40,
606+
details: { reasoningTokens: 15 },
607+
})
608+
609+
const stream = handler.createMessage(systemPrompt, mockMessages)
610+
const chunks = []
611+
for await (const chunk of stream) {
612+
chunks.push(chunk)
613+
}
614+
615+
const usageChunk = chunks.find((c) => c.type === "usage")
616+
expect(usageChunk).toBeDefined()
617+
expect(usageChunk!.reasoningTokens).toBe(40)
618+
})
619+
620+
it("should fall back to legacy details.reasoningTokens", async () => {
621+
setupStream({
622+
inputTokens: 100,
623+
outputTokens: 50,
624+
details: { reasoningTokens: 15 },
625+
})
626+
627+
const stream = handler.createMessage(systemPrompt, mockMessages)
628+
const chunks = []
629+
for await (const chunk of stream) {
630+
chunks.push(chunk)
631+
}
632+
633+
const usageChunk = chunks.find((c) => c.type === "usage")
634+
expect(usageChunk).toBeDefined()
635+
expect(usageChunk!.reasoningTokens).toBe(15)
636+
})
637+
})
638+
})
475639
})

‎src/api/providers/__tests__/native-ollama.spec.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,13 @@ describe("NativeOllamaHandler", () => {
8484
expect(results).toHaveLength(3)
8585
expect(results[0]).toEqual({ type: "text", text: "Hello" })
8686
expect(results[1]).toEqual({ type: "text", text: " world" })
87-
expect(results[2]).toEqual({ type: "usage", inputTokens: 10, outputTokens: 2 })
87+
expect(results[2]).toEqual({
88+
type: "usage",
89+
inputTokens: 10,
90+
outputTokens: 2,
91+
totalInputTokens: 10,
92+
totalOutputTokens: 2,
93+
})
8894
})
8995

9096
it("should not include providerOptions by default (no num_ctx)", async () => {

0 commit comments

Comments
 (0)