From 7af577546d8adeadbda9c94c1b8df5750e780236 Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Sat, 31 Jan 2026 16:31:31 +0200 Subject: [PATCH] fix(oauth): clean up orphaned tokens on server removal and startup This fixes the issue where OAuth tokens accumulate in the database for servers that have been removed, causing unnecessary refresh attempts and confusing error messages like "OAuth token refresh failed - token expired too long ago". Changes: 1. RemoveServer() now calls ClearOAuthState() to remove OAuth tokens when a server is deleted 2. RemoveServer() notifies RefreshManager to stop tracking the removed server's token refresh schedule 3. Added CleanupOrphanedOAuthTokens() to storage manager that removes tokens for servers no longer in configuration 4. Startup now cleans up orphaned tokens before starting RefreshManager Key storage design notes: - Tokens use serverName_hash(serverName|serverURL) as storage key - This ensures multiple servers can share the same URL (different names) - URL changes for a server name are properly handled via ClearOAuthState which deletes all tokens with the serverName_ prefix Tested: Verified 7 orphaned tokens were cleaned up on startup Co-Authored-By: Claude Opus 4.5 --- internal/runtime/lifecycle.go | 18 +++++++ internal/server/server.go | 14 +++++ internal/storage/manager.go | 65 ++++++++++++++++++++++ internal/storage/manager_oauth_test.go | 75 ++++++++++++++++++++++++++ 4 files changed, 172 insertions(+) diff --git a/internal/runtime/lifecycle.go b/internal/runtime/lifecycle.go index d78db99b..0295938d 100644 --- a/internal/runtime/lifecycle.go +++ b/internal/runtime/lifecycle.go @@ -30,6 +30,24 @@ func (r *Runtime) StartBackgroundInitialization() { r.logger.Info("Update checker background process started") } + // Clean up orphaned OAuth tokens before starting the refresh manager + // This removes tokens for servers that were deleted while mcpproxy was not running + if r.storageManager != nil { + r.mu.RLock() + validServerNames := make([]string, 0, len(r.cfg.Servers)) + for _, server := range r.cfg.Servers { + validServerNames = append(validServerNames, server.Name) + } + r.mu.RUnlock() + + if deleted, err := r.storageManager.CleanupOrphanedOAuthTokens(validServerNames); err != nil { + r.logger.Warn("Failed to cleanup orphaned OAuth tokens", zap.Error(err)) + } else if deleted > 0 { + r.logger.Info("Cleaned up orphaned OAuth tokens during startup", + zap.Int("deleted", deleted)) + } + } + // Start proactive OAuth token refresh manager if r.refreshManager != nil { r.refreshManager.SetRuntime(r) diff --git a/internal/server/server.go b/internal/server/server.go index d077f0e6..ca809303 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -922,6 +922,20 @@ func (s *Server) RemoveServer(ctx context.Context, serverName string) error { return fmt.Errorf("failed to remove server from storage: %w", err) } + // Clear OAuth state (tokens, client registration) for the removed server + // This prevents orphaned tokens from accumulating in the database + if err := storageManager.ClearOAuthState(serverName); err != nil { + s.logger.Warn("Failed to clear OAuth state for removed server", + zap.String("server", serverName), + zap.Error(err)) + // Continue - this is cleanup, not critical for removal + } + + // Notify RefreshManager to stop tracking this server's token refresh + if refreshManager := s.runtime.RefreshManager(); refreshManager != nil { + refreshManager.OnTokenCleared(serverName) + } + // Remove from search index if err := s.runtime.IndexManager().DeleteServerTools(serverName); err != nil { s.logger.Warn("Failed to remove server tools from index", diff --git a/internal/storage/manager.go b/internal/storage/manager.go index 4153e3f8..9604f466 100644 --- a/internal/storage/manager.go +++ b/internal/storage/manager.go @@ -1383,3 +1383,68 @@ func (m *Manager) ClearOAuthState(serverName string) error { } return nil } + +// CleanupOrphanedOAuthTokens removes OAuth tokens for servers that no longer exist in the configuration. +// This should be called during startup to clean up tokens left behind when servers were removed +// while mcpproxy was not running, or from older versions that didn't clean up properly. +func (m *Manager) CleanupOrphanedOAuthTokens(validServerNames []string) (int, error) { + m.mu.Lock() + defer m.mu.Unlock() + + if m.db == nil { + return 0, fmt.Errorf("storage not initialized") + } + + // Create a set of valid server names for fast lookup + validServers := make(map[string]bool, len(validServerNames)) + for _, name := range validServerNames { + validServers[name] = true + } + + // Get all OAuth tokens + tokens, err := m.db.ListOAuthTokens() + if err != nil { + return 0, fmt.Errorf("failed to list OAuth tokens: %w", err) + } + + // Find orphaned tokens (tokens whose server no longer exists) + var orphanedKeys []string + for _, token := range tokens { + serverName := token.GetServerName() + if serverName == "" { + // Token has no server name, consider it orphaned + orphanedKeys = append(orphanedKeys, token.ServerName) + continue + } + if !validServers[serverName] { + orphanedKeys = append(orphanedKeys, token.ServerName) + m.logger.Infow("Found orphaned OAuth token", + "storage_key", token.ServerName, + "display_name", serverName) + } + } + + if len(orphanedKeys) == 0 { + m.logger.Debug("No orphaned OAuth tokens found") + return 0, nil + } + + // Delete orphaned tokens + deleted := 0 + for _, key := range orphanedKeys { + if err := m.db.DeleteOAuthToken(key); err != nil { + m.logger.Warnw("Failed to delete orphaned OAuth token", + "key", key, + "error", err) + continue + } + deleted++ + } + + m.logger.Infow("Cleaned up orphaned OAuth tokens", + "total_tokens", len(tokens), + "orphaned_found", len(orphanedKeys), + "deleted", deleted) + + return deleted, nil +} diff --git a/internal/storage/manager_oauth_test.go b/internal/storage/manager_oauth_test.go index d3ea9cdf..09f0fc01 100644 --- a/internal/storage/manager_oauth_test.go +++ b/internal/storage/manager_oauth_test.go @@ -101,6 +101,81 @@ func TestBoltDB_GetOAuthClientCredentials_LegacyRecord(t *testing.T) { require.Equal(t, 0, gotPort, "Legacy records should return port 0") } +// TestManager_CleanupOrphanedOAuthTokens verifies that orphaned tokens are removed +// while tokens for valid servers are preserved. +func TestManager_CleanupOrphanedOAuthTokens(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "storage-cleanup-oauth-*") + require.NoError(t, err) + t.Cleanup(func() { _ = os.RemoveAll(tmpDir) }) + + mgr, err := storage.NewManager(tmpDir, zap.NewNop().Sugar()) + require.NoError(t, err) + t.Cleanup(func() { _ = mgr.Close() }) + + // Create tokens for 4 servers: 2 valid, 2 orphaned + validServer1 := "valid-server-1" + validServer2 := "valid-server-2" + orphanedServer1 := "orphaned-server-1" + orphanedServer2 := "orphaned-server-2" + + // Generate server keys with URLs + validKey1 := oauth.GenerateServerKey(validServer1, "https://valid1.example.com") + validKey2 := oauth.GenerateServerKey(validServer2, "https://valid2.example.com") + orphanedKey1 := oauth.GenerateServerKey(orphanedServer1, "https://orphan1.example.com") + orphanedKey2 := oauth.GenerateServerKey(orphanedServer2, "https://orphan2.example.com") + + // Save tokens with DisplayName set (as PersistentTokenStore does) + err = mgr.GetBoltDB().SaveOAuthToken(&storage.OAuthTokenRecord{ + ServerName: validKey1, + DisplayName: validServer1, + AccessToken: "valid-token-1", + }) + require.NoError(t, err) + + err = mgr.GetBoltDB().SaveOAuthToken(&storage.OAuthTokenRecord{ + ServerName: validKey2, + DisplayName: validServer2, + AccessToken: "valid-token-2", + }) + require.NoError(t, err) + + err = mgr.GetBoltDB().SaveOAuthToken(&storage.OAuthTokenRecord{ + ServerName: orphanedKey1, + DisplayName: orphanedServer1, + AccessToken: "orphan-token-1", + }) + require.NoError(t, err) + + err = mgr.GetBoltDB().SaveOAuthToken(&storage.OAuthTokenRecord{ + ServerName: orphanedKey2, + DisplayName: orphanedServer2, + AccessToken: "orphan-token-2", + }) + require.NoError(t, err) + + // Cleanup with only valid servers + validServers := []string{validServer1, validServer2} + deleted, err := mgr.CleanupOrphanedOAuthTokens(validServers) + require.NoError(t, err) + require.Equal(t, 2, deleted, "Should delete 2 orphaned tokens") + + // Verify valid tokens still exist + token1, err := mgr.GetBoltDB().GetOAuthToken(validKey1) + require.NoError(t, err) + require.Equal(t, "valid-token-1", token1.AccessToken) + + token2, err := mgr.GetBoltDB().GetOAuthToken(validKey2) + require.NoError(t, err) + require.Equal(t, "valid-token-2", token2.AccessToken) + + // Verify orphaned tokens are gone + _, err = mgr.GetBoltDB().GetOAuthToken(orphanedKey1) + require.Error(t, err, "Orphaned token 1 should be deleted") + + _, err = mgr.GetBoltDB().GetOAuthToken(orphanedKey2) + require.Error(t, err, "Orphaned token 2 should be deleted") +} + // TestBoltDB_ClearOAuthClientCredentials_PreservesToken verifies that ClearOAuthClientCredentials // clears DCR fields but preserves token data (Spec 022). func TestBoltDB_ClearOAuthClientCredentials_PreservesToken(t *testing.T) {