From bc7fb961847475c49d368f88ae05c1af3862496a Mon Sep 17 00:00:00 2001 From: Randi <55005611+rdself@users.noreply.github.com> Date: Thu, 16 Apr 2026 15:07:13 -0400 Subject: [PATCH] fix: close MCP audit SQLite connections on shutdown (#1348) Integrated into release/v3.6.7 --- docker-compose.yml | 1 + open-sse/mcp-server/__tests__/audit.test.ts | 87 +++++++++++++++++++++ open-sse/mcp-server/audit.ts | 30 +++++++ open-sse/mcp-server/server.ts | 5 +- src/lib/gracefulShutdown.ts | 8 +- 5 files changed, 129 insertions(+), 2 deletions(-) create mode 100644 open-sse/mcp-server/__tests__/audit.test.ts diff --git a/docker-compose.yml b/docker-compose.yml index dba80f29..7e24dbce 100644 --- a/docker-compose.yml +++ b/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 diff --git a/open-sse/mcp-server/__tests__/audit.test.ts b/open-sse/mcp-server/__tests__/audit.test.ts new file mode 100644 index 00000000..e1c7a7a7 --- /dev/null +++ b/open-sse/mcp-server/__tests__/audit.test.ts @@ -0,0 +1,87 @@ +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; + pragma: ReturnType; + close: ReturnType; + 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(); + 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; + 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); + }); +}); diff --git a/open-sse/mcp-server/audit.ts b/open-sse/mcp-server/audit.ts index 211c7fa1..6ac75eb0 100644 --- a/open-sse/mcp-server/audit.ts +++ b/open-sse/mcp-server/audit.ts @@ -18,6 +18,9 @@ interface StatementLike { interface AuditDatabase { prepare: (sql: string) => StatementLike; + pragma: (sql: string) => unknown; + close: () => void; + open?: boolean; } interface AuditStatsRow { @@ -171,6 +174,33 @@ async function getDb(): Promise { } } +export function closeAuditDb(): boolean { + if (!db) return false; + + const database = db; + db = 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 ============ /** diff --git a/open-sse/mcp-server/server.ts b/open-sse/mcp-server/server.ts index d716d2c3..b158bc42 100644 --- a/open-sse/mcp-server/server.ts +++ b/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 { 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); diff --git a/src/lib/gracefulShutdown.ts b/src/lib/gracefulShutdown.ts index 8bc09f5d..c84c069d 100644 --- a/src/lib/gracefulShutdown.ts +++ b/src/lib/gracefulShutdown.ts @@ -96,7 +96,13 @@ async function waitForDrain(): Promise { */ async function cleanup(): Promise { 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."); }