@vibescope/mcp-server 0.2.8 → 0.3.0

package/src/token-tracking.test.ts

@@ -1,463 +1,463 @@
-import { describe, it, expect, beforeEach, vi } 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 and log warning', () => {
-		const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
-
-		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);
-
-		// Should log a warning about the serialization failure
-		expect(warnSpy).toHaveBeenCalledTimes(1);
-		expect(warnSpy).toHaveBeenCalledWith(
-			expect.stringContaining('[Vibescope] Token estimation failed:')
-		);
-
-		warnSpy.mockRestore();
-	});
-
-	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');
-	});
-});
+import { describe, it, expect, beforeEach, vi } 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 and log warning', () => {
+		const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+
+		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);
+
+		// Should log a warning about the serialization failure
+		expect(warnSpy).toHaveBeenCalledTimes(1);
+		expect(warnSpy).toHaveBeenCalledWith(
+			expect.stringContaining('[Vibescope] Token estimation failed:')
+		);
+
+		warnSpy.mockRestore();
+	});
+
+	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');
+	});
+});