From dd6ca2deffb75756b2c781f8eb4c920c30aa4eef Mon Sep 17 00:00:00 2001 From: Evgeny Gamov Date: Tue, 16 Jun 2026 21:21:20 +0500 Subject: [PATCH] fix(esmole): truncate all tool responses; guard normalizes leading slash --- packages/esmole-mcp/src/es/guard.ts | 3 +- packages/esmole-mcp/src/tools/budget.ts | 1 + packages/esmole-mcp/src/tools/health.ts | 9 ++++- packages/esmole-mcp/src/tools/indices.ts | 12 ++++-- packages/esmole-mcp/src/tools/mappingTool.ts | 14 ++++--- packages/esmole-mcp/src/tools/request.ts | 3 +- packages/esmole-mcp/src/tools/search.ts | 3 +- .../esmole-mcp/test/unit/es/guard.test.ts | 5 +++ .../test/unit/tools/mapping.test.ts | 22 +++++++++++ .../test/unit/tools/register.test.ts | 38 +++++++++++++++++++ 10 files changed, 94 insertions(+), 16 deletions(-) create mode 100644 packages/esmole-mcp/src/tools/budget.ts diff --git a/packages/esmole-mcp/src/es/guard.ts b/packages/esmole-mcp/src/es/guard.ts index 02cac39..4187b11 100644 --- a/packages/esmole-mcp/src/es/guard.ts +++ b/packages/esmole-mcp/src/es/guard.ts @@ -23,7 +23,8 @@ const READ_DELETE_PATHS = ['_pit', '_search/scroll'] const normalize = (path: string): string => { const noQuery = path.split('?')[0] - return noQuery.endsWith('/') ? noQuery.slice(0, -1) : noQuery + const withSlash = noQuery.startsWith('/') ? noQuery : '/' + noQuery + return withSlash.endsWith('/') ? withSlash.slice(0, -1) : withSlash } const endsWithSegment = (path: string, suffix: string): boolean => diff --git a/packages/esmole-mcp/src/tools/budget.ts b/packages/esmole-mcp/src/tools/budget.ts new file mode 100644 index 0000000..25809da --- /dev/null +++ b/packages/esmole-mcp/src/tools/budget.ts @@ -0,0 +1 @@ +export const MAX_BODY_BYTES = 100_000 diff --git a/packages/esmole-mcp/src/tools/health.ts b/packages/esmole-mcp/src/tools/health.ts index df7663e..932bea3 100644 --- a/packages/esmole-mcp/src/tools/health.ts +++ b/packages/esmole-mcp/src/tools/health.ts @@ -1,8 +1,9 @@ -import { withManaged } from '@dbmole/core' +import { truncateJsonBudget, withManaged } from '@dbmole/core' import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' import * as z from 'zod' import { esFormatError } from '../esError.js' import type { EsManager } from '../manager.js' +import { MAX_BODY_BYTES } from './budget.js' export const registerHealthTool = (server: McpServer, manager: EsManager): void => { server.registerTool( @@ -16,7 +17,11 @@ export const registerHealthTool = (server: McpServer, manager: EsManager): void withManaged( manager, connection, - async (managed) => managed.backend.request('GET', '/_cluster/health'), + async (managed) => { + const res = await managed.backend.request('GET', '/_cluster/health') + const { body, truncated } = truncateJsonBudget(res.body, MAX_BODY_BYTES) + return { status: res.status, body, truncated } + }, { formatError: esFormatError } ) ) diff --git a/packages/esmole-mcp/src/tools/indices.ts b/packages/esmole-mcp/src/tools/indices.ts index a6077e0..51648a7 100644 --- a/packages/esmole-mcp/src/tools/indices.ts +++ b/packages/esmole-mcp/src/tools/indices.ts @@ -1,8 +1,9 @@ -import { withManaged } from '@dbmole/core' +import { truncateJsonBudget, withManaged } from '@dbmole/core' import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' import * as z from 'zod' import { esFormatError } from '../esError.js' import type { EsManager } from '../manager.js' +import { MAX_BODY_BYTES } from './budget.js' export const registerIndicesTool = (server: McpServer, manager: EsManager): void => { server.registerTool( @@ -16,10 +17,13 @@ export const registerIndicesTool = (server: McpServer, manager: EsManager): void withManaged( manager, connection, - async (managed) => - managed.backend.request('GET', '/_cat/indices/' + pattern, { + async (managed) => { + const res = await managed.backend.request('GET', '/_cat/indices/' + pattern, { params: { format: 'json' } - }), + }) + const { body, truncated } = truncateJsonBudget(res.body, MAX_BODY_BYTES) + return { status: res.status, body, truncated } + }, { formatError: esFormatError } ) ) diff --git a/packages/esmole-mcp/src/tools/mappingTool.ts b/packages/esmole-mcp/src/tools/mappingTool.ts index cde4d2d..07ec2d7 100644 --- a/packages/esmole-mcp/src/tools/mappingTool.ts +++ b/packages/esmole-mcp/src/tools/mappingTool.ts @@ -1,20 +1,24 @@ -import { withManaged } from '@dbmole/core' +import { truncateJsonBudget, withManaged } from '@dbmole/core' import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' import * as z from 'zod' import { flattenMapping } from '../es/mapping.js' import { esFormatError } from '../esError.js' import type { EsManaged, EsManager } from '../manager.js' +import { MAX_BODY_BYTES } from './budget.js' export const runGetMapping = async ( managed: EsManaged, index: string, - raw: boolean -): Promise<{ status: number; body?: unknown; fields?: unknown }> => { + raw: boolean, + maxBytes = MAX_BODY_BYTES +): Promise<{ status: number; body?: unknown; fields?: unknown; truncated?: boolean }> => { const res = await managed.backend.request('GET', '/' + index + '/_mapping') if (raw || res.status >= 400) { - return { status: res.status, body: res.body } + const { body, truncated } = truncateJsonBudget(res.body, maxBytes) + return { status: res.status, body, truncated } } - return { status: res.status, fields: flattenMapping(res.body) } + const { body: fields, truncated } = truncateJsonBudget(flattenMapping(res.body), maxBytes) + return { status: res.status, fields, truncated } } export const registerMappingTool = (server: McpServer, manager: EsManager): void => { diff --git a/packages/esmole-mcp/src/tools/request.ts b/packages/esmole-mcp/src/tools/request.ts index 56757cb..22b887b 100644 --- a/packages/esmole-mcp/src/tools/request.ts +++ b/packages/esmole-mcp/src/tools/request.ts @@ -4,8 +4,7 @@ import * as z from 'zod' import { assertReadAllowed, EsGuardError } from '../es/guard.js' import { esFormatError } from '../esError.js' import type { EsManaged, EsManager } from '../manager.js' - -const MAX_BODY_BYTES = 100_000 +import { MAX_BODY_BYTES } from './budget.js' export type EsToolResult = { status?: number; body?: unknown; truncated?: boolean; error?: string } diff --git a/packages/esmole-mcp/src/tools/search.ts b/packages/esmole-mcp/src/tools/search.ts index 5f6cf78..0a75610 100644 --- a/packages/esmole-mcp/src/tools/search.ts +++ b/packages/esmole-mcp/src/tools/search.ts @@ -3,8 +3,7 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' import * as z from 'zod' import { esFormatError } from '../esError.js' import type { EsManaged, EsManager } from '../manager.js' - -const MAX_BODY_BYTES = 100_000 +import { MAX_BODY_BYTES } from './budget.js' export const runEsSearch = async ( managed: EsManaged, diff --git a/packages/esmole-mcp/test/unit/es/guard.test.ts b/packages/esmole-mcp/test/unit/es/guard.test.ts index fec243d..a842212 100644 --- a/packages/esmole-mcp/test/unit/es/guard.test.ts +++ b/packages/esmole-mcp/test/unit/es/guard.test.ts @@ -45,6 +45,11 @@ describe('assertReadAllowed', () => { it('ignores a trailing slash', () => { expect(() => assertReadAllowed('POST', '/idx/_search/')).not.toThrow() }) + it('treats a missing leading slash like the slashed path', () => { + expect(() => assertReadAllowed('POST', '_search')).not.toThrow() + expect(() => assertReadAllowed('POST', 'idx/_search')).not.toThrow() + expect(() => assertReadAllowed('POST', 'idx/_doc')).toThrow(EsGuardError) + }) it('rejects an unknown method', () => { expect(() => assertReadAllowed('PATCH', '/idx/_search')).toThrow(EsGuardError) }) diff --git a/packages/esmole-mcp/test/unit/tools/mapping.test.ts b/packages/esmole-mcp/test/unit/tools/mapping.test.ts index f5e2fd2..16e3c51 100644 --- a/packages/esmole-mcp/test/unit/tools/mapping.test.ts +++ b/packages/esmole-mcp/test/unit/tools/mapping.test.ts @@ -28,4 +28,26 @@ describe('runGetMapping', () => { expect(out.body).toEqual({ error: 'not_found' }) expect(out.fields).toBeUndefined() }) + it('truncates an over-budget flattened-fields response', async () => { + const properties: Record = {} + for (let i = 0; i < 50; i++) { + properties['field_with_a_long_name_' + i] = { type: 'keyword' } + } + const out = await runGetMapping( + managed(200, { idx: { mappings: { properties } } }), + 'idx', + false, + 100 + ) + expect(out.truncated).toBe(true) + }) + it('truncates an over-budget raw response', async () => { + const big: Record = {} + for (let i = 0; i < 50; i++) { + big['key_with_a_fairly_long_name_' + i] = 'some value here ' + i + } + const out = await runGetMapping(managed(200, big), 'idx', true, 100) + expect(out.truncated).toBe(true) + expect(out.fields).toBeUndefined() + }) }) diff --git a/packages/esmole-mcp/test/unit/tools/register.test.ts b/packages/esmole-mcp/test/unit/tools/register.test.ts index 5f225a9..b969cae 100644 --- a/packages/esmole-mcp/test/unit/tools/register.test.ts +++ b/packages/esmole-mcp/test/unit/tools/register.test.ts @@ -149,4 +149,42 @@ describe('registerEsTools', () => { expect(calls[0].path).toBe('/_cluster/health') expect(isCallToolResult(res)).toBe(true) }) + + it('es_list_indices truncates an over-budget body', async () => { + const big = Array.from({ length: 8000 }, (_, i) => ({ + index: 'index_with_a_fairly_long_name_' + i, + health: 'green' + })) + const calls: RequestCall[] = [] + const { handlers, server } = fakeServer() + registerEsTools(server, fakeManager(calls, big)) + const handler = handlers.get('es_list_indices') + if (handler === undefined) { + throw new Error('es_list_indices not registered') + } + const res = await handler({ connection: 'c', pattern: '*' }) + expect(isCallToolResult(res)).toBe(true) + const payload = JSON.parse((res.content as Array<{ text: string }>)[0].text) + expect(payload.truncated).toBe(true) + expect(payload.body._truncated).toBe(true) + }) + + it('es_cluster_health truncates an over-budget body', async () => { + const big: Record = {} + for (let i = 0; i < 8000; i++) { + big['node_with_a_fairly_long_name_' + i] = 'some descriptive value ' + i + } + const calls: RequestCall[] = [] + const { handlers, server } = fakeServer() + registerEsTools(server, fakeManager(calls, big)) + const handler = handlers.get('es_cluster_health') + if (handler === undefined) { + throw new Error('es_cluster_health not registered') + } + const res = await handler({ connection: 'c' }) + expect(isCallToolResult(res)).toBe(true) + const payload = JSON.parse((res.content as Array<{ text: string }>)[0].text) + expect(payload.truncated).toBe(true) + expect(payload.body._truncated).toBe(true) + }) })