/**
* Custom Tool Handler
* Creates MCP tool handlers for custom SQL-based tools defined in TOML config
*/
import { z } from "zod";
import { ToolConfig, ParameterConfig } from "../types/config.js";
import { ConnectorManager } from "../connectors/manager.js";
import {
createToolSuccessResponse,
createToolErrorResponse,
} from "../utils/response-formatter.js";
import { mapArgumentsToArray } from "../utils/parameter-mapper.js";
import { isReadOnlySQL, allowedKeywords } from "../utils/allowed-keywords.js";
import { requestStore } from "../requests/index.js";
import { getClientIdentifier } from "../utils/client-identifier.js";
/**
* Build a Zod schema from parameter definitions
* Returns a plain object with Zod schemas (MCP SDK format)
* @param parameters Parameter configurations from TOML
* @returns Plain object with Zod type definitions
*/
export function buildZodSchemaFromParameters(
parameters: ParameterConfig[] | undefined
): Record<string, z.ZodTypeAny> {
if (!parameters || parameters.length === 0) {
return {};
}
const schemaShape: Record<string, z.ZodTypeAny> = {};
for (const param of parameters) {
let fieldSchema: z.ZodTypeAny;
// Build base schema based on type
switch (param.type) {
case "string":
fieldSchema = z.string().describe(param.description);
break;
case "integer":
fieldSchema = z.number().int().describe(param.description);
break;
case "float":
fieldSchema = z.number().describe(param.description);
break;
case "boolean":
fieldSchema = z.boolean().describe(param.description);
break;
case "array":
fieldSchema = z.array(z.unknown()).describe(param.description);
break;
default:
throw new Error(`Unsupported parameter type: ${param.type}`);
}
// Add enum constraint if allowed_values is specified
if (param.allowed_values && param.allowed_values.length > 0) {
if (param.type === "string") {
fieldSchema = z.enum(param.allowed_values as [string, ...string[]]).describe(param.description);
} else {
// For non-string types, use refine to validate against allowed values
fieldSchema = fieldSchema.refine(
(val) => param.allowed_values!.includes(val),
{
message: `Value must be one of: ${param.allowed_values.join(", ")}`,
}
);
}
}
// Make field optional if it has a default value or is explicitly marked as not required
if (param.default !== undefined || param.required === false) {
fieldSchema = fieldSchema.optional();
}
schemaShape[param.name] = fieldSchema;
}
return schemaShape;
}
/**
* Build input schema in MCP format (JSON Schema compatible)
* @param parameters Parameter configurations from TOML
* @returns JSON Schema object
*/
export function buildInputSchema(parameters: ParameterConfig[] | undefined): {
type: "object";
properties: Record<string, any>;
required?: string[];
} {
// Convert Zod schema to JSON Schema-like format for MCP
const properties: Record<string, any> = {};
const required: string[] = [];
if (parameters) {
for (const param of parameters) {
const propSchema: any = {
description: param.description,
};
// Map type to JSON Schema type
switch (param.type) {
case "string":
propSchema.type = "string";
break;
case "integer":
propSchema.type = "integer";
break;
case "float":
propSchema.type = "number";
break;
case "boolean":
propSchema.type = "boolean";
break;
case "array":
propSchema.type = "array";
break;
}
// Add enum if allowed_values specified
if (param.allowed_values && param.allowed_values.length > 0) {
propSchema.enum = param.allowed_values;
}
properties[param.name] = propSchema;
// Track required fields
if (param.required !== false && param.default === undefined) {
required.push(param.name);
}
}
}
const schema: any = {
type: "object",
properties,
};
if (required.length > 0) {
schema.required = required;
}
return schema;
}
/**
* Create a custom tool handler for a user-defined SQL tool
* @param toolConfig Tool configuration from TOML
* @returns Handler function compatible with MCP server.registerTool
*/
export function createCustomToolHandler(toolConfig: ToolConfig) {
// Build Zod schema shape for MCP registration
const zodSchemaShape = buildZodSchemaFromParameters(toolConfig.parameters);
// Wrap in z.object() for validation
const zodSchema = z.object(zodSchemaShape);
return async (args: any, extra: any) => {
const startTime = Date.now();
let success = true;
let errorMessage: string | undefined;
let paramValues: any[] = [];
try {
// 1. Validate arguments against Zod schema
const validatedArgs = zodSchema.parse(args);
// 2. Get connector and execute options for the specified source
const connector = ConnectorManager.getCurrentConnector(toolConfig.source);
const executeOptions = ConnectorManager.getCurrentExecuteOptions(
toolConfig.source
);
// 3. Check if SQL is allowed based on readonly mode
const isReadonly = executeOptions.readonly === true;
if (isReadonly && !isReadOnlySQL(toolConfig.statement, connector.id)) {
errorMessage = `Tool '${toolConfig.name}' cannot execute in readonly mode for source '${toolConfig.source}'. Only read-only SQL operations are allowed: ${allowedKeywords[connector.id]?.join(", ") || "none"}`;
success = false;
return createToolErrorResponse(errorMessage, "READONLY_VIOLATION");
}
// 4. Map parameters to array format for SQL execution
paramValues = mapArgumentsToArray(
toolConfig.parameters,
validatedArgs
);
// 5. Execute SQL with parameters
const result = await connector.executeSQL(
toolConfig.statement,
executeOptions,
paramValues
);
// 6. Build response data
const responseData = {
rows: result.rows,
count: result.rows.length,
source_id: toolConfig.source,
};
return createToolSuccessResponse(responseData);
} catch (error) {
success = false;
errorMessage = (error as Error).message;
// Provide helpful error messages for common issues
if (error instanceof z.ZodError) {
const issues = error.issues.map((i) => `${i.path.join(".")}: ${i.message}`).join("; ");
errorMessage = `Parameter validation failed: ${issues}`;
} else {
// Add SQL context to execution errors for debugging
errorMessage = `${errorMessage}\n\nSQL: ${toolConfig.statement}\nParameters: ${JSON.stringify(paramValues)}`;
}
return createToolErrorResponse(errorMessage, "EXECUTION_ERROR");
} finally {
// Track the request
requestStore.add({
id: crypto.randomUUID(),
timestamp: new Date().toISOString(),
sourceId: toolConfig.source,
toolName: toolConfig.name,
sql: toolConfig.statement,
durationMs: Date.now() - startTime,
client: getClientIdentifier(extra),
success,
error: errorMessage,
});
}
};
}