From 4a3d4cf5686fd3bc60f8852b4347a2a055a48777 Mon Sep 17 00:00:00 2001 From: matevip Date: Thu, 11 Jun 2026 09:58:40 +0800 Subject: [PATCH] fix(mcp): auto-heal stale MCP connections after server restart (#317) --- .../mcp/event/McpConnectionLostEvent.java | 25 +++++ .../runtime/CwdAwareStdioClientTransport.java | 35 +++++++ .../tool/mcp/runtime/McpClientManager.java | 99 ++++++++++++++++--- .../tool/mcp/service/McpServerService.java | 47 +++++++++ .../runtime/McpClientManagerSnapshotTest.java | 98 ++++++++++++++++++ 5 files changed, 292 insertions(+), 12 deletions(-) create mode 100644 mateclaw-server/src/main/java/vip/mate/tool/mcp/event/McpConnectionLostEvent.java create mode 100644 mateclaw-server/src/test/java/vip/mate/tool/mcp/runtime/McpClientManagerSnapshotTest.java diff --git a/mateclaw-server/src/main/java/vip/mate/tool/mcp/event/McpConnectionLostEvent.java b/mateclaw-server/src/main/java/vip/mate/tool/mcp/event/McpConnectionLostEvent.java new file mode 100644 index 00000000..ef7a0281 --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/tool/mcp/event/McpConnectionLostEvent.java @@ -0,0 +1,25 @@ +package vip.mate.tool.mcp.event; + +/** + * Published when a previously-connected MCP server's transport is detected to + * have gone dead at runtime — either a live {@code listTools()} threw because + * the held connection went stale, or a stdio subprocess exited on its own. + * + *

Unlike {@link McpServerChangedEvent} (which announces a state change that + * already happened), this is a request to heal: {@code McpServerService} + * listens, reloads the server config, and triggers an asynchronous reconnect + * (debounced so a crash-looping server can't spin the reconnect executor). + * + *

Why an event rather than a direct call: {@code McpClientManager} / + * {@code CwdAwareStdioClientTransport} are pure runtime components with no DB + * access, and {@code McpServerService} already depends on the manager. Routing + * the reconnect request through {@code ApplicationEventPublisher} keeps the + * dependency one-directional and mirrors the existing + * {@link McpServerChangedEvent} pattern. + * + * @param serverId the MCP server whose connection was lost + * @param reason short human-readable cause, for logs only + * @author MateClaw Team + */ +public record McpConnectionLostEvent(Long serverId, String reason) { +} diff --git a/mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/CwdAwareStdioClientTransport.java b/mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/CwdAwareStdioClientTransport.java index eb1e8889..0d948899 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/CwdAwareStdioClientTransport.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/CwdAwareStdioClientTransport.java @@ -49,6 +49,12 @@ public class CwdAwareStdioClientTransport extends StdioClientTransport { * parent's own {@code isClosing} signal (which our override bypasses). */ private volatile boolean closing = false; + /** Invoked when the child process exits while {@link #closing} is still + * {@code false} — i.e. the server died on its own rather than being torn + * down by {@link #closeGracefully()}. Lets the manager request a reconnect. + * Null until armed by the manager for a long-lived (non-test) client. */ + private volatile Runnable onUnexpectedExit; + /** I/O threads spawned by {@link #connect}, interrupted on shutdown. */ private final List ioThreads = new CopyOnWriteArrayList<>(); @@ -78,6 +84,17 @@ public class CwdAwareStdioClientTransport extends StdioClientTransport { this.cwd = cwd; } + /** + * Arm a callback fired when the child process exits unexpectedly (not via + * {@link #closeGracefully()}). The manager uses this to request an + * asynchronous reconnect — stdio is the one transport the MCP SDK cannot + * self-heal, because a dead subprocess can only be recovered by respawning + * it, which the SDK's lazy re-initialization never does. + */ + public void setOnUnexpectedExit(Runnable handler) { + this.onUnexpectedExit = handler; + } + /** * Override parent {@code connect()} to add resilient inbound processing. * @@ -133,6 +150,24 @@ public class CwdAwareStdioClientTransport extends StdioClientTransport { // Retain the child so closeGracefully() can terminate it. this.startedProcess = process; + // Detect an unexpected death (crash / external `kill` / the user + // restarting the MCP service). Suppressed when `closing` is set, + // which is how an intentional close/replace tears the process down. + process.onExit().thenAccept(p -> { + if (closing) { + return; + } + Runnable exitHandler = this.onUnexpectedExit; + log.warn("MCP stdio process exited unexpectedly (exit code {})", p.exitValue()); + if (exitHandler != null) { + try { + exitHandler.run(); + } catch (Exception e) { + log.warn("MCP stdio unexpected-exit handler failed: {}", e.getMessage()); + } + } + }); + // --- Resilient inbound reader (the key fix) --- Thread inboundThread = new Thread(() -> { try (BufferedReader reader = new BufferedReader( diff --git a/mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/McpClientManager.java b/mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/McpClientManager.java index ad835a4f..a83024d7 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/McpClientManager.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/McpClientManager.java @@ -13,7 +13,10 @@ import io.modelcontextprotocol.spec.McpSchema; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.mcp.SyncMcpToolCallbackProvider; import org.springframework.ai.tool.ToolCallback; +import org.springframework.context.ApplicationEventPublisher; import org.springframework.stereotype.Component; +import vip.mate.tool.mcp.event.McpConnectionLostEvent; +import vip.mate.tool.mcp.event.McpServerChangedEvent; import vip.mate.tool.mcp.model.McpServerEntity; import jakarta.annotation.PreDestroy; @@ -51,6 +54,21 @@ public class McpClientManager { /** serverId -> discovered tools metadata */ private final ConcurrentHashMap> toolsCache = new ConcurrentHashMap<>(); + /** + * serverId -> last successfully-built, prefix-wrapped tool callbacks. + * Served as a fallback when a live {@code listTools()} momentarily fails + * (e.g. the upstream server just restarted), so the agent keeps seeing the + * MCP tools instead of dropping the whole server and falling back to + * non-MCP tools. Refreshed on every successful collection. + */ + private final ConcurrentHashMap> lastGoodCallbacks = new ConcurrentHashMap<>(); + + private final ApplicationEventPublisher eventPublisher; + + public McpClientManager(ApplicationEventPublisher eventPublisher) { + this.eventPublisher = eventPublisher; + } + /** serverId -> connection result info */ private final ConcurrentHashMap connectionResults = new ConcurrentHashMap<>(); @@ -102,6 +120,7 @@ public class McpClientManager { McpSyncClient old = clients.remove(serverId); toolsCache.remove(serverId); connectionResults.remove(serverId); + lastGoodCallbacks.remove(serverId); if (old != null) { closeClientSafely(serverId, old); } @@ -119,7 +138,7 @@ public class McpClientManager { long start = System.currentTimeMillis(); McpSyncClient testClient = null; try { - testClient = buildClient(server); + testClient = buildClient(server, false); testClient.initialize(); List tools = testClient.listTools().tools(); long latency = System.currentTimeMillis() - start; @@ -170,17 +189,46 @@ public class McpClientManager { try { SyncMcpToolCallbackProvider provider = new SyncMcpToolCallbackProvider(entry.getValue()); ToolCallback[] cbs = provider.getToolCallbacks(); - if (cbs == null || cbs.length == 0) { + if (cbs != null && cbs.length > 0) { + List wrapped = wrapServerCallbacks(serverId, cbs); + lastGoodCallbacks.put(serverId, wrapped); + allCallbacks.addAll(wrapped); continue; } - allCallbacks.addAll(wrapServerCallbacks(serverId, cbs)); + // Live call succeeded but returned nothing. This can be a + // transient post-restart state while the SDK re-initializes; + // keep serving the last good snapshot rather than dropping the + // server. A server that legitimately has no tools simply has no + // snapshot and contributes nothing — same as before. + addSnapshot(allCallbacks, serverId); } catch (Exception e) { - log.warn("Failed to get tool callbacks from MCP server {}: {}", serverId, e.getMessage()); + // A live listTools() failure usually means the upstream server + // restarted and the held connection went stale. Instead of + // dropping the whole server (which makes the agent fall back to + // non-MCP tools), keep serving the last known-good callbacks and + // ask the service layer to reconnect. For streamable_http the + // stale callbacks self-heal on call via the SDK's lazy + // re-initialization; for stdio/sse the async reconnect rebuilds + // them and clears the agent cache. + log.warn("MCP server {} listTools failed; serving cached snapshot and requesting reconnect: {}", + serverId, e.getMessage()); + eventPublisher.publishEvent(new McpConnectionLostEvent(serverId, "listTools-failed")); + addSnapshot(allCallbacks, serverId); } } return allCallbacks; } + /** Append the last known-good callbacks for {@code serverId}, if any. */ + private void addSnapshot(List out, long serverId) { + List snapshot = lastGoodCallbacks.get(serverId); + if (snapshot != null && !snapshot.isEmpty()) { + log.debug("Serving {} cached MCP tool callbacks for server {} while it reconnects", + snapshot.size(), serverId); + out.addAll(snapshot); + } + } + /** * Apply per-server collision detection and wrap each surviving callback * with its prefixed name. Walks {@code cbs} and the matching decision @@ -250,6 +298,7 @@ public class McpClientManager { clients.clear(); toolsCache.clear(); connectionResults.clear(); + lastGoodCallbacks.clear(); // 不清除 serverLocks：closeAll 后 server 可能被重新 connect， // 保留 lock 对象确保后续操作仍有互斥保护 } @@ -260,7 +309,7 @@ public class McpClientManager { long start = System.currentTimeMillis(); McpSyncClient newClient = null; try { - newClient = buildClient(server); + newClient = buildClient(server, true); newClient.initialize(); // Discover tools @@ -301,9 +350,20 @@ public class McpClientManager { } } - private McpSyncClient buildClient(McpServerEntity server) { + /** + * Build a sync MCP client for {@code server}. + * + * @param managed {@code true} for long-lived clients placed in the active + * pool (connect/replace), {@code false} for throwaway clients + * (testConnection). Only managed clients arm the runtime + * self-healing hooks — a server-pushed {@code tools/list_changed} + * notification refreshes agent graphs, and a stdio subprocess + * death requests a reconnect. A throwaway test client must stay + * side-effect free. + */ + private McpSyncClient buildClient(McpServerEntity server, boolean managed) { McpClientTransport transport = switch (server.getTransport()) { - case "stdio" -> buildStdioTransport(server); + case "stdio" -> buildStdioTransport(server, managed); case "sse" -> buildSseTransport(server); case "streamable_http" -> buildStreamableHttpTransport(server); default -> throw new IllegalArgumentException("Unsupported transport: " + server.getTransport()); @@ -314,12 +374,19 @@ public class McpClientManager { Duration requestTimeout = Duration.ofSeconds( server.getReadTimeoutSeconds() != null ? server.getReadTimeoutSeconds() : 60); - return McpClient.sync(transport) - .requestTimeout(requestTimeout) - .build(); + var spec = McpClient.sync(transport).requestTimeout(requestTimeout); + if (managed) { + // Server-pushed tool-list changes (tools/list_changed) refresh the + // agent graphs without any polling — the SDK invokes this consumer + // on the client's inbound notification thread. + Long serverId = server.getId(); + spec.toolsChangeConsumer(tools -> + eventPublisher.publishEvent(new McpServerChangedEvent("mcp-tools-changed:" + serverId))); + } + return spec.build(); } - private StdioClientTransport buildStdioTransport(McpServerEntity server) { + private StdioClientTransport buildStdioTransport(McpServerEntity server, boolean managed) { String command = normalizeStdioCommand(server.getCommand()); ServerParameters.Builder builder = ServerParameters.builder(command); @@ -341,11 +408,19 @@ public class McpClientManager { builder.env(expandedEnv); } - StdioClientTransport transport = new CwdAwareStdioClientTransport( + CwdAwareStdioClientTransport transport = new CwdAwareStdioClientTransport( builder.build(), McpJsonMapper.createDefault(), expandEnvVars(server.getCwd())); transport.setStdErrorHandler(line -> log.info("MCP stdio stderr [{}]: {}", server.getName(), line)); + if (managed) { + // stdio is the one transport the MCP SDK cannot self-heal: a dead + // subprocess is only recoverable by respawning it, which lazy + // re-initialization never does. Request a reconnect on unexpected exit. + Long serverId = server.getId(); + transport.setOnUnexpectedExit(() -> + eventPublisher.publishEvent(new McpConnectionLostEvent(serverId, "stdio-process-exited"))); + } return transport; } diff --git a/mateclaw-server/src/main/java/vip/mate/tool/mcp/service/McpServerService.java b/mateclaw-server/src/main/java/vip/mate/tool/mcp/service/McpServerService.java index 382be817..e79bf753 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/mcp/service/McpServerService.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/mcp/service/McpServerService.java @@ -8,8 +8,10 @@ import jakarta.annotation.PreDestroy; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.context.ApplicationEventPublisher; +import org.springframework.context.event.EventListener; import org.springframework.stereotype.Service; import vip.mate.exception.MateClawException; +import vip.mate.tool.mcp.event.McpConnectionLostEvent; import vip.mate.tool.mcp.event.McpServerChangedEvent; import vip.mate.tool.mcp.model.McpServerEntity; import vip.mate.tool.mcp.repository.McpServerMapper; @@ -20,6 +22,7 @@ import java.time.LocalDateTime; import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.regex.Pattern; @@ -61,6 +64,50 @@ public class McpServerService { connectExecutor.shutdownNow(); } + /** + * Debounce window for runtime-triggered reconnects (issue #317). A live + * tool call and an agent rebuild can both detect the same dead connection + * within milliseconds, and a crash-looping server would otherwise respawn + * on every {@code listTools()} miss. We collapse repeated reconnect requests + * for the same server inside this window. + */ + private static final long RECONNECT_DEBOUNCE_MS = 10_000; + + /** serverId -> last runtime reconnect attempt epoch millis. */ + private final ConcurrentHashMap lastRuntimeReconnectAt = new ConcurrentHashMap<>(); + + /** + * Heal a connection that died at runtime: a stale {@code listTools()} or a + * stdio subprocess that exited on its own (e.g. the user restarted the MCP + * service). Reloads the server config and reconnects asynchronously, + * debounced so a flapping server can't saturate the reconnect pool. The + * reconnect publishes {@link McpServerChangedEvent} on success, which clears + * the agent cache so the next turn rebuilds against the live tools. + */ + @EventListener + public void onConnectionLost(McpConnectionLostEvent event) { + Long serverId = event.serverId(); + if (serverId == null) { + return; + } + McpServerEntity server = mcpServerMapper.selectById(serverId); + if (server == null || !Boolean.TRUE.equals(server.getEnabled())) { + // Removed or disabled in the meantime — nothing to heal. + return; + } + + long now = System.currentTimeMillis(); + Long previous = lastRuntimeReconnectAt.get(serverId); + if (previous != null && now - previous < RECONNECT_DEBOUNCE_MS) { + log.debug("Skipping MCP reconnect for '{}' ({}): within debounce window", server.getName(), event.reason()); + return; + } + lastRuntimeReconnectAt.put(serverId, now); + + log.warn("MCP server '{}' connection lost ({}); reconnecting", server.getName(), event.reason()); + reconnectAsync(server); + } + /** Publish a connection-state change so AgentService rebuilds its agent cache (issue #289). */ private void publishChanged(String reason) { try { diff --git a/mateclaw-server/src/test/java/vip/mate/tool/mcp/runtime/McpClientManagerSnapshotTest.java b/mateclaw-server/src/test/java/vip/mate/tool/mcp/runtime/McpClientManagerSnapshotTest.java new file mode 100644 index 00000000..e62a57fd --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/tool/mcp/runtime/McpClientManagerSnapshotTest.java @@ -0,0 +1,98 @@ +package vip.mate.tool.mcp.runtime; + +import io.modelcontextprotocol.client.McpSyncClient; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.model.ToolContext; +import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.definition.DefaultToolDefinition; +import org.springframework.ai.tool.definition.ToolDefinition; +import org.springframework.ai.tool.metadata.ToolMetadata; +import org.springframework.context.ApplicationEventPublisher; +import vip.mate.tool.mcp.event.McpConnectionLostEvent; + +import java.lang.reflect.Field; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Regression guard for issue #317: a stale {@code listTools()} (the upstream MCP + * server restarted while we held the connection) must NOT drop the whole server + * — that is what made the agent fall back to non-MCP tools. Instead the manager + * serves the last known-good callbacks and asks the service layer to reconnect. + */ +class McpClientManagerSnapshotTest { + + @Test + @DisplayName("stale listTools serves cached snapshot and requests a reconnect") + @SuppressWarnings("unchecked") + void staleListToolsServesSnapshotAndRequestsReconnect() throws Exception { + ApplicationEventPublisher publisher = mock(ApplicationEventPublisher.class); + McpClientManager manager = new McpClientManager(publisher); + + // A client whose connection went stale: every listTools() throws. + McpSyncClient deadClient = mock(McpSyncClient.class); + when(deadClient.listTools()).thenThrow(new RuntimeException("session not found")); + + long serverId = 77L; + List snapshot = List.of(new PrefixedNameToolCallback("mcp_77_search_abc123", stub("search"))); + + // Pre-seed the private state as if a previous successful collection ran. + ((Map) field(manager, "clients")).put(serverId, deadClient); + ((Map>) field(manager, "lastGoodCallbacks")).put(serverId, snapshot); + + List result = manager.getAllToolCallbacks(); + + // The cached snapshot is served verbatim — the server is NOT dropped. + assertEquals(1, result.size()); + assertSame(snapshot.get(0), result.get(0)); + + // A reconnect was requested for exactly this server. + verify(publisher, times(1)).publishEvent(any(McpConnectionLostEvent.class)); + } + + private static Object field(McpClientManager manager, String name) throws Exception { + Field f = McpClientManager.class.getDeclaredField(name); + f.setAccessible(true); + Object value = f.get(manager); + if (value == null) { + ConcurrentHashMap created = new ConcurrentHashMap<>(); + f.set(manager, created); + return created; + } + return value; + } + + private static ToolCallback stub(String rawName) { + return new ToolCallback() { + @Override + public ToolDefinition getToolDefinition() { + return DefaultToolDefinition.builder().name(rawName).description("").inputSchema("{}").build(); + } + + @Override + public ToolMetadata getToolMetadata() { + return ToolMetadata.builder().build(); + } + + @Override + public String call(String toolInput) { + return ""; + } + + @Override + public String call(String toolInput, ToolContext toolContext) { + return ""; + } + }; + } +}