token-guard 0.1.0

package/src/display.ts

@@ -0,0 +1,160 @@
+import type { UsageStats, TokenGuardOptions } from './types.js';
+import { formatTokens } from './tokenizer.js';
+
+const COLORS = {
+  reset: '\x1b[0m',
+  bold: '\x1b[1m',
+  dim: '\x1b[2m',
+  red: '\x1b[31m',
+  green: '\x1b[32m',
+  yellow: '\x1b[33m',
+  blue: '\x1b[34m',
+  magenta: '\x1b[35m',
+  cyan: '\x1b[36m',
+  white: '\x1b[37m',
+  bgRed: '\x1b[41m',
+  bgYellow: '\x1b[43m',
+};
+
+export function clearLine(): void {
+  process.stderr.write('\x1b[2K\x1b[G');
+}
+
+export function moveCursorUp(lines: number): void {
+  process.stderr.write(`\x1b[${lines}A`);
+}
+
+export function progressBar(current: number, max: number, width: number = 20): string {
+  const percent = Math.min(current / max, 1);
+  const filled = Math.round(width * percent);
+  const empty = width - filled;
+
+  let color = COLORS.green;
+  if (percent >= 0.9) color = COLORS.red;
+  else if (percent >= 0.8) color = COLORS.yellow;
+
+  return `${color}${'█'.repeat(filled)}${COLORS.dim}${'░'.repeat(empty)}${COLORS.reset}`;
+}
+
+export function formatCost(cost: number): string {
+  return `$${cost.toFixed(4)}`;
+}
+
+export function renderStatus(
+  stats: UsageStats,
+  options: TokenGuardOptions,
+  inline: boolean = true
+): string {
+  const { budget, costLimit, warnPercent } = options;
+  const { inputTokens, outputTokens, totalTokens, estimatedCost, model } = stats;
+
+  if (options.quiet) {
+    return '';
+  }
+
+  const lines: string[] = [];
+
+  if (inline) {
+    // Compact inline display
+    let status = `${COLORS.cyan}◆${COLORS.reset} `;
+    status += `${COLORS.bold}${formatTokens(totalTokens)}${COLORS.reset} tokens `;
+    status += `${COLORS.dim}(in:${formatTokens(inputTokens)} out:${formatTokens(outputTokens)})${COLORS.reset} `;
+    status += `${COLORS.green}${formatCost(estimatedCost)}${COLORS.reset} `;
+
+    if (budget) {
+      const percent = (totalTokens / budget) * 100;
+      if (percent >= 100) {
+        status += `${COLORS.bgRed}${COLORS.white} BUDGET EXCEEDED ${COLORS.reset}`;
+      } else if (percent >= warnPercent) {
+        status += `${COLORS.yellow}⚠ ${percent.toFixed(0)}% of budget${COLORS.reset}`;
+      }
+    }
+
+    if (costLimit) {
+      const percent = (estimatedCost / costLimit) * 100;
+      if (percent >= 100) {
+        status += `${COLORS.bgRed}${COLORS.white} COST LIMIT EXCEEDED ${COLORS.reset}`;
+      } else if (percent >= warnPercent) {
+        status += `${COLORS.yellow}⚠ ${percent.toFixed(0)}% of cost limit${COLORS.reset}`;
+      }
+    }
+
+    return status;
+  }
+
+  // Box display
+  const width = 47;
+  const border = '─'.repeat(width - 2);
+
+  lines.push(`${COLORS.dim}┌${border}┐${COLORS.reset}`);
+  lines.push(`${COLORS.dim}│${COLORS.reset}  ${COLORS.bold}${COLORS.cyan}token-guard${COLORS.reset} v0.1.0${' '.repeat(width - 24)}${COLORS.dim}│${COLORS.reset}`);
+
+  if (budget) {
+    const budgetLine = `  Budget: ${formatTokens(budget)} tokens | Warn: ${warnPercent}%`;
+    lines.push(`${COLORS.dim}│${COLORS.reset}${budgetLine}${' '.repeat(width - budgetLine.length - 2)}${COLORS.dim}│${COLORS.reset}`);
+  }
+
+  lines.push(`${COLORS.dim}├${border}┤${COLORS.reset}`);
+
+  // Progress bar
+  if (budget) {
+    const bar = progressBar(totalTokens, budget, 20);
+    const counts = `${formatTokens(totalTokens)} / ${formatTokens(budget)}`;
+    lines.push(`${COLORS.dim}│${COLORS.reset}  ${bar}  ${counts}${' '.repeat(width - 30 - counts.length)}${COLORS.dim}│${COLORS.reset}`);
+  }
+
+  // Token breakdown
+  const breakdown = `  Input: ${formatTokens(inputTokens)} | Output: ${formatTokens(outputTokens)}`;
+  lines.push(`${COLORS.dim}│${COLORS.reset}${breakdown}${' '.repeat(width - breakdown.length - 2)}${COLORS.dim}│${COLORS.reset}`);
+
+  // Cost
+  const costLine = `  Est. Cost: ${formatCost(estimatedCost)} (${model})`;
+  lines.push(`${COLORS.dim}│${COLORS.reset}${costLine}${' '.repeat(width - costLine.length - 2)}${COLORS.dim}│${COLORS.reset}`);
+
+  lines.push(`${COLORS.dim}└${border}┘${COLORS.reset}`);
+
+  return lines.join('\n');
+}
+
+export function renderFinalReport(stats: UsageStats, options: TokenGuardOptions): string {
+  const lines: string[] = [];
+  const duration = stats.endTime
+    ? ((stats.endTime.getTime() - stats.startTime.getTime()) / 1000).toFixed(1)
+    : '?';
+
+  lines.push('');
+  lines.push(`${COLORS.bold}${COLORS.cyan}═══ token-guard Report ═══${COLORS.reset}`);
+  lines.push('');
+  lines.push(`${COLORS.bold}Command:${COLORS.reset} ${stats.command}`);
+  lines.push(`${COLORS.bold}Model:${COLORS.reset}   ${stats.model}`);
+  lines.push(`${COLORS.bold}Duration:${COLORS.reset} ${duration}s`);
+  lines.push('');
+  lines.push(`${COLORS.bold}Tokens:${COLORS.reset}`);
+  lines.push(`  Input:  ${formatTokens(stats.inputTokens).padStart(10)}`);
+  lines.push(`  Output: ${formatTokens(stats.outputTokens).padStart(10)}`);
+  lines.push(`  Total:  ${formatTokens(stats.totalTokens).padStart(10)}`);
+  lines.push('');
+  lines.push(`${COLORS.bold}Estimated Cost:${COLORS.reset} ${COLORS.green}${formatCost(stats.estimatedCost)}${COLORS.reset}`);
+
+  if (stats.budgetExceeded) {
+    lines.push('');
+    lines.push(`${COLORS.bgRed}${COLORS.white} BUDGET EXCEEDED - Process terminated ${COLORS.reset}`);
+  }
+
+  if (stats.costExceeded) {
+    lines.push('');
+    lines.push(`${COLORS.bgRed}${COLORS.white} COST LIMIT EXCEEDED - Process terminated ${COLORS.reset}`);
+  }
+
+  lines.push('');
+
+  return lines.join('\n');
+}
+
+export function warn(message: string): void {
+  console.error(`${COLORS.yellow}⚠ token-guard:${COLORS.reset} ${message}`);
+}
+
+export function error(message: string): void {
+  console.error(`${COLORS.red}✖ token-guard:${COLORS.reset} ${message}`);
+}

