优化连接监控机制

Author: relat-ivity

framework/fel/java/plugins/tool-mcp-server/src/main/java/modelengine/fel/tool/mcp/server/DefaultMcpServer.java

@@ -90,9 +90,13 @@ public void onToolAdded(String name, String description, Map<String, Object> par
             log.warn("Tool addition is ignored: tool schema is null or empty. [toolName={}]", name);
             return;
         }
+        Object props = parameters.get(PROPERTIES);
+        Object reqs = parameters.get(REQUIRED);
         if (!(parameters.get(TYPE) instanceof String)
-                || parameters.get(PROPERTIES) != null && !(parameters.get(PROPERTIES) instanceof Map)
-                || parameters.get(REQUIRED) != null && !(parameters.get(REQUIRED) instanceof List)) {
+                || (props != null && (!(props instanceof Map<?, ?>)
+                || ((Map<?, ?>) props).keySet().stream().anyMatch(k -> !(k instanceof String))))
+                || (reqs != null && (!(reqs instanceof List<?>)
+                || ((List<?>) reqs).stream().anyMatch(v -> !(v instanceof String))))) {
             log.warn("Invalid parameter schema. [toolName={}]", name);
             return;
         }
@@ -112,13 +116,19 @@ public void onToolAdded(String name, String description, Map<String, Object> par
                     return new McpSchema.CallToolResult(result, false);
                 })
                 .build();
-        this.mcpSyncServer.addTool(toolSpecification);
-
         Tool tool = new Tool();
         tool.setName(name);
         tool.setDescription(description);
         tool.setInputSchema(parameters);
-        this.tools.put(name, tool);
+
+        try {
+            this.mcpSyncServer.addTool(toolSpecification);
+            this.tools.put(name, tool);
+        } catch (Exception e) {
+            log.error("Failed to add tool: {}", name, e);
+            this.tools.remove(name);
+            return;
+        }
         log.info("Tool added to MCP server. [toolName={}, description={}, schema={}]", name, description, parameters);
         this.toolsChangedObservers.forEach(ToolsChangedObserver::onToolsChanged);
     }

framework/fel/java/plugins/tool-mcp-server/src/main/java/modelengine/fel/tool/mcp/server/DefaultMcpServerBean.java

@@ -22,6 +22,8 @@
  */
 @Component
 public class DefaultMcpServerBean {
+    private final static Duration requestTimeout = Duration.ofSeconds(10);
+
     @Bean
     public DefaultMcpStreamableServerTransportProvider defaultMcpStreamableServerTransportProvider() {
         return DefaultMcpStreamableServerTransportProvider.builder()
@@ -37,7 +39,7 @@ public McpSyncServer mcpSyncServer(DefaultMcpStreamableServerTransportProvider t
                         .tools(true)             // Enable tool support
                         .logging()               // Enable logging support
                         .build())
-                .requestTimeout(Duration.ofSeconds(10))
+                .requestTimeout(requestTimeout)
                 .build();
     }
 }

framework/fel/java/plugins/tool-mcp-server/src/main/java/modelengine/fel/tool/mcp/server/DefaultMcpStreamableServerTransportProvider.java

@@ -235,7 +235,7 @@ public Object handleGet(HttpClassicServerRequest request, HttpClassicServerRespo
                 // TODO emitter.onTimeout() logger.info()
 
                 DefaultStreamableMcpSessionTransport sessionTransport = new DefaultStreamableMcpSessionTransport(
-                        sessionId, emitter);
+                        sessionId, emitter, response);
 
                 // Check if this is a replay request
                 if (request.headers().contains(HttpHeaders.LAST_EVENT_ID)) {
@@ -278,17 +278,23 @@ public void onEmittedData(TextEvent data) {
                         @Override
                         public void onCompleted() {
                             logger.info("[SSE] Completed SSE emitting for session: {}", sessionId);
+                            try {
+                                listeningStream.close();
+                            } catch (Exception e) {
+                                logger.warn("[SSE] Error closing listeningStream on complete: {}", e.getMessage());
+                            }
                         }
 
                         @Override
                         public void onFailed(Exception cause) {
                             logger.warn("[SSE] SSE failed for session: {}, cause: {}", sessionId, cause.getMessage());
+                            try {
+                                listeningStream.close();
+                            } catch (Exception e) {
+                                logger.warn("[SSE] Error closing listeningStream on failure: {}", e.getMessage());
+                            }
                         }
                     });
