Skip to content

Commit 497c926

Browse files
chore(internal): make MCP code execution location configurable via a flag
1 parent 4178900 commit 497c926

8 files changed

Lines changed: 620 additions & 70 deletions

File tree

packages/mcp-server/Dockerfile

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,21 @@ COPY . .
3737
RUN yarn install --frozen-lockfile && \
3838
yarn build
3939

40-
# Production stage
40+
FROM denoland/deno:bin-2.6.10 AS deno_installer
41+
FROM gcr.io/distroless/cc@sha256:66d87e170bc2c5e2b8cf853501141c3c55b4e502b8677595c57534df54a68cc5 AS cc
42+
4143
FROM node:24-alpine
4244

45+
# Install deno
46+
COPY --from=deno_installer /deno /usr/local/bin/deno
47+
48+
# Add in shared libraries needed by Deno
49+
COPY --from=cc --chown=root:root --chmod=755 /lib/*-linux-gnu/* /usr/local/lib/
50+
COPY --from=cc --chown=root:root --chmod=755 /lib/ld-linux-* /lib/
51+
52+
RUN mkdir /lib64 && ln -s /usr/local/lib/ld-linux-* /lib64/
53+
ENV LD_LIBRARY_PATH=/usr/lib:/usr/local/lib
54+
4355
# Add non-root user
4456
RUN addgroup -g 1001 -S nodejs && adduser -S nodejs -u 1001
4557

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
2+
3+
export const workerPath = require.resolve('./code-tool-worker.mjs');

packages/mcp-server/src/code-tool-types.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ export type WorkerInput = {
88
client_opts: ClientOptions;
99
intent?: string | undefined;
1010
};
11+
1112
export type WorkerOutput = {
1213
is_error: boolean;
1314
result: unknown | null;
Lines changed: 319 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,319 @@
1+
// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
2+
3+
import path from 'node:path';
4+
import util from 'node:util';
5+
import Fuse from 'fuse.js';
6+
import ts from 'typescript';
7+
import { WorkerOutput } from './code-tool-types';
8+
import { ImageKit, ClientOptions } from '@imagekit/nodejs';
9+
10+
function getRunFunctionSource(code: string): {
11+
type: 'declaration' | 'expression';
12+
client: string | undefined;
13+
code: string;
14+
} | null {
15+
const sourceFile = ts.createSourceFile('code.ts', code, ts.ScriptTarget.Latest, true);
16+
const printer = ts.createPrinter();
17+
18+
for (const statement of sourceFile.statements) {
19+
// Check for top-level function declarations
20+
if (ts.isFunctionDeclaration(statement)) {
21+
if (statement.name?.text === 'run') {
22+
return {
23+
type: 'declaration',
24+
client: statement.parameters[0]?.name.getText(),
25+
code: printer.printNode(ts.EmitHint.Unspecified, statement.body!, sourceFile),
26+
};
27+
}
28+
}
29+
30+
// Check for variable declarations: const run = () => {} or const run = function() {}
31+
if (ts.isVariableStatement(statement)) {
32+
for (const declaration of statement.declarationList.declarations) {
33+
if (
34+
ts.isIdentifier(declaration.name) &&
35+
declaration.name.text === 'run' &&
36+
// Check if it's initialized with a function
37+
declaration.initializer &&
38+
(ts.isFunctionExpression(declaration.initializer) || ts.isArrowFunction(declaration.initializer))
39+
) {
40+
return {
41+
type: 'expression',
42+
client: declaration.initializer.parameters[0]?.name.getText(),
43+
code: printer.printNode(ts.EmitHint.Unspecified, declaration.initializer, sourceFile),
44+
};
45+
}
46+
}
47+
}
48+
}
49+
50+
return null;
51+
}
52+
53+
function getTSDiagnostics(code: string): string[] {
54+
const functionSource = getRunFunctionSource(code)!;
55+
const codeWithImport = [
56+
'import { ImageKit } from "@imagekit/nodejs";',
57+
functionSource.type === 'declaration' ?
58+
`async function run(${functionSource.client}: ImageKit)`
59+
: `const run: (${functionSource.client}: ImageKit) => Promise<unknown> =`,
60+
functionSource.code,
61+
].join('\n');
62+
const sourcePath = path.resolve('code.ts');
63+
const ast = ts.createSourceFile(sourcePath, codeWithImport, ts.ScriptTarget.Latest, true);
64+
const options = ts.getDefaultCompilerOptions();
65+
options.target = ts.ScriptTarget.Latest;
66+
options.module = ts.ModuleKind.NodeNext;
67+
options.moduleResolution = ts.ModuleResolutionKind.NodeNext;
68+
const host = ts.createCompilerHost(options, true);
69+
const newHost: typeof host = {
70+
...host,
71+
getSourceFile: (...args) => {
72+
if (path.resolve(args[0]) === sourcePath) {
73+
return ast;
74+
}
75+
return host.getSourceFile(...args);
76+
},
77+
readFile: (...args) => {
78+
if (path.resolve(args[0]) === sourcePath) {
79+
return codeWithImport;
80+
}
81+
return host.readFile(...args);
82+
},
83+
fileExists: (...args) => {
84+
if (path.resolve(args[0]) === sourcePath) {
85+
return true;
86+
}
87+
return host.fileExists(...args);
88+
},
89+
};
90+
const program = ts.createProgram({
91+
options,
92+
rootNames: [sourcePath],
93+
host: newHost,
94+
});
95+
const diagnostics = ts.getPreEmitDiagnostics(program, ast);
96+
return diagnostics.map((d) => {
97+
const message = ts.flattenDiagnosticMessageText(d.messageText, '\n');
98+
if (!d.file || !d.start) return `- ${message}`;
99+
const { line: lineNumber } = ts.getLineAndCharacterOfPosition(d.file, d.start);
100+
const line = codeWithImport.split('\n').at(lineNumber)?.trim();
101+
return line ? `- ${message}\n ${line}` : `- ${message}`;
102+
});
103+
}
104+
105+
const fuse = new Fuse(
106+
[
107+
'client.customMetadataFields.create',
108+
'client.customMetadataFields.delete',
109+
'client.customMetadataFields.list',
110+
'client.customMetadataFields.update',
111+
'client.files.copy',
112+
'client.files.delete',
113+
'client.files.get',
114+
'client.files.move',
115+
'client.files.rename',
116+
'client.files.update',
117+
'client.files.upload',
118+
'client.files.bulk.addTags',
119+
'client.files.bulk.delete',
120+
'client.files.bulk.removeAITags',
121+
'client.files.bulk.removeTags',
122+
'client.files.versions.delete',
123+
'client.files.versions.get',
124+
'client.files.versions.list',
125+
'client.files.versions.restore',
126+
'client.files.metadata.get',
127+
'client.files.metadata.getFromURL',
128+
'client.savedExtensions.create',
129+
'client.savedExtensions.delete',
130+
'client.savedExtensions.get',
131+
'client.savedExtensions.list',
132+
'client.savedExtensions.update',
133+
'client.assets.list',
134+
'client.cache.invalidation.create',
135+
'client.cache.invalidation.get',
136+
'client.folders.copy',
137+
'client.folders.create',
138+
'client.folders.delete',
139+
'client.folders.move',
140+
'client.folders.rename',
141+
'client.folders.job.get',
142+
'client.accounts.usage.get',
143+
'client.accounts.origins.create',
144+
'client.accounts.origins.delete',
145+
'client.accounts.origins.get',
146+
'client.accounts.origins.list',
147+
'client.accounts.origins.update',
148+
'client.accounts.urlEndpoints.create',
149+
'client.accounts.urlEndpoints.delete',
150+
'client.accounts.urlEndpoints.get',
151+
'client.accounts.urlEndpoints.list',
152+
'client.accounts.urlEndpoints.update',
153+
'client.beta.v2.files.upload',
154+
'client.webhooks.unsafeUnwrap',
155+
'client.webhooks.unwrap',
156+
],
157+
{ threshold: 1, shouldSort: true },
158+
);
159+
160+
function getMethodSuggestions(fullyQualifiedMethodName: string): string[] {
161+
return fuse
162+
.search(fullyQualifiedMethodName)
163+
.map(({ item }) => item)
164+
.slice(0, 5);
165+
}
166+
167+
const proxyToObj = new WeakMap<any, any>();
168+
const objToProxy = new WeakMap<any, any>();
169+
170+
type ClientProxyConfig = {
171+
path: string[];
172+
isBelievedBad?: boolean;
173+
};
174+
175+
function makeSdkProxy<T extends object>(obj: T, { path, isBelievedBad = false }: ClientProxyConfig): T {
176+
let proxy: T = objToProxy.get(obj);
177+
178+
if (!proxy) {
179+
proxy = new Proxy(obj, {
180+
get(target, prop, receiver) {
181+
const propPath = [...path, String(prop)];
182+
const value = Reflect.get(target, prop, receiver);
183+
184+
if (isBelievedBad || (!(prop in target) && value === undefined)) {
185+
// If we're accessing a path that doesn't exist, it will probably eventually error.
186+
// Let's proxy it and mark it bad so that we can control the error message.
187+
// We proxy an empty class so that an invocation or construction attempt is possible.
188+
return makeSdkProxy(class {}, { path: propPath, isBelievedBad: true });
189+
}
190+
191+
if (value !== null && (typeof value === 'object' || typeof value === 'function')) {
192+
return makeSdkProxy(value, { path: propPath, isBelievedBad });
193+
}
194+
195+
return value;
196+
},
197+
198+
apply(target, thisArg, args) {
199+
if (isBelievedBad || typeof target !== 'function') {
200+
const fullyQualifiedMethodName = path.join('.');
201+
const suggestions = getMethodSuggestions(fullyQualifiedMethodName);
202+
throw new Error(
203+
`${fullyQualifiedMethodName} is not a function. Did you mean: ${suggestions.join(', ')}`,
204+
);
205+
}
206+
207+
return Reflect.apply(target, proxyToObj.get(thisArg) ?? thisArg, args);
208+
},
209+
210+
construct(target, args, newTarget) {
211+
if (isBelievedBad || typeof target !== 'function') {
212+
const fullyQualifiedMethodName = path.join('.');
213+
const suggestions = getMethodSuggestions(fullyQualifiedMethodName);
214+
throw new Error(
215+
`${fullyQualifiedMethodName} is not a constructor. Did you mean: ${suggestions.join(', ')}`,
216+
);
217+
}
218+
219+
return Reflect.construct(target, args, newTarget);
220+
},
221+
});
222+
223+
objToProxy.set(obj, proxy);
224+
proxyToObj.set(proxy, obj);
225+
}
226+
227+
return proxy;
228+
}
229+
230+
function parseError(code: string, error: unknown): string | undefined {
231+
if (!(error instanceof Error)) return;
232+
const message = error.name ? `${error.name}: ${error.message}` : error.message;
233+
try {
234+
// Deno uses V8; the first "<anonymous>:LINE:COLUMN" is the top of stack.
235+
const lineNumber = error.stack?.match(/<anonymous>:([0-9]+):[0-9]+/)?.[1];
236+
// -1 for the zero-based indexing
237+
const line =
238+
lineNumber &&
239+
code
240+
.split('\n')
241+
.at(parseInt(lineNumber, 10) - 1)
242+
?.trim();
243+
return line ? `${message}\n at line ${lineNumber}\n ${line}` : message;
244+
} catch {
245+
return message;
246+
}
247+
}
248+
249+
const fetch = async (req: Request): Promise<Response> => {
250+
const { opts, code } = (await req.json()) as { opts: ClientOptions; code: string };
251+
252+
const runFunctionSource = code ? getRunFunctionSource(code) : null;
253+
if (!runFunctionSource) {
254+
const message =
255+
code ?
256+
'The code is missing a top-level `run` function.'
257+
: 'The code argument is missing. Provide one containing a top-level `run` function.';
258+
return Response.json(
259+
{
260+
is_error: true,
261+
result: `${message} Write code within this template:\n\n\`\`\`\nasync function run(client) {\n // Fill this out\n}\n\`\`\``,
262+
log_lines: [],
263+
err_lines: [],
264+
} satisfies WorkerOutput,
265+
{ status: 400, statusText: 'Code execution error' },
266+
);
267+
}
268+
269+
const diagnostics = getTSDiagnostics(code);
270+
if (diagnostics.length > 0) {
271+
return Response.json(
272+
{
273+
is_error: true,
274+
result: `The code contains TypeScript diagnostics:\n${diagnostics.join('\n')}`,
275+
log_lines: [],
276+
err_lines: [],
277+
} satisfies WorkerOutput,
278+
{ status: 400, statusText: 'Code execution error' },
279+
);
280+
}
281+
282+
const client = new ImageKit({
283+
...opts,
284+
});
285+
286+
const log_lines: string[] = [];
287+
const err_lines: string[] = [];
288+
const console = {
289+
log: (...args: unknown[]) => {
290+
log_lines.push(util.format(...args));
291+
},
292+
error: (...args: unknown[]) => {
293+
err_lines.push(util.format(...args));
294+
},
295+
};
296+
try {
297+
let run_ = async (client: any) => {};
298+
eval(`${code}\nrun_ = run;`);
299+
const result = await run_(makeSdkProxy(client, { path: ['client'] }));
300+
return Response.json({
301+
is_error: false,
302+
result,
303+
log_lines,
304+
err_lines,
305+
} satisfies WorkerOutput);
306+
} catch (e) {
307+
return Response.json(
308+
{
309+
is_error: true,
310+
result: parseError(code, e),
311+
log_lines,
312+
err_lines,
313+
} satisfies WorkerOutput,
314+
{ status: 400, statusText: 'Code execution error' },
315+
);
316+
}
317+
};
318+
319+
export default { fetch };

0 commit comments

Comments
 (0)