From 8595b3515602778a7364df103aaf1d87ea3dadc4 Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Mon, 17 Aug 2026 22:45:59 +0100
Subject: [PATCH] Fix PaperBench rubric pruning at subtree boundaries

---
 project/paperbench/paperbench/rubric/tasks.py |  2 +-
 .../tests/unit/test_rubric_tasks.py           | 24 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)
 create mode 100644 project/paperbench/tests/unit/test_rubric_tasks.py

diff --git a/project/paperbench/paperbench/rubric/tasks.py b/project/paperbench/paperbench/rubric/tasks.py
index 887e6984..4dae6226 100644
--- a/project/paperbench/paperbench/rubric/tasks.py
+++ b/project/paperbench/paperbench/rubric/tasks.py
@@ -312,7 +312,7 @@ def prune_to_depth(self, max_depth: int, current_depth: int = 0) -> TaskNode:
                 requirements=self.requirements,
                 weight=self.weight,
                 sub_tasks=[],
-                task_category=self.task_category,
+                task_category=self.task_category or "Subtree",
                 finegrained_task_category=self.finegrained_task_category,
             )
 
diff --git a/project/paperbench/tests/unit/test_rubric_tasks.py b/project/paperbench/tests/unit/test_rubric_tasks.py
new file mode 100644
index 00000000..3853e61b
--- /dev/null
+++ b/project/paperbench/tests/unit/test_rubric_tasks.py
@@ -0,0 +1,24 @@
+from paperbench.rubric.tasks import TaskNode
+
+
+def test_prune_to_depth_marks_collapsed_internal_node_as_subtree() -> None:
+    leaf = TaskNode(
+        id="leaf",
+        requirements="leaf requirement",
+        weight=1,
+        task_category="Code Development",
+    )
+    root = TaskNode(
+        id="root",
+        requirements="root requirement",
+        weight=1,
+        sub_tasks=[leaf],
+    )
+
+    pruned = root.prune_to_depth(0)
+
+    assert pruned.is_leaf()
+    assert pruned.task_category == "Subtree"
+    assert pruned.id == root.id
+    assert pruned.requirements == root.requirements
+    assert pruned.weight == root.weight