From 7fdf53d9f4f71ab64023ff8fa96a5001669cca2e Mon Sep 17 00:00:00 2001 From: root Date: Thu, 19 Mar 2026 09:57:35 +0100 Subject: [PATCH] Add _build_resume_context_message to prevent LLM from finishing on resume Defines the missing helper function called in cli.py and adds the equivalent to tui.py. Injects a user message into the restored AgentState so the LLM knows the scan was interrupted and must continue rather than call finish_scan or agent_finish due to an abruptly-ended message history. Co-Authored-By: Claude Sonnet 4.6 --- strix/interface/cli.py | 30 +++++++++++++++++++++++++++--- strix/interface/tui.py | 22 ++++++++++++++++++++++ 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/strix/interface/cli.py b/strix/interface/cli.py index 260acd65..ad7b0a34 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -97,6 +97,26 @@ def _replay_previous_output( console.print() +def _build_resume_context_message(state: Any, iteration: int) -> None: + """Inject a user message telling the LLM it was interrupted and must continue. + + Added for Resume Feature — prevents the model from calling finish_scan or + agent_finish just because the message history ends abruptly (e.g. a dangling + sub-agent tool call that never got a result). + """ + msg = ( + f"[SYSTEM - SCAN RESUMED]\n" + f"This penetration test was interrupted at iteration {iteration}. " + f"All sub-agents that were running have been terminated along with their sandbox. " + f"A fresh sandbox will be created automatically. " + f"Review the conversation history above to understand what has already been done, " + f"then CONTINUE the penetration test from where it left off. " + f"Do NOT call finish_scan or agent_finish unless all testing is genuinely complete. " + f"Re-spawn any sub-agents needed to continue uncompleted work." + ) + state.add_message("user", msg) + + async def run_cli(args: Any) -> None: # noqa: PLR0915 console = Console() @@ -125,15 +145,19 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 resumed_state.sandbox_token = None resumed_state.sandbox_info = None - # Fix: reset any blocking flags that were set when the scan was interrupted. - # If the scan was paused mid-wait the restored state would still have - # waiting_for_input=True / stop_requested=True and the loop would freeze. + # Reset any blocking flags set at the moment of interruption resumed_state.waiting_for_input = False resumed_state.waiting_start_time = None resumed_state.stop_requested = False resumed_state.completed = False resumed_state.llm_failed = False + # Inject a resume-context message so the LLM does NOT call finish_scan + # or agent_finish just because the history ended abruptly. + # Without this the model sees a dangling tool call (sub-agent that was + # killed mid-execution) and may decide the task is complete or broken. + _build_resume_context_message(resumed_state, checkpoint_data.iteration) + start_text = Text() if is_resuming: start_text.append("Penetration test resumed", style="bold #22c55e") diff --git a/strix/interface/tui.py b/strix/interface/tui.py index 157705cf..4580feac 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -41,6 +41,25 @@ from strix.telemetry.tracer import Tracer, set_global_tracer logger = logging.getLogger(__name__) +def _inject_resume_context_message(state: Any, iteration: int) -> None: + """Inject a user message telling the LLM it was interrupted and must continue. + + Added for Resume Feature — prevents the model from calling finish_scan or + agent_finish just because the message history ends abruptly. + """ + msg = ( + f"[SYSTEM - SCAN RESUMED]\n" + f"This penetration test was interrupted at iteration {iteration}. " + f"All sub-agents that were running have been terminated along with their sandbox. " + f"A fresh sandbox will be created automatically. " + f"Review the conversation history above to understand what has already been done, " + f"then CONTINUE the penetration test from where it left off. " + f"Do NOT call finish_scan or agent_finish unless all testing is genuinely complete. " + f"Re-spawn any sub-agents needed to continue uncompleted work." + ) + state.add_message("user", msg) + + def get_package_version() -> str: try: return pkg_version("strix-agent") @@ -790,6 +809,9 @@ class StrixTUIApp(App): # type: ignore[misc] resumed_state.stop_requested = False resumed_state.completed = False resumed_state.llm_failed = False + # Inject resume-context message so the LLM does NOT call finish_scan + # or agent_finish just because the history ended abruptly. + _inject_resume_context_message(resumed_state, _cp.iteration) config["state"] = resumed_state _mgr = getattr(args, "_checkpoint_manager", None)