From 2bf2936561c95525e9925e722a5f11bcb7e68c86 Mon Sep 17 00:00:00 2001 From: Eugene Date: Fri, 19 Jun 2026 16:01:50 -0400 Subject: [PATCH 1/3] fix: driver-aware MCP toolset so MongoDB connections work (#180) The MCP server always registered SQL-shaped tools (execute_query, execute_write, ...) regardless of the connected driver. On a MongoDB connection the SQL keyword gate and the MQL driver imposed mutually exclusive requirements, so no query could return a single row. buildMcpServer() now branches on getDriver().queryMode and registers one of two mutually-exclusive tool sets: - sql -> execute_query / execute_write / list_tables / describe_table - mql -> find_documents / aggregate_documents / count_documents / list_collections / describe_collection (reads) and insert_document / update_document / delete_document (writes, gated by "Allow MCP to modify data") The Mongo tools expose MQL fields as typed params and assemble the JSON-encoded request the driver already parses, so the model never hand-authors MQL JSON. Server instructions are now mode-aware and name the connected DB type, so the agent knows upfront whether to use SQL or MongoDB tools and never attempts the wrong query language. Adds server/src/mcp.test.ts covering toolset selection per mode, mode-aware instructions, MQL request assembly, and write gating. Co-Authored-By: Claude Opus 4.8 (1M context) --- server/src/mcp.test.ts | 156 +++++++++++++++++ server/src/mcp.ts | 375 +++++++++++++++++++++++++++++++++++++++-- 2 files changed, 518 insertions(+), 13 deletions(-) create mode 100644 server/src/mcp.test.ts diff --git a/server/src/mcp.test.ts b/server/src/mcp.test.ts new file mode 100644 index 0000000..2499f6b --- /dev/null +++ b/server/src/mcp.test.ts @@ -0,0 +1,156 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +vi.mock('./db.js', () => ({ + getDriver: vi.fn(), + getActiveConfig: vi.fn(), + isConnected: vi.fn(), +})); +vi.mock('./mcp-state.js', () => ({ + isMcpWritesAllowed: vi.fn(), +})); + +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js'; +import { getDriver, getActiveConfig, isConnected } from './db.js'; +import { isMcpWritesAllowed } from './mcp-state.js'; +import { buildMcpServer } from './mcp.js'; + +async function connectClient() { + const server = buildMcpServer(); + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const client = new Client({ name: 'test', version: '0' }); + await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]); + return { client, server }; +} + +function mockMongoDriver(query = vi.fn()) { + const driver = { + queryMode: 'mql' as const, + query, + getSchemas: vi.fn().mockResolvedValue(['shop']), + getSchema: vi.fn().mockResolvedValue({ tables: [{ name: 'users', rows: 3 }], views: [] }), + getTable: vi.fn().mockResolvedValue({ name: 'users', rows: 3, columns: [] }), + getCollectionInfo: vi.fn().mockResolvedValue({ validator: null, indexes: [] }), + }; + vi.mocked(getDriver).mockReturnValue(driver as any); + return driver; +} + +describe('MCP server – driver-aware toolset', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(isConnected).mockReturnValue(true); + vi.mocked(getActiveConfig).mockReturnValue({ type: 'mongodb', database: 'shop' } as any); + vi.mocked(isMcpWritesAllowed).mockReturnValue(false); + }); + + it('exposes MongoDB tools (not SQL tools) on a Mongo connection', async () => { + mockMongoDriver(); + const { client } = await connectClient(); + const names = (await client.listTools()).tools.map(t => t.name).sort(); + + expect(names).toContain('find_documents'); + expect(names).toContain('aggregate_documents'); + expect(names).toContain('count_documents'); + expect(names).toContain('list_collections'); + expect(names).not.toContain('execute_query'); + expect(names).not.toContain('execute_write'); + }); + + it('exposes SQL tools on a SQL connection', async () => { + vi.mocked(getActiveConfig).mockReturnValue({ type: 'postgres', database: 'shop' } as any); + vi.mocked(getDriver).mockReturnValue({ queryMode: 'sql' } as any); + const { client } = await connectClient(); + const names = (await client.listTools()).tools.map(t => t.name).sort(); + + expect(names).toContain('execute_query'); + expect(names).toContain('execute_write'); + expect(names).not.toContain('find_documents'); + }); + + it('server instructions tell the model it is MongoDB and not to write SQL', async () => { + mockMongoDriver(); + const { client } = await connectClient(); + const instructions = client.getInstructions() ?? ''; + + expect(instructions).toMatch(/mongodb/i); + expect(instructions).toMatch(/do not write sql/i); + }); + + it('find_documents assembles a JSON-encoded MQL request for the driver', async () => { + const query = vi.fn().mockResolvedValue({ + rows: [{ _id: 'a1', email: 'x@y.z' }], + columnMeta: [{ name: '_id' }, { name: 'email' }], + }); + mockMongoDriver(query); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'find_documents', + arguments: { collection: 'users', filter: { email: 'x@y.z' }, schema: 'shop' }, + }); + + expect(res.isError).toBeFalsy(); + const [sql, params, schema] = query.mock.calls[0]; + expect(JSON.parse(sql)).toMatchObject({ + collection: 'users', + operation: 'find', + filter: { email: 'x@y.z' }, + limit: 100, + }); + expect(params).toEqual([]); + expect(schema).toBe('shop'); + expect(JSON.parse(res.content[0].text).documents).toEqual([{ _id: 'a1', email: 'x@y.z' }]); + }); + + it('blocks insert_document when writes are disabled', async () => { + const query = vi.fn(); + mockMongoDriver(query); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'insert_document', + arguments: { collection: 'users', document: { email: 'x@y.z' } }, + }); + + expect(res.isError).toBe(true); + expect(res.content[0].text).toMatch(/Allow MCP to modify data/); + expect(query).not.toHaveBeenCalled(); + }); + + it('allows insert_document when writes are enabled', async () => { + vi.mocked(isMcpWritesAllowed).mockReturnValue(true); + const query = vi.fn().mockResolvedValue({ rows: [], columnMeta: [], affectedRows: 1 }); + mockMongoDriver(query); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'insert_document', + arguments: { collection: 'users', document: { email: 'x@y.z' } }, + }); + + expect(res.isError).toBeFalsy(); + expect(JSON.parse(query.mock.calls[0][0])).toMatchObject({ + collection: 'users', + operation: 'insertOne', + document: { email: 'x@y.z' }, + }); + expect(JSON.parse(res.content[0].text).insertedCount).toBe(1); + }); + + it('update_document requires a filter or id', async () => { + vi.mocked(isMcpWritesAllowed).mockReturnValue(true); + const query = vi.fn(); + mockMongoDriver(query); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'update_document', + arguments: { collection: 'users', update: { $set: { active: true } } }, + }); + + expect(res.isError).toBe(true); + expect(res.content[0].text).toMatch(/requires either "filter" or "id"/); + expect(query).not.toHaveBeenCalled(); + }); +}); diff --git a/server/src/mcp.ts b/server/src/mcp.ts index a6e54cf..4b3b52f 100644 --- a/server/src/mcp.ts +++ b/server/src/mcp.ts @@ -4,6 +4,7 @@ import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/ import { z } from 'zod'; import { getDriver, getActiveConfig, isConnected } from './db.js'; import { isMcpWritesAllowed } from './mcp-state.js'; +import type { QueryResult } from './drivers/interface.js'; const DEFAULT_ROW_LIMIT = 100; const MAX_ROW_LIMIT = 10_000; @@ -32,20 +33,11 @@ function requireConnection(): string | null { return null; } -export function buildMcpServer(): McpServer { - const server = new McpServer( - { name: 'helix-mcp', version: '0.1.0' }, - { - capabilities: { tools: {} }, - instructions: [ - 'Helix MCP exposes the database currently connected in the Helix UI.', - 'Reads are always allowed. Writes (INSERT/UPDATE/DELETE/REPLACE) are only allowed', - 'when the user has explicitly enabled them in the Helix UI top-right menu.', - 'DDL (CREATE/DROP/ALTER/TRUNCATE/RENAME) is not supported in this version.', - ].join(' '), - }, - ); +// --------------------------------------------------------------------------- +// SQL tools (MySQL / Postgres — queryMode 'sql') +// --------------------------------------------------------------------------- +function registerSqlTools(server: McpServer): void { server.registerTool( 'list_tables', { @@ -202,6 +194,363 @@ export function buildMcpServer(): McpServer { } }, ); +} + +// --------------------------------------------------------------------------- +// MongoDB tools (queryMode 'mql') +// +// The Mongo driver speaks a JSON-encoded MQL request (see drivers/mongodb.ts). +// Rather than make the model hand-author that JSON, each tool below exposes the +// individual MQL fields as typed parameters and assembles the request itself. +// --------------------------------------------------------------------------- + +const docSchema = z.record(z.string(), z.unknown()); + +/** Build the JSON-encoded MQL request the Mongo driver's query() expects and run it. */ +function runMql(req: Record, schema?: string): Promise { + return getDriver().query(JSON.stringify(req), [], schema); +} + +function registerMongoTools(server: McpServer): void { + server.registerTool( + 'list_collections', + { + description: + 'List collections (and views) in a MongoDB database. ' + + 'If no database is given, lists the available databases.', + inputSchema: { + schema: z.string().optional().describe('Database name. Omit to list available databases.'), + }, + }, + async ({ schema }) => { + const err = requireConnection(); + if (err) return toolError(err); + + try { + const driver = getDriver(); + if (!schema) { + const schemas = await driver.getSchemas(); + return toolJson({ databases: schemas }); + } + const info = await driver.getSchema(schema); + return toolJson({ + database: schema, + collections: info.tables.map(t => ({ name: t.name, approxDocs: t.rows })), + views: info.views, + }); + } catch (e) { + return toolError(e instanceof Error ? e.message : String(e)); + } + }, + ); + + server.registerTool( + 'describe_collection', + { + description: + 'Describe a MongoDB collection: field names and types inferred from a sample of documents ' + + '(schemaless — fields may vary per document), plus its indexes and JSON Schema validator if set.', + inputSchema: { + schema: z.string().describe('Database name.'), + collection: z.string().describe('Collection name.'), + }, + }, + async ({ schema, collection }) => { + const err = requireConnection(); + if (err) return toolError(err); + + try { + const driver = getDriver(); + const info = await driver.getTable(schema, collection); + if (!info) { + return toolError(`Collection "${schema}"."${collection}" not found.`); + } + const collInfo = driver.getCollectionInfo + ? await driver.getCollectionInfo(schema, collection) + : null; + return toolJson({ + database: schema, + collection, + approxDocs: info.rows, + fields: info.columns, + indexes: collInfo?.indexes ?? [], + validator: collInfo?.validator ?? null, + }); + } catch (e) { + return toolError(e instanceof Error ? e.message : String(e)); + } + }, + ); + + server.registerTool( + 'find_documents', + { + description: + 'Find documents in a MongoDB collection (db.collection.find). ' + + `Returns up to ${DEFAULT_ROW_LIMIT} documents by default; pass "limit" (up to ${MAX_ROW_LIMIT}) to change. ` + + 'Note: a string "_id" filter is matched literally and will not match an ObjectId — omit it or filter on other fields.', + inputSchema: { + collection: z.string().describe('Collection name.'), + filter: docSchema.optional().describe('MongoDB query filter, e.g. {"status": "active"}. Omit for all documents.'), + projection: docSchema.optional().describe('Fields to include/exclude, e.g. {"email": 1, "_id": 0}.'), + sort: docSchema.optional().describe('Sort spec, e.g. {"createdAt": -1}.'), + limit: z.number().int().positive().max(MAX_ROW_LIMIT).optional() + .describe(`Max documents returned (default ${DEFAULT_ROW_LIMIT}, max ${MAX_ROW_LIMIT}).`), + skip: z.number().int().nonnegative().optional().describe('Number of documents to skip.'), + schema: z.string().optional().describe('Database name. Omit to use the connected database.'), + }, + }, + async ({ collection, filter, projection, sort, limit, skip, schema }) => { + const err = requireConnection(); + if (err) return toolError(err); + + const cap = limit ?? DEFAULT_ROW_LIMIT; + try { + const start = Date.now(); + const result = await runMql( + { collection, operation: 'find', filter, projection, sort, limit: cap, skip }, + schema, + ); + const executionTime = Date.now() - start; + return toolJson({ + fields: result.columnMeta.map(c => c.name), + documents: result.rows, + docCount: result.rows.length, + limitApplied: cap, + executionTime, + }); + } catch (e) { + return toolError(e instanceof Error ? e.message : String(e)); + } + }, + ); + + server.registerTool( + 'aggregate_documents', + { + description: + 'Run an aggregation pipeline on a MongoDB collection (db.collection.aggregate). ' + + `Results are capped at ${DEFAULT_ROW_LIMIT} documents by default; pass "limit" (up to ${MAX_ROW_LIMIT}) to change. ` + + 'For large pipelines, add an explicit {"$limit": N} stage to bound work server-side.', + inputSchema: { + collection: z.string().describe('Collection name.'), + pipeline: z.array(docSchema).describe('Aggregation pipeline stages, e.g. [{"$match": {...}}, {"$group": {...}}].'), + limit: z.number().int().positive().max(MAX_ROW_LIMIT).optional() + .describe(`Max documents returned (default ${DEFAULT_ROW_LIMIT}, max ${MAX_ROW_LIMIT}).`), + schema: z.string().optional().describe('Database name. Omit to use the connected database.'), + }, + }, + async ({ collection, pipeline, limit, schema }) => { + const err = requireConnection(); + if (err) return toolError(err); + + const cap = limit ?? DEFAULT_ROW_LIMIT; + try { + const start = Date.now(); + const result = await runMql({ collection, operation: 'aggregate', pipeline }, schema); + const executionTime = Date.now() - start; + const totalRows = result.rows.length; + const capped = result.rows.slice(0, cap); + return toolJson({ + fields: result.columnMeta.map(c => c.name), + documents: capped, + docCount: capped.length, + totalDocs: totalRows, + truncated: totalRows > capped.length, + limitApplied: cap, + executionTime, + }); + } catch (e) { + return toolError(e instanceof Error ? e.message : String(e)); + } + }, + ); + + server.registerTool( + 'count_documents', + { + description: 'Count documents matching a filter in a MongoDB collection (db.collection.countDocuments).', + inputSchema: { + collection: z.string().describe('Collection name.'), + filter: docSchema.optional().describe('MongoDB query filter. Omit to count all documents.'), + schema: z.string().optional().describe('Database name. Omit to use the connected database.'), + }, + }, + async ({ collection, filter, schema }) => { + const err = requireConnection(); + if (err) return toolError(err); + + try { + const start = Date.now(); + const result = await runMql({ collection, operation: 'count', filter }, schema); + const executionTime = Date.now() - start; + return toolJson({ count: result.rows[0]?.count ?? 0, executionTime }); + } catch (e) { + return toolError(e instanceof Error ? e.message : String(e)); + } + }, + ); + + server.registerTool( + 'insert_document', + { + description: + 'Insert a single document into a MongoDB collection (db.collection.insertOne). ' + + 'Requires the user to have enabled "Allow MCP to modify data" in the Helix UI.', + inputSchema: { + collection: z.string().describe('Collection name.'), + document: docSchema.describe('The document to insert.'), + schema: z.string().optional().describe('Database name. Omit to use the connected database.'), + }, + }, + async ({ collection, document, schema }) => { + const err = requireConnection(); + if (err) return toolError(err); + if (!isMcpWritesAllowed()) { + return toolError( + 'Writes are disabled. Ask the user to enable "Allow MCP to modify data" in the Helix UI (top-right menu).', + ); + } + + try { + const start = Date.now(); + const result = await runMql({ collection, operation: 'insertOne', document }, schema); + const executionTime = Date.now() - start; + return toolJson({ insertedCount: result.affectedRows ?? 0, executionTime }); + } catch (e) { + return toolError(e instanceof Error ? e.message : String(e)); + } + }, + ); + + server.registerTool( + 'update_document', + { + description: + 'Update a single document in a MongoDB collection (db.collection.updateOne). ' + + 'Provide either "filter" or "id" to select the document, and "update" with operators like {"$set": {...}}. ' + + 'Requires the user to have enabled "Allow MCP to modify data" in the Helix UI.', + inputSchema: { + collection: z.string().describe('Collection name.'), + update: docSchema.describe('Update document with operators, e.g. {"$set": {"status": "active"}}.'), + filter: docSchema.optional().describe('Filter selecting the document to update.'), + id: z.string().optional().describe('Convenience selector by _id (24-hex strings are matched as ObjectId, falling back to the raw value).'), + schema: z.string().optional().describe('Database name. Omit to use the connected database.'), + }, + }, + async ({ collection, update, filter, id, schema }) => { + const err = requireConnection(); + if (err) return toolError(err); + if (!isMcpWritesAllowed()) { + return toolError( + 'Writes are disabled. Ask the user to enable "Allow MCP to modify data" in the Helix UI (top-right menu).', + ); + } + if (filter === undefined && id === undefined) { + return toolError('update_document requires either "filter" or "id" to select the document.'); + } + + try { + const start = Date.now(); + const result = await runMql({ collection, operation: 'updateOne', update, filter, id }, schema); + const executionTime = Date.now() - start; + return toolJson({ matchedCount: result.affectedRows ?? 0, executionTime }); + } catch (e) { + return toolError(e instanceof Error ? e.message : String(e)); + } + }, + ); + + server.registerTool( + 'delete_document', + { + description: + 'Delete a single document from a MongoDB collection (db.collection.deleteOne). ' + + 'Provide either "filter" or "id" to select the document. ' + + 'Requires the user to have enabled "Allow MCP to modify data" in the Helix UI.', + inputSchema: { + collection: z.string().describe('Collection name.'), + filter: docSchema.optional().describe('Filter selecting the document to delete.'), + id: z.string().optional().describe('Convenience selector by _id (24-hex strings are matched as ObjectId, falling back to the raw value).'), + schema: z.string().optional().describe('Database name. Omit to use the connected database.'), + }, + }, + async ({ collection, filter, id, schema }) => { + const err = requireConnection(); + if (err) return toolError(err); + if (!isMcpWritesAllowed()) { + return toolError( + 'Writes are disabled. Ask the user to enable "Allow MCP to modify data" in the Helix UI (top-right menu).', + ); + } + if (filter === undefined && id === undefined) { + return toolError('delete_document requires either "filter" or "id" to select the document.'); + } + + try { + const start = Date.now(); + const result = await runMql({ collection, operation: 'deleteOne', filter, id }, schema); + const executionTime = Date.now() - start; + return toolJson({ deletedCount: result.affectedRows ?? 0, executionTime }); + } catch (e) { + return toolError(e instanceof Error ? e.message : String(e)); + } + }, + ); +} + +// --------------------------------------------------------------------------- +// Server assembly +// --------------------------------------------------------------------------- + +/** + * The active query language for the connected database. Defaults to 'sql' when + * nothing is connected so the (backward-compatible) SQL toolset is advertised. + * The transport is stateless, so `tools/list` is re-evaluated on every request + * and always reflects the database connected at that moment. + */ +function activeQueryMode(): 'sql' | 'mql' { + return isConnected() ? getDriver().queryMode : 'sql'; +} + +function buildInstructions(mode: 'sql' | 'mql'): string { + const cfg = getActiveConfig(); + const dbType = cfg?.type ?? 'none'; + const base = [ + 'Helix MCP exposes the database currently connected in the Helix UI.', + `The connected database is ${isConnected() ? `a ${dbType} database` : 'not connected yet'}.`, + ]; + if (mode === 'mql') { + base.push( + 'This is a MongoDB (document) database — use the MongoDB tools (find_documents, aggregate_documents, count_documents, ...).', + 'Do NOT write SQL: there are no SELECT/SHOW statements here, only collections and documents.', + 'Writes (insert/update/delete) are only allowed when the user has enabled "Allow MCP to modify data" in the Helix UI top-right menu.', + ); + } else { + base.push( + 'This is a SQL database — use SQL via execute_query (reads) and execute_write (writes).', + 'Writes (INSERT/UPDATE/DELETE/REPLACE) are only allowed when the user has enabled "Allow MCP to modify data" in the Helix UI top-right menu.', + 'DDL (CREATE/DROP/ALTER/TRUNCATE/RENAME) is not supported in this version.', + ); + } + return base.join(' '); +} + +export function buildMcpServer(): McpServer { + const mode = activeQueryMode(); + const server = new McpServer( + { name: 'helix-mcp', version: '0.1.0' }, + { + capabilities: { tools: {} }, + instructions: buildInstructions(mode), + }, + ); + + if (mode === 'mql') { + registerMongoTools(server); + } else { + registerSqlTools(server); + } return server; } From 8c890da0adb341aca5286245e23599e9fb334429 Mon Sep 17 00:00:00 2001 From: Eugene Date: Fri, 19 Jun 2026 16:10:47 -0400 Subject: [PATCH 2/3] fix: invariant MCP toolset + per-call mode guards for DB switches (#180) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit registered tools per queryMode. That left a gap: a client that connected before a database was selected (or while a different DB type was connected) caches tools/list, and the stateless transport can't push tools/list_changed — so after the user connects/switches to MongoDB the agent still holds the SQL tools and calling execute_query hits a server with no such tool ("tool not found"), with no hint why. Flip the design so the tool list is invariant across connections, which removes the stale-cache problem entirely: - Always register both the SQL and MongoDB tool sets plus a new connection_info tool. tools/list never changes with the DB type, so a cached list is never wrong. - Each domain tool validates the live queryMode at call time (requireMode). A tool from the wrong family returns an actionable error naming the connected DB type and the tools that apply, and telling the agent to re-check connection_info / re-list / reconnect. - connection_info reports the connected DB type, query language, write state, and recommended tools so the agent can orient up front and detect a mid-session switch. - Instructions are generic but name the currently-connected DB and point at connection_info. Updates mcp.test.ts: asserts the invariant toolset, connection_info output, and that cross-family calls return actionable mismatch errors instead of "tool not found". Co-Authored-By: Claude Opus 4.8 (1M context) --- server/src/mcp.test.ts | 88 +++++++++++++++----- server/src/mcp.ts | 180 +++++++++++++++++++++++++++++++---------- 2 files changed, 207 insertions(+), 61 deletions(-) diff --git a/server/src/mcp.test.ts b/server/src/mcp.test.ts index 2499f6b..659c9d4 100644 --- a/server/src/mcp.test.ts +++ b/server/src/mcp.test.ts @@ -36,7 +36,13 @@ function mockMongoDriver(query = vi.fn()) { return driver; } -describe('MCP server – driver-aware toolset', () => { +function mockSqlDriver(query = vi.fn()) { + const driver = { queryMode: 'sql' as const, query }; + vi.mocked(getDriver).mockReturnValue(driver as any); + return driver; +} + +describe('MCP server – invariant toolset with per-call mode guards', () => { beforeEach(() => { vi.clearAllMocks(); vi.mocked(isConnected).mockReturnValue(true); @@ -44,37 +50,81 @@ describe('MCP server – driver-aware toolset', () => { vi.mocked(isMcpWritesAllowed).mockReturnValue(false); }); - it('exposes MongoDB tools (not SQL tools) on a Mongo connection', async () => { + it('advertises the same full toolset regardless of connected DB (no stale-cache problem)', async () => { + // Mongo connection + mockMongoDriver(); + const mongoNames = (await connectClient()).client.listTools().then(r => r.tools.map(t => t.name).sort()); + // SQL connection + vi.mocked(getActiveConfig).mockReturnValue({ type: 'postgres', database: 'shop' } as any); + mockSqlDriver(); + const sqlNames = (await connectClient()).client.listTools().then(r => r.tools.map(t => t.name).sort()); + + const expected = [ + 'aggregate_documents', 'connection_info', 'count_documents', 'delete_document', + 'describe_collection', 'describe_table', 'execute_query', 'execute_write', + 'find_documents', 'insert_document', 'list_collections', 'list_tables', 'update_document', + ]; + expect(await mongoNames).toEqual(expected); + expect(await sqlNames).toEqual(expected); + }); + + it('connection_info reports the connected database type and which tools to use', async () => { mockMongoDriver(); const { client } = await connectClient(); - const names = (await client.listTools()).tools.map(t => t.name).sort(); - - expect(names).toContain('find_documents'); - expect(names).toContain('aggregate_documents'); - expect(names).toContain('count_documents'); - expect(names).toContain('list_collections'); - expect(names).not.toContain('execute_query'); - expect(names).not.toContain('execute_write'); + const res: any = await client.callTool({ name: 'connection_info', arguments: {} }); + const info = JSON.parse(res.content[0].text); + + expect(info).toMatchObject({ + connected: true, + databaseType: 'mongodb', + queryLanguage: 'MongoDB (MQL)', + }); + expect(info.useTools.read).toContain('find_documents'); + }); + + it('a SQL tool on a Mongo connection returns an actionable mismatch error, not "tool not found"', async () => { + const query = vi.fn(); + mockMongoDriver(query); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'execute_query', + arguments: { sql: 'SELECT 1' }, + }); + + expect(res.isError).toBe(true); + const text = res.content[0].text; + expect(text).toMatch(/MongoDB database/); + expect(text).toMatch(/find_documents/); + expect(text).toMatch(/connection_info/); + expect(query).not.toHaveBeenCalled(); }); - it('exposes SQL tools on a SQL connection', async () => { + it('a MongoDB tool on a SQL connection returns an actionable mismatch error', async () => { vi.mocked(getActiveConfig).mockReturnValue({ type: 'postgres', database: 'shop' } as any); - vi.mocked(getDriver).mockReturnValue({ queryMode: 'sql' } as any); + const query = vi.fn(); + mockSqlDriver(query); const { client } = await connectClient(); - const names = (await client.listTools()).tools.map(t => t.name).sort(); - expect(names).toContain('execute_query'); - expect(names).toContain('execute_write'); - expect(names).not.toContain('find_documents'); + const res: any = await client.callTool({ + name: 'find_documents', + arguments: { collection: 'users' }, + }); + + expect(res.isError).toBe(true); + const text = res.content[0].text; + expect(text).toMatch(/postgres \(SQL\) database/); + expect(text).toMatch(/execute_query/); + expect(query).not.toHaveBeenCalled(); }); - it('server instructions tell the model it is MongoDB and not to write SQL', async () => { + it('server instructions name the connected DB and point at connection_info', async () => { mockMongoDriver(); const { client } = await connectClient(); const instructions = client.getInstructions() ?? ''; - expect(instructions).toMatch(/mongodb/i); - expect(instructions).toMatch(/do not write sql/i); + expect(instructions).toMatch(/connection_info/); + expect(instructions).toMatch(/mongodb database/i); }); it('find_documents assembles a JSON-encoded MQL request for the driver', async () => { diff --git a/server/src/mcp.ts b/server/src/mcp.ts index 4b3b52f..aa2fbb1 100644 --- a/server/src/mcp.ts +++ b/server/src/mcp.ts @@ -12,6 +12,11 @@ const MAX_ROW_LIMIT = 10_000; const READ_KEYWORDS = new Set(['SELECT', 'SHOW', 'DESCRIBE', 'DESC', 'EXPLAIN', 'WITH']); const WRITE_KEYWORDS = new Set(['INSERT', 'UPDATE', 'DELETE', 'REPLACE']); +// Tool names grouped by database family, reused in guard messages and connection_info. +const SQL_TOOLS = 'execute_query, execute_write, list_tables, describe_table'; +const MONGO_READ_TOOLS = 'find_documents, aggregate_documents, count_documents, list_collections, describe_collection'; +const MONGO_WRITE_TOOLS = 'insert_document, update_document, delete_document'; + function firstKeyword(sql: string): string { const stripped = sql.replace(/\/\*[\s\S]*?\*\//g, '').replace(/--.*$/gm, '').trimStart(); const match = stripped.match(/^([A-Za-z]+)/); @@ -33,6 +38,75 @@ function requireConnection(): string | null { return null; } +/** + * Guard a tool to the query language of the connected database. + * + * The tool list is intentionally invariant across connections — every tool is + * always advertised — so a client that cached `tools/list` never ends up with a + * stale set when the user connects or switches databases mid-session. Instead, + * a tool from the wrong family returns this actionable message at call time, + * naming the connected DB type and the tools that actually apply, so the agent + * can self-correct (or ask the user to reconnect) instead of seeing a cryptic + * "tool not found". Caller must have already passed requireConnection(). + */ +function requireMode(expected: 'sql' | 'mql'): string | null { + if (getDriver().queryMode === expected) return null; + if (expected === 'sql') { + // A SQL tool was invoked, but the live connection is MongoDB. + return ( + 'This MCP connection is a MongoDB database, so the SQL tools do not apply here. ' + + `Use the MongoDB tools instead — reads: ${MONGO_READ_TOOLS}; writes: ${MONGO_WRITE_TOOLS}. ` + + 'Call connection_info to confirm the connected database; if your client cached an older ' + + 'tool list, re-list tools or reconnect the MCP server.' + ); + } + // A MongoDB tool was invoked, but the live connection is SQL. + const dbType = getActiveConfig()?.type ?? 'SQL'; + return ( + `This MCP connection is a ${dbType} (SQL) database, so the MongoDB tools do not apply here. ` + + `Use the SQL tools instead: ${SQL_TOOLS}. ` + + 'Call connection_info to confirm the connected database; if your client cached an older ' + + 'tool list, re-list tools or reconnect the MCP server.' + ); +} + +// --------------------------------------------------------------------------- +// connection_info — always available, regardless of database family. +// --------------------------------------------------------------------------- + +function registerConnectionInfo(server: McpServer): void { + server.registerTool( + 'connection_info', + { + description: + 'Report the database the MCP server is currently connected to: its type, query language, ' + + 'and which tools to use. Call this first to decide between the SQL and MongoDB tools, and ' + + 'again whenever a tool reports a database-type mismatch (the connection can change mid-session).', + inputSchema: {}, + }, + async () => { + if (!isConnected()) { + return toolJson({ + connected: false, + message: 'No database is connected. Ask the user to connect one via the Helix UI.', + }); + } + const isMongo = getDriver().queryMode === 'mql'; + const cfg = getActiveConfig(); + return toolJson({ + connected: true, + databaseType: cfg?.type ?? 'unknown', + database: cfg?.database ?? null, + queryLanguage: isMongo ? 'MongoDB (MQL)' : 'SQL', + writesEnabled: isMcpWritesAllowed(), + useTools: isMongo + ? { read: MONGO_READ_TOOLS.split(', '), write: MONGO_WRITE_TOOLS.split(', ') } + : { read: ['execute_query', 'list_tables', 'describe_table'], write: ['execute_write'] }, + }); + }, + ); +} + // --------------------------------------------------------------------------- // SQL tools (MySQL / Postgres — queryMode 'sql') // --------------------------------------------------------------------------- @@ -41,7 +115,7 @@ function registerSqlTools(server: McpServer): void { server.registerTool( 'list_tables', { - description: 'List tables, views, procedures, and triggers in a schema. If no schema is given, lists schemas (databases).', + description: '[SQL databases] List tables, views, procedures, and triggers in a schema. If no schema is given, lists schemas (databases).', inputSchema: { schema: z.string().optional().describe('Schema/database name. Omit to list available schemas.'), }, @@ -49,6 +123,8 @@ function registerSqlTools(server: McpServer): void { async ({ schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('sql'); + if (modeErr) return toolError(modeErr); try { const driver = getDriver(); @@ -77,7 +153,7 @@ function registerSqlTools(server: McpServer): void { server.registerTool( 'describe_table', { - description: 'Describe columns of a table: name, type, nullability, default, primary key, auto-increment.', + description: '[SQL databases] Describe columns of a table: name, type, nullability, default, primary key, auto-increment.', inputSchema: { schema: z.string().describe('Schema/database name.'), table: z.string().describe('Table name.'), @@ -86,6 +162,8 @@ function registerSqlTools(server: McpServer): void { async ({ schema, table }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('sql'); + if (modeErr) return toolError(modeErr); try { const tableInfo = await getDriver().getTable(schema, table); @@ -103,7 +181,7 @@ function registerSqlTools(server: McpServer): void { 'execute_query', { description: - 'Run a read-only SQL query (SELECT / SHOW / DESCRIBE / EXPLAIN / WITH). ' + + '[SQL databases] Run a read-only SQL query (SELECT / SHOW / DESCRIBE / EXPLAIN / WITH). ' + `Results are capped at ${DEFAULT_ROW_LIMIT} rows by default; pass "limit" (up to ${MAX_ROW_LIMIT}) to change.`, inputSchema: { sql: z.string().min(1).describe('Read-only SQL statement.'), @@ -115,6 +193,8 @@ function registerSqlTools(server: McpServer): void { async ({ sql, schema, limit }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('sql'); + if (modeErr) return toolError(modeErr); const kw = firstKeyword(sql); if (!READ_KEYWORDS.has(kw)) { @@ -152,7 +232,7 @@ function registerSqlTools(server: McpServer): void { 'execute_write', { description: - 'Run an INSERT, UPDATE, DELETE, or REPLACE statement. ' + + '[SQL databases] Run an INSERT, UPDATE, DELETE, or REPLACE statement. ' + 'Requires the user to have enabled "Allow MCP to modify data" in the Helix UI. ' + 'DDL (CREATE/DROP/ALTER/TRUNCATE) is not supported.', inputSchema: { @@ -163,6 +243,8 @@ function registerSqlTools(server: McpServer): void { async ({ sql, schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('sql'); + if (modeErr) return toolError(modeErr); if (!isMcpWritesAllowed()) { return toolError( @@ -216,7 +298,7 @@ function registerMongoTools(server: McpServer): void { 'list_collections', { description: - 'List collections (and views) in a MongoDB database. ' + + '[MongoDB] List collections (and views) in a database. ' + 'If no database is given, lists the available databases.', inputSchema: { schema: z.string().optional().describe('Database name. Omit to list available databases.'), @@ -225,6 +307,8 @@ function registerMongoTools(server: McpServer): void { async ({ schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); try { const driver = getDriver(); @@ -248,7 +332,7 @@ function registerMongoTools(server: McpServer): void { 'describe_collection', { description: - 'Describe a MongoDB collection: field names and types inferred from a sample of documents ' + + '[MongoDB] Describe a collection: field names and types inferred from a sample of documents ' + '(schemaless — fields may vary per document), plus its indexes and JSON Schema validator if set.', inputSchema: { schema: z.string().describe('Database name.'), @@ -258,6 +342,8 @@ function registerMongoTools(server: McpServer): void { async ({ schema, collection }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); try { const driver = getDriver(); @@ -286,7 +372,7 @@ function registerMongoTools(server: McpServer): void { 'find_documents', { description: - 'Find documents in a MongoDB collection (db.collection.find). ' + + '[MongoDB] Find documents in a collection (db.collection.find). ' + `Returns up to ${DEFAULT_ROW_LIMIT} documents by default; pass "limit" (up to ${MAX_ROW_LIMIT}) to change. ` + 'Note: a string "_id" filter is matched literally and will not match an ObjectId — omit it or filter on other fields.', inputSchema: { @@ -303,6 +389,8 @@ function registerMongoTools(server: McpServer): void { async ({ collection, filter, projection, sort, limit, skip, schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); const cap = limit ?? DEFAULT_ROW_LIMIT; try { @@ -329,7 +417,7 @@ function registerMongoTools(server: McpServer): void { 'aggregate_documents', { description: - 'Run an aggregation pipeline on a MongoDB collection (db.collection.aggregate). ' + + '[MongoDB] Run an aggregation pipeline on a collection (db.collection.aggregate). ' + `Results are capped at ${DEFAULT_ROW_LIMIT} documents by default; pass "limit" (up to ${MAX_ROW_LIMIT}) to change. ` + 'For large pipelines, add an explicit {"$limit": N} stage to bound work server-side.', inputSchema: { @@ -343,6 +431,8 @@ function registerMongoTools(server: McpServer): void { async ({ collection, pipeline, limit, schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); const cap = limit ?? DEFAULT_ROW_LIMIT; try { @@ -369,7 +459,7 @@ function registerMongoTools(server: McpServer): void { server.registerTool( 'count_documents', { - description: 'Count documents matching a filter in a MongoDB collection (db.collection.countDocuments).', + description: '[MongoDB] Count documents matching a filter in a collection (db.collection.countDocuments).', inputSchema: { collection: z.string().describe('Collection name.'), filter: docSchema.optional().describe('MongoDB query filter. Omit to count all documents.'), @@ -379,6 +469,8 @@ function registerMongoTools(server: McpServer): void { async ({ collection, filter, schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); try { const start = Date.now(); @@ -395,7 +487,7 @@ function registerMongoTools(server: McpServer): void { 'insert_document', { description: - 'Insert a single document into a MongoDB collection (db.collection.insertOne). ' + + '[MongoDB] Insert a single document into a collection (db.collection.insertOne). ' + 'Requires the user to have enabled "Allow MCP to modify data" in the Helix UI.', inputSchema: { collection: z.string().describe('Collection name.'), @@ -406,6 +498,8 @@ function registerMongoTools(server: McpServer): void { async ({ collection, document, schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); if (!isMcpWritesAllowed()) { return toolError( 'Writes are disabled. Ask the user to enable "Allow MCP to modify data" in the Helix UI (top-right menu).', @@ -427,7 +521,7 @@ function registerMongoTools(server: McpServer): void { 'update_document', { description: - 'Update a single document in a MongoDB collection (db.collection.updateOne). ' + + '[MongoDB] Update a single document in a collection (db.collection.updateOne). ' + 'Provide either "filter" or "id" to select the document, and "update" with operators like {"$set": {...}}. ' + 'Requires the user to have enabled "Allow MCP to modify data" in the Helix UI.', inputSchema: { @@ -441,6 +535,8 @@ function registerMongoTools(server: McpServer): void { async ({ collection, update, filter, id, schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); if (!isMcpWritesAllowed()) { return toolError( 'Writes are disabled. Ask the user to enable "Allow MCP to modify data" in the Helix UI (top-right menu).', @@ -465,7 +561,7 @@ function registerMongoTools(server: McpServer): void { 'delete_document', { description: - 'Delete a single document from a MongoDB collection (db.collection.deleteOne). ' + + '[MongoDB] Delete a single document from a collection (db.collection.deleteOne). ' + 'Provide either "filter" or "id" to select the document. ' + 'Requires the user to have enabled "Allow MCP to modify data" in the Helix UI.', inputSchema: { @@ -478,6 +574,8 @@ function registerMongoTools(server: McpServer): void { async ({ collection, filter, id, schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); if (!isMcpWritesAllowed()) { return toolError( 'Writes are disabled. Ask the user to enable "Allow MCP to modify data" in the Helix UI (top-right menu).', @@ -504,53 +602,51 @@ function registerMongoTools(server: McpServer): void { // --------------------------------------------------------------------------- /** - * The active query language for the connected database. Defaults to 'sql' when - * nothing is connected so the (backward-compatible) SQL toolset is advertised. - * The transport is stateless, so `tools/list` is re-evaluated on every request - * and always reflects the database connected at that moment. + * Build the server instructions. The tool list is invariant — both the SQL and + * MongoDB tools are always advertised — so the instructions explain that each + * family only works against its database type, point at connection_info, and + * name the database connected right now (which can change during a session). */ -function activeQueryMode(): 'sql' | 'mql' { - return isConnected() ? getDriver().queryMode : 'sql'; -} - -function buildInstructions(mode: 'sql' | 'mql'): string { +function buildInstructions(): string { + const connected = isConnected(); const cfg = getActiveConfig(); - const dbType = cfg?.type ?? 'none'; - const base = [ + const lines = [ 'Helix MCP exposes the database currently connected in the Helix UI.', - `The connected database is ${isConnected() ? `a ${dbType} database` : 'not connected yet'}.`, + 'The same tools are always listed, but each only works against its database family: ' + + `the SQL tools (${SQL_TOOLS}) require a SQL database, and the MongoDB tools ` + + `(${MONGO_READ_TOOLS}, ${MONGO_WRITE_TOOLS}) require a MongoDB database.`, + 'Call connection_info first to learn which database is connected and which tools to use. ' + + 'The connection can change during a session, so re-check connection_info if a tool reports ' + + 'a database-type mismatch.', ]; - if (mode === 'mql') { - base.push( - 'This is a MongoDB (document) database — use the MongoDB tools (find_documents, aggregate_documents, count_documents, ...).', - 'Do NOT write SQL: there are no SELECT/SHOW statements here, only collections and documents.', - 'Writes (insert/update/delete) are only allowed when the user has enabled "Allow MCP to modify data" in the Helix UI top-right menu.', + if (connected) { + const isMongo = getDriver().queryMode === 'mql'; + lines.push( + `Currently connected: a ${cfg?.type ?? 'unknown'} database — ` + + `${isMongo ? 'use the MongoDB tools' : 'use the SQL tools'}.`, ); } else { - base.push( - 'This is a SQL database — use SQL via execute_query (reads) and execute_write (writes).', - 'Writes (INSERT/UPDATE/DELETE/REPLACE) are only allowed when the user has enabled "Allow MCP to modify data" in the Helix UI top-right menu.', - 'DDL (CREATE/DROP/ALTER/TRUNCATE/RENAME) is not supported in this version.', - ); + lines.push('Currently no database is connected.'); } - return base.join(' '); + lines.push( + 'Writes are only allowed when the user has enabled "Allow MCP to modify data" in the Helix UI ' + + 'top-right menu. DDL (CREATE/DROP/ALTER/TRUNCATE/RENAME) is not supported.', + ); + return lines.join(' '); } export function buildMcpServer(): McpServer { - const mode = activeQueryMode(); const server = new McpServer( { name: 'helix-mcp', version: '0.1.0' }, { capabilities: { tools: {} }, - instructions: buildInstructions(mode), + instructions: buildInstructions(), }, ); - if (mode === 'mql') { - registerMongoTools(server); - } else { - registerSqlTools(server); - } + registerConnectionInfo(server); + registerSqlTools(server); + registerMongoTools(server); return server; } From 73ab9a332c337cf41532b38f91ec44fcf7e4e79a Mon Sep 17 00:00:00 2001 From: Eugene Date: Sat, 20 Jun 2026 09:34:45 -0400 Subject: [PATCH 3/3] =?UTF-8?q?fix:=20address=20PR=20review=20=E2=80=94=20?= =?UTF-8?q?bound=20aggregate=5Fdocuments=20server-side=20(#180)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - aggregate_documents now appends {$limit: cap + 1} to the pipeline so the MongoDB server bounds the result set instead of loading the whole aggregation into Node and slicing — matching find_documents. The +1 still lets us report truncation without a full scan. - Tidy describe_collection's optional getCollectionInfo call. - Add MCP-layer tests for aggregate_documents (limit + truncation), count_documents, list_collections, describe_collection, and delete_document; close client/server pairs in the invariant-toolset test. Co-Authored-By: Claude Opus 4.8 (1M context) --- server/src/mcp.test.ts | 117 +++++++++++++++++++++++++++++++++++++---- server/src/mcp.ts | 24 +++++---- 2 files changed, 122 insertions(+), 19 deletions(-) diff --git a/server/src/mcp.test.ts b/server/src/mcp.test.ts index 659c9d4..0172346 100644 --- a/server/src/mcp.test.ts +++ b/server/src/mcp.test.ts @@ -51,21 +51,27 @@ describe('MCP server – invariant toolset with per-call mode guards', () => { }); it('advertises the same full toolset regardless of connected DB (no stale-cache problem)', async () => { + const expected = [ + 'aggregate_documents', 'connection_info', 'count_documents', 'delete_document', + 'describe_collection', 'describe_table', 'execute_query', 'execute_write', + 'find_documents', 'insert_document', 'list_collections', 'list_tables', 'update_document', + ]; + // Mongo connection mockMongoDriver(); - const mongoNames = (await connectClient()).client.listTools().then(r => r.tools.map(t => t.name).sort()); + const mongo = await connectClient(); + const mongoNames = (await mongo.client.listTools()).tools.map(t => t.name).sort(); + await Promise.all([mongo.client.close(), mongo.server.close()]); + // SQL connection vi.mocked(getActiveConfig).mockReturnValue({ type: 'postgres', database: 'shop' } as any); mockSqlDriver(); - const sqlNames = (await connectClient()).client.listTools().then(r => r.tools.map(t => t.name).sort()); + const sql = await connectClient(); + const sqlNames = (await sql.client.listTools()).tools.map(t => t.name).sort(); + await Promise.all([sql.client.close(), sql.server.close()]); - const expected = [ - 'aggregate_documents', 'connection_info', 'count_documents', 'delete_document', - 'describe_collection', 'describe_table', 'execute_query', 'execute_write', - 'find_documents', 'insert_document', 'list_collections', 'list_tables', 'update_document', - ]; - expect(await mongoNames).toEqual(expected); - expect(await sqlNames).toEqual(expected); + expect(mongoNames).toEqual(expected); + expect(sqlNames).toEqual(expected); }); it('connection_info reports the connected database type and which tools to use', async () => { @@ -153,6 +159,99 @@ describe('MCP server – invariant toolset with per-call mode guards', () => { expect(JSON.parse(res.content[0].text).documents).toEqual([{ _id: 'a1', email: 'x@y.z' }]); }); + it('aggregate_documents appends a server-side $limit (cap+1) and reports truncation', async () => { + // Driver returns cap+1 rows (101) -> tool must slice to 100 and flag truncated. + const rows = Array.from({ length: 101 }, (_, i) => ({ _id: i })); + const query = vi.fn().mockResolvedValue({ rows, columnMeta: [{ name: '_id' }] }); + mockMongoDriver(query); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'aggregate_documents', + arguments: { collection: 'orders', pipeline: [{ $match: { paid: true } }] }, + }); + + expect(res.isError).toBeFalsy(); + const req = JSON.parse(query.mock.calls[0][0]); + expect(req).toMatchObject({ collection: 'orders', operation: 'aggregate' }); + expect(req.pipeline).toEqual([{ $match: { paid: true } }, { $limit: 101 }]); + const body = JSON.parse(res.content[0].text); + expect(body.docCount).toBe(100); + expect(body.truncated).toBe(true); + }); + + it('count_documents returns the count from the driver result', async () => { + const query = vi.fn().mockResolvedValue({ rows: [{ count: 42 }], columnMeta: [{ name: 'count' }] }); + mockMongoDriver(query); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'count_documents', + arguments: { collection: 'users', filter: { active: true } }, + }); + + expect(res.isError).toBeFalsy(); + expect(JSON.parse(query.mock.calls[0][0])).toMatchObject({ + collection: 'users', operation: 'count', filter: { active: true }, + }); + expect(JSON.parse(res.content[0].text).count).toBe(42); + }); + + it('list_collections returns collections for a database', async () => { + mockMongoDriver(); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'list_collections', + arguments: { schema: 'shop' }, + }); + + expect(res.isError).toBeFalsy(); + const body = JSON.parse(res.content[0].text); + expect(body.database).toBe('shop'); + expect(body.collections).toEqual([{ name: 'users', approxDocs: 3 }]); + }); + + it('describe_collection surfaces fields, indexes, and validator', async () => { + const driver = mockMongoDriver(); + driver.getTable.mockResolvedValue({ + name: 'users', rows: 3, columns: [{ name: '_id', dataType: 'objectId' }], + }); + driver.getCollectionInfo.mockResolvedValue({ + validator: { $jsonSchema: {} }, indexes: [{ name: '_id_' }], + }); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'describe_collection', + arguments: { schema: 'shop', collection: 'users' }, + }); + + expect(res.isError).toBeFalsy(); + const body = JSON.parse(res.content[0].text); + expect(body.fields).toEqual([{ name: '_id', dataType: 'objectId' }]); + expect(body.indexes).toEqual([{ name: '_id_' }]); + expect(body.validator).toEqual({ $jsonSchema: {} }); + }); + + it('delete_document assembles a deleteOne request and returns deletedCount', async () => { + vi.mocked(isMcpWritesAllowed).mockReturnValue(true); + const query = vi.fn().mockResolvedValue({ rows: [], columnMeta: [], affectedRows: 1 }); + mockMongoDriver(query); + const { client } = await connectClient(); + + const res: any = await client.callTool({ + name: 'delete_document', + arguments: { collection: 'users', id: '507f1f77bcf86cd799439011' }, + }); + + expect(res.isError).toBeFalsy(); + expect(JSON.parse(query.mock.calls[0][0])).toMatchObject({ + collection: 'users', operation: 'deleteOne', id: '507f1f77bcf86cd799439011', + }); + expect(JSON.parse(res.content[0].text).deletedCount).toBe(1); + }); + it('blocks insert_document when writes are disabled', async () => { const query = vi.fn(); mockMongoDriver(query); diff --git a/server/src/mcp.ts b/server/src/mcp.ts index aa2fbb1..3c00713 100644 --- a/server/src/mcp.ts +++ b/server/src/mcp.ts @@ -351,9 +351,7 @@ function registerMongoTools(server: McpServer): void { if (!info) { return toolError(`Collection "${schema}"."${collection}" not found.`); } - const collInfo = driver.getCollectionInfo - ? await driver.getCollectionInfo(schema, collection) - : null; + const collInfo = await driver.getCollectionInfo?.(schema, collection) ?? null; return toolJson({ database: schema, collection, @@ -437,16 +435,22 @@ function registerMongoTools(server: McpServer): void { const cap = limit ?? DEFAULT_ROW_LIMIT; try { const start = Date.now(); - const result = await runMql({ collection, operation: 'aggregate', pipeline }, schema); + // Bound work server-side instead of fetching the whole result and + // slicing in Node. Fetching one doc past the cap lets us still report + // truncation without loading a potentially huge result set into memory, + // mirroring how find_documents forwards its limit to the driver. + const result = await runMql( + { collection, operation: 'aggregate', pipeline: [...pipeline, { $limit: cap + 1 }] }, + schema, + ); const executionTime = Date.now() - start; - const totalRows = result.rows.length; - const capped = result.rows.slice(0, cap); + const truncated = result.rows.length > cap; + const documents = truncated ? result.rows.slice(0, cap) : result.rows; return toolJson({ fields: result.columnMeta.map(c => c.name), - documents: capped, - docCount: capped.length, - totalDocs: totalRows, - truncated: totalRows > capped.length, + documents, + docCount: documents.length, + truncated, limitApplied: cap, executionTime, });