junit-mcp-server/src/index.js

A Model Context Protocol (MCP) server for generating JUnit tests for Java classes.

#!/usr/bin/env node
// Entry point for the MCP server using the MCP tool registration pattern
const { z } = require('zod');
const http = require('http');
const junitTestGenerator = require('../build/generators/junitTestGenerator');

// Define the input schema using zod
const generateJunitTestInputSchema = z.object({
  javaSource: z.string().describe('Java class source code'),
  className: z.string().describe('Name of the Java class')
});

// Define the output schema using zod
const generateJunitTestOutputSchema = z.object({
  junitTest: z.string().describe('Generated JUnit test code')
});

// Tool handler function
async function generateJunitTestHandler(params) {
  // Validate input
  const input = generateJunitTestInputSchema.parse(params);
  // Call the generator
  const result = await junitTestGenerator.run(input);
  // Validate output
  return generateJunitTestOutputSchema.parse(result);
}

// Register tools (for now, just one)
const tools = {
  generate_junit_test: {
    name: 'generate_junit_test',
    description: 'Generate JUnit tests for a given Java class source',
    inputSchema: generateJunitTestInputSchema,
    outputSchema: generateJunitTestOutputSchema,
    handler: generateJunitTestHandler
  }
};

// Simple HTTP server to expose the tool (MCP protocol can be added later)
const server = http.createServer(async (req, res) => {
  if (req.method === 'POST' && req.url === '/generate-junit-test') {
    let body = '';
    req.on('data', chunk => { body += chunk.toString(); });
    req.on('end', async () => {
      try {
        const params = JSON.parse(body);
        const result = await tools.generate_junit_test.handler(params);
        res.writeHead(200, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify(result));
      } catch (error) {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: error.message }));
      }
    });
  } else if (req.method === 'GET' && req.url === '/tools') {
    // List available tools in MCP format
    const toolList = Object.values(tools).map(tool => ({
      name: tool.name,
      description: tool.description,
      inputSchema: tool.inputSchema.describe(),
      outputSchema: tool.outputSchema.describe()
    }));
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({ tools: toolList }));
  } else {
    res.writeHead(200, { 'Content-Type': 'text/plain' });
    res.end('JUnit Test Generator API\n\nPOST /generate-junit-test\nBody: { "javaSource": "...", "className": "..." }\nGET /tools for available tools');
  }
});



// Only start HTTP server if explicitly requested
if (process.env.MCP_HTTP === '1') {
  let port = process.env.PORT ? parseInt(process.env.PORT, 10) : 49152;
  let serverInstance = null;
  function startServer(p) {
    serverInstance = server.listen(p, () => {
      console.log(`JUnit test generator server running on port ${p}`);
      console.log('Send POST requests to: http://localhost:' + p + '/generate-junit-test');
    });
  }
  server.on('error', (err) => {
    if (err.code === 'EADDRINUSE') {
      console.warn(`Port ${port} in use, trying port ${port + 1}...`);
      port++;
      setTimeout(() => startServer(port), 100);
    } else {
      throw err;
    }
  });
  function cleanup() {
    if (serverInstance) {
      serverInstance.close(() => {
        console.log('Server closed and port released.');
        process.exit(0);
      });
    } else {
      process.exit(0);
    }
  }
  process.on('SIGINT', cleanup);
  process.on('SIGTERM', cleanup);
  startServer(port);
}

// --- MCP stdio support ---
if (process.stdin.isTTY === false) {
  let inputBuffer = '';
  let pendingConfirmation = null;
  let pendingParams = null;
  process.stdin.setEncoding('utf8');
  process.stdin.on('data', chunk => {
    inputBuffer += chunk;
    let lines = inputBuffer.split('\n');
    inputBuffer = lines.pop(); // keep incomplete line
    for (const line of lines) {
      if (!line.trim()) continue;
      // If waiting for confirmation, handle yes/no
      if (pendingConfirmation) {
        const answer = line.trim().toLowerCase();
        if (answer === 'yes' || answer === 'y') {
          tools.generate_junit_test.handler(pendingParams)
            .then(output => {
              process.stdout.write(JSON.stringify({ result: output }) + '\n');
              pendingConfirmation = null;
              pendingParams = null;
            })
            .catch(error => {
              process.stdout.write(JSON.stringify({ error: error.message }) + '\n');
              pendingConfirmation = null;
              pendingParams = null;
            });
        } else {
          process.stdout.write(JSON.stringify({ cancelled: true, message: 'JUnit test generation cancelled by user.' }) + '\n');
          pendingConfirmation = null;
          pendingParams = null;
        }
        continue;
      }

      // --- Intelligence: interpret novice prompts ---
      // If the input is not valid JSON, try to parse as a simple command
      let req;
      try {
        req = JSON.parse(line);
      } catch (err) {
        // Try to interpret simple human text
        const lower = line.trim().toLowerCase();
        // Examples: "generate junit tests", "create junit for class Calculator"
        const match = lower.match(/generate|create.*junit.*(for class|for|tests?)(.*)/);
        if (match) {
          let className = null;
          // Try to extract class name
          const classMatch = lower.match(/class ([a-zA-Z0-9_]+)/);
          if (classMatch) className = classMatch[1];
          pendingConfirmation = true;
          pendingParams = { javaSource: '', className: className || 'the provided class' };
          process.stdout.write(
            `Generate JUnit tests for ${pendingParams.className}? (Type 'yes' to confirm or 'no' to cancel)\n`
          );
          continue;
        }
        // If not recognized, show help
        process.stdout.write(
          'Unrecognized input. To generate JUnit tests, type: "generate junit tests for class Calculator" or send a valid JSON request.\n'
        );
        continue;
      }

      if (req.tool === 'generate_junit_test') {
        // Ask for confirmation before generating
        const className = req.parameters.className || 'the provided class';
        pendingConfirmation = true;
        pendingParams = req.parameters;
        process.stdout.write(
          `Generate JUnit tests for ${className}? (Type 'yes' to confirm or 'no' to cancel)\n`
        );
      } else if (req.tool === 'list_tools') {
        const toolList = Object.values(tools).map(tool => ({
          name: tool.name,
          description: tool.description,
          inputSchema: tool.inputSchema.describe(),
          outputSchema: tool.outputSchema.describe()
        }));
        process.stdout.write(JSON.stringify({ tools: toolList }) + '\n');
      } else {
        process.stdout.write(JSON.stringify({ error: 'Unknown tool: ' + req.tool }) + '\n');
      }
    }
  });
}