import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { z } from "zod";
import {
ComponentsApi,
FuncsApi,
ManagementFuncsApi,
SchemasApi,
} from "@systeminit/api-client";
import { apiConfig, WORKSPACE_ID } from "../si_client.ts";
import {
errorResponse,
generateDescription,
successResponse,
withAnalytics,
} from "./commonBehavior.ts";
import { AttributesSchema } from "../data/components.ts";
import { validateSchemaPrereqs } from "../data/schemaHints.ts";
const name = "component-import";
const title = "Import a component of a given schema and resource id";
const description =
`<description>Import a component for a given schema name and resourceId into a change set.</description><usage>Use this tool to bring an existing resource into System Initiative. For example, if the user asks to import the AWS VPC vpc-0abcdef1234567890, then 'vpc-0abcdef1234567890' is the resourceId and AWS::EC2::VPC is the schemaName. After importing a component, you should ask the user if they want you to update the attributes of the imported component to use subscriptions to any existing components attributes - for example, an imported AWS::EC2::Subnet would be updated to have a subscription to the /resource_value/VpcId of the AWS::EC2::VPC that matches the imported VpcId attribute of the subnet.</usage>`;
const ImportComponentInputSchemaRaw = {
changeSetId: z.string().describe(
"The change set to import the component in; components cannot be imported on the HEAD change set",
),
schemaName: z.string().describe("the schema name of the component to import"),
resourceId: z.string().describe(
"the resource id to import; should map to a unique id generated by the component on creation",
),
attributes: AttributesSchema.describe(
"attributes to set on the component before import; for AWS resources, this *must* include setting a subscription for the AWS Credential, and either a raw value or a subscription for /domain/extra/Region as well",
),
};
const ImportComponentOutputSchemaRaw = {
status: z.enum(["success", "failure"]),
errorMessage: z.string().optional().describe(
"If the status is failure, the error message will contain information about what went wrong",
),
data: z.object({
componentId: z.string().describe("the component id"),
componentName: z.string().describe("the components name"),
schemaName: z.string().describe("the schema for the component"),
funcRunId: z.string().nullable().optional().describe(
"the function run id for this management function; useful for debugging failure",
),
}),
};
const ImportComponentOutputSchema = z.object(
ImportComponentOutputSchemaRaw,
);
type ImportComponentResult = z.infer<
typeof ImportComponentOutputSchema
>["data"];
export function componentImportTool(server: McpServer) {
server.registerTool(
name,
{
title,
description: generateDescription(
description,
"componentDiscoverResponse",
ImportComponentOutputSchema,
),
inputSchema: ImportComponentInputSchemaRaw,
outputSchema: ImportComponentOutputSchemaRaw,
},
async (
{ changeSetId, schemaName, resourceId, attributes },
): Promise<CallToolResult> => {
return await withAnalytics(name, async () => {
const prereqError = validateSchemaPrereqs(schemaName, attributes);
if (prereqError) {
return prereqError;
}
const siApi = new ComponentsApi(apiConfig);
const siSchemasApi = new SchemasApi(apiConfig);
const siFuncsApi = new FuncsApi(apiConfig);
try {
const findSchemaResponse = await siSchemasApi.findSchema({
workspaceId: WORKSPACE_ID,
changeSetId: changeSetId,
schema: schemaName,
});
const schemaId = findSchemaResponse.data.schemaId;
const response = await siApi.createComponent({
workspaceId: WORKSPACE_ID,
changeSetId: changeSetId,
createComponentV1Request: {
name: resourceId,
resourceId,
schemaName,
attributes,
},
});
const result: Record<string, string> = {
componentId: response.data.component.id,
componentName: response.data.component.name,
schemaName: schemaName,
};
// Now get the variantFuncs so we can decide on the import function
const defaultVariantResponse = await siSchemasApi.getDefaultVariant({
workspaceId: WORKSPACE_ID,
changeSetId: changeSetId,
schemaId,
});
const variantFunctions = defaultVariantResponse.data.variantFuncs;
const importFunc = variantFunctions.find(
(func) =>
func.funcKind.kind === "management" &&
func.funcKind.managementFuncKind === "import",
);
if (!importFunc) {
return errorResponse({
message: `The schema ${schemaName} doesn't support import`,
});
}
const importFuncResponse = await siFuncsApi.getFunc({
workspaceId: WORKSPACE_ID,
changeSetId: changeSetId,
funcId: importFunc.id,
});
const funcName = importFuncResponse.data.name;
try {
const importResponse = await siApi.executeManagementFunction({
workspaceId: WORKSPACE_ID,
changeSetId,
componentId: result["componentId"],
executeManagementFunctionV1Request: {
managementFunction: { function: funcName },
},
});
let importState = "Pending";
const retrySleepInMs = 1000;
const retryMaxCount = 120;
let currentCount = 0;
const mgmtApi = new ManagementFuncsApi(apiConfig);
while (
(importState == "Pending" || importState == "Executing" ||
importState == "Operating") && currentCount <= retryMaxCount
) {
if (currentCount != 0) {
sleep(retrySleepInMs);
}
try {
const status = await mgmtApi.getManagementFuncRunState({
workspaceId: WORKSPACE_ID,
changeSetId,
managementFuncJobStateId:
importResponse.data.managementFuncJobStateId,
});
importState = status.data.state;
if (status.data.funcRunId) {
result["funcRunId"] = status.data.funcRunId;
}
currentCount += 1;
} catch (error) {
return errorResponse({
message: `error fetching management function state: ${
JSON.stringify(error, null, 2)
}`,
});
}
}
if (currentCount > retryMaxCount) {
return successResponse(
result,
"The import function is still in progress; use the funcRunId to find out more",
);
} else if (importState == "Failure") {
return errorResponse({
response: {
status: "failed",
data: result,
},
message:
`failed to import ${schemaName} with resourceId ${resourceId}; see funcRunId ${
result["funcRunId"]
} with the func-run-get tool for more information`,
});
} else {
return successResponse(result);
}
} catch (error) {
return errorResponse(error);
}
} catch (error) {
return errorResponse(error);
}
});
},
);
}
function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
