fix: fully close MCP audit SQLite connections on shutdown

docker-compose.yml

@@ -20,6 +20,7 @@
 
 x-common: &common
   restart: unless-stopped
+  stop_grace_period: 40s
   env_file: .env
   environment:
     - DATA_DIR=/app/data # Must match the volume mount below

open-sse/mcp-server/__tests__/audit.test.ts

@@ -0,0 +1,89 @@
+import fs from "node:fs";
+import os from "node:os";
+import path from "node:path";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+type MockAuditDb = {
+  prepare: ReturnType<typeof vi.fn>;
+  pragma: ReturnType<typeof vi.fn>;
+  close: ReturnType<typeof vi.fn>;
+  open?: boolean;
+};
+
+function createStatementMock() {
+  return {
+    get: vi.fn(),
+    all: vi.fn(),
+    run: vi.fn(),
+  };
+}
+
+describe("MCP audit shutdown", () => {
+  let dataDir: string;
+  let dbFile: string;
+
+  beforeEach(() => {
+    vi.resetModules();
+    globalThis.__omnirouteMcpAuditDb = undefined;
+    dataDir = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-mcp-audit-"));
+    dbFile = path.join(dataDir, "storage.sqlite");
+    fs.writeFileSync(dbFile, "");
+    process.env.DATA_DIR = dataDir;
+  });
+
+  afterEach(() => {
+    delete process.env.DATA_DIR;
+    globalThis.__omnirouteMcpAuditDb = undefined;
+    vi.restoreAllMocks();
+  });
+
+  it("checkpoints and closes the audit database during shutdown", async () => {
+    const mockDb: MockAuditDb = {
+      prepare: vi.fn(() => createStatementMock()),
+      pragma: vi.fn(),
+      close: vi.fn(),
+      open: true,
+    };
+    const MockDatabase = vi.fn(function MockDatabase() {
+      return mockDb;
+    });
+
+    vi.doMock("better-sqlite3", () => ({
+      default: MockDatabase,
+    }));
+
+    const audit = await import("../audit.ts");
+
+    await audit.logToolCall("omniroute_get_health", { ok: true }, { ok: true }, 12, true);
+    expect(mockDb.prepare).toHaveBeenCalledTimes(1);
+
+    expect(audit.closeAuditDb()).toBe(true);
+    expect(mockDb.pragma).toHaveBeenCalledWith("wal_checkpoint(TRUNCATE)");
+    expect(mockDb.close).toHaveBeenCalledTimes(1);
+    expect(audit.closeAuditDb()).toBe(false);
+  });
+
+  it("still closes the audit database when checkpoint fails", async () => {
+    const mockDb: MockAuditDb = {
+      prepare: vi.fn(() => createStatementMock()),
+      pragma: vi.fn(() => {
+        throw new Error("database is busy");
+      }),
+      close: vi.fn(),
+      open: true,
+    };
+    const MockDatabase = vi.fn(function MockDatabase() {
+      return mockDb;
+    });
+
+    vi.doMock("better-sqlite3", () => ({
+      default: MockDatabase,
+    }));
+
+    const audit = await import("../audit.ts");
+
+    await audit.logToolCall("omniroute_get_health", {}, {}, 5, true);
+    expect(audit.closeAuditDb()).toBe(true);
+    expect(mockDb.close).toHaveBeenCalledTimes(1);
+  });
+});

open-sse/mcp-server/audit.ts

@@ -18,6 +18,13 @@ interface StatementLike<TRow = unknown> {
 
 interface AuditDatabase {
   prepare: <TRow = unknown>(sql: string) => StatementLike<TRow>;
+  pragma: (sql: string) => unknown;
+  close: () => void;
+  open?: boolean;
+}
+
+declare global {
+  var __omnirouteMcpAuditDb: AuditDatabase | null | undefined;
 }
 
 interface AuditStatsRow {
@@ -121,7 +128,13 @@ function buildAuditFilterSql(filters: McpAuditQuery): { whereSql: string; params
   };
 }
 
-let db: AuditDatabase | null = null;
+function getCachedAuditDb(): AuditDatabase | null {
+  return globalThis.__omnirouteMcpAuditDb ?? null;
+}
+
+function setCachedAuditDb(database: AuditDatabase | null): void {
+  globalThis.__omnirouteMcpAuditDb = database;
+}
 
 function toNumber(value: unknown, fallback = 0): number {
   const parsed =
@@ -142,7 +155,8 @@ function toString(value: unknown): string {
  * Uses the same SQLite database as the main OmniRoute app.
  */
 async function getDb(): Promise<AuditDatabase | null> {
-  if (db) return db;
+  const cachedDb = getCachedAuditDb();
+  if (cachedDb) return cachedDb;
 
   try {
     // Try importing the db module from the main app
@@ -162,15 +176,43 @@ async function getDb(): Promise<AuditDatabase | null> {
     const Database = (await import("better-sqlite3")).default as unknown as new (
       dbPath: string
     ) => AuditDatabase;
-    db = new Database(dbPath);
-    return db;
+    const database = new Database(dbPath);
+    setCachedAuditDb(database);
+    return database;
   } catch (err: unknown) {
     const message = err instanceof Error ? err.message : String(err);
     console.error("[MCP Audit] Failed to connect to database:", message);
     return null;
   }
 }
 
+export function closeAuditDb(): boolean {
+  const database = getCachedAuditDb();
+  if (!database) return false;
+
+  setCachedAuditDb(null);
+
+  try {
+    try {
+      if (database.open !== false) {
+        database.pragma("wal_checkpoint(TRUNCATE)");
+      }
+    } catch (err: unknown) {
+      const message = err instanceof Error ? err.message : String(err);
+      console.warn("[MCP Audit] WAL checkpoint failed during close:", message);
+    }
+  } finally {
+    try {
+      database.close();
+    } catch (err: unknown) {
+      const message = err instanceof Error ? err.message : String(err);
+      console.warn("[MCP Audit] Failed to close database:", message);
+    }
+  }
+
+  return true;
+}
+
 // ============ Audit Logger ============
 
 /**

open-sse/mcp-server/server.ts

@@ -43,7 +43,7 @@ import {
 } from "./schemas/tools.ts";
 import { startMcpHeartbeat } from "./runtimeHeartbeat.ts";
 
-import { logToolCall } from "./audit.ts";
+import { closeAuditDb, logToolCall } from "./audit.ts";
 import {
   evaluateToolScopes,
   resolveCallerScopeContext,
@@ -876,6 +876,9 @@ export async function startMcpStdio(): Promise<void> {
     await server.connect(transport);
     console.error("[MCP] OmniRoute MCP Server connected and ready.");
   } finally {
+    if (closeAuditDb()) {
+      console.error("[MCP] Audit database checkpointed and closed.");
+    }
     stopHeartbeatOnce();
     process.off("exit", stopHeartbeatOnce);
     process.off("SIGINT", stopHeartbeatOnce);

src/lib/gracefulShutdown.ts

@@ -96,7 +96,13 @@ async function waitForDrain(): Promise<void> {
  */
 async function cleanup(): Promise<void> {
   try {
-    const { closeDbInstance } = await import("@/lib/db/core");
+    const [{ closeAuditDb }, { closeDbInstance }] = await Promise.all([
+      import("@omniroute/open-sse/mcp-server/audit.ts"),
+      import("@/lib/db/core"),
+    ]);
+    if (closeAuditDb()) {
+      console.log("[Shutdown] MCP audit database checkpointed and closed.");
+    }
     if (closeDbInstance()) {
       console.log("[Shutdown] SQLite database checkpointed and closed.");
     }