Use openrouter stream_options include_usage (#1905)

Author: mrubens

src/api/providers/__tests__/openrouter.test.ts

@@ -112,6 +112,16 @@ describe("OpenRouterHandler", () => {
 						},
 					],
 				}
+				// Add usage information in the stream response
+				yield {
+					id: "test-id",
+					choices: [{ delta: {} }],
+					usage: {
+						prompt_tokens: 10,
+						completion_tokens: 20,
+						cost: 0.001,
+					},
+				}
 			},
 		}
 
@@ -121,17 +131,6 @@ describe("OpenRouterHandler", () => {
 			completions: { create: mockCreate },
 		} as any
 
-		// Mock axios.get for generation details
-		;(axios.get as jest.Mock).mockResolvedValue({
-			data: {
-				data: {
-					native_tokens_prompt: 10,
-					native_tokens_completion: 20,
-					total_cost: 0.001,
-				},
-			},
-		})
-
 		const systemPrompt = "test system prompt"
 		const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "test message" }]
 
@@ -153,7 +152,6 @@ describe("OpenRouterHandler", () => {
 			inputTokens: 10,
 			outputTokens: 20,
 			totalCost: 0.001,
-			fullResponseText: "test response",
 		})
 
 		// Verify OpenAI client was called with correct parameters

src/api/providers/openrouter.ts

@@ -24,11 +24,6 @@ type OpenRouterChatCompletionParams = OpenAI.Chat.ChatCompletionCreateParams & {
 	thinking?: BetaThinkingConfigParam
 }
 
-// Add custom interface for OpenRouter usage chunk.
-interface OpenRouterApiStreamUsageChunk extends ApiStreamUsageChunk {
-	fullResponseText: string
-}
-
 export class OpenRouterHandler extends BaseProvider implements SingleCompletionHandler {
 	protected options: ApiHandlerOptions
 	private client: OpenAI
@@ -110,7 +105,7 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
 			top_p: topP,
 			messages: openAiMessages,
 			stream: true,
-			include_reasoning: true,
+			stream_options: { include_usage: true },
 			// Only include provider if openRouterSpecificProvider is not "[default]".
 			...(this.options.openRouterSpecificProvider &&
 				this.options.openRouterSpecificProvider !== OPENROUTER_DEFAULT_PROVIDER_NAME && {
@@ -122,7 +117,7 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
 
 		const stream = await this.client.chat.completions.create(completionParams)
 
-		let genId: string | undefined
+		let lastUsage
 
 		for await (const chunk of stream as unknown as AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk>) {
 			// OpenRouter returns an error object instead of the OpenAI SDK throwing an error.
@@ -132,10 +127,6 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
 				throw new Error(`OpenRouter API Error ${error?.code}: ${error?.message}`)
 			}
 
-			if (!genId && chunk.id) {
-				genId = chunk.id
-			}
-
 			const delta = chunk.choices[0]?.delta
 
 			if ("reasoning" in delta && delta.reasoning) {
@@ -146,47 +137,23 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
 				fullResponseText += delta.content
 				yield { type: "text", text: delta.content } as ApiStreamChunk
 			}
-		}
-
-		const endpoint = `${this.client.baseURL}/generation?id=${genId}`
 
-		const config: AxiosRequestConfig = {
-			headers: { Authorization: `Bearer ${this.options.openRouterApiKey}` },
-			timeout: 3_000,
+			if (chunk.usage) {
+				lastUsage = chunk.usage
+			}
 		}
 
-		let attempt = 0
-		let lastError: Error | undefined
-		const startTime = Date.now()
-
-		while (attempt++ < 10) {
-			await delay(attempt * 100) // Give OpenRouter some time to produce the generation metadata.
-
-			try {
-				const response = await axios.get(endpoint, config)
-				const generation = response.data?.data
-
-				yield {
-					type: "usage",
-					inputTokens: generation?.native_tokens_prompt || 0,
-					outputTokens: generation?.native_tokens_completion || 0,
-					totalCost: generation?.total_cost || 0,
-					fullResponseText,
-				} as OpenRouterApiStreamUsageChunk
-
-				break
-			} catch (error: unknown) {
-				if (error instanceof Error) {
-					lastError = error
-				}
-			}
+		if (lastUsage) {
+			yield this.processUsageMetrics(lastUsage)
 		}
+	}
 
-		if (lastError) {
-			console.error(
-				`Failed to fetch OpenRouter generation details after attempt #${attempt} (${Date.now() - startTime}ms) [${genId}]`,
-				lastError,
-			)
+	processUsageMetrics(usage: any): ApiStreamUsageChunk {
+		return {
+			type: "usage",
+			inputTokens: usage?.prompt_tokens || 0,
+			outputTokens: usage?.completion_tokens || 0,
+			totalCost: usage?.cost || 0,
 		}
 	}