From 489794ae3df6dff436692175cfbf018f611afb28 Mon Sep 17 00:00:00 2001
From: Bartok9 <183148895+Bartok9@users.noreply.github.com>
Date: Sat, 1 Aug 2026 01:51:54 -0400
Subject: [PATCH] fix(bench): reject negative TaskArgs.extra_tool_calls

Fail-closed validation so pass budgets cannot be inflated with negative
extra tool-call allowances. Reject bool (int subclass) before coercion.

Signed-off-by: Bartok9 <183148895+Bartok9@users.noreply.github.com>
---
 .../tool_calling_agent/interfaces.py          | 12 ++++++-
 .../tool_calling_agent/test_task_args.py      | 36 +++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 tests/rai_bench/tool_calling_agent/test_task_args.py

diff --git a/src/rai_bench/rai_bench/tool_calling_agent/interfaces.py b/src/rai_bench/rai_bench/tool_calling_agent/interfaces.py
index a92fbd57b..3d40ac6d8 100644
--- a/src/rai_bench/rai_bench/tool_calling_agent/interfaces.py
+++ b/src/rai_bench/rai_bench/tool_calling_agent/interfaces.py
@@ -19,7 +19,7 @@
 from langchain_core.messages import AIMessage, BaseMessage, ToolCall
 from langchain_core.runnables.config import DEFAULT_RECURSION_LIMIT
 from langchain_core.tools import BaseTool
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 
 from rai_bench.tool_calling_agent.results_tracking import SubTaskResult, ValidatorResult
 
@@ -488,6 +488,16 @@ class TaskArgs(BaseModel):
     prompt_detail: Literal["brief", "descriptive"] = "brief"
     examples_in_system_prompt: Literal[0, 2, 5] = 0
 
+    @field_validator("extra_tool_calls", mode="before")
+    @classmethod
+    def _extra_tool_calls_non_negative(cls, v: object) -> int:
+        # bool is a subclass of int; reject it explicitly (before coercion).
+        if isinstance(v, bool) or not isinstance(v, int):
+            raise ValueError("extra_tool_calls must be a non-negative int")
+        if v < 0:
+            raise ValueError("extra_tool_calls must be >= 0")
+        return v
+
 
 class Task(ABC):
     complexity: Literal["easy", "medium", "hard"]
diff --git a/tests/rai_bench/tool_calling_agent/test_task_args.py b/tests/rai_bench/tool_calling_agent/test_task_args.py
new file mode 100644
index 000000000..90198e038
--- /dev/null
+++ b/tests/rai_bench/tool_calling_agent/test_task_args.py
@@ -0,0 +1,36 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+from pydantic import ValidationError
+
+from rai_bench.tool_calling_agent.interfaces import TaskArgs
+
+
+def test_task_args_extra_tool_calls_default_and_positive() -> None:
+    assert TaskArgs().extra_tool_calls == 0
+    assert TaskArgs(extra_tool_calls=0).extra_tool_calls == 0
+    assert TaskArgs(extra_tool_calls=3).extra_tool_calls == 3
+
+
+def test_task_args_extra_tool_calls_rejects_negative() -> None:
+    with pytest.raises(ValidationError):
+        TaskArgs(extra_tool_calls=-1)
+
+
+def test_task_args_extra_tool_calls_rejects_bool() -> None:
+    with pytest.raises(ValidationError):
+        TaskArgs(extra_tool_calls=True)  # type: ignore[arg-type]
+    with pytest.raises(ValidationError):
+        TaskArgs(extra_tool_calls=False)  # type: ignore[arg-type]