diff --git a/docs/features/mcp-server.mdx b/docs/features/mcp-server.mdx index 6ac549b..540d518 100644 --- a/docs/features/mcp-server.mdx +++ b/docs/features/mcp-server.mdx @@ -85,7 +85,7 @@ One tool per IAM permission. Each appears only if the agent holds that permissio | Tool | Permission | Accepts | |------|------------|---------| | `explain_query` | `explain` | A single `SELECT` to plan (options: `analyze`, `buffers`, `format`) | -| `query` | `read` | Read-only statements (`SELECT`, `SHOW`, …) | +| `query` | `read` | Read-only statements (`SELECT`, `SHOW`, …). Results are capped at 1000 rows (override down with `maxRows`) | | `write_data` | `write` | `INSERT` / `UPDATE` / `DELETE` / `COPY` | | `run_ddl` | `ddl` | `CREATE` / `ALTER` / `DROP` / `GRANT` / `REVOKE` / … | @@ -99,6 +99,8 @@ Every execution tool runs the submitted SQL through pgconsole's parser-based per `explain_query` only accepts a single `SELECT` (Postgres `EXPLAIN` rejects other statement kinds); with `analyze` (which actually executes the statement) it additionally requires every permission running the statement would require. +`query` caps its result at 1000 rows so a broad `SELECT` can't flood an agent's context. When capped, the response sets `truncated: true` and reports the full `rowCount` alongside the returned rows — narrow with `LIMIT`/`WHERE`, or pass a smaller `maxRows`. This cap applies to the MCP route only; the web UI is uncapped (it needs the full result set for CSV export and inline editing). + Least privilege: a pure agent gets only what its `agent:` IAM rules grant; a delegated agent can never exceed the user it acts for, and is further narrowed by its `permissions`/`connections` caps. Give each agent the narrowest grant it needs. diff --git a/server/mcp.ts b/server/mcp.ts index a667580..65502ff 100644 --- a/server/mcp.ts +++ b/server/mcp.ts @@ -23,6 +23,12 @@ declare const __APP_VERSION__: string export const MCP_PATH = '/mcp' const PAGE_SIZE = 100 +// Hard cap on rows returned by the execution tools, so a broad SELECT can't flood an agent's +// context. The full result is fetched, then capped; a `truncated` flag tells the agent to narrow +// (LIMIT/WHERE or a smaller maxRows). MCP only — the UI route is intentionally uncapped because it +// needs the full result set for CSV export and inline editing. +export const MAX_RESULT_ROWS = 1000 + // A resolved MCP caller — identity + audit actor for one agent. Permission resolution // itself lives in iam.ts (getAgentPermissions), the single home for that decision. export class Principal { @@ -176,10 +182,19 @@ const TOOLS: ToolDef[] = [ }, { name: 'query', - description: 'Run a read-only statement (SELECT, SHOW, …) and return the rows.', + description: `Run a read-only statement (SELECT, SHOW, …) and return the rows. Results are capped at ${MAX_RESULT_ROWS} rows; when capped, \`truncated\` is true and \`rowCount\` is the full total — narrow with LIMIT/WHERE or a smaller \`maxRows\`.`, inputSchema: { type: 'object', - properties: { ...connectionProp, sql: { type: 'string', description: 'A read-only statement, e.g. SELECT or SHOW.' } }, + properties: { + ...connectionProp, + sql: { type: 'string', description: 'A read-only statement, e.g. SELECT or SHOW.' }, + maxRows: { + type: 'integer', + minimum: 1, + maximum: MAX_RESULT_ROWS, + description: `Max rows to return (default and hard cap ${MAX_RESULT_ROWS}). Rows beyond this are dropped and \`truncated\` is set.`, + }, + }, required: ['connection', 'sql'], additionalProperties: false, }, @@ -271,6 +286,22 @@ function optStr(args: Record, name: string): string | undefined return trimmed === '' ? undefined : trimmed } +// The effective row cap for an execution call: the optional `maxRows` arg clamped to the hard +// ceiling, defaulting to the ceiling when absent. +export function readMaxRows(args: Record): number { + const v = args['maxRows'] + if (v === undefined || v === null) return MAX_RESULT_ROWS + if (typeof v !== 'number' || !Number.isInteger(v) || v < 1) { + throw new Error("'maxRows' must be a positive integer") + } + return Math.min(v, MAX_RESULT_ROWS) +} + +// Cap a materialized result set to `cap` rows, reporting whether any were dropped. +export function capRows(rows: T[], cap: number): { rows: T[]; truncated: boolean } { + return rows.length > cap ? { rows: rows.slice(0, cap), truncated: true } : { rows, truncated: false } +} + // ---- Tool implementations ---- async function listConnections(principal: Principal) { @@ -490,11 +521,32 @@ async function execute(principal: Principal, tool: string, expectedPerm: Permiss const finalSql = buildExecutableSql(rawSql, analysis) const result = await runAndAudit(principal, tool, connection, details, finalSql, rawSql) + // rowCount is the true total: result.count is the server's CommandComplete tag; the + // result.rows.length fallback is the true total only because the full result is materialized + // before capping. Revisit this if a cursor-based fetch is ever introduced (see follow-ups). const rowCount = result.count ?? result.rows.length + // maxRows is a `query`-only knob; other tools just get the hard cap (tool schemas aren't + // enforced at runtime, so don't honor a maxRows smuggled into write_data/run_ddl). + const cap = expectedPerm === 'read' ? readMaxRows(args) : MAX_RESULT_ROWS + const { rows, truncated } = capRows(result.rows, cap) if (expectedPerm === 'read') { - return { rowCount, columns: result.columns, rows: result.rows } + return { + rowCount, + returnedRows: rows.length, + truncated, + ...(truncated + ? { note: `Showing the first ${rows.length} of ${rowCount} rows. Refine with WHERE/ORDER BY/LIMIT to target the rows you need.` } + : {}), + columns: result.columns, + rows, + } + } + // write/ddl: RETURNING rows (usually few); surface the cap only when it actually bit. + return { + rowCount, + ...(truncated ? { returnedRows: rows.length, truncated: true } : {}), + rows: rows.length ? rows : undefined, } - return { rowCount, rows: result.rows.length ? result.rows : undefined } } // Shared handler for the execution tools (query/write_data/run_ddl). The disjoint permission to diff --git a/tests/mcp.test.ts b/tests/mcp.test.ts index 85492b9..ddccbf4 100644 --- a/tests/mcp.test.ts +++ b/tests/mcp.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect } from 'vitest' import { loadConfigFromString, getAgents, getAgentById, getAgentByToken } from '../server/lib/config' -import { selectToolNames, Principal, dispatchTool } from '../server/mcp' +import { selectToolNames, Principal, dispatchTool, capRows, readMaxRows, MAX_RESULT_ROWS } from '../server/mcp' import type { Permission } from '../server/lib/config' const BASE = ` @@ -253,3 +253,42 @@ describe('selectToolNames', () => { expect(selectToolNames(true, perms('ddl'))).toContain('run_ddl') }) }) + +describe('result cap (capRows / readMaxRows)', () => { + const mk = (n: number) => Array.from({ length: n }, (_, i) => i) + + it('passes through when under the cap', () => { + expect(capRows(mk(10), MAX_RESULT_ROWS)).toEqual({ rows: mk(10), truncated: false }) + }) + + it('does not truncate exactly at the cap', () => { + const r = capRows(mk(MAX_RESULT_ROWS), MAX_RESULT_ROWS) + expect(r.truncated).toBe(false) + expect(r.rows).toHaveLength(MAX_RESULT_ROWS) + }) + + it('truncates and flags when over the cap', () => { + const r = capRows(mk(MAX_RESULT_ROWS + 5), MAX_RESULT_ROWS) + expect(r.truncated).toBe(true) + expect(r.rows).toHaveLength(MAX_RESULT_ROWS) + }) + + it('readMaxRows defaults to the hard cap when absent', () => { + expect(readMaxRows({})).toBe(MAX_RESULT_ROWS) + }) + + it('readMaxRows clamps a larger request to the hard cap', () => { + expect(readMaxRows({ maxRows: MAX_RESULT_ROWS * 10 })).toBe(MAX_RESULT_ROWS) + }) + + it('readMaxRows honors a smaller request', () => { + expect(readMaxRows({ maxRows: 25 })).toBe(25) + }) + + it('readMaxRows rejects non-positive / non-integer values', () => { + expect(() => readMaxRows({ maxRows: 0 })).toThrow() + expect(() => readMaxRows({ maxRows: -1 })).toThrow() + expect(() => readMaxRows({ maxRows: 1.5 })).toThrow() + expect(() => readMaxRows({ maxRows: 'all' })).toThrow() + }) +})