chore(internal): allow setting x-stainless-api-key header on mcp server requests

Author: stainless-app[bot]

packages/mcp-server/src/auth.ts

@@ -2,9 +2,24 @@
 
 import { IncomingMessage } from 'node:http';
 import { ClientOptions } from 'cas-parser-node';
+import { McpOptions } from './options';
 
-export const parseAuthHeaders = (req: IncomingMessage, required?: boolean): Partial<ClientOptions> => {
+export const parseClientAuthHeaders = (req: IncomingMessage, required?: boolean): Partial<ClientOptions> => {
   const apiKey =
     Array.isArray(req.headers['x-api-key']) ? req.headers['x-api-key'][0] : req.headers['x-api-key'];
   return { apiKey };
 };
+
+export const getStainlessApiKey = (req: IncomingMessage, mcpOptions: McpOptions): string | undefined => {
+  // Try to get the key from the x-stainless-api-key header
+  const headerKey =
+    Array.isArray(req.headers['x-stainless-api-key']) ?
+      req.headers['x-stainless-api-key'][0]
+    : req.headers['x-stainless-api-key'];
+  if (headerKey && typeof headerKey === 'string') {
+    return headerKey;
+  }
+
+  // Fall back to value set in the mcpOptions (e.g. from environment variable), if provided
+  return mcpOptions.stainlessApiKey;
+};

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

@@ -1,11 +1,17 @@
 // File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-import { McpTool, Metadata, ToolCallResult, asErrorResult, asTextContentResult } from './types';
+import {
+  McpRequestContext,
+  McpTool,
+  Metadata,
+  ToolCallResult,
+  asErrorResult,
+  asTextContentResult,
+} from './types';
 import { Tool } from '@modelcontextprotocol/sdk/types.js';
 import { readEnv, requireValue } from './util';
 import { WorkerInput, WorkerOutput } from './code-tool-types';
 import { SdkMethod } from './methods';
-import { CasParser } from 'cas-parser-node';
 
 const prompt = `Runs JavaScript code to interact with the Cas Parser API.
 
@@ -36,7 +42,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(params: { blockedMethods: SdkMethod[] | undefined }): McpTool {
+export function codeTool({ blockedMethods }: { blockedMethods: SdkMethod[] | undefined }): McpTool {
   const metadata: Metadata = { resource: 'all', operation: 'write', tags: [] };
   const tool: Tool = {
     name: 'execute',
@@ -56,19 +62,24 @@ export function codeTool(params: { blockedMethods: SdkMethod[] | undefined }): M
       required: ['code'],
     },
   };
-  const handler = async (client: CasParser, args: any): Promise<ToolCallResult> => {
+  const handler = 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;
 
     // 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 (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
@@ -78,16 +89,14 @@ export function codeTool(params: { blockedMethods: SdkMethod[] | undefined }): M
       }
     }
 
-    // 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';
 
+    // Setting a Stainless API key authenticates requests to the code tool endpoint.
     const res = await fetch(codeModeEndpoint, {
       method: 'POST',
       headers: {
-        ...(stainlessAPIKey && { Authorization: stainlessAPIKey }),
+        ...(reqContext.stainlessApiKey && { Authorization: reqContext.stainlessApiKey }),
         'Content-Type': 'application/json',
         client_envs: JSON.stringify({
           CAS_PARSER_API_KEY: requireValue(

packages/mcp-server/src/docs-search-tool.ts

@@ -1,8 +1,6 @@
 // File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-import { Metadata, asTextContentResult } from './types';
-import { readEnv } from './util';
-
+import { Metadata, McpRequestContext, asTextContentResult } from './types';
 import { Tool } from '@modelcontextprotocol/sdk/types.js';
 
 export const metadata: Metadata = {
@@ -43,13 +41,18 @@ export const tool: Tool = {
 const docsSearchURL =
   process.env['DOCS_SEARCH_URL'] || 'https://api.stainless.com/api/projects/cas-parser/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 stainlessAPIKey = readEnv('STAINLESS_API_KEY');
   const result = await fetch(`${docsSearchURL}?${query}`, {
     headers: {
-      ...(stainlessAPIKey && { Authorization: stainlessAPIKey }),
+      ...(reqContext.stainlessApiKey && { Authorization: reqContext.stainlessApiKey }),
     },
   });
 

packages/mcp-server/src/http.ts

@@ -6,7 +6,7 @@ import { ClientOptions } from 'cas-parser-node';
 import express from 'express';
 import morgan from 'morgan';
 import morganBody from 'morgan-body';
-import { parseAuthHeaders } from './auth';
+import { getStainlessApiKey, parseClientAuthHeaders } from './auth';
 import { McpOptions } from './options';
 import { initMcpServer, newMcpServer } from './server';
 
@@ -21,17 +21,20 @@ const newServer = async ({
   req: express.Request;
   res: express.Response;
 }): Promise<McpServer | null> => {
-  const server = await newMcpServer();
+  const stainlessApiKey = getStainlessApiKey(req, mcpOptions);
+  const server = await newMcpServer(stainlessApiKey);
 
   try {
-    const authOptions = parseAuthHeaders(req, false);
+    const authOptions = parseClientAuthHeaders(req, false);
+
     await initMcpServer({
       server: server,
       mcpOptions: mcpOptions,
       clientOptions: {
         ...clientOptions,
         ...authOptions,
       },
+      stainlessApiKey: stainlessApiKey,
     });
   } catch (error) {
     res.status(401).json({
@@ -112,20 +115,24 @@ export const streamableHTTPApp = ({
   return app;
 };
 
-export const launchStreamableHTTPServer = async (params: {
+export const launchStreamableHTTPServer = async ({
+  mcpOptions,
+  debug,
+  port,
+}: {
   mcpOptions: McpOptions;
   debug: boolean;
   port: number | string | undefined;
 }) => {
-  const app = streamableHTTPApp({ mcpOptions: params.mcpOptions, debug: params.debug });
-  const server = app.listen(params.port);
+  const app = streamableHTTPApp({ mcpOptions, debug });
+  const server = app.listen(port);
   const address = server.address();
 
   if (typeof address === 'string') {
     console.error(`MCP Server running on streamable HTTP at ${address}`);
   } else if (address !== null) {
     console.error(`MCP Server running on streamable HTTP on port ${address.port}`);
   } else {
-    console.error(`MCP Server running on streamable HTTP on port ${params.port}`);
+    console.error(`MCP Server running on streamable HTTP on port ${port}`);
   }
 };

packages/mcp-server/src/options.ts

@@ -4,6 +4,7 @@ import qs from 'qs';
 import yargs from 'yargs';
 import { hideBin } from 'yargs/helpers';
 import z from 'zod';
+import { readEnv } from './util';
 
 export type CLIOptions = McpOptions & {
   debug: boolean;
@@ -14,6 +15,7 @@ export type CLIOptions = McpOptions & {
 
 export type McpOptions = {
   includeDocsTools?: boolean | undefined;
+  stainlessApiKey?: string | undefined;
   codeAllowHttpGets?: boolean | undefined;
   codeAllowedMethods?: string[] | undefined;
   codeBlockedMethods?: string[] | undefined;
@@ -51,6 +53,12 @@ export function parseCLIOptions(): CLIOptions {
       description: 'Port to serve on if using http transport',
     })
     .option('socket', { type: 'string', description: 'Unix socket to serve on if using http transport' })
+    .option('stainless-api-key', {
+      type: 'string',
+      default: readEnv('STAINLESS_API_KEY'),
+      description:
+        'API key for Stainless. Used to authenticate requests to Stainless-hosted tools endpoints.',
+    })
     .option('tools', {
       type: 'string',
       array: true,
@@ -81,6 +89,7 @@ export function parseCLIOptions(): CLIOptions {
   return {
     ...(includeDocsTools !== undefined && { includeDocsTools }),
     debug: !!argv.debug,
+    stainlessApiKey: argv.stainlessApiKey,
     codeAllowHttpGets: argv.codeAllowHttpGets,
     codeAllowedMethods: argv.codeAllowedMethods,
     codeBlockedMethods: argv.codeBlockedMethods,

packages/mcp-server/src/server.ts

@@ -13,17 +13,17 @@ import { codeTool } from './code-tool';
 import docsSearchTool from './docs-search-tool';
 import { McpOptions } from './options';
 import { blockedMethodsForCodeTool } from './methods';
-import { HandlerFunction, McpTool } from './types';
+import { HandlerFunction, McpRequestContext, ToolCallResult, McpTool } from './types';
 import { readEnv } from './util';
 
-async function getInstructions() {
-  // This API key is optional; providing it allows the server to fetch instructions for unreleased versions.
-  const stainlessAPIKey = readEnv('STAINLESS_API_KEY');
+async function getInstructions(stainlessApiKey: string | undefined): Promise<string> {
+  // Setting the stainless API key is optional, but may be required
+  // to authenticate requests to the Stainless API.
   const response = await fetch(
     readEnv('CODE_MODE_INSTRUCTIONS_URL') ?? 'https://api.stainless.com/api/ai/instructions/cas-parser',
     {
       method: 'GET',
-      headers: { ...(stainlessAPIKey && { Authorization: stainlessAPIKey }) },
+      headers: { ...(stainlessApiKey && { Authorization: stainlessApiKey }) },
     },
   );
 
@@ -52,14 +52,14 @@ async function getInstructions() {
   return instructions;
 }
 
-export const newMcpServer = async () =>
+export const newMcpServer = async (stainlessApiKey: string | undefined) =>
   new McpServer(
     {
       name: 'cas_parser_node_api',
       version: '1.7.1',
     },
     {
-      instructions: await getInstructions(),
+      instructions: await getInstructions(stainlessApiKey),
       capabilities: { tools: {}, logging: {} },
     },
   );
@@ -72,6 +72,7 @@ export async function initMcpServer(params: {
   server: Server | McpServer;
   clientOptions?: ClientOptions;
   mcpOptions?: McpOptions;
+  stainlessApiKey?: string | undefined;
 }) {
   const server = params.server instanceof McpServer ? params.server.server : params.server;
 
@@ -116,7 +117,14 @@ export async function initMcpServer(params: {
       throw new Error(`Unknown tool: ${name}`);
     }
 
-    return executeHandler(mcpTool.handler, client, args);
+    return executeHandler({
+      handler: mcpTool.handler,
+      reqContext: {
+        client,
+        stainlessApiKey: params.stainlessApiKey ?? params.mcpOptions?.stainlessApiKey,
+      },
+      args,
+    });
   });
 
   server.setRequestHandler(SetLevelRequestSchema, async (request) => {
@@ -161,10 +169,14 @@ export function selectTools(options?: McpOptions): McpTool[] {
 /**
  * Runs the provided handler with the given client and arguments.
  */
-export async function executeHandler(
-  handler: HandlerFunction,
-  client: CasParser,
-  args: Record<string, unknown> | undefined,
-) {
-  return await handler(client, args || {});
+export async function executeHandler({
+  handler,
+  reqContext,
+  args,
+}: {
+  handler: HandlerFunction;
+  reqContext: McpRequestContext;
+  args: Record<string, unknown> | undefined;
+}): Promise<ToolCallResult> {
+  return await handler({ reqContext, args: args || {} });
 }

packages/mcp-server/src/stdio.ts

@@ -3,9 +3,9 @@ import { McpOptions } from './options';
 import { initMcpServer, newMcpServer } from './server';
 
 export const launchStdioServer = async (mcpOptions: McpOptions) => {
-  const server = await newMcpServer();
+  const server = await newMcpServer(mcpOptions.stainlessApiKey);
 
-  await initMcpServer({ server, mcpOptions });
+  await initMcpServer({ server, mcpOptions, stainlessApiKey: mcpOptions.stainlessApiKey });
 
   const transport = new StdioServerTransport();
   await server.connect(transport);

packages/mcp-server/src/types.ts

@@ -42,10 +42,18 @@ export type ToolCallResult = {
   isError?: boolean;
 };
 
-export type HandlerFunction = (
-  client: CasParser,
-  args: Record<string, unknown> | undefined,
-) => Promise<ToolCallResult>;
+export type McpRequestContext = {
+  client: CasParser;
+  stainlessApiKey?: string | undefined;
+};
+
+export type HandlerFunction = ({
+  reqContext,
+  args,
+}: {
+  reqContext: McpRequestContext;
+  args: Record<string, unknown> | undefined;
+}) => Promise<ToolCallResult>;
 
 export function asTextContentResult(result: unknown): ToolCallResult {
   return {