chore(internal): allow basic filtering of methods allowed for MCP code mode

Author: stainless-app[bot]

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

@@ -4,6 +4,7 @@ import { McpTool, Metadata, ToolCallResult, asErrorResult, asTextContentResult }
 import { Tool } from '@modelcontextprotocol/sdk/types.js';
 import { readEnv, requireValue } from './server';
 import { WorkerInput, WorkerOutput } from './code-tool-types';
+import { SdkMethod } from './methods';
 import { ImageKit } from '@imagekit/nodejs';
 
 const prompt = `Runs JavaScript code to interact with the Image Kit API.
@@ -35,7 +36,7 @@ Variables will not persist between calls, so make sure to return or log any data
  *
  * @param endpoints - The endpoints to include in the list.
  */
-export function codeTool(): McpTool {
+export function codeTool(params: { blockedMethods: SdkMethod[] | undefined }): McpTool {
   const metadata: Metadata = { resource: 'all', operation: 'write', tags: [] };
   const tool: Tool = {
     name: 'execute',
@@ -59,6 +60,24 @@ export function codeTool(): McpTool {
     const code = args.code as string;
     const intent = args.intent as string | undefined;
 
+    // Do very basic blocking of code that includes forbidden method names.
+    //
+    // WARNING: This is not secure against obfuscation and other evasion methods. If
+    // stronger security blocks are required, then these should be enforced in the downstream
+    // API (e.g., by having users call the MCP server with API keys with limited permissions).
+    if (params.blockedMethods) {
+      const blockedMatches = params.blockedMethods.filter((method) =>
+        code.includes(method.fullyQualifiedName),
+      );
+      if (blockedMatches.length > 0) {
+        return asErrorResult(
+          `The following methods have been blocked by the MCP server and cannot be used in code execution: ${blockedMatches
+            .map((m) => m.fullyQualifiedName)
+            .join(', ')}`,
+        );
+      }
+    }
+
     // this is not required, but passing a Stainless API key for the matching project_name
     // will allow you to run code-mode queries against non-published versions of your SDK.
     const stainlessAPIKey = readEnv('STAINLESS_API_KEY');

packages/mcp-server/src/http.ts

@@ -11,10 +11,12 @@ import { parseAuthHeaders } from './headers';
 
 const newServer = async ({
   clientOptions,
+  mcpOptions,
   req,
   res,
 }: {
   clientOptions: ClientOptions;
+  mcpOptions: McpOptions;
   req: express.Request;
   res: express.Response;
 }): Promise<McpServer | null> => {
@@ -24,6 +26,7 @@ const newServer = async ({
     const authOptions = parseAuthHeaders(req, false);
     await initMcpServer({
       server: server,
+      mcpOptions: mcpOptions,
       clientOptions: {
         ...clientOptions,
         ...authOptions,

packages/mcp-server/src/index.ts

@@ -18,7 +18,7 @@ async function main() {
 
   switch (options.transport) {
     case 'stdio':
-      await launchStdioServer();
+      await launchStdioServer(options);
       break;
     case 'http':
       await launchStreamableHTTPServer({