/** * BrainBank — MCP Server * * Exposes BrainBank as an MCP server via stdio transport. * Works with Google Antigravity, Claude Code, and any MCP-compatible client. * * Tools: * brainbank_query — AI-pruned semantic context for a task * brainbank_index — Index/re-index a repository */ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; import { z } from 'zod/v3'; import { WorkspacePool } from './workspace-pool.ts'; import { createWorkspaceBrain, resolveRepoPath } from './workspace-factory.ts'; // ── Multi-Workspace BrainBank Pool ───────────────────── const pool = new WorkspacePool({ factory: createWorkspaceBrain, maxMemoryMB: parseInt(process.env.BRAINBANK_MAX_MEMORY_MB ?? '2048', 10), ttlMinutes: parseInt(process.env.BRAINBANK_TTL_MINUTES ?? '30', 10), onError: (repo, err) => { const msg = err instanceof Error ? err.message : String(err); console.error(`BrainBank pool error [${repo}]: ${msg}`); }, }); /** Resolve repo and get a BrainBank from the pool. */ async function getBrainBank(targetRepo?: string) { return pool.get(resolveRepoPath(targetRepo)); } // ── MCP Server Setup ──────────────────────────────── const server = new McpServer({ name: 'brainbank', version: '0.1.5', }); // ── Tool: brainbank_query ─────────────────────────── server.registerTool( 'brainbank_query', { title: 'BrainBank Context', description: 'Get AI-pruned semantic context for a task. Returns search hits + full call tree ' + 'with `called by` annotations, topologically ordered. Pass `context` for general ' + 'background and `pruner` for specific filtering instructions.', inputSchema: z.object({ task: z.string().describe('Description of the task you need context for'), affectedFiles: z.array(z.string()).optional().default([]).describe('Files you plan to modify (improves co-edit suggestions)'), codeResults: z.number().optional().default(20).describe('Max code results'), gitResults: z.number().optional().default(5).describe('Max git commit results'), docsResults: z.number().optional().describe('Max document results (omit to skip docs)'), sources: z.record(z.number()).optional().describe('Per-source result limits, overrides codeResults/gitResults/docsResults (e.g. { code: 10, git: 0, docs: 5 })'), path: z.union([z.string(), z.array(z.string())]).optional().describe('Filter results to files under these path prefixes. Pass a single string or array (e.g. ["src/services/", "lib/"])'), ignore: z.array(z.string()).optional().describe('Exclude results whose filePath starts with any of these prefixes (e.g. ["src/tests/", "src/mocks/"])'), repo: z.string().describe('Repository path (default: BRAINBANK_REPO)'), // BrainBankQL context fields lines: z.boolean().optional().describe('Prefix each code line with its source line number (e.g. 127| code)'), symbols: z.boolean().optional().describe('Append symbol index (all functions, classes, interfaces) for matched files'), compact: z.boolean().optional().describe('Show only function/class signatures, skip bodies'), callTree: z.union([z.boolean(), z.object({ depth: z.number() })]).optional().describe('Include call tree expansion. Pass { depth: N } to control depth'), imports: z.boolean().optional().describe('Include dependency/import summary section'), context: z.string().optional().describe('General task context for the pruner — e.g. GitHub issue body, requirements doc, conversation context. Helps the AI pruner understand the broader task'), pruner: z.string().optional().describe('Specific pruning instructions — tells the pruner exactly what to focus on (e.g. "only keep files related to WebSocket JWT refresh")'), }), }, async ({ task, affectedFiles, codeResults, gitResults, docsResults, sources, path, ignore, repo, lines, symbols, compact, callTree, imports, context: ctx, pruner: prunerCtx }) => { const repoPath = resolveRepoPath(repo); const brainbank = await getBrainBank(repo); // Build sources from explicit params, then let `sources` override const base: Record = { code: codeResults, git: gitResults }; if (docsResults !== undefined) base.docs = docsResults; const resolvedSources = sources ? { ...base, ...sources } : base; // Build fields from explicit params (only include defined values) const fields: Record = {}; if (lines !== undefined) fields.lines = lines; if (symbols !== undefined) fields.symbols = symbols; if (compact !== undefined) fields.compact = compact; if (callTree !== undefined) fields.callTree = callTree; if (imports !== undefined) fields.imports = imports; const context = await brainbank.getContext(task, { affectedFiles, sources: resolvedSources, pathPrefix: path, ignorePaths: ignore, source: 'mcp', fields: Object.keys(fields).length > 0 ? fields : undefined, context: ctx, prunerContext: prunerCtx, }); return { content: [{ type: 'text' as const, text: context }] }; }, ); // ── Tool: brainbank_index ─────────────────────────── server.registerTool( 'brainbank_index', { title: 'BrainBank Index', description: 'Index (or re-index) a repository. Scans code files, git history, and documents, ' + 'then builds vector embeddings and FTS5 indices. Run this after significant code changes ' + 'or when search results feel stale. Returns a summary of what changed.', inputSchema: z.object({ repo: z.string().describe('Repository path to index'), forceReindex: z.boolean().optional().default(false).describe('Force re-index all files (ignore cache). Use when embeddings changed or index is corrupted'), modules: z.array(z.string()).optional().describe('Which plugins to index (e.g. ["code", "git", "docs"]). Omit to use config defaults'), }), }, async ({ repo, forceReindex, modules }) => { const brainbank = await getBrainBank(repo); const result = await brainbank.index({ forceReindex, modules: modules && modules.length > 0 ? modules : undefined, }); // Format results summary const lines: string[] = ['# Index Results\n']; let hasChanges = false; for (const [name, value] of Object.entries(result)) { if (!value) continue; const v = value as Record; if (typeof v.indexed !== 'number') { lines.push(`✓ ${name}: done`); continue; } const indexed = v.indexed as number; const skipped = (v.skipped ?? 0) as number; const removed = (v.removed ?? 0) as number; const chunks = (v.chunks ?? 0) as number; if (indexed > 0 || removed > 0) hasChanges = true; const parts: string[] = []; if (indexed > 0) parts.push(`+${indexed} files (${chunks} chunks)`); if (removed > 0) parts.push(`−${removed} files`); if (skipped > 0) parts.push(`${skipped} unchanged`); lines.push(`${name}: ${parts.join(' ')}`); } if (!hasChanges) { lines.push('No changes — everything up to date'); } // Stats const stats = brainbank.stats(); lines.push('\n## Totals\n'); for (const [name, s] of Object.entries(stats)) { if (!s || typeof s !== 'object') continue; const entries = Object.entries(s as Record) .map(([k, v]) => `${k}: ${v}`) .join(', '); lines.push(`${name}: ${entries}`); } return { content: [{ type: 'text' as const, text: lines.join('\n') }] }; }, ); // ── Start Server ──────────────────────────────────── async function main() { const transport = new StdioServerTransport(); await server.connect(transport); } main().catch(err => { const message = err instanceof Error ? err.message : String(err); console.error(`BrainBank MCP Server Error: ${message}`); process.exit(1); });