Make PaperCommandRegistration#unregister respect unregisterBukkit parameter

Author: willkroboth

commandapi-platforms/commandapi-bukkit/commandapi-bukkit-core/src/main/java/dev/jorel/commandapi/PaperCommandRegistration.java

@@ -7,23 +7,40 @@
 import com.mojang.brigadier.tree.RootCommandNode;
 
 import java.util.List;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 
 /**
  * Handles logic for registering commands after Paper build 65, where <a href="https://github.com/PaperMC/Paper/pull/8235">https://github.com/PaperMC/Paper/pull/8235</a>
  * changed a bunch of the behind-the-scenes logic.
  */
 public class PaperCommandRegistration<Source> extends CommandRegistrationStrategy<Source> {
-	// References to necessary objects
+	// References to necessary methods
 	private final Supplier<CommandDispatcher<Source>> getBrigadierDispatcher;
+	private final Predicate<CommandNode<Source>> isBukkitCommand;
 
 	// Store registered commands nodes for eventual reloads
 	private final RootCommandNode<Source> registeredNodes = new RootCommandNode<>();
 
-	public PaperCommandRegistration(Supplier<CommandDispatcher<Source>> getBrigadierDispatcher) {
+	public PaperCommandRegistration(
+		Supplier<CommandDispatcher<Source>> getBrigadierDispatcher, Predicate<CommandNode<Source>> isBukkitCommand
+	) {
 		this.getBrigadierDispatcher = getBrigadierDispatcher;
+		this.isBukkitCommand = isBukkitCommand;
 	}
 
+	// Provide access to internal functions that may be useful to developers
+	/**
+	 * Checks if a Brigadier command node came from wrapping a Bukkit command
+	 *
+	 * @param node The CommandNode to check
+	 * @return true if the CommandNode is being handled by Paper's BukkitCommandNode
+	 */
+	public boolean isBukkitCommand(CommandNode<Source> node) {
+		return isBukkitCommand.test(node);
+	}
+
+	// Implement CommandRegistrationStrategy methods
 	@Override
 	public CommandDispatcher<Source> getBrigadierDispatcher() {
 		return getBrigadierDispatcher.get();
@@ -57,10 +74,16 @@ public LiteralCommandNode<Source> registerCommandNode(LiteralArgumentBuilder<Sou
 	@Override
 	public void unregister(String commandName, boolean unregisterNamespaces, boolean unregisterBukkit) {
 		// Remove nodes from the  dispatcher
-		removeBrigadierCommands(getBrigadierDispatcher.get().getRoot(), commandName, unregisterNamespaces, c -> true);
-
-		// Don't add nodes back after a reload
-		removeBrigadierCommands(registeredNodes, commandName, unregisterNamespaces, c -> true);
+		removeBrigadierCommands(getBrigadierDispatcher.get().getRoot(), commandName, unregisterNamespaces,
+			// If we are unregistering a Bukkit command, ONLY unregister BukkitCommandNodes
+			// If we are unregistering a Vanilla command, DO NOT unregister BukkitCommandNodes
+			c -> !unregisterBukkit ^ isBukkitCommand.test(c));
+
+		// CommandAPI commands count as non-Bukkit
+		if (!unregisterBukkit) {
+			// Don't add nodes back after a reload
+			removeBrigadierCommands(registeredNodes, commandName, unregisterNamespaces, c -> true);
+		}
 
 		// Update the dispatcher file
 		CommandAPIHandler.getInstance().writeDispatcherToFile();

commandapi-platforms/commandapi-bukkit/commandapi-bukkit-nms/commandapi-bukkit-1.20.5/src/main/java/dev/jorel/commandapi/nms/NMS_1_20_R4.java

@@ -1059,7 +1059,19 @@ public CommandRegistrationStrategy<CommandSourceStack> createCommandRegistration
 				node -> node.getCommand() instanceof BukkitCommandWrapper
 			);
 		} else {
-			return new PaperCommandRegistration<>(() -> this.<MinecraftServer>getMinecraftServer().getCommands().getDispatcher());
+			// This class is Paper-server specific, so we need to use paper's userdev plugin to
+			//  access it directly. That might need gradle, but there might also be a maven version?
+			//  https://discord.com/channels/289587909051416579/1121227200277004398/1246910745761812480
+			Class<?> bukkitCommandNode_bukkitBrigCommand;
+			try {
+				bukkitCommandNode_bukkitBrigCommand = Class.forName("io.papermc.paper.command.brigadier.bukkit.BukkitCommandNode$BukkitBrigCommand");
+			} catch (ClassNotFoundException e) {
+				throw new IllegalStateException("Expected to find class", e);
+			}
+			return new PaperCommandRegistration<>(
+				() -> this.<MinecraftServer>getMinecraftServer().getCommands().getDispatcher(),
+				node -> bukkitCommandNode_bukkitBrigCommand.isInstance(node.getCommand())
+			);
 		}
 	}
 }