mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 11:13:43 +08:00
fix(mcp): auto-heal stale MCP connections after server restart (#317)
This commit is contained in:
matevip
parent
ec4bb9061c
commit
4a3d4cf568
@ -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.
|
||||
*
|
||||
* <p>Unlike {@link McpServerChangedEvent} (which announces a state change that
|
||||
* already happened), this is a <em>request to heal</em>: {@code McpServerService}
|
||||
* listens, reloads the server config, and triggers an asynchronous reconnect
|
||||
* (debounced so a crash-looping server can't spin the reconnect executor).
|
||||
*
|
||||
* <p>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) {
|
||||
}
|
||||
@ -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<Thread> 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(
|
||||
|
||||
@ -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<Long, List<McpSchema.Tool>> 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<Long, List<ToolCallback>> lastGoodCallbacks = new ConcurrentHashMap<>();
|
||||
|
||||
private final ApplicationEventPublisher eventPublisher;
|
||||
|
||||
public McpClientManager(ApplicationEventPublisher eventPublisher) {
|
||||
this.eventPublisher = eventPublisher;
|
||||
}
|
||||
|
||||
/** serverId -> connection result info */
|
||||
private final ConcurrentHashMap<Long, ConnectionResult> 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<McpSchema.Tool> 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<ToolCallback> 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<ToolCallback> out, long serverId) {
|
||||
List<ToolCallback> 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;
|
||||
}
|
||||
|
||||
|
||||
@ -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<Long, Long> 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 {
|
||||
|
||||
@ -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<ToolCallback> snapshot = List.of(new PrefixedNameToolCallback("mcp_77_search_abc123", stub("search")));
|
||||
|
||||
// Pre-seed the private state as if a previous successful collection ran.
|
||||
((Map<Long, McpSyncClient>) field(manager, "clients")).put(serverId, deadClient);
|
||||
((Map<Long, List<ToolCallback>>) field(manager, "lastGoodCallbacks")).put(serverId, snapshot);
|
||||
|
||||
List<ToolCallback> 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<Object, Object> 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 "";
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user