fix(mcp): 隔离慢连接清理对其他服务的阻塞

2026-10-13 08:50:22 +08:00 · 2026-10-10 20:12:58 +08:00
parent 98e9dc0786
commit 54e8d7ffde
2 changed files with 248 additions and 8 deletions
@@ -30,6 +30,7 @@ type pendingMCPConnection struct {
 	client  MCPClient
 	err     error
 	version time.Time
+	retired MCPClient // closed by the connection worker before opening its replacement
 }

 type managedMCPClient struct {
@@ -111,6 +112,7 @@ func (m *MCPManager) GetOrCreateClient(ctx context.Context, service *types.MCPSe
 	}
 	key := cacheKey(service, principal)

+	var retired MCPClient
 	m.clientsMu.Lock()
 	if err := m.ctx.Err(); err != nil {
 		m.clientsMu.Unlock()
@@ -122,7 +124,7 @@ func (m *MCPManager) GetOrCreateClient(ctx context.Context, service *types.MCPSe
 			m.clientsMu.Unlock()
 			return client, nil
 		}
-		_ = client.Disconnect()
+		retired = client
 		delete(m.clients, key)
 	}
 	pending := m.connecting[key]
@@ -133,7 +135,9 @@ func (m *MCPManager) GetOrCreateClient(ctx context.Context, service *types.MCPSe
 	}
 	if pending == nil {
 		lifeCtx, cancel := context.WithCancel(m.ctx)
-		pending = &pendingMCPConnection{done: make(chan struct{}), cancel: cancel, version: service.UpdatedAt}
+		pending = &pendingMCPConnection{
+			done: make(chan struct{}), cancel: cancel, version: service.UpdatedAt, retired: retired,
+		}
 		m.connecting[key] = pending
 		config := &ClientConfig{Service: service, TenantID: tenantID, Principal: principal, OAuthRepo: m.oauthRepo}
 		go m.connectClient(lifeCtx, key, config, pending)
@@ -150,7 +154,16 @@ func (m *MCPManager) GetOrCreateClient(ctx context.Context, service *types.MCPSe
 func (m *MCPManager) connectClient(
 	ctx context.Context, key string, config *ClientConfig, pending *pendingMCPConnection,
 ) {
-	client, err := NewMCPClient(config)
+	// Closing a session can issue a remote DELETE. Keep it off the manager
+	// lock and the caller's goroutine, while preserving replacement order.
+	if pending.retired != nil {
+		_ = pending.retired.Disconnect()
+	}
+	var client MCPClient
+	err := ctx.Err()
+	if err == nil {
+		client, err = NewMCPClient(config)
+	}
 	if err == nil {
 		// SSE needs the connection lifetime, not the requesting turn's deadline.
 		err = client.Connect(ctx)
@@ -171,15 +184,15 @@ func (m *MCPManager) connectClient(
 	} else {
 		pending.err = err
 		pending.cancel()
-		if client != nil {
-			_ = client.Disconnect()
-		}
 	}
 	if m.connecting[key] == pending {
 		delete(m.connecting, key)
 	}
 	close(pending.done)
 	m.clientsMu.Unlock()
+	if err != nil && client != nil {
+		_ = client.Disconnect()
+	}
 }

 // initializeClient handles the shared initialization flow with timeout enforcement.
@@ -222,7 +235,6 @@ func (m *MCPManager) GetClient(serviceID string) (MCPClient, bool) {
 // the service ID).
 func (m *MCPManager) CloseClient(serviceID string) error {
 	m.clientsMu.Lock()
-	defer m.clientsMu.Unlock()
 	for key, pending := range m.connecting {
 		if key == serviceID || strings.HasPrefix(key, serviceID+"\x00") {
 			pending.cancel()
@@ -230,16 +242,24 @@ func (m *MCPManager) CloseClient(serviceID string) error {
 		}
 	}

+	retired := make(map[string]MCPClient)
 	for key, client := range m.clients {
 		// Match the plain service-ID key as well as per-principal OAuth keys
 		// ("<serviceID>\x00<principal>").
 		if key != serviceID && !strings.HasPrefix(key, serviceID+"\x00") {
 			continue
 		}
+		retired[key] = client
+		delete(m.clients, key)
+	}
+	m.clientsMu.Unlock()
+
+	// Remove only the captured clients. A replacement may be installed while
+	// a remote server is still acknowledging the old session's closure.
+	for key, client := range retired {
 		if err := client.Disconnect(); err != nil {
 			logger.GetLogger(m.ctx).Errorf("Failed to disconnect MCP client %s: %v", key, err)
 		}
-		delete(m.clients, key)
 		logger.GetLogger(m.ctx).Infof("MCP client closed: %s", key)
 	}
 	return nil
@@ -0,0 +1,220 @@
+package mcp
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/utils"
+	"github.com/stretchr/testify/require"
+)
+
+type slowDisconnectServer struct {
+	started chan struct{}
+	release chan struct{}
+	once    sync.Once
+	starts  atomic.Int32
+}
+
+func (s *slowDisconnectServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.Method == http.MethodDelete {
+		if r.URL.Path == "/slow" {
+			s.once.Do(func() { close(s.started) })
+			select {
+			case <-s.release:
+			case <-r.Context().Done():
+			}
+		}
+		w.WriteHeader(http.StatusNoContent)
+		return
+	}
+	if r.Method != http.MethodPost {
+		w.WriteHeader(http.StatusMethodNotAllowed)
+		return
+	}
+	var request struct {
+		ID     json.RawMessage `json:"id"`
+		Method string          `json:"method"`
+	}
+	if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+	if request.ID == nil {
+		w.WriteHeader(http.StatusAccepted)
+		return
+	}
+	w.Header().Set("Content-Type", "application/json")
+	response := map[string]any{"jsonrpc": "2.0", "id": request.ID}
+	switch request.Method {
+	case "server/discover":
+		response["error"] = map[string]any{"code": -32601, "message": "Method not found"}
+	case "initialize":
+		if r.URL.Path == "/slow" {
+			s.starts.Add(1)
+		}
+		w.Header().Set("Mcp-Session-Id", "session"+r.URL.Path)
+		response["result"] = map[string]any{
+			"protocolVersion": "2025-06-18", "capabilities": map[string]any{"tools": map[string]any{}},
+			"serverInfo": map[string]any{"name": "disconnect-test", "version": "1"},
+		}
+	default:
+		response["result"] = map[string]any{"tools": []any{}}
+	}
+	_ = json.NewEncoder(w).Encode(response)
+}
+
+type disconnectFixture struct {
+	server  *slowDisconnectServer
+	manager *MCPManager
+	slow    *types.MCPService
+	fast    *types.MCPService
+	unblock func()
+}
+
+func newDisconnectFixture(t *testing.T) *disconnectFixture {
+	t.Helper()
+	utils.SetSSRFWhitelistFromRaw("127.0.0.1")
+	t.Cleanup(utils.ResetSSRFWhitelistForTest)
+	s := &slowDisconnectServer{started: make(chan struct{}), release: make(chan struct{})}
+	server := httptest.NewServer(s)
+	manager := NewMCPManager(nil)
+	unblock := sync.OnceFunc(func() { close(s.release) })
+	t.Cleanup(func() {
+		unblock()
+		manager.Shutdown()
+		server.Close()
+	})
+	slowURL, fastURL := server.URL+"/slow", server.URL+"/fast"
+	f := &disconnectFixture{
+		server: s, manager: manager, unblock: unblock,
+		slow: &types.MCPService{
+			ID: "slow", Enabled: true, URL: &slowURL,
+			TransportType: types.MCPTransportHTTPStreamable,
+		},
+		fast: &types.MCPService{
+			ID: "fast", Enabled: true, URL: &fastURL,
+			TransportType: types.MCPTransportHTTPStreamable,
+		},
+	}
+	_, err := manager.GetOrCreateClient(t.Context(), f.slow)
+	require.NoError(t, err)
+	_, err = manager.GetOrCreateClient(t.Context(), f.fast)
+	require.NoError(t, err)
+	return f
+}
+
+func TestManagerSlowDisconnectDoesNotBlockOtherServices(t *testing.T) {
+	for _, operation := range []string{"close", "replace"} {
+		t.Run(operation, func(t *testing.T) { testSlowDisconnectIsolation(t, operation) })
+	}
+}
+
+func testSlowDisconnectIsolation(t *testing.T, operation string) {
+	t.Helper()
+	f := newDisconnectFixture(t)
+	defer f.unblock()
+	done := make(chan error, 1)
+	go func() {
+		if operation == "close" {
+			done <- f.manager.CloseClient(f.slow.ID)
+			return
+		}
+		snapshot := *f.slow
+		snapshot.UpdatedAt = time.Now()
+		_, err := f.manager.GetOrCreateClient(t.Context(), &snapshot)
+		done <- err
+	}()
+	select {
+	case <-f.server.started:
+	case <-time.After(2 * time.Second):
+		t.Fatal("disconnect did not reach the HTTP server")
+	}
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
+	defer cancel()
+	fastDone := make(chan error, 1)
+	go func() {
+		client, err := f.manager.GetOrCreateClient(ctx, f.fast)
+		if err == nil {
+			_, err = client.ListTools(ctx)
+		}
+		fastDone <- err
+	}()
+	select {
+	case err := <-fastDone:
+		require.NoError(t, err)
+	case <-ctx.Done():
+		t.Fatal("an unrelated healthy MCP service is blocked by remote disconnect I/O")
+	}
+	f.unblock()
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("disconnect did not complete after the server recovered")
+	}
+}
+
+func TestManagerReplacementWaiterCanCancelDuringDisconnect(t *testing.T) {
+	f := newDisconnectFixture(t)
+	defer f.unblock()
+	snapshot := *f.slow
+	snapshot.UpdatedAt = time.Now()
+	ctx, cancel := context.WithCancel(t.Context())
+	defer cancel()
+	done := make(chan error, 1)
+	go func() { _, err := f.manager.GetOrCreateClient(ctx, &snapshot); done <- err }()
+	select {
+	case <-f.server.started:
+	case <-time.After(2 * time.Second):
+		t.Fatal("replacement did not start closing the old connection")
+	}
+	cancel()
+	select {
+	case err := <-done:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("caller cancellation is blocked by remote disconnect I/O")
+	}
+	require.EqualValues(t, 1, f.server.starts.Load(), "replacement must wait for old session cleanup")
+	f.unblock()
+	readyCtx, readyCancel := context.WithTimeout(t.Context(), 2*time.Second)
+	defer readyCancel()
+	client, err := f.manager.GetOrCreateClient(readyCtx, &snapshot)
+	require.NoError(t, err, "one cancelled waiter must not cancel the shared replacement")
+	require.True(t, client.IsConnected())
+	require.EqualValues(t, 2, f.server.starts.Load())
+}
+
+func TestManagerCloseDoesNotRemoveConcurrentReplacement(t *testing.T) {
+	f := newDisconnectFixture(t)
+	defer f.unblock()
+	done := make(chan error, 1)
+	go func() { done <- f.manager.CloseClient(f.slow.ID) }()
+	select {
+	case <-f.server.started:
+	case <-time.After(2 * time.Second):
+		t.Fatal("close did not reach the HTTP server")
+	}
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
+	defer cancel()
+	// The old client is already detached, so a later request can reconnect.
+	client, err := f.manager.GetOrCreateClient(ctx, f.slow)
+	require.NoError(t, err)
+	f.unblock()
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("close did not finish after the server recovered")
+	}
+	current, ok := f.manager.GetClient(f.slow.ID)
+	require.True(t, ok)
+	require.Same(t, client, current, "old cleanup must not remove the newly cached connection")
+}