@vibescope/mcp-server 0.2.0 → 0.2.2

package/src/token-tracking.test.ts

@@ -0,0 +1,453 @@
+import { describe, it, expect, beforeEach } from 'vitest';
+import {
+	estimateTokens,
+	createTokenUsage,
+	trackTokenUsage,
+	setCurrentModel,
+	resetTokenUsage,
+	getTokenUsageSummary,
+	type TokenUsage,
+} from './token-tracking.js';
+
+// ============================================================================
+// estimateTokens Tests
+// ============================================================================
+
+describe('estimateTokens', () => {
+	it('should return 1 for empty object', () => {
+		// "{}" is 2 chars, ceil(2/4) = 1
+		expect(estimateTokens({})).toBe(1);
+	});
+
+	it('should return 1 for empty array', () => {
+		// "[]" is 2 chars, ceil(2/4) = 1
+		expect(estimateTokens([])).toBe(1);
+	});
+
+	it('should return 1 for empty string', () => {
+		// '""' is 2 chars, ceil(2/4) = 1
+		expect(estimateTokens('')).toBe(1);
+	});
+
+	it('should return 1 for null', () => {
+		// "null" is 4 chars, ceil(4/4) = 1
+		expect(estimateTokens(null)).toBe(1);
+	});
+
+	it('should return 1 for boolean', () => {
+		// "true" is 4 chars, ceil(4/4) = 1
+		expect(estimateTokens(true)).toBe(1);
+		// "false" is 5 chars, ceil(5/4) = 2
+		expect(estimateTokens(false)).toBe(2);
+	});
+
+	it('should estimate tokens for simple object', () => {
+		const obj = { name: 'test' };
+		// {"name":"test"} is 15 chars, ceil(15/4) = 4
+		expect(estimateTokens(obj)).toBe(4);
+	});
+
+	it('should estimate tokens for array of strings', () => {
+		const arr = ['one', 'two', 'three'];
+		// ["one","two","three"] is 21 chars, ceil(21/4) = 6
+		expect(estimateTokens(arr)).toBe(6);
+	});
+
+	it('should estimate tokens for nested object', () => {
+		const obj = {
+			user: {
+				name: 'John',
+				age: 30,
+			},
+			active: true,
+		};
+		// Complex object - just verify it returns a reasonable positive number
+		const tokens = estimateTokens(obj);
+		expect(tokens).toBeGreaterThan(0);
+		expect(tokens).toBeLessThan(100); // Sanity check
+	});
+
+	it('should estimate tokens for large object', () => {
+		const obj = {
+			tasks: Array(100)
+				.fill(null)
+				.map((_, i) => ({
+					id: `task-${i}`,
+					title: `Task number ${i}`,
+					status: 'pending',
+				})),
+		};
+		const tokens = estimateTokens(obj);
+		// Should be a large number for 100 tasks
+		expect(tokens).toBeGreaterThan(500);
+	});
+
+	it('should handle numbers', () => {
+		// "12345" is 5 chars, ceil(5/4) = 2
+		expect(estimateTokens(12345)).toBe(2);
+		// "3.14159" is 7 chars, ceil(7/4) = 2
+		expect(estimateTokens(3.14159)).toBe(2);
+	});
+
+	it('should handle undefined by treating as null', () => {
+		// JSON.stringify(undefined) returns undefined, not a string
+		// Our function handles this gracefully
+		const tokens = estimateTokens(undefined);
+		expect(tokens).toBeGreaterThanOrEqual(1);
+	});
+
+	it('should handle circular reference gracefully', () => {
+		const obj: Record<string, unknown> = { name: 'test' };
+		obj.self = obj; // Create circular reference
+
+		// Should not throw, should return minimal estimate
+		const tokens = estimateTokens(obj);
+		expect(tokens).toBe(1);
+	});
+
+	it('should handle objects with toJSON method', () => {
+		const obj = {
+			data: 'test',
+			toJSON() {
+				return { serialized: true };
+			},
+		};
+		// toJSON returns {"serialized":true} which is 18 chars, ceil(18/4) = 5
+		expect(estimateTokens(obj)).toBe(5);
+	});
+});
+
+// ============================================================================
+// createTokenUsage Tests
+// ============================================================================
+
+describe('createTokenUsage', () => {
+	it('should create fresh token usage object', () => {
+		const usage = createTokenUsage();
+
+		expect(usage.callCount).toBe(0);
+		expect(usage.totalTokens).toBe(0);
+		expect(usage.byTool).toEqual({});
+		expect(usage.byModel).toEqual({});
+		expect(usage.currentModel).toBeNull();
+	});
+
+	it('should create independent instances', () => {
+		const usage1 = createTokenUsage();
+		const usage2 = createTokenUsage();
+
+		usage1.callCount = 5;
+		usage1.byTool['test'] = { calls: 1, tokens: 10 };
+
+		expect(usage2.callCount).toBe(0);
+		expect(usage2.byTool).toEqual({});
+	});
+});
+
+// ============================================================================
+// trackTokenUsage Tests
+// ============================================================================
+
+describe('trackTokenUsage', () => {
+	let usage: TokenUsage;
+
+	beforeEach(() => {
+		usage = createTokenUsage();
+	});
+
+	it('should increment call count', () => {
+		trackTokenUsage(usage, 'test_tool', {}, {});
+		expect(usage.callCount).toBe(1);
+
+		trackTokenUsage(usage, 'test_tool', {}, {});
+		expect(usage.callCount).toBe(2);
+	});
+
+	it('should accumulate total tokens', () => {
+		trackTokenUsage(usage, 'tool1', { key: 'value' }, { result: 'ok' });
+		const firstTotal = usage.totalTokens;
+		expect(firstTotal).toBeGreaterThan(0);
+
+		trackTokenUsage(usage, 'tool2', { key: 'value' }, { result: 'ok' });
+		expect(usage.totalTokens).toBeGreaterThan(firstTotal);
+	});
+
+	it('should track by tool name', () => {
+		trackTokenUsage(usage, 'add_task', { title: 'Test' }, { success: true });
+		trackTokenUsage(usage, 'add_task', { title: 'Test 2' }, { success: true });
+		trackTokenUsage(usage, 'complete_task', { id: '123' }, { success: true });
+
+		expect(usage.byTool['add_task'].calls).toBe(2);
+		expect(usage.byTool['complete_task'].calls).toBe(1);
+		expect(usage.byTool['add_task'].tokens).toBeGreaterThan(0);
+	});
+
+	it('should track by model when set', () => {
+		setCurrentModel(usage, 'opus');
+		trackTokenUsage(usage, 'tool1', { data: 'input' }, { data: 'output' });
+
+		expect(usage.byModel['opus']).toBeDefined();
+		expect(usage.byModel['opus'].input).toBeGreaterThan(0);
+		expect(usage.byModel['opus'].output).toBeGreaterThan(0);
+	});
+
+	it('should not track by model when not set', () => {
+		trackTokenUsage(usage, 'tool1', { data: 'input' }, { data: 'output' });
+
+		expect(Object.keys(usage.byModel)).toHaveLength(0);
+	});
+
+	it('should track multiple models separately', () => {
+		setCurrentModel(usage, 'opus');
+		trackTokenUsage(usage, 'tool1', { x: 1 }, { y: 2 });
+
+		setCurrentModel(usage, 'sonnet');
+		trackTokenUsage(usage, 'tool2', { x: 1 }, { y: 2 });
+		trackTokenUsage(usage, 'tool3', { x: 1 }, { y: 2 });
+
+		expect(Object.keys(usage.byModel)).toHaveLength(2);
+		expect(usage.byModel['opus']).toBeDefined();
+		expect(usage.byModel['sonnet']).toBeDefined();
+	});
+
+	it('should handle empty args and response', () => {
+		trackTokenUsage(usage, 'empty_tool', {}, {});
+
+		expect(usage.callCount).toBe(1);
+		expect(usage.totalTokens).toBeGreaterThanOrEqual(2); // At least 2 for {} + {}
+		expect(usage.byTool['empty_tool'].calls).toBe(1);
+	});
+
+	it('should handle large args and response', () => {
+		const largeArgs = { items: Array(1000).fill('item') };
+		const largeResponse = { results: Array(1000).fill({ ok: true }) };
+
+		trackTokenUsage(usage, 'large_tool', largeArgs, largeResponse);
+
+		expect(usage.totalTokens).toBeGreaterThan(1000);
+	});
+});
+
+// ============================================================================
+// setCurrentModel Tests
+// ============================================================================
+
+describe('setCurrentModel', () => {
+	it('should set the current model', () => {
+		const usage = createTokenUsage();
+
+		setCurrentModel(usage, 'opus');
+		expect(usage.currentModel).toBe('opus');
+
+		setCurrentModel(usage, 'sonnet');
+		expect(usage.currentModel).toBe('sonnet');
+	});
+
+	it('should allow clearing the model with null', () => {
+		const usage = createTokenUsage();
+
+		setCurrentModel(usage, 'opus');
+		expect(usage.currentModel).toBe('opus');
+
+		setCurrentModel(usage, null);
+		expect(usage.currentModel).toBeNull();
+	});
+});
+
+// ============================================================================
+// resetTokenUsage Tests
+// ============================================================================
+
+describe('resetTokenUsage', () => {
+	it('should reset all tracking data', () => {
+		const usage = createTokenUsage();
+
+		// Add some data
+		setCurrentModel(usage, 'opus');
+		trackTokenUsage(usage, 'tool1', { a: 1 }, { b: 2 });
+		trackTokenUsage(usage, 'tool2', { c: 3 }, { d: 4 });
+
+		expect(usage.callCount).toBe(2);
+		expect(usage.totalTokens).toBeGreaterThan(0);
+		expect(Object.keys(usage.byTool)).toHaveLength(2);
+
+		// Reset
+		resetTokenUsage(usage);
+
+		expect(usage.callCount).toBe(0);
+		expect(usage.totalTokens).toBe(0);
+		expect(usage.byTool).toEqual({});
+		expect(usage.byModel).toEqual({});
+		expect(usage.currentModel).toBeNull();
+	});
+
+	it('should allow tracking after reset', () => {
+		const usage = createTokenUsage();
+
+		trackTokenUsage(usage, 'tool1', {}, {});
+		resetTokenUsage(usage);
+
+		trackTokenUsage(usage, 'tool2', { x: 1 }, { y: 2 });
+
+		expect(usage.callCount).toBe(1);
+		expect(usage.byTool['tool1']).toBeUndefined();
+		expect(usage.byTool['tool2']).toBeDefined();
+	});
+});
+
+// ============================================================================
+// getTokenUsageSummary Tests
+// ============================================================================
+
+describe('getTokenUsageSummary', () => {
+	it('should return empty summary for new usage', () => {
+		const usage = createTokenUsage();
+		const summary = getTokenUsageSummary(usage);
+
+		expect(summary.total_calls).toBe(0);
+		expect(summary.total_tokens).toBe(0);
+		expect(summary.average_tokens_per_call).toBe(0);
+		expect(summary.by_tool).toEqual({});
+		expect(summary.by_model).toEqual({});
+		expect(summary.current_model).toBeNull();
+	});
+
+	it('should calculate average tokens per call', () => {
+		const usage = createTokenUsage();
+
+		// Track 3 calls with roughly similar sizes
+		trackTokenUsage(usage, 'tool1', { a: 1 }, { b: 2 });
+		trackTokenUsage(usage, 'tool1', { a: 1 }, { b: 2 });
+		trackTokenUsage(usage, 'tool1', { a: 1 }, { b: 2 });
+
+		const summary = getTokenUsageSummary(usage);
+
+		expect(summary.total_calls).toBe(3);
+		expect(summary.average_tokens_per_call).toBeGreaterThan(0);
+		expect(summary.average_tokens_per_call).toBe(
+			Math.round(summary.total_tokens / summary.total_calls)
+		);
+	});
+
+	it('should include per-tool averages', () => {
+		const usage = createTokenUsage();
+
+		trackTokenUsage(usage, 'small_tool', { x: 1 }, { y: 2 });
+		trackTokenUsage(usage, 'small_tool', { x: 1 }, { y: 2 });
+
+		const summary = getTokenUsageSummary(usage);
+
+		expect(summary.by_tool['small_tool'].calls).toBe(2);
+		expect(summary.by_tool['small_tool'].avg).toBeGreaterThan(0);
+		expect(summary.by_tool['small_tool'].avg).toBe(
+			Math.round(summary.by_tool['small_tool'].tokens / summary.by_tool['small_tool'].calls)
+		);
+	});
+
+	it('should include model breakdown', () => {
+		const usage = createTokenUsage();
+
+		setCurrentModel(usage, 'opus');
+		trackTokenUsage(usage, 'tool1', { a: 1 }, { b: 2 });
+
+		setCurrentModel(usage, 'haiku');
+		trackTokenUsage(usage, 'tool2', { c: 3 }, { d: 4 });
+
+		const summary = getTokenUsageSummary(usage);
+
+		expect(summary.by_model['opus']).toBeDefined();
+		expect(summary.by_model['haiku']).toBeDefined();
+		expect(summary.current_model).toBe('haiku');
+	});
+
+	it('should return a copy of byModel to prevent mutation', () => {
+		const usage = createTokenUsage();
+
+		setCurrentModel(usage, 'opus');
+		trackTokenUsage(usage, 'tool1', { a: 1 }, { b: 2 });
+
+		const summary = getTokenUsageSummary(usage);
+
+		// Modify the summary
+		summary.by_model['opus'].input = 999;
+
+		// Original should be unchanged
+		expect(usage.byModel['opus'].input).not.toBe(999);
+	});
+});
+
+// ============================================================================
+// Integration Tests
+// ============================================================================
+
+describe('Token Tracking Integration', () => {
+	it('should track a realistic session workflow', () => {
+		const usage = createTokenUsage();
+
+		// Agent starts with opus model
+		setCurrentModel(usage, 'opus');
+
+		// Start session
+		trackTokenUsage(
+			usage,
+			'start_work_session',
+			{ git_url: 'https://github.com/org/repo', model: 'opus' },
+			{
+				session_id: '123',
+				persona: 'Atlas',
+				next_task: { id: 'task-1', title: 'Fix bug' },
+			}
+		);
+
+		// Update task
+		trackTokenUsage(
+			usage,
+			'update_task',
+			{ task_id: 'task-1', status: 'in_progress', progress_percentage: 25 },
+			{ success: true }
+		);
+
+		// Complete task
+		trackTokenUsage(
+			usage,
+			'complete_task',
+			{ task_id: 'task-1', summary: 'Fixed the bug' },
+			{ success: true, next_task: null }
+		);
+
+		const summary = getTokenUsageSummary(usage);
+
+		expect(summary.total_calls).toBe(3);
+		expect(summary.by_tool['start_work_session'].calls).toBe(1);
+		expect(summary.by_tool['update_task'].calls).toBe(1);
+		expect(summary.by_tool['complete_task'].calls).toBe(1);
+		expect(summary.by_model['opus']).toBeDefined();
+		expect(summary.by_model['opus'].input).toBeGreaterThan(0);
+		expect(summary.by_model['opus'].output).toBeGreaterThan(0);
+	});
+
+	it('should handle session reset', () => {
+		const usage = createTokenUsage();
+
+		// First session
+		setCurrentModel(usage, 'sonnet');
+		trackTokenUsage(usage, 'tool1', {}, {});
+		trackTokenUsage(usage, 'tool2', {}, {});
+
+		const firstSummary = getTokenUsageSummary(usage);
+		expect(firstSummary.total_calls).toBe(2);
+
+		// Reset for new session
+		resetTokenUsage(usage);
+
+		// Second session
+		setCurrentModel(usage, 'haiku');
+		trackTokenUsage(usage, 'tool3', {}, {});
+
+		const secondSummary = getTokenUsageSummary(usage);
+		expect(secondSummary.total_calls).toBe(1);
+		expect(secondSummary.by_tool['tool1']).toBeUndefined();
+		expect(secondSummary.by_tool['tool3']).toBeDefined();
+		expect(secondSummary.current_model).toBe('haiku');
+	});
+});

package/src/token-tracking.ts

@@ -0,0 +1,164 @@
+/**
+ * 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,
+	};
+}