|
3 | 3 | import { Tool } from '@modelcontextprotocol/sdk/types.js'; |
4 | 4 | import { Metadata, McpRequestContext, asTextContentResult } from './types'; |
5 | 5 | import { getLogger } from './logger'; |
| 6 | +import type { LocalDocsSearch } from './local-docs-search'; |
6 | 7 |
|
7 | 8 | export const metadata: Metadata = { |
8 | 9 | resource: 'all', |
@@ -43,20 +44,41 @@ export const tool: Tool = { |
43 | 44 | const docsSearchURL = |
44 | 45 | process.env['DOCS_SEARCH_URL'] || 'https://api.stainless.com/api/projects/cas-parser/docs/search'; |
45 | 46 |
|
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 | +async function searchLocal(args: Record<string, unknown>): Promise<unknown> { |
| 54 | + if (!_localSearch) { |
| 55 | + throw new Error('Local search not initialized'); |
| 56 | + } |
| 57 | + |
| 58 | + const query = (args['query'] as string) ?? ''; |
| 59 | + const language = (args['language'] as string) ?? 'typescript'; |
| 60 | + const detail = (args['detail'] as string) ?? 'default'; |
| 61 | + |
| 62 | + return _localSearch.search({ |
| 63 | + query, |
| 64 | + language, |
| 65 | + detail, |
| 66 | + maxResults: 5, |
| 67 | + }).results; |
| 68 | +} |
| 69 | + |
| 70 | +async function searchRemote(args: Record<string, unknown>, reqContext: McpRequestContext): Promise<unknown> { |
53 | 71 | const body = args as any; |
54 | 72 | const query = new URLSearchParams(body).toString(); |
55 | 73 |
|
56 | 74 | const startTime = Date.now(); |
57 | 75 | const result = await fetch(`${docsSearchURL}?${query}`, { |
58 | 76 | headers: { |
59 | 77 | ...(reqContext.stainlessApiKey && { Authorization: reqContext.stainlessApiKey }), |
| 78 | + ...(reqContext.mcpSessionId && { 'x-stainless-mcp-session-id': reqContext.mcpSessionId }), |
| 79 | + ...(reqContext.mcpClientInfo && { |
| 80 | + 'x-stainless-mcp-client-info': JSON.stringify(reqContext.mcpClientInfo), |
| 81 | + }), |
60 | 82 | }, |
61 | 83 | }); |
62 | 84 |
|
@@ -94,7 +116,23 @@ export const handler = async ({ |
94 | 116 | }, |
95 | 117 | 'Got docs search result', |
96 | 118 | ); |
97 | | - return asTextContentResult(resultBody); |
| 119 | + return resultBody; |
| 120 | +} |
| 121 | + |
| 122 | +export const handler = async ({ |
| 123 | + reqContext, |
| 124 | + args, |
| 125 | +}: { |
| 126 | + reqContext: McpRequestContext; |
| 127 | + args: Record<string, unknown> | undefined; |
| 128 | +}) => { |
| 129 | + const body = args ?? {}; |
| 130 | + |
| 131 | + if (_localSearch) { |
| 132 | + return asTextContentResult(await searchLocal(body)); |
| 133 | + } |
| 134 | + |
| 135 | + return asTextContentResult(await searchRemote(body, reqContext)); |
98 | 136 | }; |
99 | 137 |
|
100 | 138 | export default { metadata, tool, handler }; |
0 commit comments