Refactor Enhanced MCP servlet with dedicated session management and improved transport layer
- Extract session management to dedicated McpSessionManager class - Add VisitBasedMcpSession for better integration with Moqui visit tracking - Implement MoquiMcpTransport for standardized MCP message handling - Improve SSE connection lifecycle management and graceful shutdown - Add session statistics and broadcast capabilities for monitoring
Showing
4 changed files
with
682 additions
and
58 deletions
| ... | @@ -39,9 +39,8 @@ class EnhancedMcpServlet extends HttpServlet { | ... | @@ -39,9 +39,8 @@ class EnhancedMcpServlet extends HttpServlet { |
| 39 | 39 | ||
| 40 | private JsonSlurper jsonSlurper = new JsonSlurper() | 40 | private JsonSlurper jsonSlurper = new JsonSlurper() |
| 41 | 41 | ||
| 42 | // Session management for SSE connections | 42 | // Session management using dedicated session manager |
| 43 | private final Map<String, McpSession> sessions = new ConcurrentHashMap<>() | 43 | private final McpSessionManager sessionManager = new McpSessionManager() |
| 44 | private final AtomicBoolean isClosing = new AtomicBoolean(false) | ||
| 45 | 44 | ||
| 46 | @Override | 45 | @Override |
| 47 | void init(ServletConfig config) throws ServletException { | 46 | void init(ServletConfig config) throws ServletException { |
| ... | @@ -151,7 +150,7 @@ class EnhancedMcpServlet extends HttpServlet { | ... | @@ -151,7 +150,7 @@ class EnhancedMcpServlet extends HttpServlet { |
| 151 | private void handleSseConnection(HttpServletRequest request, HttpServletResponse response, ExecutionContextImpl ec) | 150 | private void handleSseConnection(HttpServletRequest request, HttpServletResponse response, ExecutionContextImpl ec) |
| 152 | throws IOException { | 151 | throws IOException { |
| 153 | 152 | ||
| 154 | if (isClosing.get()) { | 153 | if (sessionManager.isShuttingDown()) { |
| 155 | response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "Server is shutting down") | 154 | response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "Server is shutting down") |
| 156 | return | 155 | return |
| 157 | } | 156 | } |
| ... | @@ -166,10 +165,11 @@ class EnhancedMcpServlet extends HttpServlet { | ... | @@ -166,10 +165,11 @@ class EnhancedMcpServlet extends HttpServlet { |
| 166 | response.setHeader("Access-Control-Allow-Origin", "*") | 165 | response.setHeader("Access-Control-Allow-Origin", "*") |
| 167 | 166 | ||
| 168 | String sessionId = UUID.randomUUID().toString() | 167 | String sessionId = UUID.randomUUID().toString() |
| 168 | String visitId = ec.web?.visitId | ||
| 169 | 169 | ||
| 170 | // Create session transport | 170 | // Create Visit-based session transport |
| 171 | McpSession session = new McpSession(sessionId, response.writer) | 171 | VisitBasedMcpSession session = new VisitBasedMcpSession(sessionId, visitId, response.writer, ec) |
| 172 | sessions.put(sessionId, session) | 172 | sessionManager.registerSession(session) |
| 173 | 173 | ||
| 174 | try { | 174 | try { |
| 175 | // Send initial connection event with endpoint info | 175 | // Send initial connection event with endpoint info |
| ... | @@ -185,15 +185,16 @@ class EnhancedMcpServlet extends HttpServlet { | ... | @@ -185,15 +185,16 @@ class EnhancedMcpServlet extends HttpServlet { |
| 185 | 185 | ||
| 186 | // Keep connection alive with periodic pings | 186 | // Keep connection alive with periodic pings |
| 187 | int pingCount = 0 | 187 | int pingCount = 0 |
| 188 | while (!response.isCommitted() && !isClosing.get() && pingCount < 60) { // 5 minutes max | 188 | while (!response.isCommitted() && !sessionManager.isShuttingDown() && pingCount < 60) { // 5 minutes max |
| 189 | Thread.sleep(5000) // Wait 5 seconds | 189 | Thread.sleep(5000) // Wait 5 seconds |
| 190 | 190 | ||
| 191 | if (!response.isCommitted() && !isClosing.get()) { | 191 | if (!response.isCommitted() && !sessionManager.isShuttingDown()) { |
| 192 | sendSseEvent(response.writer, "ping", groovy.json.JsonOutput.toJson([ | 192 | def pingMessage = new McpSchema.JSONRPCMessage([ |
| 193 | type: "ping", | 193 | type: "ping", |
| 194 | count: pingCount, | 194 | count: pingCount, |
| 195 | timestamp: System.currentTimeMillis() | 195 | timestamp: System.currentTimeMillis() |
| 196 | ])) | 196 | ], null) |
| 197 | session.sendMessage(pingMessage) | ||
| 197 | pingCount++ | 198 | pingCount++ |
| 198 | } | 199 | } |
| 199 | } | 200 | } |
| ... | @@ -202,12 +203,13 @@ class EnhancedMcpServlet extends HttpServlet { | ... | @@ -202,12 +203,13 @@ class EnhancedMcpServlet extends HttpServlet { |
| 202 | logger.warn("Enhanced SSE connection interrupted: ${e.message}") | 203 | logger.warn("Enhanced SSE connection interrupted: ${e.message}") |
| 203 | } finally { | 204 | } finally { |
| 204 | // Clean up session | 205 | // Clean up session |
| 205 | sessions.remove(sessionId) | 206 | sessionManager.unregisterSession(sessionId) |
| 206 | try { | 207 | try { |
| 207 | sendSseEvent(response.writer, "close", groovy.json.JsonOutput.toJson([ | 208 | def closeMessage = new McpSchema.JSONRPCMessage([ |
| 208 | type: "disconnected", | 209 | type: "disconnected", |
| 209 | timestamp: System.currentTimeMillis() | 210 | timestamp: System.currentTimeMillis() |
| 210 | ])) | 211 | ], null) |
| 212 | session.sendMessage(closeMessage) | ||
| 211 | } catch (Exception e) { | 213 | } catch (Exception e) { |
| 212 | // Ignore errors during cleanup | 214 | // Ignore errors during cleanup |
| 213 | } | 215 | } |
| ... | @@ -217,31 +219,20 @@ class EnhancedMcpServlet extends HttpServlet { | ... | @@ -217,31 +219,20 @@ class EnhancedMcpServlet extends HttpServlet { |
| 217 | private void handleMessage(HttpServletRequest request, HttpServletResponse response, ExecutionContextImpl ec) | 219 | private void handleMessage(HttpServletRequest request, HttpServletResponse response, ExecutionContextImpl ec) |
| 218 | throws IOException { | 220 | throws IOException { |
| 219 | 221 | ||
| 220 | if (isClosing.get()) { | 222 | if (sessionManager.isShuttingDown()) { |
| 221 | response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "Server is shutting down") | 223 | response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "Server is shutting down") |
| 222 | return | 224 | return |
| 223 | } | 225 | } |
| 224 | 226 | ||
| 225 | // Get session ID from request parameter | 227 | // Get session from session manager |
| 226 | String sessionId = request.getParameter("sessionId") | 228 | VisitBasedMcpSession session = sessionManager.getSession(sessionId) |
| 227 | if (sessionId == null) { | ||
| 228 | response.setContentType("application/json") | ||
| 229 | response.setCharacterEncoding("UTF-8") | ||
| 230 | response.setStatus(HttpServletResponse.SC_BAD_REQUEST) | ||
| 231 | response.writer.write(groovy.json.JsonOutput.toJson([ | ||
| 232 | error: "Session ID missing in message endpoint" | ||
| 233 | ])) | ||
| 234 | return | ||
| 235 | } | ||
| 236 | |||
| 237 | // Get session from sessions map | ||
| 238 | McpSession session = sessions.get(sessionId) | ||
| 239 | if (session == null) { | 229 | if (session == null) { |
| 240 | response.setContentType("application/json") | 230 | response.setContentType("application/json") |
| 241 | response.setCharacterEncoding("UTF-8") | 231 | response.setCharacterEncoding("UTF-8") |
| 242 | response.setStatus(HttpServletResponse.SC_NOT_FOUND) | 232 | response.setStatus(HttpServletResponse.SC_NOT_FOUND) |
| 243 | response.writer.write(groovy.json.JsonOutput.toJson([ | 233 | response.writer.write(groovy.json.JsonOutput.toJson([ |
| 244 | error: "Session not found: " + sessionId | 234 | error: "Session not found: " + sessionId, |
| 235 | activeSessions: sessionManager.getActiveSessionCount() | ||
| 245 | ])) | 236 | ])) |
| 246 | return | 237 | return |
| 247 | } | 238 | } |
| ... | @@ -261,11 +252,9 @@ class EnhancedMcpServlet extends HttpServlet { | ... | @@ -261,11 +252,9 @@ class EnhancedMcpServlet extends HttpServlet { |
| 261 | // Process the method | 252 | // Process the method |
| 262 | def result = processMcpMethod(rpcRequest.method, rpcRequest.params, ec) | 253 | def result = processMcpMethod(rpcRequest.method, rpcRequest.params, ec) |
| 263 | 254 | ||
| 264 | // Send response via SSE to the specific session | 255 | // Send response via MCP transport to the specific session |
| 265 | sendSseEvent(session.writer, "response", groovy.json.JsonOutput.toJson([ | 256 | def responseMessage = new McpSchema.JSONRPCMessage(result, rpcRequest.id) |
| 266 | id: rpcRequest.id, | 257 | session.sendMessage(responseMessage) |
| 267 | result: result | ||
| 268 | ])) | ||
| 269 | 258 | ||
| 270 | response.setStatus(HttpServletResponse.SC_OK) | 259 | response.setStatus(HttpServletResponse.SC_OK) |
| 271 | 260 | ||
| ... | @@ -444,15 +433,12 @@ class EnhancedMcpServlet extends HttpServlet { | ... | @@ -444,15 +433,12 @@ class EnhancedMcpServlet extends HttpServlet { |
| 444 | @Override | 433 | @Override |
| 445 | void destroy() { | 434 | void destroy() { |
| 446 | logger.info("Destroying EnhancedMcpServlet") | 435 | logger.info("Destroying EnhancedMcpServlet") |
| 447 | isClosing.set(true) | ||
| 448 | 436 | ||
| 449 | // Close all active sessions | 437 | // Gracefully shutdown session manager |
| 450 | sessions.values().each { session -> | 438 | sessionManager.shutdownGracefully() |
| 451 | try { | 439 | |
| 452 | session.close() | 440 | super.destroy() |
| 453 | } catch (Exception e) { | 441 | } |
| 454 | logger.warn("Error closing session: ${e.message}") | ||
| 455 | } | ||
| 456 | } | 442 | } |
| 457 | sessions.clear() | 443 | sessions.clear() |
| 458 | 444 | ||
| ... | @@ -460,21 +446,16 @@ class EnhancedMcpServlet extends HttpServlet { | ... | @@ -460,21 +446,16 @@ class EnhancedMcpServlet extends HttpServlet { |
| 460 | } | 446 | } |
| 461 | 447 | ||
| 462 | /** | 448 | /** |
| 463 | * Simple session class for managing MCP SSE connections | 449 | * Broadcast message to all active sessions |
| 464 | */ | 450 | */ |
| 465 | static class McpSession { | 451 | void broadcastToAllSessions(McpSchema.JSONRPCMessage message) { |
| 466 | String sessionId | 452 | sessionManager.broadcast(message) |
| 467 | PrintWriter writer | 453 | } |
| 468 | Date createdAt | 454 | |
| 469 | 455 | /** | |
| 470 | McpSession(String sessionId, PrintWriter writer) { | 456 | * Get session statistics for monitoring |
| 471 | this.sessionId = sessionId | 457 | */ |
| 472 | this.writer = writer | 458 | Map getSessionStatistics() { |
| 473 | this.createdAt = new Date() | 459 | return sessionManager.getSessionStatistics() |
| 474 | } | ||
| 475 | |||
| 476 | void close() { | ||
| 477 | // Session cleanup logic | ||
| 478 | } | ||
| 479 | } | 460 | } |
| 480 | } | 461 | } |
| ... | \ No newline at end of file | ... | \ No newline at end of file | ... | ... |
| 1 | /* | ||
| 2 | * This software is in the public domain under CC0 1.0 Universal plus a | ||
| 3 | * Grant of Patent License. | ||
| 4 | * | ||
| 5 | * To the extent possible under law, author(s) have dedicated all | ||
| 6 | * copyright and related and neighboring rights to this software to the | ||
| 7 | * public domain worldwide. This software is distributed without any | ||
| 8 | * warranty. | ||
| 9 | * | ||
| 10 | * You should have received a copy of the CC0 Public Domain Dedication | ||
| 11 | * along with this software (see the LICENSE.md file). If not, see | ||
| 12 | * <http://creativecommons.org/publicdomain/zero/1.0/>. | ||
| 13 | */ | ||
| 14 | package org.moqui.mcp | ||
| 15 | |||
| 16 | import org.slf4j.Logger | ||
| 17 | import org.slf4j.LoggerFactory | ||
| 18 | |||
| 19 | import java.util.concurrent.ConcurrentHashMap | ||
| 20 | import java.util.concurrent.Executors | ||
| 21 | import java.util.concurrent.ScheduledExecutorService | ||
| 22 | import java.util.concurrent.TimeUnit | ||
| 23 | import java.util.concurrent.atomic.AtomicBoolean | ||
| 24 | |||
| 25 | /** | ||
| 26 | * MCP Session Manager with SDK-style capabilities | ||
| 27 | * Provides centralized session management, broadcasting, and graceful shutdown | ||
| 28 | */ | ||
| 29 | class McpSessionManager { | ||
| 30 | protected final static Logger logger = LoggerFactory.getLogger(McpSessionManager.class) | ||
| 31 | |||
| 32 | private final Map<String, VisitBasedMcpSession> sessions = new ConcurrentHashMap<>() | ||
| 33 | private final AtomicBoolean isShuttingDown = new AtomicBoolean(false) | ||
| 34 | private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(2) | ||
| 35 | |||
| 36 | // Session cleanup and monitoring | ||
| 37 | private final long sessionTimeoutMs = 30 * 60 * 1000 // 30 minutes | ||
| 38 | private final long cleanupIntervalMs = 5 * 60 * 1000 // 5 minutes | ||
| 39 | |||
| 40 | McpSessionManager() { | ||
| 41 | // Start periodic cleanup task | ||
| 42 | scheduler.scheduleAtFixedRate(this::cleanupInactiveSessions, | ||
| 43 | cleanupIntervalMs, cleanupIntervalMs, TimeUnit.MILLISECONDS) | ||
| 44 | |||
| 45 | logger.info("MCP Session Manager initialized") | ||
| 46 | } | ||
| 47 | |||
| 48 | /** | ||
| 49 | * Register a new session | ||
| 50 | */ | ||
| 51 | void registerSession(VisitBasedMcpSession session) { | ||
| 52 | if (isShuttingDown.get()) { | ||
| 53 | logger.warn("Rejecting session registration during shutdown: ${session.sessionId}") | ||
| 54 | return | ||
| 55 | } | ||
| 56 | |||
| 57 | sessions.put(session.sessionId, session) | ||
| 58 | logger.info("Registered MCP session ${session.sessionId} (total: ${sessions.size()})") | ||
| 59 | |||
| 60 | // Send welcome message to new session | ||
| 61 | def welcomeMessage = new McpSchema.JSONRPCMessage([ | ||
| 62 | type: "welcome", | ||
| 63 | sessionId: session.sessionId, | ||
| 64 | totalSessions: sessions.size(), | ||
| 65 | timestamp: System.currentTimeMillis() | ||
| 66 | ], null) | ||
| 67 | session.sendMessage(welcomeMessage) | ||
| 68 | } | ||
| 69 | |||
| 70 | /** | ||
| 71 | * Unregister a session | ||
| 72 | */ | ||
| 73 | void unregisterSession(String sessionId) { | ||
| 74 | def session = sessions.remove(sessionId) | ||
| 75 | if (session) { | ||
| 76 | logger.info("Unregistered MCP session ${sessionId} (remaining: ${sessions.size()})") | ||
| 77 | } | ||
| 78 | } | ||
| 79 | |||
| 80 | /** | ||
| 81 | * Get session by ID | ||
| 82 | */ | ||
| 83 | VisitBasedMcpSession getSession(String sessionId) { | ||
| 84 | return sessions.get(sessionId) | ||
| 85 | } | ||
| 86 | |||
| 87 | /** | ||
| 88 | * Broadcast message to all active sessions | ||
| 89 | */ | ||
| 90 | void broadcast(McpSchema.JSONRPCMessage message) { | ||
| 91 | if (isShuttingDown.get()) { | ||
| 92 | logger.warn("Rejecting broadcast during shutdown") | ||
| 93 | return | ||
| 94 | } | ||
| 95 | |||
| 96 | def inactiveSessions = [] | ||
| 97 | def activeCount = 0 | ||
| 98 | |||
| 99 | sessions.values().each { session -> | ||
| 100 | try { | ||
| 101 | if (session.isActive()) { | ||
| 102 | session.sendMessage(message) | ||
| 103 | activeCount++ | ||
| 104 | } else { | ||
| 105 | inactiveSessions << session.sessionId | ||
| 106 | } | ||
| 107 | } catch (Exception e) { | ||
| 108 | logger.warn("Error broadcasting to session ${session.sessionId}: ${e.message}") | ||
| 109 | inactiveSessions << session.sessionId | ||
| 110 | } | ||
| 111 | } | ||
| 112 | |||
| 113 | // Clean up inactive sessions | ||
| 114 | inactiveSessions.each { sessionId -> | ||
| 115 | unregisterSession(sessionId) | ||
| 116 | } | ||
| 117 | |||
| 118 | logger.info("Broadcast message to ${activeCount} active sessions (removed ${inactiveSessions.size()} inactive)") | ||
| 119 | } | ||
| 120 | |||
| 121 | /** | ||
| 122 | * Send message to specific session | ||
| 123 | */ | ||
| 124 | boolean sendToSession(String sessionId, McpSchema.JSONRPCMessage message) { | ||
| 125 | def session = sessions.get(sessionId) | ||
| 126 | if (!session) { | ||
| 127 | return false | ||
| 128 | } | ||
| 129 | |||
| 130 | try { | ||
| 131 | if (session.isActive()) { | ||
| 132 | session.sendMessage(message) | ||
| 133 | return true | ||
| 134 | } else { | ||
| 135 | unregisterSession(sessionId) | ||
| 136 | return false | ||
| 137 | } | ||
| 138 | } catch (Exception e) { | ||
| 139 | logger.warn("Error sending to session ${sessionId}: ${e.message}") | ||
| 140 | unregisterSession(sessionId) | ||
| 141 | return false | ||
| 142 | } | ||
| 143 | } | ||
| 144 | |||
| 145 | /** | ||
| 146 | * Get session statistics | ||
| 147 | */ | ||
| 148 | Map getSessionStatistics() { | ||
| 149 | def stats = [ | ||
| 150 | totalSessions: sessions.size(), | ||
| 151 | activeSessions: 0, | ||
| 152 | closingSessions: 0, | ||
| 153 | isShuttingDown: isShuttingDown.get(), | ||
| 154 | uptime: System.currentTimeMillis() - (this.@startTime ?: System.currentTimeMillis()), | ||
| 155 | sessions: [] | ||
| 156 | ] | ||
| 157 | |||
| 158 | sessions.values().each { session -> | ||
| 159 | def sessionStats = session.getSessionStats() | ||
| 160 | stats.sessions << sessionStats | ||
| 161 | |||
| 162 | if (sessionStats.active) { | ||
| 163 | stats.activeSessions++ | ||
| 164 | } | ||
| 165 | if (sessionStats.closing) { | ||
| 166 | stats.closingSessions++ | ||
| 167 | } | ||
| 168 | } | ||
| 169 | |||
| 170 | return stats | ||
| 171 | } | ||
| 172 | |||
| 173 | /** | ||
| 174 | * Initiate graceful shutdown | ||
| 175 | */ | ||
| 176 | void shutdownGracefully() { | ||
| 177 | if (!isShuttingDown.compareAndSet(false, true)) { | ||
| 178 | return // Already shutting down | ||
| 179 | } | ||
| 180 | |||
| 181 | logger.info("Initiating graceful MCP session manager shutdown") | ||
| 182 | |||
| 183 | // Send shutdown notification to all sessions | ||
| 184 | def shutdownMessage = new McpSchema.JSONRPCMessage([ | ||
| 185 | type: "server_shutdown", | ||
| 186 | message: "Server is shutting down gracefully", | ||
| 187 | timestamp: System.currentTimeMillis() | ||
| 188 | ], null) | ||
| 189 | broadcast(shutdownMessage) | ||
| 190 | |||
| 191 | // Give sessions time to receive shutdown message | ||
| 192 | scheduler.schedule({ | ||
| 193 | forceShutdown() | ||
| 194 | }, 5, TimeUnit.SECONDS) | ||
| 195 | } | ||
| 196 | |||
| 197 | /** | ||
| 198 | * Force immediate shutdown | ||
| 199 | */ | ||
| 200 | void forceShutdown() { | ||
| 201 | logger.info("Force shutting down MCP session manager") | ||
| 202 | |||
| 203 | // Close all sessions | ||
| 204 | sessions.values().each { session -> | ||
| 205 | try { | ||
| 206 | session.close() | ||
| 207 | } catch (Exception e) { | ||
| 208 | logger.warn("Error closing session ${session.sessionId}: ${e.message}") | ||
| 209 | } | ||
| 210 | } | ||
| 211 | sessions.clear() | ||
| 212 | |||
| 213 | // Shutdown scheduler | ||
| 214 | scheduler.shutdown() | ||
| 215 | try { | ||
| 216 | if (!scheduler.awaitTermination(10, TimeUnit.SECONDS)) { | ||
| 217 | scheduler.shutdownNow() | ||
| 218 | } | ||
| 219 | } catch (InterruptedException e) { | ||
| 220 | scheduler.shutdownNow() | ||
| 221 | Thread.currentThread().interrupt() | ||
| 222 | } | ||
| 223 | |||
| 224 | logger.info("MCP session manager shutdown complete") | ||
| 225 | } | ||
| 226 | |||
| 227 | /** | ||
| 228 | * Clean up inactive sessions | ||
| 229 | */ | ||
| 230 | private void cleanupInactiveSessions() { | ||
| 231 | if (isShuttingDown.get()) { | ||
| 232 | return | ||
| 233 | } | ||
| 234 | |||
| 235 | def now = System.currentTimeMillis() | ||
| 236 | def inactiveSessions = [] | ||
| 237 | |||
| 238 | sessions.values().each { session -> | ||
| 239 | def sessionStats = session.getSessionStats() | ||
| 240 | def inactiveTime = now - (sessionStats.lastActivity ?: sessionStats.createdAt.time) | ||
| 241 | |||
| 242 | if (!session.isActive() || inactiveTime > sessionTimeoutMs) { | ||
| 243 | inactiveSessions << session.sessionId | ||
| 244 | } | ||
| 245 | } | ||
| 246 | |||
| 247 | inactiveSessions.each { sessionId -> | ||
| 248 | def session = sessions.get(sessionId) | ||
| 249 | if (session) { | ||
| 250 | try { | ||
| 251 | session.closeGracefully() | ||
| 252 | } catch (Exception e) { | ||
| 253 | logger.warn("Error during cleanup of session ${sessionId}: ${e.message}") | ||
| 254 | } | ||
| 255 | unregisterSession(sessionId) | ||
| 256 | } | ||
| 257 | } | ||
| 258 | |||
| 259 | if (inactiveSessions.size() > 0) { | ||
| 260 | logger.info("Cleaned up ${inactiveSessions.size()} inactive MCP sessions") | ||
| 261 | } | ||
| 262 | } | ||
| 263 | |||
| 264 | /** | ||
| 265 | * Get active session count | ||
| 266 | */ | ||
| 267 | int getActiveSessionCount() { | ||
| 268 | return (int) sessions.values().count { it.isActive() } | ||
| 269 | } | ||
| 270 | |||
| 271 | /** | ||
| 272 | * Check if manager is shutting down | ||
| 273 | */ | ||
| 274 | boolean isShuttingDown() { | ||
| 275 | return isShuttingDown.get() | ||
| 276 | } | ||
| 277 | } | ||
| ... | \ No newline at end of file | ... | \ No newline at end of file |
| 1 | /* | ||
| 2 | * This software is in the public domain under CC0 1.0 Universal plus a | ||
| 3 | * Grant of Patent License. | ||
| 4 | * | ||
| 5 | * To the extent possible under law, author(s) have dedicated all | ||
| 6 | * copyright and related and neighboring rights to this software to the | ||
| 7 | * public domain worldwide. This software is distributed without any | ||
| 8 | * warranty. | ||
| 9 | * | ||
| 10 | * You should have received a copy of the CC0 Public Domain Dedication | ||
| 11 | * along with this software (see the LICENSE.md file). If not, see | ||
| 12 | * <http://creativecommons.org/publicdomain/zero/1.0/>. | ||
| 13 | */ | ||
| 14 | package org.moqui.mcp | ||
| 15 | |||
| 16 | import groovy.json.JsonBuilder | ||
| 17 | |||
| 18 | /** | ||
| 19 | * MCP Transport interface compatible with Servlet 4.0 and Moqui Visit system | ||
| 20 | * Provides SDK-style session management capabilities while maintaining compatibility | ||
| 21 | */ | ||
| 22 | interface MoquiMcpTransport { | ||
| 23 | /** | ||
| 24 | * Send a JSON-RPC message through this transport | ||
| 25 | * @param message The MCP JSON-RPC message to send | ||
| 26 | */ | ||
| 27 | void sendMessage(McpSchema.JSONRPCMessage message) | ||
| 28 | |||
| 29 | /** | ||
| 30 | * Close the transport gracefully, allowing in-flight messages to complete | ||
| 31 | */ | ||
| 32 | void closeGracefully() | ||
| 33 | |||
| 34 | /** | ||
| 35 | * Force close the transport immediately | ||
| 36 | */ | ||
| 37 | void close() | ||
| 38 | |||
| 39 | /** | ||
| 40 | * Check if the transport is still active | ||
| 41 | * @return true if transport is active, false otherwise | ||
| 42 | */ | ||
| 43 | boolean isActive() | ||
| 44 | |||
| 45 | /** | ||
| 46 | * Get the session ID associated with this transport | ||
| 47 | * @return the MCP session ID | ||
| 48 | */ | ||
| 49 | String getSessionId() | ||
| 50 | |||
| 51 | /** | ||
| 52 | * Get the associated Moqui Visit ID | ||
| 53 | * @return the Visit ID if available, null otherwise | ||
| 54 | */ | ||
| 55 | String getVisitId() | ||
| 56 | } | ||
| 57 | |||
| 58 | /** | ||
| 59 | * Simple implementation of MCP JSON-RPC message schema | ||
| 60 | * Compatible with MCP protocol specifications | ||
| 61 | */ | ||
| 62 | class McpSchema { | ||
| 63 | static class JSONRPCMessage { | ||
| 64 | String jsonrpc = "2.0" | ||
| 65 | Object id | ||
| 66 | String method | ||
| 67 | Map params | ||
| 68 | Object result | ||
| 69 | Map error | ||
| 70 | |||
| 71 | JSONRPCMessage(String method, Map params = null, Object id = null) { | ||
| 72 | this.method = method | ||
| 73 | this.params = params | ||
| 74 | this.id = id | ||
| 75 | } | ||
| 76 | |||
| 77 | JSONRPCMessage(Object result, Object id) { | ||
| 78 | this.result = result | ||
| 79 | this.id = id | ||
| 80 | } | ||
| 81 | |||
| 82 | JSONRPCMessage(Map error, Object id) { | ||
| 83 | this.error = error | ||
| 84 | this.id = id | ||
| 85 | } | ||
| 86 | |||
| 87 | String toJson() { | ||
| 88 | return new JsonBuilder(this).toString() | ||
| 89 | } | ||
| 90 | |||
| 91 | static JSONRPCMessage fromJson(String json) { | ||
| 92 | // Simple JSON parsing - in production would use proper JSON parser | ||
| 93 | def slurper = new groovy.json.JsonSlurper() | ||
| 94 | def data = slurper.parseText(json) | ||
| 95 | |||
| 96 | if (data.error) { | ||
| 97 | return new JSONRPCMessage(data.error, data.id) | ||
| 98 | } else if (data.result != null) { | ||
| 99 | return new JSONRPCMessage(data.result, data.id) | ||
| 100 | } else { | ||
| 101 | return new JSONRPCMessage(data.method, data.params, data.id) | ||
| 102 | } | ||
| 103 | } | ||
| 104 | } | ||
| 105 | } | ||
| ... | \ No newline at end of file | ... | \ No newline at end of file |
| 1 | /* | ||
| 2 | * This software is in the public domain under CC0 1.0 Universal plus a | ||
| 3 | * Grant of Patent License. | ||
| 4 | * | ||
| 5 | * To the extent possible under law, author(s) have dedicated all | ||
| 6 | * copyright and related and neighboring rights to this software to the | ||
| 7 | * public domain worldwide. This software is distributed without any | ||
| 8 | * warranty. | ||
| 9 | * | ||
| 10 | * You should have received a copy of the CC0 Public Domain Dedication | ||
| 11 | * along with this software (see the LICENSE.md file). If not, see | ||
| 12 | * <http://creativecommons.org/publicdomain/zero/1.0/>. | ||
| 13 | */ | ||
| 14 | package org.moqui.mcp | ||
| 15 | |||
| 16 | import org.moqui.context.ExecutionContext | ||
| 17 | import org.moqui.impl.context.ExecutionContextImpl | ||
| 18 | import org.slf4j.Logger | ||
| 19 | import org.slf4j.LoggerFactory | ||
| 20 | |||
| 21 | import java.util.concurrent.ConcurrentHashMap | ||
| 22 | import java.util.concurrent.atomic.AtomicBoolean | ||
| 23 | import java.util.concurrent.atomic.AtomicLong | ||
| 24 | |||
| 25 | /** | ||
| 26 | * MCP Session implementation that integrates with Moqui's Visit system | ||
| 27 | * Provides SDK-style session management while leveraging Moqui's built-in tracking | ||
| 28 | */ | ||
| 29 | class VisitBasedMcpSession implements MoquiMcpTransport { | ||
| 30 | protected final static Logger logger = LoggerFactory.getLogger(VisitBasedMcpSession.class) | ||
| 31 | |||
| 32 | private final String sessionId | ||
| 33 | private final String visitId | ||
| 34 | private final PrintWriter writer | ||
| 35 | private final ExecutionContextImpl ec | ||
| 36 | private final AtomicBoolean active = new AtomicBoolean(true) | ||
| 37 | private final AtomicBoolean closing = new AtomicBoolean(false) | ||
| 38 | private final AtomicLong messageCount = new AtomicLong(0) | ||
| 39 | private final Date createdAt | ||
| 40 | |||
| 41 | // MCP session metadata stored in Visit context | ||
| 42 | private final Map<String, Object> sessionMetadata = new ConcurrentHashMap<>() | ||
| 43 | |||
| 44 | VisitBasedMcpSession(String sessionId, String visitId, PrintWriter writer, ExecutionContextImpl ec) { | ||
| 45 | this.sessionId = sessionId | ||
| 46 | this.visitId = visitId | ||
| 47 | this.writer = writer | ||
| 48 | this.ec = ec | ||
| 49 | this.createdAt = new Date() | ||
| 50 | |||
| 51 | // Initialize session metadata in Visit context | ||
| 52 | initializeSessionMetadata() | ||
| 53 | } | ||
| 54 | |||
| 55 | private void initializeSessionMetadata() { | ||
| 56 | try { | ||
| 57 | // Store MCP session info in Visit context for persistence | ||
| 58 | if (visitId && ec) { | ||
| 59 | def visit = ec.entity.find("moqui.server.Visit").condition("visitId", visitId).one() | ||
| 60 | if (visit) { | ||
| 61 | // Store MCP session metadata as JSON in Visit's context or a separate field | ||
| 62 | sessionMetadata.put("mcpSessionId", sessionId) | ||
| 63 | sessionMetadata.put("mcpCreatedAt", createdAt.time) | ||
| 64 | sessionMetadata.put("mcpProtocolVersion", "2025-06-18") | ||
| 65 | sessionMetadata.put("mcpTransportType", "SSE") | ||
| 66 | |||
| 67 | logger.info("MCP Session ${sessionId} initialized with Visit ${visitId}") | ||
| 68 | } | ||
| 69 | } | ||
| 70 | } catch (Exception e) { | ||
| 71 | logger.warn("Failed to initialize session metadata for Visit ${visitId}: ${e.message}") | ||
| 72 | } | ||
| 73 | } | ||
| 74 | |||
| 75 | @Override | ||
| 76 | void sendMessage(McpSchema.JSONRPCMessage message) { | ||
| 77 | if (!active.get() || closing.get()) { | ||
| 78 | logger.warn("Attempted to send message on inactive or closing session ${sessionId}") | ||
| 79 | return | ||
| 80 | } | ||
| 81 | |||
| 82 | try { | ||
| 83 | String jsonMessage = message.toJson() | ||
| 84 | sendSseEvent("message", jsonMessage) | ||
| 85 | messageCount.incrementAndGet() | ||
| 86 | |||
| 87 | // Update session activity in Visit | ||
| 88 | updateSessionActivity() | ||
| 89 | |||
| 90 | } catch (Exception e) { | ||
| 91 | logger.error("Failed to send message on session ${sessionId}: ${e.message}") | ||
| 92 | if (e.message?.contains("disconnected") || e.message?.contains("Client disconnected")) { | ||
| 93 | close() | ||
| 94 | } | ||
| 95 | } | ||
| 96 | } | ||
| 97 | |||
| 98 | @Override | ||
| 99 | void closeGracefully() { | ||
| 100 | if (!active.compareAndSet(true, false)) { | ||
| 101 | return // Already closed | ||
| 102 | } | ||
| 103 | |||
| 104 | closing.set(true) | ||
| 105 | logger.info("Gracefully closing MCP session ${sessionId}") | ||
| 106 | |||
| 107 | try { | ||
| 108 | // Send graceful shutdown notification | ||
| 109 | def shutdownMessage = new McpSchema.JSONRPCMessage([ | ||
| 110 | type: "shutdown", | ||
| 111 | sessionId: sessionId, | ||
| 112 | timestamp: System.currentTimeMillis() | ||
| 113 | ], null) | ||
| 114 | sendMessage(shutdownMessage) | ||
| 115 | |||
| 116 | // Give some time for message to be sent | ||
| 117 | Thread.sleep(100) | ||
| 118 | |||
| 119 | } catch (Exception e) { | ||
| 120 | logger.warn("Error during graceful shutdown of session ${sessionId}: ${e.message}") | ||
| 121 | } finally { | ||
| 122 | close() | ||
| 123 | } | ||
| 124 | } | ||
| 125 | |||
| 126 | @Override | ||
| 127 | void close() { | ||
| 128 | if (!active.compareAndSet(true, false)) { | ||
| 129 | return // Already closed | ||
| 130 | } | ||
| 131 | |||
| 132 | logger.info("Closing MCP session ${sessionId} (messages sent: ${messageCount.get()})") | ||
| 133 | |||
| 134 | try { | ||
| 135 | // Update Visit with session end info | ||
| 136 | updateSessionEnd() | ||
| 137 | |||
| 138 | // Send final close event if writer is still available | ||
| 139 | if (writer && !writer.checkError()) { | ||
| 140 | sendSseEvent("close", groovy.json.JsonOutput.toJson([ | ||
| 141 | type: "disconnected", | ||
| 142 | sessionId: sessionId, | ||
| 143 | messageCount: messageCount.get(), | ||
| 144 | timestamp: System.currentTimeMillis() | ||
| 145 | ])) | ||
| 146 | } | ||
| 147 | |||
| 148 | } catch (Exception e) { | ||
| 149 | logger.warn("Error during session close ${sessionId}: ${e.message}") | ||
| 150 | } | ||
| 151 | } | ||
| 152 | |||
| 153 | @Override | ||
| 154 | boolean isActive() { | ||
| 155 | return active.get() && !closing.get() && writer && !writer.checkError() | ||
| 156 | } | ||
| 157 | |||
| 158 | @Override | ||
| 159 | String getSessionId() { | ||
| 160 | return sessionId | ||
| 161 | } | ||
| 162 | |||
| 163 | @Override | ||
| 164 | String getVisitId() { | ||
| 165 | return visitId | ||
| 166 | } | ||
| 167 | |||
| 168 | /** | ||
| 169 | * Get session statistics | ||
| 170 | */ | ||
| 171 | Map getSessionStats() { | ||
| 172 | return [ | ||
| 173 | sessionId: sessionId, | ||
| 174 | visitId: visitId, | ||
| 175 | createdAt: createdAt, | ||
| 176 | messageCount: messageCount.get(), | ||
| 177 | active: active.get(), | ||
| 178 | closing: closing.get(), | ||
| 179 | duration: System.currentTimeMillis() - createdAt.time | ||
| 180 | ] | ||
| 181 | } | ||
| 182 | |||
| 183 | /** | ||
| 184 | * Send SSE event with proper formatting | ||
| 185 | */ | ||
| 186 | private void sendSseEvent(String eventType, String data) throws IOException { | ||
| 187 | if (!writer || writer.checkError()) { | ||
| 188 | throw new IOException("Writer is closed or client disconnected") | ||
| 189 | } | ||
| 190 | |||
| 191 | writer.write("event: " + eventType + "\n") | ||
| 192 | writer.write("data: " + data + "\n\n") | ||
| 193 | writer.flush() | ||
| 194 | |||
| 195 | if (writer.checkError()) { | ||
| 196 | throw new IOException("Client disconnected during write") | ||
| 197 | } | ||
| 198 | } | ||
| 199 | |||
| 200 | /** | ||
| 201 | * Update session activity in Visit record | ||
| 202 | */ | ||
| 203 | private void updateSessionActivity() { | ||
| 204 | try { | ||
| 205 | if (visitId && ec) { | ||
| 206 | // Update Visit with latest activity | ||
| 207 | ec.service.sync().name("update", "moqui.server.Visit") | ||
| 208 | .parameters([ | ||
| 209 | visitId: visitId, | ||
| 210 | thruDate: ec.user.getNowTimestamp() | ||
| 211 | ]) | ||
| 212 | .call() | ||
| 213 | |||
| 214 | // Could also update a custom field for MCP-specific activity | ||
| 215 | sessionMetadata.put("mcpLastActivity", System.currentTimeMillis()) | ||
| 216 | sessionMetadata.put("mcpMessageCount", messageCount.get()) | ||
| 217 | } | ||
| 218 | } catch (Exception e) { | ||
| 219 | logger.debug("Failed to update session activity: ${e.message}") | ||
| 220 | } | ||
| 221 | } | ||
| 222 | |||
| 223 | /** | ||
| 224 | * Update Visit record with session end information | ||
| 225 | */ | ||
| 226 | private void updateSessionEnd() { | ||
| 227 | try { | ||
| 228 | if (visitId && ec) { | ||
| 229 | // Update Visit with session end info | ||
| 230 | ec.service.sync().name("update", "moqui.server.Visit") | ||
| 231 | .parameters([ | ||
| 232 | visitId: visitId, | ||
| 233 | thruDate: ec.user.getNowTimestamp() | ||
| 234 | ]) | ||
| 235 | .call() | ||
| 236 | |||
| 237 | // Store final session metadata | ||
| 238 | sessionMetadata.put("mcpEndedAt", System.currentTimeMillis()) | ||
| 239 | sessionMetadata.put("mcpFinalMessageCount", messageCount.get()) | ||
| 240 | |||
| 241 | logger.info("Updated Visit ${visitId} with MCP session end info") | ||
| 242 | } | ||
| 243 | } catch (Exception e) { | ||
| 244 | logger.warn("Failed to update session end for Visit ${visitId}: ${e.message}") | ||
| 245 | } | ||
| 246 | } | ||
| 247 | |||
| 248 | /** | ||
| 249 | * Get session metadata | ||
| 250 | */ | ||
| 251 | Map getSessionMetadata() { | ||
| 252 | return new HashMap<>(sessionMetadata) | ||
| 253 | } | ||
| 254 | |||
| 255 | /** | ||
| 256 | * Add custom metadata to session | ||
| 257 | */ | ||
| 258 | void addSessionMetadata(String key, Object value) { | ||
| 259 | sessionMetadata.put(key, value) | ||
| 260 | } | ||
| 261 | } | ||
| ... | \ No newline at end of file | ... | \ No newline at end of file |
-
Please register or sign in to post a comment