mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
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:
+14
-1
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user