From 618e676928d0e9c8cf7a297a73ae186526f7e279 Mon Sep 17 00:00:00 2001 From: Guy Basson Date: Mon, 12 Oct 2026 21:07:07 +0300 Subject: [PATCH] fix(agents): defer saved completion and deliver sessionless wake messages --- strix/core/execution.py | 10 +- strix/report/state.py | 18 ++- strix/tools/finish/tool.py | 13 +- strix/tools/wait_for_user/tool.py | 3 +- tests/test_mailbox_review_regressions.py | 185 +++++++++++++++++++++++ 5 files changed, 223 insertions(+), 6 deletions(-) create mode 100644 tests/test_mailbox_review_regressions.py diff --git a/strix/core/execution.py b/strix/core/execution.py index ae9b8fa5..9f56cbc3 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -393,13 +393,19 @@ async def _run_agent_loop( interrupt=False, ) - await coordinator.consume_pending(agent_id) + _, items = await coordinator.consume_pending(agent_id, include_items=True) + history = ( + result.to_input_list() if result is not None and session is None else first_cycle_input + ) + next_input = ( + [] if session is not None else ItemHelpers.input_to_new_input_list(history) + items + ) with contextlib.suppress(BudgetPausedError): result = await _run_until_lifecycle( agent, coordinator, agent_id, - initial_input=[], + initial_input=next_input, run_config=run_config, context=context, max_turns=max_turns, diff --git a/strix/report/state.py b/strix/report/state.py index fd0ebd94..4b6c22d7 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -718,20 +718,34 @@ class ReportState: methodology: str, technical_analysis: str, recommendations: str, + *, + mark_complete: bool = True, ) -> None: + """Persist report prose, optionally deferring completion until the inbox is settled.""" self.scan_results = { - "scan_completed": True, + "scan_completed": False, "executive_summary": executive_summary.strip(), "methodology": methodology.strip(), "technical_analysis": technical_analysis.strip(), "recommendations": recommendations.strip(), - "success": True, + "success": False, } self.final_scan_result = self._format_final_scan_result(self.scan_results) self.run_record["scan_results"] = self.scan_results logger.info("Updated scan final fields") + if mark_complete: + self.complete_scan() + else: + self.save_run_data() + + def complete_scan(self) -> None: + """Commit an accepted finish and emit completion telemetry only then.""" + if self.scan_results is None: + raise ValueError("Cannot complete a scan without a report summary") + self.scan_results["scan_completed"] = True + self.scan_results["success"] = True self.save_run_data(mark_complete=True) posthog.end(self, exit_reason="finished_by_tool") scarf.end(self, exit_reason="finished_by_tool") diff --git a/strix/tools/finish/tool.py b/strix/tools/finish/tool.py index 0148c2a6..ac727d31 100644 --- a/strix/tools/finish/tool.py +++ b/strix/tools/finish/tool.py @@ -62,6 +62,7 @@ def _do_finish( methodology=methodology.strip(), technical_analysis=technical_analysis.strip(), recommendations=recommendations.strip(), + mark_complete=False, ) vuln_count = len(report_state.vulnerability_reports) coverage_summary = _coverage_summary(agent_graph) @@ -70,7 +71,7 @@ def _do_finish( return {"success": False, "error": f"Failed to complete scan: {e!s}"} else: logger.info( - "finish_scan: completed scan with %d vulnerability report(s)", + "finish_scan: prepared summary with %d vulnerability report(s)", vuln_count, ) result: dict[str, Any] = { @@ -83,6 +84,14 @@ def _do_finish( return result +def _commit_finish() -> None: + from strix.report.state import get_global_report_state + + report_state = get_global_report_state() + if report_state is not None: + report_state.complete_scan() + + def _coverage_summary(agent_graph: dict[str, Any]) -> dict[str, Any]: """Coverage counts, unresolved surfaces, and gaps the runtime can see. @@ -362,4 +371,6 @@ async def finish_scan( # A message arrived while the draft summary was being persisted. # Keep running so the root can reconcile it and finalize again. return json.dumps(pending_error) + if result.get("success") and result.get("scan_completed"): + await asyncio.to_thread(_commit_finish) return json.dumps(result, ensure_ascii=False, default=str) diff --git a/strix/tools/wait_for_user/tool.py b/strix/tools/wait_for_user/tool.py index 970acb9d..94452c6f 100644 --- a/strix/tools/wait_for_user/tool.py +++ b/strix/tools/wait_for_user/tool.py @@ -76,7 +76,7 @@ async def wait_for_user(ctx: RunContextWrapper) -> str: # A message that arrived while this turn was running is the user already # talking: take it now instead of parking for one they have sent. - pending, _ = await coordinator.consume_pending(me) + pending, items = await coordinator.consume_pending(me, include_items=True) if pending > 0: await coordinator.mark_running(me) return json.dumps( @@ -84,6 +84,7 @@ async def wait_for_user(ctx: RunContextWrapper) -> str: "success": True, "wait_outcome": "message_arrived", "pending_messages": pending, + "messages": items, "note": "The user had already sent a new message; keep going.", }, ensure_ascii=False, diff --git a/tests/test_mailbox_review_regressions.py b/tests/test_mailbox_review_regressions.py new file mode 100644 index 00000000..0081582c --- /dev/null +++ b/tests/test_mailbox_review_regressions.py @@ -0,0 +1,185 @@ +"""Offline regressions for completion persistence and interactive mailbox delivery.""" + +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any +from unittest.mock import Mock + +import pytest +from agents import RunConfig +from agents.memory import SQLiteSession +from agents.tool_context import ToolContext + +from strix.core import execution +from strix.core.agents import AgentCoordinator +from strix.report import state as state_module +from strix.report.state import ReportState +from strix.tools.finish.tool import finish_scan +from strix.tools.wait_for_user.tool import wait_for_user + + +if TYPE_CHECKING: + from pathlib import Path + + +async def _call(tool: Any, context: dict[str, Any], arguments: dict[str, Any]) -> dict[str, Any]: + raw = json.dumps(arguments) + ctx = ToolContext(context=context, tool_name=tool.name, tool_call_id="call", tool_arguments=raw) + return json.loads(await tool.on_invoke_tool(ctx, raw)) + + +@pytest.mark.parametrize("after_rejection", ["retry", "failed", "interrupted"]) +async def test_rejected_finish_does_not_persist_completion_or_send_telemetry( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, after_rejection: str +) -> None: + coordinator = AgentCoordinator() + await coordinator.register("root", "Root", parent_id=None) + await coordinator.attach_runtime("root", resumable=False) + state = ReportState("offline-finish-race") + state._run_dir = tmp_path + monkeypatch.setattr(state_module, "get_global_report_state", lambda: state) + posthog_end, scarf_end = Mock(), Mock() + monkeypatch.setattr(state_module.posthog, "end", posthog_end) + monkeypatch.setattr(state_module.scarf, "end", scarf_end) + loop = asyncio.get_running_loop() + original_save = state.save_run_data + inject = True + + def save(*args: Any, **kwargs: Any) -> None: + nonlocal inject + original_save(*args, **kwargs) + if inject: + inject = False + accepted = asyncio.run_coroutine_threadsafe( + coordinator.send("root", {"from": "reviewer", "content": "Recheck summary"}), + loop, + ).result(timeout=5) + assert accepted + + monkeypatch.setattr(state, "save_run_data", save) + context = {"coordinator": coordinator, "agent_id": "root", "parent_id": None} + arguments = dict.fromkeys( + ("executive_summary", "methodology", "technical_analysis", "recommendations"), "Reviewed" + ) + result = await _call(finish_scan, context, arguments) + assert result["success"] is False + assert result["scan_completed"] is False + record = json.loads((tmp_path / "run.json").read_text()) + assert record["status"] == "running" + assert record["end_time"] is None + assert record["scan_results"]["scan_completed"] is False + assert coordinator.statuses["root"] == "running" + posthog_end.assert_not_called() + scarf_end.assert_not_called() + + if after_rejection == "retry": + await coordinator.consume_pending("root", include_items=True) + result = await _call(finish_scan, context, arguments) + assert result["success"] and result["scan_completed"] + assert coordinator.statuses["root"] == "completed" + record = json.loads((tmp_path / "run.json").read_text()) + assert record["status"] == "completed" + assert record["scan_results"]["scan_completed"] is True + assert record["end_time"] is not None + posthog_end.assert_called_once_with(state, exit_reason="finished_by_tool") + scarf_end.assert_called_once_with(state, exit_reason="finished_by_tool") + else: + state.save_run_data(status=after_rejection) + record = json.loads((tmp_path / "run.json").read_text()) + assert record["status"] == after_rejection + assert record["scan_results"]["scan_completed"] is False + posthog_end.assert_not_called() + scarf_end.assert_not_called() + + +async def test_wait_for_user_delivers_sessionless_messages_in_tool_result() -> None: + coordinator = AgentCoordinator() + await coordinator.register("root", "Root", parent_id=None) + await coordinator.send("root", {"from": "user", "content": "Recheck the result"}) + result = await _call( + wait_for_user, + {"coordinator": coordinator, "agent_id": "root", "interactive": True}, + {}, + ) + assert result["wait_outcome"] == "message_arrived" + assert result["messages"] == [{"role": "user", "content": "Recheck the result"}] + assert not await coordinator.has_pending("root") + assert coordinator.statuses["root"] == "running" + + +@pytest.mark.parametrize("with_session", [False, True]) +@pytest.mark.parametrize("start_parked", [False, True]) +async def test_interactive_wake_delivers_messages_once_and_waits_again( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, with_session: bool, start_parked: bool +) -> None: + coordinator = AgentCoordinator() + await coordinator.register("root", "Root", parent_id=None) + session = SQLiteSession("root", tmp_path / "agents.db") if with_session else None + history = [{"role": "user", "content": "Original task"}] + calls: list[Any] = [] + woke = asyncio.Event() + waiting_again = asyncio.Event() + waits = 0 + original_wait = coordinator.wait_for_message + + async def wait(*args: Any, **kwargs: Any) -> bool: + nonlocal waits + waits += 1 + if waits == 2: + waiting_again.set() + assert waits <= 2, "An acknowledged message must not wake the agent again" + return await original_wait(*args, **kwargs) + + async def cycle(*_args: Any, **kwargs: Any) -> Any: + items = await session.get_items() if session is not None else kwargs["initial_input"] + calls.append(items) + is_wake = start_parked or len(calls) > 1 + if is_wake: + rendered = json.dumps(items) + assert rendered.count("New instruction") == 1 + assert rendered.count("Original task") == 1 + assert not await coordinator.has_pending("root") + woke.set() + await coordinator.park_waiting("root", wait_kind="user") + return SimpleNamespace(to_input_list=lambda: history) + + monkeypatch.setattr(execution, "_run_until_lifecycle", cycle) + monkeypatch.setattr(coordinator, "wait_for_message", wait) + await coordinator.send("root", {"from": "user", "content": "New instruction"}) + # A restored session already holds its opening task when starting parked. + if start_parked and session is not None: + await session.add_items(history) + task = asyncio.create_task( + execution.run_agent_loop( + agent=SimpleNamespace(name="Root"), + initial_input=history, + run_config=RunConfig(tracing_disabled=True), + context={"agent_id": "root"}, + max_turns=4, + coordinator=coordinator, + agent_id="root", + interactive=True, + session=session, + start_parked=start_parked, + ) + ) + try: + done, _ = await asyncio.wait( + [task, waiter := asyncio.create_task(waiting_again.wait())], + timeout=5, + return_when=asyncio.FIRST_COMPLETED, + ) + if task in done: + await task + assert waiting_again.is_set() and woke.is_set() + assert not task.done() + assert len(calls) == (1 if start_parked else 2) + finally: + task.cancel() + waiter.cancel() + await asyncio.gather(task, waiter, return_exceptions=True) + if session is not None: + session.close()