From 94de941b445f212c67cdc6bf454de64bc32652b3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 4 Aug 2026 23:24:51 +0000 Subject: [PATCH 1/2] Persist plugin MCP registrations and rehydrate on restart Co-authored-by: nick.misasi --- api/api.go | 1 + api/api_admin.go | 2 +- api/api_admin_test.go | 68 ++++++---- api/api_test.go | 12 +- external/pluginmcp/README.md | 11 +- mcp/client_manager.go | 108 +++++++++++++++- mcp/client_manager_test.go | 234 ++++++++++++++++++++++++++++++++--- 7 files changed, 385 insertions(+), 51 deletions(-) diff --git a/api/api.go b/api/api.go index 484e5bcec..f0968a2b8 100644 --- a/api/api.go +++ b/api/api.go @@ -69,6 +69,7 @@ type MCPClientManager interface { GetConfig() mcp.Config RegisterPluginServer(cfg mcp.PluginServerConfig) + UpdatePluginServer(cfg mcp.PluginServerConfig) UnregisterPluginServer(pluginID string) ListPluginServers() []mcp.PluginServerConfig GetPluginServer(pluginID string) (mcp.PluginServerConfig, bool) diff --git a/api/api_admin.go b/api/api_admin.go index 0b0a0528a..7a7d2b448 100644 --- a/api/api_admin.go +++ b/api/api_admin.go @@ -509,7 +509,7 @@ func (a *API) handleUpdatePluginServer(c *gin.Context) { return } - a.mcpClientManager.RegisterPluginServer(updated) + a.mcpClientManager.UpdatePluginServer(updated) a.configUpdater.Update(cfg) // Rebuild when either old or new state was external so removed tools diff --git a/api/api_admin_test.go b/api/api_admin_test.go index e1ff629c2..1f80fb7d1 100644 --- a/api/api_admin_test.go +++ b/api/api_admin_test.go @@ -547,11 +547,12 @@ func TestHandleUpdatePluginServer(t *testing.T) { body string hasAdminPerm bool expectStatus int - expectRegisterCalls int + expectRegistryCalls int expectEnabledAfter bool expectExposeAfter bool expectToolConfigsAfter []mcp.ToolConfig expectRebuildCalls int + orphanPluginIDs map[string]bool }{ { name: "happy path: flips Enabled true->false", @@ -562,7 +563,7 @@ func TestHandleUpdatePluginServer(t *testing.T) { body: `{"enabled": false}`, hasAdminPerm: true, expectStatus: http.StatusOK, - expectRegisterCalls: 1, + expectRegistryCalls: 1, expectEnabledAfter: false, expectExposeAfter: false, expectRebuildCalls: 0, @@ -577,7 +578,7 @@ func TestHandleUpdatePluginServer(t *testing.T) { body: `{"enabled": false}`, hasAdminPerm: true, expectStatus: http.StatusOK, - expectRegisterCalls: 1, + expectRegistryCalls: 1, expectEnabledAfter: false, expectExposeAfter: true, expectRebuildCalls: 1, @@ -592,7 +593,7 @@ func TestHandleUpdatePluginServer(t *testing.T) { body: `{"expose_external": true}`, hasAdminPerm: true, expectStatus: http.StatusOK, - expectRegisterCalls: 1, + expectRegistryCalls: 1, expectEnabledAfter: true, expectExposeAfter: false, expectRebuildCalls: 0, @@ -607,11 +608,24 @@ func TestHandleUpdatePluginServer(t *testing.T) { body: `{}`, hasAdminPerm: true, expectStatus: http.StatusOK, - expectRegisterCalls: 1, + expectRegistryCalls: 1, expectEnabledAfter: true, expectExposeAfter: true, expectRebuildCalls: 1, }, + { + name: "admin update keeps config-only orphan unregistered", + pluginID: "com.mattermost.demo", + preRegistered: []mcp.PluginServerConfig{{ + PluginID: "com.mattermost.demo", Name: "Demo", Path: "/mcp", Enabled: true, + }}, + orphanPluginIDs: map[string]bool{"com.mattermost.demo": true}, + body: `{"enabled": false}`, + hasAdminPerm: true, + expectStatus: http.StatusOK, + expectRegistryCalls: 1, + expectEnabledAfter: false, + }, { name: "404 when pluginID not registered", pluginID: "com.missing", @@ -638,7 +652,7 @@ func TestHandleUpdatePluginServer(t *testing.T) { body: `{"enabled": false}`, hasAdminPerm: false, expectStatus: http.StatusForbidden, - expectRegisterCalls: 0, + expectRegistryCalls: 0, }, { name: "tool_configs partial PUT sets policy, preserves enabled", @@ -650,7 +664,7 @@ func TestHandleUpdatePluginServer(t *testing.T) { body: `{"tool_configs": [{"name": "echo", "policy": "ask", "enabled": false}]}`, hasAdminPerm: true, expectStatus: http.StatusOK, - expectRegisterCalls: 1, + expectRegistryCalls: 1, expectEnabledAfter: true, expectExposeAfter: false, expectToolConfigsAfter: []mcp.ToolConfig{ @@ -670,7 +684,7 @@ func TestHandleUpdatePluginServer(t *testing.T) { body: `{"tool_configs": []}`, hasAdminPerm: true, expectStatus: http.StatusOK, - expectRegisterCalls: 1, + expectRegistryCalls: 1, expectEnabledAfter: true, expectExposeAfter: false, expectToolConfigsAfter: []mcp.ToolConfig{}, @@ -689,7 +703,7 @@ func TestHandleUpdatePluginServer(t *testing.T) { body: `{"enabled": false}`, hasAdminPerm: true, expectStatus: http.StatusOK, - expectRegisterCalls: 1, + expectRegistryCalls: 1, expectEnabledAfter: false, expectExposeAfter: false, expectToolConfigsAfter: []mcp.ToolConfig{ @@ -709,6 +723,7 @@ func TestHandleUpdatePluginServer(t *testing.T) { mgr := api.mcpClientManager.(*mockMCPClientManager) mgr.pluginServers = tt.preRegistered + mgr.orphanPluginIDs = tt.orphanPluginIDs // Seed a baseline persisted config so the handler can clone it // instead of treating the store's nil as a 500. @@ -727,15 +742,19 @@ func TestHandleUpdatePluginServer(t *testing.T) { resp := recorder.Result() require.Equal(t, tt.expectStatus, resp.StatusCode) - require.Len(t, mgr.registerCalls, tt.expectRegisterCalls) + require.Empty(t, mgr.registerCalls) + require.Len(t, mgr.updateCalls, tt.expectRegistryCalls) if tt.expectStatus == http.StatusOK { - require.Equal(t, tt.expectEnabledAfter, mgr.registerCalls[0].Enabled) - require.Equal(t, tt.expectExposeAfter, mgr.registerCalls[0].ExposeExternal) - require.Equal(t, "Demo", mgr.registerCalls[0].Name) - require.Equal(t, "/mcp", mgr.registerCalls[0].Path) - require.Equal(t, "com.mattermost.demo", mgr.registerCalls[0].PluginID) + require.Equal(t, tt.expectEnabledAfter, mgr.updateCalls[0].Enabled) + require.Equal(t, tt.expectExposeAfter, mgr.updateCalls[0].ExposeExternal) + require.Equal(t, "Demo", mgr.updateCalls[0].Name) + require.Equal(t, "/mcp", mgr.updateCalls[0].Path) + require.Equal(t, "com.mattermost.demo", mgr.updateCalls[0].PluginID) if tt.expectToolConfigsAfter != nil { - require.Equal(t, tt.expectToolConfigsAfter, mgr.registerCalls[0].ToolConfigs, "ToolConfigs assertion") + require.Equal(t, tt.expectToolConfigsAfter, mgr.updateCalls[0].ToolConfigs, "ToolConfigs assertion") + } + if tt.orphanPluginIDs[tt.pluginID] { + require.False(t, mgr.IsPluginRegistered(tt.pluginID)) } } require.Equal(t, tt.expectRebuildCalls, spy.callCount) @@ -757,7 +776,7 @@ func TestHandleUpdatePluginServer_PersistsToConfig(t *testing.T) { expectSaveCalls int expectUpdateCalls int expectPublishCalls int - expectRegisterCalls int + expectRegistryCalls int expectUnregisterCalls int assertPersistedState func(t *testing.T, savedCfg *config.Config) }{ @@ -774,7 +793,7 @@ func TestHandleUpdatePluginServer_PersistsToConfig(t *testing.T) { expectSaveCalls: 1, expectUpdateCalls: 1, expectPublishCalls: 1, - expectRegisterCalls: 1, + expectRegistryCalls: 1, expectUnregisterCalls: 0, assertPersistedState: func(t *testing.T, savedCfg *config.Config) { require.Len(t, savedCfg.MCP.PluginServers, 1) @@ -819,7 +838,7 @@ func TestHandleUpdatePluginServer_PersistsToConfig(t *testing.T) { expectSaveCalls: 1, expectUpdateCalls: 1, expectPublishCalls: 1, - expectRegisterCalls: 1, + expectRegistryCalls: 1, expectUnregisterCalls: 0, assertPersistedState: func(t *testing.T, savedCfg *config.Config) { require.Len(t, savedCfg.MCP.PluginServers, 2, @@ -855,7 +874,7 @@ func TestHandleUpdatePluginServer_PersistsToConfig(t *testing.T) { expectSaveCalls: 0, expectUpdateCalls: 0, expectPublishCalls: 0, - expectRegisterCalls: 0, + expectRegistryCalls: 0, expectUnregisterCalls: 0, }, { @@ -869,7 +888,7 @@ func TestHandleUpdatePluginServer_PersistsToConfig(t *testing.T) { expectSaveCalls: 1, expectUpdateCalls: 0, expectPublishCalls: 0, - expectRegisterCalls: 0, + expectRegistryCalls: 0, expectUnregisterCalls: 0, }, { @@ -883,7 +902,7 @@ func TestHandleUpdatePluginServer_PersistsToConfig(t *testing.T) { expectSaveCalls: 1, expectUpdateCalls: 0, expectPublishCalls: 1, - expectRegisterCalls: 0, + expectRegistryCalls: 0, expectUnregisterCalls: 0, }, { @@ -900,7 +919,7 @@ func TestHandleUpdatePluginServer_PersistsToConfig(t *testing.T) { expectSaveCalls: 0, expectUpdateCalls: 0, expectPublishCalls: 0, - expectRegisterCalls: 0, + expectRegistryCalls: 0, expectUnregisterCalls: 0, }, } @@ -954,7 +973,8 @@ func TestHandleUpdatePluginServer_PersistsToConfig(t *testing.T) { require.Equal(t, tt.expectUpdateCalls, stores.configUpdater.callCount) require.Equal(t, tt.expectPublishCalls, stores.clusterNotifier.callCount) - require.Len(t, mgr.registerCalls, tt.expectRegisterCalls, "live plugin registry must not be mutated on failure paths") + require.Empty(t, mgr.registerCalls) + require.Len(t, mgr.updateCalls, tt.expectRegistryCalls, "live plugin registry must not be mutated on failure paths") require.Len(t, mgr.unregisterCalls, tt.expectUnregisterCalls, "live plugin registry must not be mutated on failure paths") if tt.assertPersistedState != nil { diff --git a/api/api_test.go b/api/api_test.go index ebaf0a65e..e774398ea 100644 --- a/api/api_test.go +++ b/api/api_test.go @@ -126,6 +126,7 @@ type mockMCPClientManager struct { ensureSessionErr error registerCalls []mcp.PluginServerConfig + updateCalls []mcp.PluginServerConfig unregisterCalls []string pluginServers []mcp.PluginServerConfig // orphanPluginIDs simulates entries present in pluginServers but with @@ -206,7 +207,16 @@ func (m *mockMCPClientManager) GetConfig() mcp.Config { func (m *mockMCPClientManager) RegisterPluginServer(cfg mcp.PluginServerConfig) { m.registerCalls = append(m.registerCalls, cfg) - // Mirror real ClientManager: same PluginID replaces existing entry. + delete(m.orphanPluginIDs, cfg.PluginID) + m.storePluginServer(cfg) +} + +func (m *mockMCPClientManager) UpdatePluginServer(cfg mcp.PluginServerConfig) { + m.updateCalls = append(m.updateCalls, cfg) + m.storePluginServer(cfg) +} + +func (m *mockMCPClientManager) storePluginServer(cfg mcp.PluginServerConfig) { for i, existing := range m.pluginServers { if existing.PluginID == cfg.PluginID { m.pluginServers[i] = cfg diff --git a/external/pluginmcp/README.md b/external/pluginmcp/README.md index 39e41b7da..02c31b9ca 100644 --- a/external/pluginmcp/README.md +++ b/external/pluginmcp/README.md @@ -213,12 +213,11 @@ External callers can't inject one. Don't add a second auth gate in your outer `ServeHTTP`, and don't read `X-Mattermost-UserID` directly from headers in handlers; always go through `GetUserID`. -**Registration is one-shot per `OnActivate`.** The retry goroutine -exits on success or after 15 attempts. If the Agents plugin restarts -later, the in-memory registration is lost. The Agents plugin restores -admin-persisted entries on its own restart, but a never-saved -registration only comes back when your plugin is re-activated. Permanent -non-retriable errors (4xx other than 404/429) log +**Registration retries are bounded.** The retry goroutine exits on +success or after 15 attempts. After a successful registration, the +Agents plugin persists it and restores it across its own restarts while +your plugin remains enabled. Permanent non-retriable errors (4xx other +than 404/429) log `registration with Agents plugin failed permanently` and stop. **Tool-count budget.** Each tool costs ~20-200 schema tokens in every diff --git a/mcp/client_manager.go b/mcp/client_manager.go index 3b0fb67f1..8043f6b25 100644 --- a/mcp/client_manager.go +++ b/mcp/client_manager.go @@ -5,6 +5,7 @@ package mcp import ( "context" + "encoding/json" "errors" "fmt" "net/http" @@ -15,9 +16,13 @@ import ( "github.com/mattermost/mattermost-plugin-agents/v2/llm" "github.com/mattermost/mattermost-plugin-agents/v2/mmapi" + "github.com/mattermost/mattermost/server/public/model" "github.com/mattermost/mattermost/server/public/pluginapi" ) +// Only PluginID, Name, Path, and ExposeExternal are authoritative here; Enabled and ToolConfigs come from admin config. +const pluginRegistrationsKVKey = "mcp_plugin_registrations_v1" + var ErrOAuthNotConfigured = errors.New("oauth not configured") func cacheableContext(ctx context.Context) context.Context { @@ -46,8 +51,8 @@ type ClientManager struct { // pluginServersMu must not be held across PluginHTTP round trips. pluginServersMu sync.RWMutex pluginServers map[string]PluginServerConfig // keyed by PluginID - // pluginRegistered marks entries with a live RegisterPluginServer call; - // orphan entries hydrated only from persisted config are absent. + // pluginRegistered marks entries backed by a source-plugin registration; + // config-only orphan entries are absent. pluginRegistered map[string]bool // sourcePluginAPI is the agents-plugin mmapi.Client; used by // PluginHTTPRoundTripper to dispatch to source plugins. @@ -67,6 +72,9 @@ func NewClientManager(config Config, log pluginapi.LogService, pluginAPI *plugin pluginRegistered: make(map[string]bool), sourcePluginAPI: sourcePluginAPI, } + manager.hydratePluginRegistrations() + // PluginMCPHandlers is constructed later and builds the external aggregate + // from this hydrated registry. manager.ReInit(config, embeddedServer) return manager } @@ -462,6 +470,16 @@ func (m *ClientManager) RegisterPluginServer(cfg PluginServerConfig) { defer m.pluginServersMu.Unlock() m.pluginServers[cfg.PluginID] = cfg m.pluginRegistered[cfg.PluginID] = true + m.mutatePersistedPluginRegistrations(func(registrations map[string]PluginServerConfig) { + registrations[cfg.PluginID] = cfg + }) +} + +// UpdatePluginServer applies admin-owned fields without changing registration state. +func (m *ClientManager) UpdatePluginServer(cfg PluginServerConfig) { + m.pluginServersMu.Lock() + defer m.pluginServersMu.Unlock() + m.pluginServers[cfg.PluginID] = cfg } func (m *ClientManager) UnregisterPluginServer(pluginID string) { @@ -469,6 +487,9 @@ func (m *ClientManager) UnregisterPluginServer(pluginID string) { defer m.pluginServersMu.Unlock() delete(m.pluginServers, pluginID) delete(m.pluginRegistered, pluginID) + m.mutatePersistedPluginRegistrations(func(registrations map[string]PluginServerConfig) { + delete(registrations, pluginID) + }) } func (m *ClientManager) ListPluginServers() []PluginServerConfig { @@ -489,15 +510,92 @@ func (m *ClientManager) GetPluginServer(pluginID string) (PluginServerConfig, bo return cfg, ok } -// IsPluginRegistered reports whether the source plugin currently has a live -// in-process registration. Returns false for entries hydrated only from -// persisted config. +// IsPluginRegistered reports whether an entry is backed by a source-plugin +// registration, including one restored from the KV store. func (m *ClientManager) IsPluginRegistered(pluginID string) bool { m.pluginServersMu.RLock() defer m.pluginServersMu.RUnlock() return m.pluginRegistered[pluginID] } +func (m *ClientManager) hydratePluginRegistrations() { + m.pluginServersMu.Lock() + defer m.pluginServersMu.Unlock() + + registrations, ok := m.loadPersistedPluginRegistrationsLocked() + if !ok { + return + } + + verifyPluginStates := len(registrations) > 0 + var pluginStates map[string]*model.PluginState + if verifyPluginStates { + serverConfig := m.pluginAPI.Configuration.GetConfig() + if serverConfig == nil { + m.log.Warn("Unable to verify plugin states while restoring MCP registrations; keeping all registrations") + verifyPluginStates = false + } else { + pluginStates = serverConfig.PluginSettings.PluginStates + } + } + + prunedPluginIDs := make([]string, 0) + restored := 0 + for pluginID, cfg := range registrations { + if verifyPluginStates { + state := pluginStates[pluginID] + if state == nil || !state.Enable { + prunedPluginIDs = append(prunedPluginIDs, pluginID) + continue + } + } + + m.pluginServers[pluginID] = cfg + m.pluginRegistered[pluginID] = true + restored++ + } + + if len(prunedPluginIDs) > 0 { + m.mutatePersistedPluginRegistrations(func(registrations map[string]PluginServerConfig) { + for _, pluginID := range prunedPluginIDs { + delete(registrations, pluginID) + } + }) + } + m.log.Debug("Restored plugin MCP registrations from KV store", "count", restored, "pruned", len(prunedPluginIDs)) +} + +func (m *ClientManager) mutatePersistedPluginRegistrations(update func(map[string]PluginServerConfig)) { + err := m.pluginAPI.KV.SetAtomicWithRetries(pluginRegistrationsKVKey, func(oldValue []byte) (any, error) { + registrations := make(map[string]PluginServerConfig) + if len(oldValue) > 0 { + if err := json.Unmarshal(oldValue, ®istrations); err != nil { + return nil, fmt.Errorf("unmarshal plugin MCP registrations: %w", err) + } + } + if registrations == nil { + registrations = make(map[string]PluginServerConfig) + } + update(registrations) + return registrations, nil + }) + if err != nil { + m.log.Error("Failed to persist plugin MCP registrations to KV store", "error", err) + } +} + +func (m *ClientManager) loadPersistedPluginRegistrationsLocked() (map[string]PluginServerConfig, bool) { + var registrations map[string]PluginServerConfig + if err := m.pluginAPI.KV.Get(pluginRegistrationsKVKey, ®istrations); err != nil { + m.log.Error("Failed to load plugin MCP registrations from KV store", "error", err) + return nil, false + } + if registrations == nil { + registrations = make(map[string]PluginServerConfig) + } + return registrations, true +} + // syncPluginServersFromConfig merges persisted admin-owned plugin-server fields // onto live plugin registrations. Callers must not hold pluginServersMu. func (m *ClientManager) syncPluginServersFromConfig(cfg Config) { diff --git a/mcp/client_manager_test.go b/mcp/client_manager_test.go index 7576ed5bb..99008106e 100644 --- a/mcp/client_manager_test.go +++ b/mcp/client_manager_test.go @@ -4,7 +4,9 @@ package mcp import ( + "bytes" "context" + "encoding/json" "net/http" "net/http/httptest" "strings" @@ -50,6 +52,75 @@ func (c *recordKVSetWithExpiryClient) KVSetWithExpiry(key string, value interfac return c.setErr } +type pluginRegistrationKVFixture struct { + mu sync.Mutex + data []byte + writes int +} + +func setupPluginRegistrationKV(t *testing.T, pluginTestAPI *plugintest.API, registrations map[string]PluginServerConfig) *pluginRegistrationKVFixture { + t.Helper() + + fixture := &pluginRegistrationKVFixture{} + if registrations != nil { + var err error + fixture.data, err = json.Marshal(registrations) + require.NoError(t, err) + } + + pluginTestAPI.On("KVGet", pluginRegistrationsKVKey). + Return(func(string) []byte { + fixture.mu.Lock() + defer fixture.mu.Unlock() + return append([]byte(nil), fixture.data...) + }, (*model.AppError)(nil)). + Maybe() + pluginTestAPI.On( + "KVSetWithOptions", + pluginRegistrationsKVKey, + mock.AnythingOfType("[]uint8"), + mock.AnythingOfType("model.PluginKVSetOptions"), + ).Return(func(_ string, value []byte, options model.PluginKVSetOptions) (bool, *model.AppError) { + fixture.mu.Lock() + defer fixture.mu.Unlock() + if options.Atomic && !bytes.Equal(fixture.data, options.OldValue) { + return false, nil + } + fixture.data = append([]byte(nil), value...) + fixture.writes++ + return true, nil + }).Maybe() + + return fixture +} + +func setupClientManagerTestAPI(t *testing.T, pluginTestAPI *plugintest.API) { + t.Helper() + setupTestLogger(pluginTestAPI) + setupPluginRegistrationKV(t, pluginTestAPI, nil) + pluginTestAPI.On("GetConfig").Return(&model.Config{}).Maybe() +} + +func (f *pluginRegistrationKVFixture) registrations(t *testing.T) map[string]PluginServerConfig { + t.Helper() + + f.mu.Lock() + defer f.mu.Unlock() + + if len(f.data) == 0 { + return map[string]PluginServerConfig{} + } + var registrations map[string]PluginServerConfig + require.NoError(t, json.Unmarshal(f.data, ®istrations)) + return registrations +} + +func (f *pluginRegistrationKVFixture) writeCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.writes +} + func TestClientManagerReInitIdleTimeoutDefaulting(t *testing.T) { testCases := []struct { name string @@ -93,7 +164,10 @@ func TestClientManagerReInitIdleTimeoutDefaulting(t *testing.T) { } func TestClientManager_PluginServerRegistry_RegisterUnregisterList(t *testing.T) { - m := &ClientManager{pluginServers: map[string]PluginServerConfig{}, pluginRegistered: map[string]bool{}} + pluginTestAPI := &plugintest.API{} + setupClientManagerTestAPI(t, pluginTestAPI) + client := pluginapi.NewClient(pluginTestAPI, nil) + m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, nil) t.Cleanup(m.Close) cfgA := PluginServerConfig{PluginID: "a", Name: "A", Path: "/mcp", Enabled: true} @@ -129,8 +203,137 @@ func TestClientManager_PluginServerRegistry_RegisterUnregisterList(t *testing.T) require.Len(t, m.ListPluginServers(), 1) } +func TestClientManager_PluginRegistrationPersistence(t *testing.T) { + pluginTestAPI := &plugintest.API{} + setupTestLogger(pluginTestAPI) + fixture := setupPluginRegistrationKV(t, pluginTestAPI, nil) + client := pluginapi.NewClient(pluginTestAPI, nil) + + m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, nil) + t.Cleanup(m.Close) + + first := PluginServerConfig{PluginID: "com.example.first", Name: "First", Path: "/mcp", Enabled: true} + second := PluginServerConfig{PluginID: "com.example.second", Name: "Second", Path: "/mcp", Enabled: true} + m.RegisterPluginServer(first) + m.RegisterPluginServer(second) + + require.Equal(t, map[string]PluginServerConfig{ + first.PluginID: first, + second.PluginID: second, + }, fixture.registrations(t)) + + m.UnregisterPluginServer(first.PluginID) + require.Equal(t, map[string]PluginServerConfig{ + second.PluginID: second, + }, fixture.registrations(t)) +} + +func TestClientManager_HydratesLivePluginRegistrations(t *testing.T) { + pluginTestAPI := &plugintest.API{} + setupTestLogger(pluginTestAPI) + + live := PluginServerConfig{ + PluginID: "com.example.live", + Name: "Live Name", + Path: "/live", + Enabled: false, + ExposeExternal: true, + } + disabled := PluginServerConfig{PluginID: "com.example.disabled", Name: "Disabled", Path: "/mcp", Enabled: true} + absent := PluginServerConfig{PluginID: "com.example.absent", Name: "Absent", Path: "/mcp", Enabled: true} + fixture := setupPluginRegistrationKV(t, pluginTestAPI, map[string]PluginServerConfig{ + live.PluginID: live, + disabled.PluginID: disabled, + absent.PluginID: absent, + }) + pluginTestAPI.On("GetConfig").Return(&model.Config{ + PluginSettings: model.PluginSettings{ + PluginStates: map[string]*model.PluginState{ + live.PluginID: {Enable: true}, + disabled.PluginID: {Enable: false}, + }, + }, + }) + client := pluginapi.NewClient(pluginTestAPI, nil) + + adminToolConfig := ToolConfig{Name: "echo", Policy: ToolPolicyAsk, Enabled: true} + m := NewClientManager(Config{ + IdleTimeoutMinutes: 30, + PluginServers: []PluginServerConfig{{ + PluginID: live.PluginID, + Name: "Stale Admin Name", + Path: "/stale", + Enabled: true, + ExposeExternal: false, + ToolConfigs: []ToolConfig{adminToolConfig}, + }}, + }, client.Log, client, nil, nil, nil, nil) + t.Cleanup(m.Close) + + got, ok := m.GetPluginServer(live.PluginID) + require.True(t, ok) + require.True(t, m.IsPluginRegistered(live.PluginID)) + require.Equal(t, live.Name, got.Name) + require.Equal(t, live.Path, got.Path) + require.True(t, got.ExposeExternal) + require.True(t, got.Enabled) + require.Equal(t, []ToolConfig{adminToolConfig}, got.ToolConfigs) + + require.False(t, m.IsPluginRegistered(disabled.PluginID)) + require.False(t, m.IsPluginRegistered(absent.PluginID)) + require.Equal(t, map[string]PluginServerConfig{ + live.PluginID: live, + }, fixture.registrations(t)) +} + +func TestClientManager_HydrationKeepsRegistrationsWhenServerConfigUnavailable(t *testing.T) { + pluginTestAPI := &plugintest.API{} + setupTestLogger(pluginTestAPI) + + first := PluginServerConfig{PluginID: "com.example.first", Name: "First", Path: "/mcp", Enabled: true} + second := PluginServerConfig{PluginID: "com.example.second", Name: "Second", Path: "/mcp", Enabled: true} + persisted := map[string]PluginServerConfig{ + first.PluginID: first, + second.PluginID: second, + } + fixture := setupPluginRegistrationKV(t, pluginTestAPI, persisted) + pluginTestAPI.On("GetConfig").Return((*model.Config)(nil)) + client := pluginapi.NewClient(pluginTestAPI, nil) + + m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, nil) + t.Cleanup(m.Close) + + require.True(t, m.IsPluginRegistered(first.PluginID)) + require.True(t, m.IsPluginRegistered(second.PluginID)) + require.Equal(t, persisted, fixture.registrations(t)) + require.Zero(t, fixture.writeCount()) +} + +func TestClientManager_UpdatePluginServerPreservesRegistrationStateAndKV(t *testing.T) { + pluginTestAPI := &plugintest.API{} + setupTestLogger(pluginTestAPI) + fixture := setupPluginRegistrationKV(t, pluginTestAPI, nil) + client := pluginapi.NewClient(pluginTestAPI, nil) + + orphan := PluginServerConfig{PluginID: "com.example.orphan", Name: "Orphan", Path: "/mcp", Enabled: true} + m := NewClientManager(Config{IdleTimeoutMinutes: 30, PluginServers: []PluginServerConfig{orphan}}, client.Log, client, nil, nil, nil, nil) + t.Cleanup(m.Close) + + orphan.Enabled = false + m.UpdatePluginServer(orphan) + + got, ok := m.GetPluginServer(orphan.PluginID) + require.True(t, ok) + require.Equal(t, orphan, got) + require.False(t, m.IsPluginRegistered(orphan.PluginID)) + require.Zero(t, fixture.writeCount()) +} + func TestClientManager_GetPluginServer(t *testing.T) { - m := &ClientManager{pluginServers: map[string]PluginServerConfig{}, pluginRegistered: map[string]bool{}} + pluginTestAPI := &plugintest.API{} + setupClientManagerTestAPI(t, pluginTestAPI) + client := pluginapi.NewClient(pluginTestAPI, nil) + m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, nil) t.Cleanup(m.Close) cfg, ok := m.GetPluginServer("missing") @@ -160,7 +363,7 @@ func TestClientManager_GetPluginServer(t *testing.T) { func TestClientManager_HydratesPluginServersFromConfig(t *testing.T) { pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) persisted := []PluginServerConfig{ @@ -226,7 +429,7 @@ func TestClientManager_HydratesPluginServersFromConfig(t *testing.T) { // plugin. func TestClientManager_ReInitSyncsPluginServerAdminFields(t *testing.T) { pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, nil) @@ -271,7 +474,7 @@ func TestClientManager_ReInitSyncsPluginServerAdminFields(t *testing.T) { func TestClientManager_ReInitInsertsConfigOnlyEntries(t *testing.T) { pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, nil) @@ -302,7 +505,7 @@ func TestClientManager_ReInitInsertsConfigOnlyEntries(t *testing.T) { // Live registrations absent from config must survive config broadcasts. func TestClientManager_ReInitPreservesUnpersistedRuntimeEntries(t *testing.T) { pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, nil) @@ -335,7 +538,7 @@ func TestClientManager_ReInitPreservesUnpersistedRuntimeEntries(t *testing.T) { func TestClientManager_IsPluginRegistered(t *testing.T) { pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) cfg := Config{ @@ -377,7 +580,7 @@ func TestClientManager_IsPluginRegistered(t *testing.T) { func TestClientManager_SyncPluginServersFromConfig_SkipsEmptyPluginID(t *testing.T) { pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) cfg := Config{ @@ -403,7 +606,7 @@ func TestClientManager_GetToolsForUser_PluginEnabled(t *testing.T) { mockAPI := newPluginHTTPForwarder(t, target) pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, mockAPI) @@ -433,7 +636,7 @@ func TestClientManager_GetToolsForUser_PluginDisabled_ZeroTools(t *testing.T) { }).Maybe() pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, mockAPI) @@ -484,7 +687,7 @@ func TestClientManager_GetToolsForUser_PluginEnabled_HTTPFailure(t *testing.T) { } pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, mockAPI) @@ -529,7 +732,7 @@ func TestClientManager_GetToolsForUser_PluginConnectErrorsAreRequestScoped(t *te } pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, mockAPI) @@ -558,7 +761,7 @@ func TestClientManager_GetToolsForUser_MultiplePluginServers(t *testing.T) { t.Cleanup(targetB.Close) pluginTestAPI := &plugintest.API{} - setupTestLogger(pluginTestAPI) + setupClientManagerTestAPI(t, pluginTestAPI) client := pluginapi.NewClient(pluginTestAPI, nil) // PluginHTTPRoundTripper rewrites paths to "//mcp"; route accordingly. @@ -599,7 +802,10 @@ func TestClientManager_GetToolsForUser_MultiplePluginServers(t *testing.T) { // Run with -race. Concurrent Register/Unregister/List/snapshot must not // deadlock or race. func TestClientManager_PluginServerRegistry_RaceSafe(t *testing.T) { - m := &ClientManager{pluginServers: map[string]PluginServerConfig{}, pluginRegistered: map[string]bool{}} + pluginTestAPI := &plugintest.API{} + setupClientManagerTestAPI(t, pluginTestAPI) + client := pluginapi.NewClient(pluginTestAPI, nil) + m := NewClientManager(Config{IdleTimeoutMinutes: 30}, client.Log, client, nil, nil, nil, nil) t.Cleanup(m.Close) const writers = 8 From c065a76646cf6ead20c0a27afc6353d10777bc32 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 4 Aug 2026 23:51:53 +0000 Subject: [PATCH 2/2] Retrigger CI after e2e-shard-4 container-boot flake Co-authored-by: nick.misasi