clean

Author: arnaudroques

src/main/java/net/sourceforge/plantuml/servlet/mcp/McpServlet.java

@@ -35,12 +35,14 @@
 import jakarta.servlet.http.HttpServletResponse;
 
 import java.io.IOException;
+import java.util.LinkedHashMap;
+import java.util.Map;
 
 /**
  * MCP (Model Context Protocol) servlet for PlantUML server.
  * Handles JSON-RPC 2.0 messages over HTTP POST.
  *
- * Configuration via environment variables:
+ * Configuration via environment variables / system properties:
  * - PLANTUML_MCP_ENABLED: "true" to enable the endpoint (default: false)
  * - PLANTUML_MCP_API_KEY: Optional API key for Bearer authentication
  *
@@ -55,67 +57,77 @@
 @SuppressWarnings("serial")
 public class McpServlet extends HttpServlet {
 
+    /** Exposed for unit tests. */
     public JsonRpcServer jsonRpcServer;
+
     private String apiKey;
     private boolean mcpEnabled;
+    private ObjectMapper objectMapper;
 
     @Override
     public void init() throws ServletException {
         super.init();
 
+        // Shared ObjectMapper for both JSON-RPC and HTTP responses
+        this.objectMapper = new ObjectMapper();
+
         // Read configuration
-        mcpEnabled = isMcpEnabled();
-        apiKey = getConfigString("PLANTUML_MCP_API_KEY", "");
+        this.mcpEnabled = isMcpEnabled();
+        this.apiKey = getConfigString("PLANTUML_MCP_API_KEY", "");
 
         if (mcpEnabled) {
-            // Initialize JSON-RPC server
-            ObjectMapper mapper = new ObjectMapper();
-            McpServiceImpl service = new McpServiceImpl(mapper);
-            this.jsonRpcServer = new JsonRpcServer(mapper, service, McpService.class);
+            // Initialize JSON-RPC server with the MCP service implementation
+            McpServiceImpl service = new McpServiceImpl(objectMapper);
+            this.jsonRpcServer = new JsonRpcServer(objectMapper, service, McpService.class);
 
             log("MCP endpoint initialized" +
                 (apiKey.isEmpty() ? " (no authentication)" : " (with API key authentication)"));
         } else {
+            this.jsonRpcServer = null;
             log("MCP endpoint disabled (set PLANTUML_MCP_ENABLED=true to enable)");
         }
     }
 
     @Override
     protected void doPost(HttpServletRequest request, HttpServletResponse response)
-            throws ServletException, IOException {
+        throws ServletException, IOException {
 
         // Check if MCP is enabled
-        if (!mcpEnabled) {
+        if (!mcpEnabled || jsonRpcServer == null) {
             sendJsonError(response, HttpServletResponse.SC_NOT_FOUND,
-                         "MCP API is not enabled. Set PLANTUML_MCP_ENABLED=true to enable.");
+                "MCP API is not enabled. Set PLANTUML_MCP_ENABLED=true to enable.");
             return;
         }
 
         // Authenticate if API key is configured
         if (!authenticate(request, response)) {
-            return;
+            return; // sendJsonError already called
         }
 
-        // Set response headers
+        // CORS + JSON response headers
         response.setContentType("application/json");
         response.setCharacterEncoding("UTF-8");
         response.setHeader("Access-Control-Allow-Origin", "*");
         response.setHeader("Access-Control-Allow-Methods", "POST, OPTIONS");
         response.setHeader("Access-Control-Allow-Headers", "Content-Type, Authorization");
 
         try {
-            // Laisser jsonrpc4j gérer la requête Jakarta directement
+            // Delegate full JSON-RPC handling to jsonrpc4j (Jakarta Servlet API)
             jsonRpcServer.handle(request, response);
         } catch (Exception e) {
             log("Error handling JSON-RPC request: " + e.getMessage(), e);
-            sendJsonError(response, HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
-                "Internal server error: " + e.getMessage());
-        }    }
+
+            if (!response.isCommitted()) {
+                sendJsonError(response, HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
+                    "Internal server error: " + e.getMessage());
+            }
+        }
+    }
 
     @Override
     protected void doOptions(HttpServletRequest request, HttpServletResponse response)
