@metronome/mcp 2.2.0 → 3.4.1

package/src/code-tool.ts

@@ -1,9 +1,57 @@
 // File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-import { Metadata, ToolCallResult, asTextContentResult } from './tools/types';
+import fs from 'node:fs';
+import path from 'node:path';
+import url from 'node:url';
+import { newDenoHTTPWorker } from '@valtown/deno-http-worker';
+import { workerPath } from './code-tool-paths.cjs';
+import {
+  ContentBlock,
+  McpRequestContext,
+  McpTool,
+  Metadata,
+  ToolCallResult,
+  asErrorResult,
+  asTextContentResult,
+} from './types';
 import { Tool } from '@modelcontextprotocol/sdk/types.js';
-import { readEnv } from './server';
-import { WorkerSuccess } from './code-tool-types';
+import { readEnv, requireValue } from './util';
+import { WorkerInput, WorkerOutput } from './code-tool-types';
+import { getLogger } from './logger';
+import { SdkMethod } from './methods';
+import { McpCodeExecutionMode } from './options';
+import { ClientOptions } from '@metronome/sdk';
+
+const prompt = `Runs JavaScript code to interact with the Metronome API.
+
+You are a skilled TypeScript programmer writing code to interface with the service.
+Define an async function named "run" that takes a single parameter of an initialized SDK client and it will be run.
+For example:
+
+\`\`\`
+async function run(client) {
+  await client.v1.usage.ingest({ usage: [{
+    transaction_id: '90e9401f-0f8c-4cd3-9a9f-d6beb56d8d72',
+    customer_id: 'team@example.com',
+    event_type: 'heartbeat',
+    timestamp: '2024-01-01T00:00:00Z',
+    properties: {
+    cluster_id: '42',
+    cpu_seconds: 60,
+    region: 'Europe',
+  },
+  }] })
+}
+\`\`\`
+
+You will be returned anything that your function returns, plus the results of any console.log statements.
+Do not add try-catch blocks for single API calls. The tool will handle errors for you.
+Do not add comments unless necessary for generating better code.
+Code will run in a container, and cannot interact with the network outside of the given SDK client.
+Variables will not persist between calls, so make sure to return or log any data you might need later.
+Remember that you are writing TypeScript code, so you need to be careful with your types.
+Always type dynamic key-value stores explicitly as Record<string, YourValueType> instead of {}.`;
+
 /**
  * A tool that runs code against a copy of the SDK.
  *
@@ -11,53 +59,335 @@ import { WorkerSuccess } from './code-tool-types';
  * we expose a single tool that can be used to search for endpoints by name, resource, operation, or tag, and then
  * a generic endpoint that can be used to invoke any endpoint with the provided arguments.
  *
- * @param endpoints - The endpoints to include in the list.
+ * @param blockedMethods - The methods to block for code execution. Blocking is done by simple string
+ * matching, so it is not secure against obfuscation. For stronger security, block in the downstream API
+ * with limited API keys.
+ * @param codeExecutionMode - Whether to execute code in a local Deno environment or in a remote
+ * sandbox environment hosted by Stainless.
  */
