@gguf/coder 0.2.9 → 0.3.1

package/source/commands/usage.spec.tsx

@@ -0,0 +1,495 @@
+import test from 'ava';
+import {render} from 'ink-testing-library';
+import React from 'react';
+import {themes} from '../config/themes.js';
+import {ThemeContext} from '../hooks/useTheme.js';
+import type {Message} from '../types/core.js';
+import {usageCommand} from './usage.js';
+
+console.log(`\nusage.spec.tsx – ${React.version}`);
+
+/**
+ * Note: These tests make real API calls to models.dev via getModelContextLimit().
+ * The API has a cache and fallback mechanisms, so it should be reliable.
+ * If CI fails due to network issues, we may need to add a proper mocking solution.
+ * All tests have reasonable timeouts (default 10s) to prevent hanging.
+ */
+
+// Mock ThemeProvider for testing
+const MockThemeProvider = ({children}: {children: React.ReactNode}) => {
+	const mockTheme = {
+		currentTheme: 'tokyo-night' as const,
+		colors: themes['tokyo-night'].colors,
+		setCurrentTheme: () => {},
+	};
+
+	return (
+		<ThemeContext.Provider value={mockTheme}>{children}</ThemeContext.Provider>
+	);
+};
+
+// Helper to create mock metadata
+function createMockMetadata(provider = 'test-provider', model = 'test-model') {
+	return {
+		provider,
+		model,
+		tokens: 100,
+		getMessageTokens: (message: Message) => {
+			// Simple mock: count characters as rough token estimate
+			return Math.ceil(message.content.length / 4);
+		},
+	};
+}
+
+// Helper to create messages
+function createMessages(): Message[] {
+	return [
+		{role: 'user', content: 'Hello, how are you?'},
+		{role: 'assistant', content: 'I am doing well, thank you!'},
+		{role: 'user', content: 'Can you help me with a task?'},
+	];
+}
+
+// ============================================================================
+// Command Metadata Tests
+// ============================================================================
+
+test('usage command has correct name', t => {
+	t.is(usageCommand.name, 'usage');
+});
+
+test('usage command has description', t => {
+	t.truthy(usageCommand.description);
+	t.is(typeof usageCommand.description, 'string');
+	t.true(usageCommand.description.length > 0);
+});
+
+test('usage command handler is a function', t => {
+	t.is(typeof usageCommand.handler, 'function');
+});
+
+// ============================================================================
+// Command Handler Tests
+// ============================================================================
+
+test('usage command handler returns valid React element', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command handler with empty messages', async t => {
+	const messages: Message[] = [];
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command handler with different providers', async t => {
+	const messages = createMessages();
+
+	// Test with different provider names
+	const providers = ['openai', 'anthropic', 'ollama'];
+
+	for (const provider of providers) {
+		const metadata = createMockMetadata(provider, 'test-model');
+		const result = await usageCommand.handler([], messages, metadata);
+
+		t.truthy(result);
+		t.true(React.isValidElement(result));
+	}
+});
+
+test('usage command handler with different models', async t => {
+	const messages = createMessages();
+
+	// Test with a realistic model name that might be found or use fallback
+	const metadata = createMockMetadata('test-provider', 'gpt-4');
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command handler with system messages', async t => {
+	const messages: Message[] = [
+		{role: 'system', content: 'You are a helpful assistant.'},
+		{role: 'user', content: 'Hello'},
+		{role: 'assistant', content: 'Hi there!'},
+	];
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command handler with tool messages', async t => {
+	const messages: Message[] = [
+		{role: 'user', content: 'List files'},
+		{
+			role: 'assistant',
+			content: '',
+			tool_calls: [
+				{
+					id: 'call_1',
+					function: {name: 'search_files', arguments: {pattern: '**/*.ts'}},
+				},
+			],
+		},
+		{role: 'tool', content: 'file1.ts\nfile2.ts', tool_call_id: 'call_1'},
+	];
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command handler with multiple message types', async t => {
+	const messages: Message[] = [
+		{role: 'system', content: 'You are a helpful assistant.'},
+		{role: 'user', content: 'Hello'},
+		{role: 'assistant', content: 'Hi! How can I help?'},
+		{role: 'user', content: 'Search for files'},
+		{
+			role: 'assistant',
+			content: '',
+			tool_calls: [
+				{
+					id: 'call_1',
+					function: {name: 'search_files', arguments: {pattern: '*.ts'}},
+				},
+			],
+		},
+		{role: 'tool', content: 'Results here', tool_call_id: 'call_1'},
+		{role: 'assistant', content: 'Here are the files I found.'},
+	];
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+// ============================================================================
+// Component Rendering Tests
+// ============================================================================
+
+test('usage command renders without crashing', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	t.truthy(lastFrame());
+});
+
+test('usage command displays provider name', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata('openai', 'gpt-4');
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /openai/i);
+});
+
+test('usage command displays model name', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata('anthropic', 'claude-3-opus');
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /claude-3-opus/);
+});
+
+test('usage command displays context usage header', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /Context Usage/);
+});
+
+test('usage command displays overall usage section', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /Overall Usage/);
+});
+
+test('usage command displays category breakdown', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /Breakdown by Category/);
+});
+
+test('usage command displays system prompt category', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /System Prompt/);
+});
+
+test('usage command displays user messages category', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /User Messages/);
+});
+
+test('usage command displays assistant messages category', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /Assistant Messages/);
+});
+
+test('usage command displays tool messages category', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /Tool Messages/);
+});
+
+test('usage command displays tool definitions category', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /Tool Definitions/);
+});
+
+test('usage command displays model information section', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /Model Information/);
+});
+
+test('usage command displays recent activity section', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /Recent Activity/);
+});
+
+test('usage command displays tokenizer name', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	// Should display "Tokenizer:" label
+	t.regex(output!, /Tokenizer:/);
+});
+
+test('usage command displays available tokens', async t => {
+	const messages = createMessages();
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+	const {lastFrame} = render(<MockThemeProvider>{result}</MockThemeProvider>);
+
+	const output = lastFrame();
+	t.truthy(output);
+	t.regex(output!, /Available:/);
+});
+
+// ============================================================================
+// Edge Cases
+// ============================================================================
+
+test('usage command handles very long messages', async t => {
+	const longContent = 'a'.repeat(10000);
+	const messages: Message[] = [
+		{role: 'user', content: longContent},
+		{role: 'assistant', content: 'Short response'},
+	];
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command handles messages with special characters', async t => {
+	const messages: Message[] = [
+		{role: 'user', content: 'Hello 你好 مرحبا'},
+		{role: 'assistant', content: 'Response with émojis 🚀 and symbols ®™©'},
+	];
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command handles empty message content', async t => {
+	const messages: Message[] = [
+		{role: 'user', content: ''},
+		{role: 'assistant', content: ''},
+	];
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command handles single message', async t => {
+	const messages: Message[] = [{role: 'user', content: 'Only message'}];
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command handles many messages', async t => {
+	// Create 100 messages
+	const messages: Message[] = [];
+	for (let i = 0; i < 100; i++) {
+		messages.push({
+			role: i % 2 === 0 ? 'user' : 'assistant',
+			content: `Message ${i}`,
+		});
+	}
+	const metadata = createMockMetadata();
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+// ============================================================================
+// Custom Token Counter Tests
+// ============================================================================
+
+test('usage command uses custom getMessageTokens function', async t => {
+	const messages = createMessages();
+	let callCount = 0;
+
+	const metadata = {
+		provider: 'test-provider',
+		model: 'test-model',
+		tokens: 100,
+		getMessageTokens: (_message: Message) => {
+			callCount++;
+			return 10; // Fixed token count for testing
+		},
+	};
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+	// Should have been called for each non-system message
+	t.true(callCount > 0);
+});
+
+test('usage command with custom token counter returning zero', async t => {
+	const messages = createMessages();
+
+	const metadata = {
+		provider: 'test-provider',
+		model: 'test-model',
+		tokens: 0,
+		getMessageTokens: () => 0,
+	};
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});
+
+test('usage command with custom token counter returning large values', async t => {
+	const messages = createMessages();
+
+	const metadata = {
+		provider: 'test-provider',
+		model: 'test-model',
+		tokens: 1000000,
+		getMessageTokens: () => 100000,
+	};
+
+	const result = await usageCommand.handler([], messages, metadata);
+
+	t.truthy(result);
+	t.true(React.isValidElement(result));
+});

package/source/commands/usage.tsx

@@ -0,0 +1,100 @@
+/**
+ * /usage command
+ * Displays token usage statistics
+ */
+
+import {UsageDisplay} from '@/components/usage/usage-display';
+import {getToolManager} from '@/message-handler';
+import {getModelContextLimit} from '@/models/index';
+import {createTokenizer} from '@/tokenization/index';
+import type {Command} from '@/types/commands';
+import type {Message} from '@/types/core';
+import {
+	calculateTokenBreakdown,
+	calculateToolDefinitionsTokens,
+} from '@/usage/calculator';
+import {processPromptTemplate} from '@/utils/prompt-processor';
+import React from 'react';
+
+export const usageCommand: Command = {
+	name: 'usage',
+	description: 'Display token usage statistics',
+	handler: async (
+		_args: string[],
+		messages: Message[],
+		metadata: {
+			provider: string;
+			model: string;
+			tokens: number;
+			getMessageTokens: (message: Message) => number;
+		},
+	) => {
+		const {provider, model, getMessageTokens} = metadata;
+
+		// Create tokenizer for accurate breakdown
+		const tokenizer = createTokenizer(provider, model);
+
+		// Generate the system prompt to include in token calculation
+		const toolManager = getToolManager();
+		const systemPrompt = processPromptTemplate();
+
+		// Create system message to include in token calculation
+		const systemMessage: Message = {
+			role: 'system',
+			content: systemPrompt,
+		};
+
+		// Calculate token breakdown from messages including system prompt
+		// Note: We don't use getMessageTokens for the system message since it's freshly generated
+		// and won't be in the cache. Instead, we use the tokenizer directly for accurate counting.
+		const baseBreakdown = calculateTokenBreakdown(
+			[systemMessage, ...messages],
+			tokenizer,
+			message => {
+				// For system message, always use tokenizer directly to avoid cache misses
+				if (message.role === 'system') {
+					return tokenizer.countTokens(message);
+				}
+				// For other messages, use cached token counts
+				return getMessageTokens(message);
+			},
+		);
+
+		// Extract tokenizer name before cleanup
+		const tokenizerName = tokenizer.getName();
+
+		// Clean up tokenizer resources
+		if (tokenizer.free) {
+			tokenizer.free();
+		}
+
+		// Calculate tool definitions tokens and create final breakdown (immutable)
+		// Note: Tool definitions are sent separately to the API and add token overhead
+		const toolDefinitions = toolManager
+			? calculateToolDefinitionsTokens(
+					Object.keys(toolManager.getToolRegistry()).length,
+				)
+			: 0;
+
+		const breakdown = {
+			...baseBreakdown,
+			toolDefinitions,
+			total: baseBreakdown.total + toolDefinitions,
+		};
+
+		// Get context limit from models.dev
+		const contextLimit = await getModelContextLimit(model);
+
+		return React.createElement(UsageDisplay, {
+			key: `usage-${Date.now()}`,
+			provider,
+			model,
+			contextLimit,
+			currentTokens: breakdown.total,
+			breakdown,
+			messages,
+			tokenizerName,
+			getMessageTokens,
+		});
+	},
+};