diff --git a/packages/core/src/format.ts b/packages/core/src/format.ts new file mode 100644 index 0000000..8fbf719 --- /dev/null +++ b/packages/core/src/format.ts @@ -0,0 +1,54 @@ +export const DEFAULT_ROW_LIMIT = 100 +export const MAX_ROW_LIMIT = 1000 + +export const clampLimit = ( + requested: number | undefined, + fallback = DEFAULT_ROW_LIMIT, + max = MAX_ROW_LIMIT +): number => { + if (requested === undefined) { + return fallback + } + return Math.max(1, Math.min(Math.trunc(requested), max)) +} + +export const truncateRows = (rows: T[], limit: number): { rows: T[]; truncated: boolean } => + rows.length > limit + ? { rows: rows.slice(0, limit), truncated: true } + : { rows, truncated: false } + +// Slice a string to at most `maxBytes` UTF-8 bytes without splitting a +// multi-byte sequence (back off to a code-point boundary). +export const sliceUtf8 = (s: string, maxBytes: number): string => { + const buf = Buffer.from(s, 'utf8') + if (buf.length <= maxBytes) { + return s + } + let end = maxBytes + while (end > 0 && (buf[end] & 0xc0) === 0x80) { + end -= 1 + } + return buf.subarray(0, end).toString('utf8') +} + +// Serialize `value` to JSON and, if it exceeds `maxBytes` (UTF-8), return a +// truncated marker object whose own serialization stays within `maxBytes`. +export const truncateJsonBudget = ( + value: unknown, + maxBytes: number +): { body: unknown; truncated: boolean } => { + const text = JSON.stringify(value) + if (text === undefined) { + return { body: value, truncated: false } + } + const bytes = Buffer.byteLength(text, 'utf8') + if (bytes <= maxBytes) { + return { body: value, truncated: false } + } + const RESERVE = 256 + const preview = sliceUtf8(text, Math.max(0, maxBytes - RESERVE)) + return { + body: { _truncated: true, _originalBytes: bytes, _maxBytes: maxBytes, preview }, + truncated: true + } +} diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 937adbc..0e0b39d 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -1,5 +1,13 @@ export type { ConnectionSource, ResolvedConnection, SshConfig } from './config/baseSchema.js' export { baseConnectionShape, sshConfigSchema } from './config/baseSchema.js' +export { + clampLimit, + DEFAULT_ROW_LIMIT, + MAX_ROW_LIMIT, + sliceUtf8, + truncateJsonBudget, + truncateRows +} from './format.js' export { errorMessage, fail, ok } from './mcp/respond.js' export type { Tunnel } from './net/tunnel.js' export { buildSshAuth, expandHome, openTunnel } from './net/tunnel.js' diff --git a/packages/core/test/unit/format.test.ts b/packages/core/test/unit/format.test.ts new file mode 100644 index 0000000..dd7d59b --- /dev/null +++ b/packages/core/test/unit/format.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from 'vitest' +import { clampLimit, sliceUtf8, truncateJsonBudget, truncateRows } from '../../src/format.js' + +describe('clampLimit', () => { + it('returns fallback when undefined', () => { + expect(clampLimit(undefined)).toBe(100) + expect(clampLimit(undefined, 7)).toBe(7) + }) + it('clamps to [1, max]', () => { + expect(clampLimit(0)).toBe(1) + expect(clampLimit(99999)).toBe(1000) + expect(clampLimit(5000, 100, 2000)).toBe(2000) + }) +}) + +describe('truncateRows', () => { + it('does not truncate when under limit', () => { + expect(truncateRows([1, 2], 5)).toEqual({ rows: [1, 2], truncated: false }) + }) + it('truncates and flags when over limit', () => { + expect(truncateRows([1, 2, 3], 2)).toEqual({ rows: [1, 2], truncated: true }) + }) +}) + +describe('truncateJsonBudget', () => { + it('passes small payloads through untouched', () => { + const v = { a: 1 } + expect(truncateJsonBudget(v, 1000)).toEqual({ body: v, truncated: false }) + }) + it('marks oversized payloads truncated', () => { + const big = { s: 'x'.repeat(500) } + const out = truncateJsonBudget(big, 100) + expect(out.truncated).toBe(true) + expect((out.body as { _truncated: boolean })._truncated).toBe(true) + }) +}) + +describe('sliceUtf8', () => { + it('returns the string unchanged when it fits', () => { + expect(sliceUtf8('hello', 100)).toBe('hello') + }) + it('caps to a real UTF-8 byte budget', () => { + const out = sliceUtf8('abcdef', 3) + expect(Buffer.byteLength(out, 'utf8')).toBeLessThanOrEqual(3) + expect(out).toBe('abc') + }) + it('does not split a multi-byte code point at the boundary', () => { + const out = sliceUtf8('aéé', 3) + expect(Buffer.byteLength(out, 'utf8')).toBeLessThanOrEqual(3) + expect(out).toBe('aé') + }) +}) diff --git a/packages/dbmole-mcp/src/format.ts b/packages/dbmole-mcp/src/format.ts index 1d63cf6..7d9c872 100644 --- a/packages/dbmole-mcp/src/format.ts +++ b/packages/dbmole-mcp/src/format.ts @@ -1,17 +1,7 @@ -export const DEFAULT_ROW_LIMIT = 100 -export const MAX_ROW_LIMIT = 1000 +import { clampLimit, truncateRows } from '@dbmole/core' -export const clampRowLimit = (requested?: number): number => { - if (requested === undefined) { - return DEFAULT_ROW_LIMIT - } - return Math.max(1, Math.min(Math.trunc(requested), MAX_ROW_LIMIT)) -} - -export const truncateRows = (rows: T[], limit: number): { rows: T[]; truncated: boolean } => - rows.length > limit - ? { rows: rows.slice(0, limit), truncated: true } - : { rows, truncated: false } +export { truncateRows } +export const clampRowLimit = (requested?: number): number => clampLimit(requested) export const normalizeCell = (value: unknown): unknown => { if (typeof value === 'bigint') { diff --git a/packages/dbmole-mcp/src/tools/query.ts b/packages/dbmole-mcp/src/tools/query.ts index 3090861..bbf0267 100644 --- a/packages/dbmole-mcp/src/tools/query.ts +++ b/packages/dbmole-mcp/src/tools/query.ts @@ -1,7 +1,8 @@ +import { MAX_ROW_LIMIT } from '@dbmole/core' import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' import * as z from 'zod' import type { Manager } from '../db/manager.js' -import { clampRowLimit, MAX_ROW_LIMIT } from '../format.js' +import { clampRowLimit } from '../format.js' import { withManaged } from './managed.js' export const registerQueryTools = (server: McpServer, manager: Manager): void => { diff --git a/packages/dbmole-mcp/test/unit/format.test.ts b/packages/dbmole-mcp/test/unit/format.test.ts index fc5c4aa..68856ad 100644 --- a/packages/dbmole-mcp/test/unit/format.test.ts +++ b/packages/dbmole-mcp/test/unit/format.test.ts @@ -1,20 +1,13 @@ import { describe, expect, it } from 'vitest' -import { - clampRowLimit, - DEFAULT_ROW_LIMIT, - formatDbError, - MAX_ROW_LIMIT, - normalizeCell, - truncateRows -} from '../../src/format.js' +import { clampRowLimit, formatDbError, normalizeCell, truncateRows } from '../../src/format.js' describe('clampRowLimit', () => { it('defaults when undefined', () => { - expect(clampRowLimit(undefined)).toBe(DEFAULT_ROW_LIMIT) + expect(clampRowLimit(undefined)).toBe(100) }) it('clamps to max and min and truncates fractions', () => { - expect(clampRowLimit(99999)).toBe(MAX_ROW_LIMIT) + expect(clampRowLimit(99999)).toBe(1000) expect(clampRowLimit(0)).toBe(1) expect(clampRowLimit(-5)).toBe(1) expect(clampRowLimit(10.9)).toBe(10)