@sema-agent/core
Version:
Stateless, task-oriented AI agent core
105 lines (104 loc) • 4.63 kB
JavaScript
import { Type } from "typebox";
import { thrownText } from "../core/thrown-value.js";
const FORBIDDEN = /\b(insert|update|delete|drop|alter|create|truncate|grant|revoke|merge|replace|attach|detach|pragma|vacuum|exec|execute|call|copy|lock|reindex)\b/i;
class SqlValidationError extends Error {
}
function stripTrailingSemicolon(sql) {
return sql.replace(/;\s*$/, "");
}
function stripStringLiterals(sql) {
return sql.replace(/'(?:[^']|'')*'/g, "''").replace(/"(?:[^"]|"")*"/g, '""');
}
function extractReferencedTables(sql) {
const tables = [];
const re = /\b(?:from|join)\s+([a-zA-Z_][\w.$"`[\]]*)/gi;
let m;
while ((m = re.exec(sql)) !== null) {
const ident = m[1].replace(/["`[\]]/g, "");
const bare = ident.includes(".") ? ident.split(".").pop() : ident;
tables.push(bare.toLowerCase());
}
return tables;
}
export function validateReadOnlySql(sql, allowTables) {
const trimmed = stripTrailingSemicolon(sql.trim());
if (!trimmed) {
throw new SqlValidationError("empty query");
}
const scrubbed = stripStringLiterals(trimmed);
if (scrubbed.includes(";")) {
throw new SqlValidationError("multiple statements are not allowed");
}
if (scrubbed.includes("--") || scrubbed.includes("/*")) {
throw new SqlValidationError("SQL comments are not allowed");
}
if (!/^\s*(select|with)\b/i.test(scrubbed)) {
throw new SqlValidationError("only SELECT / WITH (read-only) queries are allowed");
}
if (/\binto\b/i.test(scrubbed)) {
throw new SqlValidationError("SELECT ... INTO is not allowed");
}
const forbidden = FORBIDDEN.exec(scrubbed);
if (forbidden) {
throw new SqlValidationError(`forbidden keyword "${forbidden[1].toLowerCase()}"`);
}
if (allowTables && allowTables.length > 0) {
if (/\bunion\b|\bintersect\b|\bexcept\b/i.test(scrubbed)) {
throw new SqlValidationError("set operations (UNION/INTERSECT/EXCEPT) are not allowed when allowTables is set — table references cannot be reliably enumerated");
}
if (/\b(?:from|join)\s*\(/i.test(scrubbed)) {
throw new SqlValidationError("sub-queries in FROM/JOIN are not allowed when allowTables is set — table references cannot be reliably enumerated");
}
if (/\bfrom\s+[a-zA-Z_][\w.$"`[\]]*\s*,/i.test(scrubbed)) {
throw new SqlValidationError("comma-joins (`FROM a, b`) are not allowed when allowTables is set — use an explicit JOIN so every table is checked");
}
const allow = new Set(allowTables.map((t) => t.toLowerCase()));
const cteRe = /([a-zA-Z_]\w*)\s+as\s*\(/gi;
let cm;
while ((cm = cteRe.exec(scrubbed)) !== null) {
allow.add(cm[1].toLowerCase());
}
const referenced = extractReferencedTables(scrubbed);
const denied = referenced.filter((t) => !allow.has(t));
if (denied.length > 0) {
throw new SqlValidationError(`table(s) not allowed: ${[...new Set(denied)].join(", ")}. Allowed: ${[...allow].join(", ")}`);
}
}
}
export function createSqlTool(opts) {
const maxRows = opts.maxRows ?? 100;
return {
name: opts.name ?? "query_sql",
effect: "read",
description: opts.description ??
`Run a READ-ONLY SQL query (SELECT/WITH only) and get rows back (max ${maxRows}).` +
(opts.allowTables?.length ? ` Allowed tables: ${opts.allowTables.join(", ")}.` : "") +
" Do not attempt to modify data.",
parameters: Type.Object({
sql: Type.String({ description: "A single read-only SELECT/WITH query." }),
}),
execute: async (args) => {
const sql = String(args.sql ?? "");
try {
validateReadOnlySql(sql, opts.allowTables);
}
catch (e) {
const msg = thrownText(e);
return { content: `SQL rejected: ${msg}`, details: { error: msg }, terminate: false };
}
const rows = await opts.query(stripTrailingSemicolon(sql.trim()));
const truncated = rows.length > maxRows;
const shown = truncated ? rows.slice(0, maxRows) : rows;
const payload = {
rowCount: rows.length,
returned: shown.length,
truncated,
rows: shown,
};
return {
content: JSON.stringify(payload, null, 2),
details: { rowCount: rows.length, truncated },
};
},
};
}