@vibescope/mcp-server 0.2.5 → 0.2.6

package/src/token-tracking.ts

@@ -1,164 +1,166 @@
-/**
- * Token Tracking Utilities
- *
- * Functions for estimating and tracking token usage across MCP tool calls.
- * Extracted from index.ts to enable unit testing.
- */
-
-// ============================================================================
-// Types
-// ============================================================================
-
-export interface ModelTokens {
-	input: number;
-	output: number;
-}
-
-export interface TokenUsage {
-	callCount: number;
-	totalTokens: number;
-	byTool: Record<string, { calls: number; tokens: number }>;
-	byModel: Record<string, ModelTokens>;
-	currentModel: string | null;
-}
-
-// ============================================================================
-// Token Estimation
-// ============================================================================
-
-/**
- * Estimate tokens from a JSON-serializable object.
- * Uses a rough heuristic of ~4 characters per token.
- *
- * @param obj - Any JSON-serializable object
- * @returns Estimated token count (always >= 1)
- */
-export function estimateTokens(obj: unknown): number {
-	try {
-		const json = JSON.stringify(obj);
-		return Math.max(1, Math.ceil(json.length / 4));
-	} catch {
-		// If JSON serialization fails, return a minimal estimate
-		return 1;
-	}
-}
-
-// ============================================================================
-// Token Usage Tracking
-// ============================================================================
-
-/**
- * Create a fresh token usage tracker.
- */
-export function createTokenUsage(): TokenUsage {
-	return {
-		callCount: 0,
-		totalTokens: 0,
-		byTool: {},
-		byModel: {},
-		currentModel: null,
-	};
-}
-
-/**
- * Track token usage for a tool call.
- * Updates the usage object in-place with input/output token estimates.
- *
- * @param usage - The token usage object to update
- * @param toolName - Name of the tool being called
- * @param args - Input arguments to the tool
- * @param response - Response from the tool
- */
-export function trackTokenUsage(
-	usage: TokenUsage,
-	toolName: string,
-	args: unknown,
-	response: unknown
-): void {
-	const inputTokens = estimateTokens(args);
-	const outputTokens = estimateTokens(response);
-	const totalTokens = inputTokens + outputTokens;
-
-	usage.callCount++;
-	usage.totalTokens += totalTokens;
-
-	if (!usage.byTool[toolName]) {
-		usage.byTool[toolName] = { calls: 0, tokens: 0 };
-	}
-	usage.byTool[toolName].calls++;
-	usage.byTool[toolName].tokens += totalTokens;
-
-	// Track by model if a model is set
-	const model = usage.currentModel;
-	if (model) {
-		if (!usage.byModel[model]) {
-			usage.byModel[model] = { input: 0, output: 0 };
-		}
-		usage.byModel[model].input += inputTokens;
-		usage.byModel[model].output += outputTokens;
-	}
-}
-
-/**
- * Set the current model for token tracking.
- * Subsequent calls to trackTokenUsage will be attributed to this model.
- *
- * @param usage - The token usage object to update
- * @param model - Model name (e.g., "opus", "sonnet", "haiku")
- */
-export function setCurrentModel(usage: TokenUsage, model: string | null): void {
-	usage.currentModel = model;
-}
-
-/**
- * Reset token usage tracking to initial state.
- *
- * @param usage - The token usage object to reset
- */
-export function resetTokenUsage(usage: TokenUsage): void {
-	usage.callCount = 0;
-	usage.totalTokens = 0;
-	usage.byTool = {};
-	usage.byModel = {};
-	usage.currentModel = null;
-}
-
-/**
- * Get a summary of token usage for reporting.
- *
- * @param usage - The token usage object
- * @returns Summary object with stats
- */
-export function getTokenUsageSummary(usage: TokenUsage): {
-	total_calls: number;
-	total_tokens: number;
-	average_tokens_per_call: number;
-	by_tool: Record<string, { calls: number; tokens: number; avg: number }>;
-	by_model: Record<string, ModelTokens>;
-	current_model: string | null;
-} {
-	const byTool: Record<string, { calls: number; tokens: number; avg: number }> = {};
-
-	for (const [tool, stats] of Object.entries(usage.byTool)) {
-		byTool[tool] = {
-			calls: stats.calls,
-			tokens: stats.tokens,
-			avg: stats.calls > 0 ? Math.round(stats.tokens / stats.calls) : 0,
-		};
-	}
-
-	// Deep copy byModel to prevent mutation
-	const byModel: Record<string, ModelTokens> = {};
-	for (const [model, tokens] of Object.entries(usage.byModel)) {
-		byModel[model] = { input: tokens.input, output: tokens.output };
-	}
-
-	return {
-		total_calls: usage.callCount,
-		total_tokens: usage.totalTokens,
-		average_tokens_per_call: usage.callCount > 0 ? Math.round(usage.totalTokens / usage.callCount) : 0,
-		by_tool: byTool,
-		by_model: byModel,
-		current_model: usage.currentModel,
-	};
-}
+/**
+ * Token Tracking Utilities
+ *
+ * Functions for estimating and tracking token usage across MCP tool calls.
+ * Extracted from index.ts to enable unit testing.
+ */
+
+// ============================================================================
+// Types
+// ============================================================================
+
+export interface ModelTokens {
+	input: number;
+	output: number;
+}
+
+export interface TokenUsage {
+	callCount: number;
+	totalTokens: number;
+	byTool: Record<string, { calls: number; tokens: number }>;
+	byModel: Record<string, ModelTokens>;
+	currentModel: string | null;
+}
+
+// ============================================================================
+// Token Estimation
+// ============================================================================
+
+/**
+ * Estimate tokens from a JSON-serializable object.
+ * Uses a rough heuristic of ~4 characters per token.
+ *
+ * @param obj - Any JSON-serializable object
+ * @returns Estimated token count (always >= 1)
+ */
+export function estimateTokens(obj: unknown): number {
+	try {
+		const json = JSON.stringify(obj);
+		return Math.max(1, Math.ceil(json.length / 4));
+	} catch (error) {
+		// Log warning when serialization fails (e.g., circular references, BigInt)
+		const errorMessage = error instanceof Error ? error.message : String(error);
+		console.warn(`[Vibescope] Token estimation failed: ${errorMessage}. Returning minimal estimate of 1 token.`);
+		return 1;
+	}
+}
+
+// ============================================================================
+// Token Usage Tracking
+// ============================================================================
+
+/**
+ * Create a fresh token usage tracker.
+ */
+export function createTokenUsage(): TokenUsage {
+	return {
+		callCount: 0,
+		totalTokens: 0,
+		byTool: {},
+		byModel: {},
+		currentModel: null,
+	};
+}
+
+/**
+ * Track token usage for a tool call.
+ * Updates the usage object in-place with input/output token estimates.
+ *
+ * @param usage - The token usage object to update
+ * @param toolName - Name of the tool being called
+ * @param args - Input arguments to the tool
+ * @param response - Response from the tool
+ */
+export function trackTokenUsage(
+	usage: TokenUsage,
+	toolName: string,
+	args: unknown,
+	response: unknown
+): void {
+	const inputTokens = estimateTokens(args);
+	const outputTokens = estimateTokens(response);
+	const totalTokens = inputTokens + outputTokens;
+
+	usage.callCount++;
+	usage.totalTokens += totalTokens;
+
+	if (!usage.byTool[toolName]) {
+		usage.byTool[toolName] = { calls: 0, tokens: 0 };
+	}
+	usage.byTool[toolName].calls++;
+	usage.byTool[toolName].tokens += totalTokens;
+
+	// Track by model if a model is set
+	const model = usage.currentModel;
+	if (model) {
+		if (!usage.byModel[model]) {
+			usage.byModel[model] = { input: 0, output: 0 };
+		}
+		usage.byModel[model].input += inputTokens;
+		usage.byModel[model].output += outputTokens;
+	}
+}
+
+/**
+ * Set the current model for token tracking.
+ * Subsequent calls to trackTokenUsage will be attributed to this model.
+ *
+ * @param usage - The token usage object to update
+ * @param model - Model name (e.g., "opus", "sonnet", "haiku")
+ */
+export function setCurrentModel(usage: TokenUsage, model: string | null): void {
+	usage.currentModel = model;
+}
+
+/**
+ * Reset token usage tracking to initial state.
+ *
+ * @param usage - The token usage object to reset
+ */
+export function resetTokenUsage(usage: TokenUsage): void {
+	usage.callCount = 0;
+	usage.totalTokens = 0;
+	usage.byTool = {};
+	usage.byModel = {};
+	usage.currentModel = null;
+}
+
+/**
+ * Get a summary of token usage for reporting.
+ *
+ * @param usage - The token usage object
+ * @returns Summary object with stats
+ */
+export function getTokenUsageSummary(usage: TokenUsage): {
+	total_calls: number;
+	total_tokens: number;
+	average_tokens_per_call: number;
+	by_tool: Record<string, { calls: number; tokens: number; avg: number }>;
+	by_model: Record<string, ModelTokens>;
+	current_model: string | null;
+} {
+	const byTool: Record<string, { calls: number; tokens: number; avg: number }> = {};
+
+	for (const [tool, stats] of Object.entries(usage.byTool)) {
+		byTool[tool] = {
+			calls: stats.calls,
+			tokens: stats.tokens,
+			avg: stats.calls > 0 ? Math.round(stats.tokens / stats.calls) : 0,
+		};
+	}
+
+	// Deep copy byModel to prevent mutation
+	const byModel: Record<string, ModelTokens> = {};
+	for (const [model, tokens] of Object.entries(usage.byModel)) {
+		byModel[model] = { input: tokens.input, output: tokens.output };
+	}
+
+	return {
+		total_calls: usage.callCount,
+		total_tokens: usage.totalTokens,
+		average_tokens_per_call: usage.callCount > 0 ? Math.round(usage.totalTokens / usage.callCount) : 0,
+		by_tool: byTool,
+		by_model: byModel,
+		current_model: usage.currentModel,
+	};
+}