mcp-rubber-duck 1.5.2 → 1.6.0

package/src/services/pricing.ts

@@ -0,0 +1,155 @@
+import { ModelPricing, PricingConfig } from '../config/types.js';
+import { DEFAULT_PRICING } from '../data/default-pricing.js';
+import { logger } from '../utils/logger.js';
+
+export interface CostCalculation {
+  inputCost: number;
+  outputCost: number;
+  totalCost: number;
+}
+
+/**
+ * Provider name aliases.
+ * Maps common user-provided names to canonical provider names in pricing data.
+ */
+const PROVIDER_ALIASES: Record<string, string> = {
+  gemini: 'google',
+};
+
+/**
+ * PricingService manages token pricing data.
+ *
+ * It merges hardcoded default pricing with optional user config overrides.
+ * User overrides take precedence over defaults.
+ */
+export class PricingService {
+  private pricing: PricingConfig;
+
+  constructor(configPricing?: PricingConfig) {
+    this.pricing = this.mergePricing(DEFAULT_PRICING, configPricing);
+    const providerCount = Object.keys(this.pricing).length;
+    const modelCount = Object.values(this.pricing).reduce(
+      (acc, models) => acc + Object.keys(models).length,
+      0
+    );
+    logger.debug(`PricingService initialized with ${providerCount} providers, ${modelCount} models`);
+  }
+
+  /**
+   * Deep merge pricing configs. Overrides take precedence.
+   */
+  private mergePricing(defaults: PricingConfig, overrides?: PricingConfig): PricingConfig {
+    const result: PricingConfig = {};
+
+    // Deep copy all defaults
+    for (const [provider, models] of Object.entries(defaults)) {
+      result[provider] = {};
+      for (const [model, pricing] of Object.entries(models)) {
+        result[provider][model] = { ...pricing };
+      }
+    }
+
+    if (!overrides) {
+      return result;
+    }
+
+    // Apply overrides
+    for (const [provider, models] of Object.entries(overrides)) {
+      if (!result[provider]) {
+        result[provider] = {};
+      }
+      for (const [model, pricing] of Object.entries(models)) {
+        result[provider][model] = { ...pricing };
+      }
+    }
+
+    return result;
+  }
+
+  /**
+   * Resolve provider name, checking for aliases.
+   * First checks if provider exists directly, then checks aliases.
+   */
+  private resolveProvider(provider: string): string {
+    // Direct match takes precedence
+    if (this.pricing[provider]) {
+      return provider;
+    }
+    // Check aliases
+    return PROVIDER_ALIASES[provider] || provider;
+  }
+
+  /**
+   * Get pricing for a specific provider and model.
+   * Returns undefined if pricing is not configured.
+   * Supports provider aliases (e.g., "gemini" -> "google").
+   */
+  getPricing(provider: string, model: string): ModelPricing | undefined {
+    const resolvedProvider = this.resolveProvider(provider);
+    return this.pricing[resolvedProvider]?.[model];
+  }
+
+  /**
+   * Calculate the cost for a given number of tokens.
+   * Returns null if pricing is not configured for the provider/model.
+   */
+  calculateCost(
+    provider: string,
+    model: string,
+    promptTokens: number,
+    completionTokens: number
+  ): CostCalculation | null {
+    const pricing = this.getPricing(provider, model);
+    if (!pricing) {
+      return null;
+    }
+
+    const inputCost = (promptTokens / 1_000_000) * pricing.inputPricePerMillion;
+    const outputCost = (completionTokens / 1_000_000) * pricing.outputPricePerMillion;
+
+    return {
+      inputCost,
+      outputCost,
+      totalCost: inputCost + outputCost,
+    };
+  }
+
+  /**
+   * Check if pricing is configured for a provider/model combination.
+   * Supports provider aliases.
+   */
+  hasPricingFor(provider: string, model: string): boolean {
+    return this.getPricing(provider, model) !== undefined;
+  }
+
+  /**
+   * Get all pricing data (for debugging/display).
+   * Returns a deep copy to prevent external mutation.
+   */
+  getAllPricing(): PricingConfig {
+    const result: PricingConfig = {};
+    for (const [provider, models] of Object.entries(this.pricing)) {
+      result[provider] = {};
+      for (const [model, pricing] of Object.entries(models)) {
+        result[provider][model] = { ...pricing };
+      }
+    }
+    return result;
+  }
+
+  /**
+   * Get list of all configured providers.
+   */
+  getProviders(): string[] {
+    return Object.keys(this.pricing);
+  }
+
+  /**
+   * Get list of all models for a provider.
+   * Supports provider aliases.
+   */
+  getModelsForProvider(provider: string): string[] {
+    const resolvedProvider = this.resolveProvider(provider);
+    return Object.keys(this.pricing[resolvedProvider] || {});
+  }
+}

