From d22644aadac02cf7aab1c755788180ef151b69fd Mon Sep 17 00:00:00 2001
From: Jakub Zacek <dawon@dawon.eu>
Date: Tue, 22 Apr 2025 11:30:00 +0200
Subject: [PATCH] Expand cooldown API (#12435)

---
 .../java/org/bukkit/entity/HumanEntity.java   | 31 +++++++++++++++++--
 .../craftbukkit/entity/CraftHumanEntity.java  | 24 ++++++++++++--
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/paper-api/src/main/java/org/bukkit/entity/HumanEntity.java b/paper-api/src/main/java/org/bukkit/entity/HumanEntity.java
index 8304cafc66..a4484ad3cb 100644
--- a/paper-api/src/main/java/org/bukkit/entity/HumanEntity.java
+++ b/paper-api/src/main/java/org/bukkit/entity/HumanEntity.java
@@ -3,6 +3,8 @@ package org.bukkit.entity;
 import java.util.Collection;
 import java.util.Set;
 import java.util.function.Consumer;
+import io.papermc.paper.datacomponent.item.UseCooldown;
+import net.kyori.adventure.key.Key;
 import org.bukkit.GameMode;
 import org.bukkit.Location;
 import org.bukkit.Material;
@@ -368,7 +370,7 @@ public interface HumanEntity extends LivingEntity, AnimalTamer, InventoryHolder
 
     /**
      * Set a cooldown on the specified material for a certain amount of ticks.
-     * ticks. 0 ticks will result in the removal of the cooldown.
+     * 0 ticks will result in the removal of the cooldown.
      * <p>
      * Cooldowns are used by the server for items such as ender pearls and
      * shields to prevent them from being used repeatedly.
@@ -415,7 +417,7 @@ public interface HumanEntity extends LivingEntity, AnimalTamer, InventoryHolder
 
     /**
      * Set a cooldown on the specified item for a certain amount of ticks.
-     * ticks. 0 ticks will result in the removal of the cooldown.
+     * 0 ticks will result in the removal of the cooldown.
      * <p>
      * Cooldowns are used by the server for items such as ender pearls and
      * shields to prevent them from being used repeatedly.
@@ -428,6 +430,31 @@ public interface HumanEntity extends LivingEntity, AnimalTamer, InventoryHolder
      */
     public void setCooldown(ItemStack item, int ticks);
 
+    /**
+     * Get the cooldown time in ticks remaining for the specified cooldown group.
+     *
+     * @param key the cooldown group to check
+     * @return the remaining cooldown time in ticks
+     * @see UseCooldown#cooldownGroup()
+     */
+    public int getCooldown(Key key);
+
+    /**
+     * Set a cooldown on items with the specified cooldown group for a certain amount of ticks.
+     * 0 ticks will result in the removal of the cooldown.
+     * <p>
+     * Cooldowns are used by the server for items such as ender pearls and
+     * shields to prevent them from being used repeatedly.
+     * <p>
+     * Note that cooldowns will not by themselves stop an item from being used
+     * for attacking.
+     *
+     * @param key cooldown group to set the cooldown for
+     * @param ticks the amount of ticks to set or 0 to remove
+     * @see UseCooldown#cooldownGroup()
+     */
+    public void setCooldown(Key key, int ticks);
+
     /**
      * Get the sleep ticks of the player. This value may be capped.
      *
diff --git a/paper-server/src/main/java/org/bukkit/craftbukkit/entity/CraftHumanEntity.java b/paper-server/src/main/java/org/bukkit/craftbukkit/entity/CraftHumanEntity.java
index e0e6a5087a..6ad1c2220b 100644
--- a/paper-server/src/main/java/org/bukkit/craftbukkit/entity/CraftHumanEntity.java
+++ b/paper-server/src/main/java/org/bukkit/craftbukkit/entity/CraftHumanEntity.java
@@ -8,6 +8,8 @@ import java.util.Collection;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Consumer;
+import io.papermc.paper.adventure.PaperAdventure;
+import net.kyori.adventure.key.Key;
 import net.minecraft.core.BlockPos;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.network.chat.Component;
@@ -651,14 +653,14 @@ public class CraftHumanEntity extends CraftLivingEntity implements HumanEntity {
 
     @Override
     public boolean hasCooldown(ItemStack item) {
-        Preconditions.checkArgument(item != null, "Material cannot be null");
+        Preconditions.checkArgument(item != null, "Item cannot be null");
 
         return this.getHandle().getCooldowns().isOnCooldown(CraftItemStack.asNMSCopy(item));
     }
 
     @Override
     public int getCooldown(ItemStack item) {
-        Preconditions.checkArgument(item != null, "Material cannot be null");
+        Preconditions.checkArgument(item != null, "Item cannot be null");
 
         ResourceLocation group = this.getHandle().getCooldowns().getCooldownGroup(CraftItemStack.asNMSCopy(item));
         if (group == null) {
@@ -671,12 +673,28 @@ public class CraftHumanEntity extends CraftLivingEntity implements HumanEntity {
 
     @Override
     public void setCooldown(ItemStack item, int ticks) {
-        Preconditions.checkArgument(item != null, "Material cannot be null");
+        Preconditions.checkArgument(item != null, "Item cannot be null");
         Preconditions.checkArgument(ticks >= 0, "Cannot have negative cooldown");
 
         this.getHandle().getCooldowns().addCooldown(CraftItemStack.asNMSCopy(item), ticks);
     }
 
+    @Override
+    public int getCooldown(Key key) {
+        Preconditions.checkArgument(key != null, "Key cannot be null");
+
+        ItemCooldowns.CooldownInstance cooldown = this.getHandle().getCooldowns().cooldowns.get(PaperAdventure.asVanilla(key));
+        return (cooldown == null) ? 0 : Math.max(0, cooldown.endTime() - this.getHandle().getCooldowns().tickCount);
+    }
+
+    @Override
+    public void setCooldown(Key key, int ticks) {
+        Preconditions.checkArgument(key != null, "Key cannot be null");
+        Preconditions.checkArgument(ticks >= 0, "Cannot have negative cooldown");
+
+        this.getHandle().getCooldowns().addCooldown(PaperAdventure.asVanilla(key), ticks);
+    }
+
     @Override
     public org.bukkit.entity.Entity releaseLeftShoulderEntity() {
         if (!getHandle().getShoulderEntityLeft().isEmpty()) {