fix(agents): defer saved completion and deliver sessionless wake messages

This commit is contained in:
Guy Basson 2026-10-12 21:07:07 +03:00
parent 6e56b4e329
commit 618e676928
5 changed files with 223 additions and 6 deletions

View file

@ -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,

View file

@ -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")

View file

@ -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)

View file

@ -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,

View file

@ -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()