package/src/services/usage.ts

@@ -0,0 +1,293 @@
+import { readFileSync, writeFileSync, existsSync, mkdirSync } from 'fs';
+import { join } from 'path';
+import { homedir } from 'os';
+import {
+  UsageData,
+  DailyUsage,
+  ModelUsageStats,
+  UsageTimePeriod,
+  UsageStatsResult,
+} from '../config/types.js';
+import { PricingService } from './pricing.js';
+import { logger } from '../utils/logger.js';
+
+const USAGE_DATA_VERSION = 1;
+const DEFAULT_DEBOUNCE_MS = 5000;
+
+/**
+ * UsageService tracks token usage per model per day.
+ *
+ * Data is stored in ~/.mcp-rubber-duck/data/usage.json
+ * Writes are debounced to avoid excessive disk I/O.
+ */
+export class UsageService {
+  private usagePath: string;
+  private data: UsageData;
+  private pricingService: PricingService;
+  private pendingWrites: number = 0;
+  private writeDebounceTimer: ReturnType<typeof setTimeout> | null = null;
+  private debounceMs: number;
+
+  constructor(pricingService: PricingService, options?: { dataDir?: string; debounceMs?: number }) {
+    this.pricingService = pricingService;
+    this.debounceMs = options?.debounceMs ?? DEFAULT_DEBOUNCE_MS;
+
+    const dataDir = options?.dataDir ?? join(homedir(), '.mcp-rubber-duck', 'data');
+    this.ensureDirectoryExists(dataDir);
+    this.usagePath = join(dataDir, 'usage.json');
+    this.data = this.loadUsage();
+
+    logger.debug(`UsageService initialized, data path: ${this.usagePath}`);
+  }
+
+  private ensureDirectoryExists(dir: string): void {
+    if (!existsSync(dir)) {
+      mkdirSync(dir, { recursive: true });
+      logger.debug(`Created data directory: ${dir}`);
+    }
+  }
+
+  private loadUsage(): UsageData {
+    try {
+      if (existsSync(this.usagePath)) {
+        const raw = readFileSync(this.usagePath, 'utf-8');
+        const data = JSON.parse(raw) as UsageData;
+
+        // Validate structure
+        if (typeof data !== 'object' || data === null) {
+          throw new Error('Invalid usage data: not an object');
+        }
+        if (typeof data.daily !== 'object' || data.daily === null) {
+          throw new Error('Invalid usage data: missing daily object');
+        }
+
+        logger.debug(`Loaded usage data from ${this.usagePath}`);
+        return data;
+      }
+    } catch (error) {
+      logger.warn('Failed to load usage data, starting fresh:', error);
+    }
+    return { version: USAGE_DATA_VERSION, daily: {} };
+  }
+
+  private saveUsage(): void {
+    try {
+      writeFileSync(this.usagePath, JSON.stringify(this.data, null, 2));
+      logger.debug(`Saved usage data to ${this.usagePath}`);
+    } catch (error) {
+      logger.error('Failed to save usage data:', error);
+    }
+  }
+
+  private scheduleSave(): void {
+    this.pendingWrites++;
+    if (this.writeDebounceTimer) {
+      clearTimeout(this.writeDebounceTimer);
+    }
+    this.writeDebounceTimer = setTimeout(() => {
+      this.saveUsage();
+      this.pendingWrites = 0;
+      this.writeDebounceTimer = null;
+    }, this.debounceMs);
+  }
+
+  private getTodayKey(): string {
+    // Use local date to match getStats() behavior
+    const now = new Date();
+    const year = now.getFullYear();
+    const month = String(now.getMonth() + 1).padStart(2, '0');
+    const day = String(now.getDate()).padStart(2, '0');
+    return `${year}-${month}-${day}`;
+  }
+
+  /**
+   * Record usage for a provider/model.
+   */
+  recordUsage(
+    provider: string,
+    model: string,
+    promptTokens: number,
+    completionTokens: number,
+    cacheHit: boolean = false,
+    error: boolean = false
+  ): void {
+    const today = this.getTodayKey();
+
+    // Initialize nested structure if needed
+    if (!this.data.daily[today]) {
+      this.data.daily[today] = {};
+    }
+    if (!this.data.daily[today][provider]) {
+      this.data.daily[today][provider] = {};
+    }
+    if (!this.data.daily[today][provider][model]) {
+      this.data.daily[today][provider][model] = {
+        requests: 0,
+        promptTokens: 0,
+        completionTokens: 0,
+        cacheHits: 0,
+        errors: 0,
+      };
+    }
+
+    const stats = this.data.daily[today][provider][model];
+    stats.requests++;
+    stats.promptTokens += promptTokens;
+    stats.completionTokens += completionTokens;
+    if (cacheHit) stats.cacheHits++;
+    if (error) stats.errors++;
+
+    logger.debug(
+      `Recorded usage: ${provider}/${model} +${promptTokens}/${completionTokens} tokens`
+    );
+
+    this.scheduleSave();
+  }
+
+  /**
+   * Get usage statistics for a time period.
+   */
+  getStats(period: UsageTimePeriod): UsageStatsResult {
+    const today = new Date();
+    today.setHours(0, 0, 0, 0);
+
+    let startDate: Date;
+    switch (period) {
+      case 'today':
+        startDate = today;
+        break;
+      case '7d':
+        startDate = new Date(today.getTime() - 6 * 24 * 60 * 60 * 1000);
+        break;
+      case '30d':
+        startDate = new Date(today.getTime() - 29 * 24 * 60 * 60 * 1000);
+        break;
+      case 'all':
+        startDate = new Date(0);
+        break;
+    }
+
+    const aggregated: DailyUsage = {};
+    const totals: ModelUsageStats = {
+      requests: 0,
+      promptTokens: 0,
+      completionTokens: 0,
+      cacheHits: 0,
+      errors: 0,
+    };
+    const costByProvider: Record<string, number> = {};
+    let totalCost = 0;
+    let hasCostData = false;
+
+    for (const [dateKey, dayData] of Object.entries(this.data.daily)) {
+      // Parse dateKey as local date (not UTC) to match getTodayKey() format
+      // dateKey is "YYYY-MM-DD", we need to parse it as local midnight
+      const [year, month, day] = dateKey.split('-').map(Number);
+      const date = new Date(year, month - 1, day); // month is 0-indexed
+
+      // Skip invalid dates or dates outside range
+      if (isNaN(date.getTime()) || date < startDate || date > today) continue;
+
+      for (const [provider, providerData] of Object.entries(dayData)) {
+        if (!aggregated[provider]) {
+          aggregated[provider] = {};
+          costByProvider[provider] = 0;
+        }
+
+        for (const [model, stats] of Object.entries(providerData)) {
+          if (!aggregated[provider][model]) {
+            aggregated[provider][model] = {
+              requests: 0,
+              promptTokens: 0,
+              completionTokens: 0,
+              cacheHits: 0,
+              errors: 0,
+            };
+          }
+
+          const agg = aggregated[provider][model];
+          agg.requests += stats.requests;
+          agg.promptTokens += stats.promptTokens;
+          agg.completionTokens += stats.completionTokens;
+          agg.cacheHits += stats.cacheHits;
+          agg.errors += stats.errors;
+
+          totals.requests += stats.requests;
+          totals.promptTokens += stats.promptTokens;
+          totals.completionTokens += stats.completionTokens;
+          totals.cacheHits += stats.cacheHits;
+          totals.errors += stats.errors;
+
+          // Calculate cost if pricing available
+          const cost = this.pricingService.calculateCost(
+            provider,
+            model,
+            stats.promptTokens,
+            stats.completionTokens
+          );
+          if (cost) {
+            totalCost += cost.totalCost;
+            costByProvider[provider] += cost.totalCost;
+            hasCostData = true;
+          }
+        }
+      }
+    }
+
+    const result: UsageStatsResult = {
+      period,
+      startDate: this.formatDate(startDate),
+      endDate: this.formatDate(today),
+      usage: aggregated,
+      totals: {
+        ...totals,
+      },
+    };
+
+    if (hasCostData) {
+      result.totals.estimatedCostUSD = totalCost;
+      result.costByProvider = costByProvider;
+    }
+
+    return result;
+  }
+
+  private formatDate(date: Date): string {
+    // Use local date components to match getTodayKey() and date filtering
+    const year = date.getFullYear();
+    const month = String(date.getMonth() + 1).padStart(2, '0');
+    const day = String(date.getDate()).padStart(2, '0');
+    return `${year}-${month}-${day}`;
+  }
+
+  /**
+   * Flush pending writes immediately. Call on shutdown.
+   */
+  shutdown(): void {
+    if (this.writeDebounceTimer) {
+      clearTimeout(this.writeDebounceTimer);
+      this.writeDebounceTimer = null;
+    }
+    if (this.pendingWrites > 0) {
+      this.saveUsage();
+      this.pendingWrites = 0;
+    }
+    logger.debug('UsageService shutdown complete');
+  }
+
+  /**
+   * Get raw usage data (for testing/debugging).
+   * Returns a deep copy to prevent external mutation.
+   */
+  getRawData(): UsageData {
+    return JSON.parse(JSON.stringify(this.data)) as UsageData;
+  }
+
+  /**
+   * Clear all usage data (for testing).
+   */
+  clearData(): void {
+    this.data = { version: USAGE_DATA_VERSION, daily: {} };
+    this.scheduleSave();
+  }
+}

