mcp-rubber-duck 1.5.2 → 1.6.1

package/src/providers/enhanced-manager.ts

@@ -2,6 +2,7 @@ import { EnhancedDuckProvider } from './duck-provider-enhanced.js';
 import { ProviderManager } from './manager.js';
 import { ConfigManager } from '../config/config.js';
 import { FunctionBridge } from '../services/function-bridge.js';
+import { UsageService } from '../services/usage.js';
 import { DuckResponse } from '../config/types.js';
 import { ChatOptions, MCPResult } from './types.js';
 import { logger } from '../utils/logger.js';
@@ -11,12 +12,12 @@ export class EnhancedProviderManager extends ProviderManager {
   private functionBridge?: FunctionBridge;
   private mcpEnabled: boolean = false;
 
-  constructor(configManager: ConfigManager, functionBridge?: FunctionBridge) {
-    super(configManager);
+  constructor(configManager: ConfigManager, functionBridge?: FunctionBridge, usageService?: UsageService) {
+    super(configManager, usageService);
     this.functionBridge = functionBridge;
-    this.mcpEnabled = !!functionBridge && 
+    this.mcpEnabled = !!functionBridge &&
       (configManager.getConfig().mcp_bridge?.enabled || false);
-    
+
     if (this.mcpEnabled) {
       this.initializeEnhancedProviders();
     }
@@ -91,6 +92,7 @@ export class EnhancedProviderManager extends ProviderManager {
 
     const provider = this.getEnhancedProvider(providerName);
     const startTime = Date.now();
+    const modelToUse = options?.model || provider.getInfo().model;
 
     try {
       const response = await provider.chat({
@@ -98,6 +100,18 @@ export class EnhancedProviderManager extends ProviderManager {
         ...options,
       });
 
+      // Record usage
+      if (this.usageService && response.usage) {
+        this.usageService.recordUsage(
+          provider.name,
+          response.model,
+          response.usage.promptTokens,
+          response.usage.completionTokens,
+          false,
+          false
+        );
+      }
+
       return {
         provider: provider.name,
         nickname: provider.nickname,
@@ -117,6 +131,11 @@ export class EnhancedProviderManager extends ProviderManager {
         mcpResults: response.mcpResults,
       };
     } catch (error: unknown) {
+      // Record error
+      if (this.usageService) {
+        this.usageService.recordUsage(provider.name, modelToUse, 0, 0, false, true);
+      }
+
       // Try failover if enabled
       if (this.configManager.getConfig().enable_failover && providerName === undefined) {
         const errorMessage = error instanceof Error ? error.message : String(error);
@@ -208,6 +227,7 @@ export class EnhancedProviderManager extends ProviderManager {
 
     const provider = this.getEnhancedProvider(providerName);
     const startTime = Date.now();
+    const modelToUse = options?.model || provider.getInfo().model;
 
     try {
       const response = await provider.retryWithApproval(
@@ -219,6 +239,18 @@ export class EnhancedProviderManager extends ProviderManager {
         }
       );
 
+      // Record usage
+      if (this.usageService && response.usage) {
+        this.usageService.recordUsage(
+          provider.name,
+          response.model,
+          response.usage.promptTokens,
+          response.usage.completionTokens,
+          false,
+          false
+        );
+      }
+
       return {
         provider: provider.name,
         nickname: provider.nickname,
@@ -238,6 +270,11 @@ export class EnhancedProviderManager extends ProviderManager {
         mcpResults: response.mcpResults,
       };
     } catch (error: unknown) {
+      // Record error
+      if (this.usageService) {
+        this.usageService.recordUsage(provider.name, modelToUse, 0, 0, false, true);
+      }
+
       const errorMessage = error instanceof Error ? error.message : String(error);
       throw new Error(`Failed to retry with approval: ${errorMessage}`);
     }

package/src/providers/manager.ts

@@ -2,6 +2,7 @@ import { DuckProvider } from './provider.js';
 import { ConfigManager } from '../config/config.js';
 import { ProviderHealth, DuckResponse } from '../config/types.js';
 import { ChatOptions, ModelInfo } from './types.js';
+import { UsageService } from '../services/usage.js';
 import { logger } from '../utils/logger.js';
 import { getRandomDuckMessage } from '../utils/ascii-art.js';
 
@@ -9,10 +10,12 @@ export class ProviderManager {
   private providers: Map<string, DuckProvider> = new Map();
   private healthStatus: Map<string, ProviderHealth> = new Map();
   protected configManager: ConfigManager;
+  protected usageService?: UsageService;
   private defaultProvider?: string;
 
-  constructor(configManager: ConfigManager) {
+  constructor(configManager: ConfigManager, usageService?: UsageService) {
     this.configManager = configManager;
+    this.usageService = usageService;
     this.initializeProviders();
   }
 
@@ -91,6 +94,7 @@ export class ProviderManager {
   ): Promise<DuckResponse> {
     const provider = this.getProvider(providerName);
     const startTime = Date.now();
+    const modelToUse = options?.model || provider.getInfo().model;
 
     try {
       const response = await provider.chat({
@@ -98,6 +102,18 @@ export class ProviderManager {
         ...options,
       });
 
+      // Record usage
+      if (this.usageService && response.usage) {
+        this.usageService.recordUsage(
+          provider.name,
+          response.model,
+          response.usage.promptTokens,
+          response.usage.completionTokens,
+          false,
+          false
+        );
+      }
+
       return {
         provider: provider.name,
         nickname: provider.nickname,
@@ -115,6 +131,11 @@ export class ProviderManager {
         cached: false,
       };
     } catch (error: unknown) {
+      // Record error
+      if (this.usageService) {
+        this.usageService.recordUsage(provider.name, modelToUse, 0, 0, false, true);
+      }
+
       // Try failover if enabled
       if (this.configManager.getConfig().enable_failover && providerName === undefined) {
         const errorMessage = error instanceof Error ? error.message : String(error);

package/src/server.ts

@@ -13,6 +13,8 @@ import { ConversationManager } from './services/conversation.js';
 import { ResponseCache } from './services/cache.js';
 import { HealthMonitor } from './services/health.js';
 import { MCPClientManager } from './services/mcp-client-manager.js';
+import { PricingService } from './services/pricing.js';
+import { UsageService } from './services/usage.js';
 import { DuckResponse } from './config/types.js';
 import { ApprovalService } from './services/approval.js';
 import { FunctionBridge } from './services/function-bridge.js';
@@ -37,9 +39,14 @@ import { getPendingApprovalsTool } from './tools/get-pending-approvals.js';
 import { approveMCPRequestTool } from './tools/approve-mcp-request.js';
 import { mcpStatusTool } from './tools/mcp-status.js';
 
+// Import usage stats tool
+import { getUsageStatsTool } from './tools/get-usage-stats.js';
+
 export class RubberDuckServer {
   private server: Server;
   private configManager: ConfigManager;
+  private pricingService: PricingService;
+  private usageService: UsageService;
   private providerManager: ProviderManager;
   private enhancedProviderManager?: EnhancedProviderManager;
   private conversationManager: ConversationManager;
@@ -67,9 +74,16 @@ export class RubberDuckServer {
 
     // Initialize managers
     this.configManager = new ConfigManager();
-    this.providerManager = new ProviderManager(this.configManager);
+    const config = this.configManager.getConfig();
+
+    // Initialize pricing and usage services
+    this.pricingService = new PricingService(config.pricing);
+    this.usageService = new UsageService(this.pricingService);
+
+    // Initialize provider manager with usage tracking
+    this.providerManager = new ProviderManager(this.configManager, this.usageService);
     this.conversationManager = new ConversationManager();
-    this.cache = new ResponseCache(this.configManager.getConfig().cache_ttl);
+    this.cache = new ResponseCache(config.cache_ttl);
     this.healthMonitor = new HealthMonitor(this.providerManager);
 
     // Initialize MCP bridge if enabled
@@ -105,10 +119,11 @@ export class RubberDuckServer {
         mcpConfig.trusted_tools_by_server || {}
       );
 
-      // Initialize enhanced provider manager
+      // Initialize enhanced provider manager with usage tracking
       this.enhancedProviderManager = new EnhancedProviderManager(
         this.configManager,
-        this.functionBridge
+        this.functionBridge,
+        this.usageService
       );
 
       this.mcpEnabled = true;
@@ -178,6 +193,10 @@ export class RubberDuckServer {
           case 'duck_debate':
             return await duckDebateTool(this.providerManager, args || {});
 
+          // Usage stats tool
+          case 'get_usage_stats':
+            return getUsageStatsTool(this.usageService, args || {});
+
           // MCP-specific tools
           case 'get_pending_approvals':
             if (!this.approvalService) {
@@ -641,6 +660,22 @@ export class RubberDuckServer {
           required: ['prompt', 'format'],
         },
       },
+      {
+        name: 'get_usage_stats',
+        description:
+          'Get usage statistics for a time period. Shows token counts and costs (when pricing configured).',
+        inputSchema: {
+          type: 'object',
+          properties: {
+            period: {
+              type: 'string',
+              enum: ['today', '7d', '30d', 'all'],
+              default: 'today',
+              description: 'Time period for stats',
+            },
+          },
+        },
+      },
     ];
 
     // Add MCP-specific tools if enabled
@@ -732,6 +767,9 @@ export class RubberDuckServer {
   }
 
   async stop() {
+    // Cleanup usage service (flush pending writes)
+    this.usageService.shutdown();
+
     // Cleanup MCP resources
     if (this.approvalService) {
       this.approvalService.shutdown();

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();
+  }
+}