From 4535415300f1599cd83807e1d1dd7f2615aabf70 Mon Sep 17 00:00:00 2001
From: Mario Zechner <badlogicgames@gmail.com>
Date: Tue, 10 Mar 2026 16:28:09 +0100
Subject: [PATCH] fix(ai): capture usage from choice.usage for non-standard
 OpenAI-compatible providers

Some OpenAI-compatible providers (e.g., Moonshot/Kimi) return usage
data in chunk.choices[0].usage instead of the standard chunk.usage.
Extract usage parsing into a helper and check choice.usage as fallback.

closes #2017
---
 packages/ai/CHANGELOG.md                      |  1 +
 .../ai/src/providers/openai-completions.ts    | 57 ++++++++++++-------
 2 files changed, 36 insertions(+), 22 deletions(-)

diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md
index da13d16f..c0ce622d 100644
--- a/packages/ai/CHANGELOG.md
+++ b/packages/ai/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### Fixed
 
 - Fixed GitHub Copilot device-code login polling to respect OAuth slow-down intervals, wait before the first token poll, and include a clearer clock-drift hint in WSL/VM environments when repeated slow-downs lead to timeout.
+- Fixed usage statistics not being captured for OpenAI-compatible providers that return usage in `choice.usage` instead of the standard `chunk.usage` (e.g., Moonshot/Kimi) ([#2017](https://github.com/badlogic/pi-mono/issues/2017))
 
 ## [0.57.1] - 2026-03-07
 
diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts
index 437dea52..beca1132 100644
--- a/packages/ai/src/providers/openai-completions.ts
+++ b/packages/ai/src/providers/openai-completions.ts
@@ -128,33 +128,18 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
 
 			for await (const chunk of openaiStream) {
 				if (chunk.usage) {
-					const cachedTokens = chunk.usage.prompt_tokens_details?.cached_tokens || 0;
-					const reasoningTokens = chunk.usage.completion_tokens_details?.reasoning_tokens || 0;
-					const input = (chunk.usage.prompt_tokens || 0) - cachedTokens;
-					const outputTokens = (chunk.usage.completion_tokens || 0) + reasoningTokens;
-					output.usage = {
-						// OpenAI includes cached tokens in prompt_tokens, so subtract to get non-cached input
-						input,
-						output: outputTokens,
-						cacheRead: cachedTokens,
-						cacheWrite: 0,
-						// Compute totalTokens ourselves since we add reasoning_tokens to output
-						// and some providers (e.g., Groq) don't include them in total_tokens
-						totalTokens: input + outputTokens + cachedTokens,
-						cost: {
-							input: 0,
-							output: 0,
-							cacheRead: 0,
-							cacheWrite: 0,
-							total: 0,
-						},
-					};
-					calculateCost(model, output.usage);
+					output.usage = parseChunkUsage(chunk.usage, model);
 				}
 
 				const choice = chunk.choices?.[0];
 				if (!choice) continue;
 
+				// Fallback: some providers (e.g., Moonshot) return usage
+				// in choice.usage instead of the standard chunk.usage
+				if (!chunk.usage && (choice as any).usage) {
+					output.usage = parseChunkUsage((choice as any).usage, model);
+				}
+
 				if (choice.finish_reason) {
 					output.stopReason = mapStopReason(choice.finish_reason);
 				}
@@ -717,6 +702,34 @@ function convertTools(
 	}));
 }
 
+function parseChunkUsage(
+	rawUsage: {
+		prompt_tokens?: number;
+		completion_tokens?: number;
+		prompt_tokens_details?: { cached_tokens?: number };
+		completion_tokens_details?: { reasoning_tokens?: number };
+	},
+	model: Model<"openai-completions">,
+): AssistantMessage["usage"] {
+	const cachedTokens = rawUsage.prompt_tokens_details?.cached_tokens || 0;
+	const reasoningTokens = rawUsage.completion_tokens_details?.reasoning_tokens || 0;
+	// OpenAI includes cached tokens in prompt_tokens, so subtract to get non-cached input
+	const input = (rawUsage.prompt_tokens || 0) - cachedTokens;
+	// Compute totalTokens ourselves since we add reasoning_tokens to output
+	// and some providers (e.g., Groq) don't include them in total_tokens
+	const outputTokens = (rawUsage.completion_tokens || 0) + reasoningTokens;
+	const usage: AssistantMessage["usage"] = {
+		input,
+		output: outputTokens,
+		cacheRead: cachedTokens,
+		cacheWrite: 0,
+		totalTokens: input + outputTokens + cachedTokens,
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+	};
+	calculateCost(model, usage);
+	return usage;
+}
+
 function mapStopReason(reason: ChatCompletionChunk.Choice["finish_reason"]): StopReason {
 	if (reason === null) return "stop";
 	switch (reason) {