diff --git a/openhack/agents/base.py b/openhack/agents/base.py index 7d16e37..5bfcc8b 100644 --- a/openhack/agents/base.py +++ b/openhack/agents/base.py @@ -652,7 +652,16 @@ async def _agent_loop(self) -> dict: # (for example, two identical greetings). Caching it across # turns rejects a valid completion and makes the continuation # guard ask the model for a second response. - if not is_finish_task and key in self._call_cache: + if tool_call.parse_error: + result = { + "error": "tool_call_truncated", + "message": ( + f"Arguments for '{tool_call.name}' were truncated due to token limits " + f"(parse error: {tool_call.parse_error}). " + "Please retry and call the tool with more concise arguments." + ), + } + elif not is_finish_task and key in self._call_cache: # Anti-repetition: this exact call already ran — don't # re-execute (saves tokens) and nudge the agent to move on. prev_turn, prev_summ = self._call_cache[key] diff --git a/openhack/agents/browser_verifier.py b/openhack/agents/browser_verifier.py index a31daca..5f3b876 100644 --- a/openhack/agents/browser_verifier.py +++ b/openhack/agents/browser_verifier.py @@ -493,7 +493,16 @@ async def run(self, task: str, context: Optional[dict] = None) -> dict: ) try: - if tool_call.name == "browser_navigate": + if tool_call.parse_error: + result = { + "error": "tool_call_truncated", + "message": ( + f"Arguments for '{tool_call.name}' were truncated due to token limits " + f"(parse error: {tool_call.parse_error}). " + "Please retry and call the tool with more concise arguments." + ), + } + elif tool_call.name == "browser_navigate": result = await self._handle_browser_navigate(tool_call.arguments) elif tool_call.name == "browser_click": result = await self._handle_browser_click(tool_call.arguments) diff --git a/openhack/agents/hunter.py b/openhack/agents/hunter.py index b53ef39..39093c1 100644 --- a/openhack/agents/hunter.py +++ b/openhack/agents/hunter.py @@ -200,10 +200,29 @@ def _correct_line_number(self, file_path: str, line_number: int, code_snippet: s return line_number def _handle_report_finding(self, args: dict) -> dict: - file_path = args.get("file_path", "") + raw_file_path = args.get("file_path") + raw_desc = args.get("description") + + if not isinstance(raw_file_path, str) or not isinstance(raw_desc, str): + logger.warning(f"Finding rejected: missing or non-string file_path/description (args={args})") + return { + "status": "rejected", + "reason": "Missing or invalid required fields: 'file_path' and 'description' must be non-empty strings.", + } + + file_path = raw_file_path.strip() + description = raw_desc.strip() + + # Reject findings missing core required fields + if not file_path or not description: + logger.warning(f"Finding rejected: missing file_path or description (args={args})") + return { + "status": "rejected", + "reason": "Missing required fields: 'file_path' and 'description' must not be empty.", + } # Pre-filter: reject findings in excluded paths - if file_path and self._is_excluded_path(file_path): + if self._is_excluded_path(file_path): logger.info(f"Finding rejected (excluded path): {args.get('category', '?')} in {file_path}") return {"status": "rejected", "reason": f"File is in an excluded path (test/CLI/docs/examples): {file_path}"} @@ -402,13 +421,31 @@ async def run(self, task: str, context: Optional[dict] = None) -> dict: }, ) - if tool_call.name == "report_finding": - result = self._handle_report_finding(tool_call.arguments) + if tool_call.parse_error: + result = { + "error": "tool_call_truncated", + "message": ( + f"Arguments for '{tool_call.name}' were truncated due to token limits " + f"(parse error: {tool_call.parse_error}). " + "Please retry and call the tool with more concise arguments, " + "or report findings one at a time." + ), + } + elif tool_call.name == "report_finding": + try: + result = self._handle_report_finding(tool_call.arguments) + except Exception as e: + logger.warning(f"[{self.name}] Failed to handle report_finding: {e}") + result = {"error": f"Failed to record finding: {e}"} iterations_since_finding = 0 consecutive_reads = 0 analysis_mode = False elif tool_call.name == "finish_hunt": - result = self._handle_finish_hunt(tool_call.arguments) + try: + result = self._handle_finish_hunt(tool_call.arguments) + except Exception as e: + logger.warning(f"[{self.name}] Failed to handle finish_hunt: {e}") + result = {"error": f"Failed to finish hunt: {e}"} should_finish = True else: try: diff --git a/openhack/agents/sandbox_verifier.py b/openhack/agents/sandbox_verifier.py index 66506a9..e579493 100644 --- a/openhack/agents/sandbox_verifier.py +++ b/openhack/agents/sandbox_verifier.py @@ -355,7 +355,16 @@ async def run(self, task: str, context: Optional[dict] = None) -> dict: }, ) - if tool_call.name == "sandbox_http_request": + if tool_call.parse_error: + result = { + "error": "tool_call_truncated", + "message": ( + f"Arguments for '{tool_call.name}' were truncated due to token limits " + f"(parse error: {tool_call.parse_error}). " + "Please retry and call the tool with more concise arguments." + ), + } + elif tool_call.name == "sandbox_http_request": result = await self._handle_sandbox_http_request(tool_call.arguments) elif tool_call.name == "sandbox_multi_step": result = await self._handle_sandbox_multi_step(tool_call.arguments) diff --git a/openhack/agents/validator.py b/openhack/agents/validator.py index 1195c7a..525010c 100644 --- a/openhack/agents/validator.py +++ b/openhack/agents/validator.py @@ -91,7 +91,10 @@ def get_tools(self) -> list[dict]: return super().get_tools() + [VALIDATE_FINDING_TOOL, FINISH_VALIDATION_TOOL] def _handle_validate_finding(self, args: dict) -> dict: - status = args.get("status", "").lower() + raw_status = args.get("status") + if not isinstance(raw_status, str) or not raw_status.strip(): + return {"error": "Missing or invalid required field: 'status' must be a non-empty string."} + status = raw_status.lower().strip() if self.original_finding_index is not None: original_index = self.original_finding_index else: @@ -189,10 +192,27 @@ async def run(self, task: str, context: Optional[dict] = None) -> dict: }, ) - if tool_call.name == "validate_finding": - result = self._handle_validate_finding(tool_call.arguments) + if tool_call.parse_error: + result = { + "error": "tool_call_truncated", + "message": ( + f"Arguments for '{tool_call.name}' were truncated due to token limits " + f"(parse error: {tool_call.parse_error}). " + "Please retry and call the tool with more concise arguments." + ), + } + elif tool_call.name == "validate_finding": + try: + result = self._handle_validate_finding(tool_call.arguments) + except Exception as e: + logger.warning(f"[{self.name}] Failed to handle validate_finding: {e}") + result = {"error": f"Failed to validate finding: {e}"} elif tool_call.name == "finish_validation": - result = self._handle_finish_validation(tool_call.arguments) + try: + result = self._handle_finish_validation(tool_call.arguments) + except Exception as e: + logger.warning(f"[{self.name}] Failed to handle finish_validation: {e}") + result = {"error": f"Failed to finish validation: {e}"} should_finish = True else: result = self.tools.execute_tool(tool_call.name, tool_call.arguments) diff --git a/tests/test_tool_call_truncation.py b/tests/test_tool_call_truncation.py new file mode 100644 index 0000000..3d7fbc2 --- /dev/null +++ b/tests/test_tool_call_truncation.py @@ -0,0 +1,344 @@ +"""Tests for handling truncated tool calls (finish_reason: length) across agents. + +Verifies that when token limits or JSON truncation cause parse_error on tool calls, +agents do not execute the call with empty arguments or record ghost/empty findings. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock + +from openhack.agents.hunter import HunterAgent +from openhack.agents.validator import ValidatorAgent +from openhack.agents.base import BaseAgent +from openhack.agents.sandbox_verifier import SandboxVerifierAgent +from openhack.agents.browser_verifier import BrowserVerifierAgent +from openhack.agents.llm import ToolCall, LLMResponse, Message +from openhack.agents.session import Session +from openhack.tools.registry import ToolRegistry + + +class _TestAgent(BaseAgent): + name = "test-agent" + + def get_system_prompt(self, context): + return "Test." + + +@pytest.fixture +def session(tmp_path): + return Session(str(tmp_path), event_log_path=tmp_path / "events.jsonl") + + +@pytest.fixture +def tools(tmp_path, session): + return ToolRegistry(tmp_path, include_agent_tools=True, session=session) + + +def test_hunter_handle_report_finding_rejects_empty_args(session, tools): + llm = MagicMock() + hunter = HunterAgent(llm, tools, session, vuln_categories=["xss"]) + result = hunter._handle_report_finding({}) + assert result["status"] == "rejected" + assert "Missing" in result["reason"] + assert len(hunter.findings) == 0 + + +def test_hunter_handle_report_finding_rejects_missing_file_or_desc(session, tools): + llm = MagicMock() + hunter = HunterAgent(llm, tools, session, vuln_categories=["xss"]) + # Missing description + res1 = hunter._handle_report_finding({"file_path": "app.py"}) + assert res1["status"] == "rejected" + assert len(hunter.findings) == 0 + + # Missing file_path + res2 = hunter._handle_report_finding({"description": "Some bug"}) + assert res2["status"] == "rejected" + assert len(hunter.findings) == 0 + + +def test_hunter_handle_report_finding_rejects_malformed_non_string_fields(session, tools): + llm = MagicMock() + hunter = HunterAgent(llm, tools, session, vuln_categories=["xss"]) + # Dict description instead of string + res = hunter._handle_report_finding({ + "file_path": "app.py", + "description": {"summary": "SQL injection"}, + }) + assert res["status"] == "rejected" + assert "Missing or invalid required fields" in res["reason"] + assert len(hunter.findings) == 0 + + # List file_path instead of string + res2 = hunter._handle_report_finding({ + "file_path": ["app.py"], + "description": "SQL injection", + }) + assert res2["status"] == "rejected" + assert len(hunter.findings) == 0 + + +def test_hunter_handle_report_finding_accepts_valid_finding(session, tools): + llm = MagicMock() + hunter = HunterAgent(llm, tools, session, vuln_categories=["xss"]) + result = hunter._handle_report_finding({ + "file_path": "app.py", + "description": "Cross-site scripting via query param", + "category": "xss", + }) + assert result["status"] == "recorded" + assert len(hunter.findings) == 1 + assert hunter.findings[0]["file_path"] == "app.py" + + +@pytest.mark.asyncio +async def test_hunter_run_handles_truncated_tool_call(session, tools): + mock_llm = MagicMock() + hunter = HunterAgent(mock_llm, tools, session, vuln_categories=["sql_injection"]) + + # Turn 1: Model emits a truncated report_finding tool call (hit max_tokens) + truncated_tc = ToolCall( + id="call_1", + name="report_finding", + arguments={}, + raw_arguments='{"category": "sql_injection", "file_path": "db.py", "descript', + parse_error="Unterminated string starting at line 1 column 54", + ) + turn1_resp = LLMResponse( + content="", + tool_calls=[truncated_tc], + finish_reason="length", + ) + + # Turn 2: After getting truncated error, model sends complete report_finding and finish_hunt + valid_tc = ToolCall( + id="call_2", + name="report_finding", + arguments={ + "file_path": "db.py", + "description": "SQL injection in query", + "category": "sql_injection", + }, + raw_arguments='{"category": "sql_injection", "file_path": "db.py", "description": "SQL injection in query"}', + parse_error=None, + ) + finish_tc = ToolCall( + id="call_3", + name="finish_hunt", + arguments={}, + raw_arguments="{}", + parse_error=None, + ) + turn2_resp = LLMResponse( + content="Hunt completed", + tool_calls=[valid_tc, finish_tc], + finish_reason="stop", + ) + + mock_llm.chat = AsyncMock(side_effect=[turn1_resp, turn2_resp]) + + await hunter.run("Hunt for SQL injection") + + # Only the valid finding should have been recorded, not the truncated one + assert len(hunter.findings) == 1 + assert hunter.findings[0]["file_path"] == "db.py" + + # Verify that the truncated tool call fed back an error message to the LLM + tool_result_messages = [m for m in hunter.messages if m.role == "tool"] + assert any("tool_call_truncated" in m.content for m in tool_result_messages) + + +def test_validator_handle_validate_finding_rejects_empty_status(session, tools): + llm = MagicMock() + validator = ValidatorAgent(llm, tools, session) + result = validator._handle_validate_finding({}) + assert "error" in result + assert "Missing or invalid required field" in result["error"] + assert len(validator.validated_findings) == 0 + assert len(validator.false_positives) == 0 + + # Non-string status + result2 = validator._handle_validate_finding({"status": {"confirmed": True}}) + assert "error" in result2 + assert len(validator.validated_findings) == 0 + assert len(validator.false_positives) == 0 + + +@pytest.mark.asyncio +async def test_validator_run_handles_truncated_tool_call(session, tools): + mock_llm = MagicMock() + validator = ValidatorAgent(mock_llm, tools, session) + + truncated_tc = ToolCall( + id="call_val_1", + name="validate_finding", + arguments={}, + raw_arguments='{"status": "confirmed", "evid', + parse_error="Unterminated string", + ) + turn1_resp = LLMResponse( + content="", + tool_calls=[truncated_tc], + finish_reason="length", + ) + finish_tc = ToolCall( + id="call_val_2", + name="finish_validation", + arguments={}, + raw_arguments="{}", + parse_error=None, + ) + turn2_resp = LLMResponse( + content="", + tool_calls=[finish_tc], + finish_reason="stop", + ) + + mock_llm.chat = AsyncMock(side_effect=[turn1_resp, turn2_resp]) + + await validator.run("Validate findings") + + # Should not have recorded any false positive due to truncation + assert len(validator.false_positives) == 0 + assert len(validator.validated_findings) == 0 + + +@pytest.mark.asyncio +async def test_base_agent_handles_truncated_tool_call(session, tools): + mock_llm = MagicMock() + agent = _TestAgent(mock_llm, tools, session) + + truncated_tc = ToolCall( + id="call_base_1", + name="read_file", + arguments={}, + raw_arguments='{"path": "secret', + parse_error="Unterminated string", + ) + turn1_resp = LLMResponse( + content="", + tool_calls=[truncated_tc], + finish_reason="length", + ) + finish_tc = ToolCall( + id="call_base_2", + name="finish_task", + arguments={"summary": "done", "reason": "completed", "verification": "checked"}, + raw_arguments='{"summary": "done", "reason": "completed", "verification": "checked"}', + parse_error=None, + ) + turn2_resp = LLMResponse( + content="All done", + tool_calls=[finish_tc], + finish_reason="stop", + ) + + mock_llm.chat = AsyncMock(side_effect=[turn1_resp, turn2_resp]) + + await agent.run("Do something") + + tool_result_messages = [m for m in agent.messages if m.role == "tool"] + assert any("tool_call_truncated" in m.content for m in tool_result_messages) + + +@pytest.mark.asyncio +async def test_sandbox_verifier_handles_truncated_tool_call(session, tools): + mock_llm = MagicMock() + verifier = SandboxVerifierAgent(mock_llm, tools, session, sandbox_url="http://localhost:8080") + + # Turn 1: Truncated sandbox action call + truncated_tc = ToolCall( + id="call_sb_1", + name="sandbox_http_request", + arguments={}, + raw_arguments='{"method": "POST", "path": "/login", "body": "pass', + parse_error="Unterminated string", + ) + turn1_resp = LLMResponse( + content="", + tool_calls=[truncated_tc], + finish_reason="length", + ) + + # Turn 2: Retried completion call + finish_tc = ToolCall( + id="call_sb_2", + name="report_exploit_result", + arguments={ + "status": "exploitable", + "confidence": "high", + "evidence": "Logged in as admin", + "attempts_made": 1, + }, + raw_arguments='{"status": "exploitable", "confidence": "high", "evidence": "Logged in as admin", "attempts_made": 1}', + parse_error=None, + ) + turn2_resp = LLMResponse( + content="Exploit confirmed", + tool_calls=[finish_tc], + finish_reason="stop", + ) + + mock_llm.chat = AsyncMock(side_effect=[turn1_resp, turn2_resp]) + + await verifier.run("Verify finding") + + # Check error message reached model + tool_result_messages = [m for m in verifier.messages if m.role == "tool"] + assert any("tool_call_truncated" in m.content for m in tool_result_messages) + assert verifier.exploit_result is not None + assert verifier.exploit_result["status"] == "exploitable" + + +@pytest.mark.asyncio +async def test_browser_verifier_handles_truncated_tool_call(session, tools): + mock_llm = MagicMock() + mock_runner = MagicMock() + mock_ctx = AsyncMock() + mock_runner.create_context = AsyncMock(return_value=mock_ctx) + verifier = BrowserVerifierAgent( + mock_llm, tools, session, sandbox_url="http://localhost:8080", browser_runner=mock_runner + ) + + # Turn 1: Truncated browser navigation call + truncated_tc = ToolCall( + id="call_br_1", + name="browser_navigate", + arguments={}, + raw_arguments='{"url": "http://localhost:8080/search?q=