From 2b4fa0c690fb8c3ed3e12b85a7866eabd7e7ad98 Mon Sep 17 00:00:00 2001
From: OpenClaw <openclawd@cloonar.com>
Date: Sat, 7 Mar 2026 08:03:56 +0100
Subject: [PATCH] fix: await flushDirtyEntries during shutdown to prevent usage
 data loss

Remove fire-and-forget SIGTERM/SIGINT handlers from usage.ts (race condition
with pool.end() in index.ts). Instead, await flushDirtyEntries() in the
index.ts shutdown orchestrator between stopping the server and closing the
DB pool.
---
 src/__tests__/usage-shutdown.test.ts | 30 ++++++++++++++++++++++++++++
 src/index.ts                         | 10 +++++++++-
 src/middleware/usage.ts              |  5 ++---
 3 files changed, 41 insertions(+), 4 deletions(-)
 create mode 100644 src/__tests__/usage-shutdown.test.ts

diff --git a/src/__tests__/usage-shutdown.test.ts b/src/__tests__/usage-shutdown.test.ts
new file mode 100644
index 0000000..7d5032e
--- /dev/null
+++ b/src/__tests__/usage-shutdown.test.ts
@@ -0,0 +1,30 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+
+vi.unmock("../middleware/usage.js");
+
+describe("usage shutdown race condition fix", () => {
+  it("should NOT register SIGTERM/SIGINT handlers as module-level side effects", async () => {
+    const onSpy = vi.spyOn(process, "on");
+    vi.resetModules();
+    
+    // Track which signals usage.ts registers
+    const signalsBefore = onSpy.mock.calls.map(c => c[0]);
+    
+    await import("../middleware/usage.js");
+    
+    const signalsAfter = onSpy.mock.calls.map(c => c[0]);
+    const newSignals = signalsAfter.slice(signalsBefore.length);
+    
+    // usage.ts should NOT register SIGTERM or SIGINT handlers
+    const usageSignals = newSignals.filter(s => s === "SIGTERM" || s === "SIGINT");
+    expect(usageSignals).toEqual([]);
+    
+    onSpy.mockRestore();
+  });
+
+  it("should export flushDirtyEntries for external shutdown orchestration", async () => {
+    vi.resetModules();
+    const usageMod = await import("../middleware/usage.js");
+    expect(typeof usageMod.flushDirtyEntries).toBe("function");
+  });
+});
diff --git a/src/index.ts b/src/index.ts
index ca6152c..6bbb998 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -18,7 +18,7 @@ import { recoverRouter } from "./routes/recover.js";
 import { emailChangeRouter } from "./routes/email-change.js";
 import { billingRouter } from "./routes/billing.js";
 import { authMiddleware } from "./middleware/auth.js";
-import { usageMiddleware, loadUsageData } from "./middleware/usage.js";
+import { usageMiddleware, loadUsageData, flushDirtyEntries } from "./middleware/usage.js";
 import { getUsageStats } from "./middleware/usage.js";
 import { pdfRateLimitMiddleware, getConcurrencyStats } from "./middleware/pdfRateLimit.js";
 import { initBrowser, closeBrowser } from "./services/browser.js";
@@ -406,6 +406,14 @@ async function start() {
       });
     });
 
+    // 1.5. Flush dirty usage entries while DB pool is still alive
+    try {
+      await flushDirtyEntries();
+      logger.info("Usage data flushed");
+    } catch (err) {
+      logger.error({ err }, "Error flushing usage data during shutdown");
+    }
+
     // 2. Close Puppeteer browser pool
     try {
       await closeBrowser();
diff --git a/src/middleware/usage.ts b/src/middleware/usage.ts
index 3f9a428..040fc7a 100644
--- a/src/middleware/usage.ts
+++ b/src/middleware/usage.ts
@@ -73,9 +73,8 @@ export async function flushDirtyEntries(): Promise<void> {
 // Periodic flush
 setInterval(flushDirtyEntries, FLUSH_INTERVAL_MS);
 
-// Flush on process exit
-process.on("SIGTERM", () => { flushDirtyEntries().catch(() => {}); });
-process.on("SIGINT", () => { flushDirtyEntries().catch(() => {}); });
+// Note: SIGTERM/SIGINT flush is handled by the shutdown orchestrator in index.ts
+// to avoid race conditions with pool.end().
 
 export function usageMiddleware(req: any, res: any, next: any): void {
   const keyInfo = req.apiKeyInfo;