From f3deacb74e3d4b961720f66a4c93e975b5821702 Mon Sep 17 00:00:00 2001
From: Ralph Khreish <35776126+Crunchyman-ralph@users.noreply.github.com>
Date: Thu, 2 Oct 2025 16:47:45 +0200
Subject: [PATCH] fix: pricing show 0 when it is defined

---
 scripts/modules/ai-services-unified.js | 43 ++++++++++++++++++++------
 scripts/modules/ui.js                  |  8 +++--
 src/ai-providers/base-provider.js      | 12 +++----
 3 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/scripts/modules/ai-services-unified.js b/scripts/modules/ai-services-unified.js
index e1ba0b6b..ac64d101 100644
--- a/scripts/modules/ai-services-unified.js
+++ b/scripts/modules/ai-services-unified.js
@@ -93,31 +93,55 @@ function _getProvider(providerName) {
 
 // Helper function to get cost for a specific model
 function _getCostForModel(providerName, modelId) {
-	const DEFAULT_COST = { inputCost: 0, outputCost: 0, currency: 'USD' };
+	const DEFAULT_COST = {
+		inputCost: 0,
+		outputCost: 0,
+		currency: 'USD',
+		isUnknown: false
+	};
 
 	if (!MODEL_MAP || !MODEL_MAP[providerName]) {
 		log(
 			'warn',
 			`Provider "${providerName}" not found in MODEL_MAP. Cannot determine cost for model ${modelId}.`
 		);
-		return DEFAULT_COST;
+		return { ...DEFAULT_COST, isUnknown: true };
 	}
 
 	const modelData = MODEL_MAP[providerName].find((m) => m.id === modelId);
 
-	if (!modelData?.cost_per_1m_tokens) {
+	if (!modelData) {
 		log(
 			'debug',
-			`Cost data not found for model "${modelId}" under provider "${providerName}". Assuming zero cost.`
+			`Model "${modelId}" not found under provider "${providerName}". Assuming unknown cost.`
 		);
-		return DEFAULT_COST;
+		return { ...DEFAULT_COST, isUnknown: true };
+	}
+
+	// Check if cost_per_1m_tokens is explicitly null (unknown pricing)
+	if (modelData.cost_per_1m_tokens === null) {
+		log(
+			'debug',
+			`Cost data is null for model "${modelId}" under provider "${providerName}". Pricing unknown.`
+		);
+		return { ...DEFAULT_COST, isUnknown: true };
+	}
+
+	// Check if cost_per_1m_tokens is missing/undefined (also unknown)
+	if (modelData.cost_per_1m_tokens === undefined) {
+		log(
+			'debug',
+			`Cost data not found for model "${modelId}" under provider "${providerName}". Pricing unknown.`
+		);
+		return { ...DEFAULT_COST, isUnknown: true };
 	}
 
 	const costs = modelData.cost_per_1m_tokens;
 	return {
 		inputCost: costs.input || 0,
 		outputCost: costs.output || 0,
-		currency: costs.currency || 'USD'
+		currency: costs.currency || 'USD',
+		isUnknown: false
 	};
 }
 
@@ -867,8 +891,8 @@ async function logAiUsage({
 		const timestamp = new Date().toISOString();
 		const totalTokens = (inputTokens || 0) + (outputTokens || 0);
 
-		// Destructure currency along with costs
-		const { inputCost, outputCost, currency } = _getCostForModel(
+		// Destructure currency along with costs and unknown flag
+		const { inputCost, outputCost, currency, isUnknown } = _getCostForModel(
 			providerName,
 			modelId
 		);
@@ -890,7 +914,8 @@ async function logAiUsage({
 			outputTokens: outputTokens || 0,
 			totalTokens,
 			totalCost,
-			currency // Add currency to the telemetry data
+			currency, // Add currency to the telemetry data
+			isUnknownCost: isUnknown // Flag to indicate if pricing is unknown
 		};
 
 		if (getDebugFlag()) {
diff --git a/scripts/modules/ui.js b/scripts/modules/ui.js
index 65450d0b..a82dc530 100644
--- a/scripts/modules/ui.js
+++ b/scripts/modules/ui.js
@@ -2310,7 +2310,8 @@ function displayAiUsageSummary(telemetryData, outputType = 'cli') {
 		outputTokens,
 		totalTokens,
 		totalCost,
-		commandName
+		commandName,
+		isUnknownCost
 	} = telemetryData;
 
 	let summary = chalk.bold.blue('AI Usage Summary:') + '\n';
@@ -2320,7 +2321,10 @@ function displayAiUsageSummary(telemetryData, outputType = 'cli') {
 	summary += chalk.gray(
 		`  Tokens: ${totalTokens} (Input: ${inputTokens}, Output: ${outputTokens})\n`
 	);
-	summary += chalk.gray(`  Est. Cost: $${totalCost.toFixed(6)}`);
+
+	// Show "Unknown" if pricing data is not available, otherwise show the cost
+	const costDisplay = isUnknownCost ? 'Unknown' : `$${totalCost.toFixed(6)}`;
+	summary += chalk.gray(`  Est. Cost: ${costDisplay}`);
 
 	console.log(
 		boxen(summary, {
diff --git a/src/ai-providers/base-provider.js b/src/ai-providers/base-provider.js
index 7409ad20..a81843e1 100644
--- a/src/ai-providers/base-provider.js
+++ b/src/ai-providers/base-provider.js
@@ -179,9 +179,9 @@ export class BaseAIProvider {
 			return {
 				text: result.text,
 				usage: {
-					inputTokens: result.usage?.promptTokens,
-					outputTokens: result.usage?.completionTokens,
-					totalTokens: result.usage?.totalTokens
+					inputTokens: result.usage?.inputTokens || 0,
+					outputTokens: result.usage?.outputTokens || 0,
+					totalTokens: result.usage?.totalTokens || 0
 				}
 			};
 		} catch (error) {
@@ -299,9 +299,9 @@ export class BaseAIProvider {
 			return {
 				object: result.object,
 				usage: {
-					inputTokens: result.usage?.promptTokens,
-					outputTokens: result.usage?.completionTokens,
-					totalTokens: result.usage?.totalTokens
+					inputTokens: result.usage?.inputTokens || 0,
+					outputTokens: result.usage?.outputTokens || 0,
+					totalTokens: result.usage?.totalTokens || 0
 				}
 			};
 		} catch (error) {