Skip to content

Commit ccc0fcd

Browse files
chore(internal): support local docs search in MCP servers
1 parent 8dea8bd commit ccc0fcd

5 files changed

Lines changed: 634 additions & 10 deletions

File tree

packages/mcp-server/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
"cors": "^2.8.5",
4242
"express": "^5.1.0",
4343
"fuse.js": "^7.1.0",
44+
"minisearch": "^7.2.0",
4445
"jq-web": "https://github.com/stainless-api/jq-web/releases/download/v0.8.8/jq-web.tar.gz",
4546
"pino": "^10.3.1",
4647
"pino-http": "^11.0.0",

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

Lines changed: 56 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import { Tool } from '@modelcontextprotocol/sdk/types.js';
44
import { Metadata, McpRequestContext, asTextContentResult } from './types';
55
import { getLogger } from './logger';
6+
import type { LocalDocsSearch } from './local-docs-search';
67

78
export const metadata: Metadata = {
89
resource: 'all',
@@ -43,20 +44,49 @@ export const tool: Tool = {
4344
const docsSearchURL =
4445
process.env['DOCS_SEARCH_URL'] || 'https://api.stainless.com/api/projects/scan-documents/docs/search';
4546

46-
export const handler = async ({
47-
reqContext,
48-
args,
49-
}: {
50-
reqContext: McpRequestContext;
51-
args: Record<string, unknown> | undefined;
52-
}) => {
47+
let _localSearch: LocalDocsSearch | undefined;
48+
49+
export function setLocalSearch(search: LocalDocsSearch): void {
50+
_localSearch = search;
51+
}
52+
53+
const SUPPORTED_LANGUAGES = new Set(['http', 'typescript', 'javascript']);
54+
55+
async function searchLocal(args: Record<string, unknown>): Promise<unknown> {
56+
if (!_localSearch) {
57+
throw new Error('Local search not initialized');
58+
}
59+
60+
const query = (args['query'] as string) ?? '';
61+
const language = (args['language'] as string) ?? 'typescript';
62+
const detail = (args['detail'] as string) ?? 'verbose';
63+
64+
if (!SUPPORTED_LANGUAGES.has(language)) {
65+
throw new Error(
66+
`Local docs search only supports HTTP, TypeScript, and JavaScript. Got language="${language}". ` +
67+
`Use --docs-search-mode stainless-api for other languages, or set language to "http", "typescript", or "javascript".`,
68+
);
69+
}
70+
71+
return _localSearch.search({
72+
query,
73+
language,
74+
detail,
75+
maxResults: 10,
76+
}).results;
77+
}
78+
79+
async function searchRemote(
80+
args: Record<string, unknown>,
81+
stainlessApiKey: string | undefined,
82+
): Promise<unknown> {
5383
const body = args as any;
5484
const query = new URLSearchParams(body).toString();
5585

5686
const startTime = Date.now();
5787
const result = await fetch(`${docsSearchURL}?${query}`, {
5888
headers: {
59-
...(reqContext.stainlessApiKey && { Authorization: reqContext.stainlessApiKey }),
89+
...(stainlessApiKey && { Authorization: stainlessApiKey }),
6090
},
6191
});
6292

@@ -75,7 +105,7 @@ export const handler = async ({
75105
'Got error response from docs search tool',
76106
);
77107

78-
if (result.status === 404 && !reqContext.stainlessApiKey) {
108+
if (result.status === 404 && !stainlessApiKey) {
79109
throw new Error(
80110
'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.',
81111
);
@@ -94,7 +124,23 @@ export const handler = async ({
94124
},
95125
'Got docs search result',
96126
);
97-
return asTextContentResult(resultBody);
127+
return resultBody;
128+
}
129+
130+
export const handler = async ({
131+
reqContext,
132+
args,
133+
}: {
134+
reqContext: McpRequestContext;
135+
args: Record<string, unknown> | undefined;
136+
}) => {
137+
const body = args ?? {};
138+
139+
if (_localSearch) {
140+
return asTextContentResult(await searchLocal(body));
141+
}
142+
143+
return asTextContentResult(await searchRemote(body, reqContext.stainlessApiKey));
98144
};
99145

100146
export default { metadata, tool, handler };

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

Lines changed: 551 additions & 0 deletions
Large diffs are not rendered by default.

packages/mcp-server/src/options.ts

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ export type McpOptions = {
1818
includeCodeTool?: boolean | undefined;
1919
includeDocsTools?: boolean | undefined;
2020
stainlessApiKey?: string | undefined;
21+
docsSearchMode?: 'stainless-api' | 'local' | undefined;
22+
docsDir?: string | undefined;
2123
codeAllowHttpGets?: boolean | undefined;
2224
codeAllowedMethods?: string[] | undefined;
2325
codeBlockedMethods?: string[] | undefined;
@@ -58,6 +60,18 @@ export function parseCLIOptions(): CLIOptions {
5860
description: 'Path to custom instructions for the MCP server',
5961
})
6062
.option('debug', { type: 'boolean', description: 'Enable debug logging' })
63+
.option('docs-dir', {
64+
type: 'string',
65+
description:
66+
'Path to a directory of local documentation files (markdown/JSON) to include in local docs search.',
67+
})
68+
.option('docs-search-mode', {
69+
type: 'string',
70+
choices: ['stainless-api', 'local'],
71+
default: 'stainless-api',
72+
description:
73+
"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.",
74+
})
6175
.option('log-format', {
6276
type: 'string',
6377
choices: ['json', 'pretty'],
@@ -118,6 +132,8 @@ export function parseCLIOptions(): CLIOptions {
118132
...(includeDocsTools !== undefined && { includeDocsTools }),
119133
debug: !!argv.debug,
120134
stainlessApiKey: argv.stainlessApiKey,
135+
docsSearchMode: argv.docsSearchMode as 'stainless-api' | 'local' | undefined,
136+
docsDir: argv.docsDir,
121137
codeAllowHttpGets: argv.codeAllowHttpGets,
122138
codeAllowedMethods: argv.codeAllowedMethods,
123139
codeBlockedMethods: argv.codeBlockedMethods,
@@ -163,5 +179,7 @@ export function parseQueryOptions(defaultOptions: McpOptions, query: unknown): M
163179
...(codeTool !== undefined && { includeCodeTool: codeTool }),
164180
...(docsTools !== undefined && { includeDocsTools: docsTools }),
165181
codeExecutionMode: defaultOptions.codeExecutionMode,
182+
docsSearchMode: defaultOptions.docsSearchMode,
183+
docsDir: defaultOptions.docsDir,
166184
};
167185
}

packages/mcp-server/src/server.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import { ClientOptions } from 'scan-documents';
1111
import ScanDocuments from 'scan-documents';
1212
import { codeTool } from './code-tool';
1313
import docsSearchTool from './docs-search-tool';
14+
import { setLocalSearch } from './docs-search-tool';
15+
import { LocalDocsSearch } from './local-docs-search';
1416
import { getInstructions } from './instructions';
1517
import { McpOptions } from './options';
1618
import { blockedMethodsForCodeTool } from './methods';
@@ -62,6 +64,12 @@ export async function initMcpServer(params: {
6264
error: logAtLevel('error'),
6365
};
6466

67+
if (params.mcpOptions?.docsSearchMode === 'local') {
68+
const docsDir = params.mcpOptions?.docsDir;
69+
const localSearch = await LocalDocsSearch.create(docsDir ? { docsDir } : undefined);
70+
setLocalSearch(localSearch);
71+
}
72+
6573
let _client: ScanDocuments | undefined;
6674
let _clientError: Error | undefined;
6775
let _logLevel: 'debug' | 'info' | 'warn' | 'error' | 'off' | undefined;

0 commit comments

Comments
 (0)