package/src/tools/get-usage-stats.ts

@@ -0,0 +1,109 @@
+import { UsageService } from '../services/usage.js';
+import { UsageTimePeriod } from '../config/types.js';
+import { duckArt } from '../utils/ascii-art.js';
+import { logger } from '../utils/logger.js';
+
+const VALID_PERIODS: UsageTimePeriod[] = ['today', '7d', '30d', 'all'];
+
+export function getUsageStatsTool(
+  usageService: UsageService,
+  args: Record<string, unknown>
+) {
+  const { period = 'today' } = args as { period?: string };
+
+  // Validate period
+  if (!VALID_PERIODS.includes(period as UsageTimePeriod)) {
+    throw new Error(`Invalid period "${period}". Valid options: ${VALID_PERIODS.join(', ')}`);
+  }
+
+  const stats = usageService.getStats(period as UsageTimePeriod);
+
+  // Format output
+  let output = `${duckArt.panel}\n\n`;
+  output += `📊 Usage Statistics: ${formatPeriodLabel(period as UsageTimePeriod)}\n`;
+  output += `━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n`;
+  output += `Period: ${stats.startDate} to ${stats.endDate}\n\n`;
+
+  // Totals section
+  output += `📈 TOTALS\n`;
+  output += `─────────────────────\n`;
+  output += `Requests: ${stats.totals.requests.toLocaleString()}\n`;
+  output += `Prompt Tokens: ${stats.totals.promptTokens.toLocaleString()}\n`;
+  output += `Completion Tokens: ${stats.totals.completionTokens.toLocaleString()}\n`;
+  output += `Total Tokens: ${(stats.totals.promptTokens + stats.totals.completionTokens).toLocaleString()}\n`;
+  output += `Cache Hits: ${stats.totals.cacheHits.toLocaleString()}\n`;
+  output += `Errors: ${stats.totals.errors.toLocaleString()}\n`;
+
+  if (stats.totals.estimatedCostUSD !== undefined) {
+    output += `💰 Estimated Cost: $${formatCost(stats.totals.estimatedCostUSD)} USD\n`;
+  }
+
+  // Per-provider breakdown
+  const providers = Object.keys(stats.usage);
+  if (providers.length > 0) {
+    output += `\n🦆 BY PROVIDER\n`;
+    output += `─────────────────────\n`;
+
+    for (const provider of providers) {
+      const models = stats.usage[provider];
+      output += `\n**${provider}**\n`;
+
+      for (const [model, modelStats] of Object.entries(models)) {
+        output += `  ${model}:\n`;
+        output += `    Requests: ${modelStats.requests.toLocaleString()}\n`;
+        output += `    Tokens: ${modelStats.promptTokens.toLocaleString()} in / ${modelStats.completionTokens.toLocaleString()} out\n`;
+        if (modelStats.cacheHits > 0) {
+          output += `    Cache Hits: ${modelStats.cacheHits.toLocaleString()}\n`;
+        }
+        if (modelStats.errors > 0) {
+          output += `    Errors: ${modelStats.errors.toLocaleString()}\n`;
+        }
+      }
+
+      if (stats.costByProvider && stats.costByProvider[provider] !== undefined) {
+        output += `  💰 Provider Cost: $${formatCost(stats.costByProvider[provider])}\n`;
+      }
+    }
+  } else {
+    output += `\nNo usage data for this period.\n`;
+  }
+
+  // Footer note about cost
+  if (stats.totals.estimatedCostUSD === undefined && stats.totals.requests > 0) {
+    output += `\n💡 Cost estimates not available. Configure pricing in config.json or update to latest version.\n`;
+  }
+
+  logger.info(`Retrieved usage stats for period: ${period}`);
+
+  return {
+    content: [
+      {
+        type: 'text',
+        text: output,
+      },
+    ],
+  };
+}
+
+function formatPeriodLabel(period: UsageTimePeriod): string {
+  switch (period) {
+    case 'today':
+      return 'Today';
+    case '7d':
+      return 'Last 7 Days';
+    case '30d':
+      return 'Last 30 Days';
+    case 'all':
+      return 'All Time';
+  }
+}
+
+function formatCost(cost: number): string {
+  if (cost < 0.01) {
+    return cost.toFixed(6);
+  } else if (cost < 1) {
+    return cost.toFixed(4);
+  } else {
+    return cost.toFixed(2);
+  }
+}