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