package/src/index.ts

@@ -0,0 +1,7 @@
+// token-guard - Monitor and limit token usage for AI coding assistants
+
+export { TokenMonitor } from './monitor.js';
+export { countTokens, formatTokens } from './tokenizer.js';
+export { calculateCost, getPricing, detectModel, MODEL_PRICING } from './pricing.js';
+export { loadConfig, getConfigPath } from './config.js';
+export type { TokenGuardOptions, UsageStats, ModelPricing, Config } from './types.js';

package/src/monitor.ts

@@ -0,0 +1,189 @@
+import { spawn, type ChildProcess } from 'node:child_process';
+import { writeFileSync } from 'node:fs';
+import type { TokenGuardOptions, UsageStats } from './types.js';
+import { countTokens } from './tokenizer.js';
+import { calculateCost, detectModel } from './pricing.js';
+import { renderStatus, renderFinalReport, clearLine, warn, error } from './display.js';
+
+export class TokenMonitor {
+  private stats: UsageStats;
+  private options: TokenGuardOptions;
+  private process: ChildProcess | null = null;
+  private updateInterval: NodeJS.Timeout | null = null;
+  private lastStatusLength = 0;
+
+  constructor(command: string, options: TokenGuardOptions) {
+    const model = options.model || detectModel(command);
+
+    this.options = options;
+    this.stats = {
+      inputTokens: 0,
+      outputTokens: 0,
+      totalTokens: 0,
+      estimatedCost: 0,
+      model,
+      startTime: new Date(),
+      command,
+      budgetExceeded: false,
+      costExceeded: false,
+    };
+  }
+
+  async run(args: string[]): Promise<number> {
+    return new Promise((resolve) => {
+      const [cmd, ...cmdArgs] = args;
+
+      this.process = spawn(cmd, cmdArgs, {
+        stdio: ['inherit', 'pipe', 'pipe'],
+        shell: true,
+      });
+
+      // Track input (what we send to the process)
+      // For now, we count the command itself as input
+      this.addInputTokens(args.join(' '));
+
+      // Track stdout
+      this.process.stdout?.on('data', (data: Buffer) => {
+        const text = data.toString();
+        this.addOutputTokens(text);
+        process.stdout.write(data);
+        this.checkLimits();
+      });
+
+      // Track stderr
+      this.process.stderr?.on('data', (data: Buffer) => {
+        const text = data.toString();
+        this.addOutputTokens(text);
+        process.stderr.write(data);
+        this.checkLimits();
+      });
+
+      // Periodic status update
+      if (!this.options.quiet) {
+        this.updateInterval = setInterval(() => {
+          this.showStatus();
+        }, 1000);
+      }
+
+      this.process.on('close', (code) => {
+        this.cleanup();
+        this.stats.endTime = new Date();
+
+        if (!this.options.quiet) {
+          console.error(renderFinalReport(this.stats, this.options));
+        }
+
+        if (this.options.output) {
+          this.saveReport();
+        }
+
+        resolve(code ?? 0);
+      });
+
+      this.process.on('error', (err) => {
+        this.cleanup();
+        error(`Failed to start command: ${err.message}`);
+        resolve(1);
+      });
+    });
+  }
+
+  private addInputTokens(text: string): void {
+    const tokens = countTokens(text);
+    this.stats.inputTokens += tokens;
+    this.updateTotals();
+  }
+
+  private addOutputTokens(text: string): void {
+    const tokens = countTokens(text);
+    this.stats.outputTokens += tokens;
+    this.updateTotals();
+  }
+
+  private updateTotals(): void {
+    this.stats.totalTokens = this.stats.inputTokens + this.stats.outputTokens;
+    this.stats.estimatedCost = calculateCost(
+      this.stats.inputTokens,
+      this.stats.outputTokens,
+      this.stats.model
+    );
+  }
+
+  private checkLimits(): void {
+    const { budget, costLimit, warnPercent } = this.options;
+
+    // Check token budget
+    if (budget) {
+      const percent = (this.stats.totalTokens / budget) * 100;
+
+      if (percent >= 100) {
+        this.stats.budgetExceeded = true;
+        this.terminate('Token budget exceeded');
+        return;
+      }
+
+      if (percent >= warnPercent && percent < warnPercent + 5) {
+        warn(`${percent.toFixed(0)}% of token budget used`);
+      }
+    }
+
+    // Check cost limit
+    if (costLimit) {
+      const percent = (this.stats.estimatedCost / costLimit) * 100;
+
+      if (percent >= 100) {
+        this.stats.costExceeded = true;
+        this.terminate('Cost limit exceeded');
+        return;
+      }
+
+      if (percent >= warnPercent && percent < warnPercent + 5) {
+        warn(`${percent.toFixed(0)}% of cost limit used`);
+      }
+    }
+  }
+
+  private terminate(reason: string): void {
+    error(reason);
+    if (this.process) {
+      this.process.kill('SIGTERM');
+    }
+  }
+
+  private showStatus(): void {
+    if (this.options.quiet) return;
+
+    const status = renderStatus(this.stats, this.options, true);
+    clearLine();
+    process.stderr.write(status);
+  }
+
+  private cleanup(): void {
+    if (this.updateInterval) {
+      clearInterval(this.updateInterval);
+      this.updateInterval = null;
+    }
+    clearLine();
+  }
+
+  private saveReport(): void {
+    if (!this.options.output) return;
+
+    const report = {
+      ...this.stats,
+      startTime: this.stats.startTime.toISOString(),
+      endTime: this.stats.endTime?.toISOString(),
+      options: this.options,
+    };
+
+    try {
+      writeFileSync(this.options.output, JSON.stringify(report, null, 2));
+    } catch (err) {
+      error(`Failed to save report: ${err}`);
+    }
+  }
+
+  getStats(): UsageStats {
+    return { ...this.stats };
+  }
+}

