From a406381fad1f4fb685fb6c3565eaaff01276b4e9 Mon Sep 17 00:00:00 2001
From: Deepanshu <deepanshukumar9650@gmail.com>
Date: Wed, 23 Sep 2026 14:40:52 +0700
Subject: [PATCH] fix(usage): preserve API key usage attribution in live stats

The 24h/today branch of getUsageStats keyed byApiKey buckets on the
masked key, while the daily rollup and the lastUsed overlay key on the
full key. Every key minted by one instance shares the sk-{machineId}
prefix, so the mask collapsed all of them into a single bucket and
attributed one key's usage to another.

Key the live branch by the full api key too; the masked value is still
carried on the bucket for display.
---
 src/lib/db/repos/usageRepo.js                |  2 +-
 tests/unit/usage-api-key-attribution.test.js | 57 ++++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit/usage-api-key-attribution.test.js

diff --git a/src/lib/db/repos/usageRepo.js b/src/lib/db/repos/usageRepo.js
index 2f86ff14..41199937 100644
--- a/src/lib/db/repos/usageRepo.js
+++ b/src/lib/db/repos/usageRepo.js
@@ -634,7 +634,7 @@ export async function getUsageStats(period = "all") {
         const keyInfo = apiKeyMap[r.apiKey];
         const keyName = keyInfo?.name || r.apiKey.slice(0, 8) + "...";
         const apiKeyMasked = maskApiKey(r.apiKey);
-        const akKey = `${apiKeyMasked}|${r.model}|${r.provider || "unknown"}`;
+        const akKey = `${r.apiKey}|${r.model}|${r.provider || "unknown"}`;
         if (!stats.byApiKey[akKey]) {
           stats.byApiKey[akKey] = { requests: 0, promptTokens: 0, completionTokens: 0, cachedTokens: 0, cost: 0, rawModel: r.model, provider: providerDisplayName, apiKeyMasked, keyName, apiKeyKey: apiKeyMasked, lastUsed: r.timestamp };
         }
diff --git a/tests/unit/usage-api-key-attribution.test.js b/tests/unit/usage-api-key-attribution.test.js
new file mode 100644
index 00000000..049c2898
--- /dev/null
+++ b/tests/unit/usage-api-key-attribution.test.js
@@ -0,0 +1,57 @@
+import fs from "node:fs";
+import os from "node:os";
+import path from "node:path";
+import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
+
+let tempDir;
+let db;
+
+beforeEach(async () => {
+  tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "9router-api-key-"));
+  process.env.DATA_DIR = tempDir;
+  vi.resetModules();
+  db = await import("@/lib/db/index.js");
+  await db.initDb();
+});
+
+afterEach(() => {
+  delete process.env.DATA_DIR;
+});
+
+describe("Usage stats API key attribution", () => {
+  it("keeps API keys with the same masked prefix in separate buckets", async () => {
+    const apiKeyA = "sk-machine-aaaaaa-11111111";
+    const apiKeyB = "sk-machine-bbbbbb-22222222";
+
+    await db.saveRequestUsage({
+      provider: "openai",
+      model: "gpt-4",
+      connectionId: "c1",
+      apiKey: apiKeyA,
+      tokens: { prompt_tokens: 10, completion_tokens: 5 },
+      endpoint: "/v1/chat",
+      status: "ok",
+    });
+
+    await db.saveRequestUsage({
+      provider: "openai",
+      model: "gpt-4",
+      connectionId: "c1",
+      apiKey: apiKeyB,
+      tokens: { prompt_tokens: 20, completion_tokens: 10 },
+      endpoint: "/v1/chat",
+      status: "ok",
+    });
+
+    const stats = await db.getUsageStats("24h");
+    const apiKeyEntries = Object.values(stats.byApiKey);
+
+    expect(apiKeyEntries).toHaveLength(2);
+
+    expect(
+      apiKeyEntries
+        .map((entry) => entry.promptTokens)
+        .sort((a, b) => a - b)
+    ).toEqual([10, 20]);
+  });
+});