Add local token usage dashboard with SQLite storage

src/lib/paths.ts

@@ -5,10 +5,12 @@ import path from "node:path"
 const APP_DIR = path.join(os.homedir(), ".local", "share", "copilot-api")
 
 const GITHUB_TOKEN_PATH = path.join(APP_DIR, "github_token")
+const TOKEN_USAGE_DB_PATH = path.join(APP_DIR, "token-usage.db")
 
 export const PATHS = {
   APP_DIR,
   GITHUB_TOKEN_PATH,
+  TOKEN_USAGE_DB_PATH,
 }
 
 export async function ensurePaths(): Promise<void> {

src/lib/token-store.ts

@@ -0,0 +1,130 @@
+import { Database } from "bun:sqlite"
+
+import { PATHS } from "./paths"
+
+interface TokenUsageRow {
+  timestamp_min: number
+  model: string
+  input_tokens: number
+  output_tokens: number
+  request_count: number
+}
+
+let db: Database | null = null
+
+const getDb = (): Database => {
+  if (!db)
+    throw new Error("Token store not initialized. Call initTokenStore() first.")
+  return db
+}
+
+export const initTokenStore = (): void => {
+  db = new Database(PATHS.TOKEN_USAGE_DB_PATH)
+
+  db.run(`
+    CREATE TABLE IF NOT EXISTS token_usage (
+      timestamp_min INTEGER NOT NULL,
+      model TEXT NOT NULL,
+      input_tokens INTEGER NOT NULL DEFAULT 0,
+      output_tokens INTEGER NOT NULL DEFAULT 0,
+      request_count INTEGER NOT NULL DEFAULT 0,
+      PRIMARY KEY (timestamp_min, model)
+    )
+  `)
+
+  db.run(`
+    CREATE INDEX IF NOT EXISTS idx_token_usage_timestamp
+    ON token_usage (timestamp_min)
+  `)
+}
+
+const currentMinuteBucket = (): number =>
+  Math.floor(Date.now() / 1000 / 60) * 60
+
+const pruneOldData = (): void => {
+  const thirtyDaysAgo = Math.floor(Date.now() / 1000) - 30 * 24 * 60 * 60
+  getDb().run("DELETE FROM token_usage WHERE timestamp_min < ?", [
+    thirtyDaysAgo,
+  ])
+}
+
+export const recordTokenUsage = (
+  model: string,
+  inputTokens: number,
+  outputTokens: number,
+): void => {
+  try {
+    const bucket = currentMinuteBucket()
+
+    getDb().run(
+      `INSERT INTO token_usage (timestamp_min, model, input_tokens, output_tokens, request_count)
+       VALUES (?, ?, ?, ?, 1)
+       ON CONFLICT (timestamp_min, model) DO UPDATE SET
+         input_tokens = input_tokens + excluded.input_tokens,
+         output_tokens = output_tokens + excluded.output_tokens,
+         request_count = request_count + 1`,
+      [bucket, model, inputTokens, outputTokens],
+    )
+
+    pruneOldData()
+  } catch (error) {
+    // Never let storage errors surface to callers
+    console.error("[token-store] Failed to record token usage:", error)
+  }
+}
+
+export interface TokenUsageSummary {
+  total_input: number
+  total_output: number
+  total_requests: number
+  models: Array<string>
+}
+
+export interface TokenUsageResponse {
+  range: string
+  data: Array<TokenUsageRow>
+  summary: TokenUsageSummary
+}
+
+const RANGE_LABELS: Record<string, string> = {
+  "3600": "1h",
+  "21600": "6h",
+  "86400": "24h",
+  "604800": "7d",
+  "2592000": "30d",
+}
+
+export const getTokenUsageData = (rangeSeconds: number): TokenUsageResponse => {
+  const since = Math.floor(Date.now() / 1000) - rangeSeconds
+
+  const rows = getDb()
+    .query<TokenUsageRow, [number]>(
+      `SELECT timestamp_min, model, input_tokens, output_tokens, request_count
+       FROM token_usage
+       WHERE timestamp_min >= ?
+       ORDER BY timestamp_min ASC`,
+    )
+    .all(since)
+
+  const summary: TokenUsageSummary = {
+    total_input: 0,
+    total_output: 0,
+    total_requests: 0,
+    models: [],
+  }
+
+  const modelSet = new Set<string>()
+  for (const row of rows) {
+    summary.total_input += row.input_tokens
+    summary.total_output += row.output_tokens
+    summary.total_requests += row.request_count
+    modelSet.add(row.model)
+  }
+  summary.models = [...modelSet].sort()
+
+  return {
+    range: RANGE_LABELS[String(rangeSeconds)] ?? `${rangeSeconds}s`,
+    data: rows,
+    summary,
+  }
+}

src/routes/chat-completions/handler.ts

@@ -6,39 +6,59 @@ import { streamSSE, type SSEMessage } from "hono/streaming"
 import { awaitApproval } from "~/lib/approval"
 import { checkRateLimit } from "~/lib/rate-limit"
 import { state } from "~/lib/state"
+import { recordTokenUsage } from "~/lib/token-store"
 import { getTokenCount } from "~/lib/tokenizer"
 import { isNullish } from "~/lib/utils"
 import {
   createChatCompletions,
+  type ChatCompletionChunk,
   type ChatCompletionResponse,
   type ChatCompletionsPayload,
 } from "~/services/copilot/create-chat-completions"
 
+type TokenEstimate = { input: number; output: number } | null
+
+const estimateTokens = async (
+  payload: ChatCompletionsPayload,
+): Promise<TokenEstimate> => {
+  const model = state.models?.data.find((m) => m.id === payload.model)
+  if (!model) {
+    consola.warn("No model selected, skipping token count calculation")
+    return null
+  }
+  try {
+    const count = await getTokenCount(payload, model)
+    consola.info("Current token count:", count)
+    return count
+  } catch (error) {
+    consola.warn("Failed to calculate token count:", error)
+    return null
+  }
+}
+
+const parseChunkUsage = (chunk: SSEMessage): ChatCompletionChunk["usage"] => {
+  if (!chunk.data || chunk.data === "[DONE]") return undefined
+  if (typeof chunk.data !== "string") return undefined
+  try {
+    const parsed = JSON.parse(chunk.data) as ChatCompletionChunk
+    return parsed.usage
+  } catch {
+    return undefined
+  }
+}
+
 export async function handleCompletion(c: Context) {
   await checkRateLimit(state)
 
   let payload = await c.req.json<ChatCompletionsPayload>()
   consola.debug("Request payload:", JSON.stringify(payload).slice(-400))
 
-  // Find the selected model
-  const selectedModel = state.models?.data.find(
-    (model) => model.id === payload.model,
-  )
-
-  // Calculate and display token count
-  try {
-    if (selectedModel) {
-      const tokenCount = await getTokenCount(payload, selectedModel)
-      consola.info("Current token count:", tokenCount)
-    } else {
-      consola.warn("No model selected, skipping token count calculation")
-    }
-  } catch (error) {
-    consola.warn("Failed to calculate token count:", error)
-  }
+  const estimated = await estimateTokens(payload)
 
   if (state.manualApprove) await awaitApproval()
 
+  const selectedModel = state.models?.data.find((m) => m.id === payload.model)
+
   if (isNullish(payload.max_tokens)) {
     payload = {
       ...payload,
@@ -47,19 +67,36 @@ export async function handleCompletion(c: Context) {
     consola.debug("Set max_tokens to:", JSON.stringify(payload.max_tokens))
   }
 
+  if (payload.stream) {
+    payload = { ...payload, stream_options: { include_usage: true } }
+  }
+
   const response = await createChatCompletions(payload)
 
   if (isNonStreaming(response)) {
     consola.debug("Non-streaming response:", JSON.stringify(response))
+    const inputTokens = response.usage?.prompt_tokens ?? estimated?.input ?? 0
+    const outputTokens =
+      response.usage?.completion_tokens ?? estimated?.output ?? 0
+    recordTokenUsage(payload.model, inputTokens, outputTokens)
     return c.json(response)
   }
 
   consola.debug("Streaming response")
   return streamSSE(c, async (stream) => {
+    let lastUsage: ChatCompletionChunk["usage"] | undefined
+
     for await (const chunk of response) {
       consola.debug("Streaming chunk:", JSON.stringify(chunk))
       await stream.writeSSE(chunk as SSEMessage)
+      lastUsage = parseChunkUsage(chunk as SSEMessage) ?? lastUsage
     }
+
+    recordTokenUsage(
+      payload.model,
+      lastUsage?.prompt_tokens ?? estimated?.input ?? 0,
+      lastUsage?.completion_tokens ?? estimated?.output ?? 0,
+    )
   })
 }
 

src/routes/dashboard/budget.ts

@@ -0,0 +1,37 @@
+export interface QuotaDetail {
+  entitlement: number
+  remaining: number
+  percent_remaining: number
+  unlimited: boolean
+  over_limit?: boolean
+}
+
+/** Returns CSS hex color based on % used (100 - percent_remaining). */
+export function getBudgetColor(percentUsed: number): string {
+  if (percentUsed >= 95) return "#ef4444"
+  if (percentUsed >= 80) return "#f59e0b"
+  return "#22c55e"
+}
+
+/** Computes percentage used (0–100) from a QuotaDetail. Returns 0 for unlimited. */
+export function getPercentUsed(quota: QuotaDetail): number {
+  if (quota.unlimited) return 0
+  if (quota.entitlement === 0) return 0
+  return Math.min(
+    100,
+    ((quota.entitlement - quota.remaining) / quota.entitlement) * 100,
+  )
+}
+
+/** Returns a human-readable label for a quota key. */
+export function getQuotaLabel(key: string): string {
+  const labels: Record<string, string> = {
+    premium_interactions: "Premium Interactions",
+    chat: "Chat",
+    completions: "Completions",
+  }
+  return (
+    labels[key]
+    ?? key.replaceAll("_", " ").replaceAll(/\b\w/g, (c) => c.toUpperCase())
+  )
+}