-export async function codeTool() {
+export function codeTool({
+  blockedMethods,
+  codeExecutionMode,
+}: {
+  blockedMethods: SdkMethod[] | undefined;
+  codeExecutionMode: McpCodeExecutionMode;
+}): McpTool {
   const metadata: Metadata = { resource: 'all', operation: 'write', tags: [] };
   const tool: Tool = {
     name: 'execute',
-    description:
-      'Runs JavaScript code to interact with the API.\n\nYou are a skilled programmer writing code to interface with the service.\nDefine an async function named "run" that takes a single parameter of an initialized SDK client and it will be run.\nWrite code within this template:\n\n```\nasync function run(client) {\n  // Fill this out\n}\n```\n\nYou will be returned anything that your function returns, plus the results of any console.log statements.\nIf any code triggers an error, the tool will return an error response, so you do not need to add error handling unless you want to output something more helpful than the raw error.\nIt is not necessary to add comments to code, unless by adding those comments you believe that you can generate better code.\nThis code will run in a container, and you will not be able to use fetch or otherwise interact with the network calls other than through the client you are given.\nAny variables you define won\'t live between successive uses of this call, so make sure to return or log any data you might need later.',
-    inputSchema: { type: 'object', properties: { code: { type: 'string' } } },
+    description: prompt,
+    inputSchema: {
+      type: 'object',
+      properties: {
+        code: {
+          type: 'string',
+          description: 'Code to execute.',
+        },
+        intent: {
+          type: 'string',
+          description: 'Task you are trying to perform. Used for improving the service.',
+        },
+      },
+      required: ['code'],
+    },
   };
-  const handler = async (_: unknown, args: any): Promise<ToolCallResult> => {
+
+  const logger = getLogger();
+
+  const handler = async ({
+    reqContext,
+    args,
+  }: {
+    reqContext: McpRequestContext;
+    args: any;
+  }): Promise<ToolCallResult> => {
     const code = args.code as string;
+    // 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 (blockedMethods) {
+      const blockedMatches = 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(', ')}`,
+        );
+      }
+    }
+
+    let result: ToolCallResult;
+    const startTime = Date.now();
 
-    // 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');
-    const codeModeEndpoint =
-      readEnv('CODE_MODE_ENDPOINT_URL') ?? 'https://api.stainless.com/api/ai/code-tool';
-
-    const res = await fetch(codeModeEndpoint, {
-      method: 'POST',
-      headers: {
-        ...(stainlessAPIKey && { Authorization: stainlessAPIKey }),
-        'Content-Type': 'application/json',
-        client_envs: JSON.stringify({
-          METRONOME_BEARER_TOKEN: readEnv('METRONOME_BEARER_TOKEN'),
-          METRONOME_WEBHOOK_SECRET: readEnv('METRONOME_WEBHOOK_SECRET'),
-          METRONOME_BASE_URL: readEnv('METRONOME_BASE_URL'),
-        }),
+    if (codeExecutionMode === 'local') {
+      logger.debug('Executing code in local Deno environment');
+      result = await localDenoHandler({ reqContext, args });
+    } else {
+      logger.debug('Executing code in remote Stainless environment');
+      result = await remoteStainlessHandler({ reqContext, args });
+    }
+
+    logger.info(
+      {
+        codeExecutionMode,
+        durationMs: Date.now() - startTime,
+        isError: result.isError,
+        contentRows: result.content?.length ?? 0,
       },
-      body: JSON.stringify({
-        project_name: 'metronome',
-        client_opts: {},
-        code,
+      'Got code tool execution result',
+    );
+    return result;
+  };
+
+  return { metadata, tool, handler };
+}
+
+const remoteStainlessHandler = async ({
+  reqContext,
+  args,
+}: {
+  reqContext: McpRequestContext;
+  args: any;
+}): Promise<ToolCallResult> => {
+  const code = args.code as string;
+  const intent = args.intent as string | undefined;
+  const client = reqContext.client;
+
+  const codeModeEndpoint = readEnv('CODE_MODE_ENDPOINT_URL') ?? 'https://api.stainless.com/api/ai/code-tool';
+
+  // Setting a Stainless API key authenticates requests to the code tool endpoint.
+  const res = await fetch(codeModeEndpoint, {
+    method: 'POST',
+    headers: {
+      ...(reqContext.stainlessApiKey && { Authorization: reqContext.stainlessApiKey }),
+      'Content-Type': 'application/json',
+      'x-stainless-mcp-client-envs': JSON.stringify({
+        METRONOME_BEARER_TOKEN: requireValue(
+          readEnv('METRONOME_BEARER_TOKEN') ?? client.bearerToken,
+          'set METRONOME_BEARER_TOKEN environment variable or provide bearerToken client option',
+        ),
+        METRONOME_WEBHOOK_SECRET: readEnv('METRONOME_WEBHOOK_SECRET') ?? client.webhookSecret ?? undefined,
+        METRONOME_BASE_URL: readEnv('METRONOME_BASE_URL') ?? client.baseURL ?? undefined,
       }),
-    });
+    },
+    body: JSON.stringify({
+      project_name: 'metronome',
+      code,
+      intent,
+      client_opts: {},
+    } satisfies WorkerInput),
+  });
 
-    if (!res.ok) {
+  if (!res.ok) {
+    if (res.status === 404 && !reqContext.stainlessApiKey) {
       throw new Error(
-        `${res.status}: ${
-          res.statusText
-        } error when trying to contact Code Tool server. Details: ${await res.text()}`,
+        'Could not access code tool for this project. You may need to provide a Stainless API key via the STAINLESS_API_KEY environment variable, the --stainless-api-key flag, or the x-stainless-api-key HTTP header.',
       );
     }
+    throw new Error(
+      `${res.status}: ${
+        res.statusText
+      } error when trying to contact Code Tool server. Details: ${await res.text()}`,
+    );
+  }
 
-    return asTextContentResult((await res.json()) as WorkerSuccess);
+  const { is_error, result, log_lines, err_lines } = (await res.json()) as WorkerOutput;
+  const hasLogs = log_lines.length > 0 || err_lines.length > 0;
+  const output = {
+    result,
+    ...(log_lines.length > 0 && { log_lines }),
+    ...(err_lines.length > 0 && { err_lines }),
   };
+  if (is_error) {
+    return asErrorResult(typeof result === 'string' && !hasLogs ? result : JSON.stringify(output, null, 2));
+  }
+  return asTextContentResult(output);
+};
 
-  return { metadata, tool, handler };
-}
+const localDenoHandler = async ({
+  reqContext,
+  args,
+}: {
+  reqContext: McpRequestContext;
+  args: unknown;
+}): Promise<ToolCallResult> => {
+  const client = reqContext.client;
+  const baseURLHostname = new URL(client.baseURL).hostname;
+  const { code } = args as { code: string };
+
+  let denoPath: string;
+
+  const packageRoot = path.resolve(path.dirname(workerPath), '..');
+  const packageNodeModulesPath = path.resolve(packageRoot, 'node_modules');
+
+  // Check if deno is in PATH
+  const { execSync } = await import('node:child_process');
+  try {
+    execSync('command -v deno', { stdio: 'ignore' });
+    denoPath = 'deno';
+  } catch {
+    try {
+      // Use deno binary in node_modules if it's found
+      const denoNodeModulesPath = path.resolve(packageNodeModulesPath, 'deno', 'bin.cjs');
+      await fs.promises.access(denoNodeModulesPath, fs.constants.X_OK);
+      denoPath = denoNodeModulesPath;
+    } catch {
+      return asErrorResult(
+        'Deno is required for code execution but was not found. ' +
+          'Install it from https://deno.land or run: npm install deno',
+      );
+    }
+  }
+
+  const allowReadPaths = [
+    'code-tool-worker.mjs',
+    `${workerPath.replace(/([\/\\]node_modules)[\/\\].+$/, '$1')}/`,
+    packageRoot,
+  ];
+
+  // Follow symlinks in node_modules to allow read access to workspace-linked packages
+  try {
+    const sdkPkgName = '@metronome/sdk';
+    const sdkDir = path.resolve(packageNodeModulesPath, sdkPkgName);
+    const realSdkDir = fs.realpathSync(sdkDir);
+    if (realSdkDir !== sdkDir) {
+      allowReadPaths.push(realSdkDir);
+    }
+  } catch {
+    // Ignore if symlink resolution fails
+  }
+
+  const allowRead = allowReadPaths.join(',');
+
+  const worker = await newDenoHTTPWorker(url.pathToFileURL(workerPath), {
+    denoExecutable: denoPath,
+    runFlags: [
+      `--node-modules-dir=manual`,
+      `--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: path.dirname(workerPath),
+    },
+  });
+
+  try {
+    const resp = await new Promise<Response>((resolve, reject) => {
+      worker.addEventListener('exit', (exitCode) => {
+        reject(new Error(`Worker exited with code ${exitCode}`));
+      });
+
+      const opts: ClientOptions = {
+        baseURL: client.baseURL,
+        bearerToken: client.bearerToken,
+        webhookSecret: client.webhookSecret,
+        defaultHeaders: {
+          'X-Stainless-MCP': 'true',
+        },
+      };
+
+      const req = worker.request(
+        'http://localhost',
+        {
+          headers: {
+            'content-type': 'application/json',
+          },
+          method: 'POST',
+        },
+        (resp) => {
+          const body: Uint8Array[] = [];
+          resp.on('error', (err) => {
+            reject(err);
+          });
+          resp.on('data', (chunk) => {
+            body.push(chunk);
+          });
+          resp.on('end', () => {
+            resolve(
+              new Response(Buffer.concat(body).toString(), {
+                status: resp.statusCode ?? 200,
+                headers: resp.headers as any,
+              }),
+            );
+          });
+        },
+      );
+
+      const body = JSON.stringify({
+        opts,
+        code,
+      });
+
+      req.write(body, (err) => {
+        if (err != null) {
+          reject(err);
+        }
+      });
+
+      req.end();
+    });
+
+    if (resp.status === 200) {
+      const { result, log_lines, err_lines } = (await resp.json()) as WorkerOutput;
+      const returnOutput: ContentBlock | null =
+        result == null ? null : (
+          {
+            type: 'text',
+            text: typeof result === 'string' ? result : JSON.stringify(result),
+          }
+        );
+      const logOutput: ContentBlock | null =
+        log_lines.length === 0 ?
+          null
+        : {
+            type: 'text',
+            text: log_lines.join('\n'),
+          };
+      const errOutput: ContentBlock | null =
+        err_lines.length === 0 ?
+          null
+        : {
+            type: 'text',
+            text: 'Error output:\n' + err_lines.join('\n'),
+          };
+      return {
+        content: [returnOutput, logOutput, errOutput].filter((block) => block !== null),
+      };
+    } else {
+      const { result, log_lines, err_lines } = (await resp.json()) as WorkerOutput;
+      const messageOutput: ContentBlock | null =
+        result == null ? null : (
+          {
+            type: 'text',
+            text: typeof result === 'string' ? result : JSON.stringify(result),
+          }
+        );
+      const logOutput: ContentBlock | null =
+        log_lines.length === 0 ?
+          null
+        : {
+            type: 'text',
+            text: log_lines.join('\n'),
+          };
+      const errOutput: ContentBlock | null =
+        err_lines.length === 0 ?
+          null
+        : {
+            type: 'text',
+            text: 'Error output:\n' + err_lines.join('\n'),
+          };
+      return {
+        content: [messageOutput, logOutput, errOutput].filter((block) => block !== null),
+        isError: true,
+      };
+    }
+  } finally {
+    worker.terminate();
+  }
+};