package/src/pricing.ts

@@ -0,0 +1,66 @@
+import type { ModelPricing } from './types.js';
+
+// Pricing per 1M tokens (as of Jan 2026)
+export const MODEL_PRICING: Record<string, ModelPricing> = {
+  // Claude models
+  'claude-opus-4-5': { input: 15.00, output: 75.00 },
+  'claude-opus-4': { input: 15.00, output: 75.00 },
+  'claude-sonnet-4': { input: 3.00, output: 15.00 },
+  'claude-3-opus': { input: 15.00, output: 75.00 },
+  'claude-3-sonnet': { input: 3.00, output: 15.00 },
+  'claude-3-haiku': { input: 0.25, output: 1.25 },
+
+  // OpenAI models
+  'gpt-4o': { input: 2.50, output: 10.00 },
+  'gpt-4-turbo': { input: 10.00, output: 30.00 },
+  'gpt-4': { input: 30.00, output: 60.00 },
+  'gpt-3.5-turbo': { input: 0.50, output: 1.50 },
+  'o1': { input: 15.00, output: 60.00 },
+  'o1-mini': { input: 3.00, output: 12.00 },
+
+  // Default fallback (sonnet-level pricing)
+  'default': { input: 3.00, output: 15.00 },
+};
+
+export function getPricing(model: string): ModelPricing {
+  const normalizedModel = model.toLowerCase().replace(/[_-]/g, '-');
+
+  for (const [key, pricing] of Object.entries(MODEL_PRICING)) {
+    if (normalizedModel.includes(key.replace(/[_-]/g, '-'))) {
+      return pricing;
+    }
+  }
+
+  return MODEL_PRICING['default'];
+}
+
+export function calculateCost(
+  inputTokens: number,
+  outputTokens: number,
+  model: string
+): number {
+  const pricing = getPricing(model);
+  const inputCost = (inputTokens / 1_000_000) * pricing.input;
+  const outputCost = (outputTokens / 1_000_000) * pricing.output;
+  return inputCost + outputCost;
+}
+
+export function detectModel(command: string): string {
+  const lowerCommand = command.toLowerCase();
+
+  if (lowerCommand.includes('claude')) {
+    if (lowerCommand.includes('opus')) return 'claude-opus-4';
+    if (lowerCommand.includes('haiku')) return 'claude-3-haiku';
+    return 'claude-sonnet-4'; // default for claude commands
+  }
+
+  if (lowerCommand.includes('aider')) {
+    // Aider defaults vary, assume sonnet
+    return 'claude-sonnet-4';
+  }
+
+  if (lowerCommand.includes('gpt-4')) return 'gpt-4o';
+  if (lowerCommand.includes('openai')) return 'gpt-4o';
+
+  return 'default';
+}

