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) {