import { z } from "zod";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import {
Template,
DataMakerResponse,
Connection,
Endpoint,
Scenario,
} from "./types.js";
import {
fetchAPI,
storeToR2AndSummarize,
isSapEndpoint,
parseResponseData,
countTokens,
convertToObjectArray,
} from "./helpers.js";
import { config } from "dotenv";
import { RESPONSE_TOKEN_THRESHOLD } from "./lib/config.js";
config();
const tokenThreshold = RESPONSE_TOKEN_THRESHOLD;
export function registerTools(server: McpServer) {
server.tool(
"generate_from_id",
"Generate data from a datamaker template id",
{
template_id: z.string().describe("A valid datamaker template id"),
quantity: z
.number()
.default(10)
.describe("Number of records to generate"),
},
async ({ template_id, quantity }, context) => {
const ctx = context as any;
try {
// Get all templates
const templates = await fetchAPI<Template[]>(
"/templates",
"GET",
undefined,
ctx?.jwtToken
);
// Find the template with the given id
const template = templates.find((t) => t.id === template_id);
if (!template) {
throw new Error(`Template with id ${template_id} not found`);
}
// Filter to ensure inactive fields are not returned
const activeFields = template.fields.filter(
(field: any) => field.active === undefined || field.active === true
);
// Generate data from the template
const response = await fetchAPI<DataMakerResponse>(
"/datamaker",
"POST",
{
fields: activeFields,
quantity,
},
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(response.live_data, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"get_templates",
"Get all templates",
{
projectId: z.string().optional().describe("A valid datamaker project id"),
},
async ({ projectId }, context) => {
const ctx = context as any;
try {
const queryParams = projectId ? `?projectId=${projectId}` : "";
const endpoint = `/templates${queryParams}`;
const templates = await fetchAPI<Template[]>(
endpoint,
"GET",
undefined,
ctx?.jwtToken
);
const tokens = countTokens(JSON.stringify(templates, null, 2));
if (tokens > tokenThreshold) {
const result = await storeToR2AndSummarize(templates, "templates");
return {
content: [
{
type: "text",
text: `Available templates (${
result.totalCount
} templates) are too large to show. Showing first ${
result.summary.length
} templates:\n\n${JSON.stringify(
result.summary,
null,
2
)}\n\nFull dataset stored to R2\nš **View all templates in a link that opens in a new tab: ${
result.viewUrl
}\n\nThis link expires in 24 hours.`,
},
],
};
}
const simplifiedTemplateObjects = templates.map((template) => ({
id: template.id,
name: template.name,
}));
return {
content: [
{
type: "text",
text: JSON.stringify(simplifiedTemplateObjects, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"get_template_by_id",
"Get a template by id",
{
template_id: z.string().describe("A valid datamaker template id"),
},
async ({ template_id }, context) => {
const ctx = context as any;
try {
const template = await fetchAPI<Template>(
`/templates/${template_id}`,
"GET",
undefined,
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(template, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"get_connections",
"Get all connections",
{
projectId: z.string().optional().describe("A valid datamaker project id"),
},
async ({ projectId }, context) => {
const ctx = context as any;
try {
const queryParams = projectId ? `?projectId=${projectId}` : "";
const endpoint = `/connections${queryParams}`;
const connections = await fetchAPI<Connection[]>(
endpoint,
"GET",
undefined,
ctx?.jwtToken
);
const simplifiedConnectionObjects = connections.map((connection) => ({
id: connection.id,
name: connection.name,
}));
return {
content: [
{
type: "text",
text: JSON.stringify(simplifiedConnectionObjects, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"delete_connection",
"Delete a connection by id",
{
connection_id: z.string().describe("A valid datamaker connection id"),
},
async ({ connection_id }, context) => {
const ctx = context as any;
try {
const response = await fetchAPI<{ message: string }>(
`/connections/${connection_id}`,
"DELETE",
undefined,
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(response, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"create_connection",
"Create a new connection",
{
name: z.string().describe("Name of the connection"),
connectionString: z
.string()
.describe("Connection string for the database"),
type: z
.string()
.describe(
"Type of the connection (e.g., 'postgresql', 'mysql', 'mongodb')"
),
},
async ({ name, connectionString, type }, context) => {
const ctx = context as any;
try {
const connection = await fetchAPI<Connection>(
"/connections",
"POST",
{
name,
connectionString,
type,
},
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(connection, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"update_connection",
"Update a connection by id",
{
connection_id: z.string().describe("A valid datamaker connection id"),
name: z.string().optional().describe("Name of the connection"),
connectionString: z
.string()
.optional()
.describe("Connection string for the database"),
type: z
.string()
.optional()
.describe(
"Type of the connection (e.g., 'postgresql', 'mysql', 'mongodb')"
),
},
async ({ connection_id, name, connectionString, type }, context) => {
const ctx = context as any;
try {
// Build update data object with only provided fields
const updateData: any = {};
if (name !== undefined) updateData.name = name;
if (connectionString !== undefined)
updateData.connectionString = connectionString;
if (type !== undefined) updateData.type = type;
const connection = await fetchAPI<Connection>(
`/connections/${connection_id}`,
"PUT",
updateData,
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(connection, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"get_connection_by_id",
"Get a connection by id",
{
connection_id: z.string().describe("A valid datamaker connection id"),
},
async ({ connection_id }, context) => {
const ctx = context as any;
try {
const connection = await fetchAPI<Connection>(
`/connections/${connection_id}`,
"GET",
undefined,
ctx?.jwtToken
);
const simplifiedConnectionObject = {
id: connection.id,
name: connection.name,
conectionString: connection.connectionString,
type: connection.type,
};
return {
content: [
{
type: "text",
text: JSON.stringify(simplifiedConnectionObject, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"get_endpoints",
"Get all endpoints",
{
projectId: z.string().optional().describe("A valid datamaker project id"),
},
async ({ projectId }, context) => {
const ctx = context as any;
try {
const queryParams = projectId ? `?projectId=${projectId}` : "";
const endpoint = `/endpoints${queryParams}`;
const endpoints = await fetchAPI<Endpoint[]>(
endpoint,
"GET",
undefined,
ctx?.jwtToken
);
const tokens = countTokens(JSON.stringify(endpoints, null, 2));
if (tokens > tokenThreshold) {
const result = await storeToR2AndSummarize(endpoints, "endpoints");
return {
content: [
{
type: "text",
text: `Available endpoints (${
result.totalCount
} endpoints) are too large to show. Showing first ${
result.summary.length
} endpoints:\n\n${JSON.stringify(
result.summary,
null,
2
)}\n\nFull dataset stored to R2\nš **View all endpoints in a link that opens in a new tab: ${
result.viewUrl
}\n\nThis link expires in 24 hours.`,
},
],
};
}
const simplifiedEndpointObjects = endpoints.map((endpoint) => ({
id: endpoint.id,
name: endpoint.name,
}));
return {
content: [
{
type: "text",
text: JSON.stringify(simplifiedEndpointObjects, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
});
server.tool(
"get_endpoint_by_id",
"Get an endpoint by id",
{
endpoint_id: z.string().describe("A valid datamaker endpoint id"),
},
async ({ endpoint_id }, context) => {
const ctx = context as any;
try {
const endpoint = await fetchAPI<Endpoint>(
`/endpoints/${endpoint_id}`,
"GET",
undefined,
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(endpoint, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"export_to_endpoint",
"Export data to an endpoint",
{
endpoint_id: z.string().describe("A valid datamaker endpoint id"),
data: z.any().describe("The data to export"),
},
async ({ endpoint_id, data }, context) => {
const ctx = context as any;
try {
// get the endpoint info
const endpoint = await fetchAPI<Endpoint>(
`/endpoints/${endpoint_id}`,
"GET",
undefined,
ctx?.jwtToken
);
// Check if this is a SAP endpoint and return early error
if (isSapEndpoint(endpoint.url)) {
return {
content: [
{
type: "text",
text: "Error: This tool was incorrectly used for a SAP endpoint. Please use the appropriate tool for SAP endpoints instead.",
},
],
};
}
// Convert headers object to array format expected by /export/rest
const headersArray = Object.entries(endpoint.headers || {}).map(
([headerKey, headerValue]) => ({
headerKey,
headerValue,
})
);
// Use /export/rest endpoint which handles authentication
const responseData = await fetchAPI<any>(
"/export/rest",
"POST",
{
id: endpoint_id,
method: endpoint.method,
url: endpoint.url,
body: data,
headers: headersArray,
},
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(responseData, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"fetch_from_endpoint",
"Fetch data from one of the user defined endpoints from datamaker.",
{
endpoint_id: z.string().describe("A valid endpoint id"),
filter: z
.any()
.optional()
.describe(
"Filter object with field names as keys and objects with value and operator properties"
),
},
async ({ endpoint_id, filter }, context) => {
const ctx = context as any;
try {
const endpoint = await fetchAPI<Endpoint>(
`/endpoints/${endpoint_id}`,
"GET",
undefined,
ctx?.jwtToken
);
// Check if this is a SAP endpoint and return early error
if (isSapEndpoint(endpoint?.url)) {
return {
content: [
{
type: "text",
text: "Error: This tool was incorrectly used for a SAP endpoint. Please use the appropriate tool for SAP endpoints instead.",
},
],
};
}
// Handle non-SAP endpoints
const response = await fetch(endpoint.url, {
method: endpoint.method as "GET" | "POST" | "PUT" | "DELETE",
headers: endpoint.headers,
});
if (!response.ok) {
const errorText = await response.text();
throw new Error(
`HTTP error! status: ${response.status}, response: ${JSON.stringify(
errorText
)}`
);
}
const responseData = await parseResponseData(response);
const tokens = countTokens(JSON.stringify(responseData, null, 2));
if (tokens > tokenThreshold) {
// Convert the response data to an array of objects for R2 storage
const objectArray = convertToObjectArray(responseData);
const result = await storeToR2AndSummarize(
objectArray,
"endpoint-responses"
);
return {
content: [
{
type: "text",
text: `Response data (${
result.totalCount
} items) is too large to show. Showing first ${
result.summary.length
} items:\n\n${JSON.stringify(
result.summary,
null,
2
)}\n\nFull dataset stored to R2\nš **View all data in a link that opens in a new tab: ${
result.viewUrl
}\n\nThis link expires in 24 hours.`,
},
],
};
}
return {
content: [
{
type: "text",
text: JSON.stringify(responseData, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"flatten_json",
"Flatten a nested JSON object so that nested keys are joined by dots (e.g., {user: {name: 'John Doe'}} becomes {'user.name': 'John Doe'}).",
{
data: z.any().describe("The nested JSON object to flatten"),
},
async ({ data }) => {
function flatten(
obj: Record<string, any>,
prefix = "",
res: Record<string, any> = {}
): Record<string, any> {
for (const key in obj) {
if (Object.prototype.hasOwnProperty.call(obj, key)) {
const value = obj[key];
const newKey = prefix ? `${prefix}.${key}` : key;
if (
typeof value === "object" &&
value !== null &&
!Array.isArray(value)
) {
flatten(value, newKey, res);
} else {
res[newKey] = value;
}
}
}
return res;
}
try {
const flattened = flatten(data);
return {
content: [
{
type: "text",
text: JSON.stringify(flattened, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"save_scenario",
"Save Datamaker scenario into database",
{
code: z.string().describe("Python code for the scenario"),
name: z.string().describe("Name of the scenario"),
projectId: z.string().optional().describe("A valid datamaker project id"),
},
async ({ code, name, projectId }, context) => {
const ctx = context as any;
try {
const scenario = await fetchAPI<Promise<Scenario>>(
`/scenarios/save`,
"POST",
{ code, name, projectId },
ctx?.jwtToken
);
if (!scenario) {
throw new Error("Failed to save scenario.");
}
return {
content: [
{
type: "text" as const,
text: JSON.stringify(scenario, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"execute_scenario",
"Execute a Datamaker scenario",
{
scenarioId: z.string().describe("A valid datamaker scenario id"),
projectId: z.string().describe("A valid datamaker project id"),
},
async ({ scenarioId, projectId }, context) => {
const ctx = context as any;
try {
const response = await fetchAPI<any>(
`/scenarios/execute`,
"POST",
{ projectId, scenarioId },
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(response, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool(
"get_scenario_by_id",
"Get a scenario by id",
{
scenarioId: z.string().describe("A valid datamaker scenario id"),
},
async ({ scenarioId }, context) => {
const ctx = context as any;
try {
const scenario = await fetchAPI<Scenario>(
`/scenarios/${scenarioId}`,
"GET",
undefined,
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(scenario, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
server.tool("get_scenarios", "Get all scenarios", {}, async ({}, context) => {
const ctx = context as any;
try {
const scenarios = await fetchAPI<Scenario[]>(
"/scenarios",
"GET",
undefined,
ctx?.jwtToken
);
const tokens = countTokens(JSON.stringify(scenarios, null, 2));
if (tokens > tokenThreshold) {
const result = await storeToR2AndSummarize(scenarios, "scenarios");
return {
content: [
{
type: "text",
text: `Available scenarios (${
result.totalCount
} scenarios) are too large to show. Showing first ${
result.summary.length
} scenarios:\n\n${JSON.stringify(
result.summary,
null,
2
)}\n\nFull dataset stored to R2\nš **View all scenarios in a link that opens in a new tab: ${
result.viewUrl
}\n\nThis link expires in 24 hours.`,
},
],
};
}
const simplifiedScenarioObjects = scenarios.map((scenario) => ({
id: scenario.id,
name: scenario.name,
}));
return {
content: [
{
type: "text",
text: JSON.stringify(simplifiedScenarioObjects, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
});
server.tool(
"execute_python_script",
"Execute a Python script by first saving it to R2 and then running it using the datamaker runner",
{
script: z.string().describe("The Python script code to execute"),
filename: z
.string()
.describe("The filename for the script (e.g., 'script.py')"),
},
async ({ script, filename }, context) => {
const ctx = context as any;
try {
// Step 1: Save the script to R2
const saveResponse = await fetchAPI<{ url: string }>(
"/upload-text",
"POST",
{
text: script,
filename: filename,
},
ctx?.jwtToken
);
if (!saveResponse.url) {
throw new Error("Failed to save script: No URL returned");
}
// Step 2: Execute the script using the returned URL
const runResponse = await fetchAPI<any>(
"/execute-python",
"POST",
{
url: saveResponse.url,
},
ctx?.jwtToken
);
return {
content: [
{
type: "text",
text: JSON.stringify(runResponse, null, 2),
},
],
};
} catch (error) {
return {
content: [
{
type: "text",
text: `Error: ${
error instanceof Error ? error.message : "Unknown error"
}`,
},
],
};
}
}
);
}
