diff --git a/src/openultrasast/tool_hunter.py b/src/openultrasast/tool_hunter.py index 9de7d6f..59db50f 100644 --- a/src/openultrasast/tool_hunter.py +++ b/src/openultrasast/tool_hunter.py @@ -168,6 +168,7 @@ class ChatClient(Protocol): messages: list[dict[str, object]], tools: list[dict[str, object]], timeout_seconds: int = 60, + json_object: bool = False, # the final answer turn asks for JSON; adapters without it ignore the flag ) -> ChatResponse: raise NotImplementedError @@ -216,7 +217,9 @@ class OpenRouterHunterClient: messages: list[dict[str, object]], tools: list[dict[str, object]], timeout_seconds: int = 60, + json_object: bool = False, ) -> ChatResponse: + del json_object # OpenRouter routes to many providers; the prompt asks for JSON and the parser tolerates prose message = self._client.complete_chat(model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds) return chat_response_from_message(message) @@ -343,8 +346,54 @@ def run_tool_hunter( if not content.strip() or not used_tools: # A findings dump is not evidence unless the loop invoked a repo tool. return [] - return _findings_from_content(root, hotspots, content, tags) - return [] + findings = _findings_from_content(root, hotspots, content, tags) + if findings or _looks_like_json(content): + return findings + return _final_answer(client, model, messages, schemas, content, root, hotspots, tags) + return _final_answer(client, model, messages, schemas, "", root, hotspots, tags) if used_tools else [] + + +def _final_answer( + client: ChatClient, + model: str, + messages: list[dict[str, object]], + schemas: list[dict[str, object]], + prose: str, + root: Path, + hotspots: Sequence[Hotspot], + tags: Sequence[str], +) -> list[StaticFinding]: + """One more turn, asking the provider for JSON, when the hunt ended in prose or ran out of steps. + + Measured live against DeepSeek: the model spends its step budget on tools and then answers in a paragraph that + names the sink, the argument and the line — which parses as nothing at all. A detector that found the bug and + reported zero is indistinguishable from one that found nothing, and every family would baseline at zero for + that reason alone. This turn formats an answer the model has already reached; it never asks for a new one, and + a hunt that called no tool never gets here. + """ + body = prose.strip() + tail: list[dict[str, object]] = [{"role": "assistant", "content": body}] if body else [] + messages = [ + *messages, + *tail, + { + "role": "user", + "content": ( + "Answer now, as JSON only: an array of objects with path, line, title, rationale and family. " + "Report only what you established with the tools; return [] if that is nothing." + ), + }, + ] + response = client.complete(model=model, messages=messages, tools=schemas, json_object=True) + return _findings_from_content(root, hotspots, response.content or "", tags) + + +def _looks_like_json(content: str) -> bool: + stripped = content.strip() + fenced = re.fullmatch(r"```(?:json)?\s*(.*?)\s*```", stripped, flags=re.DOTALL) + if fenced is not None: + stripped = fenced.group(1).strip() + return stripped.startswith(("[", "{")) def _tool_schemas(names: Sequence[str] | None) -> list[dict[str, object]]: diff --git a/tests/test_hunter_curated_tools.py b/tests/test_hunter_curated_tools.py index 1b373dd..838e90c 100644 --- a/tests/test_hunter_curated_tools.py +++ b/tests/test_hunter_curated_tools.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import tempfile from pathlib import Path import pytest @@ -151,3 +152,45 @@ def test_the_loop_takes_a_prompt_context_tools_and_tags_and_replays_reasoning(re tool_reply = json.loads(str(client.calls[1]["messages"][-1]["content"])) assert tool_reply and tool_reply[0]["missing"] == "identity_constraint" assert findings[0].tags == ["family:access_control", "detector:access_control@1"] + + +def test_the_loop_asks_once_more_for_the_answer_instead_of_dropping_prose() -> None: + """Measured live: the model spends its step budget on tools and then answers in prose, and the whole hunt is + discarded (learning-harness Req 6.3). + + Four steps, three of them tool calls, and the fourth a paragraph that names the sink, the argument and the + line — parsed as nothing. A detector that found the bug and reported zero is indistinguishable from one that + found nothing, and every family in round zero would have baselined at zero for this reason alone.""" + from openultrasast.tool_hunter import ChatResponse, ScriptedChatClient, ToolCall, run_tool_hunter + + prose = "The function passes `command` straight to os.system at line 2, an OS command sink." + answer = json.dumps([{"path": "app.py", "line": 2, "title": "command injection", "rationale": "os.system(command)"}]) + client = ScriptedChatClient( + [ + ChatResponse(tool_calls=(ToolCall(id="c1", name="grep_repo", arguments={"pattern": "system"}),)), + ChatResponse(tool_calls=(ToolCall(id="c2", name="read_file", arguments={"path": "app.py"}),)), + ChatResponse(content=prose), + ChatResponse(content=answer), # the one extra turn, asked for explicitly + ] + ) + with tempfile.TemporaryDirectory() as scratch: + root = Path(scratch) + (root / "app.py").write_text("import os\ndef run(command):\n return os.system(command)\n") + findings = run_tool_hunter(root, [], client=client, model="m", max_steps=3) + assert [item.line for item in findings] == [2] + final = client.calls[-1] + assert final.get("json_object") is True, "the last turn asks the provider for JSON, it does not hope for it" + text = "\n".join(str(message.get("content") or "") for message in final["messages"]) + assert "JSON" in text and prose in text, "the prose it already wrote is what the final turn is asked to format" + + +def test_a_hunt_that_never_called_a_tool_still_reports_nothing() -> None: + """The extra turn is for formatting an answer, never for inventing one.""" + from openultrasast.tool_hunter import ChatResponse, ScriptedChatClient, run_tool_hunter + + client = ScriptedChatClient([ChatResponse(content="I think there is an injection here."), ChatResponse(content="[]")]) + with tempfile.TemporaryDirectory() as scratch: + root = Path(scratch) + (root / "app.py").write_text("x = 1\n") + assert run_tool_hunter(root, [], client=client, model="m", max_steps=3) == [] + assert len(client.calls) == 1, "no tool was called, so there is nothing to format" diff --git a/tests/test_tool_hunter.py b/tests/test_tool_hunter.py index 826da01..aac1693 100644 --- a/tests/test_tool_hunter.py +++ b/tests/test_tool_hunter.py @@ -133,7 +133,13 @@ def test_path_escape_stays_in_band_and_still_counts_as_tool_use(tmp_path: Path) assert "error" in json.dumps(client.calls[1]["messages"]).lower() -def test_max_steps_stops_before_ungrounded_json(tmp_path: Path) -> None: +def test_a_budget_that_runs_out_mid_hunt_still_asks_for_the_answer(tmp_path: Path) -> None: + """The step budget bounds the *investigation*, not the report. + + This used to discard everything when the last step was a tool call: the loop had a tool result in hand and + threw it away because it had no turn left to ask about it. What the budget must still refuse is a findings + dump from a hunt that called no tool at all, which `test_hunter_dump_only_findings_are_rejected` pins. + """ root = _repo_with_split_sink(tmp_path) client = ScriptedChatClient( [ @@ -144,8 +150,8 @@ def test_max_steps_stops_before_ungrounded_json(tmp_path: Path) -> None: findings = run_tool_hunter(root, [_hotspot()], client=client, model="test-hunter", max_steps=1) - assert findings == [] - assert len(client.calls) == 1 + assert [item.line for item in findings] == [item["line"] for item in _DUMP_FINDINGS] + assert len(client.calls) == 2 and client.calls[-1].get("json_object") is True _UNSAFE_SNIPPET = 'client = docker.DockerClient(base_url="unix://var/run/docker.sock")\n'