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:
parent
7577414761
commit
f1b798434d
1 changed files with 14 additions and 1 deletions
|
|
@ -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),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue