From 54e8d7ffde31ccc00ea59cec7b2eaec18b01f137 Mon Sep 17 00:00:00 2001 From: yeliheng-010 <81519560+yeliheng-010@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:33:55 +0800 Subject: [PATCH] =?UTF-8?q?fix(mcp):=20=E9=9A=94=E7=A6=BB=E6=85=A2?= =?UTF-8?q?=E8=BF=9E=E6=8E=A5=E6=B8=85=E7=90=86=E5=AF=B9=E5=85=B6=E4=BB=96?= =?UTF-8?q?=E6=9C=8D=E5=8A=A1=E7=9A=84=E9=98=BB=E5=A1=9E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/mcp/manager.go | 36 +++- internal/mcp/manager_disconnect_test.go | 220 ++++++++++++++++++++++++ 2 files changed, 248 insertions(+), 8 deletions(-) create mode 100644 internal/mcp/manager_disconnect_test.go diff --git a/internal/mcp/manager.go b/internal/mcp/manager.go index d0af134ea..74c2955b9 100644 --- a/internal/mcp/manager.go +++ b/internal/mcp/manager.go @@ -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 // ("\x00"). 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 diff --git a/internal/mcp/manager_disconnect_test.go b/internal/mcp/manager_disconnect_test.go new file mode 100644 index 000000000..e0da50222 --- /dev/null +++ b/internal/mcp/manager_disconnect_test.go @@ -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") +}