feat(completions): ensure prompt tokens are always calculated

apps/gateway/src/api.e2e.ts

@@ -991,6 +991,85 @@ describe("e2e", () => {
 		expect(logs[0].usedModel).toBe("custom");
 	});
 
+	test("Prompt tokens are never zero even when provider returns 0", async () => {
+		const res = await app.request("/v1/chat/completions", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+				Authorization: `Bearer real-token`,
+			},
+			body: JSON.stringify({
+				model: "openai/gpt-4o-mini",
+				messages: [
+					{
+						role: "user",
+						content: "ZERO_TOKENS test message",
+					},
+				],
+			}),
+		});
+
+		expect(res.status).toBe(200);
+
+		const json = await res.json();
+
+		// Verify we have usage information
+		expect(json).toHaveProperty("usage");
+		expect(json.usage).toHaveProperty("prompt_tokens");
+		expect(json.usage).toHaveProperty("completion_tokens");
+		expect(json.usage).toHaveProperty("total_tokens");
+
+		// Verify types are numbers
+		expect(typeof json.usage.prompt_tokens).toBe("number");
+		expect(typeof json.usage.completion_tokens).toBe("number");
+		expect(typeof json.usage.total_tokens).toBe("number");
+
+		// Most importantly: prompt_tokens should never be 0, even if provider returns 0
+		expect(json.usage.prompt_tokens).toBeGreaterThan(0);
+
+		// Completion tokens can be non-zero as set by mock
+		expect(json.usage.completion_tokens).toBeGreaterThan(0);
+
+		// Total tokens should be at least as large as prompt tokens
+		expect(json.usage.total_tokens).toBeGreaterThanOrEqual(
+			json.usage.prompt_tokens,
+		);
+	});
+
+	test("Prompt tokens are calculated for streaming when provider returns 0", async () => {
+		const res = await app.request("/v1/chat/completions", {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+				Authorization: `Bearer real-token`,
+			},
+			body: JSON.stringify({
+				model: "openai/gpt-4o-mini",
+				messages: [
+					{
+						role: "user",
+						content: "ZERO_TOKENS streaming test message",
+					},
+				],
+				stream: true,
+			}),
+		});
+
+		expect(res.status).toBe(200);
+
+		const result = await readAll(res.body);
+
+		// Find a usage chunk
+		const usageChunk = result.chunks.find((chunk: any) => chunk.usage);
+		expect(usageChunk).toBeDefined();
+
+		if (usageChunk) {
+			// Verify prompt tokens are calculated and greater than 0
+			expect(usageChunk.usage.prompt_tokens).toBeGreaterThan(0);
+			expect(typeof usageChunk.usage.prompt_tokens).toBe("number");
+		}
+	});
+
 	test("Success when requesting multi-provider model without prefix", async () => {
 		const multiProviderModel = models.find((m) => m.providers.length > 1);
 		if (!multiProviderModel) {

apps/gateway/src/chat/chat.ts

@@ -324,7 +324,7 @@
 					try {
 						const parsed = JSON.parse(content);
 						content = JSON.stringify(parsed);
 					} catch (_e) {}
 				}
 			}
 
@@ -648,9 +648,12 @@
 					},
 				],
 				usage: {
-					prompt_tokens: promptTokens,
-					completion_tokens: completionTokens,
-					total_tokens: totalTokens,
+					prompt_tokens: Math.max(1, promptTokens || 1),
+					completion_tokens: completionTokens || 0,
+					total_tokens: Math.max(
+						1,
+						totalTokens || Math.max(1, promptTokens || 1),
+					),
 					...(reasoningTokens !== null && {
 						reasoning_tokens: reasoningTokens,
 					}),
@@ -689,9 +692,12 @@
 					},
 				],
 				usage: {
-					prompt_tokens: promptTokens,
-					completion_tokens: completionTokens,
-					total_tokens: totalTokens,
+					prompt_tokens: Math.max(1, promptTokens || 1),
+					completion_tokens: completionTokens || 0,
+					total_tokens: Math.max(
+						1,
+						totalTokens || Math.max(1, promptTokens || 1),
+					),
 					...(reasoningTokens !== null && {
 						reasoning_tokens: reasoningTokens,
 					}),
@@ -727,9 +733,12 @@
 						},
 					],
 					usage: {
-						prompt_tokens: promptTokens,
-						completion_tokens: completionTokens,
-						total_tokens: totalTokens,
+						prompt_tokens: Math.max(1, promptTokens || 1),
+						completion_tokens: completionTokens || 0,
+						total_tokens: Math.max(
+							1,
+							totalTokens || Math.max(1, promptTokens || 1),
+						),
 						...(reasoningTokens !== null && {
 							reasoning_tokens: reasoningTokens,
 						}),
@@ -746,10 +755,34 @@
 /**
  * Transforms streaming chunk to OpenAI format for non-OpenAI providers
  */
+// Helper function to calculate prompt tokens when missing or 0
+function calculatePromptTokensFromMessages(messages: any[]): number {
+	try {
+		const chatMessages: ChatMessage[] = messages.map((m: any) => ({
+			role: m.role,
+			content:
+				typeof m.content === "string" ? m.content : JSON.stringify(m.content),
+			name: m.name,
+		}));
+		return encodeChat(chatMessages, DEFAULT_TOKENIZER_MODEL).length;
+	} catch (_error) {
+		return Math.max(
+			1,
+			Math.round(
+				messages.reduce(
+					(acc: number, m: any) => acc + (m.content?.length || 0),
+					0,
+				) / 4,
+			),
+		);
+	}
+}
+
 function transformStreamingChunkToOpenAIFormat(
 	usedProvider: Provider,
 	usedModel: string,
 	data: any,
+	messages: any[],
 ): any {
 	let transformedData = data;
 
@@ -942,11 +975,16 @@
 					],
 					usage: data.usageMetadata
 						? {
-								prompt_tokens: data.usageMetadata.promptTokenCount || 0,
+								prompt_tokens:
+									data.usageMetadata.promptTokenCount > 0
+										? data.usageMetadata.promptTokenCount
+										: calculatePromptTokensFromMessages(messages),
 								completion_tokens: data.usageMetadata.candidatesTokenCount || 0,
 								// Calculate total including reasoning tokens for Google models
 								total_tokens:
-									(data.usageMetadata.promptTokenCount || 0) +
+									(data.usageMetadata.promptTokenCount > 0
+										? data.usageMetadata.promptTokenCount
+										: calculatePromptTokensFromMessages(messages)) +
 									(data.usageMetadata.candidatesTokenCount || 0) +
 									(data.usageMetadata.thoughtsTokenCount || 0),
 								...(data.usageMetadata.thoughtsTokenCount && {
@@ -982,11 +1020,16 @@
 					],
 					usage: data.usageMetadata
 						? {
-								prompt_tokens: data.usageMetadata.promptTokenCount || 0,
+								prompt_tokens:
+									data.usageMetadata.promptTokenCount > 0
+										? data.usageMetadata.promptTokenCount
+										: calculatePromptTokensFromMessages(messages),
 								completion_tokens: data.usageMetadata.candidatesTokenCount || 0,
 								// Calculate total including reasoning tokens for Google models
 								total_tokens:
-									(data.usageMetadata.promptTokenCount || 0) +
+									(data.usageMetadata.promptTokenCount > 0
+										? data.usageMetadata.promptTokenCount
+										: calculatePromptTokensFromMessages(messages)) +
 									(data.usageMetadata.candidatesTokenCount || 0) +
 									(data.usageMetadata.thoughtsTokenCount || 0),
 								...(data.usageMetadata.thoughtsTokenCount && {
@@ -1248,7 +1291,7 @@
 	let rawBody: unknown;
 	try {
 		rawBody = await c.req.json();
 	} catch (_error) {
 		return c.json(
 			{
 				error: {
@@ -2493,7 +2536,7 @@
 									JSON.parse(jsonCandidate);
 									// JSON is valid - end at first newline to exclude SSE fields
 									eventEnd = dataIndex + 6 + firstNewline;
 								} catch (_e) {
 									// JSON is not complete, use the full segment to next data event
 									eventEnd = nextEventIndex;
 								}
@@ -2590,7 +2633,7 @@
 							let finalTotalTokens = totalTokens;
 
 							// Estimate missing tokens if needed using helper function
-							if (finalPromptTokens === null) {
+							if (finalPromptTokens === null || finalPromptTokens === 0) {
 								const estimation = estimateTokens(
 									usedProvider,
 									messages,
@@ -2629,9 +2672,12 @@
 										},
 									],
 									usage: {
-										prompt_tokens: finalPromptTokens || 0,
+										prompt_tokens: Math.max(1, finalPromptTokens || 1),
 										completion_tokens: finalCompletionTokens || 0,
-										total_tokens: finalTotalTokens || 0,
+										total_tokens: Math.max(
+											1,
+											finalTotalTokens || Math.max(1, finalPromptTokens || 1),
+										),
 									},
 								};
 
@@ -2679,6 +2725,7 @@
 								usedProvider,
 								usedModel,
 								data,
+								messages,
 							);
 
 							// For Anthropic, if we have partial usage data, complete it
@@ -2973,14 +3020,26 @@
 								},
 							],
 							usage: {
-								prompt_tokens: Math.round(
-									promptTokens || calculatedPromptTokens || 0,
+								prompt_tokens: Math.max(
+									1,
+									Math.round(
+										promptTokens && promptTokens > 0
+											? promptTokens
+											: calculatedPromptTokens || 1,
+									),
 								),
 								completion_tokens: Math.round(
 									completionTokens || calculatedCompletionTokens || 0,
 								),
 								total_tokens: Math.round(
-									totalTokens || calculatedTotalTokens || 0,
+									totalTokens ||
+										calculatedTotalTokens ||
+										Math.max(
+											1,
+											promptTokens && promptTokens > 0
+												? promptTokens
+												: calculatedPromptTokens || 1,
+										),
 								),
 								...(cachedTokens !== null && {
 									prompt_tokens_details: {

apps/gateway/src/lib/prompt-tokens.spec.ts

@@ -0,0 +1,79 @@
+import { describe, it, expect } from "vitest";
+
+describe("Prompt token calculation", () => {
+	it("should calculate prompt tokens when provider returns 0", () => {
+		// Mock message
+		const messages = [
+			{
+				role: "user",
+				content: "This is a test message to calculate tokens",
+				name: undefined,
+			},
+		];
+
+		// Simulate calculation logic (similar to what we implemented)
+		const calculatePromptTokens = (
+			promptTokenCount: number,
+			messages: any[],
+		): number => {
+			if (promptTokenCount > 0) {
+				return promptTokenCount;
+			}
+
+			// Calculate prompt tokens if missing or 0
+			try {
+				// Simple estimation fallback (as in our implementation)
+				const totalChars = messages.reduce(
+					(acc, m) => acc + (m.content?.length || 0),
+					0,
+				);
+				return Math.max(1, Math.round(totalChars / 4));
+			} catch (_error) {
+				return 1; // Minimum fallback
+			}
+		};
+
+		// Test that 0 prompt tokens are calculated
+		const result = calculatePromptTokens(0, messages);
+		expect(result).toBeGreaterThan(0);
+		expect(typeof result).toBe("number");
+
+		// Test that existing prompt tokens are preserved
+		const result2 = calculatePromptTokens(50, messages);
+		expect(result2).toBe(50);
+	});
+
+	it("should always return at least 1 token", () => {
+		const calculateMinTokens = (promptTokens: number | null): number => {
+			return Math.max(1, promptTokens || 1);
+		};
+
+		expect(calculateMinTokens(0)).toBe(1);
+		expect(calculateMinTokens(null)).toBe(1);
+		expect(calculateMinTokens(undefined as any)).toBe(1);
+		expect(calculateMinTokens(10)).toBe(10);
+	});
+
+	it("should handle empty messages gracefully", () => {
+		const calculatePromptTokens = (
+			promptTokenCount: number,
+			messages: any[],
+		): number => {
+			if (promptTokenCount > 0) {
+				return promptTokenCount;
+			}
+
+			const totalChars = messages.reduce(
+				(acc, m) => acc + (m.content?.length || 0),
+				0,
+			);
+			return Math.max(1, Math.round(totalChars / 4));
+		};
+
+		const result = calculatePromptTokens(0, []);
+		expect(result).toBe(1); // Should return minimum of 1
+
+		const result2 = calculatePromptTokens(0, [{ role: "user", content: "" }]);
+		expect(result2).toBe(1); // Should return minimum of 1 for empty content
+	});
+});

apps/gateway/src/test-utils/mock-openai-server.ts

@@ -53,6 +53,11 @@ mockOpenAIServer.post("/v1/chat/completions", async (c) => {
 		return c.json(sampleErrorResponse);
 	}
 
+	// Check if this request should trigger zero tokens response
+	const shouldReturnZeroTokens = body.messages.some(
+		(msg: any) => msg.role === "user" && msg.content.includes("ZERO_TOKENS"),
+	);
+
 	// Get the user's message to include in the response
 	const userMessage =
 		body.messages.find((msg: any) => msg.role === "user")?.content || "";
@@ -69,6 +74,13 @@ mockOpenAIServer.post("/v1/chat/completions", async (c) => {
 				},
 			},
 		],
+		usage: shouldReturnZeroTokens
+			? {
+					prompt_tokens: 0,
+					completion_tokens: 20,
+					total_tokens: 20,
+				}
+			: sampleChatCompletionResponse.usage,
 	};
 
 	return c.json(response);