mirror of
https://github.com/usestrix/strix.git
synced 2026-10-16 04:28:23 +00:00
fix(agents): defer saved completion and deliver sessionless wake messages
This commit is contained in:
parent
6e56b4e329
commit
618e676928
5 changed files with 223 additions and 6 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
185
tests/test_mailbox_review_regressions.py
Normal file
185
tests/test_mailbox_review_regressions.py
Normal 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()
|
||||
Loading…
Add table
Reference in a new issue