diff --git a/secator/ai/utils.py b/secator/ai/utils.py index 6f1eb92fb..e564ece3a 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -40,6 +40,39 @@ def _strip_leading_orphan_tools(messages: List[Dict]) -> int: return removed +def _dedupe_tool_results(messages: List[Dict]) -> int: + """Drop duplicate tool_result messages sharing a tool_call_id. + + Anthropic (and OpenRouter's providers) fold consecutive 'tool' messages into a + single user turn and reject more than one tool_result per tool_use id + ("each tool_use must have a single result. Found multiple tool_result blocks + with id X") — a NON-retryable 400. Duplicates arise when batch results are + grouped out of order (itertools.groupby only groups *consecutive* keys), or + when history trim/compaction restructures the window. Within each run of + consecutive 'tool' messages, keep the first result for each id and drop the + rest (in place). Returns the number removed. + """ + removed = 0 + i = 0 + while i < len(messages): + if messages[i].get("role") != "tool": + i += 1 + continue + seen = set() + j = i + while j < len(messages) and messages[j].get("role") == "tool": + tc_id = messages[j].get("tool_call_id") + if tc_id is not None and tc_id in seen: + del messages[j] + removed += 1 + continue # a message shifted into j; re-check without advancing + if tc_id is not None: + seen.add(tc_id) + j += 1 + i = j + return removed + + def _repair_orphan_tool_uses(messages: List[Dict]) -> int: """Repair orphan tool_use/tool_result pairing for Anthropic/OpenAI. @@ -58,6 +91,9 @@ def _repair_orphan_tool_uses(messages: List[Dict]) -> int: """ # Leading orphan tool_results have no parent in this window — drop them. repaired = _strip_leading_orphan_tools(messages) + # Duplicate tool_results for one id are rejected as a non-retryable 400 — drop + # extras so the request is valid (and, when hit as a 400, so the retry repairs it). + repaired += _dedupe_tool_results(messages) inserted = 0 i = 0 while i < len(messages): diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 9ff60abcf..1a0b2c4fd 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -2,7 +2,6 @@ """AI-powered penetration testing task.""" import json import uuid -from itertools import groupby from pathlib import Path from typing import Generator @@ -1023,11 +1022,18 @@ def _dispatch_and_collect(self, actions, ctx): collected.append(result) ctx.results.append(result) - # Group results by tool_call_id and add to history + # Group results by tool_call_id and add to history. Use an order-preserving + # dict, NOT itertools.groupby: batch results (_run_batch) interleave by id, and + # groupby only groups *consecutive* keys — so an interleaved id yielded several + # groups and thus several tool_result messages for one tool_use, which the + # provider rejects ("multiple tool_result blocks with id X"). A dict groups all + # of an id's results together regardless of arrival order → exactly one result. budget = self.history.get_action_budget(self.model) fallback_path = Path(self.reports_folder) / "report.json" if self.reports_folder else None - for tc_id, group in groupby(collected, key=lambda r: r["_context"]['tool_call_id']): - group_results = list(group) + grouped = {} + for r in collected: + grouped.setdefault(r["_context"]['tool_call_id'], []).append(r) + for tc_id, group_results in grouped.items(): tc_name = group_results[0]["_context"]['tool_call_name'] has_errors = any(r["_type"] == "error" for r in group_results) serialized = [ diff --git a/tests/unit/test_ai_utils.py b/tests/unit/test_ai_utils.py index 8140bde4e..70b587de8 100644 --- a/tests/unit/test_ai_utils.py +++ b/tests/unit/test_ai_utils.py @@ -296,6 +296,45 @@ def side_effect(**kwargs): self.assertEqual(mock_completion.call_count, 2) # repaired then succeeded mock_sleep.assert_not_called() # repair skips the backoff + @patch('time.sleep') + @patch('litellm.completion') + def test_call_llm_duplicate_tool_result_400_repairs_and_retries(self, mock_completion, mock_sleep): + """A 'multiple tool_result blocks with id' 400 is now deduped and retried, + instead of failing fast as non-retryable.""" + import litellm + from secator.ai.utils import call_llm + + ok_response = MagicMock() + ok_response.choices = [MagicMock(message=MagicMock(content="ok", tool_calls=None))] + ok_response.usage = None + + err = litellm.BadRequestError( + message=("messages.24.content.3: each tool_use must have a single result. " + "Found multiple tool_result blocks with id: toolu_dup"), + model="claude", llm_provider="anthropic", + ) + calls = [] + + def side_effect(**kwargs): + if not calls: # first call: inject an assistant + duplicate tool_results, then raise + kwargs["messages"][:0] = [ + {"role": "assistant", "content": None, + "tool_calls": [{"id": "toolu_dup", "type": "function", + "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "toolu_dup", "name": "f", "content": "r1"}, + {"role": "tool", "tool_call_id": "toolu_dup", "name": "f", "content": "r2"}, + ] + calls.append(1) + raise err + return ok_response + + mock_completion.side_effect = side_effect + result = call_llm([{"role": "user", "content": "hi"}], "claude", max_retries=3) + + self.assertEqual(result["content"], "ok") + self.assertEqual(mock_completion.call_count, 2) # deduped then succeeded + mock_sleep.assert_not_called() + @patch('time.sleep') @patch('litellm.completion') @patch('litellm.completion_cost') @@ -560,5 +599,63 @@ def completion_side_effect(**kwargs): mock_sleep.assert_not_called() # repair branch should skip the backoff sleep +class TestDedupeToolResults(unittest.TestCase): + """Providers reject >1 tool_result per tool_use id ('multiple tool_result blocks + with id X') — a non-retryable 400. Duplicates arise from batch results grouped + out of order or history trim/compaction; drop the extras, keep the first.""" + + def test_drops_consecutive_duplicate_same_id(self): + from secator.ai.utils import _dedupe_tool_results + messages = [ + {"role": "assistant", "content": None, + "tool_calls": [{"id": "x", "type": "function", "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "first"}, + {"role": "tool", "tool_call_id": "y", "name": "f", "content": "other"}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "DUP"}, + {"role": "user", "content": "next"}, + ] + removed = _dedupe_tool_results(messages) + self.assertEqual(removed, 1) + tool_ids = [m["tool_call_id"] for m in messages if m.get("role") == "tool"] + self.assertEqual(tool_ids, ["x", "y"]) # first x kept, dup dropped, y intact + # the kept x is the FIRST result + self.assertEqual(next(m for m in messages if m.get("tool_call_id") == "x")["content"], "first") + + def test_no_op_when_unique(self): + from secator.ai.utils import _dedupe_tool_results + messages = [ + {"role": "tool", "tool_call_id": "a", "content": "1"}, + {"role": "tool", "tool_call_id": "b", "content": "2"}, + ] + before = [dict(m) for m in messages] + self.assertEqual(_dedupe_tool_results(messages), 0) + self.assertEqual(messages, before) + + def test_dedupe_scoped_per_consecutive_run(self): + """The same id in two SEPARATE tool runs (own assistant each) is not a dup.""" + from secator.ai.utils import _dedupe_tool_results + messages = [ + {"role": "tool", "tool_call_id": "x", "content": "r1"}, + {"role": "assistant", "content": "thinking"}, + {"role": "tool", "tool_call_id": "x", "content": "r2"}, + ] + self.assertEqual(_dedupe_tool_results(messages), 0) # separated by a non-tool msg + + def test_repair_dedupes_duplicate_tool_results(self): + """_repair_orphan_tool_uses now removes duplicates as part of its pass.""" + from secator.ai.utils import _repair_orphan_tool_uses + messages = [ + {"role": "assistant", "content": None, + "tool_calls": [{"id": "x", "type": "function", "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "first"}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "DUP"}, + {"role": "user", "content": "next"}, + ] + changed = _repair_orphan_tool_uses(messages) + self.assertGreaterEqual(changed, 1) + tool_ids = [m["tool_call_id"] for m in messages if m.get("role") == "tool"] + self.assertEqual(tool_ids, ["x"]) # exactly one result for x + + if __name__ == '__main__': unittest.main()