fix(mcp): prevent race condition between CallTool and Close

Add a closed flag to the Manager struct to prevent CallTool from
accessing server connections after Close has been called. The flag
is checked within the RLock in CallTool to ensure thread-safety.

Previously, CallTool could obtain a server reference using RLock,
then that reference could be closed by Close() running concurrently,
leading to use-after-close errors. Now:

1. CallTool checks the closed flag before accessing servers
2. Close sets the closed flag before closing connections
3. CallTool directly accesses m.servers within the lock instead
   of using GetServer() to avoid releasing the lock prematurely

This ensures CallTool will not use a server connection that is
being closed or has been closed.

Addresses Copilot code review feedback.
This commit is contained in:
yuchou87
2026-02-19 19:47:05 +08:00
parent 7577414761
commit f1b798434d
+14 -1
View File
@@ -103,6 +103,7 @@ type ServerConnection struct {
type Manager struct {
servers map[string]*ServerConnection
mu sync.RWMutex
closed bool
}
// NewManager creates a new MCP manager
@@ -403,7 +404,14 @@ func (m *Manager) GetServer(name string) (*ServerConnection, bool) {
// CallTool calls a tool on a specific server
func (m *Manager) CallTool(ctx context.Context, serverName, toolName string, arguments map[string]interface{}) (*mcp.CallToolResult, error) {
conn, ok := m.GetServer(serverName)
m.mu.RLock()
if m.closed {
m.mu.RUnlock()
return nil, fmt.Errorf("manager is closed")
}
conn, ok := m.servers[serverName]
m.mu.RUnlock()
if !ok {
return nil, fmt.Errorf("server %s not found", serverName)
}
@@ -426,6 +434,11 @@ func (m *Manager) Close() error {
m.mu.Lock()
defer m.mu.Unlock()
if m.closed {
return nil
}
m.closed = true
logger.InfoCF("mcp", "Closing all MCP server connections",
map[string]interface{}{
"count": len(m.servers),