From d9cbc5381a8f1f1500418ca8a245274b4693d7b1 Mon Sep 17 00:00:00 2001
From: samanhappy <samanhappy@gmail.com>
Date: Sat, 7 Jun 2025 20:41:51 +0800
Subject: [PATCH] feat: implement keep-alive functionality for SSE connections
 (#166)

Co-authored-by: samanhappy@qq.com <my6051199>
---
 .eslintrc.json                      |  2 +-
 src/controllers/serverController.ts | 10 +++++++
 src/services/mcpService.ts          | 42 +++++++++++++++++++++++++++++
 src/types/index.ts                  |  2 ++
 4 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/.eslintrc.json b/.eslintrc.json
index f742495..89871a4 100644
--- a/.eslintrc.json
+++ b/.eslintrc.json
@@ -20,6 +20,6 @@
       }
     ],
     "@typescript-eslint/no-explicit-any": "off",
-    "no-undef": "off",
+    "no-undef": "off"
   }
 }
diff --git a/src/controllers/serverController.ts b/src/controllers/serverController.ts
index caa23a2..c528332 100644
--- a/src/controllers/serverController.ts
+++ b/src/controllers/serverController.ts
@@ -107,6 +107,11 @@ export const createServer = async (req: Request, res: Response): Promise<void> =
       return;
     }
 
+    // Set default keep-alive interval for SSE servers if not specified
+    if ((config.type === 'sse' || (!config.type && config.url)) && !config.keepAliveInterval) {
+      config.keepAliveInterval = 60000; // Default 60 seconds for SSE servers
+    }
+
     const result = await addServer(name, config);
     if (result.success) {
       notifyToolChanged();
@@ -224,6 +229,11 @@ export const updateServer = async (req: Request, res: Response): Promise<void> =
       return;
     }
 
+    // Set default keep-alive interval for SSE servers if not specified
+    if ((config.type === 'sse' || (!config.type && config.url)) && !config.keepAliveInterval) {
+      config.keepAliveInterval = 60000; // Default 60 seconds for SSE servers
+    }
+
     const result = await updateMcpServer(name, config);
     if (result.success) {
       notifyToolChanged();
diff --git a/src/services/mcpService.ts b/src/services/mcpService.ts
index e4f1709..2fde34c 100644
--- a/src/services/mcpService.ts
+++ b/src/services/mcpService.ts
@@ -13,6 +13,38 @@ import { saveToolsAsVectorEmbeddings, searchToolsByVector } from './vectorSearch
 
 const servers: { [sessionId: string]: Server } = {};
 
+// Helper function to set up keep-alive ping for SSE connections
+const setupKeepAlive = (serverInfo: ServerInfo, serverConfig: ServerConfig): void => {
+  // Only set up keep-alive for SSE connections
+  if (!(serverInfo.transport instanceof SSEClientTransport)) {
+    return;
+  }
+
+  // Clear any existing interval first
+  if (serverInfo.keepAliveIntervalId) {
+    clearInterval(serverInfo.keepAliveIntervalId);
+  }
+
+  // Use configured interval or default to 60 seconds for SSE
+  const interval = serverConfig.keepAliveInterval || 60000;
+
+  serverInfo.keepAliveIntervalId = setInterval(async () => {
+    try {
+      if (serverInfo.client && serverInfo.status === 'connected') {
+        await serverInfo.client.ping();
+        console.log(`Keep-alive ping successful for server: ${serverInfo.name}`);
+      }
+    } catch (error) {
+      console.warn(`Keep-alive ping failed for server ${serverInfo.name}:`, error);
+      // TODO Consider handling reconnection logic here if needed
+    }
+  }, interval);
+
+  console.log(
+    `Keep-alive ping set up for server ${serverInfo.name} with interval ${interval / 1000} seconds`,
+  );
+};
+
 export const initUpstreamServers = async (): Promise<void> => {
   await registerAllTools(true);
 };
@@ -210,6 +242,9 @@ export const initializeClientsFromSettings = (isInit: boolean): ServerInfo[] =>
             serverInfo.status = 'connected';
             serverInfo.error = null;
 
+            // Set up keep-alive ping for SSE connections
+            setupKeepAlive(serverInfo, conf);
+
             // Save tools as vector embeddings for search
             saveToolsAsVectorEmbeddings(name, serverInfo.tools);
           })
@@ -389,6 +424,13 @@ export const updateMcpServer = async (
 function closeServer(name: string) {
   const serverInfo = serverInfos.find((serverInfo) => serverInfo.name === name);
   if (serverInfo && serverInfo.client && serverInfo.transport) {
+    // Clear keep-alive interval if exists
+    if (serverInfo.keepAliveIntervalId) {
+      clearInterval(serverInfo.keepAliveIntervalId);
+      serverInfo.keepAliveIntervalId = undefined;
+      console.log(`Cleared keep-alive interval for server: ${serverInfo.name}`);
+    }
+
     serverInfo.client.close();
     serverInfo.transport.close();
     console.log(`Closed client and transport for server: ${serverInfo.name}`);
diff --git a/src/types/index.ts b/src/types/index.ts
index 2ba454a..7007e02 100644
--- a/src/types/index.ts
+++ b/src/types/index.ts
@@ -105,6 +105,7 @@ export interface ServerConfig {
   env?: Record<string, string>; // Environment variables
   headers?: Record<string, string>; // HTTP headers for SSE/streamable-http servers
   enabled?: boolean; // Flag to enable/disable the server
+  keepAliveInterval?: number; // Keep-alive ping interval in milliseconds (default: 60000ms for SSE servers)
   tools?: Record<string, { enabled: boolean; description?: string }>; // Tool-specific configurations with enable/disable state and custom descriptions
 }
 
@@ -118,6 +119,7 @@ export interface ServerInfo {
   transport?: SSEClientTransport | StdioClientTransport | StreamableHTTPClientTransport; // Transport mechanism used
   createTime: number; // Timestamp of when the server was created
   enabled?: boolean; // Flag to indicate if the server is enabled
+  keepAliveIntervalId?: NodeJS.Timeout; // Timer ID for keep-alive ping interval
 }
 
 // Details about a tool available on the server