426 lines
13 KiB
JavaScript
426 lines
13 KiB
JavaScript
#!/usr/bin/env node
|
|
|
|
import { spawn } from 'node:child_process';
|
|
import { existsSync } from 'node:fs';
|
|
import { resolve } from 'node:path';
|
|
|
|
const TIMEOUT_MS = parseInt(process.env.MEMTRACE_TIMEOUT_MS || '10000', 10);
|
|
const TIMEOUT_TOKEN = 'MEMTRACE_MCP_ERROR_TIMEOUT';
|
|
const SUMMARIZE_TOKEN_LIMIT = 2000;
|
|
|
|
function parseArgs() {
|
|
const args = process.argv.slice(2);
|
|
|
|
if (args.length === 0 || args.includes('--help') || args.includes('-h')) {
|
|
console.log(`Usage: node memtrace-adapter.mjs --target <symbol> --query <type> [--repo <repo_id>] [--summarize]
|
|
|
|
Arguments:
|
|
--target <symbol> Symbol name or file path to query (required for get_impact, find_dead_code)
|
|
--query <type> Query type: get_impact, find_dead_code, list_repos (required)
|
|
--repo <repo_id> Repository ID (optional — auto-detected from .memtrace-workspace if omitted)
|
|
--summarize (Optional) Apply token-budgeted hierarchical summarization for --query get_impact (output ≤ 2000 tokens)
|
|
|
|
Query types:
|
|
get_impact Fetch structural blast radius for a target symbol
|
|
find_dead_code Find dead code in a target module
|
|
list_repos List indexed repositories with freshness timestamps
|
|
|
|
Examples:
|
|
node memtrace-adapter.mjs --target "validateToken" --query get_impact
|
|
node memtrace-adapter.mjs --target "validateToken" --query get_impact --summarize
|
|
node memtrace-adapter.mjs --query list_repos
|
|
node memtrace-adapter.mjs --target "src/auth" --query find_dead_code
|
|
node memtrace-adapter.mjs --help`);
|
|
process.exit(0);
|
|
}
|
|
|
|
const result = { target: null, query: null, repo: null, summarize: false };
|
|
for (let i = 0; i < args.length; i++) {
|
|
if (args[i] === '--target' && i + 1 < args.length) {
|
|
result.target = args[++i];
|
|
} else if (args[i] === '--query' && i + 1 < args.length) {
|
|
result.query = args[++i];
|
|
} else if (args[i] === '--repo' && i + 1 < args.length) {
|
|
result.repo = args[++i];
|
|
} else if (args[i] === '--summarize') {
|
|
result.summarize = true;
|
|
} else {
|
|
fail(`Unknown argument: ${args[i]}`);
|
|
process.exit(1);
|
|
}
|
|
}
|
|
|
|
if (!result.query) {
|
|
fail('Missing required argument: --query');
|
|
process.exit(1);
|
|
}
|
|
|
|
const validQueries = ['get_impact', 'find_dead_code', 'list_repos'];
|
|
if (!validQueries.includes(result.query)) {
|
|
fail(`Invalid query type: "${result.query}". Valid: ${validQueries.join(', ')}`);
|
|
process.exit(1);
|
|
}
|
|
|
|
if ((result.query === 'get_impact' || result.query === 'find_dead_code')) {
|
|
if (result.target === null) {
|
|
fail(`Missing required argument: --target is required for --query ${result.query}`);
|
|
process.exit(1);
|
|
}
|
|
if (result.target.trim() === '') {
|
|
fail('--target must be a non-empty string');
|
|
process.exit(1);
|
|
}
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
function fail(msg) {
|
|
console.error(`ERROR: ${msg}`);
|
|
console.log(TIMEOUT_TOKEN);
|
|
}
|
|
|
|
class McpClient {
|
|
constructor() {
|
|
this.child = null;
|
|
this.stdoutBuffer = '';
|
|
this.requestId = 0;
|
|
}
|
|
|
|
spawn() {
|
|
return new Promise((resolvePromise, reject) => {
|
|
try {
|
|
this.child = spawn('memtrace', ['mcp'], {
|
|
stdio: ['pipe', 'pipe', 'pipe'],
|
|
shell: process.platform === 'win32',
|
|
windowsHide: true
|
|
});
|
|
} catch (err) {
|
|
reject(new Error(`Failed to spawn memtrace: ${err.message}`));
|
|
return;
|
|
}
|
|
|
|
const onError = (err) => {
|
|
cleanup();
|
|
const msg = err.code === 'ENOENT'
|
|
? `memtrace binary not found on PATH. Ensure memtrace is installed (npm install -g memtrace) and available.`
|
|
: `memtrace spawn error: ${err.message}`;
|
|
reject(new Error(msg));
|
|
};
|
|
|
|
const onExit = (code, signal) => {
|
|
cleanup();
|
|
if (signal) {
|
|
reject(new Error(`memtrace process terminated by signal ${signal}`));
|
|
} else if (code !== 0) {
|
|
reject(new Error(`memtrace process exited with code ${code}`));
|
|
}
|
|
};
|
|
|
|
const cleanup = () => {
|
|
if (this.child) {
|
|
this.child.removeListener('error', onError);
|
|
this.child.removeListener('exit', onExit);
|
|
this.child.stderr.removeListener('data', onStderr);
|
|
}
|
|
};
|
|
|
|
const onStderr = () => {
|
|
// MCP servers may log diagnostics to stderr — capture but don't reject
|
|
};
|
|
|
|
this.child.on('error', onError);
|
|
this.child.on('exit', onExit);
|
|
this.child.stderr.on('data', onStderr);
|
|
|
|
resolvePromise();
|
|
});
|
|
}
|
|
|
|
sendRequest(method, params = {}) {
|
|
const id = ++this.requestId;
|
|
const request = JSON.stringify({ jsonrpc: '2.0', id, method, params }) + '\n';
|
|
return new Promise((resolvePromise, reject) => {
|
|
const listener = (data) => {
|
|
this.stdoutBuffer += data.toString();
|
|
const lines = this.stdoutBuffer.split('\n');
|
|
this.stdoutBuffer = lines.pop() || '';
|
|
|
|
for (const line of lines) {
|
|
if (!line.trim()) continue;
|
|
try {
|
|
const response = JSON.parse(line);
|
|
if (response.id === id) {
|
|
this.child.stdout.removeListener('data', listener);
|
|
if (response.error) {
|
|
reject(new Error(`MCP error: ${response.error.message || JSON.stringify(response.error)}`));
|
|
} else {
|
|
resolvePromise(response.result);
|
|
}
|
|
return;
|
|
}
|
|
} catch (err) {
|
|
// Partial JSON in buffer — wait for more data
|
|
}
|
|
}
|
|
};
|
|
|
|
this.child.stdout.on('data', listener);
|
|
this.child.stdin.write(request);
|
|
});
|
|
}
|
|
|
|
async handshake() {
|
|
const capabilities = await this.sendRequest('initialize', {
|
|
protocolVersion: '2024-11-05',
|
|
capabilities: {},
|
|
clientInfo: { name: 'bmad-memtrace-adapter', version: '1.0.0' }
|
|
});
|
|
// Send initialized notification (fire-and-forget, no response expected)
|
|
this.child.stdin.write(JSON.stringify({ jsonrpc: '2.0', method: 'notifications/initialized', params: {} }) + '\n');
|
|
return capabilities;
|
|
}
|
|
|
|
async callTool(name, args) {
|
|
return this.sendRequest('tools/call', { name, arguments: args });
|
|
}
|
|
|
|
async shutdown() {
|
|
try {
|
|
await this.sendRequest('shutdown', {});
|
|
} catch (err) {
|
|
// Shutdown errors are non-fatal
|
|
}
|
|
}
|
|
|
|
kill() {
|
|
if (this.child) {
|
|
try { this.child.stdin.end(); } catch (e) {}
|
|
try { this.child.kill('SIGTERM'); } catch (e) {}
|
|
// SIGTERM is sufficient on modern systems; no need for SIGKILL timer
|
|
}
|
|
}
|
|
}
|
|
|
|
function resolveRepoId(args) {
|
|
if (args.repo) return args.repo;
|
|
|
|
// Try to auto-detect from .memtrace-workspace
|
|
let cwd = process.cwd();
|
|
const parts = cwd.split(/[\\/]/);
|
|
for (let i = parts.length; i > 0; i--) {
|
|
const dir = parts.slice(0, i).join('/');
|
|
if (existsSync(resolve('/', ...(process.platform === 'win32' ? [dir.split(':')[0] + ':', ...dir.split(':')[1]?.split('/').filter(Boolean) || []] : []), '.memtrace-workspace'))) {
|
|
// Found workspace anchor — use basename
|
|
return parts[i - 1] || parts[parts.length - 1];
|
|
}
|
|
}
|
|
|
|
// Fallback: use CWD basename
|
|
return parts[parts.length - 1];
|
|
}
|
|
|
|
async function queryGetImpact(client, target, repoId) {
|
|
const result = await client.callTool('get_impact', { target, repo_id: repoId });
|
|
return {
|
|
target,
|
|
risk_level: result.risk || 'Low',
|
|
affected_symbols: (result.affected_symbols || []).map(s => ({
|
|
name: s.name,
|
|
file: s.file || '',
|
|
depth: s.depth || 1
|
|
})),
|
|
affected_files: result.affected_files || [],
|
|
total_count: result.total_affected || result.affected_symbols?.length || 0,
|
|
elapsed_ms: 0
|
|
};
|
|
}
|
|
|
|
async function queryFindDeadCode(client, target, repoId) {
|
|
const result = await client.callTool('find_dead_code', {
|
|
repo_id: repoId,
|
|
file_path: target
|
|
});
|
|
const symbols = (result?.symbols || []).map(s => ({
|
|
name: s.name || '<unknown>',
|
|
kind: s.kind || 'Function',
|
|
file: s.file || '',
|
|
line: s.line || 0
|
|
}));
|
|
return {
|
|
query: 'find_dead_code',
|
|
target,
|
|
symbols,
|
|
total_count: symbols.length,
|
|
elapsed_ms: 0
|
|
};
|
|
}
|
|
|
|
async function queryListRepos(client) {
|
|
const result = await client.callTool('list_indexed_repositories', {});
|
|
const repos = Array.isArray(result?.repos) ? result.repos : [];
|
|
return {
|
|
query: 'list_repos',
|
|
repositories: repos.map(r => ({
|
|
repo_id: r.repo_id,
|
|
last_indexed: r.last_indexed_at || r.last_indexed || null,
|
|
total_nodes: r.total_nodes || r.nodes || 0
|
|
})),
|
|
elapsed_ms: 0
|
|
};
|
|
}
|
|
|
|
function estimateTokens(obj) {
|
|
try {
|
|
return Math.ceil(JSON.stringify(obj).length / 4 * 1.15);
|
|
} catch {
|
|
return Infinity;
|
|
}
|
|
}
|
|
|
|
function summarizeBlastRadius(result) {
|
|
const raw = result.affected_symbols;
|
|
const symbols = Array.isArray(raw) ? raw : [];
|
|
|
|
const modules = new Map();
|
|
for (const s of symbols) {
|
|
if (typeof s !== 'object' || s === null) continue;
|
|
const file = s.file || '';
|
|
const parts = file.split(/[\\/]/);
|
|
const dir = parts.slice(0, -1).join('/');
|
|
const prefix = dir ? dir.split('/').slice(0, 2).join('/') + '/' : '/';
|
|
if (!modules.has(prefix)) modules.set(prefix, []);
|
|
modules.get(prefix).push(s);
|
|
}
|
|
|
|
const isFiniteDepth = (s) => typeof s.depth === 'number' && isFinite(s.depth);
|
|
|
|
const crit = symbols
|
|
.filter(s => typeof s === 'object' && s !== null && isFiniteDepth(s) && s.depth <= 2)
|
|
.sort((a, b) => (a.depth ?? 99) - (b.depth ?? 99) || (a.name || '').localeCompare(b.name || ''))
|
|
.slice(0, 20)
|
|
.map(s => ({ name: s.name, file: s.file || '', depth: s.depth }));
|
|
|
|
const moduleImpact = {};
|
|
for (const [prefix, syms] of modules) {
|
|
const valid = syms.filter(s => typeof s === 'object' && s !== null);
|
|
const sorted = [...valid].sort((a, b) => (a.depth ?? 99) - (b.depth ?? 99) || (a.name || '').localeCompare(b.name || ''));
|
|
moduleImpact[prefix] = {
|
|
count: syms.length,
|
|
top_symbols: sorted.slice(0, 3).map(s => ({ name: s.name, file: s.file || '', depth: s.depth }))
|
|
};
|
|
}
|
|
|
|
const MAX_CRITICAL = 20;
|
|
const STAGE_CRITICAL = 10;
|
|
const MIN_CRITICAL = 5;
|
|
const MAX_MODULES = 50;
|
|
|
|
let summarized = {
|
|
total_affected: symbols.length,
|
|
total_critical: crit.length,
|
|
critical_dependents: crit,
|
|
module_impact: moduleImpact
|
|
};
|
|
summarized.token_estimate = estimateTokens(summarized);
|
|
|
|
while (summarized.token_estimate > SUMMARIZE_TOKEN_LIMIT) {
|
|
const cur = summarized.critical_dependents.length;
|
|
if (cur > STAGE_CRITICAL) {
|
|
summarized.critical_dependents = summarized.critical_dependents.slice(0, STAGE_CRITICAL);
|
|
summarized.total_critical = summarized.critical_dependents.length;
|
|
} else if (cur > MIN_CRITICAL) {
|
|
summarized.critical_dependents = summarized.critical_dependents.slice(0, MIN_CRITICAL);
|
|
summarized.total_critical = summarized.critical_dependents.length;
|
|
} else if (Object.keys(summarized.module_impact).some(k => summarized.module_impact[k].top_symbols)) {
|
|
for (const key of Object.keys(summarized.module_impact)) {
|
|
delete summarized.module_impact[key].top_symbols;
|
|
}
|
|
} else {
|
|
const entries = Object.entries(summarized.module_impact)
|
|
.sort((a, b) => b[1].count - a[1].count)
|
|
.slice(0, MAX_MODULES);
|
|
summarized.module_impact = Object.fromEntries(entries);
|
|
}
|
|
summarized.token_estimate = estimateTokens(summarized);
|
|
}
|
|
|
|
return summarized;
|
|
}
|
|
|
|
function withTimeout(promise, ms) {
|
|
let timer;
|
|
const timeout = new Promise((_, reject) => {
|
|
timer = setTimeout(() => {
|
|
reject(new TimeoutError(`Query timed out after ${ms}ms`));
|
|
}, ms);
|
|
});
|
|
return Promise.race([promise, timeout]).finally(() => clearTimeout(timer));
|
|
}
|
|
|
|
class TimeoutError extends Error {
|
|
constructor(message) {
|
|
super(message);
|
|
this.name = 'TimeoutError';
|
|
}
|
|
}
|
|
|
|
async function main() {
|
|
const args = parseArgs();
|
|
const start = Date.now();
|
|
|
|
if (args.summarize && args.query !== 'get_impact') {
|
|
console.error('WARNING: --summarize is only applicable to --query get_impact. Ignored.');
|
|
args.summarize = false;
|
|
}
|
|
|
|
const repoId = resolveRepoId(args);
|
|
const client = new McpClient();
|
|
|
|
try {
|
|
await client.spawn();
|
|
await withTimeout(client.handshake(), TIMEOUT_MS);
|
|
|
|
let queryFn;
|
|
if (args.query === 'get_impact') {
|
|
queryFn = queryGetImpact(client, args.target, repoId);
|
|
} else if (args.query === 'find_dead_code') {
|
|
queryFn = queryFindDeadCode(client, args.target, repoId);
|
|
} else if (args.query === 'list_repos') {
|
|
queryFn = queryListRepos(client);
|
|
} else {
|
|
throw new Error(`Unhandled query type: ${args.query}`);
|
|
}
|
|
|
|
let result = await withTimeout(queryFn, TIMEOUT_MS);
|
|
result.elapsed_ms = Date.now() - start;
|
|
|
|
if (args.summarize && args.query === 'get_impact') {
|
|
result.summarized = summarizeBlastRadius(result);
|
|
}
|
|
|
|
await withTimeout(client.shutdown(), 5000);
|
|
|
|
try {
|
|
console.log(JSON.stringify(result, null, 2));
|
|
} catch (serializeErr) {
|
|
fail(`Failed to serialize result: ${serializeErr.message}`);
|
|
process.exit(1);
|
|
}
|
|
process.exit(0);
|
|
} catch (err) {
|
|
client.kill();
|
|
const elapsed = Date.now() - start;
|
|
|
|
if (err instanceof TimeoutError) {
|
|
console.log(TIMEOUT_TOKEN);
|
|
console.error(`ERROR: Query timed out after ${elapsed}ms`);
|
|
} else {
|
|
fail(err.message);
|
|
}
|
|
process.exit(1);
|
|
}
|
|
}
|
|
|
|
main();
|