chore(internal): support local docs search in MCP servers

Author: stainless-app[bot]

packages/mcp-server/package.json

@@ -41,6 +41,7 @@
     "cors": "^2.8.5",
     "express": "^5.1.0",
     "fuse.js": "^7.1.0",
+    "minisearch": "^7.2.0",
     "jq-web": "https://github.com/stainless-api/jq-web/releases/download/v0.8.8/jq-web.tar.gz",
     "pino": "^10.3.1",
     "pino-http": "^11.0.0",

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

@@ -3,6 +3,7 @@
 import { Tool } from '@modelcontextprotocol/sdk/types.js';
 import { Metadata, McpRequestContext, asTextContentResult } from './types';
 import { getLogger } from './logger';
+import type { LocalDocsSearch } from './local-docs-search';
 
 export const metadata: Metadata = {
   resource: 'all',
@@ -43,20 +44,49 @@ export const tool: Tool = {
 const docsSearchURL =
   process.env['DOCS_SEARCH_URL'] || 'https://api.stainless.com/api/projects/scan-documents/docs/search';
 
-export const handler = async ({
-  reqContext,
-  args,
-}: {
-  reqContext: McpRequestContext;
-  args: Record<string, unknown> | undefined;
-}) => {
+let _localSearch: LocalDocsSearch | undefined;
+
+export function setLocalSearch(search: LocalDocsSearch): void {
+  _localSearch = search;
+}
+
+const SUPPORTED_LANGUAGES = new Set(['http', 'typescript', 'javascript']);
+
+async function searchLocal(args: Record<string, unknown>): Promise<unknown> {
+  if (!_localSearch) {
+    throw new Error('Local search not initialized');
+  }
+
+  const query = (args['query'] as string) ?? '';
+  const language = (args['language'] as string) ?? 'typescript';
+  const detail = (args['detail'] as string) ?? 'verbose';
+
+  if (!SUPPORTED_LANGUAGES.has(language)) {
+    throw new Error(
+      `Local docs search only supports HTTP, TypeScript, and JavaScript. Got language="${language}". ` +
+        `Use --docs-search-mode stainless-api for other languages, or set language to "http", "typescript", or "javascript".`,
+    );
+  }
+
+  return _localSearch.search({
+    query,
+    language,
+    detail,
+    maxResults: 10,
+  }).results;
+}
+
+async function searchRemote(
+  args: Record<string, unknown>,
+  stainlessApiKey: string | undefined,
+): Promise<unknown> {
   const body = args as any;
   const query = new URLSearchParams(body).toString();
 
   const startTime = Date.now();
   const result = await fetch(`${docsSearchURL}?${query}`, {
     headers: {
-      ...(reqContext.stainlessApiKey && { Authorization: reqContext.stainlessApiKey }),
+      ...(stainlessApiKey && { Authorization: stainlessApiKey }),
     },
   });
 
@@ -75,7 +105,7 @@ export const handler = async ({
       'Got error response from docs search tool',
     );
 
-    if (result.status === 404 && !reqContext.stainlessApiKey) {
+    if (result.status === 404 && !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.',
       );
@@ -94,7 +124,23 @@ export const handler = async ({
     },
     'Got docs search result',
   );
-  return asTextContentResult(resultBody);
+  return resultBody;
+}
+
+export const handler = async ({
+  reqContext,
+  args,
+}: {
+  reqContext: McpRequestContext;
+  args: Record<string, unknown> | undefined;
+}) => {
+  const body = args ?? {};
+
+  if (_localSearch) {
+    return asTextContentResult(await searchLocal(body));
+  }
+
+  return asTextContentResult(await searchRemote(body, reqContext.stainlessApiKey));
 };
 
 export default { metadata, tool, handler };

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

@@ -18,6 +18,8 @@ export type McpOptions = {
   includeCodeTool?: boolean | undefined;
   includeDocsTools?: boolean | undefined;
   stainlessApiKey?: string | undefined;
+  docsSearchMode?: 'stainless-api' | 'local' | undefined;
+  docsDir?: string | undefined;
   codeAllowHttpGets?: boolean | undefined;
   codeAllowedMethods?: string[] | undefined;
   codeBlockedMethods?: string[] | undefined;
@@ -58,6 +60,18 @@ export function parseCLIOptions(): CLIOptions {
       description: 'Path to custom instructions for the MCP server',
     })
     .option('debug', { type: 'boolean', description: 'Enable debug logging' })
+    .option('docs-dir', {
+      type: 'string',
+      description:
+        'Path to a directory of local documentation files (markdown/JSON) to include in local docs search.',
+    })
+    .option('docs-search-mode', {
+      type: 'string',
+      choices: ['stainless-api', 'local'],
+      default: 'stainless-api',
+      description:
+        "Where to search documentation; 'stainless-api' uses the Stainless-hosted search API whereas 'local' uses an in-memory search index built from embedded SDK method data and optional local docs files.",
+    })
     .option('log-format', {
       type: 'string',
       choices: ['json', 'pretty'],
@@ -118,6 +132,8 @@ export function parseCLIOptions(): CLIOptions {
     ...(includeDocsTools !== undefined && { includeDocsTools }),
     debug: !!argv.debug,
     stainlessApiKey: argv.stainlessApiKey,
+    docsSearchMode: argv.docsSearchMode as 'stainless-api' | 'local' | undefined,
+    docsDir: argv.docsDir,
     codeAllowHttpGets: argv.codeAllowHttpGets,
     codeAllowedMethods: argv.codeAllowedMethods,
     codeBlockedMethods: argv.codeBlockedMethods,
@@ -163,5 +179,7 @@ export function parseQueryOptions(defaultOptions: McpOptions, query: unknown): M
     ...(codeTool !== undefined && { includeCodeTool: codeTool }),
     ...(docsTools !== undefined && { includeDocsTools: docsTools }),
     codeExecutionMode: defaultOptions.codeExecutionMode,
+    docsSearchMode: defaultOptions.docsSearchMode,
+    docsDir: defaultOptions.docsDir,
   };
 }

packages/mcp-server/src/options.ts

@@ -11,6 +11,8 @@ import { ClientOptions } from 'scan-documents';
 import ScanDocuments from 'scan-documents';
 import { codeTool } from './code-tool';
 import docsSearchTool from './docs-search-tool';
+import { setLocalSearch } from './docs-search-tool';
+import { LocalDocsSearch } from './local-docs-search';
 import { getInstructions } from './instructions';
 import { McpOptions } from './options';
 import { blockedMethodsForCodeTool } from './methods';
@@ -62,6 +64,12 @@ export async function initMcpServer(params: {
     error: logAtLevel('error'),
   };
 
+  if (params.mcpOptions?.docsSearchMode === 'local') {
+    const docsDir = params.mcpOptions?.docsDir;
+    const localSearch = await LocalDocsSearch.create(docsDir ? { docsDir } : undefined);
+    setLocalSearch(localSearch);
+  }
+
   let _client: ScanDocuments | undefined;
   let _clientError: Error | undefined;
   let _logLevel: 'debug' | 'info' | 'warn' | 'error' | 'off' | undefined;