-            throws ServletException, IOException {
-        // Handle CORS preflight
+        throws ServletException, IOException {
+        // CORS preflight
         response.setHeader("Access-Control-Allow-Origin", "*");
         response.setHeader("Access-Control-Allow-Methods", "POST, OPTIONS");
         response.setHeader("Access-Control-Allow-Headers", "Content-Type, Authorization");
@@ -127,15 +139,23 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
         throws ServletException, IOException {
 
         if (!mcpEnabled) {
-            response.setStatus(HttpServletResponse.SC_NOT_FOUND);
+            sendJsonError(response, HttpServletResponse.SC_NOT_FOUND, "MCP API is not enabled");
             return;
         }
 
-        // éventuellement : authenticate(request, response)
-
-        response.setContentType("text/plain");
-        response.setCharacterEncoding("UTF-8");
-        response.getWriter().println("PlantUML MCP endpoint is up.");
+        // Small informational endpoint about the MCP service
+        Map<String, Object> info = new LinkedHashMap<>();
+        info.put("service", "PlantUML MCP Server");
+        info.put("version", "1.0.0");
+        info.put("protocol", "JSON-RPC 2.0");
+        info.put("transport", "HTTP POST");
+        info.put("methods", new String[] { "initialize", "tools/list", "tools/call" });
+        info.put("tools", new String[] { "diagram_type" });
+        info.put("authentication", apiKey.isEmpty()
+            ? "none"
+            : "Bearer token required");
+
+        writeJson(response, HttpServletResponse.SC_OK, info);
     }
 
     // ============= AUTHENTICATION =============
@@ -145,17 +165,17 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
      * Returns true if authentication succeeds or is not required.
      */
     private boolean authenticate(HttpServletRequest request, HttpServletResponse response)
-            throws IOException {
+        throws IOException {
 
         // No authentication required if no API key is set
-        if (apiKey.isEmpty()) {
+        if (apiKey == null || apiKey.isEmpty()) {
             return true;
         }
 
         String authHeader = request.getHeader("Authorization");
         if (authHeader == null || !authHeader.startsWith("Bearer ")) {
             sendJsonError(response, HttpServletResponse.SC_UNAUTHORIZED,
-                         "Missing or invalid Authorization header. Expected: Bearer <token>");
+                "Missing or invalid Authorization header. Expected: Bearer <token>");
             return false;
         }
 
@@ -186,31 +206,31 @@ private String getConfigString(String key, String defaultValue) {
         return defaultValue;
     }
 
-    // ============= ERROR HANDLING =============
+    // ============= JSON HELPERS =============
 
     private void sendJsonError(HttpServletResponse response, int status, String message)
-            throws IOException {
+        throws IOException {
+
+        Map<String, Object> errorBody = new LinkedHashMap<>();
+        Map<String, Object> error = new LinkedHashMap<>();
+        error.put("code", status);
+        error.put("message", message);
+        errorBody.put("error", error);
+
+        writeJson(response, status, errorBody);
+    }
+
+    private void writeJson(HttpServletResponse response, int status, Object body)
+        throws IOException {
+
         response.setStatus(status);
         response.setContentType("application/json");
         response.setCharacterEncoding("UTF-8");
 
-        String json = "{\n" +
-                "  \"error\": {\n" +
-                "    \"code\": " + status + ",\n" +
-                "    \"message\": \"" + escapeJson(message) + "\"\n" +
-                "  }\n" +
-                "}";
+        // Basic CORS header so clients can read error/info responses as well
+        response.setHeader("Access-Control-Allow-Origin", "*");
 
-        response.getWriter().print(json);
+        objectMapper.writeValue(response.getWriter(), body);
         response.getWriter().flush();
     }
-
-    private String escapeJson(String text) {
-        if (text == null) return "";
-        return text.replace("\\", "\\\\")
-                   .replace("\"", "\\\"")
-                   .replace("\n", "\\n")
-                   .replace("\r", "\\r")
-                   .replace("\t", "\\t");
-    }
 }