server.ts 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971
  1. /**
  2. * QMD MCP Server - Model Context Protocol server for QMD
  3. *
  4. * Exposes QMD search and document retrieval as MCP tools and resources.
  5. * Documents are accessible via qmd:// URIs.
  6. *
  7. * Follows MCP spec 2025-06-18 for proper response types.
  8. */
  9. import { createServer, type IncomingMessage, type ServerResponse } from "node:http";
  10. import { randomUUID } from "node:crypto";
  11. import { readFileSync } from "node:fs";
  12. import { join, dirname } from "node:path";
  13. import { fileURLToPath } from "url";
  14. import { McpServer, ResourceTemplate } from "@modelcontextprotocol/sdk/server/mcp.js";
  15. import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
  16. import { WebStandardStreamableHTTPServerTransport }
  17. from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
  18. import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
  19. import { z } from "zod";
  20. import { existsSync } from "fs";
  21. import {
  22. createStore,
  23. extractSnippet,
  24. addLineNumbers,
  25. getDefaultDbPath,
  26. DEFAULT_MULTI_GET_MAX_BYTES,
  27. createEmbeddingProvider,
  28. type QMDStore,
  29. type ExpandedQuery,
  30. type IndexStatus,
  31. type EmbeddingProvider,
  32. } from "../index.js";
  33. import { getConfigPath } from "../collections.js";
  34. /**
  35. * Resolve the commercial provider only when a semantic operation first uses
  36. * it. MCP startup and deterministic BM25/document tools remain available when
  37. * commercial credentials are absent; semantic operations fail with typed HOLD.
  38. */
  39. class LazyMcpEmbeddingProvider implements EmbeddingProvider {
  40. readonly kind = "openai" as const;
  41. private provider: EmbeddingProvider | undefined;
  42. private resolve(): EmbeddingProvider {
  43. this.provider ??= createEmbeddingProvider({});
  44. return this.provider;
  45. }
  46. getModelId(): string {
  47. return this.resolve().getModelId();
  48. }
  49. getDimensions(): number | undefined {
  50. return this.provider?.getDimensions();
  51. }
  52. healthcheck(signal?: AbortSignal) {
  53. return this.resolve().healthcheck(signal);
  54. }
  55. embed(text: string, options?: Parameters<EmbeddingProvider["embed"]>[1]) {
  56. return this.resolve().embed(text, options);
  57. }
  58. embedBatch(texts: string[], options?: Parameters<EmbeddingProvider["embedBatch"]>[1]) {
  59. return this.resolve().embedBatch(texts, options);
  60. }
  61. getLastError(): string | undefined {
  62. return this.provider?.getLastError?.();
  63. }
  64. async dispose(): Promise<void> {
  65. await this.provider?.dispose();
  66. }
  67. }
  68. function buildMcpEmbedProvider(): EmbeddingProvider {
  69. return new LazyMcpEmbeddingProvider();
  70. }
  71. // =============================================================================
  72. // Types for structured content
  73. // =============================================================================
  74. type SearchResultItem = {
  75. docid: string; // Short docid (#abc123) for quick reference
  76. file: string;
  77. title: string;
  78. score: number;
  79. context: string | null;
  80. snippet: string;
  81. };
  82. type StatusResult = {
  83. totalDocuments: number;
  84. needsEmbedding: number;
  85. hasVectorIndex: boolean;
  86. collections: {
  87. name: string;
  88. path: string | null;
  89. pattern: string | null;
  90. documents: number;
  91. lastUpdated: string;
  92. }[];
  93. };
  94. // =============================================================================
  95. // Helper functions
  96. // =============================================================================
  97. /**
  98. * Encode a path for use in qmd:// URIs.
  99. * Encodes special characters but preserves forward slashes for readability.
  100. */
  101. function encodeQmdPath(path: string): string {
  102. // Encode each path segment separately to preserve slashes
  103. return path.split('/').map(segment => encodeURIComponent(segment)).join('/');
  104. }
  105. /**
  106. * Format search results as human-readable text summary
  107. */
  108. function formatSearchSummary(results: SearchResultItem[], query: string): string {
  109. if (results.length === 0) {
  110. return `No results found for "${query}"`;
  111. }
  112. const lines = [`Found ${results.length} result${results.length === 1 ? '' : 's'} for "${query}":\n`];
  113. for (const r of results) {
  114. lines.push(`${r.docid} ${Math.round(r.score * 100)}% ${r.file} - ${r.title}`);
  115. }
  116. return lines.join('\n');
  117. }
  118. function getPackageVersion(): string {
  119. try {
  120. const pkgPath = join(dirname(fileURLToPath(import.meta.url)), "../../package.json");
  121. const pkg = JSON.parse(readFileSync(pkgPath, "utf-8"));
  122. return pkg.version ?? "unknown";
  123. } catch {
  124. return "unknown";
  125. }
  126. }
  127. // =============================================================================
  128. // MCP Server
  129. // =============================================================================
  130. /**
  131. * Build dynamic server instructions from actual index state.
  132. * Injected into the LLM's system prompt via MCP initialize response —
  133. * gives the LLM immediate context about what's searchable without a tool call.
  134. */
  135. async function buildInstructions(store: QMDStore): Promise<string> {
  136. const status = await store.getStatus();
  137. const contexts = await store.listContexts();
  138. const globalCtx = await store.getGlobalContext();
  139. const lines: string[] = [];
  140. // --- What is this? ---
  141. lines.push(`QMD is your local search engine over ${status.totalDocuments} markdown documents.`);
  142. if (globalCtx) lines.push(`Context: ${globalCtx}`);
  143. // --- What's searchable? ---
  144. if (status.collections.length > 0) {
  145. lines.push("");
  146. lines.push("Collections (scope with `collection` parameter):");
  147. for (const col of status.collections) {
  148. // Find root context for this collection
  149. const rootCtx = contexts.find(c => c.collection === col.name && (c.path === "" || c.path === "/"));
  150. const desc = rootCtx ? ` — ${rootCtx.context}` : "";
  151. lines.push(` - "${col.name}" (${col.documents} docs)${desc}`);
  152. }
  153. }
  154. // --- Capability gaps ---
  155. if (!status.hasVectorIndex) {
  156. lines.push("");
  157. lines.push("Note: No vector embeddings yet. Run `qmd embed` to enable semantic search (vec/hyde).");
  158. } else if (status.needsEmbedding > 0) {
  159. lines.push("");
  160. lines.push(`Note: ${status.needsEmbedding} documents need embedding. Run \`qmd embed\` to update.`);
  161. }
  162. // --- Search tool ---
  163. lines.push("");
  164. lines.push("Search: Use `query` with sub-queries (lex/vec/hyde):");
  165. lines.push(" - type:'lex' — BM25 keyword search (exact terms, fast)");
  166. lines.push(" - type:'vec' — semantic vector search (meaning-based)");
  167. lines.push(" - type:'hyde' — hypothetical document (write what the answer looks like)");
  168. lines.push("");
  169. lines.push(" Always provide `intent` on every search call to disambiguate and improve snippets.");
  170. lines.push("");
  171. lines.push("Examples:");
  172. lines.push(" Quick keyword lookup: [{type:'lex', query:'error handling'}]");
  173. lines.push(" Semantic search: [{type:'vec', query:'how to handle errors gracefully'}]");
  174. lines.push(" Best results: [{type:'lex', query:'error'}, {type:'vec', query:'error handling best practices'}]");
  175. lines.push(" With intent: searches=[{type:'lex', query:'performance'}], intent='web page load times'");
  176. // --- Retrieval workflow ---
  177. lines.push("");
  178. lines.push("Retrieval:");
  179. lines.push(" - `get` — single document by path or docid (#abc123). Supports line offset (`file.md:100`).");
  180. lines.push(" - `multi_get` — batch retrieve by glob (`journals/2025-05*.md`) or comma-separated list.");
  181. // --- Non-obvious things that prevent mistakes ---
  182. lines.push("");
  183. lines.push("Tips:");
  184. lines.push(" - File paths in results are relative to their collection.");
  185. lines.push(" - Use `minScore: 0.5` to filter low-confidence results.");
  186. lines.push(" - Results include a `context` field describing the content type.");
  187. return lines.join("\n");
  188. }
  189. /**
  190. * Create an MCP server with all QMD tools, resources, and prompts registered.
  191. * Shared by both stdio and HTTP transports.
  192. */
  193. async function createMcpServer(store: QMDStore): Promise<McpServer> {
  194. const server = new McpServer(
  195. { name: "qmd", version: getPackageVersion() },
  196. { instructions: await buildInstructions(store) },
  197. );
  198. // Pre-fetch default collection names for search tools
  199. const defaultCollectionNames = await store.getDefaultCollectionNames();
  200. // ---------------------------------------------------------------------------
  201. // Resource: qmd://{path} - read-only access to documents by path
  202. // Note: No list() - documents are discovered via search tools
  203. // ---------------------------------------------------------------------------
  204. server.registerResource(
  205. "document",
  206. new ResourceTemplate("qmd://{+path}", { list: undefined }),
  207. {
  208. title: "QMD Document",
  209. description: "A markdown document from your QMD knowledge base. Use search tools to discover documents.",
  210. mimeType: "text/markdown",
  211. },
  212. async (uri, { path }) => {
  213. // Decode URL-encoded path (MCP clients send encoded URIs)
  214. const pathStr = Array.isArray(path) ? path.join('/') : (path || '');
  215. const decodedPath = decodeURIComponent(pathStr);
  216. // Use SDK to find document — findDocument handles collection/path resolution
  217. const result = await store.get(decodedPath, { includeBody: true });
  218. if ("error" in result) {
  219. return { contents: [{ uri: uri.href, text: `Document not found: ${decodedPath}` }] };
  220. }
  221. let text = addLineNumbers(result.body || ""); // Default to line numbers
  222. if (result.context) {
  223. text = `<!-- Context: ${result.context} -->\n\n` + text;
  224. }
  225. return {
  226. contents: [{
  227. uri: uri.href,
  228. name: result.displayPath,
  229. title: result.title || result.displayPath,
  230. mimeType: "text/markdown",
  231. text,
  232. }],
  233. };
  234. }
  235. );
  236. // ---------------------------------------------------------------------------
  237. // Tool: query (Primary search tool)
  238. // ---------------------------------------------------------------------------
  239. const subSearchSchema = z.object({
  240. type: z.enum(['lex', 'vec', 'hyde']).describe(
  241. "lex = BM25 keywords (supports \"phrase\" and -negation); " +
  242. "vec = semantic question; hyde = hypothetical answer passage"
  243. ),
  244. query: z.string().describe(
  245. "The query text. For lex: use keywords, \"quoted phrases\", and -negation. " +
  246. "For vec: natural language question. For hyde: 50-100 word answer passage."
  247. ),
  248. });
  249. server.registerTool(
  250. "query",
  251. {
  252. title: "Query",
  253. description: `Search the knowledge base using a query document — one or more typed sub-queries combined for best recall.
  254. ## Query Types
  255. **lex** — BM25 keyword search. Fast, exact, no LLM needed.
  256. Full lex syntax:
  257. - \`term\` — prefix match ("perf" matches "performance")
  258. - \`"exact phrase"\` — phrase must appear verbatim
  259. - \`-term\` or \`-"phrase"\` — exclude documents containing this
  260. Good lex examples:
  261. - \`"connection pool" timeout -redis\`
  262. - \`"machine learning" -sports -athlete\`
  263. - \`handleError async typescript\`
  264. **vec** — Semantic vector search. Write a natural language question. Finds documents by meaning, not exact words.
  265. - \`how does the rate limiter handle burst traffic?\`
  266. - \`what is the tradeoff between consistency and availability?\`
  267. **hyde** — Hypothetical document. Write 50-100 words that look like the answer. Often the most powerful for nuanced topics.
  268. - \`The rate limiter uses a token bucket algorithm. When a client exceeds 100 req/min, subsequent requests return 429 until the window resets.\`
  269. ## Strategy
  270. Combine types for best results. First sub-query gets 2× weight — put your strongest signal first.
  271. | Goal | Approach |
  272. |------|----------|
  273. | Know exact term/name | \`lex\` only |
  274. | Concept search | \`vec\` only |
  275. | Best recall | \`lex\` + \`vec\` |
  276. | Complex/nuanced | \`lex\` + \`vec\` + \`hyde\` |
  277. | Unknown vocabulary | Use a standalone natural-language query (no typed lines) so the server can auto-expand it |
  278. ## Examples
  279. Simple lookup:
  280. \`\`\`json
  281. [{ "type": "lex", "query": "CAP theorem" }]
  282. \`\`\`
  283. Best recall on a technical topic:
  284. \`\`\`json
  285. [
  286. { "type": "lex", "query": "\\"connection pool\\" timeout -redis" },
  287. { "type": "vec", "query": "why do database connections time out under load" },
  288. { "type": "hyde", "query": "Connection pool exhaustion occurs when all connections are in use and new requests must wait. This typically happens under high concurrency when queries run longer than expected." }
  289. ]
  290. \`\`\`
  291. Intent-aware lex (C++ performance, not sports):
  292. \`\`\`json
  293. [
  294. { "type": "lex", "query": "\\"C++ performance\\" optimization -sports -athlete" },
  295. { "type": "vec", "query": "how to optimize C++ program performance" }
  296. ]
  297. \`\`\``,
  298. annotations: { readOnlyHint: true, openWorldHint: false },
  299. inputSchema: {
  300. searches: z.array(subSearchSchema).min(1).max(10).describe(
  301. "Typed sub-queries to execute (lex/vec/hyde). First gets 2x weight."
  302. ),
  303. limit: z.number().optional().default(10).describe("Max results (default: 10)"),
  304. minScore: z.number().optional().default(0).describe("Min relevance 0-1 (default: 0)"),
  305. candidateLimit: z.number().optional().describe(
  306. "Maximum candidates to rerank (default: 40, lower = faster but may miss results)"
  307. ),
  308. collections: z.array(z.string()).optional().describe("Filter to collections (OR match)"),
  309. intent: z.string().optional().describe(
  310. "Background context to disambiguate the query. Example: query='performance', intent='web page load times and Core Web Vitals'. Does not search on its own."
  311. ),
  312. rerank: z.boolean().optional().default(true).describe(
  313. "Rerank results using LLM (default: true). Set to false for faster results on CPU-only machines."
  314. ),
  315. },
  316. },
  317. async ({ searches, limit, minScore, candidateLimit, collections, intent, rerank }) => {
  318. // Map to internal format
  319. const queries: ExpandedQuery[] = searches.map(s => ({
  320. type: s.type,
  321. query: s.query,
  322. }));
  323. // Use default collections if none specified
  324. const effectiveCollections = collections ?? defaultCollectionNames;
  325. const results = await store.search({
  326. queries,
  327. collections: effectiveCollections.length > 0 ? effectiveCollections : undefined,
  328. limit,
  329. minScore,
  330. rerank,
  331. intent,
  332. });
  333. // Use first lex or vec query for snippet extraction
  334. const primaryQuery = searches.find(s => s.type === 'lex')?.query
  335. || searches.find(s => s.type === 'vec')?.query
  336. || searches[0]?.query || "";
  337. const filtered: SearchResultItem[] = results.map(r => {
  338. const { line, snippet } = extractSnippet(r.bestChunk, primaryQuery, 300, undefined, undefined, intent);
  339. return {
  340. docid: `#${r.docid}`,
  341. file: r.displayPath,
  342. title: r.title,
  343. score: Math.round(r.score * 100) / 100,
  344. context: r.context,
  345. snippet: addLineNumbers(snippet, line),
  346. };
  347. });
  348. return {
  349. content: [{ type: "text", text: formatSearchSummary(filtered, primaryQuery) }],
  350. structuredContent: { results: filtered },
  351. };
  352. }
  353. );
  354. // ---------------------------------------------------------------------------
  355. // Tool: qmd_get (Retrieve document)
  356. // ---------------------------------------------------------------------------
  357. server.registerTool(
  358. "get",
  359. {
  360. title: "Get Document",
  361. description: "Retrieve the full content of a document by its file path or docid. Use paths or docids (#abc123) from search results. Suggests similar files if not found.",
  362. annotations: { readOnlyHint: true, openWorldHint: false },
  363. inputSchema: {
  364. file: z.string().describe("File path or docid from search results (e.g., 'pages/meeting.md', '#abc123', or 'pages/meeting.md:100' to start at line 100)"),
  365. fromLine: z.number().optional().describe("Start from this line number (1-indexed)"),
  366. maxLines: z.number().optional().describe("Maximum number of lines to return"),
  367. lineNumbers: z.boolean().optional().default(false).describe("Add line numbers to output (format: 'N: content')"),
  368. },
  369. },
  370. async ({ file, fromLine, maxLines, lineNumbers }) => {
  371. // Support :line suffix in `file` (e.g. "foo.md:120") when fromLine isn't provided
  372. let parsedFromLine = fromLine;
  373. let lookup = file;
  374. const colonMatch = lookup.match(/:(\d+)$/);
  375. if (colonMatch && colonMatch[1] && parsedFromLine === undefined) {
  376. parsedFromLine = parseInt(colonMatch[1], 10);
  377. lookup = lookup.slice(0, -colonMatch[0].length);
  378. }
  379. const result = await store.get(lookup, { includeBody: false });
  380. if ("error" in result) {
  381. let msg = `Document not found: ${file}`;
  382. if (result.similarFiles.length > 0) {
  383. msg += `\n\nDid you mean one of these?\n${result.similarFiles.map(s => ` - ${s}`).join('\n')}`;
  384. }
  385. return {
  386. content: [{ type: "text", text: msg }],
  387. isError: true,
  388. };
  389. }
  390. const body = await store.getDocumentBody(result.filepath, { fromLine: parsedFromLine, maxLines }) ?? "";
  391. let text = body;
  392. if (lineNumbers) {
  393. const startLine = parsedFromLine || 1;
  394. text = addLineNumbers(text, startLine);
  395. }
  396. if (result.context) {
  397. text = `<!-- Context: ${result.context} -->\n\n` + text;
  398. }
  399. return {
  400. content: [{
  401. type: "resource",
  402. resource: {
  403. uri: `qmd://${encodeQmdPath(result.displayPath)}`,
  404. name: result.displayPath,
  405. title: result.title,
  406. mimeType: "text/markdown",
  407. text,
  408. },
  409. }],
  410. };
  411. }
  412. );
  413. // ---------------------------------------------------------------------------
  414. // Tool: qmd_multi_get (Retrieve multiple documents)
  415. // ---------------------------------------------------------------------------
  416. server.registerTool(
  417. "multi_get",
  418. {
  419. title: "Multi-Get Documents",
  420. description: "Retrieve multiple documents by glob pattern (e.g., 'journals/2025-05*.md') or comma-separated list. Skips files larger than maxBytes.",
  421. annotations: { readOnlyHint: true, openWorldHint: false },
  422. inputSchema: {
  423. pattern: z.string().describe("Glob pattern or comma-separated list of file paths"),
  424. maxLines: z.number().optional().describe("Maximum lines per file"),
  425. maxBytes: z.number().optional().default(10240).describe("Skip files larger than this (default: 10240 = 10KB)"),
  426. lineNumbers: z.boolean().optional().default(false).describe("Add line numbers to output (format: 'N: content')"),
  427. },
  428. },
  429. async ({ pattern, maxLines, maxBytes, lineNumbers }) => {
  430. const { docs, errors } = await store.multiGet(pattern, { includeBody: true, maxBytes: maxBytes || DEFAULT_MULTI_GET_MAX_BYTES });
  431. if (docs.length === 0 && errors.length === 0) {
  432. return {
  433. content: [{ type: "text", text: `No files matched pattern: ${pattern}` }],
  434. isError: true,
  435. };
  436. }
  437. const content: ({ type: "text"; text: string } | { type: "resource"; resource: { uri: string; name: string; title?: string; mimeType: string; text: string } })[] = [];
  438. if (errors.length > 0) {
  439. content.push({ type: "text", text: `Errors:\n${errors.join('\n')}` });
  440. }
  441. for (const result of docs) {
  442. if (result.skipped) {
  443. content.push({
  444. type: "text",
  445. text: `[SKIPPED: ${result.doc.displayPath} - ${result.skipReason}. Use 'qmd_get' with file="${result.doc.displayPath}" to retrieve.]`,
  446. });
  447. continue;
  448. }
  449. let text = result.doc.body || "";
  450. if (maxLines !== undefined) {
  451. const lines = text.split("\n");
  452. text = lines.slice(0, maxLines).join("\n");
  453. if (lines.length > maxLines) {
  454. text += `\n\n[... truncated ${lines.length - maxLines} more lines]`;
  455. }
  456. }
  457. if (lineNumbers) {
  458. text = addLineNumbers(text);
  459. }
  460. if (result.doc.context) {
  461. text = `<!-- Context: ${result.doc.context} -->\n\n` + text;
  462. }
  463. content.push({
  464. type: "resource",
  465. resource: {
  466. uri: `qmd://${encodeQmdPath(result.doc.displayPath)}`,
  467. name: result.doc.displayPath,
  468. title: result.doc.title,
  469. mimeType: "text/markdown",
  470. text,
  471. },
  472. });
  473. }
  474. return { content };
  475. }
  476. );
  477. // ---------------------------------------------------------------------------
  478. // Tool: qmd_status (Index status)
  479. // ---------------------------------------------------------------------------
  480. server.registerTool(
  481. "status",
  482. {
  483. title: "Index Status",
  484. description: "Show the status of the QMD index: collections, document counts, and health information.",
  485. annotations: { readOnlyHint: true, openWorldHint: false },
  486. inputSchema: {},
  487. },
  488. async () => {
  489. const status: StatusResult = await store.getStatus();
  490. const summary = [
  491. `QMD Index Status:`,
  492. ` Total documents: ${status.totalDocuments}`,
  493. ` Needs embedding: ${status.needsEmbedding}`,
  494. ` Vector index: ${status.hasVectorIndex ? 'yes' : 'no'}`,
  495. ` Collections: ${status.collections.length}`,
  496. ];
  497. for (const col of status.collections) {
  498. summary.push(` - ${col.name}: ${col.path} (${col.documents} docs)`);
  499. }
  500. return {
  501. content: [{ type: "text", text: summary.join('\n') }],
  502. structuredContent: status,
  503. };
  504. }
  505. );
  506. return server;
  507. }
  508. // =============================================================================
  509. // Process supervision — RSS ceiling self-restart (i-6sw24v09)
  510. // =============================================================================
  511. /**
  512. * Periodically check this MCP process's RSS and exit cleanly when it
  513. * crosses a configurable ceiling, so the parent (Claude Code stdio
  514. * client, pm2, systemd, etc.) can respawn a fresh process. Contains
  515. * the blast radius of memory leaks in the search/expansion path
  516. * without requiring a full re-architecture.
  517. *
  518. * Configuration:
  519. * QMD_MCP_RSS_LIMIT_BYTES — ceiling in bytes. Default 0 (disabled).
  520. * Set to e.g. `2147483648` (2 GiB) to opt in.
  521. * QMD_MCP_RSS_CHECK_INTERVAL_MS — poll interval. Default 60000 (60s).
  522. *
  523. * Default-off so we can ship the diagnostic + pragma fix safely and
  524. * graduate the supervisor on once we have soak data showing no
  525. * false positives. Operators can opt in immediately by exporting
  526. * `QMD_MCP_RSS_LIMIT_BYTES=2147483648` in their MCP launcher env.
  527. */
  528. export type RssSupervisorHandle = {
  529. stop: () => void;
  530. /** Snapshot the most recent RSS reading (test hook). */
  531. lastRss: () => number;
  532. };
  533. export interface RssSupervisorOptions {
  534. /** RSS ceiling in bytes. 0 disables. */
  535. limitBytes?: number;
  536. /** Polling cadence in ms. Default 60000. */
  537. intervalMs?: number;
  538. /** Override RSS reader for tests. */
  539. readRss?: () => number;
  540. /** Override exit hook for tests. */
  541. onExceeded?: (rss: number, limit: number) => void;
  542. /** Override stderr writer for tests. */
  543. log?: (line: string) => void;
  544. }
  545. export function startRssSupervisor(opts: RssSupervisorOptions = {}): RssSupervisorHandle | null {
  546. const env = process.env;
  547. const limit = opts.limitBytes ?? parseInt(env.QMD_MCP_RSS_LIMIT_BYTES ?? "0", 10);
  548. if (!Number.isFinite(limit) || limit <= 0) return null; // disabled
  549. const interval = opts.intervalMs ?? parseInt(env.QMD_MCP_RSS_CHECK_INTERVAL_MS ?? "60000", 10);
  550. const safeInterval = Number.isFinite(interval) && interval > 0 ? interval : 60000;
  551. const readRss = opts.readRss ?? (() => process.memoryUsage().rss);
  552. const log = opts.log ?? ((line: string) => process.stderr.write(line));
  553. const onExceeded = opts.onExceeded ?? ((rss: number, lim: number) => {
  554. log(`[qmd mcp] RSS_LIMIT_EXCEEDED rss=${rss} limit=${lim} pid=${process.pid} — exiting for parent respawn\n`);
  555. process.exit(1);
  556. });
  557. let lastRss = 0;
  558. const timer = setInterval(() => {
  559. try {
  560. lastRss = readRss();
  561. if (lastRss > limit) {
  562. clearInterval(timer);
  563. onExceeded(lastRss, limit);
  564. }
  565. } catch (err) {
  566. // Defensive — never let the supervisor crash the server.
  567. const msg = err instanceof Error ? err.message : String(err);
  568. log(`[qmd mcp] WARN rss supervisor check failed: ${msg}\n`);
  569. }
  570. }, safeInterval);
  571. // Don't keep the event loop alive just for the supervisor.
  572. if (typeof timer.unref === "function") timer.unref();
  573. return {
  574. stop: () => clearInterval(timer),
  575. lastRss: () => lastRss,
  576. };
  577. }
  578. // =============================================================================
  579. // Transport: stdio (default)
  580. // =============================================================================
  581. export async function startMcpServer(): Promise<void> {
  582. const configPath = getConfigPath();
  583. const embedProvider = buildMcpEmbedProvider();
  584. const store = await createStore({
  585. dbPath: getDefaultDbPath(),
  586. ...(existsSync(configPath) ? { configPath } : {}),
  587. ...(embedProvider ? { embedProvider } : {}),
  588. });
  589. startRssSupervisor();
  590. const server = await createMcpServer(store);
  591. const transport = new StdioServerTransport();
  592. await server.connect(transport);
  593. }
  594. // =============================================================================
  595. // Transport: Streamable HTTP
  596. // =============================================================================
  597. export type HttpServerHandle = {
  598. httpServer: import("http").Server;
  599. port: number;
  600. stop: () => Promise<void>;
  601. };
  602. /**
  603. * Start MCP server over Streamable HTTP (JSON responses, no SSE).
  604. * Binds to localhost only. Returns a handle for shutdown and port discovery.
  605. */
  606. export async function startMcpHttpServer(port: number, options?: { quiet?: boolean }): Promise<HttpServerHandle> {
  607. const configPath = getConfigPath();
  608. const embedProvider = buildMcpEmbedProvider();
  609. const store = await createStore({
  610. dbPath: getDefaultDbPath(),
  611. ...(existsSync(configPath) ? { configPath } : {}),
  612. ...(embedProvider ? { embedProvider } : {}),
  613. });
  614. const rssSupervisor = startRssSupervisor();
  615. // Pre-fetch default collection names for REST endpoint
  616. const defaultCollectionNames = await store.getDefaultCollectionNames();
  617. // Session map: each client gets its own McpServer + Transport pair (MCP spec requirement).
  618. // The store is shared — it's stateless SQLite, safe for concurrent access.
  619. const sessions = new Map<string, WebStandardStreamableHTTPServerTransport>();
  620. async function createSession(): Promise<WebStandardStreamableHTTPServerTransport> {
  621. const transport = new WebStandardStreamableHTTPServerTransport({
  622. sessionIdGenerator: () => randomUUID(),
  623. enableJsonResponse: true,
  624. onsessioninitialized: (sessionId: string) => {
  625. sessions.set(sessionId, transport);
  626. log(`${ts()} New session ${sessionId} (${sessions.size} active)`);
  627. },
  628. });
  629. const server = await createMcpServer(store);
  630. await server.connect(transport);
  631. transport.onclose = () => {
  632. if (transport.sessionId) {
  633. sessions.delete(transport.sessionId);
  634. }
  635. };
  636. return transport;
  637. }
  638. const startTime = Date.now();
  639. const quiet = options?.quiet ?? false;
  640. /** Format timestamp for request logging */
  641. function ts(): string {
  642. return new Date().toISOString().slice(11, 23); // HH:mm:ss.SSS
  643. }
  644. /** Extract a human-readable label from a JSON-RPC body */
  645. function describeRequest(body: any): string {
  646. const method = body?.method ?? "unknown";
  647. if (method === "tools/call") {
  648. const tool = body.params?.name ?? "?";
  649. const args = body.params?.arguments;
  650. // Show query string if present, truncated
  651. if (args?.query) {
  652. const q = String(args.query).slice(0, 80);
  653. return `tools/call ${tool} "${q}"`;
  654. }
  655. if (args?.path) return `tools/call ${tool} ${args.path}`;
  656. if (args?.pattern) return `tools/call ${tool} ${args.pattern}`;
  657. return `tools/call ${tool}`;
  658. }
  659. return method;
  660. }
  661. function log(msg: string): void {
  662. if (!quiet) console.error(msg);
  663. }
  664. // Helper to collect request body
  665. async function collectBody(req: IncomingMessage): Promise<string> {
  666. const chunks: Buffer[] = [];
  667. for await (const chunk of req) chunks.push(chunk as Buffer);
  668. return Buffer.concat(chunks).toString();
  669. }
  670. const httpServer = createServer(async (nodeReq: IncomingMessage, nodeRes: ServerResponse) => {
  671. const reqStart = Date.now();
  672. const pathname = nodeReq.url || "/";
  673. try {
  674. if (pathname === "/health" && nodeReq.method === "GET") {
  675. const body = JSON.stringify({ status: "ok", uptime: Math.floor((Date.now() - startTime) / 1000) });
  676. nodeRes.writeHead(200, { "Content-Type": "application/json" });
  677. nodeRes.end(body);
  678. log(`${ts()} GET /health (${Date.now() - reqStart}ms)`);
  679. return;
  680. }
  681. // REST endpoint: POST /search — structured search without MCP protocol
  682. // REST endpoint: POST /query (alias: /search) — structured search without MCP protocol
  683. if ((pathname === "/query" || pathname === "/search") && nodeReq.method === "POST") {
  684. const rawBody = await collectBody(nodeReq);
  685. const params = JSON.parse(rawBody);
  686. // Validate required fields
  687. if (!params.searches || !Array.isArray(params.searches)) {
  688. nodeRes.writeHead(400, { "Content-Type": "application/json" });
  689. nodeRes.end(JSON.stringify({ error: "Missing required field: searches (array)" }));
  690. return;
  691. }
  692. // Map to internal format
  693. const queries: ExpandedQuery[] = params.searches.map((s: any) => ({
  694. type: s.type as 'lex' | 'vec' | 'hyde',
  695. query: String(s.query || ""),
  696. }));
  697. // Use default collections if none specified
  698. const effectiveCollections = params.collections ?? defaultCollectionNames;
  699. const results = await store.search({
  700. queries,
  701. collections: effectiveCollections.length > 0 ? effectiveCollections : undefined,
  702. limit: params.limit ?? 10,
  703. minScore: params.minScore ?? 0,
  704. intent: params.intent,
  705. });
  706. // Use first lex or vec query for snippet extraction
  707. const primaryQuery = params.searches.find((s: any) => s.type === 'lex')?.query
  708. || params.searches.find((s: any) => s.type === 'vec')?.query
  709. || params.searches[0]?.query || "";
  710. const formatted = results.map(r => {
  711. const { line, snippet } = extractSnippet(r.bestChunk, primaryQuery, 300);
  712. return {
  713. docid: `#${r.docid}`,
  714. file: r.displayPath,
  715. title: r.title,
  716. score: Math.round(r.score * 100) / 100,
  717. context: r.context,
  718. snippet: addLineNumbers(snippet, line),
  719. };
  720. });
  721. nodeRes.writeHead(200, { "Content-Type": "application/json" });
  722. nodeRes.end(JSON.stringify({ results: formatted }));
  723. log(`${ts()} POST /query ${params.searches.length} queries (${Date.now() - reqStart}ms)`);
  724. return;
  725. }
  726. if (pathname === "/mcp" && nodeReq.method === "POST") {
  727. const rawBody = await collectBody(nodeReq);
  728. const body = JSON.parse(rawBody);
  729. const label = describeRequest(body);
  730. const url = `http://localhost:${port}${pathname}`;
  731. const headers: Record<string, string> = {};
  732. for (const [k, v] of Object.entries(nodeReq.headers)) {
  733. if (typeof v === "string") headers[k] = v;
  734. }
  735. // Route to existing session or create new one on initialize
  736. const sessionId = headers["mcp-session-id"];
  737. let transport: WebStandardStreamableHTTPServerTransport;
  738. if (sessionId) {
  739. const existing = sessions.get(sessionId);
  740. if (!existing) {
  741. nodeRes.writeHead(404, { "Content-Type": "application/json" });
  742. nodeRes.end(JSON.stringify({
  743. jsonrpc: "2.0",
  744. error: { code: -32001, message: "Session not found" },
  745. id: body?.id ?? null,
  746. }));
  747. return;
  748. }
  749. transport = existing;
  750. } else if (isInitializeRequest(body)) {
  751. transport = await createSession();
  752. } else {
  753. nodeRes.writeHead(400, { "Content-Type": "application/json" });
  754. nodeRes.end(JSON.stringify({
  755. jsonrpc: "2.0",
  756. error: { code: -32000, message: "Bad Request: Missing session ID" },
  757. id: body?.id ?? null,
  758. }));
  759. return;
  760. }
  761. const request = new Request(url, { method: "POST", headers, body: rawBody });
  762. const response = await transport.handleRequest(request, { parsedBody: body });
  763. nodeRes.writeHead(response.status, Object.fromEntries(response.headers));
  764. nodeRes.end(Buffer.from(await response.arrayBuffer()));
  765. log(`${ts()} POST /mcp ${label} (${Date.now() - reqStart}ms)`);
  766. return;
  767. }
  768. if (pathname === "/mcp") {
  769. const headers: Record<string, string> = {};
  770. for (const [k, v] of Object.entries(nodeReq.headers)) {
  771. if (typeof v === "string") headers[k] = v;
  772. }
  773. // GET/DELETE must have a valid session
  774. const sessionId = headers["mcp-session-id"];
  775. if (!sessionId) {
  776. nodeRes.writeHead(400, { "Content-Type": "application/json" });
  777. nodeRes.end(JSON.stringify({
  778. jsonrpc: "2.0",
  779. error: { code: -32000, message: "Bad Request: Missing session ID" },
  780. id: null,
  781. }));
  782. return;
  783. }
  784. const transport = sessions.get(sessionId);
  785. if (!transport) {
  786. nodeRes.writeHead(404, { "Content-Type": "application/json" });
  787. nodeRes.end(JSON.stringify({
  788. jsonrpc: "2.0",
  789. error: { code: -32001, message: "Session not found" },
  790. id: null,
  791. }));
  792. return;
  793. }
  794. const url = `http://localhost:${port}${pathname}`;
  795. const rawBody = nodeReq.method !== "GET" && nodeReq.method !== "HEAD" ? await collectBody(nodeReq) : undefined;
  796. const request = new Request(url, { method: nodeReq.method || "GET", headers, ...(rawBody ? { body: rawBody } : {}) });
  797. const response = await transport.handleRequest(request);
  798. nodeRes.writeHead(response.status, Object.fromEntries(response.headers));
  799. nodeRes.end(Buffer.from(await response.arrayBuffer()));
  800. return;
  801. }
  802. nodeRes.writeHead(404);
  803. nodeRes.end("Not Found");
  804. } catch (err) {
  805. console.error("HTTP handler error:", err);
  806. nodeRes.writeHead(500);
  807. nodeRes.end("Internal Server Error");
  808. }
  809. });
  810. await new Promise<void>((resolve, reject) => {
  811. httpServer.on("error", reject);
  812. httpServer.listen(port, "localhost", () => resolve());
  813. });
  814. const actualPort = (httpServer.address() as import("net").AddressInfo).port;
  815. let stopping = false;
  816. const stop = async () => {
  817. if (stopping) return;
  818. stopping = true;
  819. for (const transport of sessions.values()) {
  820. await transport.close();
  821. }
  822. sessions.clear();
  823. if (rssSupervisor) rssSupervisor.stop();
  824. httpServer.close();
  825. await store.close();
  826. // Dispose the query-side embedding provider (if any) — releases
  827. // HTTP keep-alive sockets in OpenAIEmbeddingsProvider (i-loazq6ze).
  828. if (embedProvider) {
  829. try { await embedProvider.dispose(); } catch { /* ignore */ }
  830. }
  831. };
  832. process.on("SIGTERM", async () => {
  833. console.error("Shutting down (SIGTERM)...");
  834. await stop();
  835. process.exit(0);
  836. });
  837. process.on("SIGINT", async () => {
  838. console.error("Shutting down (SIGINT)...");
  839. await stop();
  840. process.exit(0);
  841. });
  842. log(`QMD MCP server listening on http://localhost:${actualPort}/mcp`);
  843. return { httpServer, port: actualPort, stop };
  844. }
  845. // Run if this is the main module
  846. if (fileURLToPath(import.meta.url) === process.argv[1] || process.argv[1]?.endsWith("/server.ts") || process.argv[1]?.endsWith("/server.js")) {
  847. startMcpServer().catch(console.error);
  848. }