Skip to content

Commit b02e615

Browse files
committed
fix(mistral): drop cached_tokens that exceed prompt_tokens
What: readCachedPromptTokens now takes the reported prompt total and rejects a cached_tokens value larger than it, mirroring the subset guard the OpenAI sibling applies to the identical nested wire field (src/api/providers/openai.ts:304-313, `reportedCacheReadTokens <= inputTokens`). Why: without the bound, a malformed upstream count larger than prompt_tokens would bill reads at 0.1x and surface as cacheReads > tokensIn. The sibling already treats that shape as invalid rather than trusted, so the two readers of `prompt_tokens_details.cached_tokens` now agree. The bound is inclusive (`>` is the rejection test), matching the sibling: a fully cached prompt where cached_tokens equals prompt_tokens stays valid. Impact: malformed usage is dropped instead of mis-billed; valid usage is unchanged. Two spec cases pin the behavior: an over-total value is dropped, and an equal-to-total value is kept.
1 parent a248927 commit b02e615

2 files changed

Lines changed: 70 additions & 4 deletions

File tree

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -550,6 +550,57 @@ describe("MistralHandler", () => {
550550
expect(results[0]).toEqual({ type: "usage", inputTokens: 1013, outputTokens: 30 })
551551
})
552552

553+
it("drops cached_tokens that exceed prompt_tokens instead of billing an impossible read", async () => {
554+
// Subset guard, mirroring the OpenAI sibling that reads the same nested
555+
// field (openai.ts:304-313): a malformed upstream count larger than the
556+
// prompt total must not bill reads at 0.1x or surface cacheReads > tokensIn.
557+
mockCreate.mockImplementationOnce(async (_options) =>
558+
asyncStreamFrom([
559+
{
560+
data: {
561+
choices: [],
562+
usage: {
563+
promptTokens: 100,
564+
completionTokens: 30,
565+
additionalProperties: { prompt_tokens_details: { cached_tokens: 1008 } },
566+
},
567+
},
568+
},
569+
]),
570+
)
571+
572+
const results = await collectStream(handler.createMessage(systemPrompt, messages))
573+
574+
expect(results[0]).toEqual({ type: "usage", inputTokens: 100, outputTokens: 30 })
575+
})
576+
577+
it("keeps a cached_tokens value equal to prompt_tokens (inclusive subset bound)", async () => {
578+
// The sibling guard is `<=`, not `<`: a fully cached prompt is valid.
579+
mockCreate.mockImplementationOnce(async (_options) =>
580+
asyncStreamFrom([
581+
{
582+
data: {
583+
choices: [],
584+
usage: {
585+
promptTokens: 1008,
586+
completionTokens: 30,
587+
additionalProperties: { prompt_tokens_details: { cached_tokens: 1008 } },
588+
},
589+
},
590+
},
591+
]),
592+
)
593+
594+
const results = await collectStream(handler.createMessage(systemPrompt, messages))
595+
596+
expect(results[0]).toEqual({
597+
type: "usage",
598+
inputTokens: 1008,
599+
outputTokens: 30,
600+
cacheReadTokens: 1008,
601+
})
602+
})
603+
553604
it("bills a cached read at 0.1x and the uncached remainder at full input price", async () => {
554605
// End-to-end cost proof for the coupled change: the reported usage
555606
// chunk from the handler feeds calculateApiCostOpenAI, which subtracts

‎src/api/providers/mistral.ts‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,14 +56,28 @@ type MistralTool = {
5656
// SDK parse but only there. Mistral bills these tokens at 10% of the input
5757
// price, and `prompt_tokens` already includes them, so the count is reported as
5858
// `cacheReadTokens` and never added to `inputTokens`.
59-
function readCachedPromptTokens(usage: { additionalProperties?: { [key: string]: unknown } }): number | undefined {
59+
//
60+
// The `cachedTokens <= promptTokens` subset guard mirrors the OpenAI sibling
61+
// that reads the identical nested wire field (src/api/providers/openai.ts:304-313,
62+
// `reportedCacheReadTokens <= inputTokens`): a malformed upstream value larger
63+
// than the prompt total is dropped instead of billing reads at 0.1x and
64+
// surfacing as cacheReads > tokensIn.
65+
function readCachedPromptTokens(
66+
usage: { additionalProperties?: { [key: string]: unknown } },
67+
promptTokens: number,
68+
): number | undefined {
6069
const details = usage.additionalProperties?.["prompt_tokens_details"]
6170
if (typeof details !== "object" || details === null) {
6271
return undefined
6372
}
6473

6574
const cachedTokens = (details as { cached_tokens?: unknown }).cached_tokens
66-
if (typeof cachedTokens !== "number" || !Number.isInteger(cachedTokens) || cachedTokens <= 0) {
75+
if (
76+
typeof cachedTokens !== "number" ||
77+
!Number.isInteger(cachedTokens) ||
78+
cachedTokens <= 0 ||
79+
cachedTokens > promptTokens
80+
) {
6781
return undefined
6882
}
6983

@@ -199,11 +213,12 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand
199213

200214
if (event.data.usage) {
201215
throwIfAborted(requestSignal)
216+
const promptTokens = event.data.usage.promptTokens || 0
202217
yield {
203218
type: "usage",
204-
inputTokens: event.data.usage.promptTokens || 0,
219+
inputTokens: promptTokens,
205220
outputTokens: event.data.usage.completionTokens || 0,
206-
cacheReadTokens: readCachedPromptTokens(event.data.usage),
221+
cacheReadTokens: readCachedPromptTokens(event.data.usage, promptTokens),
207222
}
208223
}
209224
}

0 commit comments

Comments
 (0)