-                    
-                    // Add connection monitoring to detect client disconnection
-                    // This is a workaround to ensure listeningStream.close() is called when client disconnects
-                    startConnectionMonitoring(sessionId, listeningStream, response);
                 }
             });
         }
@@ -299,43 +305,6 @@ public void onFailed(Exception cause) {
         }
     }
 
-    /**
-     * Starts connection monitoring to detect client disconnection and ensure proper cleanup.
-     * This is a workaround to ensure listeningStream.close() is called when client disconnects.
-     *
-     * @param sessionId The session ID
-     * @param listeningStream The listening stream to close when connection is lost
-     * @param response The HTTP response to check for connection status
-     */
-    private void startConnectionMonitoring(String sessionId,
-                                        McpStreamableServerSession.McpStreamableServerSessionStream listeningStream,
-                                        HttpClassicServerResponse response) {
-        // Use a separate thread to periodically check connection status
-        Thread monitoringThread = new Thread(() -> {
-            try {
-                while (!Thread.currentThread().isInterrupted()) {
-                    Thread.sleep(1000); // Check every second
-                    
-                    // Check if the HTTP response is still active
-                    if (!response.isActive()) {
-                        logger.info("[SSE] Connection lost for session, completing emitter to trigger cleanup");
-                        listeningStream.close();
-                        break;
-                    }
-                }
-            } catch (InterruptedException e) {
-                logger.debug("[SSE] Connection monitoring interrupted for session");
-                Thread.currentThread().interrupt();
-            } catch (Exception e) {
-                logger.warn("[SSE] Error in connection monitoring: {}", e.getMessage());
-            }
-        });
-        
-        monitoringThread.setDaemon(true);
-        monitoringThread.setName("sse-connection-monitor-" + sessionId);
-        monitoringThread.start();
-    }
-
     /**
      * Handles POST requests for incoming JSON-RPC messages from clients.
      *
@@ -449,11 +418,11 @@ public void onCompleted() {
 
                         @Override
                         public void onFailed(Exception e) {
-                            // No action needed
+                            logger.warn("[SSE] SSE failed for session: {}, cause: {}", sessionId, e.getMessage());
                         }
                     });
 
-                    DefaultStreamableMcpSessionTransport sessionTransport = new DefaultStreamableMcpSessionTransport(sessionId, emitter);
+                    DefaultStreamableMcpSessionTransport sessionTransport = new DefaultStreamableMcpSessionTransport(sessionId, emitter, response);
 
                     try {
                         session.responseStream(jsonrpcRequest, sessionTransport)
@@ -571,15 +540,20 @@ private class DefaultStreamableMcpSessionTransport implements McpStreamableServe
         private final ReentrantLock lock = new ReentrantLock();
 
         private volatile boolean closed = false;
+        
+        private final HttpClassicServerResponse response;
 
         /**
          * Creates a new session transport with the specified ID and SSE builder.
          *
          * @param sessionId The unique identifier for this session
+         * @param emitter The emitter for sending events
+         * @param response The HTTP response for checking connection status
          */
-        DefaultStreamableMcpSessionTransport(String sessionId, Emitter<TextEvent> emitter) {
+        DefaultStreamableMcpSessionTransport(String sessionId, Emitter<TextEvent> emitter, HttpClassicServerResponse response) {
             this.sessionId = sessionId;
             this.emitter = emitter;
+            this.response = response;
             logger.info("[SSE] Building SSE for session: {} ", sessionId);
         }
 
@@ -616,6 +590,13 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, String messageId
                         logger.info("Session {} was closed during message send attempt", this.sessionId);
                         return;
                     }
+                    
+                    // Check if connection is still active before sending
+                    if (!this.response.isActive()) {
+                        logger.warn("[SSE] Connection inactive detected while sending message for session: {}", this.sessionId);
+                        this.close();
+                        return;
+                    }
 
                     String jsonText = objectMapper.writeValueAsString(message);
                     TextEvent textEvent = TextEvent.custom()