package/src/tokenizer.ts

@@ -0,0 +1,57 @@
+// Simple token estimation without external dependencies
+// Uses GPT-4 style tokenization rules (~4 chars per token)
+
+const CHARS_PER_TOKEN = 4;
+
+export function countTokens(text: string): number {
+  if (!text) return 0;
+
+  // More accurate estimation:
+  // - Split on whitespace and punctuation
+  // - Count special tokens for code
+
+  // Count words
+  const words = text.split(/\s+/).filter(w => w.length > 0);
+
+  // Estimate tokens (rough approximation)
+  // - Short words (1-4 chars) = 1 token
+  // - Medium words (5-8 chars) = 1-2 tokens
+  // - Long words (9+ chars) = 2+ tokens
+  // - Code symbols often get their own tokens
+
+  let tokens = 0;
+
+  for (const word of words) {
+    if (word.length <= 4) {
+      tokens += 1;
+    } else if (word.length <= 8) {
+      tokens += Math.ceil(word.length / 4);
+    } else {
+      tokens += Math.ceil(word.length / 3);
+    }
+  }
+
+  // Add tokens for punctuation and special characters
+  const specialChars = text.match(/[{}()\[\]<>:;,."'`~!@#$%^&*+=|\\/?-]/g);
+  if (specialChars) {
+    tokens += Math.ceil(specialChars.length * 0.5);
+  }
+
+  // Add tokens for newlines (often separate tokens)
+  const newlines = text.match(/\n/g);
+  if (newlines) {
+    tokens += newlines.length;
+  }
+
+  return Math.max(1, Math.round(tokens));
+}
+
+export function formatTokens(count: number): string {
+  if (count >= 1_000_000) {
+    return `${(count / 1_000_000).toFixed(2)}M`;
+  }
+  if (count >= 1_000) {
+    return `${(count / 1_000).toFixed(1)}k`;
+  }
+  return count.toString();
+}

package/src/types.ts

@@ -0,0 +1,33 @@
+export interface TokenGuardOptions {
+  budget?: number;
+  costLimit?: number;
+  model?: string;
+  warnPercent: number;
+  quiet: boolean;
+  output?: string;
+}
+
+export interface ModelPricing {
+  input: number;  // per 1M tokens
+  output: number; // per 1M tokens
+}
+
+export interface UsageStats {
+  inputTokens: number;
+  outputTokens: number;
+  totalTokens: number;
+  estimatedCost: number;
+  model: string;
+  startTime: Date;
+  endTime?: Date;
+  command: string;
+  budgetExceeded: boolean;
+  costExceeded: boolean;
+}
+
+export interface Config {
+  defaultBudget?: number;
+  defaultCostLimit?: number;
+  warnPercent: number;
+  models: Record<string, ModelPricing>;
+}

package/tsconfig.json

@@ -0,0 +1,18 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "module": "NodeNext",
+    "moduleResolution": "NodeNext",
+    "outDir": "dist",
+    "rootDir": "src",
+    "strict": true,
+    "esModuleInterop": true,
+    "skipLibCheck": true,
+    "forceConsistentCasingInFileNames": true,
+    "declaration": true,
+    "declarationMap": true,
+    "sourceMap": true
+  },
+  "include": ["src/**/*"],
+  "exclude": ["node_modules", "dist"]
+}