|
| 1 | +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; |
| 2 | +import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; |
| 3 | +import { z } from "zod"; |
| 4 | +import { |
| 5 | + DebugFuncsApi, |
| 6 | + DebugFuncsApiExecDebugFuncRequest, |
| 7 | + DebugFuncsApiGetDebugFuncStateRequest, |
| 8 | +} from "@systeminit/api-client"; |
| 9 | +import { |
| 10 | + errorResponse, |
| 11 | + generateDescription, |
| 12 | + successResponse, |
| 13 | + withAnalytics, |
| 14 | +} from "./commonBehavior.ts"; |
| 15 | +import { apiConfig, WORKSPACE_ID } from "../si_client.ts"; |
| 16 | + |
| 17 | +const name = "exec-debug-func"; |
| 18 | +const title = "Run a debug func inside a change set"; |
| 19 | +const description = |
| 20 | + `<description>Runs a one-off function inside a change set intended to troubleshoot problems building infrastructure in the cloud. Returns arbitrary json data that should provide information about problems encountered when building infrastructure. Requires a component that has credentials and other settings necessary for making API calls within the cloud environment. The component SHOULD be connected to READ ONLY CREDENTIALS to prevent making changes to infrastructure. Optionally accepts arbitrary JSON debugInput with any additional context required for troubleshooting the problem.</description>`; |
| 21 | + |
| 22 | +const debugFuncDescription = |
| 23 | + "<description>A typescript debug function. The function should *ALWAYS* have the form of `async function debug({ component, debugInput }) { /* debug code here */ }`. The function name should *ALWAYS* be `debug`</description>"; |
| 24 | + |
| 25 | +const execDebugFuncInputSchemaRaw = { |
| 26 | + changeSetId: z |
| 27 | + .string() |
| 28 | + .describe("The id of the change set in which this function should be run"), |
| 29 | + componentId: z |
| 30 | + .string() |
| 31 | + .describe( |
| 32 | + "The id of the component that is connected to the necessary credentials and other context for making the API calls needed for debugging.", |
| 33 | + ), |
| 34 | + debugInput: z |
| 35 | + .any() |
| 36 | + .optional() |
| 37 | + .describe( |
| 38 | + "Arbitrary JSON object input that will be passed to the debug function", |
| 39 | + ), |
| 40 | + debugFunc: z.string().describe(debugFuncDescription), |
| 41 | +}; |
| 42 | + |
| 43 | +const execDebugFuncInputSchema = z.object(execDebugFuncInputSchemaRaw); |
| 44 | + |
| 45 | +type DebugFuncInputSchema = z.infer<typeof execDebugFuncInputSchema>; |
| 46 | + |
| 47 | +const execDebugFuncOutputSchemaRaw = { |
| 48 | + status: z.enum(["success", "failure"]), |
| 49 | + errorMessage: z |
| 50 | + .string() |
| 51 | + .optional() |
| 52 | + .describe( |
| 53 | + "If the status is `failure`, the error message will contain information about what went wrong", |
| 54 | + ), |
| 55 | + output: z.any().describe("The debug output produced by the debug function"), |
| 56 | + funcRunId: z |
| 57 | + .string() |
| 58 | + .optional() |
| 59 | + .describe( |
| 60 | + "The func run id of the function, to be passed to the func-run-get tool if more details are necessary, such as the function logs.", |
| 61 | + ), |
| 62 | +}; |
| 63 | + |
| 64 | +const execDebugFuncOutputSchema = z.object(execDebugFuncOutputSchemaRaw); |
| 65 | + |
| 66 | +const POLL_WAIT_START_MS = 100; |
| 67 | +const MAX_BASE_WAIT_MS = 5000; |
| 68 | +const MAX_JITTERED_WAIT_MS = 15000; |
| 69 | +const MAX_POLLS = 1000; |
| 70 | +const JITTER = 0.3; // 30% jitter |
| 71 | + |
| 72 | +const SUCCESS = "Success"; |
| 73 | +const FAILURE = "Failure"; |
| 74 | + |
| 75 | +// Honk, shoo, mimimimimi |
| 76 | +const zzz = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); |
| 77 | + |
| 78 | +export function execDebugFunc(server: McpServer) { |
| 79 | + server.registerTool( |
| 80 | + name, |
| 81 | + { |
| 82 | + title, |
| 83 | + description: generateDescription( |
| 84 | + description, |
| 85 | + "execDebugFuncResponse", |
| 86 | + execDebugFuncOutputSchema, |
| 87 | + ), |
| 88 | + inputSchema: execDebugFuncInputSchemaRaw, |
| 89 | + outputSchema: execDebugFuncOutputSchemaRaw, |
| 90 | + }, |
| 91 | + async ({ |
| 92 | + changeSetId, |
| 93 | + componentId, |
| 94 | + debugInput, |
| 95 | + debugFunc, |
| 96 | + }: DebugFuncInputSchema): Promise<CallToolResult> => |
| 97 | + await withAnalytics(name, async () => { |
| 98 | + const debugFuncPattern = /async\s+function\s+debug/; |
| 99 | + if (!debugFuncPattern.test(debugFunc)) { |
| 100 | + return errorResponse({ |
| 101 | + message: |
| 102 | + "Debug function must have the form 'async function debug({ component, debugInput }) { ... }'", |
| 103 | + }); |
| 104 | + } |
| 105 | + |
| 106 | + const debugFuncsApi = new DebugFuncsApi(apiConfig); |
| 107 | + const execFuncRequest: DebugFuncsApiExecDebugFuncRequest = { |
| 108 | + workspaceId: WORKSPACE_ID, |
| 109 | + changeSetId, |
| 110 | + execDebugFuncV1Request: { |
| 111 | + code: debugFunc, |
| 112 | + name: `debug-func-${componentId}-${changeSetId}`, |
| 113 | + handler: "debug", |
| 114 | + componentId, |
| 115 | + debugInput, |
| 116 | + }, |
| 117 | + }; |
| 118 | + |
| 119 | + // Wait for the function to have a chance to run |
| 120 | + await zzz(POLL_WAIT_START_MS); |
| 121 | + |
| 122 | + try { |
| 123 | + const execResponse = await debugFuncsApi.execDebugFunc( |
| 124 | + execFuncRequest, |
| 125 | + ); |
| 126 | + const debugFuncJobStateId = execResponse.data.debugFuncJobStateId; |
| 127 | + |
| 128 | + let pollCount = 0; |
| 129 | + let currentWaitMs = POLL_WAIT_START_MS; |
| 130 | + let funcRunId: string | null | undefined; |
| 131 | + |
| 132 | + while (pollCount < MAX_POLLS) { |
| 133 | + const jobStateRequest: DebugFuncsApiGetDebugFuncStateRequest = { |
| 134 | + workspaceId: WORKSPACE_ID, |
| 135 | + changeSetId, |
| 136 | + debugFuncJobStateId, |
| 137 | + }; |
| 138 | + |
| 139 | + const { data } = await debugFuncsApi.getDebugFuncState( |
| 140 | + jobStateRequest, |
| 141 | + ); |
| 142 | + funcRunId = data.funcRunId; |
| 143 | + |
| 144 | + if (data.state === SUCCESS || data.state === FAILURE) { |
| 145 | + const errorMessage = data.state === FAILURE |
| 146 | + ? data.failure || "Debug function execution failed" |
| 147 | + : undefined; |
| 148 | + return successResponse({ |
| 149 | + status: data.state, |
| 150 | + output: data.result, |
| 151 | + errorMessage, |
| 152 | + funcRunId, |
| 153 | + }); |
| 154 | + } |
| 155 | + |
| 156 | + pollCount++; |
| 157 | + |
| 158 | + if (pollCount >= MAX_POLLS) { |
| 159 | + const message = funcRunId |
| 160 | + ? `Function execution timed out but you can check for results or logs with the FuncRunId: ${funcRunId}` |
| 161 | + : "Function execution timed out"; |
| 162 | + return errorResponse({ |
| 163 | + message, |
| 164 | + }); |
| 165 | + } |
| 166 | + |
| 167 | + const jitter = Math.random() * JITTER; |
| 168 | + const waitMs = Math.min( |
| 169 | + currentWaitMs * (1 + jitter), |
| 170 | + MAX_JITTERED_WAIT_MS, |
| 171 | + ); |
| 172 | + await zzz(waitMs); |
| 173 | + |
| 174 | + currentWaitMs = Math.min(currentWaitMs * 2, MAX_BASE_WAIT_MS); |
| 175 | + } |
| 176 | + |
| 177 | + return errorResponse({ |
| 178 | + message: `Function execution timed out after ${MAX_POLLS} polls`, |
| 179 | + }); |
| 180 | + } catch (error) { |
| 181 | + return errorResponse({ |
| 182 | + message: `Failed to execute debug function: ${error}`, |
| 183 | + }); |
| 184 | + } |
| 185 | + }), |
| 186 | + ); |
| 187 | +} |
0 commit comments