mcp-cost-tracker 1.0.0

package/src/storage/database.ts

@@ -0,0 +1,525 @@
+import Database from 'better-sqlite3';
+import path from 'path';
+import os from 'os';
+import fs from 'fs';
+
+export interface UsageRecord {
+  id?: number;
+  timestamp: string;
+  provider: string;
+  model: string;
+  input_tokens: number;
+  output_tokens: number;
+  cached_input_tokens: number;
+  total_tokens: number;
+  input_cost: number;
+  output_cost: number;
+  cached_input_cost: number;
+  total_cost: number;
+  session_id: string;
+  project: string;
+  description: string;
+  metadata: string;
+}
+
+export interface BudgetRecord {
+  id?: number;
+  name: string;
+  limit_amount: number;
+  period: 'daily' | 'weekly' | 'monthly' | 'total';
+  provider_filter: string;
+  model_filter: string;
+  project_filter: string;
+  created_at: string;
+}
+
+export interface CustomPricingRecord {
+  id?: number;
+  provider: string;
+  model: string;
+  input_price_per_mtok: number;
+  output_price_per_mtok: number;
+  cached_input_price_per_mtok: number;
+  created_at: string;
+}
+
+export class CostDatabase {
+  private db: Database.Database;
+
+  constructor(dbPath?: string) {
+    const defaultDir = path.join(os.homedir(), '.mcp-cost-tracker');
+    if (!fs.existsSync(defaultDir)) {
+      fs.mkdirSync(defaultDir, { recursive: true });
+    }
+    const finalPath = dbPath || path.join(defaultDir, 'costs.db');
+    this.db = new Database(finalPath);
+    this.db.pragma('journal_mode = WAL');
+    this.db.pragma('foreign_keys = ON');
+    this.initialize();
+  }
+
+  private initialize(): void {
+    this.db.exec(`
+      CREATE TABLE IF NOT EXISTS usage_records (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        timestamp TEXT NOT NULL DEFAULT (datetime('now')),
+        provider TEXT NOT NULL,
+        model TEXT NOT NULL,
+        input_tokens INTEGER NOT NULL DEFAULT 0,
+        output_tokens INTEGER NOT NULL DEFAULT 0,
+        cached_input_tokens INTEGER NOT NULL DEFAULT 0,
+        total_tokens INTEGER NOT NULL DEFAULT 0,
+        input_cost REAL NOT NULL DEFAULT 0,
+        output_cost REAL NOT NULL DEFAULT 0,
+        cached_input_cost REAL NOT NULL DEFAULT 0,
+        total_cost REAL NOT NULL DEFAULT 0,
+        session_id TEXT DEFAULT '',
+        project TEXT DEFAULT 'default',
+        description TEXT DEFAULT '',
+        metadata TEXT DEFAULT '{}'
+      );
+
+      CREATE TABLE IF NOT EXISTS budgets (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        name TEXT NOT NULL UNIQUE,
+        limit_amount REAL NOT NULL,
+        period TEXT NOT NULL DEFAULT 'monthly',
+        provider_filter TEXT DEFAULT '',
+        model_filter TEXT DEFAULT '',
+        project_filter TEXT DEFAULT '',
+        created_at TEXT NOT NULL DEFAULT (datetime('now'))
+      );
+
+      CREATE TABLE IF NOT EXISTS custom_pricing (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        provider TEXT NOT NULL,
+        model TEXT NOT NULL,
+        input_price_per_mtok REAL NOT NULL,
+        output_price_per_mtok REAL NOT NULL,
+        cached_input_price_per_mtok REAL DEFAULT 0,
+        created_at TEXT NOT NULL DEFAULT (datetime('now')),
+        UNIQUE(provider, model)
+      );
+
+      CREATE INDEX IF NOT EXISTS idx_usage_timestamp ON usage_records(timestamp);
+      CREATE INDEX IF NOT EXISTS idx_usage_provider ON usage_records(provider);
+      CREATE INDEX IF NOT EXISTS idx_usage_model ON usage_records(model);
+      CREATE INDEX IF NOT EXISTS idx_usage_project ON usage_records(project);
+      CREATE INDEX IF NOT EXISTS idx_usage_session ON usage_records(session_id);
+    `);
+  }
+
+  // ============================================================
+  // Usage Records
+  // ============================================================
+
+  recordUsage(record: Omit<UsageRecord, 'id'>): UsageRecord {
+    const stmt = this.db.prepare(`
+      INSERT INTO usage_records (
+        timestamp, provider, model, input_tokens, output_tokens,
+        cached_input_tokens, total_tokens, input_cost, output_cost,
+        cached_input_cost, total_cost, session_id, project, description, metadata
+      ) VALUES (
+        @timestamp, @provider, @model, @input_tokens, @output_tokens,
+        @cached_input_tokens, @total_tokens, @input_cost, @output_cost,
+        @cached_input_cost, @total_cost, @session_id, @project, @description, @metadata
+      )
+    `);
+    const result = stmt.run(record);
+    return { ...record, id: result.lastInsertRowid as number };
+  }
+
+  getUsageRecords(options: {
+    startDate?: string;
+    endDate?: string;
+    provider?: string;
+    model?: string;
+    project?: string;
+    sessionId?: string;
+    limit?: number;
+    offset?: number;
+  } = {}): UsageRecord[] {
+    let query = 'SELECT * FROM usage_records WHERE 1=1';
+    const params: Record<string, unknown> = {};
+
+    if (options.startDate) {
+      query += ' AND timestamp >= @startDate';
+      params.startDate = options.startDate;
+    }
+    if (options.endDate) {
+      query += ' AND timestamp <= @endDate';
+      params.endDate = options.endDate;
+    }
+    if (options.provider) {
+      query += ' AND provider = @provider';
+      params.provider = options.provider;
+    }
+    if (options.model) {
+      query += ' AND model = @model';
+      params.model = options.model;
+    }
+    if (options.project) {
+      query += ' AND project = @project';
+      params.project = options.project;
+    }
+    if (options.sessionId) {
+      query += ' AND session_id = @sessionId';
+      params.sessionId = options.sessionId;
+    }
+
+    query += ' ORDER BY timestamp DESC';
+
+    if (options.limit) {
+      query += ' LIMIT @limit';
+      params.limit = options.limit;
+    }
+    if (options.offset) {
+      query += ' OFFSET @offset';
+      params.offset = options.offset;
+    }
+
+    return this.db.prepare(query).all(params) as UsageRecord[];
+  }
+
+  // ============================================================
+  // Aggregation Queries
+  // ============================================================
+
+  getSummary(options: {
+    startDate?: string;
+    endDate?: string;
+    provider?: string;
+    project?: string;
+  } = {}): {
+    total_cost: number;
+    total_input_tokens: number;
+    total_output_tokens: number;
+    total_cached_tokens: number;
+    total_requests: number;
+    avg_cost_per_request: number;
+  } {
+    let query = `
+      SELECT
+        COALESCE(SUM(total_cost), 0) as total_cost,
+        COALESCE(SUM(input_tokens), 0) as total_input_tokens,
+        COALESCE(SUM(output_tokens), 0) as total_output_tokens,
+        COALESCE(SUM(cached_input_tokens), 0) as total_cached_tokens,
+        COUNT(*) as total_requests,
+        COALESCE(AVG(total_cost), 0) as avg_cost_per_request
+      FROM usage_records WHERE 1=1
+    `;
+    const params: Record<string, unknown> = {};
+
+    if (options.startDate) {
+      query += ' AND timestamp >= @startDate';
+      params.startDate = options.startDate;
+    }
+    if (options.endDate) {
+      query += ' AND timestamp <= @endDate';
+      params.endDate = options.endDate;
+    }
+    if (options.provider) {
+      query += ' AND provider = @provider';
+      params.provider = options.provider;
+    }
+    if (options.project) {
+      query += ' AND project = @project';
+      params.project = options.project;
+    }
+
+    return this.db.prepare(query).get(params) as any;
+  }
+
+  getCostByModel(options: {
+    startDate?: string;
+    endDate?: string;
+    project?: string;
+  } = {}): Array<{
+    provider: string;
+    model: string;
+    total_cost: number;
+    total_input_tokens: number;
+    total_output_tokens: number;
+    request_count: number;
+  }> {
+    let query = `
+      SELECT
+        provider,
+        model,
+        SUM(total_cost) as total_cost,
+        SUM(input_tokens) as total_input_tokens,
+        SUM(output_tokens) as total_output_tokens,
+        COUNT(*) as request_count
+      FROM usage_records WHERE 1=1
+    `;
+    const params: Record<string, unknown> = {};
+
+    if (options.startDate) {
+      query += ' AND timestamp >= @startDate';
+      params.startDate = options.startDate;
+    }
+    if (options.endDate) {
+      query += ' AND timestamp <= @endDate';
+      params.endDate = options.endDate;
+    }
+    if (options.project) {
+      query += ' AND project = @project';
+      params.project = options.project;
+    }
+
+    query += ' GROUP BY provider, model ORDER BY total_cost DESC';
+
+    return this.db.prepare(query).all(params) as any[];
+  }
+
+  getCostByProvider(options: {
+    startDate?: string;
+    endDate?: string;
+  } = {}): Array<{
+    provider: string;
+    total_cost: number;
+    total_input_tokens: number;
+    total_output_tokens: number;
+    request_count: number;
+  }> {
+    let query = `
+      SELECT
+        provider,
+        SUM(total_cost) as total_cost,
+        SUM(input_tokens) as total_input_tokens,
+        SUM(output_tokens) as total_output_tokens,
+        COUNT(*) as request_count
+      FROM usage_records WHERE 1=1
+    `;
+    const params: Record<string, unknown> = {};
+
+    if (options.startDate) {
+      query += ' AND timestamp >= @startDate';
+      params.startDate = options.startDate;
+    }
+    if (options.endDate) {
+      query += ' AND timestamp <= @endDate';
+      params.endDate = options.endDate;
+    }
+
+    query += ' GROUP BY provider ORDER BY total_cost DESC';
+
+    return this.db.prepare(query).all(params) as any[];
+  }
+
+  getCostByDay(options: {
+    startDate?: string;
+    endDate?: string;
+    provider?: string;
+    project?: string;
+    days?: number;
+  } = {}): Array<{
+    date: string;
+    total_cost: number;
+    total_tokens: number;
+    request_count: number;
+  }> {
+    const days = options.days || 30;
+    let query = `
+      SELECT
+        DATE(timestamp) as date,
+        SUM(total_cost) as total_cost,
+        SUM(total_tokens) as total_tokens,
+        COUNT(*) as request_count
+      FROM usage_records
+      WHERE timestamp >= datetime('now', '-${days} days')
+    `;
+    const params: Record<string, unknown> = {};
+
+    if (options.startDate) {
+      query = query.replace(
+        `timestamp >= datetime('now', '-${days} days')`,
+        'timestamp >= @startDate'
+      );
+      params.startDate = options.startDate;
+    }
+    if (options.endDate) {
+      query += ' AND timestamp <= @endDate';
+      params.endDate = options.endDate;
+    }
+    if (options.provider) {
+      query += ' AND provider = @provider';
+      params.provider = options.provider;
+    }
+    if (options.project) {
+      query += ' AND project = @project';
+      params.project = options.project;
+    }
+
+    query += ' GROUP BY DATE(timestamp) ORDER BY date ASC';
+
+    return this.db.prepare(query).all(params) as any[];
+  }
+
+  getCostByProject(): Array<{
+    project: string;
+    total_cost: number;
+    total_tokens: number;
+    request_count: number;
+  }> {
+    return this.db.prepare(`
+      SELECT
+        project,
+        SUM(total_cost) as total_cost,
+        SUM(total_tokens) as total_tokens,
+        COUNT(*) as request_count
+      FROM usage_records
+      GROUP BY project
+      ORDER BY total_cost DESC
+    `).all() as any[];
+  }
+
+  getTopExpensiveRequests(limit: number = 10): UsageRecord[] {
+    return this.db.prepare(`
+      SELECT * FROM usage_records
+      ORDER BY total_cost DESC
+      LIMIT ?
+    `).all(limit) as UsageRecord[];
+  }
+
+  // ============================================================
+  // Budget Management
+  // ============================================================
+
+  setBudget(budget: Omit<BudgetRecord, 'id' | 'created_at'>): BudgetRecord {
+    const stmt = this.db.prepare(`
+      INSERT OR REPLACE INTO budgets (name, limit_amount, period, provider_filter, model_filter, project_filter)
+      VALUES (@name, @limit_amount, @period, @provider_filter, @model_filter, @project_filter)
+    `);
+    stmt.run(budget);
+    return this.db.prepare('SELECT * FROM budgets WHERE name = ?').get(budget.name) as BudgetRecord;
+  }
+
+  getBudgets(): BudgetRecord[] {
+    return this.db.prepare('SELECT * FROM budgets ORDER BY name').all() as BudgetRecord[];
+  }
+
+  deleteBudget(name: string): boolean {
+    const result = this.db.prepare('DELETE FROM budgets WHERE name = ?').run(name);
+    return result.changes > 0;
+  }
+
+  checkBudget(budgetName: string): {
+    budget: BudgetRecord;
+    spent: number;
+    remaining: number;
+    percentage: number;
+    exceeded: boolean;
+  } | null {
+    const budget = this.db.prepare('SELECT * FROM budgets WHERE name = ?').get(budgetName) as BudgetRecord | undefined;
+    if (!budget) return null;
+
+    let dateFilter = '';
+    switch (budget.period) {
+      case 'daily':
+        dateFilter = "AND timestamp >= datetime('now', 'start of day')";
+        break;
+      case 'weekly':
+        dateFilter = "AND timestamp >= datetime('now', '-7 days')";
+        break;
+      case 'monthly':
+        dateFilter = "AND timestamp >= datetime('now', 'start of month')";
+        break;
+      case 'total':
+        dateFilter = '';
+        break;
+    }
+
+    let providerFilter = '';
+    if (budget.provider_filter) {
+      providerFilter = ` AND provider = '${budget.provider_filter}'`;
+    }
+
+    let modelFilter = '';
+    if (budget.model_filter) {
+      modelFilter = ` AND model = '${budget.model_filter}'`;
+    }
+
+    let projectFilter = '';
+    if (budget.project_filter) {
+      projectFilter = ` AND project = '${budget.project_filter}'`;
+    }
+
+    const result = this.db.prepare(`
+      SELECT COALESCE(SUM(total_cost), 0) as spent
+      FROM usage_records
+      WHERE 1=1 ${dateFilter} ${providerFilter} ${modelFilter} ${projectFilter}
+    `).get() as { spent: number };
+
+    const spent = result.spent;
+    const remaining = Math.max(0, budget.limit_amount - spent);
+    const percentage = budget.limit_amount > 0 ? (spent / budget.limit_amount) * 100 : 0;
+
+    return {
+      budget,
+      spent,
+      remaining,
+      percentage,
+      exceeded: spent >= budget.limit_amount,
+    };
+  }
+
+  // ============================================================
+  // Custom Pricing
+  // ============================================================
+
+  setCustomPricing(record: Omit<CustomPricingRecord, 'id' | 'created_at'>): CustomPricingRecord {
+    const stmt = this.db.prepare(`
+      INSERT OR REPLACE INTO custom_pricing (provider, model, input_price_per_mtok, output_price_per_mtok, cached_input_price_per_mtok)
+      VALUES (@provider, @model, @input_price_per_mtok, @output_price_per_mtok, @cached_input_price_per_mtok)
+    `);
+    stmt.run(record);
+    return this.db.prepare('SELECT * FROM custom_pricing WHERE provider = ? AND model = ?')
+      .get(record.provider, record.model) as CustomPricingRecord;
+  }
+
+  getCustomPricing(): CustomPricingRecord[] {
+    return this.db.prepare('SELECT * FROM custom_pricing ORDER BY provider, model').all() as CustomPricingRecord[];
+  }
+
+  deleteCustomPricing(provider: string, model: string): boolean {
+    const result = this.db.prepare('DELETE FROM custom_pricing WHERE provider = ? AND model = ?').run(provider, model);
+    return result.changes > 0;
+  }
+
+  // ============================================================
+  // Utility
+  // ============================================================
+
+  clearAllData(): void {
+    this.db.exec('DELETE FROM usage_records');
+  }
+
+  getStats(): {
+    total_records: number;
+    first_record: string | null;
+    last_record: string | null;
+    db_size_bytes: number;
+  } {
+    const count = this.db.prepare('SELECT COUNT(*) as count FROM usage_records').get() as { count: number };
+    const first = this.db.prepare('SELECT MIN(timestamp) as ts FROM usage_records').get() as { ts: string | null };
+    const last = this.db.prepare('SELECT MAX(timestamp) as ts FROM usage_records').get() as { ts: string | null };
+
+    const dbPath = this.db.name;
+    let dbSize = 0;
+    try {
+      dbSize = fs.statSync(dbPath).size;
+    } catch { /* ignore */ }
+
+    return {
+      total_records: count.count,
+      first_record: first.ts,
+      last_record: last.ts,
+      db_size_bytes: dbSize,
+    };
+  }
+
+  close(): void {
+    this.db.close();
+  }
+}