diff --git a/packages/esmole-mcp/src/esError.ts b/packages/esmole-mcp/src/esError.ts new file mode 100644 index 0000000..74b3284 --- /dev/null +++ b/packages/esmole-mcp/src/esError.ts @@ -0,0 +1,2 @@ +export const esFormatError = (_config: unknown, error: unknown): string => + '[elasticsearch] ' + (error instanceof Error ? error.message : String(error)) diff --git a/packages/esmole-mcp/src/tools/health.ts b/packages/esmole-mcp/src/tools/health.ts new file mode 100644 index 0000000..df7663e --- /dev/null +++ b/packages/esmole-mcp/src/tools/health.ts @@ -0,0 +1,23 @@ +import { 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' + +export const registerHealthTool = (server: McpServer, manager: EsManager): void => { + server.registerTool( + 'es_cluster_health', + { + description: + 'Cluster health summary (status, node and shard counts). Wraps GET /_cluster/health.', + inputSchema: { connection: z.string() } + }, + async ({ connection }) => + withManaged( + manager, + connection, + async (managed) => managed.backend.request('GET', '/_cluster/health'), + { formatError: esFormatError } + ) + ) +} diff --git a/packages/esmole-mcp/src/tools/index.ts b/packages/esmole-mcp/src/tools/index.ts new file mode 100644 index 0000000..3552336 --- /dev/null +++ b/packages/esmole-mcp/src/tools/index.ts @@ -0,0 +1,15 @@ +import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +import type { EsManager } from '../manager.js' +import { registerHealthTool } from './health.js' +import { registerIndicesTool } from './indices.js' +import { registerMappingTool } from './mappingTool.js' +import { registerRequestTool } from './request.js' +import { registerSearchTool } from './search.js' + +export const registerEsTools = (server: McpServer, manager: EsManager): void => { + registerRequestTool(server, manager) + registerSearchTool(server, manager) + registerIndicesTool(server, manager) + registerMappingTool(server, manager) + registerHealthTool(server, manager) +} diff --git a/packages/esmole-mcp/src/tools/indices.ts b/packages/esmole-mcp/src/tools/indices.ts new file mode 100644 index 0000000..a6077e0 --- /dev/null +++ b/packages/esmole-mcp/src/tools/indices.ts @@ -0,0 +1,26 @@ +import { 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' + +export const registerIndicesTool = (server: McpServer, manager: EsManager): void => { + server.registerTool( + 'es_list_indices', + { + description: + 'List indices matching a pattern with health, doc count and store size. Wraps GET /_cat/indices.', + inputSchema: { connection: z.string(), pattern: z.string().default('*') } + }, + async ({ connection, pattern }) => + withManaged( + manager, + connection, + async (managed) => + managed.backend.request('GET', '/_cat/indices/' + pattern, { + params: { format: 'json' } + }), + { formatError: esFormatError } + ) + ) +} diff --git a/packages/esmole-mcp/src/tools/mappingTool.ts b/packages/esmole-mcp/src/tools/mappingTool.ts new file mode 100644 index 0000000..cde4d2d --- /dev/null +++ b/packages/esmole-mcp/src/tools/mappingTool.ts @@ -0,0 +1,37 @@ +import { 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' + +export const runGetMapping = async ( + managed: EsManaged, + index: string, + raw: boolean +): Promise<{ status: number; body?: unknown; fields?: unknown }> => { + const res = await managed.backend.request('GET', '/' + index + '/_mapping') + if (raw || res.status >= 400) { + return { status: res.status, body: res.body } + } + return { status: res.status, fields: flattenMapping(res.body) } +} + +export const registerMappingTool = (server: McpServer, manager: EsManager): void => { + server.registerTool( + 'es_get_mapping', + { + description: + 'Return an index mapping as a flat field:type list (compact, token-cheap). Pass raw=true for the nested ES JSON.', + inputSchema: { + connection: z.string(), + index: z.string(), + raw: z.boolean().default(false) + } + }, + async ({ connection, index, raw }) => + withManaged(manager, connection, (managed) => runGetMapping(managed, index, raw), { + formatError: esFormatError + }) + ) +} diff --git a/packages/esmole-mcp/src/tools/request.ts b/packages/esmole-mcp/src/tools/request.ts new file mode 100644 index 0000000..56757cb --- /dev/null +++ b/packages/esmole-mcp/src/tools/request.ts @@ -0,0 +1,58 @@ +import { truncateJsonBudget, withManaged } from '@dbmole/core' +import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +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 + +export type EsToolResult = { status?: number; body?: unknown; truncated?: boolean; error?: string } + +export const runEsRequest = async ( + managed: EsManaged, + method: string, + path: string, + body: unknown, + params: Record | undefined, + maxBytes = MAX_BODY_BYTES +): Promise => { + if (managed.config.readonly) { + try { + assertReadAllowed(method, path) + } catch (error) { + if (error instanceof EsGuardError) { + return { error: error.message } + } + throw error + } + } + const res = await managed.backend.request(method, path, { body, params }) + const { body: capped, truncated } = truncateJsonBudget(res.body, maxBytes) + return { status: res.status, body: capped, truncated } +} + +export const registerRequestTool = (server: McpServer, manager: EsManager): void => { + server.registerTool( + 'es_request', + { + description: + 'Generic Elasticsearch REST passthrough. Call any endpoint with any method. Returns {status, body}; never throws on 4xx/5xx. ' + + 'Use a string body for NDJSON (_bulk). On readonly connections only read operations are permitted.', + inputSchema: { + connection: z.string().describe('connection name, see list_connections'), + method: z.enum(['GET', 'POST', 'PUT', 'DELETE', 'HEAD']), + path: z.string().describe("path from the ES root, e.g. '/idx/_search'"), + body: z.union([z.record(z.unknown()), z.array(z.unknown()), z.string()]).optional(), + params: z.record(z.union([z.string(), z.number(), z.boolean()])).optional() + } + }, + async ({ connection, method, path, body, params }) => + withManaged( + manager, + connection, + (managed) => runEsRequest(managed, method, path, body, params), + { formatError: esFormatError } + ) + ) +} diff --git a/packages/esmole-mcp/src/tools/search.ts b/packages/esmole-mcp/src/tools/search.ts new file mode 100644 index 0000000..5f6cf78 --- /dev/null +++ b/packages/esmole-mcp/src/tools/search.ts @@ -0,0 +1,64 @@ +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 { EsManaged, EsManager } from '../manager.js' + +const MAX_BODY_BYTES = 100_000 + +export const runEsSearch = async ( + managed: EsManaged, + index: string, + query: Record | undefined, + aggs: Record | undefined, + size: number | undefined, + source: string[] | boolean | undefined, + maxBytes = MAX_BODY_BYTES +): Promise<{ status: number; body: unknown; truncated: boolean }> => { + const aggsOnly = aggs !== undefined && query === undefined + const searchBody: Record = {} + if (query !== undefined) { + searchBody.query = query + } + if (aggs !== undefined) { + searchBody.aggs = aggs + } + searchBody.size = aggsOnly ? 0 : (size ?? 10) + if (source !== undefined) { + searchBody._source = source + } + const res = await managed.backend.request('POST', '/' + index + '/_search', { + body: searchBody + }) + const { body, truncated } = truncateJsonBudget(res.body, maxBytes) + return { status: res.status, body, truncated } +} + +export const registerSearchTool = (server: McpServer, manager: EsManager): void => { + server.registerTool( + 'es_search', + { + description: + 'Search an index. Wraps POST /{index}/_search. Provide query and/or aggs. With aggs and no query, size is forced to 0 (aggregation-only). ' + + 'Use _source to project only needed fields and keep responses small. Returns {status, body, truncated}.', + inputSchema: { + connection: z.string(), + index: z.string(), + query: z.record(z.unknown()).optional(), + aggs: z.record(z.unknown()).optional(), + size: z.number().int().min(0).max(1000).optional(), + source: z + .union([z.array(z.string()), z.boolean()]) + .optional() + .describe('_source projection: list of fields or false') + } + }, + async ({ connection, index, query, aggs, size, source }) => + withManaged( + manager, + connection, + (managed) => runEsSearch(managed, index, query, aggs, size, source), + { formatError: esFormatError } + ) + ) +} diff --git a/packages/esmole-mcp/test/unit/tools/esError.test.ts b/packages/esmole-mcp/test/unit/tools/esError.test.ts new file mode 100644 index 0000000..ace7cad --- /dev/null +++ b/packages/esmole-mcp/test/unit/tools/esError.test.ts @@ -0,0 +1,11 @@ +import { describe, expect, it } from 'vitest' +import { esFormatError } from '../../../src/esError.js' + +describe('esFormatError', () => { + it('prefixes an Error message', () => { + expect(esFormatError({}, new Error('boom'))).toBe('[elasticsearch] boom') + }) + it('stringifies a non-Error value', () => { + expect(esFormatError(undefined, 'plain')).toBe('[elasticsearch] plain') + }) +}) diff --git a/packages/esmole-mcp/test/unit/tools/mapping.test.ts b/packages/esmole-mcp/test/unit/tools/mapping.test.ts new file mode 100644 index 0000000..f5e2fd2 --- /dev/null +++ b/packages/esmole-mcp/test/unit/tools/mapping.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, it } from 'vitest' +import { runGetMapping } from '../../../src/tools/mappingTool.js' + +const managed = (status: number, body: unknown) => + ({ + config: { readonly: false }, + backend: { request: async () => ({ status, body }), dispose: async () => {} } + }) as never + +describe('runGetMapping', () => { + it('flattens by default', async () => { + const out = await runGetMapping( + managed(200, { idx: { mappings: { properties: { y: { type: 'long' } } } } }), + 'idx', + false + ) + expect(out.fields).toContainEqual({ field: 'y', type: 'long' }) + }) + it('returns raw body when raw=true', async () => { + const raw = { idx: { mappings: { properties: { y: { type: 'long' } } } } } + const out = await runGetMapping(managed(200, raw), 'idx', true) + expect(out.body).toEqual(raw) + expect(out.fields).toBeUndefined() + }) + it('passes the raw body through on a ≥400 status', async () => { + const out = await runGetMapping(managed(404, { error: 'not_found' }), 'idx', false) + expect(out.status).toBe(404) + expect(out.body).toEqual({ error: 'not_found' }) + 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 new file mode 100644 index 0000000..5f225a9 --- /dev/null +++ b/packages/esmole-mcp/test/unit/tools/register.test.ts @@ -0,0 +1,152 @@ +import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js' +import { describe, expect, it } from 'vitest' +import type { EsBackend } from '../../../src/es/backend.js' +import type { EsManaged, EsManager } from '../../../src/manager.js' +import { registerEsTools } from '../../../src/tools/index.js' + +type RequestCall = { + method: string + path: string + opts?: { body?: unknown; params?: Record } +} + +type ToolHandler = (args: Record) => Promise + +type RecordingServer = { + handlers: Map + server: McpServer +} + +// Only registerTool is exercised; expose just that method and cast once to +// McpServer (AGENTS.md allows a single cast for the server/manager shape — +// no `as unknown as`). +const fakeServer = (): RecordingServer => { + const handlers = new Map() + const recorder = { + registerTool: (name: string, _config: unknown, handler: ToolHandler): void => { + handlers.set(name, handler) + } + } + return { handlers, server: recorder as McpServer } +} + +const fakeManager = (calls: RequestCall[], body: unknown): EsManager => { + const backend: EsBackend = { + request: async (method, path, opts) => { + calls.push({ method, path, opts }) + return { status: 200, body } + }, + dispose: async () => {} + } + const managed: EsManaged = { + backend, + config: { name: 'c', readonly: false } as EsManaged['config'], + source: 'test' + } + return { + get: async () => managed, + invalidate: async () => {}, + disposeAll: async () => {} + } +} + +const isCallToolResult = (res: CallToolResult): boolean => Array.isArray(res.content) + +describe('registerEsTools', () => { + it('registers all five tools', () => { + const { handlers, server } = fakeServer() + registerEsTools(server, fakeManager([], {})) + expect([...handlers.keys()].sort()).toEqual( + [ + 'es_cluster_health', + 'es_get_mapping', + 'es_list_indices', + 'es_request', + 'es_search' + ].sort() + ) + }) + + it('es_request handler calls backend with method/path and returns a CallToolResult', async () => { + const calls: RequestCall[] = [] + const { handlers, server } = fakeServer() + registerEsTools(server, fakeManager(calls, { ok: true })) + const handler = handlers.get('es_request') + if (handler === undefined) { + throw new Error('es_request not registered') + } + const res = await handler({ + connection: 'c', + method: 'POST', + path: '/idx/_doc', + body: { a: 1 } + }) + expect(calls[0].method).toBe('POST') + expect(calls[0].path).toBe('/idx/_doc') + expect(isCallToolResult(res)).toBe(true) + expect(res.isError).toBeFalsy() + }) + + it('es_search handler posts /{index}/_search', async () => { + const calls: RequestCall[] = [] + const { handlers, server } = fakeServer() + registerEsTools(server, fakeManager(calls, { hits: {} })) + const handler = handlers.get('es_search') + if (handler === undefined) { + throw new Error('es_search not registered') + } + const res = await handler({ connection: 'c', index: 'books', query: { match_all: {} } }) + expect(calls[0].method).toBe('POST') + expect(calls[0].path).toBe('/books/_search') + expect(isCallToolResult(res)).toBe(true) + }) + + it('es_list_indices handler GETs /_cat/indices/ with format json', async () => { + const calls: RequestCall[] = [] + const { handlers, server } = fakeServer() + registerEsTools(server, fakeManager(calls, [])) + 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: 'logs-*' }) + expect(calls[0].method).toBe('GET') + expect(calls[0].path).toBe('/_cat/indices/logs-*') + expect(calls[0].opts?.params).toEqual({ format: 'json' }) + expect(isCallToolResult(res)).toBe(true) + }) + + it('es_get_mapping handler GETs /{index}/_mapping and flattens by default', async () => { + const calls: RequestCall[] = [] + const { handlers, server } = fakeServer() + registerEsTools( + server, + fakeManager(calls, { idx: { mappings: { properties: { y: { type: 'long' } } } } }) + ) + const handler = handlers.get('es_get_mapping') + if (handler === undefined) { + throw new Error('es_get_mapping not registered') + } + const res = await handler({ connection: 'c', index: 'idx', raw: false }) + expect(calls[0].method).toBe('GET') + expect(calls[0].path).toBe('/idx/_mapping') + expect(isCallToolResult(res)).toBe(true) + const payload = JSON.parse((res.content as Array<{ text: string }>)[0].text) + expect(payload.fields).toContainEqual({ field: 'y', type: 'long' }) + }) + + it('es_cluster_health handler GETs /_cluster/health', async () => { + const calls: RequestCall[] = [] + const { handlers, server } = fakeServer() + registerEsTools(server, fakeManager(calls, { status: 'green' })) + 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(calls[0].method).toBe('GET') + expect(calls[0].path).toBe('/_cluster/health') + expect(isCallToolResult(res)).toBe(true) + }) +}) diff --git a/packages/esmole-mcp/test/unit/tools/request.test.ts b/packages/esmole-mcp/test/unit/tools/request.test.ts new file mode 100644 index 0000000..de0554a --- /dev/null +++ b/packages/esmole-mcp/test/unit/tools/request.test.ts @@ -0,0 +1,95 @@ +import { describe, expect, it, vi } from 'vitest' +import { runEsRequest } from '../../../src/tools/request.js' + +const fakeManaged = (readonly: boolean, calls: unknown[]) => + ({ + config: { name: 'c', readonly, type: 'elasticsearch' }, + backend: { + request: async (method: string, path: string, opts?: unknown) => { + calls.push({ method, path, opts }) + return { status: 200, body: { hits: { total: 1 } } } + }, + dispose: async () => {} + } + }) as never + +describe('runEsRequest', () => { + it('blocks a mutating method on a readonly connection before calling backend', async () => { + const calls: unknown[] = [] + const res = await runEsRequest( + fakeManaged(true, calls), + 'DELETE', + '/idx/_doc/1', + undefined, + undefined, + 1_000_000 + ) + expect(res.status).toBeUndefined() + expect(res.error).toMatch(/readonly/) + expect(calls).toHaveLength(0) + }) + it('allows a read POST on a readonly connection', async () => { + const calls: unknown[] = [] + const res = await runEsRequest( + fakeManaged(true, calls), + 'POST', + '/idx/_search', + { query: {} }, + undefined, + 1_000_000 + ) + expect(res.status).toBe(200) + expect(calls).toHaveLength(1) + }) + it('passes through and returns status+body when allowed', async () => { + const calls: unknown[] = [] + const res = await runEsRequest( + fakeManaged(false, calls), + 'POST', + '/idx/_doc', + { a: 1 }, + undefined, + 1_000_000 + ) + expect(res.status).toBe(200) + expect(calls).toHaveLength(1) + }) + it('truncates an over-budget body', async () => { + const calls: unknown[] = [] + const managed = fakeManaged(false, calls) as { backend: { request: unknown } } + managed.backend.request = async () => ({ status: 200, body: { big: 'x'.repeat(5000) } }) + const res = await runEsRequest( + managed as never, + 'GET', + '/idx/_search', + undefined, + undefined, + 100 + ) + expect(res.truncated).toBe(true) + }) + it('rethrows a non-guard error from the readonly check', async () => { + const calls: unknown[] = [] + vi.resetModules() + vi.doMock('../../../src/es/guard.js', () => ({ + EsGuardError: class EsGuardError extends Error {}, + assertReadAllowed: () => { + throw new Error('boom') + } + })) + const { runEsRequest: mocked } = await import('../../../src/tools/request.js') + await expect( + mocked( + fakeManaged(true, calls), + 'POST', + '/idx/_search', + undefined, + undefined, + 1_000_000 + ) + ).rejects.toThrow('boom') + expect(calls).toHaveLength(0) + vi.doUnmock('../../../src/es/guard.js') + vi.resetModules() + }) +}) diff --git a/packages/esmole-mcp/test/unit/tools/search.test.ts b/packages/esmole-mcp/test/unit/tools/search.test.ts new file mode 100644 index 0000000..b751a3c --- /dev/null +++ b/packages/esmole-mcp/test/unit/tools/search.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, it } from 'vitest' +import { runEsSearch } from '../../../src/tools/search.js' + +const fakeManaged = (calls: unknown[]) => + ({ + config: { name: 'c', readonly: false, type: 'elasticsearch' }, + backend: { + request: async (method: string, path: string, opts?: { body?: unknown }) => { + calls.push({ method, path, body: opts?.body }) + return { status: 200, body: { hits: { total: { value: 0 } }, aggregations: {} } } + }, + dispose: async () => {} + } + }) as never + +describe('runEsSearch', () => { + it('posts query + size and forwards _source projection', async () => { + const calls: { path: string; body: { query: unknown; size: number; _source: unknown } }[] = + [] + await runEsSearch(fakeManaged(calls), 'books', { match: { t: 'x' } }, undefined, 5, [ + 'title' + ]) + expect(calls[0].path).toBe('/books/_search') + expect(calls[0].body.query).toEqual({ match: { t: 'x' } }) + expect(calls[0].body.size).toBe(5) + expect(calls[0].body._source).toEqual(['title']) + }) + it('forces size 0 in aggs-only mode (aggs present, no query)', async () => { + const calls: { body: { size: number; aggs: unknown; query?: unknown } }[] = [] + await runEsSearch( + fakeManaged(calls), + 'books', + undefined, + { byYear: { terms: { field: 'year' } } }, + undefined, + undefined + ) + expect(calls[0].body.size).toBe(0) + expect(calls[0].body.aggs).toEqual({ byYear: { terms: { field: 'year' } } }) + expect(calls[0].body.query).toBeUndefined() + }) + it('defaults size to 10 when query given without size', async () => { + const calls: { body: { size: number } }[] = [] + await runEsSearch( + fakeManaged(calls), + 'books', + { match_all: {} }, + undefined, + undefined, + undefined + ) + expect(calls[0].body.size).toBe(10) + }) +})