Loosen robotrumble validation constraints

Author: john-b-yang

codeclash/arenas/robotrumble/robotrumble.py

@@ -1,4 +1,5 @@
 import json
+import re
 import shlex
 import subprocess
 from collections import Counter
@@ -13,6 +14,14 @@
     "js": ["function robot(state, unit) {"],
     "py": ["def robot(state, unit):", "def robot(state: State, unit: Obj)"],
 }
+# Recognizes a valid `robot(state, unit)` entry point regardless of type annotations, return
+# hints, or extra whitespace (e.g. `def robot(state, unit) -> Action:` or
+# `def robot(state: State, unit: Obj) -> Optional[Action]:`). The MAP_EXT_TO_HEADER strings
+# above stay as the human-readable examples shown in the validation error message.
+MAP_EXT_TO_SIGNATURE_RE = {
+    "js": re.compile(r"function\s+robot\s*\(\s*state\b[^)]*,\s*unit\b[^)]*\)"),
+    "py": re.compile(r"def\s+robot\s*\(\s*state\b[^)]*,\s*unit\b[^)]*\)"),
+}
 ROBOTRUMBLE_HIDDEN_EXEC = ".codeclash_exec"
 
 
@@ -155,10 +164,10 @@ def validate_code(self, agent: Player) -> tuple[bool, str | None]:
             return False, "There should be a `robot.js` or `robot.py` file"
         agent.environment.execute(f'echo "robot.{ext}" > {ROBOTRUMBLE_HIDDEN_EXEC}')
 
-        # Check that the robot function is defined
-        if not any(
-            [header in agent.environment.execute(f"cat robot.{ext}")["output"] for header in MAP_EXT_TO_HEADER[ext]]
-        ):
+        # Check that the robot function is defined. Match on a signature regex (tolerant of type
+        # annotations, return hints, and whitespace) rather than an exact substring, so a valid
+        # entry point like `def robot(state, unit) -> Action:` is not spuriously rejected.
+        if not MAP_EXT_TO_SIGNATURE_RE[ext].search(agent.environment.execute(f"cat robot.{ext}")["output"]):
             headers = "\n- ".join(MAP_EXT_TO_HEADER[ext])
             return (
                 False,

tests/arenas/test_robotrumble.py

@@ -86,6 +86,28 @@ def test_valid_py_submission(self, arena, mock_player_factory):
         assert is_valid is True
         assert error is None
 
+    def test_valid_py_submission_annotated(self, arena, mock_player_factory):
+        """A robot() with type annotations / a return hint must still validate.
+
+        Regression: the old exact-substring check rejected `def robot(state, unit) -> Action:`
+        (matched neither whitelisted header), auto-failing otherwise-valid submissions.
+        """
+        annotated = '\ndef robot(state, unit) -> "Action":\n    return Action.move(Direction.East)\n'
+        player = mock_player_factory(
+            name="test_player",
+            files={"robot.py": annotated},
+            command_outputs={
+                "test -f robot.js && echo 'exists'": {"output": "", "returncode": 1},
+                "test -f robot.py && echo 'exists'": {"output": "exists", "returncode": 0},
+                "cat robot.py": {"output": annotated, "returncode": 0},
+                f'echo "robot.py" > {ROBOTRUMBLE_HIDDEN_EXEC}': {"output": "", "returncode": 0},
+                "./rumblebot run term --raw robot.py robot.py -t 1": {"output": "Blue won", "returncode": 0},
+            },
+        )
+        is_valid, error = arena.validate_code(player)
+        assert is_valid is True
+        assert error is None
+
     def test_missing_robot_file(self, arena, mock_player_factory):
         """Test that missing robot file fails validation."""
         player = mock_player_factory(