fix: flush usage entries independently to prevent batch poisoning (BUG-100)

src/__tests__/usage-flush.test.ts

@@ -0,0 +1,57 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+
+// Unmock usage middleware — we want to test the real implementation
+vi.unmock("../middleware/usage.js");
+
+import { connectWithRetry } from "../services/db.js";
+
+describe("flushDirtyEntries – independent key flushing", () => {
+  let usageMod: typeof import("../middleware/usage.js");
+
+  beforeEach(async () => {
+    vi.clearAllMocks();
+    vi.resetModules();
+    // Re-import to get fresh state
+    usageMod = await import("../middleware/usage.js");
+  });
+
+  it("should flush remaining keys even if one key fails", async () => {
+    // Set up two keys in the in-memory cache via the middleware
+    const next = vi.fn();
+    const res = { status: vi.fn(() => ({ json: vi.fn() })) };
+
+    usageMod.usageMiddleware({ apiKeyInfo: { key: "key-good" } }, res, next);
+    usageMod.usageMiddleware({ apiKeyInfo: { key: "key-bad" } }, res, next);
+
+    // Track which keys were successfully upserted
+    const flushedKeys: string[] = [];
+    let callCount = 0;
+
+    const mockQuery = vi.fn().mockImplementation((sql: string, params?: any[]) => {
+      if (sql.includes("INSERT INTO usage")) {
+        callCount++;
+        if (params && params[0] === "key-bad") {
+          throw new Error("simulated constraint violation");
+        }
+        flushedKeys.push(params![0]);
+      }
+      return { rows: [], rowCount: 0 };
+    });
+
+    const mockRelease = vi.fn();
+
+    // Each call to connectWithRetry returns a fresh client
+    vi.mocked(connectWithRetry).mockImplementation(async () => ({
+      query: mockQuery,
+      release: mockRelease,
+    }) as any);
+
+    // Access the flush function (exported for testing)
+    await usageMod.flushDirtyEntries();
+
+    // The good key should have been flushed despite the bad key failing
+    expect(flushedKeys).toContain("key-good");
+    // Release should be called for each key (independent clients)
+    expect(mockRelease.mock.calls.length).toBeGreaterThanOrEqual(2);
+  });
+});

src/middleware/usage.ts

@@ -36,45 +36,37 @@ export async function loadUsageData(): Promise<void> {
 }
 
 // Batch flush dirty entries to DB (Audit #10 + #12)
-async function flushDirtyEntries(): Promise<void> {
+export async function flushDirtyEntries(): Promise<void> {
   if (dirtyKeys.size === 0) return;
 
   const keysToFlush = [...dirtyKeys];
-  
+
-  const client = await connectWithRetry();
+  for (const key of keysToFlush) {
-  try {
+    const record = usage.get(key);
-    await client.query("BEGIN");
+    if (!record) continue;
-    for (const key of keysToFlush) {
+    const client = await connectWithRetry();
-      const record = usage.get(key);
+    try {
-      if (!record) continue;
+      await client.query(
-      try {
+        `INSERT INTO usage (key, count, month_key) VALUES ($1, $2, $3)
-        await client.query(
+         ON CONFLICT (key) DO UPDATE SET count = $2, month_key = $3`,
-          `INSERT INTO usage (key, count, month_key) VALUES ($1, $2, $3)
+        [key, record.count, record.monthKey]
-           ON CONFLICT (key) DO UPDATE SET count = $2, month_key = $3`,
+      );
-          [key, record.count, record.monthKey]
+      dirtyKeys.delete(key);
-        );
+      retryCount.delete(key);
+    } catch (error) {
+      // Audit #12: retry logic for failed writes
+      const retries = (retryCount.get(key) || 0) + 1;
+      if (retries >= MAX_RETRIES) {
+        logger.error({ key: key.slice(0, 8) + "...", retries }, "CRITICAL: Usage write failed after max retries, data may diverge");
         dirtyKeys.delete(key);
         retryCount.delete(key);
-      } catch (error) {
+      } else {
-        // Audit #12: retry logic for failed writes
+        retryCount.set(key, retries);
-        const retries = (retryCount.get(key) || 0) + 1;
+        logger.warn({ key: key.slice(0, 8) + "...", retries }, "Usage write failed, will retry");
-        if (retries >= MAX_RETRIES) {
-          logger.error({ key: key.slice(0, 8) + "...", retries }, "CRITICAL: Usage write failed after max retries, data may diverge");
-          dirtyKeys.delete(key);
-          retryCount.delete(key);
-        } else {
-          retryCount.set(key, retries);
-          logger.warn({ key: key.slice(0, 8) + "...", retries }, "Usage write failed, will retry");
-        }
       }
+    } finally {
+      client.release();
     }
-    await client.query("COMMIT");
-  } catch (error) {
-    await client.query("ROLLBACK").catch(() => {});
-    logger.error({ err: error }, "Failed to flush usage batch");
-    // Keep all keys dirty for retry
-  } finally {
-    client.release();
   }
 }