@adonis-agora/agent/dist/src/data/data-tool.js

Governed, durable-backed AI agent (chat + tool-calling + governance) for AdonisJS — part of the Agora ecosystem.

import { z } from 'zod';
import { defineTool } from '../ai-tool-ref.js';
import { injectLimit } from './limit.js';
import { loadSqlParser } from './parser.js';
import { SqlValidator } from './sql-validator.js';
import { GroupTableAccessPolicy, } from './table-access.js';
import { TenantScopeRewriter } from './tenant-scope.js';
const DEFAULT_MAX_ROWS = 100;
const DEFAULT_TOOL_NAME = 'executeSql';
/** Serialized rows above this size are truncated so a single tool result can't blow the context. */
const MAX_RESULT_BYTES = 256 * 1024;
const inputSchema = z.object({
    sql: z.string().min(1).describe('A single read-only MySQL SELECT statement'),
});
const DESCRIPTION = 'Execute a single read-only MySQL SELECT statement. Use this to answer questions about real ' +
    'data. Only SELECT is allowed (no INSERT/UPDATE/DELETE/DDL). Access is restricted to the ' +
    'tables your role is permitted to read, results are capped, and tenant-scoped tables are ' +
    'automatically constrained to your current tenant.';
function isTableAccessPolicy(value) {
    return typeof value.canAccess === 'function';
}
/** Normalize whatever a Lucid `rawQuery` returns (SQLite array, PG `{rows}`, MySQL `[rows,fields]`). */
function normalizeRows(raw) {
    if (Array.isArray(raw)) {
        // mysql2 returns `[rows, fields]`; the first element is itself the rows array.
        if (raw.length > 0 && Array.isArray(raw[0])) {
            return raw[0];
        }
        return raw;
    }
    if (raw !== null && typeof raw === 'object' && Array.isArray(raw.rows)) {
        return raw.rows;
    }
    return [];
}
function withTimeout(promise, ms) {
    if (ms === undefined)
        return promise;
    return new Promise((resolve, reject) => {
        const timer = setTimeout(() => {
            reject(new Error(`Query exceeded the ${ms}ms statement timeout.`));
        }, ms);
        promise.then((value) => {
            clearTimeout(timer);
            resolve(value);
        }, (err) => {
            clearTimeout(timer);
            reject(err);
        });
    });
}
/**
 * Builds the governed `executeSql` read tool as a branded functional tool ({@link defineTool}) the agent
 * registry auto-discovers. Register it in `config/agent.ts` (`defineConfig({ tools: [dataTool({ ... })] })`)
 * or export it from an `app/agent_tools/` module.
 *
 * The pipeline per call (identical to the reference satellite) is:
 * validate (single SELECT) → assert table access for every referenced table →
 * tenant-scope rewrite (if configured) → inject LIMIT → run → return rows,
 * truncating the payload if it would exceed ~256KB. `node-sql-parser` is imported lazily on the first
 * call, so it stays an optional peer.
 */
export function dataTool(config) {
    const maxRows = config.maxRows ?? DEFAULT_MAX_ROWS;
    const tableAccess = isTableAccessPolicy(config.tableAccess)
        ? config.tableAccess
        : new GroupTableAccessPolicy(config.tableAccess);
    const runner = resolveRunner(config);
    // Lazily built once from the lazily-imported parser, then reused across calls.
    let pipeline = null;
    const getPipeline = () => {
        if (pipeline === null) {
            pipeline = loadSqlParser().then((parser) => ({
                validator: new SqlValidator(parser),
                ...(config.tenant !== undefined
                    ? { tenantScope: new TenantScopeRewriter(config.tenant, parser) }
                    : {}),
            }));
        }
        return pipeline;
    };
    return defineTool({
        name: config.name ?? DEFAULT_TOOL_NAME,
        kind: 'read',
        description: DESCRIPTION,
        input: inputSchema,
        ...(config.roles !== undefined ? { roles: config.roles } : {}),
        ...(config.ability !== undefined ? { ability: config.ability } : {}),
    }, async (input, ctx) => {
        const { validator, tenantScope } = await getPipeline();
        const { tables } = validator.validate(input.sql);
        const roles = ctx.actor.roles ?? [];
        const forbidden = tables.filter((table) => !tableAccess.canAccess(roles, table));
        if (forbidden.length > 0) {
            const formatted = forbidden.map((table) => `\`${table}\``).join(', ');
            const rolesLabel = roles.length > 0 ? roles.join(', ') : 'none';
            throw new Error(`Your roles (${rolesLabel}) are not allowed to query ${formatted}.`);
        }
        let sql = input.sql;
        if (tenantScope) {
            // CRITICAL: pass ctx.actor.tenantRef through UNTOUCHED. A strictly-`undefined` tenantRef is the
            // privileged pass-through; a `null`/empty one must NOT be coerced to undefined (that would leak
            // every tenant's rows), so never `?? undefined` here.
            sql = tenantScope.rewrite(sql, ctx.actor.tenantRef);
        }
        const parser = await loadSqlParser();
        sql = injectLimit(parser, sql, maxRows);
        const rows = await withTimeout(runner.run(sql), config.statementTimeoutMs);
        return buildResult(rows, sql);
    });
}
function resolveRunner(config) {
    if (config.runner !== undefined)
        return config.runner;
    const db = config.db;
    if (db === undefined) {
        throw new Error('dataTool: provide either `db` (a Lucid database handle) or a custom `runner`.');
    }
    return {
        async run(sql) {
            const raw = await db.rawQuery(sql);
            return normalizeRows(raw);
        },
    };
}
function buildResult(rows, sql) {
    const rowCount = rows.length;
    if (byteLength(rows) <= MAX_RESULT_BYTES) {
        return { rows, rowCount, sql };
    }
    // Keep prefix rows until the serialized payload would exceed the cap.
    const kept = [];
    let size = 2; // opening + closing bracket of the JSON array
    for (const row of rows) {
        const rowSize = byteLength(row) + 1; // +1 for the joining comma
        if (size + rowSize > MAX_RESULT_BYTES)
            break;
        kept.push(row);
        size += rowSize;
    }
    return { rows: kept, rowCount, sql, truncated: true };
}
function byteLength(value) {
    return Buffer.byteLength(JSON.stringify(value) ?? '', 'utf8');
}
//# sourceMappingURL=data-tool.js.map