@nifrajs/mcp-db 3.3.0 → 4.0.0

package/dist/index.js

@@ -23,6 +23,8 @@
  *   app.mount("/mcp", mcp.fetch)
  */
 import { createMcpServer, defineMcpTool } from "@nifrajs/mcp";
+import { boundedSqliteQuery, countSqliteQuery, fitToBytes, gateSqliteStatement, unexposedPlanRelation, } from "./internal/shared.js";
+import { openQuerySession, QueryTimeoutError, } from "./internal/sqlite-session.js";
 export class McpDbConfigError extends Error {
     constructor(message) {
         super(`@nifrajs/mcp-db: ${message}`);
@@ -42,498 +44,11 @@ function sqlIdentifiers(strings, ...identifiers) {
     }
     return sql;
 }
-const stripSqlNoise = (sql) => sql
-    // Strings first so comment markers inside literals don't count.
-    .replace(/'(?:[^']|'')*'/g, "''")
-    .replace(/"(?:[^"]|"")*"/g, '""')
-    .replace(/--[^\n]*/g, " ")
-    .replace(/\/\*[\s\S]*?\*\//g, " ");
-// SQLite's bare-identifier alphabet. Non-ASCII is included because SQLite treats any byte >= 0x80 as
-// an identifier character; leaving it out would end a word early and split one table name into two.
-const WORD_START = /[A-Za-z_\u0080-\uffff]/;
-const WORD_PART = /[A-Za-z0-9_$\u0080-\uffff]/;
-/**
- * Tokenize far enough to name every relation. A regex cannot do this job: SQLite needs no separator
- * between a keyword and a quoted identifier (`FROM"users"` is legal), and it accepts four identifier
- * quotings - `"x"`, `[x]`, `` `x` ``, and bare. A `\s+`-anchored pattern silently reads none of
- * those as a table reference, which is exactly how an aliased `FROM"users"AS habits` slipped past
- * the allowlist and returned a non-exposed table's rows.
- *
- * String literals and comments become nothing, so a `FROM` inside attacker-controlled text is inert.
- */
-function tokenizeSql(sql) {
-    const tokens = [];
-    let i = 0;
-    /** Read a quoted identifier, collapsing the doubled-delimiter escape SQLite uses for `"` and `` ` ``. */
-    const readQuoted = (close, escapable) => {
-        let value = "";
-        i++; // past the opening delimiter
-        while (i < sql.length) {
-            if (sql[i] === close) {
-                if (escapable && sql[i + 1] === close) {
-                    value += close;
-                    i += 2;
-                    continue;
-                }
-                i++;
-                tokens.push({ kind: "word", value, bare: false });
-                return;
-            }
-            value += sql[i];
-            i++;
-        }
-        // Unterminated - emit what we have. The caller's single-statement and plan checks still run.
-        tokens.push({ kind: "word", value, bare: false });
-    };
-    while (i < sql.length) {
-        const char = sql[i];
-        if (char === " " || char === "\t" || char === "\n" || char === "\r" || char === "\f") {
-            i++;
-        }
-        else if (char === "-" && sql[i + 1] === "-") {
-            const end = sql.indexOf("\n", i);
-            i = end === -1 ? sql.length : end + 1;
-        }
-        else if (char === "/" && sql[i + 1] === "*") {
-            const end = sql.indexOf("*/", i + 2);
-            i = end === -1 ? sql.length : end + 2;
-        }
-        else if (char === "'") {
-            // Literal: consumed and dropped, so its contents can never be read as SQL.
-            i++;
-            while (i < sql.length) {
-                if (sql[i] === "'") {
-                    if (sql[i + 1] === "'") {
-                        i += 2;
-                        continue;
-                    }
-                    i++;
-                    break;
-                }
-                i++;
-            }
-        }
-        else if (char === '"') {
-            readQuoted('"', true);
-        }
-        else if (char === "`") {
-            readQuoted("`", true);
-        }
-        else if (char === "[") {
-            readQuoted("]", false); // MSSQL-style bracket quoting: no escape form
-        }
-        else if (WORD_START.test(char)) {
-            const start = i;
-            while (i < sql.length && WORD_PART.test(sql[i]))
-                i++;
-            tokens.push({ kind: "word", value: sql.slice(start, i), bare: true });
-        }
-        else if (char >= "0" && char <= "9") {
-            while (i < sql.length && /[0-9A-Za-z._]/.test(sql[i]))
-                i++;
-        }
-        else {
-            tokens.push({ kind: "punct", value: char, bare: false });
-            i++;
-        }
-    }
-    return tokens;
-}
-/** True when `token` is the unquoted keyword `word` (case-insensitive). */
-const isKeyword = (token, word) => token?.bare === true && token.value.toLowerCase() === word;
-/**
- * Every relation the statement names after `FROM`/`JOIN`, lowercased, with CTE names excluded.
- * A schema qualifier resolves to its table (`main.habits` -> `habits`); a subquery contributes
- * nothing of its own because its inner `FROM` is tokenized alongside everything else.
- */
-function relationNames(sql) {
-    const tokens = tokenizeSql(sql);
-    // `name AS (` only ever introduces a CTE - a column alias cannot be followed by a paren - so this
-    // needs no `WITH` tracking to be exact.
-    const ctes = new Set();
-    for (let i = 1; i < tokens.length - 1; i++) {
-        const name = tokens[i - 1];
-        if (isKeyword(tokens[i], "as") && tokens[i + 1]?.value === "(" && name.kind === "word") {
-            ctes.add(name.value.toLowerCase());
-        }
-    }
-    const names = [];
-    for (let i = 0; i < tokens.length; i++) {
-        if (!isKeyword(tokens[i], "from") && !isKeyword(tokens[i], "join"))
-            continue;
-        const first = tokens[i + 1];
-        if (first === undefined || first.kind !== "word")
-            continue; // `FROM (subquery)` - nothing to name here
-        // A qualified reference names the table second; an unqualified one names it first.
-        const qualified = tokens[i + 2]?.value === "." && tokens[i + 3]?.kind === "word";
-        const table = qualified ? tokens[i + 3].value : first.value;
-        if (table !== "" && !ctes.has(table.toLowerCase()))
-            names.push(table.toLowerCase());
-    }
-    return names;
-}
-/** Reject multi-statement input: allow one terminator only when it is the final character. */
-const isSingleStatement = (sql) => {
-    let quote = null;
-    let lineComment = false;
-    let blockComment = false;
-    for (let i = 0; i < sql.length; i++) {
-        const char = sql[i];
-        const next = sql[i + 1];
-        if (lineComment) {
-            if (char === "\n")
-                lineComment = false;
-            continue;
-        }
-        if (blockComment) {
-            if (char === "*" && next === "/") {
-                blockComment = false;
-                i++;
-            }
-            continue;
-        }
-        if (quote !== null) {
-            if (char === quote) {
-                if (next === quote)
-                    i++;
-                else
-                    quote = null;
-            }
-            continue;
-        }
-        if (char === "'" || char === '"') {
-            quote = char;
-            continue;
-        }
-        if (char === "-" && next === "-") {
-            lineComment = true;
-            i++;
-            continue;
-        }
-        if (char === "/" && next === "*") {
-            blockComment = true;
-            i++;
-            continue;
-        }
-        if (char === ";") {
-            return sql.slice(i + 1).trim() === "";
-        }
-    }
-    return true;
-};
-const isReadStatement = (sql) => {
-    const head = stripSqlNoise(sql).trim().toLowerCase();
-    return head.startsWith("select") || head.startsWith("with");
-};
-/** Remove one trailing statement terminator (and anything after it) while preserving literals and
- * identifiers. Wrapping the query in a bounded subquery must not let a trailing comment consume the
- * wrapper's closing parenthesis. */
-function executableSql(sql) {
-    let quote = null;
-    let lineComment = false;
-    let lineCommentStart = -1;
-    let blockComment = false;
-    for (let i = 0; i < sql.length; i++) {
-        const char = sql[i];
-        const next = sql[i + 1];
-        if (lineComment) {
-            if (char === "\n") {
-                lineComment = false;
-                lineCommentStart = -1;
-            }
-            continue;
-        }
-        if (blockComment) {
-            if (char === "*" && next === "/") {
-                blockComment = false;
-                i++;
-            }
-            continue;
-        }
-        if (quote !== null) {
-            if (char === quote) {
-                if (next === quote)
-                    i++;
-                else
-                    quote = null;
-            }
-            continue;
-        }
-        if (char === "'" || char === '"') {
-            quote = char;
-            continue;
-        }
-        if (char === "-" && next === "-") {
-            lineComment = true;
-            lineCommentStart = i;
-            i++;
-            continue;
-        }
-        if (char === "/" && next === "*") {
-            blockComment = true;
-            i++;
-            continue;
-        }
-        if (char === ";")
-            return sql.slice(0, i).trim();
-    }
-    if (lineComment && lineCommentStart >= 0)
-        return sql.slice(0, lineCommentStart).trim();
-    return sql.trim();
-}
 function assertResultLimit(value, name) {
     if (!Number.isSafeInteger(value) || value <= 0 || value >= Number.MAX_SAFE_INTEGER) {
         throw new McpDbConfigError(`${name} must be a finite positive safe integer`);
     }
 }
-class QueryTimeoutError extends Error {
-    constructor() {
-        super("query exceeded the time limit");
-        this.name = "QueryTimeoutError";
-    }
-}
-/**
- * The isolated lane. The worker reopens the SAME database FILE read-only and answers one statement
- * at a time; it is spawned once and reused for the life of the server. It deliberately does not
- * take a `serialize()` snapshot - that copies the whole database into memory on EVERY call, which
- * is a larger availability problem than the slow query the deadline exists to bound.
- */
-const SQLITE_WORKER_SOURCE = `
-const { Database } = require("bun:sqlite")
-let db
-self.onmessage = (event) => {
-  const { id, filename, sql, type } = event.data
-  if (type === "close") {
-    let response
-    try {
-      // prepare().all() leaves a statement alive until it is finalized or collected. Bun's
-      // default close(false) preserves those statements and can keep the database file locked;
-      // close(true) finalizes every outstanding statement before the worker exits.
-      db?.close(true)
-      db = undefined
-      response = { id, closed: true, ok: true }
-    } catch (error) {
-      response = { id, closed: true, ok: false, error: String(error) }
-    }
-    // A forced parent-side terminate can leave native SQLite handles locked on Windows even after
-    // close() returns. Let the worker exit its own event turn so the runtime can finish releasing
-    // those handles before the parent considers graceful cleanup complete.
-    self.postMessage(response)
-    self.close()
-    return
-  }
-  try {
-    if (db === undefined) {
-      db = new Database(filename, { readonly: true })
-      db.run("PRAGMA query_only = ON")
-    }
-    self.postMessage({ id, ok: true, rows: db.prepare(sql).all() })
-  } catch (error) {
-    self.postMessage({ id, ok: false, error: String(error) })
-  }
-}
-`;
-function isWorkerAvailable() {
-    return typeof Worker !== "undefined" && typeof Blob !== "undefined" && typeof URL !== "undefined";
-}
-/** Only a file-backed database can be reopened in a worker. `:memory:` and the anonymous temp
- * database live in this process alone, so they take the in-process lane instead. */
-function reopenableFilename(db) {
-    const filename = db.filename;
-    if (typeof filename !== "string" || filename === "" || filename === ":memory:")
-        return undefined;
-    return filename.startsWith("file::memory:") ? undefined : filename;
-}
-const WORKER_CLOSE_TIMEOUT_MS = 1_000;
-function workerSession(filename) {
-    let worker;
-    let sourceUrl;
-    let nextId = 0;
-    const pending = new Map();
-    let lane = Promise.resolve();
-    let closed = false;
-    let closePromise;
-    let pendingClose;
-    const detachWorker = (active = worker) => {
-        if (active === undefined)
-            return;
-        if (worker !== active)
-            return;
-        worker = undefined;
-        if (sourceUrl !== undefined)
-            URL.revokeObjectURL(sourceUrl);
-        sourceUrl = undefined;
-    };
-    const releaseWorker = (active = worker) => {
-        if (active === undefined)
-            return;
-        active.terminate();
-        detachWorker(active);
-    };
-    const rejectPending = (error) => {
-        const waiters = [...pending.values()];
-        pending.clear();
-        for (const waiter of waiters)
-            waiter.reject(error);
-    };
-    /** Drop the worker and fail everything riding on it. The only bound this side can enforce over a
-     * statement already running inside native SQLite is to stop owning the connection it runs on. */
-    const discard = (error) => {
-        const closing = pendingClose;
-        pendingClose = undefined;
-        releaseWorker();
-        rejectPending(error);
-        closing?.resolve();
-    };
-    const ensureWorker = () => {
-        if (worker !== undefined)
-            return worker;
-        sourceUrl = URL.createObjectURL(new Blob([SQLITE_WORKER_SOURCE], { type: "text/javascript" }));
-        const spawned = new Worker(sourceUrl);
-        spawned.onmessage = (event) => {
-            const { id, ok, rows, error } = event.data;
-            if (typeof id !== "number")
-                return;
-            if (event.data.closed === true) {
-                const closing = pendingClose;
-                if (closing?.id !== id)
-                    return;
-                pendingClose = undefined;
-                closing.resolve(event.data.ok === true);
-                return;
-            }
-            const waiter = pending.get(id);
-            if (waiter === undefined)
-                return;
-            if (Date.now() >= waiter.deadline) {
-                discard(new QueryTimeoutError());
-                return;
-            }
-            pending.delete(id);
-            if (ok === true)
-                waiter.resolve(rows ?? []);
-            else
-                waiter.reject(new Error(error ?? "query failed in worker"));
-        };
-        spawned.onerror = () => discard(new Error("query failed in worker"));
-        spawned.unref?.();
-        worker = spawned;
-        return spawned;
-    };
-    const dispatch = (sql, deadline) => {
-        if (closed)
-            return Promise.reject(new Error("query session is closed"));
-        const active = ensureWorker();
-        const remaining = deadline - Date.now();
-        if (remaining <= 0) {
-            discard(new QueryTimeoutError());
-            return Promise.reject(new QueryTimeoutError());
-        }
-        const id = ++nextId;
-        return new Promise((resolve, reject) => {
-            const timer = setTimeout(() => discard(new QueryTimeoutError()), remaining);
-            pending.set(id, {
-                deadline,
-                resolve: (rows) => {
-                    clearTimeout(timer);
-                    resolve(rows);
-                },
-                reject: (error) => {
-                    clearTimeout(timer);
-                    reject(error);
-                },
-            });
-            try {
-                active.postMessage({ id, filename, sql });
-            }
-            catch (error) {
-                pending.delete(id);
-                clearTimeout(timer);
-                reject(error instanceof Error ? error : new Error(String(error)));
-            }
-        });
-    };
-    return {
-        run(sql, deadline) {
-            // One worker owns the connection, so overlapping calls queue behind each other rather than
-            // interleaving statements on it.
-            const result = lane.then(() => dispatch(sql, deadline));
-            lane = result.catch(() => undefined);
-            return result;
-        },
-        close() {
-            if (closePromise !== undefined)
-                return closePromise;
-            closed = true;
-            closePromise = (async () => {
-                const active = worker;
-                if (active !== undefined) {
-                    await new Promise((resolve) => {
-                        const id = ++nextId;
-                        let settled = false;
-                        const finish = (error, graceful = false) => {
-                            if (settled)
-                                return;
-                            settled = true;
-                            clearTimeout(timer);
-                            if (pendingClose?.id === id)
-                                pendingClose = undefined;
-                            if (error !== undefined)
-                                rejectPending(error);
-                            if (graceful)
-                                detachWorker(active);
-                            else
-                                releaseWorker(active);
-                            resolve();
-                        };
-                        const timer = setTimeout(() => {
-                            finish(new Error("query session is closed"));
-                        }, WORKER_CLOSE_TIMEOUT_MS);
-                        pendingClose = { id, resolve: (graceful) => finish(undefined, graceful) };
-                        try {
-                            active.postMessage({ id, type: "close" });
-                        }
-                        catch {
-                            finish(new Error("query session is closed"));
-                        }
-                    });
-                }
-                await lane;
-            })();
-            return closePromise;
-        },
-    };
-}
-/**
- * The in-process lane, used when the database cannot be reopened elsewhere. The deadline is checked
- * around the statement and nothing more: `prepare().all()` is synchronous, so it holds the only
- * thread there is and no timer of ours can fire while it runs. That makes this lane BEST EFFORT -
- * it reports an overrun, it does not cut one short.
- */
-function inProcessSession(db) {
-    return {
-        async run(sql, deadline) {
-            if (Date.now() >= deadline)
-                throw new QueryTimeoutError();
-            const rows = db.prepare(sql).all();
-            if (Date.now() > deadline)
-                throw new QueryTimeoutError();
-            return rows;
-        },
-        async close() { },
-    };
-}
-/**
- * Pick the lane once, at construction. A file-backed database is isolated in a reusable read-only
- * worker, which is the only arrangement that can actually stop a running statement. Anything else
- * runs in process - every database keeps working, none is refused.
- */
-function openQuerySession(db) {
-    const filename = reopenableFilename(db);
-    if (filename !== undefined && isWorkerAvailable())
-        return workerSession(filename);
-    return inProcessSession(db);
-}
 /**
  * Serve `db` as a mountable MCP server (`mcp.fetch` at `POST /mcp`). See module docs for the
  * security model. Throws {@link McpDbConfigError} on any unsafe configuration - always at
@@ -622,31 +137,13 @@ export function serveDatabaseAsMcp(db, options) {
                 }
                 activeQueries += 1;
                 try {
-                    const sql = String(args.sql ?? "").trim();
-                    if (sql === "")
-                        return { isError: true, text: "empty query" };
-                    if (!isSingleStatement(sql)) {
-                        return { isError: true, text: "only a single statement is allowed" };
-                    }
-                    if (!isReadStatement(sql)) {
-                        return { isError: true, text: "only SELECT (or WITH…SELECT) statements are allowed" };
-                    }
+                    const gate = gateSqliteStatement(String(args.sql ?? ""), (relation) => allowlist.has(relation));
+                    if (!gate.ok)
+                        return { isError: true, text: gate.text };
                     const deadline = Date.now() + queryTimeoutMs;
-                    // Verify SQL relation tokens before the plan check. SQLite's plan output uses aliases as
-                    // scan targets (e.g. `users AS habits` becomes `SCAN habits`), so plan-only validation can
-                    // mistake an attacker-controlled alias for an allowlisted table.
-                    const executable = executableSql(sql);
-                    for (const relation of relationNames(executable)) {
-                        if (!allowlist.has(relation)) {
-                            return {
-                                isError: true,
-                                text: `query touches ${JSON.stringify(relation)}, which is not exposed`,
-                            };
-                        }
-                    }
                     // Verify via the query plan that only allowlisted tables are touched. SQLite names every
                     // scanned/searched relation in EXPLAIN QUERY PLAN detail rows.
-                    const query = executable;
+                    const query = gate.query;
                     let planRows;
                     try {
                         planRows = (await querySession.run(`EXPLAIN QUERY PLAN ${query}`, deadline));
@@ -656,34 +153,18 @@ export function serveDatabaseAsMcp(db, options) {
                             return { isError: true, text: error.message };
                         return { isError: true, text: `query failed to plan: ${String(error)}` };
                     }
-                    for (const row of planRows) {
-                        const detail = typeof row.detail === "string" ? row.detail : "";
-                        // SQLite also emits non-table nodes such as `SCAN CONSTANT ROW` for a
-                        // constant-only SELECT. Only treat a scan/search target as a relation
-                        // when it is not one of those planner pseudo-nodes.
-                        if (/^SCAN\s+CONSTANT\s+ROW$/i.test(detail.trim()) ||
-                            /^SCAN\s+SUBQUERY\s+\d+$/i.test(detail.trim())) {
-                            continue;
-                        }
-                        // A schema-qualified reference plans as `SCAN main.habits`. The table is the last dotted
-                        // segment; reading the first would reject `main.habits` as a table literally named "main".
-                        const match = /(?:SCAN|SEARCH)\s+(?:TABLE\s+)?(?:[A-Za-z_][A-Za-z0-9_]*\.)?([A-Za-z_][A-Za-z0-9_]*)/i.exec(detail);
-                        if (match !== null) {
-                            const relation = match[1]?.toLowerCase() ?? "";
-                            if (!allowlist.has(relation)) {
-                                return {
-                                    isError: true,
-                                    text: `query touches ${JSON.stringify(match[1])}, which is not exposed`,
-                                };
-                            }
-                        }
+                    const unexposed = unexposedPlanRelation(planRows, (relation) => allowlist.has(relation));
+                    if (unexposed !== undefined) {
+                        return {
+                            isError: true,
+                            text: `query touches ${JSON.stringify(unexposed)}, which is not exposed`,
+                        };
                     }
                     let limitedRows;
                     try {
                         // Fetch at most one row beyond the advertised cap. This keeps the database driver's
                         // materialization bounded even when the caller submits `SELECT * FROM a_huge_table`.
-                        // nifra-expect sql-dynamic: wraps the caller's already-allowlisted read-only query (guarded above) in a bounding subselect; the query text is the tool's input, run on a read-only session
-                        limitedRows = await querySession.run(`SELECT * FROM (${query}) AS "__nifra_result" LIMIT ${maxRows + 1}`, deadline);
+                        limitedRows = await querySession.run(boundedSqliteQuery(query, maxRows), deadline);
                     }
                     catch (error) {
                         if (error instanceof QueryTimeoutError)
@@ -694,8 +175,7 @@ export function serveDatabaseAsMcp(db, options) {
                     let total = limitedRows.length;
                     if (wasLimited && options.runQuery?.exactTotal === true) {
                         try {
-                            // nifra-expect sql-dynamic: wraps the caller's already-allowlisted read-only query (guarded above) in a counting subselect; the query text is the tool's input, run on a read-only session
-                            const countRows = (await querySession.run(`SELECT count(*) AS "__nifra_total" FROM (${query}) AS "__nifra_count"`, deadline));
+                            const countRows = (await querySession.run(countSqliteQuery(query), deadline));
                             const count = countRows[0]?.__nifra_total;
                             if (typeof count !== "number" || !Number.isSafeInteger(count) || count < maxRows) {
                                 return { isError: true, text: "query failed: invalid row count" };
@@ -711,27 +191,14 @@ export function serveDatabaseAsMcp(db, options) {
                     else if (wasLimited) {
                         total = null;
                     }
-                    let shown = Math.min(limitedRows.length, maxRows);
-                    let payload = limitedRows.slice(0, shown);
-                    let result = {
-                        rows: payload,
+                    const fit = fitToBytes(Math.min(limitedRows.length, maxRows), (shown) => ({
+                        rows: limitedRows.slice(0, shown),
                         ...(wasLimited || shown < limitedRows.length ? { truncated: true, total } : {}),
-                    };
-                    let serialized = JSON.stringify(result);
-                    const encoded = new TextEncoder();
-                    while (encoded.encode(serialized).byteLength > maxResultBytes && shown > 0) {
-                        shown = Math.max(0, Math.floor(shown / 2));
-                        payload = limitedRows.slice(0, shown);
-                        result = {
-                            rows: payload,
-                            ...(wasLimited || shown < limitedRows.length ? { truncated: true, total } : {}),
-                        };
-                        serialized = JSON.stringify(result);
-                    }
-                    if (encoded.encode(serialized).byteLength > maxResultBytes) {
+                    }), maxResultBytes);
+                    if (fit === undefined) {
                         return { isError: true, text: "query result exceeds maxResultBytes" };
                     }
-                    return { text: serialized, structuredContent: result };
+                    return { text: fit.serialized, structuredContent: fit.value };
                 }
                 finally {
                     activeQueries -= 1;
@@ -744,6 +211,7 @@ export function serveDatabaseAsMcp(db, options) {
         name: options.name ?? "nifra-db",
         version: options.version ?? "1.0.0",
         tools,
+        ...(options.allowedHosts !== undefined ? { allowedHosts: options.allowedHosts } : {}),
     });
     const runQueryOptions = options.runQuery;
     if (session === undefined || runQueryOptions === undefined)