package/src/docs-search-tool.ts

@@ -1,8 +1,8 @@
 // File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-import { Metadata, asTextContentResult } from './tools/types';
-
 import { Tool } from '@modelcontextprotocol/sdk/types.js';
+import { Metadata, McpRequestContext, asTextContentResult } from './types';
+import { getLogger } from './logger';
 
 export const metadata: Metadata = {
   resource: 'all',
@@ -13,7 +13,8 @@ export const metadata: Metadata = {
 
 export const tool: Tool = {
   name: 'search_docs',
-  description: 'Search for documentation for how to use the client to interact with the API.',
+  description:
+    'Search SDK documentation to find methods, parameters, and usage examples for interacting with the API. Use this before writing code when you need to discover the right approach.',
   inputSchema: {
     type: 'object',
     properties: {
@@ -42,18 +43,58 @@ export const tool: Tool = {
 const docsSearchURL =
   process.env['DOCS_SEARCH_URL'] || 'https://api.stainless.com/api/projects/metronome/docs/search';
 
-export const handler = async (_: unknown, args: Record<string, unknown> | undefined) => {
+export const handler = async ({
+  reqContext,
+  args,
+}: {
+  reqContext: McpRequestContext;
+  args: Record<string, unknown> | undefined;
+}) => {
   const body = args as any;
   const query = new URLSearchParams(body).toString();
-  const result = await fetch(`${docsSearchURL}?${query}`);
+
+  const startTime = Date.now();
+  const result = await fetch(`${docsSearchURL}?${query}`, {
+    headers: {
+      ...(reqContext.stainlessApiKey && { Authorization: reqContext.stainlessApiKey }),
+    },
+  });
+
+  const logger = getLogger();
 
   if (!result.ok) {
+    const errorText = await result.text();
+    logger.warn(
+      {
+        durationMs: Date.now() - startTime,
+        query: body.query,
+        status: result.status,
+        statusText: result.statusText,
+        errorText,
+      },
+      'Got error response from docs search tool',
+    );
+
+    if (result.status === 404 && !reqContext.stainlessApiKey) {
+      throw new Error(
+        'Could not find docs for this project. You may need to provide a Stainless API key via the STAINLESS_API_KEY environment variable, the --stainless-api-key flag, or the x-stainless-api-key HTTP header.',
+      );
+    }
+
     throw new Error(
-      `${result.status}: ${result.statusText} when using doc search tool. Details: ${await result.text()}`,
+      `${result.status}: ${result.statusText} when using doc search tool. Details: ${errorText}`,
     );
   }
 
-  return asTextContentResult(await result.json());
+  const resultBody = await result.json();
+  logger.info(
+    {
+      durationMs: Date.now() - startTime,
+      query: body.query,
+    },
+    'Got docs search result',
+  );
+  return asTextContentResult(resultBody);
 };
 
 export default { metadata, tool, handler };