@akiojin/unity-mcp-server
Version:
MCP server and Unity Editor bridge — enables AI assistants to control Unity for AI-assisted workflows
164 lines (152 loc) • 6.52 kB
JavaScript
import { BaseToolHandler } from '../base/BaseToolHandler.js';
import { CodeIndex } from '../../core/codeIndex.js';
import fs from 'fs';
import path from 'path';
import { ProjectInfoProvider } from '../../core/projectInfo.js';
import { LspRpcClient } from '../../lsp/LspRpcClient.js';
import { logger } from '../../core/config.js';
export class BuildCodeIndexToolHandler extends BaseToolHandler {
constructor(unityConnection) {
super(
'build_code_index',
'Build a persistent SQLite symbol index by scanning document symbols via the C# LSP. Stores DB under Library/UnityMCP/CodeIndex/code-index.db.',
{
type: 'object',
properties: {},
required: []
}
);
this.unityConnection = unityConnection;
this.index = new CodeIndex(unityConnection);
this.projectInfo = new ProjectInfoProvider(unityConnection);
this.lsp = null; // lazy init with projectRoot
}
async execute(params = {}) {
try {
const info = await this.projectInfo.get();
const roots = [
path.resolve(info.projectRoot, 'Assets'),
path.resolve(info.projectRoot, 'Packages'),
path.resolve(info.projectRoot, 'Library/PackageCache'),
];
const files = [];
const seen = new Set();
for (const r of roots) this.walkCs(r, files, seen);
if (!this.lsp) this.lsp = new LspRpcClient(info.projectRoot);
const lsp = this.lsp;
// Incremental detection based on size-mtime signature
const makeSig = (abs) => {
try { const st = fs.statSync(abs); return `${st.size}-${Math.floor(st.mtimeMs)}`; } catch { return '0-0'; }
};
const wanted = new Map(files.map(abs => [this.toRel(abs, info.projectRoot), makeSig(abs)]));
const current = await this.index.getFiles();
const changed = [];
const removed = [];
for (const [rel, sig] of wanted) {
if (current.get(rel) !== sig) changed.push(rel);
}
for (const [rel] of current) if (!wanted.has(rel)) removed.push(rel);
const toRows = (uri, symbols) => {
const rel = this.toRel(uri.replace('file://', ''), info.projectRoot);
const rows = [];
const visit = (s, container) => {
const kind = this.kindFromLsp(s.kind);
const name = s.name || '';
const start = s.range?.start || s.selectionRange?.start || {};
rows.push({ path: rel, name, kind, container: container || null, ns: null, line: (start.line ?? 0) + 1, column: (start.character ?? 0) + 1 });
if (Array.isArray(s.children)) for (const c of s.children) visit(c, name || container);
};
if (Array.isArray(symbols)) for (const s of symbols) visit(s, null);
return rows;
};
// Remove vanished files
for (const rel of removed) await this.index.removeFile(rel);
// Update changed files
const absList = changed.map(rel => path.resolve(info.projectRoot, rel));
const concurrency = Math.max(1, Math.min(64, Number(params?.concurrency ?? 8)));
const reportEvery = Math.max(1, Number(params?.reportEvery ?? 100));
const startAt = Date.now();
let i = 0; let updated = 0; let processed = 0;
// LSP request with small retry/backoff
const requestWithRetry = async (uri, maxRetries = Math.max(0, Math.min(5, Number(params?.retry ?? 2)))) => {
let lastErr = null;
for (let attempt = 0; attempt <= maxRetries; attempt++) {
try {
const res = await lsp.request('textDocument/documentSymbol', { textDocument: { uri } });
return res?.result ?? res;
} catch (err) {
lastErr = err;
await new Promise(r => setTimeout(r, 200 * (attempt + 1)));
}
}
throw lastErr || new Error('documentSymbol failed');
};
const worker = async () => {
while (true) {
const idx = i++;
if (idx >= absList.length) break;
const abs = absList[idx];
try {
const uri = 'file://' + abs.replace(/\\/g, '/');
const docSymbols = await requestWithRetry(uri, 2);
const rows = toRows(uri, docSymbols);
const rel = this.toRel(abs, info.projectRoot);
await this.index.replaceSymbolsForPath(rel, rows);
await this.index.upsertFile(rel, wanted.get(rel));
updated += 1;
} catch {}
finally {
processed += 1;
if (processed % reportEvery === 0 || processed === absList.length) {
const elapsed = Math.max(1, Date.now() - startAt);
const rate = (processed * 1000 / elapsed).toFixed(1);
logger.info(`[index] progress ${processed}/${absList.length} (removed:${removed.length}) rate:${rate} f/s`);
}
}
}
};
const workers = Array.from({ length: Math.min(concurrency, absList.length) }, () => worker());
await Promise.all(workers);
const stats = await this.index.getStats();
return { success: true, updatedFiles: updated, removedFiles: removed.length, totalIndexedSymbols: stats.total, lastIndexedAt: stats.lastIndexedAt };
} catch (e) {
return {
success: false,
error: 'build_index_failed',
message: e.message,
hint: 'C# LSP not ready. Ensure manifest/auto-download and workspace paths are valid.'
};
}
}
walkCs(root, files, seen) {
try {
if (!fs.existsSync(root)) return;
const st = fs.statSync(root);
if (st.isFile()) {
if (root.endsWith('.cs') && !seen.has(root)) { files.push(root); seen.add(root); }
return;
}
const entries = fs.readdirSync(root, { withFileTypes: true });
for (const e of entries) {
if (e.name === 'obj' || e.name === 'bin' || e.name.startsWith('.')) continue;
this.walkCs(path.join(root, e.name), files, seen);
}
} catch {}
}
toRel(full, projectRoot) {
const normFull = String(full).replace(/\\/g, '/');
const normRoot = String(projectRoot).replace(/\\/g, '/').replace(/\/$/, '');
return normFull.startsWith(normRoot) ? normFull.substring(normRoot.length + 1) : normFull;
}
kindFromLsp(k) {
switch (k) {
case 5: return 'class';
case 23: return 'struct';
case 11: return 'interface';
case 10: return 'enum';
case 6: return 'method';
case 7: return 'property';
case 8: return 'field';
case 3: return 'namespace'; }
}
}