diff --git a/server/src/mcp.test.ts b/server/src/mcp.test.ts new file mode 100644 index 0000000..0172346 --- /dev/null +++ b/server/src/mcp.test.ts @@ -0,0 +1,305 @@ +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; +} + +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); + vi.mocked(getActiveConfig).mockReturnValue({ type: 'mongodb', database: 'shop' } as any); + vi.mocked(isMcpWritesAllowed).mockReturnValue(false); + }); + + 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 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 sql = await connectClient(); + const sqlNames = (await sql.client.listTools()).tools.map(t => t.name).sort(); + await Promise.all([sql.client.close(), sql.server.close()]); + + expect(mongoNames).toEqual(expected); + expect(sqlNames).toEqual(expected); + }); + + it('connection_info reports the connected database type and which tools to use', async () => { + mockMongoDriver(); + const { client } = await connectClient(); + 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('a MongoDB tool on a SQL connection returns an actionable mismatch error', async () => { + vi.mocked(getActiveConfig).mockReturnValue({ type: 'postgres', database: 'shop' } as any); + const query = vi.fn(); + mockSqlDriver(query); + const { client } = await connectClient(); + + 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 name the connected DB and point at connection_info', async () => { + mockMongoDriver(); + const { client } = await connectClient(); + const instructions = client.getInstructions() ?? ''; + + expect(instructions).toMatch(/connection_info/); + expect(instructions).toMatch(/mongodb database/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('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); + 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..3c00713 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; @@ -11,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]+)/); @@ -32,24 +38,84 @@ function requireConnection(): string | null { return null; } -export function buildMcpServer(): McpServer { - const server = new McpServer( - { name: 'helix-mcp', version: '0.1.0' }, +/** + * 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', { - 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(' '), + 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') +// --------------------------------------------------------------------------- +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.'), }, @@ -57,6 +123,8 @@ export function buildMcpServer(): McpServer { async ({ schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('sql'); + if (modeErr) return toolError(modeErr); try { const driver = getDriver(); @@ -85,7 +153,7 @@ export function buildMcpServer(): McpServer { 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.'), @@ -94,6 +162,8 @@ export function buildMcpServer(): McpServer { 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); @@ -111,7 +181,7 @@ export function buildMcpServer(): McpServer { '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.'), @@ -123,6 +193,8 @@ export function buildMcpServer(): McpServer { 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)) { @@ -160,7 +232,7 @@ export function buildMcpServer(): McpServer { '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: { @@ -171,6 +243,8 @@ export function buildMcpServer(): McpServer { async ({ sql, schema }) => { const err = requireConnection(); if (err) return toolError(err); + const modeErr = requireMode('sql'); + if (modeErr) return toolError(modeErr); if (!isMcpWritesAllowed()) { return toolError( @@ -202,6 +276,381 @@ 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: + '[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.'), + }, + }, + async ({ schema }) => { + const err = requireConnection(); + if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); + + 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: + '[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.'), + collection: z.string().describe('Collection name.'), + }, + }, + async ({ schema, collection }) => { + const err = requireConnection(); + if (err) return toolError(err); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); + + try { + const driver = getDriver(); + const info = await driver.getTable(schema, collection); + if (!info) { + return toolError(`Collection "${schema}"."${collection}" not found.`); + } + const collInfo = 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: + '[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: { + 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 modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); + + 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: + '[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: { + 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 modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); + + const cap = limit ?? DEFAULT_ROW_LIMIT; + try { + const start = Date.now(); + // 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 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, + docCount: documents.length, + truncated, + limitApplied: cap, + executionTime, + }); + } catch (e) { + return toolError(e instanceof Error ? e.message : String(e)); + } + }, + ); + + server.registerTool( + 'count_documents', + { + 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.'), + 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); + const modeErr = requireMode('mql'); + if (modeErr) return toolError(modeErr); + + 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: + '[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.'), + 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); + 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).', + ); + } + + 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: + '[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: { + 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); + 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).', + ); + } + 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: + '[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: { + 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); + 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).', + ); + } + 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 +// --------------------------------------------------------------------------- + +/** + * 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 buildInstructions(): string { + const connected = isConnected(); + const cfg = getActiveConfig(); + const lines = [ + 'Helix MCP exposes the database currently connected in the Helix UI.', + '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 (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 { + lines.push('Currently no database is connected.'); + } + 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 server = new McpServer( + { name: 'helix-mcp', version: '0.1.0' }, + { + capabilities: { tools: {} }, + instructions: buildInstructions(), + }, + ); + + registerConnectionInfo(server); + registerSqlTools(server); + registerMongoTools(server); return server; }