feat(mcp): add typescript check to code execution tool

Author: stainless-app[bot]

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

@@ -1,5 +1,6 @@
 // File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
+import path from 'node:path';
 import util from 'node:util';
 
 import Fuse from 'fuse.js';
@@ -8,30 +9,41 @@ import ts from 'typescript';
 import { WorkerInput, WorkerSuccess, WorkerError } from './code-tool-types';
 import { Isaacus } from 'isaacus';
 
-function getRunFunctionNode(
-  code: string,
-): ts.FunctionDeclaration | ts.FunctionExpression | ts.ArrowFunction | null {
+function getRunFunctionSource(code: string): {
+  type: 'declaration' | 'expression';
+  client: string | undefined;
+  code: string;
+} | null {
   const sourceFile = ts.createSourceFile('code.ts', code, ts.ScriptTarget.Latest, true);
+  const printer = ts.createPrinter();
 
   for (const statement of sourceFile.statements) {
     // Check for top-level function declarations
     if (ts.isFunctionDeclaration(statement)) {
       if (statement.name?.text === 'run') {
-        return statement;
+        return {
+          type: 'declaration',
+          client: statement.parameters[0]?.name.getText(),
+          code: printer.printNode(ts.EmitHint.Unspecified, statement.body!, sourceFile),
+        };
       }
     }
 
     // Check for variable declarations: const run = () => {} or const run = function() {}
     if (ts.isVariableStatement(statement)) {
       for (const declaration of statement.declarationList.declarations) {
-        if (ts.isIdentifier(declaration.name) && declaration.name.text === 'run') {
+        if (
+          ts.isIdentifier(declaration.name) &&
+          declaration.name.text === 'run' &&
           // Check if it's initialized with a function
-          if (
-            declaration.initializer &&
-            (ts.isFunctionExpression(declaration.initializer) || ts.isArrowFunction(declaration.initializer))
-          ) {
-            return declaration.initializer;
-          }
+          declaration.initializer &&
+          (ts.isFunctionExpression(declaration.initializer) || ts.isArrowFunction(declaration.initializer))
+        ) {
+          return {
+            type: 'expression',
+            client: declaration.initializer.parameters[0]?.name.getText(),
+            code: printer.printNode(ts.EmitHint.Unspecified, declaration.initializer, sourceFile),
+          };
         }
       }
     }
@@ -40,6 +52,61 @@ function getRunFunctionNode(
   return null;
 }
 
+function getTSDiagnostics(code: string): string[] {
+  const functionSource = getRunFunctionSource(code)!;
+  const codeWithImport = [
+    'import { Isaacus } from "isaacus";',
+    functionSource.type === 'declaration' ?
+      `async function run(${functionSource.client}: Isaacus)`
+    : `const run: (${functionSource.client}: Isaacus) => Promise<unknown> =`,
+    functionSource.code,
+  ].join('\n');
+  const sourcePath = path.resolve('code.ts');
+  const ast = ts.createSourceFile(sourcePath, codeWithImport, ts.ScriptTarget.Latest, true);
+  const options = ts.getDefaultCompilerOptions();
+  options.target = ts.ScriptTarget.Latest;
+  options.module = ts.ModuleKind.NodeNext;
+  options.moduleResolution = ts.ModuleResolutionKind.NodeNext;
+  const host = ts.createCompilerHost(options, true);
+  const newHost: typeof host = {
+    ...host,
+    getSourceFile: (...args) => {
+      if (path.resolve(args[0]) === sourcePath) {
+        return ast;
+      }
+      return host.getSourceFile(...args);
+    },
+    readFile: (...args) => {
+      if (path.resolve(args[0]) === sourcePath) {
+        return codeWithImport;
+      }
+      return host.readFile(...args);
+    },
+    fileExists: (...args) => {
+      if (path.resolve(args[0]) === sourcePath) {
+        return true;
+      }
+      return host.fileExists(...args);
+    },
+  };
+  const program = ts.createProgram({
+    options,
+    rootNames: [sourcePath],
+    host: newHost,
+  });
+  const diagnostics = ts.getPreEmitDiagnostics(program, ast);
+  return diagnostics.map((d) => {
+    const message = ts.flattenDiagnosticMessageText(d.messageText, '\n');
+    if (!d.file || !d.start) return `- ${message}`;
+    const { line: tsLine } = ts.getLineAndCharacterOfPosition(d.file, d.start);
+    // We add two lines in the beginning, for the client import and the function declaration.
+    // So the actual (zero-based) line number is tsLine - 2.
+    const lineNumber = tsLine - 2;
+    const line = code.split('\n').at(lineNumber)?.trim();
+    return line ? `- ${message}\n    at line ${lineNumber + 1}\n      ${line}` : `- ${message}`;
+  });
+}
+
 const fuse = new Fuse(
   [
     'client.embeddings.create',
@@ -141,24 +208,28 @@ function parseError(code: string, error: unknown): string | undefined {
 
 const fetch = async (req: Request): Promise<Response> => {
   const { opts, code } = (await req.json()) as WorkerInput;
-  if (code == null) {
+
+  const runFunctionSource = code ? getRunFunctionSource(code) : null;
+  if (!runFunctionSource) {
+    const message =
+      code ?
+        'The code is missing a top-level `run` function.'
+      : 'The code argument is missing. Provide one containing a top-level `run` function.';
     return Response.json(
       {
-        message:
-          'The code param is missing. Provide one containing a top-level `run` function. Write code within this template:\n\n```\nasync function run(client) {\n  // Fill this out\n}\n```',
+        message: `${message} Write code within this template:\n\n\`\`\`\nasync function run(client) {\n  // Fill this out\n}\n\`\`\``,
         logLines: [],
         errLines: [],
       } satisfies WorkerError,
       { status: 400, statusText: 'Code execution error' },
     );
   }
 
-  const runFunctionNode = getRunFunctionNode(code);
-  if (!runFunctionNode) {
+  const diagnostics = getTSDiagnostics(code);
+  if (diagnostics.length > 0) {
     return Response.json(
       {
-        message:
-          'The code is missing a top-level `run` function. Write code within this template:\n\n```\nasync function run(client) {\n  // Fill this out\n}\n```',
+        message: `The code contains TypeScript diagnostics:\n${diagnostics.join('\n')}`,
         logLines: [],
         errLines: [],
       } satisfies WorkerError,

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

@@ -1,7 +1,7 @@
 // File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-import { dirname } from 'node:path';
-import { pathToFileURL } from 'node:url';
+import path from 'node:path';
+import url from 'node:url';
 import Isaacus, { ClientOptions } from 'isaacus';
 import { ContentBlock, Endpoint, Metadata, ToolCallResult } from './tools/types';
 
@@ -35,18 +35,24 @@ export async function codeTool(): Promise<Endpoint> {
     const baseURLHostname = new URL(client.baseURL).hostname;
     const { code } = args as { code: string };
 
-    const worker = await newDenoHTTPWorker(pathToFileURL(workerPath), {
+    const allowRead = [
+      'code-tool-worker.mjs',
+      `${workerPath.replace(/([\/\\]node_modules)[\/\\].+$/, '$1')}/`,
+      path.resolve(path.dirname(workerPath), '..'),
+    ].join(',');
+
+    const worker = await newDenoHTTPWorker(url.pathToFileURL(workerPath), {
       runFlags: [
         `--node-modules-dir=manual`,
-        `--allow-read=code-tool-worker.mjs,${workerPath.replace(/([\/\\]node_modules)[\/\\].+$/, '$1')}/`,
+        `--allow-read=${allowRead}`,
         `--allow-net=${baseURLHostname}`,
         // Allow environment variables because instantiating the client will try to read from them,
         // even though they are not set.
         '--allow-env',
       ],
       printOutput: true,
       spawnOptions: {
-        cwd: dirname(workerPath),
+        cwd: path.dirname(workerPath),
       },
     });