Fix PaperBench rubric pruning at subtree boundaries

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

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