fix(xai): compute totalCost from usage and load dynamic pricing; update test

Author: hannesrudolph

src/api/providers/__tests__/xai.spec.ts

@@ -244,12 +244,13 @@ describe("XAIHandler", () => {
 
 		// Verify the usage data
 		expect(firstChunk.done).toBe(false)
-		expect(firstChunk.value).toEqual({
+		expect(firstChunk.value).toMatchObject({
 			type: "usage",
 			inputTokens: 10,
 			outputTokens: 20,
 			cacheReadTokens: 5,
 			cacheWriteTokens: 15,
+			totalCost: expect.any(Number),
 		})
 	})
 

src/api/providers/xai.ts

@@ -13,13 +13,17 @@ import { DEFAULT_HEADERS } from "./constants"
 import { BaseProvider } from "./base-provider"
 import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
 import { handleOpenAIError } from "./utils/openai-error-handler"
+import { calculateApiCostOpenAI } from "../../shared/cost"
+import type { ModelRecord } from "../../shared/api"
+import { getModels } from "./fetchers/modelCache"
 
 const XAI_DEFAULT_TEMPERATURE = 0
 
 export class XAIHandler extends BaseProvider implements SingleCompletionHandler {
 	protected options: ApiHandlerOptions
 	private client: OpenAI
 	private readonly providerName = "xAI"
+	protected models: ModelRecord = {}
 
 	constructor(options: ApiHandlerOptions) {
 		super()
@@ -39,12 +43,18 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler
 		const id = this.options.apiModelId ?? xaiDefaultModelId
 
 		const staticInfo = (xaiModels as Record<string, any>)[id as any]
+		const dynamicInfo = this.models?.[id as any]
 
-		// Build complete ModelInfo with required fields; dynamic data comes from router models
+		// Build complete ModelInfo using dynamic pricing/capabilities when available
 		const info: ModelInfo = {
 			contextWindow: this.options.xaiModelContextWindow ?? staticInfo?.contextWindow,
 			maxTokens: staticInfo?.maxTokens ?? undefined,
-			supportsPromptCache: false, // Placeholder - actual value comes from dynamic API call
+			supportsPromptCache: dynamicInfo?.supportsPromptCache ?? false,
+			supportsImages: dynamicInfo?.supportsImages,
+			inputPrice: dynamicInfo?.inputPrice,
+			outputPrice: dynamicInfo?.outputPrice,
+			cacheReadsPrice: dynamicInfo?.cacheReadsPrice,
+			cacheWritesPrice: dynamicInfo?.cacheWritesPrice,
 			description: staticInfo?.description,
 			supportsReasoningEffort:
 				staticInfo && "supportsReasoningEffort" in staticInfo ? staticInfo.supportsReasoningEffort : undefined,
@@ -54,11 +64,24 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler
 		return { id, info, ...params }
 	}
 
+	private async loadDynamicModels(): Promise<void> {
+		try {
+			this.models = await getModels({
+				provider: "xai",
+				apiKey: this.options.xaiApiKey,
+				baseUrl: (this.client as any).baseURL || "https://api.x.ai/v1",
+			})
+		} catch (error) {
+			console.error("[XAI] Error loading dynamic models:", error)
+		}
+	}
+
 	override async *createMessage(
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
 		metadata?: ApiHandlerCreateMessageMetadata,
 	): ApiStream {
+		await this.loadDynamicModels()
 		const { id: modelId, info: modelInfo, reasoning } = this.getModel()
 
 		// Use the OpenAI-compatible API.
@@ -107,12 +130,21 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler
 				const writeTokens =
 					"cache_creation_input_tokens" in chunk.usage ? (chunk.usage as any).cache_creation_input_tokens : 0
 
+				const totalCost = calculateApiCostOpenAI(
+					modelInfo,
+					chunk.usage.prompt_tokens || 0,
+					chunk.usage.completion_tokens || 0,
+					writeTokens || 0,
+					readTokens || 0,
+				)
+
 				yield {
 					type: "usage",
 					inputTokens: chunk.usage.prompt_tokens || 0,
 					outputTokens: chunk.usage.completion_tokens || 0,
 					cacheReadTokens: readTokens,
 					cacheWriteTokens: writeTokens,
+					totalCost,
 				}
 			}
 		}