From 5da629ef8124fb4d49c167cd4e639521cb1d8c14 Mon Sep 17 00:00:00 2001 From: Ian Date: Wed, 14 Oct 2026 15:21:33 -0400 Subject: [PATCH] feat(coverage): retry open needs_follow_up rows up to 3 times before finish_scan finish_scan used to accept open needs_follow_up rows and only mention them in its success response, after the scan was already complete. It now refuses to finish while a row has retries left, returning each row's entry_id and evidence so the root agent can hand it to a subagent. Each row is capped separately, and the gate steps aside once the budget reserve is hit. --- strix/tools/coverage/tools.py | 30 +++++++++++ strix/tools/finish/tool.py | 29 ++++++++++- tests/test_finish_coverage_gate.py | 83 ++++++++++++++++++++++++++++-- 3 files changed, 137 insertions(+), 5 deletions(-) diff --git a/strix/tools/coverage/tools.py b/strix/tools/coverage/tools.py index 94c31fb9..b804ecbc 100644 --- a/strix/tools/coverage/tools.py +++ b/strix/tools/coverage/tools.py @@ -47,6 +47,8 @@ VALID_OUTCOMES: tuple[str, ...] = ( _OUTCOMES_REQUIRING_EVIDENCE = frozenset({"ruled_out", "not_applicable", "needs_follow_up"}) +MAX_FOLLOW_UP_RETRIES = 3 + def _caller_identity(ctx: RunContextWrapper) -> tuple[str | None, str | None]: """Return the (agent_id, agent_name) of the agent invoking this tool.""" @@ -149,6 +151,34 @@ def outcome_counts() -> dict[str, int]: return {o: counts[o] for o in VALID_OUTCOMES if o in counts} +def claim_follow_up_retries(max_retries: int = MAX_FOLLOW_UP_RETRIES) -> list[dict[str, Any]]: + """Spend one retry on every open ``needs_follow_up`` row that has one left. + + Called by the finish tools, which refuse to finish while this returns rows. + Each row is capped separately, so a row that cannot be resolved here (no + credentials, wrong platform) stops blocking without starving the others. + """ + claimed: list[dict[str, Any]] = [] + with _coverage_lock: + for entry_id, entry in _coverage_storage.items(): + retries = int(entry.get("follow_up_retries", 0)) + if entry.get("outcome") != "needs_follow_up" or retries >= max_retries: + continue + entry["follow_up_retries"] = retries + 1 + claimed.append( + { + "entry_id": entry_id, + "surface": entry.get("surface", ""), + "risk_area": entry.get("risk_area", ""), + "evidence": entry.get("evidence", ""), + "retry": retries + 1, + } + ) + if claimed: + _persist_locked() + return claimed + + def _validate( *, surface: str, risk_area: str, outcome: str, evidence: str ) -> tuple[str, list[str]]: diff --git a/strix/tools/finish/tool.py b/strix/tools/finish/tool.py index e3dcb7ae..e589f8f1 100644 --- a/strix/tools/finish/tool.py +++ b/strix/tools/finish/tool.py @@ -203,8 +203,10 @@ async def finish_scan( should have a coverage entry; anything still open should be a ``needs_follow_up`` row, not a silent omission. If a significant surface has no entry at all, dispatch an agent to cover it or - record it as ``needs_follow_up`` before finishing. The response - from this tool reports coverage counts and any unresolved rows. + record it as ``needs_follow_up`` before finishing. Open + ``needs_follow_up`` rows are not a resting state: this tool refuses + to finish until each one has been retried up to 3 times. The + response from this tool reports coverage counts and any unresolved rows. **Calling this multiple times overwrites the previous report.** Make the single call comprehensive. @@ -334,6 +336,29 @@ async def finish_scan( default=str, ) + if parent_id is None and not (coordinator is not None and coordinator.reserve_stopped): + from strix.tools.coverage.tools import MAX_FOLLOW_UP_RETRIES, claim_follow_up_retries + + follow_ups = await asyncio.to_thread(claim_follow_up_retries) + if follow_ups: + return json.dumps( + { + "success": False, + "scan_completed": False, + "error": ( + f"{len(follow_ups)} coverage row(s) are still needs_follow_up. Each " + f"gets up to {MAX_FOLLOW_UP_RETRIES} retries before the scan can " + "finish. Spawn a subagent for each row with its entry_id and " + "evidence, and have it attack the gap the evidence names with a " + "different method than the one that stalled, then record the " + "result with update_coverage. Call finish_scan again afterwards." + ), + "follow_ups": follow_ups, + }, + ensure_ascii=False, + default=str, + ) + result = await asyncio.to_thread( _do_finish, parent_id=parent_id, diff --git a/tests/test_finish_coverage_gate.py b/tests/test_finish_coverage_gate.py index 52a19a7d..324bab85 100644 --- a/tests/test_finish_coverage_gate.py +++ b/tests/test_finish_coverage_gate.py @@ -2,12 +2,20 @@ from __future__ import annotations -from typing import TYPE_CHECKING +import json +from typing import TYPE_CHECKING, Any import pytest +from agents.tool_context import ToolContext -from strix.tools.coverage.tools import _record_impl, hydrate_coverage_from_disk -from strix.tools.finish.tool import _coverage_summary +from strix.core.agents import AgentCoordinator +from strix.tools.coverage.tools import ( + MAX_FOLLOW_UP_RETRIES, + _record_impl, + _update_impl, + hydrate_coverage_from_disk, +) +from strix.tools.finish.tool import _coverage_summary, finish_scan if TYPE_CHECKING: @@ -63,3 +71,72 @@ def test_an_empty_ledger_still_warns_first() -> None: assert summary["coverage_recorded"] == 0 assert "No coverage was recorded" in summary["coverage_warning"] + + +async def _finish(coordinator: AgentCoordinator) -> dict[str, Any]: + ctx = ToolContext( + context={"coordinator": coordinator, "agent_id": "root", "parent_id": None}, + tool_name="finish_scan", + tool_call_id="call-1", + tool_arguments="{}", + ) + fields = ("executive_summary", "methodology", "technical_analysis", "recommendations") + result: dict[str, Any] = json.loads( + await finish_scan.on_invoke_tool(ctx, json.dumps(dict.fromkeys(fields, "x"))) + ) + return result + + +def _record_follow_up(surface: str) -> None: + _record_impl( + surface=surface, + risk_area="RCE", + outcome="needs_follow_up", + evidence="Could not confirm which constructor fields resolve imports.", + agent_id="agent-1", + agent_name="injection-tester", + ) + + +@pytest.mark.asyncio +async def test_open_follow_ups_block_finish_until_each_has_used_its_retries() -> None: + coordinator = AgentCoordinator() + await coordinator.register("root", "strix", parent_id=None) + _record_follow_up("load_agent_from_repository") + + for retry in range(1, MAX_FOLLOW_UP_RETRIES + 1): + blocked = await _finish(coordinator) + assert blocked["scan_completed"] is False + assert blocked["follow_ups"][0]["retry"] == retry + assert blocked["follow_ups"][0]["evidence"] + + finished = await _finish(coordinator) + assert finished["scan_completed"] is True + + +@pytest.mark.asyncio +async def test_a_resolved_follow_up_stops_blocking() -> None: + coordinator = AgentCoordinator() + await coordinator.register("root", "strix", parent_id=None) + _record_follow_up("load_agent_from_repository") + + blocked = await _finish(coordinator) + _update_impl( + entry_id=blocked["follow_ups"][0]["entry_id"], + outcome="reported", + evidence="ImportString field resolves os.system.", + agent_id="agent-2", + agent_name="follow-up", + ) + + assert (await _finish(coordinator))["scan_completed"] is True + + +@pytest.mark.asyncio +async def test_follow_ups_do_not_block_once_the_budget_reserve_is_hit() -> None: + coordinator = AgentCoordinator() + await coordinator.register("root", "strix", parent_id=None) + _record_follow_up("load_agent_from_repository") + await coordinator.claim_reserve_notification() + + assert (await _finish(coordinator))["scan_completed"] is True