From 21800399efd44f66d66af59de26ffbd365f1dcbf Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Mon, 17 Aug 2026 22:53:52 +0100
Subject: [PATCH] Fix PaperBench code-only fallback tree

---
 project/paperbench/paperbench/grade.py        | 14 ++++--
 .../tests/unit/test_grade_code_only.py        | 50 +++++++++++++++++++
 2 files changed, 60 insertions(+), 4 deletions(-)
 create mode 100644 project/paperbench/tests/unit/test_grade_code_only.py

diff --git a/project/paperbench/paperbench/grade.py b/project/paperbench/paperbench/grade.py
index ae294933..344118bc 100644
--- a/project/paperbench/paperbench/grade.py
+++ b/project/paperbench/paperbench/grade.py
@@ -5,7 +5,7 @@
 import tempfile
 import time
 from contextlib import nullcontext
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Any
 
@@ -80,6 +80,14 @@ def success(self) -> bool:
         return self.num_invalid_leaf_nodes < self.num_leaf_nodes
 
 
+def _code_only_task_tree(task_tree: TaskNode) -> TaskNode:
+    code_only_tree = task_tree.code_only()
+    if code_only_tree is not None:
+        return code_only_tree
+
+    return replace(task_tree, sub_tasks=[], task_category="Code Development")
+
+
 async def run_judge(
     submission_path: Path,
     paper_id: str,
@@ -107,9 +115,7 @@ async def run_judge(
     with open(rubric_path, "r") as f:
         task_tree = TaskNode.from_dict(json.load(f))
     if code_only:
-        task_tree = task_tree.code_only() or task_tree.set_task_category(
-            "Code Development"
-        ).set_sub_tasks([])
+        task_tree = _code_only_task_tree(task_tree)
     if resources_provided:
         task_tree = task_tree.resources_provided()
 
diff --git a/project/paperbench/tests/unit/test_grade_code_only.py b/project/paperbench/tests/unit/test_grade_code_only.py
new file mode 100644
index 00000000..083f1e5e
--- /dev/null
+++ b/project/paperbench/tests/unit/test_grade_code_only.py
@@ -0,0 +1,50 @@
+from paperbench.grade import _code_only_task_tree
+from paperbench.rubric.tasks import TaskNode
+
+
+def test_code_only_fallback_collapses_non_code_rubric_to_valid_leaf() -> None:
+    leaf = TaskNode(
+        id="analysis",
+        requirements="analyse results",
+        weight=1,
+        task_category="Result Analysis",
+    )
+    root = TaskNode(
+        id="root",
+        requirements="complete the paper",
+        weight=3,
+        sub_tasks=[leaf],
+    )
+
+    code_only = _code_only_task_tree(root)
+
+    assert code_only.is_leaf()
+    assert code_only.task_category == "Code Development"
+    assert code_only.id == root.id
+    assert code_only.requirements == root.requirements
+    assert code_only.weight == root.weight
+
+
+def test_code_only_keeps_existing_code_development_subtree() -> None:
+    code_leaf = TaskNode(
+        id="code",
+        requirements="implement method",
+        weight=1,
+        task_category="Code Development",
+    )
+    analysis_leaf = TaskNode(
+        id="analysis",
+        requirements="analyse results",
+        weight=1,
+        task_category="Result Analysis",
+    )
+    root = TaskNode(
+        id="root",
+        requirements="complete the paper",
+        weight=2,
+        sub_tasks=[code_leaf, analysis_leaf],
+    )
+
+    code_only = _code_only_task_tree(root)
+
+    assert [node.id for node in code_only.get_leaf_nodes()] == ["code"]