From 651d10d90b4bd80c85c3a09e18a719e4d5dc4c13 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Thu, 18 Jun 2026 23:47:25 +0200 Subject: [PATCH] feat(ai): enable Mistral prompt caching --- package-lock.json | 36 ++++++++++-- packages/ai/CHANGELOG.md | 1 + packages/ai/package.json | 5 +- packages/ai/scripts/generate-models.ts | 2 +- packages/ai/src/models.generated.ts | 58 +++++++++---------- packages/ai/src/providers/mistral.ts | 39 +++++++++++-- .../ai/test/mistral-reasoning-mode.test.ts | 18 ++++++ packages/coding-agent/npm-shrinkwrap.json | 38 ++++++++++-- 8 files changed, 152 insertions(+), 45 deletions(-) diff --git a/package-lock.json b/package-lock.json index 821c9d30..d1abbf97 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1624,14 +1624,23 @@ } }, "node_modules/@mistralai/mistralai": { - "version": "2.2.1", - "resolved": "https://registry.npmjs.org/@mistralai/mistralai/-/mistralai-2.2.1.tgz", - "integrity": "sha512-uKU8CZmL2RzYKmplsU01hii4p3pe4HqJefpWNRWXm1Tcm0Sm4xXfwSLIy4k7ZCPlbETCGcp69E7hZs+WOJ5itQ==", + "version": "2.2.6", + "resolved": "https://registry.npmjs.org/@mistralai/mistralai/-/mistralai-2.2.6.tgz", + "integrity": "sha512-W8pX7zHxjJvMIpw8JMxeJEleapXX0Q9NPszdNzqkM3MIEoIGPObdodujj+WHteXEvGfaP/AMwlNyRfEzSY6dQQ==", "license": "Apache-2.0", "dependencies": { + "@opentelemetry/semantic-conventions": "^1.40.0", "ws": "^8.18.0", "zod": "^3.25.0 || ^4.0.0", "zod-to-json-schema": "^3.25.0" + }, + "peerDependencies": { + "@opentelemetry/api": "^1.9.0" + }, + "peerDependenciesMeta": { + "@opentelemetry/api": { + "optional": true + } } }, "node_modules/@nodable/entities": { @@ -1684,6 +1693,24 @@ "node": ">= 8" } }, + "node_modules/@opentelemetry/api": { + "version": "1.9.0", + "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz", + "integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==", + "license": "Apache-2.0", + "engines": { + "node": ">=8.0.0" + } + }, + "node_modules/@opentelemetry/semantic-conventions": { + "version": "1.41.1", + "resolved": "https://registry.npmjs.org/@opentelemetry/semantic-conventions/-/semantic-conventions-1.41.1.tgz", + "integrity": "sha512-/UhIkaZgPutTFmQ7RnIJGgDXZmtEJ7Dvi86xNTFWcnRxVRNk/aotsqDJYeEvDP+FSMB2SdW+pQzNMcWP0rwuNA==", + "license": "Apache-2.0", + "engines": { + "node": ">=14" + } + }, "node_modules/@pkgjs/parseargs": { "version": "0.11.0", "resolved": "https://registry.npmjs.org/@pkgjs/parseargs/-/parseargs-0.11.0.tgz", @@ -6219,7 +6246,8 @@ "@anthropic-ai/sdk": "0.91.1", "@aws-sdk/client-bedrock-runtime": "3.1048.0", "@google/genai": "1.52.0", - "@mistralai/mistralai": "2.2.1", + "@mistralai/mistralai": "2.2.6", + "@opentelemetry/api": "1.9.0", "@smithy/node-http-handler": "4.7.3", "http-proxy-agent": "7.0.2", "https-proxy-agent": "7.0.6", diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 2b5aeca1..c4c63d8b 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -4,6 +4,7 @@ ### Added +- Added prompt caching for Mistral requests using the pi session ID as `prompt_cache_key`, including cached-token usage and cost accounting ([#5854](https://github.com/earendil-works/pi/issues/5854)). - Added the OpenRouter Fusion alias as `openrouter/fusion` ([#5866](https://github.com/earendil-works/pi/pull/5866) by [@dannote](https://github.com/dannote)). ## [0.79.7] - 2026-06-18 diff --git a/packages/ai/package.json b/packages/ai/package.json index 9c71d156..2c0bd6cf 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -69,9 +69,10 @@ "dependencies": { "@anthropic-ai/sdk": "0.91.1", "@aws-sdk/client-bedrock-runtime": "3.1048.0", - "@smithy/node-http-handler": "4.7.3", "@google/genai": "1.52.0", - "@mistralai/mistralai": "2.2.1", + "@mistralai/mistralai": "2.2.6", + "@opentelemetry/api": "1.9.0", + "@smithy/node-http-handler": "4.7.3", "http-proxy-agent": "7.0.2", "https-proxy-agent": "7.0.6", "openai": "6.26.0", diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 0c1bcc42..c1fddc92 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -948,7 +948,7 @@ async function loadModelsDevData(): Promise[]> { cost: { input: m.cost?.input || 0, output: m.cost?.output || 0, - cacheRead: m.cost?.cache_read || 0, + cacheRead: m.cost?.cache_read ?? (m.cost?.input ? roundCost(m.cost.input * 0.1) : 0), cacheWrite: m.cost?.cache_write || 0, }, contextWindow: m.limit?.context || 4096, diff --git a/packages/ai/src/models.generated.ts b/packages/ai/src/models.generated.ts index 8e1b776f..82d0e137 100644 --- a/packages/ai/src/models.generated.ts +++ b/packages/ai/src/models.generated.ts @@ -5761,7 +5761,7 @@ export const MODELS = { cost: { input: 0.3, output: 0.9, - cacheRead: 0, + cacheRead: 0.03, cacheWrite: 0, }, contextWindow: 256000, @@ -5778,7 +5778,7 @@ export const MODELS = { cost: { input: 0.4, output: 2, - cacheRead: 0, + cacheRead: 0.04, cacheWrite: 0, }, contextWindow: 262144, @@ -5795,7 +5795,7 @@ export const MODELS = { cost: { input: 0.4, output: 2, - cacheRead: 0, + cacheRead: 0.04, cacheWrite: 0, }, contextWindow: 262144, @@ -5812,7 +5812,7 @@ export const MODELS = { cost: { input: 0.4, output: 2, - cacheRead: 0, + cacheRead: 0.04, cacheWrite: 0, }, contextWindow: 128000, @@ -5829,7 +5829,7 @@ export const MODELS = { cost: { input: 0.4, output: 2, - cacheRead: 0, + cacheRead: 0.04, cacheWrite: 0, }, contextWindow: 262144, @@ -5846,7 +5846,7 @@ export const MODELS = { cost: { input: 0.1, output: 0.3, - cacheRead: 0, + cacheRead: 0.01, cacheWrite: 0, }, contextWindow: 128000, @@ -5863,7 +5863,7 @@ export const MODELS = { cost: { input: 0.1, output: 0.3, - cacheRead: 0, + cacheRead: 0.01, cacheWrite: 0, }, contextWindow: 128000, @@ -5897,7 +5897,7 @@ export const MODELS = { cost: { input: 2, output: 5, - cacheRead: 0, + cacheRead: 0.2, cacheWrite: 0, }, contextWindow: 128000, @@ -5914,7 +5914,7 @@ export const MODELS = { cost: { input: 0.5, output: 1.5, - cacheRead: 0, + cacheRead: 0.05, cacheWrite: 0, }, contextWindow: 128000, @@ -5931,7 +5931,7 @@ export const MODELS = { cost: { input: 0.04, output: 0.04, - cacheRead: 0, + cacheRead: 0.004, cacheWrite: 0, }, contextWindow: 128000, @@ -5948,7 +5948,7 @@ export const MODELS = { cost: { input: 0.1, output: 0.1, - cacheRead: 0, + cacheRead: 0.01, cacheWrite: 0, }, contextWindow: 128000, @@ -5965,7 +5965,7 @@ export const MODELS = { cost: { input: 2, output: 6, - cacheRead: 0, + cacheRead: 0.2, cacheWrite: 0, }, contextWindow: 131072, @@ -5982,7 +5982,7 @@ export const MODELS = { cost: { input: 0.5, output: 1.5, - cacheRead: 0, + cacheRead: 0.05, cacheWrite: 0, }, contextWindow: 262144, @@ -5999,7 +5999,7 @@ export const MODELS = { cost: { input: 0.5, output: 1.5, - cacheRead: 0, + cacheRead: 0.05, cacheWrite: 0, }, contextWindow: 262144, @@ -6016,7 +6016,7 @@ export const MODELS = { cost: { input: 0.4, output: 2, - cacheRead: 0, + cacheRead: 0.04, cacheWrite: 0, }, contextWindow: 131072, @@ -6033,7 +6033,7 @@ export const MODELS = { cost: { input: 0.4, output: 2, - cacheRead: 0, + cacheRead: 0.04, cacheWrite: 0, }, contextWindow: 262144, @@ -6050,7 +6050,7 @@ export const MODELS = { cost: { input: 1.5, output: 7.5, - cacheRead: 0, + cacheRead: 0.15, cacheWrite: 0, }, contextWindow: 262144, @@ -6067,7 +6067,7 @@ export const MODELS = { cost: { input: 1.5, output: 7.5, - cacheRead: 0, + cacheRead: 0.15, cacheWrite: 0, }, contextWindow: 262144, @@ -6084,7 +6084,7 @@ export const MODELS = { cost: { input: 0.4, output: 2, - cacheRead: 0, + cacheRead: 0.04, cacheWrite: 0, }, contextWindow: 262144, @@ -6101,7 +6101,7 @@ export const MODELS = { cost: { input: 0.15, output: 0.15, - cacheRead: 0, + cacheRead: 0.015, cacheWrite: 0, }, contextWindow: 128000, @@ -6118,7 +6118,7 @@ export const MODELS = { cost: { input: 0.1, output: 0.3, - cacheRead: 0, + cacheRead: 0.01, cacheWrite: 0, }, contextWindow: 128000, @@ -6135,7 +6135,7 @@ export const MODELS = { cost: { input: 0.15, output: 0.6, - cacheRead: 0, + cacheRead: 0.015, cacheWrite: 0, }, contextWindow: 256000, @@ -6152,7 +6152,7 @@ export const MODELS = { cost: { input: 0.15, output: 0.6, - cacheRead: 0, + cacheRead: 0.015, cacheWrite: 0, }, contextWindow: 256000, @@ -6169,7 +6169,7 @@ export const MODELS = { cost: { input: 0.25, output: 0.25, - cacheRead: 0, + cacheRead: 0.025, cacheWrite: 0, }, contextWindow: 8000, @@ -6186,7 +6186,7 @@ export const MODELS = { cost: { input: 0.15, output: 0.15, - cacheRead: 0, + cacheRead: 0.015, cacheWrite: 0, }, contextWindow: 128000, @@ -6203,7 +6203,7 @@ export const MODELS = { cost: { input: 2, output: 6, - cacheRead: 0, + cacheRead: 0.2, cacheWrite: 0, }, contextWindow: 64000, @@ -6220,7 +6220,7 @@ export const MODELS = { cost: { input: 0.7, output: 0.7, - cacheRead: 0, + cacheRead: 0.07, cacheWrite: 0, }, contextWindow: 32000, @@ -6237,7 +6237,7 @@ export const MODELS = { cost: { input: 0.15, output: 0.15, - cacheRead: 0, + cacheRead: 0.015, cacheWrite: 0, }, contextWindow: 128000, @@ -6254,7 +6254,7 @@ export const MODELS = { cost: { input: 2, output: 6, - cacheRead: 0, + cacheRead: 0.2, cacheWrite: 0, }, contextWindow: 128000, diff --git a/packages/ai/src/providers/mistral.ts b/packages/ai/src/providers/mistral.ts index 1bc7d4ce..6a132236 100644 --- a/packages/ai/src/providers/mistral.ts +++ b/packages/ai/src/providers/mistral.ts @@ -226,7 +226,7 @@ function buildRequestOptions(model: Model<"mistral-conversations">, options?: Mi // Mistral infrastructure uses `x-affinity` for KV-cache reuse (prefix caching). // Respect explicit caller-provided header values. - if (options?.sessionId && !headers["x-affinity"]) { + if (shouldUsePromptCaching(options) && !headers["x-affinity"]) { headers["x-affinity"] = options.sessionId; } @@ -255,6 +255,7 @@ function buildChatPayload( if (options?.toolChoice) payload.toolChoice = mapToolChoice(options.toolChoice); if (options?.promptMode) payload.promptMode = options.promptMode; if (options?.reasoningEffort) payload.reasoningEffort = options.reasoningEffort; + if (shouldUsePromptCaching(options)) payload.promptCacheKey = options.sessionId; if (context.systemPrompt) { payload.messages.unshift({ @@ -266,6 +267,31 @@ function buildChatPayload( return payload; } +function shouldUsePromptCaching(options?: MistralOptions): options is MistralOptions & { sessionId: string } { + return options?.cacheRetention !== "none" && !!options?.sessionId; +} + +function getMistralCachedPromptTokens(usage: unknown, promptTokens: number): number { + const rawUsage = usage as { + promptTokensDetails?: { cachedTokens?: unknown } | null; + prompt_tokens_details?: { cached_tokens?: unknown } | null; + promptTokenDetails?: { cachedTokens?: unknown } | null; + prompt_token_details?: { cached_tokens?: unknown } | null; + numCachedTokens?: unknown; + num_cached_tokens?: unknown; + }; + const rawCachedTokens = + rawUsage.promptTokensDetails?.cachedTokens ?? + rawUsage.prompt_tokens_details?.cached_tokens ?? + rawUsage.promptTokenDetails?.cachedTokens ?? + rawUsage.prompt_token_details?.cached_tokens ?? + rawUsage.numCachedTokens ?? + rawUsage.num_cached_tokens ?? + 0; + const cachedTokens = typeof rawCachedTokens === "number" && Number.isFinite(rawCachedTokens) ? rawCachedTokens : 0; + return Math.min(promptTokens, Math.max(0, cachedTokens)); +} + async function consumeChatStream( model: Model<"mistral-conversations">, output: AssistantMessage, @@ -305,11 +331,16 @@ async function consumeChatStream( output.responseId ||= chunk.id; if (chunk.usage) { - output.usage.input = chunk.usage.promptTokens || 0; + const promptTokens = chunk.usage.promptTokens || 0; + const cachedPromptTokens = getMistralCachedPromptTokens(chunk.usage, promptTokens); + + output.usage.input = Math.max(0, promptTokens - cachedPromptTokens); output.usage.output = chunk.usage.completionTokens || 0; - output.usage.cacheRead = 0; + output.usage.cacheRead = cachedPromptTokens; output.usage.cacheWrite = 0; - output.usage.totalTokens = chunk.usage.totalTokens || output.usage.input + output.usage.output; + output.usage.totalTokens = + chunk.usage.totalTokens || + output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite; calculateCost(model, output.usage); } diff --git a/packages/ai/test/mistral-reasoning-mode.test.ts b/packages/ai/test/mistral-reasoning-mode.test.ts index 35a5b96b..d197a56a 100644 --- a/packages/ai/test/mistral-reasoning-mode.test.ts +++ b/packages/ai/test/mistral-reasoning-mode.test.ts @@ -6,6 +6,7 @@ import type { Context, Model, SimpleStreamOptions } from "../src/types.ts"; interface MistralPayload { promptMode?: "reasoning"; reasoningEffort?: "none" | "high"; + promptCacheKey?: string; } function makeContext(): Context { @@ -77,4 +78,21 @@ describe("Mistral reasoning mode selection", () => { expect(payload.reasoningEffort).toBeUndefined(); expect(payload.promptMode).toBeUndefined(); }); + + it("uses the session id as prompt cache key", async () => { + const payload = await capturePayload(getModel("mistral", "mistral-large-latest"), { + sessionId: "session-123", + }); + + expect(payload.promptCacheKey).toBe("session-123"); + }); + + it("omits prompt cache key when cache retention is disabled", async () => { + const payload = await capturePayload(getModel("mistral", "mistral-large-latest"), { + sessionId: "session-123", + cacheRetention: "none", + }); + + expect(payload.promptCacheKey).toBeUndefined(); + }); }); diff --git a/packages/coding-agent/npm-shrinkwrap.json b/packages/coding-agent/npm-shrinkwrap.json index c366f1a5..ad672bf2 100644 --- a/packages/coding-agent/npm-shrinkwrap.json +++ b/packages/coding-agent/npm-shrinkwrap.json @@ -494,9 +494,10 @@ "dependencies": { "@anthropic-ai/sdk": "0.91.1", "@aws-sdk/client-bedrock-runtime": "3.1048.0", - "@smithy/node-http-handler": "4.7.3", "@google/genai": "1.52.0", - "@mistralai/mistralai": "2.2.1", + "@mistralai/mistralai": "2.2.6", + "@opentelemetry/api": "1.9.0", + "@smithy/node-http-handler": "4.7.3", "http-proxy-agent": "7.0.2", "https-proxy-agent": "7.0.6", "openai": "6.26.0", @@ -741,14 +742,23 @@ "optional": true }, "node_modules/@mistralai/mistralai": { - "version": "2.2.1", - "resolved": "https://registry.npmjs.org/@mistralai/mistralai/-/mistralai-2.2.1.tgz", - "integrity": "sha512-uKU8CZmL2RzYKmplsU01hii4p3pe4HqJefpWNRWXm1Tcm0Sm4xXfwSLIy4k7ZCPlbETCGcp69E7hZs+WOJ5itQ==", + "version": "2.2.6", + "resolved": "https://registry.npmjs.org/@mistralai/mistralai/-/mistralai-2.2.6.tgz", + "integrity": "sha512-W8pX7zHxjJvMIpw8JMxeJEleapXX0Q9NPszdNzqkM3MIEoIGPObdodujj+WHteXEvGfaP/AMwlNyRfEzSY6dQQ==", "license": "Apache-2.0", "dependencies": { + "@opentelemetry/semantic-conventions": "^1.40.0", "ws": "^8.18.0", "zod": "^3.25.0 || ^4.0.0", "zod-to-json-schema": "^3.25.0" + }, + "peerDependencies": { + "@opentelemetry/api": "^1.9.0" + }, + "peerDependenciesMeta": { + "@opentelemetry/api": { + "optional": true + } } }, "node_modules/@nodable/entities": { @@ -763,6 +773,24 @@ } ] }, + "node_modules/@opentelemetry/api": { + "version": "1.9.0", + "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz", + "integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==", + "license": "Apache-2.0", + "engines": { + "node": ">=8.0.0" + } + }, + "node_modules/@opentelemetry/semantic-conventions": { + "version": "1.41.1", + "resolved": "https://registry.npmjs.org/@opentelemetry/semantic-conventions/-/semantic-conventions-1.41.1.tgz", + "integrity": "sha512-/UhIkaZgPutTFmQ7RnIJGgDXZmtEJ7Dvi86xNTFWcnRxVRNk/aotsqDJYeEvDP+FSMB2SdW+pQzNMcWP0rwuNA==", + "license": "Apache-2.0", + "engines": { + "node": ">=14" + } + }, "node_modules/@protobufjs/aspromise": { "version": "1.1.2", "resolved": "https://registry.npmjs.org/@protobufjs/aspromise/-/aspromise-1.1.2.tgz",