feat(db): add token_usage table and usage repository

src/ipc/database/cloudHandler.ts

@@ -126,6 +126,34 @@ function ensureDatabaseInitialized(dbPath: string): void {
         value TEXT NOT NULL
       );
     `);
+
+    // Create token_usage table
+    db.exec(`
+      CREATE TABLE IF NOT EXISTS token_usage (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        account_id TEXT NOT NULL,
+        model TEXT NOT NULL,
+        prompt_tokens INTEGER NOT NULL DEFAULT 0,
+        completion_tokens INTEGER NOT NULL DEFAULT 0,
+        total_tokens INTEGER NOT NULL DEFAULT 0,
+        timestamp INTEGER NOT NULL,
+        request_type TEXT
+      );
+    `);
+
+    // Migration: Check if request_type column exists on token_usage
+    const tokenUsageTableInfoRaw = db.pragma('table_info(token_usage)') as any[];
+    const tokenUsageTableInfo = parseRows(TableInfoRowSchema, tokenUsageTableInfoRaw, 'cloud.token_usage.tableInfo');
+    const hasRequestType = tokenUsageTableInfo.some((col) => col.name === 'request_type');
+    if (!hasRequestType) {
+      db.exec('ALTER TABLE token_usage ADD COLUMN request_type TEXT');
+    }
+
+    // Create indexes for token_usage
+    db.exec(`CREATE INDEX IF NOT EXISTS idx_token_usage_account ON token_usage(account_id);`);
+    db.exec(`CREATE INDEX IF NOT EXISTS idx_token_usage_timestamp ON token_usage(timestamp);`);
+    db.exec(`CREATE INDEX IF NOT EXISTS idx_token_usage_model ON token_usage(model);`);
+    db.exec(`CREATE INDEX IF NOT EXISTS idx_token_usage_account_timestamp ON token_usage(account_id, timestamp);`);
   } catch (error) {
     logger.error('Failed to initialize cloud database schema', error);
     throw error;
@@ -137,7 +165,7 @@ function ensureDatabaseInitialized(dbPath: string): void {
 /**
  * Gets a connection to the cloud accounts database.
  */
-function getCloudDb(): {
+export function getCloudDb(): {
   raw: Database.Database;
   orm: BetterSQLite3Database<typeof drizzleSchema>;
 } {

src/ipc/database/schema.ts

@@ -27,3 +27,14 @@ export const itemTable = sqliteTable('ItemTable', {
   key: text('key').primaryKey(),
   value: text('value'),
 });
+
+export const tokenUsage = sqliteTable('token_usage', {
+  id: integer('id').primaryKey({ autoIncrement: true }),
+  accountId: text('account_id').notNull(),
+  model: text('model').notNull(),
+  promptTokens: integer('prompt_tokens').notNull().default(0),
+  completionTokens: integer('completion_tokens').notNull().default(0),
+  totalTokens: integer('total_tokens').notNull().default(0),
+  timestamp: integer('timestamp').notNull(),
+  requestType: text('request_type'),
+});

src/ipc/database/usageHandler.ts

@@ -0,0 +1,244 @@
+import { logger } from '../../utils/logger';
+import { getCloudDb } from './cloudHandler';
+import { tokenUsage } from './schema';
+import {
+  HourlyUsageSchema,
+  DailyUsageSchema,
+  WeeklyUsageSchema,
+  MonthlyUsageSchema,
+  ModelUsageSchema,
+} from '../../types/db';
+
+export interface RecordUsageParams {
+  accountId: string;
+  model: string;
+  promptTokens: number;
+  completionTokens: number;
+  totalTokens: number;
+  timestamp: number;
+  requestType?: string;
+}
+
+function buildDateFilter(start?: number, end?: number): string {
+  const conditions: string[] = [];
+  if (start !== undefined) {
+    conditions.push(`timestamp >= ${start}`);
+  }
+  if (end !== undefined) {
+    conditions.push(`timestamp <= ${end}`);
+  }
+  return conditions.length > 0 ? `AND ${conditions.join(' AND ')}` : '';
+}
+
+function buildAccountFilter(accountId?: string): string {
+  return accountId ? `AND account_id = '${accountId}'` : '';
+}
+
+export class TokenUsageRepo {
+  static recordUsage(params: RecordUsageParams): void {
+    const { raw, orm } = getCloudDb();
+    try {
+      orm
+        .insert(tokenUsage)
+        .values({
+          accountId: params.accountId,
+          model: params.model,
+          promptTokens: params.promptTokens,
+          completionTokens: params.completionTokens,
+          totalTokens: params.totalTokens,
+          timestamp: params.timestamp,
+          requestType: params.requestType ?? null,
+        })
+        .run();
+    } catch (error) {
+      logger.error('Failed to record token usage', error);
+    } finally {
+      raw.close();
+    }
+  }
+
+  static getUsageByHour(
+    accountId?: string,
+    start?: number,
+    end?: number,
+  ): Array<{
+    bucket: string;
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+    requests: number;
+  }> {
+    const { raw } = getCloudDb();
+    try {
+      const sql = `
+        SELECT
+          strftime('%Y-%m-%d %H:00', timestamp / 1000, 'unixepoch', 'localtime') as bucket,
+          SUM(prompt_tokens) as promptTokens,
+          SUM(completion_tokens) as completionTokens,
+          SUM(total_tokens) as totalTokens,
+          COUNT(*) as requests
+        FROM token_usage
+        WHERE 1=1
+        ${buildAccountFilter(accountId)}
+        ${buildDateFilter(start, end)}
+        GROUP BY bucket
+        ORDER BY bucket
+      `;
+      const rows = raw.prepare(sql).all() as unknown[];
+      return rows.map((row) => HourlyUsageSchema.parse(row));
+    } catch (error) {
+      logger.error('Failed to get hourly usage', error);
+      return [];
+    } finally {
+      raw.close();
+    }
+  }
+
+  static getUsageByDay(
+    accountId?: string,
+    start?: number,
+    end?: number,
+  ): Array<{
+    bucket: string;
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+    requests: number;
+  }> {
+    const { raw } = getCloudDb();
+    try {
+      const sql = `
+        SELECT
+          strftime('%Y-%m-%d', timestamp / 1000, 'unixepoch', 'localtime') as bucket,
+          SUM(prompt_tokens) as promptTokens,
+          SUM(completion_tokens) as completionTokens,
+          SUM(total_tokens) as totalTokens,
+          COUNT(*) as requests
+        FROM token_usage
+        WHERE 1=1
+        ${buildAccountFilter(accountId)}
+        ${buildDateFilter(start, end)}
+        GROUP BY bucket
+        ORDER BY bucket
+      `;
+      const rows = raw.prepare(sql).all() as unknown[];
+      return rows.map((row) => DailyUsageSchema.parse(row));
+    } catch (error) {
+      logger.error('Failed to get daily usage', error);
+      return [];
+    } finally {
+      raw.close();
+    }
+  }
+
+  static getUsageByWeek(
+    accountId?: string,
+    start?: number,
+    end?: number,
+  ): Array<{
+    bucket: string;
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+    requests: number;
+  }> {
+    const { raw } = getCloudDb();
+    try {
+      const sql = `
+        SELECT
+          strftime('%Y-%W', timestamp / 1000, 'unixepoch', 'localtime') as bucket,
+          SUM(prompt_tokens) as promptTokens,
+          SUM(completion_tokens) as completionTokens,
+          SUM(total_tokens) as totalTokens,
+          COUNT(*) as requests
+        FROM token_usage
+        WHERE 1=1
+        ${buildAccountFilter(accountId)}
+        ${buildDateFilter(start, end)}
+        GROUP BY bucket
+        ORDER BY bucket
+      `;
+      const rows = raw.prepare(sql).all() as unknown[];
+      return rows.map((row) => WeeklyUsageSchema.parse(row));
+    } catch (error) {
+      logger.error('Failed to get weekly usage', error);
+      return [];
+    } finally {
+      raw.close();
+    }
+  }
+
+  static getUsageByMonth(
+    accountId?: string,
+    start?: number,
+    end?: number,
+  ): Array<{
+    bucket: string;
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+    requests: number;
+  }> {
+    const { raw } = getCloudDb();
+    try {
+      const sql = `
+        SELECT
+          strftime('%Y-%m', timestamp / 1000, 'unixepoch', 'localtime') as bucket,
+          SUM(prompt_tokens) as promptTokens,
+          SUM(completion_tokens) as completionTokens,
+          SUM(total_tokens) as totalTokens,
+          COUNT(*) as requests
+        FROM token_usage
+        WHERE 1=1
+        ${buildAccountFilter(accountId)}
+        ${buildDateFilter(start, end)}
+        GROUP BY bucket
+        ORDER BY bucket
+      `;
+      const rows = raw.prepare(sql).all() as unknown[];
+      return rows.map((row) => MonthlyUsageSchema.parse(row));
+    } catch (error) {
+      logger.error('Failed to get monthly usage', error);
+      return [];
+    } finally {
+      raw.close();
+    }
+  }
+
+  static getUsageByModel(
+    accountId?: string,
+    start?: number,
+    end?: number,
+  ): Array<{
+    model: string;
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+    requests: number;
+  }> {
+    const { raw } = getCloudDb();
+    try {
+      const sql = `
+        SELECT
+          model,
+          SUM(prompt_tokens) as promptTokens,
+          SUM(completion_tokens) as completionTokens,
+          SUM(total_tokens) as totalTokens,
+          COUNT(*) as requests
+        FROM token_usage
+        WHERE 1=1
+        ${buildAccountFilter(accountId)}
+        ${buildDateFilter(start, end)}
+        GROUP BY model
+        ORDER BY totalTokens DESC
+      `;
+      const rows = raw.prepare(sql).all() as unknown[];
+      return rows.map((row) => ModelUsageSchema.parse(row));
+    } catch (error) {
+      logger.error('Failed to get model usage', error);
+      return [];
+    } finally {
+      raw.close();
+    }
+  }
+}

src/types/db.ts

@@ -70,3 +70,66 @@ export const TableInfoRowSchema = z.object({
 });
 
 export type TableInfoRow = z.infer<typeof TableInfoRowSchema>;
+
+export const TokenUsageRowSchema = z.object({
+  id: z.number(),
+  account_id: z.string(),
+  model: z.string(),
+  prompt_tokens: z.number(),
+  completion_tokens: z.number(),
+  total_tokens: z.number(),
+  timestamp: z.number(),
+  request_type: z.string().nullable(),
+});
+
+export type TokenUsageRow = z.infer<typeof TokenUsageRowSchema>;
+
+export const HourlyUsageSchema = z.object({
+  bucket: z.string(),
+  promptTokens: z.number(),
+  completionTokens: z.number(),
+  totalTokens: z.number(),
+  requests: z.number(),
+});
+
+export type HourlyUsage = z.infer<typeof HourlyUsageSchema>;
+
+export const DailyUsageSchema = z.object({
+  bucket: z.string(),
+  promptTokens: z.number(),
+  completionTokens: z.number(),
+  totalTokens: z.number(),
+  requests: z.number(),
+});
+
+export type DailyUsage = z.infer<typeof DailyUsageSchema>;
+
+export const WeeklyUsageSchema = z.object({
+  bucket: z.string(),
+  promptTokens: z.number(),
+  completionTokens: z.number(),
+  totalTokens: z.number(),
+  requests: z.number(),
+});
+
+export type WeeklyUsage = z.infer<typeof WeeklyUsageSchema>;
+
+export const MonthlyUsageSchema = z.object({
+  bucket: z.string(),
+  promptTokens: z.number(),
+  completionTokens: z.number(),
+  totalTokens: z.number(),
+  requests: z.number(),
+});
+
+export type MonthlyUsage = z.infer<typeof MonthlyUsageSchema>;
+
+export const ModelUsageSchema = z.object({
+  model: z.string(),
+  promptTokens: z.number(),
+  completionTokens: z.number(),
+  totalTokens: z.number(),
+  requests: z.number(),
+});
+
+export type ModelUsage = z.infer<typeof ModelUsageSchema>;