From 669cc55ba477c19fd17ebaf22534ad127d4305b7 Mon Sep 17 00:00:00 2001
From: osjarw <62134478+osjarw@users.noreply.github.com>
Date: Tue, 30 Jul 2024 05:18:53 +0300
Subject: [PATCH] Allow the creation of recursive HTNs (#30138)

Allow recursive HTNs
---
 Content.IntegrationTests/Tests/NPC/NPCTest.cs  | 4 ++++
 Content.Server/NPC/HTN/HTNCompoundPrototype.cs | 6 ++++++
 Content.Server/NPC/HTN/HTNPlanJob.cs           | 5 +++++
 Content.Server/NPC/HTN/HTNTask.cs              | 6 ++++++
 4 files changed, 21 insertions(+)

diff --git a/Content.IntegrationTests/Tests/NPC/NPCTest.cs b/Content.IntegrationTests/Tests/NPC/NPCTest.cs
index 83321fe613..064fd6c5bf 100644
--- a/Content.IntegrationTests/Tests/NPC/NPCTest.cs
+++ b/Content.IntegrationTests/Tests/NPC/NPCTest.cs
@@ -45,6 +45,10 @@ public sealed class NPCTest
                     var count = counts.GetOrNew(compound.ID);
                     count++;
 
+                    // Compound tasks marked with AllowRecursion are only evaluated once
+                    if (counts.ContainsKey(compound.ID) && compound.AllowRecursion)
+                        continue;
+
                     Assert.That(count, Is.LessThan(50));
                     counts[compound.ID] = count;
                     Count(protoManager.Index<HTNCompoundPrototype>(compoundTask.Task), counts, htnSystem, protoManager);
diff --git a/Content.Server/NPC/HTN/HTNCompoundPrototype.cs b/Content.Server/NPC/HTN/HTNCompoundPrototype.cs
index 29cfb96f97..69f8441973 100644
--- a/Content.Server/NPC/HTN/HTNCompoundPrototype.cs
+++ b/Content.Server/NPC/HTN/HTNCompoundPrototype.cs
@@ -12,4 +12,10 @@ public sealed partial class HTNCompoundPrototype : IPrototype
 
     [DataField("branches", required: true)]
     public List<HTNBranch> Branches = new();
+
+    /// <summary>
+    /// Exclude this compound task from the CompoundRecursion integration test.
+    /// </summary>
+    [DataField]
+    public bool AllowRecursion = false;
 }
diff --git a/Content.Server/NPC/HTN/HTNPlanJob.cs b/Content.Server/NPC/HTN/HTNPlanJob.cs
index 8158303524..9c62f5840a 100644
--- a/Content.Server/NPC/HTN/HTNPlanJob.cs
+++ b/Content.Server/NPC/HTN/HTNPlanJob.cs
@@ -63,8 +63,13 @@ public sealed class HTNPlanJob : Job<HTNPlan>
         // How many primitive tasks we've added since last record.
         var primitiveCount = 0;
 
+        int tasksProcessed = 0;
+
         while (tasksToProcess.TryDequeue(out var currentTask))
         {
+            if (tasksProcessed++ > _rootTask.MaximumTasks)
+                throw new Exception("HTN Planner exceeded maximum tasks");
+
             switch (currentTask)
             {
                 case HTNCompoundTask compound:
diff --git a/Content.Server/NPC/HTN/HTNTask.cs b/Content.Server/NPC/HTN/HTNTask.cs
index 0f7c0a377e..c09849d079 100644
--- a/Content.Server/NPC/HTN/HTNTask.cs
+++ b/Content.Server/NPC/HTN/HTNTask.cs
@@ -3,4 +3,10 @@ namespace Content.Server.NPC.HTN;
 [ImplicitDataDefinitionForInheritors]
 public abstract partial class HTNTask
 {
+    /// <summary>
+    /// Limit the amount of tasks the planner considers. Exceeding this value sleeps the NPC and throws an exception.
+    /// The expected way to hit this limit is with badly written recursive tasks.
+    /// </summary>
+    [DataField]
+    public int MaximumTasks = 1000;
 }
-- 
2.52.0