@probelabs/probe
Version:
Node.js wrapper for the probe code search tool
515 lines (512 loc) • 24.9 kB
JavaScript
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { CallToolRequestSchema, ErrorCode, ListToolsRequestSchema, McpError, } from '@modelcontextprotocol/sdk/types.js';
import { exec } from 'child_process';
import { promisify } from 'util';
import path from 'path';
import fs from 'fs-extra';
import { fileURLToPath } from 'url';
// Import from parent package
import { search, query, extract } from '../index.js';
// Parse command-line arguments
function parseArgs() {
const args = process.argv.slice(2);
const config = {};
for (let i = 0; i < args.length; i++) {
if ((args[i] === '--timeout' || args[i] === '-t') && i + 1 < args.length) {
const timeout = parseInt(args[i + 1], 10);
if (!isNaN(timeout) && timeout > 0) {
config.timeout = timeout;
console.error(`Timeout set to ${timeout} seconds`);
}
else {
console.error(`Invalid timeout value: ${args[i + 1]}. Using default.`);
}
i++; // Skip the next argument
}
else if (args[i] === '--format' && i + 1 < args.length) {
config.format = args[i + 1];
console.error(`Format set to ${config.format}`);
i++; // Skip the next argument
}
else if (args[i] === '--help' || args[i] === '-h') {
console.log(`
Probe MCP Server
Usage:
probe mcp [options]
Options:
--timeout, -t <seconds> Set timeout for search operations (default: 30)
--format <format> Set output format (json, outline-xml, etc.)
--help, -h Show this help message
`);
process.exit(0);
}
}
return config;
}
const cliConfig = parseArgs();
const execAsync = promisify(exec);
// Get the package.json to determine the version
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);
// Try multiple possible locations for package.json
let packageVersion = '0.0.0';
const possiblePaths = [
path.resolve(__dirname, '..', 'package.json'), // When installed from npm: build/../package.json
path.resolve(__dirname, '..', '..', 'package.json') // In development: src/../package.json
];
for (const packageJsonPath of possiblePaths) {
try {
if (fs.existsSync(packageJsonPath)) {
console.log(`Found package.json at: ${packageJsonPath}`);
const packageJson = JSON.parse(fs.readFileSync(packageJsonPath, 'utf-8'));
if (packageJson.version) {
packageVersion = packageJson.version;
console.log(`Using version from package.json: ${packageVersion}`);
break;
}
}
}
catch (error) {
console.error(`Error reading package.json at ${packageJsonPath}:`, error);
}
}
// If we still have 0.0.0, try to get version from npm package
if (packageVersion === '0.0.0') {
try {
// Try to get version from the package name itself
const result = await execAsync('npm list -g @probelabs/probe --json');
const npmList = JSON.parse(result.stdout);
if (npmList.dependencies && npmList.dependencies['@probelabs/probe']) {
packageVersion = npmList.dependencies['@probelabs/probe'].version;
console.log(`Using version from npm list: ${packageVersion}`);
}
}
catch (error) {
console.error('Error getting version from npm:', error);
}
}
// Get the path to the bin directory
const binDir = path.resolve(__dirname, '..', 'bin');
console.log(`Bin directory: ${binDir}`);
class ProbeServer {
constructor(timeout = 30, format) {
this.defaultTimeout = timeout;
this.defaultFormat = format;
this.server = new Server({
name: '@probelabs/probe',
version: packageVersion,
}, {
capabilities: {
tools: {},
},
});
this.setupToolHandlers();
// Error handling
this.server.onerror = (error) => console.error('[MCP Error]', error);
process.on('SIGINT', async () => {
await this.server.close();
process.exit(0);
});
}
setupToolHandlers() {
// Use the tool descriptions defined at the top of the file
this.server.setRequestHandler(ListToolsRequestSchema, async () => ({
tools: [
{
name: 'search_code',
description: "Search code in the repository using ElasticSearch. Use this tool first for any code-related questions.",
inputSchema: {
type: 'object',
properties: {
path: {
type: 'string',
description: 'Absolute path to the directory to search in (e.g., "/Users/username/projects/myproject").',
},
query: {
type: 'string',
description: 'Elastic search query. Supports logical operators (AND, OR, NOT), and grouping with parentheses. Examples: "config", "(term1 OR term2) AND term3". Use quotes for exact matches, like function or type names.',
},
filesOnly: {
type: 'boolean',
description: 'Skip AST parsing and just output unique files',
},
ignore: {
type: 'array',
items: { type: 'string' },
description: 'Custom patterns to ignore (in addition to .gitignore and common patterns)'
},
excludeFilenames: {
type: 'boolean',
description: 'Exclude filenames from being used for matching'
},
exact: {
type: 'boolean',
description: 'Perform exact search without tokenization (case-insensitive)'
},
allowTests: {
type: 'boolean',
description: 'Allow test files and test code blocks in results (disabled by default)'
},
session: {
type: 'string',
description: 'Session identifier for caching. Set to "new" if unknown, or want to reset cache. Re-use session ID returned from previous searches',
default: "new",
},
timeout: {
type: 'number',
description: 'Timeout for the search operation in seconds (default: 30)',
},
noGitignore: {
type: 'boolean',
description: 'Skip .gitignore files (will use PROBE_NO_GITIGNORE environment variable if not set)',
}
},
required: ['path', 'query']
},
},
{
name: 'query_code',
description: "Search code using ast-grep structural pattern matching. Use this tool to find specific code structures like functions, classes, or methods.",
inputSchema: {
type: 'object',
properties: {
path: {
type: 'string',
description: 'Absolute path to the directory to search in (e.g., "/Users/username/projects/myproject").',
},
pattern: {
type: 'string',
description: 'The ast-grep pattern to search for. Examples: "fn $NAME($$$PARAMS) $$$BODY" for Rust functions, "def $NAME($$$PARAMS): $$$BODY" for Python functions.',
},
language: {
type: 'string',
description: 'The programming language to search in. If not specified, the tool will try to infer the language from file extensions. Supported languages: rust, javascript, typescript, python, go, c, cpp, java, ruby, php, swift, csharp, yaml.',
},
ignore: {
type: 'array',
items: { type: 'string' },
description: 'Custom patterns to ignore (in addition to common patterns)',
},
maxResults: {
type: 'number',
description: 'Maximum number of results to return'
},
format: {
type: 'string',
enum: ['markdown', 'plain', 'json', 'color'],
description: 'Output format for the query results'
},
timeout: {
type: 'number',
description: 'Timeout for the query operation in seconds (default: 30)',
},
noGitignore: {
type: 'boolean',
description: 'Skip .gitignore files (will use PROBE_NO_GITIGNORE environment variable if not set)',
}
},
required: ['path', 'pattern']
},
},
{
name: 'extract_code',
description: "Extract code blocks from files based on line number, or symbol name. Fetch full file when line number is not provided.",
inputSchema: {
type: 'object',
properties: {
path: {
type: 'string',
description: 'Absolute path to the directory to search in (e.g., "/Users/username/projects/myproject").',
},
files: {
type: 'array',
items: { type: 'string' },
description: 'Files and lines or sybmbols to extract from: /path/to/file.rs:10, /path/to/file.rs#func_name Path should be absolute.',
},
allowTests: {
type: 'boolean',
description: 'Allow test files and test code blocks in results (disabled by default)',
},
contextLines: {
type: 'number',
description: 'Number of context lines to include before and after the extracted block when AST parsing fails to find a suitable node',
default: 0
},
format: {
type: 'string',
enum: ['markdown', 'plain', 'json'],
description: 'Output format for the extracted code',
default: 'markdown'
},
timeout: {
type: 'number',
description: 'Timeout for the extract operation in seconds (default: 30)',
},
noGitignore: {
type: 'boolean',
description: 'Skip .gitignore files (will use PROBE_NO_GITIGNORE environment variable if not set)',
}
},
required: ['path', 'files'],
},
},
],
}));
this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
if (request.params.name !== 'search_code' && request.params.name !== 'query_code' && request.params.name !== 'extract_code' &&
request.params.name !== 'probe' && request.params.name !== 'query' && request.params.name !== 'extract') {
throw new McpError(ErrorCode.MethodNotFound, `Unknown tool: ${request.params.name}`);
}
try {
let result;
// Log the incoming request for debugging
console.error(`Received request for tool: ${request.params.name}`);
console.error(`Request arguments: ${JSON.stringify(request.params.arguments)}`);
// Handle both new tool names and legacy tool names
if (request.params.name === 'search_code' || request.params.name === 'probe') {
// Ensure arguments is an object
if (!request.params.arguments || typeof request.params.arguments !== 'object') {
throw new Error("Arguments must be an object");
}
const args = request.params.arguments;
// Validate required fields
if (!args.path) {
throw new Error("Path is required in arguments");
}
if (!args.query) {
throw new Error("Query is required in arguments");
}
result = await this.executeCodeSearch(args);
}
else if (request.params.name === 'query_code' || request.params.name === 'query') {
const args = request.params.arguments;
result = await this.executeCodeQuery(args);
}
else { // extract_code or extract
const args = request.params.arguments;
result = await this.executeCodeExtract(args);
}
return {
content: [
{
type: 'text',
text: result,
},
],
};
}
catch (error) {
console.error(`Error executing ${request.params.name}:`, error);
return {
content: [
{
type: 'text',
text: `Error executing ${request.params.name}: ${error instanceof Error ? error.message : String(error)}`,
},
],
isError: true,
};
}
});
}
async executeCodeSearch(args) {
try {
// Ensure path is included in the options and is a non-empty string
if (!args.path || typeof args.path !== 'string' || args.path.trim() === '') {
throw new Error("Path is required and must be a non-empty string");
}
// Ensure query is included in the options
if (!args.query) {
throw new Error("Query is required");
}
// Log the arguments we received for debugging
console.error(`Received search arguments: path=${args.path}, query=${JSON.stringify(args.query)}`);
// Create a clean options object with only the essential properties first
const options = {
path: args.path.trim(), // Ensure path is trimmed
query: args.query
};
// Add optional parameters only if they exist
if (args.filesOnly !== undefined)
options.filesOnly = args.filesOnly;
if (args.ignore !== undefined)
options.ignore = args.ignore;
if (args.excludeFilenames !== undefined)
options.excludeFilenames = args.excludeFilenames;
if (args.exact !== undefined)
options.exact = args.exact;
if (args.maxResults !== undefined)
options.maxResults = args.maxResults;
if (args.maxTokens !== undefined)
options.maxTokens = args.maxTokens;
if (args.allowTests !== undefined)
options.allowTests = args.allowTests;
// Use noGitignore from args, or fall back to PROBE_NO_GITIGNORE environment variable
if (args.noGitignore !== undefined) {
options.noGitignore = args.noGitignore;
}
else if (process.env.PROBE_NO_GITIGNORE) {
options.noGitignore = process.env.PROBE_NO_GITIGNORE === 'true';
}
if (args.session !== undefined && args.session.trim() !== '') {
options.session = args.session;
}
else {
options.session = "new";
}
// Use timeout from args, or fall back to instance default
if (args.timeout !== undefined) {
options.timeout = args.timeout;
}
else if (this.defaultTimeout !== undefined) {
options.timeout = this.defaultTimeout;
}
// Handle format options
if (this.defaultFormat === 'outline-xml') {
// For outline-xml format, we pass it as a format flag to the search command
options.format = 'outline-xml';
}
else if (this.defaultFormat === 'json') {
options.json = true;
}
console.error("Executing search with options:", JSON.stringify(options, null, 2));
// Double-check that path is still in the options object
if (!options.path) {
console.error("Path is missing from options object after construction");
throw new Error("Path is missing from options object");
}
try {
// Call search with the options object
const result = await search(options);
return result;
}
catch (searchError) {
console.error("Search function error:", searchError);
throw new Error(`Search function error: ${searchError.message || String(searchError)}`);
}
}
catch (error) {
console.error('Error executing code search:', error);
throw new McpError('MethodNotFound', `Error executing code search: ${error.message || String(error)}`);
}
}
async executeCodeQuery(args) {
try {
// Validate required parameters
if (!args.path) {
throw new Error("Path is required");
}
if (!args.pattern) {
throw new Error("Pattern is required");
}
// Create a single options object with both pattern and path
const options = {
path: args.path,
pattern: args.pattern,
language: args.language,
ignore: args.ignore,
allowTests: args.allowTests,
maxResults: args.maxResults,
format: args.format,
timeout: args.timeout || this.defaultTimeout
};
// Use noGitignore from args, or fall back to PROBE_NO_GITIGNORE environment variable
if (args.noGitignore !== undefined) {
options.noGitignore = args.noGitignore;
}
else if (process.env.PROBE_NO_GITIGNORE) {
options.noGitignore = process.env.PROBE_NO_GITIGNORE === 'true';
}
console.log("Executing query with options:", JSON.stringify({
path: options.path,
pattern: options.pattern
}));
const result = await query(options);
return result;
}
catch (error) {
console.error('Error executing code query:', error);
throw new McpError('MethodNotFound', `Error executing code query: ${error.message || String(error)}`);
}
}
async executeCodeExtract(args) {
try {
// Validate required parameters
if (!args.path) {
throw new Error("Path is required");
}
if (!args.files || !Array.isArray(args.files) || args.files.length === 0) {
throw new Error("Files array is required and must not be empty");
}
// Create a single options object with files and other parameters
const options = {
files: args.files,
path: args.path,
allowTests: args.allowTests,
contextLines: args.contextLines,
format: args.format,
timeout: args.timeout || this.defaultTimeout
};
// Use noGitignore from args, or fall back to PROBE_NO_GITIGNORE environment variable
if (args.noGitignore !== undefined) {
options.noGitignore = args.noGitignore;
}
else if (process.env.PROBE_NO_GITIGNORE) {
options.noGitignore = process.env.PROBE_NO_GITIGNORE === 'true';
}
// Call extract with the complete options object
try {
// Track request size for token usage
const requestSize = JSON.stringify(args).length;
const requestTokens = Math.ceil(requestSize / 4); // Approximate token count
// Execute the extract command
const result = await extract(options);
// Parse the result to extract token information if available
let responseTokens = 0;
let totalTokens = 0;
// Try to extract token information from the result
if (typeof result === 'string') {
const tokenMatch = result.match(/Total tokens returned: (\d+)/);
if (tokenMatch && tokenMatch[1]) {
responseTokens = parseInt(tokenMatch[1], 10);
totalTokens = requestTokens + responseTokens;
}
// Remove spinner debug output lines
const cleanedLines = result.split('\n').filter(line => !line.match(/^⠙|^⠹|^⠧|^⠇|^⠏/) &&
!line.includes('Thinking...Extract:') &&
!line.includes('Extract results:'));
// Add token usage information if not already present
if (!result.includes('Token Usage:')) {
cleanedLines.push('');
cleanedLines.push('Token Usage:');
cleanedLines.push(` Request tokens: ${requestTokens}`);
cleanedLines.push(` Response tokens: ${responseTokens}`);
cleanedLines.push(` Total tokens: ${totalTokens}`);
}
return cleanedLines.join('\n');
}
return result;
}
catch (error) {
console.error(`Error extracting:`, error);
return `Error extracting: ${error.message || String(error)}`;
}
}
catch (error) {
console.error('Error executing code extract:', error);
throw new McpError('MethodNotFound', `Error executing code extract: ${error.message || String(error)}`);
}
}
async run() {
// The @probelabs/probe package now handles binary path management internally
// We don't need to verify or download the binary in the MCP server anymore
// Just connect the server to the transport
const transport = new StdioServerTransport();
await this.server.connect(transport);
console.error('Probe MCP server running on stdio');
}
}
const server = new ProbeServer(cliConfig.timeout, cliConfig.format);
server.run().catch(console.error);