@nifrajs/mcp-db 3.5.0 → 4.0.2

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.js CHANGED
@@ -23,6 +23,8 @@
23
23
  * app.mount("/mcp", mcp.fetch)
24
24
  */
25
25
  import { createMcpServer, defineMcpTool } from "@nifrajs/mcp";
26
+ import { boundedSqliteQuery, countSqliteQuery, fitToBytes, gateSqliteStatement, unexposedPlanRelation, } from "./internal/shared.js";
27
+ import { openQuerySession, QueryTimeoutError, } from "./internal/sqlite-session.js";
26
28
  export class McpDbConfigError extends Error {
27
29
  constructor(message) {
28
30
  super(`@nifrajs/mcp-db: ${message}`);
@@ -42,556 +44,11 @@ function sqlIdentifiers(strings, ...identifiers) {
42
44
  }
43
45
  return sql;
44
46
  }
45
- const stripSqlNoise = (sql) => sql
46
- // Strings first so comment markers inside literals don't count.
47
- .replace(/'(?:[^']|'')*'/g, "''")
48
- .replace(/"(?:[^"]|"")*"/g, '""')
49
- .replace(/--[^\n]*/g, " ")
50
- .replace(/\/\*[\s\S]*?\*\//g, " ");
51
- // SQLite's bare-identifier alphabet. Non-ASCII is included because SQLite treats any byte >= 0x80 as
52
- // an identifier character; leaving it out would end a word early and split one table name into two.
53
- const WORD_START = /[A-Za-z_\u0080-\uffff]/;
54
- const WORD_PART = /[A-Za-z0-9_$\u0080-\uffff]/;
55
- /**
56
- * Tokenize far enough to name every relation. A regex cannot do this job: SQLite needs no separator
57
- * between a keyword and a quoted identifier (`FROM"users"` is legal), and it accepts four identifier
58
- * quotings - `"x"`, `[x]`, `` `x` ``, and bare. A `\s+`-anchored pattern silently reads none of
59
- * those as a table reference, which is exactly how an aliased `FROM"users"AS habits` slipped past
60
- * the allowlist and returned a non-exposed table's rows.
61
- *
62
- * String literals and comments become nothing, so a `FROM` inside attacker-controlled text is inert.
63
- */
64
- function tokenizeSql(sql) {
65
- const tokens = [];
66
- let i = 0;
67
- /** Read a quoted identifier, collapsing the doubled-delimiter escape SQLite uses for `"` and `` ` ``. */
68
- const readQuoted = (close, escapable) => {
69
- let value = "";
70
- i++; // past the opening delimiter
71
- while (i < sql.length) {
72
- if (sql[i] === close) {
73
- if (escapable && sql[i + 1] === close) {
74
- value += close;
75
- i += 2;
76
- continue;
77
- }
78
- i++;
79
- tokens.push({ kind: "word", value, bare: false });
80
- return;
81
- }
82
- value += sql[i];
83
- i++;
84
- }
85
- // Unterminated - emit what we have. The caller's single-statement and plan checks still run.
86
- tokens.push({ kind: "word", value, bare: false });
87
- };
88
- while (i < sql.length) {
89
- const char = sql[i];
90
- if (char === " " || char === "\t" || char === "\n" || char === "\r" || char === "\f") {
91
- i++;
92
- }
93
- else if (char === "-" && sql[i + 1] === "-") {
94
- const end = sql.indexOf("\n", i);
95
- i = end === -1 ? sql.length : end + 1;
96
- }
97
- else if (char === "/" && sql[i + 1] === "*") {
98
- const end = sql.indexOf("*/", i + 2);
99
- i = end === -1 ? sql.length : end + 2;
100
- }
101
- else if (char === "'") {
102
- // Literal: consumed and dropped, so its contents can never be read as SQL.
103
- i++;
104
- while (i < sql.length) {
105
- if (sql[i] === "'") {
106
- if (sql[i + 1] === "'") {
107
- i += 2;
108
- continue;
109
- }
110
- i++;
111
- break;
112
- }
113
- i++;
114
- }
115
- }
116
- else if (char === '"') {
117
- readQuoted('"', true);
118
- }
119
- else if (char === "`") {
120
- readQuoted("`", true);
121
- }
122
- else if (char === "[") {
123
- readQuoted("]", false); // MSSQL-style bracket quoting: no escape form
124
- }
125
- else if (WORD_START.test(char)) {
126
- const start = i;
127
- while (i < sql.length && WORD_PART.test(sql[i]))
128
- i++;
129
- tokens.push({ kind: "word", value: sql.slice(start, i), bare: true });
130
- }
131
- else if (char >= "0" && char <= "9") {
132
- while (i < sql.length && /[0-9A-Za-z._]/.test(sql[i]))
133
- i++;
134
- }
135
- else {
136
- tokens.push({ kind: "punct", value: char, bare: false });
137
- i++;
138
- }
139
- }
140
- return tokens;
141
- }
142
- /** True when `token` is the unquoted keyword `word` (case-insensitive). */
143
- const isKeyword = (token, word) => token?.bare === true && token.value.toLowerCase() === word;
144
- /**
145
- * Every relation the statement names after `FROM`/`JOIN`, lowercased, with CTE names excluded.
146
- * A schema qualifier resolves to its table (`main.habits` -> `habits`); a subquery contributes
147
- * nothing of its own because its inner `FROM` is tokenized alongside everything else.
148
- */
149
- function relationNames(sql) {
150
- const tokens = tokenizeSql(sql);
151
- // `name AS (` only ever introduces a CTE - a column alias cannot be followed by a paren - so this
152
- // needs no `WITH` tracking to be exact.
153
- const ctes = new Set();
154
- for (let i = 1; i < tokens.length - 1; i++) {
155
- const name = tokens[i - 1];
156
- if (isKeyword(tokens[i], "as") && tokens[i + 1]?.value === "(" && name.kind === "word") {
157
- ctes.add(name.value.toLowerCase());
158
- }
159
- }
160
- const names = [];
161
- const relationAt = (index) => {
162
- const first = tokens[index];
163
- if (first === undefined || first.kind !== "word")
164
- return undefined; // `FROM (subquery)` - nothing to name here
165
- // A qualified reference names the table second; an unqualified one names it first.
166
- const qualified = tokens[index + 1]?.value === "." && tokens[index + 2]?.kind === "word";
167
- const table = qualified ? tokens[index + 2].value : first.value;
168
- return { name: table.toLowerCase(), next: qualified ? index + 3 : index + 1 };
169
- };
170
- const addRelation = (index) => {
171
- const relation = relationAt(index);
172
- if (relation === undefined)
173
- return undefined;
174
- if (relation.name !== "" && !ctes.has(relation.name))
175
- names.push(relation.name);
176
- return relation.next;
177
- };
178
- // Comma joins are part of the FROM table-source grammar but do not have a JOIN keyword. Walk the
179
- // source list after every FROM and collect each top-level comma operand. The walk is deliberately
180
- // conservative: an uncertain token is left for SQLite's plan check, while an apparent relation is
181
- // always added so an alias cannot hide an unexposed table.
182
- const clauseBoundary = new Set([
183
- "where",
184
- "group",
185
- "order",
186
- "having",
187
- "limit",
188
- "offset",
189
- "union",
190
- "except",
191
- "intersect",
192
- "window",
193
- "returning",
194
- ]);
195
- for (let i = 0; i < tokens.length; i++) {
196
- const from = isKeyword(tokens[i], "from");
197
- const join = isKeyword(tokens[i], "join");
198
- if (!from && !join)
199
- continue;
200
- const next = addRelation(i + 1);
201
- if (join || next === undefined)
202
- continue;
203
- let depth = 0;
204
- for (let j = next; j < tokens.length; j++) {
205
- const token = tokens[j];
206
- if (token === undefined)
207
- break;
208
- if (token.value === "(") {
209
- depth++;
210
- continue;
211
- }
212
- if (token.value === ")") {
213
- if (depth === 0)
214
- break;
215
- depth--;
216
- continue;
217
- }
218
- if (depth !== 0)
219
- continue;
220
- if (token.bare && clauseBoundary.has(token.value.toLowerCase()))
221
- break;
222
- if (token.value !== ",")
223
- continue;
224
- // A comma at the source-list level is another table reference. If it is a subquery, the nested
225
- // FROM scan below still finds its base tables; no relation is inferred from the opening paren.
226
- const afterComma = addRelation(j + 1);
227
- if (afterComma !== undefined)
228
- j = afterComma - 1;
229
- }
230
- }
231
- return [...new Set(names)];
232
- }
233
- /** Reject multi-statement input: allow one terminator only when it is the final character. */
234
- const isSingleStatement = (sql) => {
235
- let quote = null;
236
- let lineComment = false;
237
- let blockComment = false;
238
- for (let i = 0; i < sql.length; i++) {
239
- const char = sql[i];
240
- const next = sql[i + 1];
241
- if (lineComment) {
242
- if (char === "\n")
243
- lineComment = false;
244
- continue;
245
- }
246
- if (blockComment) {
247
- if (char === "*" && next === "/") {
248
- blockComment = false;
249
- i++;
250
- }
251
- continue;
252
- }
253
- if (quote !== null) {
254
- if (char === quote) {
255
- if (next === quote)
256
- i++;
257
- else
258
- quote = null;
259
- }
260
- continue;
261
- }
262
- if (char === "'" || char === '"') {
263
- quote = char;
264
- continue;
265
- }
266
- if (char === "-" && next === "-") {
267
- lineComment = true;
268
- i++;
269
- continue;
270
- }
271
- if (char === "/" && next === "*") {
272
- blockComment = true;
273
- i++;
274
- continue;
275
- }
276
- if (char === ";") {
277
- return sql.slice(i + 1).trim() === "";
278
- }
279
- }
280
- return true;
281
- };
282
- const isReadStatement = (sql) => {
283
- const head = stripSqlNoise(sql).trim().toLowerCase();
284
- return head.startsWith("select") || head.startsWith("with");
285
- };
286
- /** Remove one trailing statement terminator (and anything after it) while preserving literals and
287
- * identifiers. Wrapping the query in a bounded subquery must not let a trailing comment consume the
288
- * wrapper's closing parenthesis. */
289
- function executableSql(sql) {
290
- let quote = null;
291
- let lineComment = false;
292
- let lineCommentStart = -1;
293
- let blockComment = false;
294
- for (let i = 0; i < sql.length; i++) {
295
- const char = sql[i];
296
- const next = sql[i + 1];
297
- if (lineComment) {
298
- if (char === "\n") {
299
- lineComment = false;
300
- lineCommentStart = -1;
301
- }
302
- continue;
303
- }
304
- if (blockComment) {
305
- if (char === "*" && next === "/") {
306
- blockComment = false;
307
- i++;
308
- }
309
- continue;
310
- }
311
- if (quote !== null) {
312
- if (char === quote) {
313
- if (next === quote)
314
- i++;
315
- else
316
- quote = null;
317
- }
318
- continue;
319
- }
320
- if (char === "'" || char === '"') {
321
- quote = char;
322
- continue;
323
- }
324
- if (char === "-" && next === "-") {
325
- lineComment = true;
326
- lineCommentStart = i;
327
- i++;
328
- continue;
329
- }
330
- if (char === "/" && next === "*") {
331
- blockComment = true;
332
- i++;
333
- continue;
334
- }
335
- if (char === ";")
336
- return sql.slice(0, i).trim();
337
- }
338
- if (lineComment && lineCommentStart >= 0)
339
- return sql.slice(0, lineCommentStart).trim();
340
- return sql.trim();
341
- }
342
47
  function assertResultLimit(value, name) {
343
48
  if (!Number.isSafeInteger(value) || value <= 0 || value >= Number.MAX_SAFE_INTEGER) {
344
49
  throw new McpDbConfigError(`${name} must be a finite positive safe integer`);
345
50
  }
346
51
  }
347
- class QueryTimeoutError extends Error {
348
- constructor() {
349
- super("query exceeded the time limit");
350
- this.name = "QueryTimeoutError";
351
- }
352
- }
353
- /**
354
- * The isolated lane. The worker reopens the SAME database FILE read-only and answers one statement
355
- * at a time; it is spawned once and reused for the life of the server. It deliberately does not
356
- * take a `serialize()` snapshot - that copies the whole database into memory on EVERY call, which
357
- * is a larger availability problem than the slow query the deadline exists to bound.
358
- */
359
- const SQLITE_WORKER_SOURCE = `
360
- const { Database } = require("bun:sqlite")
361
- let db
362
- self.onmessage = (event) => {
363
- const { id, filename, sql, type } = event.data
364
- if (type === "close") {
365
- let response
366
- try {
367
- // prepare().all() leaves a statement alive until it is finalized or collected. Bun's
368
- // default close(false) preserves those statements and can keep the database file locked;
369
- // close(true) finalizes every outstanding statement before the worker exits.
370
- db?.close(true)
371
- db = undefined
372
- response = { id, closed: true, ok: true }
373
- } catch (error) {
374
- response = { id, closed: true, ok: false, error: String(error) }
375
- }
376
- // A forced parent-side terminate can leave native SQLite handles locked on Windows even after
377
- // close() returns. Let the worker exit its own event turn so the runtime can finish releasing
378
- // those handles before the parent considers graceful cleanup complete.
379
- self.postMessage(response)
380
- self.close()
381
- return
382
- }
383
- try {
384
- if (db === undefined) {
385
- db = new Database(filename, { readonly: true })
386
- db.run("PRAGMA query_only = ON")
387
- }
388
- self.postMessage({ id, ok: true, rows: db.prepare(sql).all() })
389
- } catch (error) {
390
- self.postMessage({ id, ok: false, error: String(error) })
391
- }
392
- }
393
- `;
394
- function isWorkerAvailable() {
395
- return typeof Worker !== "undefined" && typeof Blob !== "undefined" && typeof URL !== "undefined";
396
- }
397
- /** Only a file-backed database can be reopened in a worker. `:memory:` and the anonymous temp
398
- * database live in this process alone, so they take the in-process lane instead. */
399
- function reopenableFilename(db) {
400
- const filename = db.filename;
401
- if (typeof filename !== "string" || filename === "" || filename === ":memory:")
402
- return undefined;
403
- return filename.startsWith("file::memory:") ? undefined : filename;
404
- }
405
- const WORKER_CLOSE_TIMEOUT_MS = 1_000;
406
- function workerSession(filename) {
407
- let worker;
408
- let sourceUrl;
409
- let nextId = 0;
410
- const pending = new Map();
411
- let lane = Promise.resolve();
412
- let closed = false;
413
- let closePromise;
414
- let pendingClose;
415
- const detachWorker = (active = worker) => {
416
- if (active === undefined)
417
- return;
418
- if (worker !== active)
419
- return;
420
- worker = undefined;
421
- if (sourceUrl !== undefined)
422
- URL.revokeObjectURL(sourceUrl);
423
- sourceUrl = undefined;
424
- };
425
- const releaseWorker = (active = worker) => {
426
- if (active === undefined)
427
- return;
428
- active.terminate();
429
- detachWorker(active);
430
- };
431
- const rejectPending = (error) => {
432
- const waiters = [...pending.values()];
433
- pending.clear();
434
- for (const waiter of waiters)
435
- waiter.reject(error);
436
- };
437
- /** Drop the worker and fail everything riding on it. The only bound this side can enforce over a
438
- * statement already running inside native SQLite is to stop owning the connection it runs on. */
439
- const discard = (error) => {
440
- const closing = pendingClose;
441
- pendingClose = undefined;
442
- releaseWorker();
443
- rejectPending(error);
444
- closing?.resolve();
445
- };
446
- const ensureWorker = () => {
447
- if (worker !== undefined)
448
- return worker;
449
- sourceUrl = URL.createObjectURL(new Blob([SQLITE_WORKER_SOURCE], { type: "text/javascript" }));
450
- const spawned = new Worker(sourceUrl);
451
- spawned.onmessage = (event) => {
452
- const { id, ok, rows, error } = event.data;
453
- if (typeof id !== "number")
454
- return;
455
- if (event.data.closed === true) {
456
- const closing = pendingClose;
457
- if (closing?.id !== id)
458
- return;
459
- pendingClose = undefined;
460
- closing.resolve(event.data.ok === true);
461
- return;
462
- }
463
- const waiter = pending.get(id);
464
- if (waiter === undefined)
465
- return;
466
- if (Date.now() >= waiter.deadline) {
467
- discard(new QueryTimeoutError());
468
- return;
469
- }
470
- pending.delete(id);
471
- if (ok === true)
472
- waiter.resolve(rows ?? []);
473
- else
474
- waiter.reject(new Error(error ?? "query failed in worker"));
475
- };
476
- spawned.onerror = () => discard(new Error("query failed in worker"));
477
- spawned.unref?.();
478
- worker = spawned;
479
- return spawned;
480
- };
481
- const dispatch = (sql, deadline) => {
482
- if (closed)
483
- return Promise.reject(new Error("query session is closed"));
484
- const active = ensureWorker();
485
- const remaining = deadline - Date.now();
486
- if (remaining <= 0) {
487
- discard(new QueryTimeoutError());
488
- return Promise.reject(new QueryTimeoutError());
489
- }
490
- const id = ++nextId;
491
- return new Promise((resolve, reject) => {
492
- const timer = setTimeout(() => discard(new QueryTimeoutError()), remaining);
493
- pending.set(id, {
494
- deadline,
495
- resolve: (rows) => {
496
- clearTimeout(timer);
497
- resolve(rows);
498
- },
499
- reject: (error) => {
500
- clearTimeout(timer);
501
- reject(error);
502
- },
503
- });
504
- try {
505
- active.postMessage({ id, filename, sql });
506
- }
507
- catch (error) {
508
- pending.delete(id);
509
- clearTimeout(timer);
510
- reject(error instanceof Error ? error : new Error(String(error)));
511
- }
512
- });
513
- };
514
- return {
515
- run(sql, deadline) {
516
- // One worker owns the connection, so overlapping calls queue behind each other rather than
517
- // interleaving statements on it.
518
- const result = lane.then(() => dispatch(sql, deadline));
519
- lane = result.catch(() => undefined);
520
- return result;
521
- },
522
- close() {
523
- if (closePromise !== undefined)
524
- return closePromise;
525
- closed = true;
526
- closePromise = (async () => {
527
- const active = worker;
528
- if (active !== undefined) {
529
- await new Promise((resolve) => {
530
- const id = ++nextId;
531
- let settled = false;
532
- const finish = (error, graceful = false) => {
533
- if (settled)
534
- return;
535
- settled = true;
536
- clearTimeout(timer);
537
- if (pendingClose?.id === id)
538
- pendingClose = undefined;
539
- if (error !== undefined)
540
- rejectPending(error);
541
- if (graceful)
542
- detachWorker(active);
543
- else
544
- releaseWorker(active);
545
- resolve();
546
- };
547
- const timer = setTimeout(() => {
548
- finish(new Error("query session is closed"));
549
- }, WORKER_CLOSE_TIMEOUT_MS);
550
- pendingClose = { id, resolve: (graceful) => finish(undefined, graceful) };
551
- try {
552
- active.postMessage({ id, type: "close" });
553
- }
554
- catch {
555
- finish(new Error("query session is closed"));
556
- }
557
- });
558
- }
559
- await lane;
560
- })();
561
- return closePromise;
562
- },
563
- };
564
- }
565
- /**
566
- * The in-process lane, used when the database cannot be reopened elsewhere. The deadline is checked
567
- * around the statement and nothing more: `prepare().all()` is synchronous, so it holds the only
568
- * thread there is and no timer of ours can fire while it runs. That makes this lane BEST EFFORT -
569
- * it reports an overrun, it does not cut one short.
570
- */
571
- function inProcessSession(db) {
572
- return {
573
- async run(sql, deadline) {
574
- if (Date.now() >= deadline)
575
- throw new QueryTimeoutError();
576
- const rows = db.prepare(sql).all();
577
- if (Date.now() > deadline)
578
- throw new QueryTimeoutError();
579
- return rows;
580
- },
581
- async close() { },
582
- };
583
- }
584
- /**
585
- * Pick the lane once, at construction. A file-backed database is isolated in a reusable read-only
586
- * worker, which is the only arrangement that can actually stop a running statement. Anything else
587
- * runs in process - every database keeps working, none is refused.
588
- */
589
- function openQuerySession(db) {
590
- const filename = reopenableFilename(db);
591
- if (filename !== undefined && isWorkerAvailable())
592
- return workerSession(filename);
593
- return inProcessSession(db);
594
- }
595
52
  /**
596
53
  * Serve `db` as a mountable MCP server (`mcp.fetch` at `POST /mcp`). See module docs for the
597
54
  * security model. Throws {@link McpDbConfigError} on any unsafe configuration - always at
@@ -680,31 +137,13 @@ export function serveDatabaseAsMcp(db, options) {
680
137
  }
681
138
  activeQueries += 1;
682
139
  try {
683
- const sql = String(args.sql ?? "").trim();
684
- if (sql === "")
685
- return { isError: true, text: "empty query" };
686
- if (!isSingleStatement(sql)) {
687
- return { isError: true, text: "only a single statement is allowed" };
688
- }
689
- if (!isReadStatement(sql)) {
690
- return { isError: true, text: "only SELECT (or WITH…SELECT) statements are allowed" };
691
- }
140
+ const gate = gateSqliteStatement(String(args.sql ?? ""), (relation) => allowlist.has(relation));
141
+ if (!gate.ok)
142
+ return { isError: true, text: gate.text };
692
143
  const deadline = Date.now() + queryTimeoutMs;
693
- // Verify SQL relation tokens before the plan check. SQLite's plan output uses aliases as
694
- // scan targets (e.g. `users AS habits` becomes `SCAN habits`), so plan-only validation can
695
- // mistake an attacker-controlled alias for an allowlisted table.
696
- const executable = executableSql(sql);
697
- for (const relation of relationNames(executable)) {
698
- if (!allowlist.has(relation)) {
699
- return {
700
- isError: true,
701
- text: `query touches ${JSON.stringify(relation)}, which is not exposed`,
702
- };
703
- }
704
- }
705
144
  // Verify via the query plan that only allowlisted tables are touched. SQLite names every
706
145
  // scanned/searched relation in EXPLAIN QUERY PLAN detail rows.
707
- const query = executable;
146
+ const query = gate.query;
708
147
  let planRows;
709
148
  try {
710
149
  planRows = (await querySession.run(`EXPLAIN QUERY PLAN ${query}`, deadline));
@@ -714,34 +153,18 @@ export function serveDatabaseAsMcp(db, options) {
714
153
  return { isError: true, text: error.message };
715
154
  return { isError: true, text: `query failed to plan: ${String(error)}` };
716
155
  }
717
- for (const row of planRows) {
718
- const detail = typeof row.detail === "string" ? row.detail : "";
719
- // SQLite also emits non-table nodes such as `SCAN CONSTANT ROW` for a
720
- // constant-only SELECT. Only treat a scan/search target as a relation
721
- // when it is not one of those planner pseudo-nodes.
722
- if (/^SCAN\s+CONSTANT\s+ROW$/i.test(detail.trim()) ||
723
- /^SCAN\s+SUBQUERY\s+\d+$/i.test(detail.trim())) {
724
- continue;
725
- }
726
- // A schema-qualified reference plans as `SCAN main.habits`. The table is the last dotted
727
- // segment; reading the first would reject `main.habits` as a table literally named "main".
728
- const match = /(?:SCAN|SEARCH)\s+(?:TABLE\s+)?(?:[A-Za-z_][A-Za-z0-9_]*\.)?([A-Za-z_][A-Za-z0-9_]*)/i.exec(detail);
729
- if (match !== null) {
730
- const relation = match[1]?.toLowerCase() ?? "";
731
- if (!allowlist.has(relation)) {
732
- return {
733
- isError: true,
734
- text: `query touches ${JSON.stringify(match[1])}, which is not exposed`,
735
- };
736
- }
737
- }
156
+ const unexposed = unexposedPlanRelation(planRows, (relation) => allowlist.has(relation));
157
+ if (unexposed !== undefined) {
158
+ return {
159
+ isError: true,
160
+ text: `query touches ${JSON.stringify(unexposed)}, which is not exposed`,
161
+ };
738
162
  }
739
163
  let limitedRows;
740
164
  try {
741
165
  // Fetch at most one row beyond the advertised cap. This keeps the database driver's
742
166
  // materialization bounded even when the caller submits `SELECT * FROM a_huge_table`.
743
- // 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
744
- limitedRows = await querySession.run(`SELECT * FROM (${query}) AS "__nifra_result" LIMIT ${maxRows + 1}`, deadline);
167
+ limitedRows = await querySession.run(boundedSqliteQuery(query, maxRows), deadline);
745
168
  }
746
169
  catch (error) {
747
170
  if (error instanceof QueryTimeoutError)
@@ -752,8 +175,7 @@ export function serveDatabaseAsMcp(db, options) {
752
175
  let total = limitedRows.length;
753
176
  if (wasLimited && options.runQuery?.exactTotal === true) {
754
177
  try {
755
- // 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
756
- const countRows = (await querySession.run(`SELECT count(*) AS "__nifra_total" FROM (${query}) AS "__nifra_count"`, deadline));
178
+ const countRows = (await querySession.run(countSqliteQuery(query), deadline));
757
179
  const count = countRows[0]?.__nifra_total;
758
180
  if (typeof count !== "number" || !Number.isSafeInteger(count) || count < maxRows) {
759
181
  return { isError: true, text: "query failed: invalid row count" };
@@ -769,27 +191,14 @@ export function serveDatabaseAsMcp(db, options) {
769
191
  else if (wasLimited) {
770
192
  total = null;
771
193
  }
772
- let shown = Math.min(limitedRows.length, maxRows);
773
- let payload = limitedRows.slice(0, shown);
774
- let result = {
775
- rows: payload,
194
+ const fit = fitToBytes(Math.min(limitedRows.length, maxRows), (shown) => ({
195
+ rows: limitedRows.slice(0, shown),
776
196
  ...(wasLimited || shown < limitedRows.length ? { truncated: true, total } : {}),
777
- };
778
- let serialized = JSON.stringify(result);
779
- const encoded = new TextEncoder();
780
- while (encoded.encode(serialized).byteLength > maxResultBytes && shown > 0) {
781
- shown = Math.max(0, Math.floor(shown / 2));
782
- payload = limitedRows.slice(0, shown);
783
- result = {
784
- rows: payload,
785
- ...(wasLimited || shown < limitedRows.length ? { truncated: true, total } : {}),
786
- };
787
- serialized = JSON.stringify(result);
788
- }
789
- if (encoded.encode(serialized).byteLength > maxResultBytes) {
197
+ }), maxResultBytes);
198
+ if (fit === undefined) {
790
199
  return { isError: true, text: "query result exceeds maxResultBytes" };
791
200
  }
792
- return { text: serialized, structuredContent: result };
201
+ return { text: fit.serialized, structuredContent: fit.value };
793
202
  }
794
203
  finally {
795
204
  activeQueries -= 1;
@@ -802,6 +211,7 @@ export function serveDatabaseAsMcp(db, options) {
802
211
  name: options.name ?? "nifra-db",
803
212
  version: options.version ?? "1.0.0",
804
213
  tools,
214
+ ...(options.allowedHosts !== undefined ? { allowedHosts: options.allowedHosts } : {}),
805
215
  });
806
216
  const runQueryOptions = options.runQuery;
807
217
  if (session === undefined || runQueryOptions === undefined)