diff --git a/strix/core/agents.py b/strix/core/agents.py index 7e738be1..59ab031f 100644 --- a/strix/core/agents.py +++ b/strix/core/agents.py @@ -50,6 +50,7 @@ class AgentRuntime: resumable: bool = True wake: asyncio.Event = field(default_factory=asyncio.Event) mailbox: list[dict[str, Any]] = field(default_factory=list) + mailbox_lock: asyncio.Lock = field(default_factory=asyncio.Lock) user_wake_required: bool = False @@ -373,6 +374,50 @@ class AgentCoordinator: logger.info("agent.status %s=%s", agent_id, status) await self._maybe_snapshot() + async def has_pending(self, agent_id: str) -> bool: + """Whether accepted messages still need to enter the agent's context.""" + async with self._lock: + runtime = self.runtimes.get(agent_id) + return self.pending_counts.get(agent_id, 0) > 0 or bool(runtime and runtime.mailbox) + + async def complete_if_no_pending( + self, agent_id: str, *, parent_message: dict[str, Any] | None = None + ) -> tuple[bool, bool]: + """Atomically settle an empty inbox and queue the parent's completion notice. + + Return (completed, parent_notified). The parent must never observe a + completed child before that child's report is available in its inbox. + """ + stream = None + notified = False + async with self._lock: + runtime = self.runtimes.get(agent_id) + if agent_id not in self.statuses: + return False, False + if self.pending_counts.get(agent_id, 0) > 0 or (runtime and runtime.mailbox): + return False, False + parent = self.parent_of.get(agent_id) + if ( + parent_message is not None + and parent is not None + and parent in self.statuses + and agent_id not in self._parent_notified + and not self._unreachable_locked(parent) + ): + recipient = self.runtimes.setdefault(parent, AgentRuntime()) + recipient.mailbox.append(dict(parent_message)) + self.pending_counts[parent] = self.pending_counts.get(parent, 0) + 1 + recipient.wake.set() + self._parent_notified.add(agent_id) + notified = True + if recipient.interrupt_on_message: + stream = recipient.stream + self._set_status_locked(agent_id, "completed") + if stream is not None: + stream.cancel(mode="immediate") + await self._maybe_snapshot() + return True, notified + def _set_status_locked( self, agent_id: str, status: Status | str, *, error: str | None = None ) -> None: @@ -483,38 +528,40 @@ class AgentCoordinator: *, include_items: bool = False, ) -> tuple[int, list[Any]]: - """Drain the agent's mailbox into its own SDK session.""" + """Persist a mailbox batch before acknowledging it; serialize concurrent drains. + + With no session, callers must request the items and insert them into + their next input/tool result themselves. Failed writes leave the batch + queued, including its order relative to messages accepted during I/O. + """ async with self._lock: runtime = self.runtimes.setdefault(agent_id, AgentRuntime()) - queued = list(runtime.mailbox) - runtime.mailbox.clear() - count = max(self.pending_counts.get(agent_id, 0), len(queued)) - self.pending_counts[agent_id] = 0 - session = runtime.session - if count <= 0: - return 0, [] - items = [self._message_to_session_item(m) for m in queued] - if items: - if session is None: - logger.warning( - "agent %s has no SDK session attached; %d queued messages were not persisted", - agent_id, - len(items), - ) - else: - try: - async with session_write_lock(session): - await session.add_items(items) - except Exception: - logger.exception( - "failed to append %d queued messages to the session of %s", - len(items), - agent_id, - ) - await self._maybe_snapshot() - if not include_items: - return count, [] - return count, items + async with runtime.mailbox_lock: + async with self._lock: + queued = list(runtime.mailbox) + count = max(self.pending_counts.get(agent_id, 0), len(queued)) + session = runtime.session + if count <= 0: + return 0, [] + items = [self._message_to_session_item(m) for m in queued] + if items: + if session is None and not include_items: + logger.warning("agent %s has no session; retaining queued messages", agent_id) + return 0, [] + if session is not None: + try: + async with session_write_lock(session): + await session.add_items(items) + except Exception: + logger.exception( + "mailbox write failed for %s; retaining messages", agent_id + ) + return 0, [] + async with self._lock: + del runtime.mailbox[: len(queued)] + self.pending_counts[agent_id] = max(0, self.pending_counts.get(agent_id, 0) - count) + await self._maybe_snapshot() + return count, items if include_items else [] async def request_stop(self, agent_id: str) -> None: async with self._lock: diff --git a/strix/core/execution.py b/strix/core/execution.py index 87a2bf02..9f56cbc3 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -9,6 +9,7 @@ import logging import re import uuid from collections.abc import Callable +from copy import copy from dataclasses import replace from functools import cache from typing import TYPE_CHECKING, Any, cast @@ -159,6 +160,32 @@ def _with_image_scrub(run_config: RunConfig, context: dict[str, Any]) -> RunConf return replace(run_config, call_model_input_filter=_scrub) +class _PendingMessagesError(Exception): + """Yield at an SDK model boundary to reload the persisted conversation.""" + + +def _with_mailbox_checkpoint( + run_config: RunConfig, coordinator: AgentCoordinator, agent_id: str +) -> RunConfig: + # A streamed SDK run can contain many model/tool turns. Updating only its + # session does not update the in-memory input for the next model request. + # Restart at this boundary, after preceding tool results have been saved, + # rather than cancelling an in-flight tool or injecting one ephemeral turn. + inner = run_config.call_model_input_filter + + async def _checkpoint(data: CallModelData[Any]) -> ModelInputData: + if await coordinator.has_pending(agent_id): + raise _PendingMessagesError + if inner is None: + return data.model_data + result = inner(data) + return await result if inspect.isawaitable(result) else result + + configured = copy(run_config) + configured.call_model_input_filter = _checkpoint + return configured + + _MAX_TRANSIENT_MODEL_RETRIES = 5 _TRANSIENT_MODEL_RETRY_BASE_DELAY_S = 2.0 _TRANSIENT_MODEL_RETRY_MAX_DELAY_S = 90.0 @@ -366,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, @@ -770,6 +803,8 @@ async def _run_cycle( # noqa: PLR0912, PLR0915 compactions = 0 model_retries = 0 request_log.set_retry_attempt(0) + if not interactive: + run_config = _with_mailbox_checkpoint(run_config, coordinator, agent_id) while True: stream: Any = None pre_run_items: list[Any] = [] @@ -829,6 +864,18 @@ async def _run_cycle( # noqa: PLR0912, PLR0915 ) finally: await coordinator.detach_stream(agent_id, stream) + except _PendingMessagesError: + # No request was sent at this checkpoint. Keep the original turn + # allowance across restarts, and never replay completed tool calls. + max_turns -= len(stream.raw_responses) + count, items = await coordinator.consume_pending(agent_id, include_items=True) + if count == 0 and await coordinator.has_pending(agent_id): + error = "Cannot persist pending agent messages" + await coordinator.set_status(agent_id, "failed", error=error) + await notify_parent_on_terminal(coordinator, agent_id, "failed") + raise RuntimeError(error) from None + input_data = [] if session is not None else stream.to_input_list() + items + continue except BudgetPausedError as exc: logger.info("agent %s paused at the scan budget limit: %s", agent_id, exc) if coordinator.budget_policy == "pause": 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/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index a26c1dc4..7f18d1e5 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -176,8 +176,10 @@ async def send_message_to_agent( ) -> str: """Send a message to another agent's inbox — sparingly. - Inter-agent messages are appended to the target's SDK session and - interrupt any active target turn so the next run cycle sees them. + Messages are queued for the target. In headless runs they enter its + context before the next model request; interactive runs also interrupt + the active turn. A successful send acknowledges queuing, not that the + target has already read or acted on the message. Use only when essential: - Sharing a discovered finding/credential another agent needs. @@ -263,7 +265,7 @@ async def send_message_to_agent( "success": True, "message_id": msg_id, "target_agent_id": target_agent_id, - "delivery_status": "delivered", + "delivery_status": "queued", }, ensure_ascii=False, default=str, @@ -693,39 +695,45 @@ async def agent_finish( filed_reports = _filed_reports_by(me) filed_report_ids = [str(r.get("id")) for r in filed_reports] - - parent_notified = False - if report_to_parent and await coordinator.claim_parent_notice(me): - async with coordinator._lock: - agent_name = coordinator.names.get(me, me) - report = _render_completion_report( - agent_name=agent_name, - agent_id=me, - task=str(inner.get("task", "")), - success=success, - result_summary=result_summary, - findings=list(findings or []), - recommendations=list(final_recommendations or []), - open_items=list(open_items or []), - filed_reports=filed_reports, - ) - await coordinator.send( - parent_id, + async with coordinator._lock: + agent_name = coordinator.names.get(me, me) + report = _render_completion_report( + agent_name=agent_name, + agent_id=me, + task=str(inner.get("task", "")), + success=success, + result_summary=result_summary, + findings=list(findings or []), + recommendations=list(final_recommendations or []), + open_items=list(open_items or []), + filed_reports=filed_reports, + ) + # Even a child that opts out of a full report owes its parent a terminal + # notice. Queue it atomically with completion, before the parent can finish. + parent_message = { + "id": f"report_{uuid.uuid4().hex[:8]}", + "from": me, + "content": report + if report_to_parent + else (f"[Agent completed] {agent_name} ({me}) finished without a completion report."), + "type": "completion" if report_to_parent else "completed", + "priority": "high", + "filed_report_ids": filed_report_ids, + } + completed, parent_notified = await coordinator.complete_if_no_pending( + me, parent_message=parent_message + ) + if not completed: + return json.dumps( { - "id": f"report_{uuid.uuid4().hex[:8]}", - "from": me, - "content": report, - "type": "completion", - "priority": "high", - "filed_report_ids": filed_report_ids, - }, + "success": False, + "agent_completed": False, + "error": ( + "Cannot finish with pending messages. Review the incoming instructions " + "before calling agent_finish again." + ), + } ) - parent_notified = True - - await coordinator.set_status(me, "completed") - if not parent_notified: - # Silence here would leave a parent waiting on a report that is never coming. - await notify_parent_on_terminal(coordinator, me, "completed") logger.info( "agent_finish: %s success=%s findings=%d filed_reports=%d parent_notified=%s", diff --git a/strix/tools/finish/tool.py b/strix/tools/finish/tool.py index e3dcb7ae..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. @@ -334,6 +343,14 @@ async def finish_scan( default=str, ) + pending_error = { + "success": False, + "scan_completed": False, + "error": "Cannot finish with pending messages. Review them before finishing the scan.", + } + if coordinator is not None and isinstance(me, str) and await coordinator.has_pending(me): + return json.dumps(pending_error) + result = await asyncio.to_thread( _do_finish, parent_id=parent_id, @@ -349,5 +366,11 @@ async def finish_scan( and coordinator is not None and isinstance(me, str) ): - await coordinator.set_status(me, "completed") + completed, _ = await coordinator.complete_if_no_pending(me) + if not completed: + # 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_agent_graph_coordination.py b/tests/test_agent_graph_coordination.py index 9694656a..ce53467a 100644 --- a/tests/test_agent_graph_coordination.py +++ b/tests/test_agent_graph_coordination.py @@ -91,7 +91,7 @@ async def test_every_terminal_non_interactive_status_is_unreachable(status: str) @pytest.mark.asyncio @pytest.mark.parametrize("status", ["running", "waiting"]) -async def test_message_to_live_child_is_delivered(status: str) -> None: +async def test_message_to_live_child_is_queued(status: str) -> None: coordinator = await _graph(interactive=False) await coordinator.set_status("child", status) @@ -103,7 +103,7 @@ async def test_message_to_live_child_is_delivered(status: str) -> None: ) assert result["success"] is True - assert result["delivery_status"] == "delivered" + assert result["delivery_status"] == "queued" assert coordinator.pending_counts["child"] == 1 diff --git a/tests/test_execution_request_log.py b/tests/test_execution_request_log.py index 426c4d47..c8ba63df 100644 --- a/tests/test_execution_request_log.py +++ b/tests/test_execution_request_log.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any import httpx import pytest @@ -68,7 +68,7 @@ async def _run( coordinator, "root", input_data="task", - run_config=cast("RunConfig", object()), + run_config=RunConfig(), context={}, max_turns=5, session=None, @@ -122,7 +122,7 @@ async def test_run_agent_loop_binds_agent_context_and_resets( await execution.run_agent_loop( agent=_Agent(), initial_input="task", - run_config=cast("RunConfig", object()), + run_config=RunConfig(), context={}, max_turns=1, coordinator=coordinator, diff --git a/tests/test_execution_transient_retry.py b/tests/test_execution_transient_retry.py index f0276932..ad528622 100644 --- a/tests/test_execution_transient_retry.py +++ b/tests/test_execution_transient_retry.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, cast +from typing import Any import httpx import pytest @@ -130,7 +130,7 @@ async def _run_once( coordinator, "root", input_data="task", - run_config=cast("RunConfig", object()), + run_config=RunConfig(), context={}, max_turns=5, session=session, diff --git a/tests/test_headless_mailbox.py b/tests/test_headless_mailbox.py new file mode 100644 index 00000000..7731336a --- /dev/null +++ b/tests/test_headless_mailbox.py @@ -0,0 +1,351 @@ +"""Offline SDK regressions for headless agents with pending inbox messages.""" + +from __future__ import annotations + +import asyncio +import copy +import json +from typing import TYPE_CHECKING, Any + +import pytest +from agents import Agent, RunConfig, function_tool +from agents.exceptions import MaxTurnsExceeded +from agents.memory import SQLiteSession +from agents.models.interface import Model +from agents.run_config import CallModelData, ModelInputData +from agents.tool_context import ToolContext +from openai.types.responses import Response, ResponseCompletedEvent, ResponseFunctionToolCall + +from strix.agents.factory import _finish_tool_use_behavior +from strix.core import execution +from strix.core.agents import AgentCoordinator +from strix.core.execution import run_agent_loop +from strix.tools.agents_graph.tools import agent_finish +from strix.tools.finish import tool as finish_module +from strix.tools.finish.tool import finish_scan + + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + from pathlib import Path + + +# Synthetic instructions sent at successive model/finish boundaries. +# No provider, target application, or historical run data is needed. +_MESSAGES = [ + {"id": "msg_hold", "from": "root", "content": "HOLD filing; reassess first."}, + {"id": "msg_reassess", "from": "root", "content": "Review the counterevidence."}, + {"id": "msg_quality", "from": "root", "content": "Apply the reporting quality gate."}, + {"id": "msg_decision", "from": "root", "content": "Final decision: do not file."}, +] + + +async def _graph() -> AgentCoordinator: + coordinator = AgentCoordinator() + await coordinator.register("root", "Root", parent_id=None) + await coordinator.register("child", "Reporter", parent_id="root") + await coordinator.attach_runtime("root", resumable=False) + await coordinator.attach_runtime("child", resumable=False) + return coordinator + + +class _MailboxModel(Model): + def __init__(self, coordinator: AgentCoordinator) -> None: + self.coordinator = coordinator + self.inputs: list[list[Any]] = [] + + async def get_response(self, *_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("Expected the streaming SDK path") + + async def stream_response(self, *_args: Any, **_kwargs: Any) -> AsyncIterator[Any]: + self.inputs.append(copy.deepcopy(_args[1])) + turn = len(self.inputs) + assert turn <= 4, "Mailbox handling must not loop or replay completed turns" + # Send after the request started. The next request must see the messages; + # finishing from this in-flight request must not strand them. + arrivals = _MESSAGES[:2] if turn == 1 else _MESSAGES[turn : turn + 1] + for message in arrivals: + assert await self.coordinator.send("child", message) + name = "work" if turn == 1 else "agent_finish" + arguments = {} if turn == 1 else {"result_summary": "done"} + response = Response( + id=f"response-{turn}", + created_at=0, + model="offline-mailbox-model", + object="response", + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + output=[ + ResponseFunctionToolCall( + id=f"fc-{turn}", + call_id=f"call-{turn}", + type="function_call", + name=name, + arguments=json.dumps(arguments), + status="completed", + ) + ], + ) + yield ResponseCompletedEvent( + response=response, sequence_number=1, type="response.completed" + ) + + +@pytest.mark.parametrize("with_session", [True, False]) +async def test_headless_messages_reach_next_request_and_block_stale_finish( + tmp_path: Path, + with_session: bool, +) -> None: + coordinator = await _graph() + model = _MailboxModel(coordinator) + work_calls: list[str] = [] + + @function_tool + async def work() -> str: + work_calls.append("work") + return "work completed" + + agent = Agent( + name="Reporter", + model=model, + tools=[work, agent_finish], + tool_use_behavior=_finish_tool_use_behavior, + ) + session = SQLiteSession("child", tmp_path / "agents.db") if with_session else None + try: + await run_agent_loop( + agent=agent, + initial_input="Review the evidence.", + run_config=RunConfig(tracing_disabled=True), + context={"coordinator": coordinator, "agent_id": "child", "parent_id": "root"}, + max_turns=4, + coordinator=coordinator, + agent_id="child", + interactive=False, + session=session, + ) + assert len(model.inputs) == 4 + assert work_calls == ["work"] + for index, items in enumerate(model.inputs[1:], start=2): + text = json.dumps(items) + for message in _MESSAGES[:index]: + assert text.count(message["content"]) == 1 + assert coordinator.statuses["child"] == "completed" + assert coordinator.pending_counts["child"] == 0 + assert coordinator.runtimes["child"].mailbox == [] + assert len(coordinator.runtimes["root"].mailbox) == 1 + if session is not None: + persisted = json.dumps(await session.get_items()) + for message in _MESSAGES: + assert persisted.count(message["content"]) == 1 + finally: + if session is not None: + session.close() + + +async def test_message_checkpoints_do_not_reset_sdk_turn_limit(tmp_path: Path) -> None: + coordinator = await _graph() + model = _MailboxModel(coordinator) + + @function_tool + async def work() -> str: + return "done" + + session = SQLiteSession("child", tmp_path / "agents.db") + try: + with pytest.raises(MaxTurnsExceeded): + await run_agent_loop( + agent=Agent( + name="Reporter", + model=model, + tools=[work, agent_finish], + tool_use_behavior=_finish_tool_use_behavior, + ), + initial_input="Review the evidence.", + run_config=RunConfig(tracing_disabled=True), + context={"coordinator": coordinator, "agent_id": "child", "parent_id": "root"}, + max_turns=2, + coordinator=coordinator, + agent_id="child", + interactive=False, + session=session, + ) + assert len(model.inputs) == 2 + assert coordinator.statuses["child"] == "stopped" + finally: + session.close() + + +async def test_finish_does_not_notify_parent_with_unread_messages() -> None: + coordinator = await _graph() + await coordinator.send("child", _MESSAGES[0]) + ctx = ToolContext( + context={"coordinator": coordinator, "agent_id": "child", "parent_id": "root"}, + tool_name="agent_finish", + tool_call_id="finish", + tool_arguments="{}", + ) + result = json.loads(await agent_finish.on_invoke_tool(ctx, '{"result_summary":"done"}')) + assert result["success"] is False + assert result["agent_completed"] is False + assert coordinator.statuses["child"] == "running" + assert coordinator.runtimes["root"].mailbox == [] + + +async def test_failed_mailbox_write_retains_messages_for_retry(tmp_path: Path) -> None: + coordinator = await _graph() + session = SQLiteSession("child", tmp_path / "agents.db") + await coordinator.attach_runtime("child", session=session) + original_add = session.add_items + + async def fail(_items: Any) -> None: + raise OSError("injected session write failure") + + try: + await coordinator.send("child", _MESSAGES[0]) + session.add_items = fail # type: ignore[method-assign] + assert await coordinator.consume_pending("child", include_items=True) == (0, []) + assert coordinator.pending_counts["child"] == 1 + assert coordinator.runtimes["child"].mailbox == [_MESSAGES[0]] + session.add_items = original_add # type: ignore[method-assign] + count, items = await coordinator.consume_pending("child", include_items=True) + assert count == 1 + assert items == await session.get_items() + finally: + session.close() + + +@pytest.mark.parametrize("send_first", [True, False]) +async def test_send_and_completion_are_serialized(send_first: bool) -> None: + coordinator = await _graph() + async with coordinator._lock: + operations = [ + lambda: coordinator.send("child", _MESSAGES[0]), + lambda: coordinator.complete_if_no_pending("child"), + ] + if not send_first: + operations.reverse() + tasks = [asyncio.create_task(operation()) for operation in operations] + await asyncio.sleep(0) + results = await asyncio.gather(*tasks) + if not send_first: + results.reverse() + accepted, (completed, _) = results + assert accepted is send_first + assert completed is not send_first + assert not (accepted and coordinator.statuses["child"] == "completed") + + +async def test_completion_report_is_queued_before_parent_sees_terminal_child( + monkeypatch: pytest.MonkeyPatch, +) -> None: + coordinator = await _graph() + observed: list[bool] = [] + + async def snapshot() -> None: + if coordinator.statuses["child"] == "completed": + observed.append(await coordinator.has_pending("root")) + + monkeypatch.setattr(coordinator, "_maybe_snapshot", snapshot) + completed, notified = await coordinator.complete_if_no_pending( + "child", parent_message={"from": "child", "content": "done"} + ) + assert completed and notified + assert observed == [True] + assert await coordinator.complete_if_no_pending("root") == (False, False) + + +async def test_concurrent_drains_preserve_arrivals_during_write(tmp_path: Path) -> None: + coordinator = await _graph() + session = SQLiteSession("child", tmp_path / "agents.db") + await coordinator.attach_runtime("child", session=session) + original_add = session.add_items + entered, release = asyncio.Event(), asyncio.Event() + + async def delayed(items: Any) -> None: + entered.set() + await release.wait() + await original_add(items) + + session.add_items = delayed # type: ignore[method-assign] + try: + await coordinator.send("child", _MESSAGES[0]) + first = asyncio.create_task(coordinator.consume_pending("child", include_items=True)) + await asyncio.wait_for(entered.wait(), 5) + await coordinator.send("child", _MESSAGES[1]) + second = asyncio.create_task(coordinator.consume_pending("child", include_items=True)) + assert await coordinator.complete_if_no_pending("child") == (False, False) + release.set() + batches = await asyncio.gather(first, second) + assert [count for count, _ in batches] == [1, 1] + persisted = await session.get_items() + assert len(persisted) == 2 + assert _MESSAGES[0]["content"] in str(persisted[0]) + assert _MESSAGES[1]["content"] in str(persisted[1]) + finally: + release.set() + session.close() + + +@pytest.mark.parametrize("during_persistence", [False, True]) +async def test_root_cannot_finish_with_unread_messages( + monkeypatch: pytest.MonkeyPatch, + during_persistence: bool, +) -> None: + coordinator = AgentCoordinator() + await coordinator.register("root", "Root", parent_id=None) + await coordinator.attach_runtime("root", resumable=False) + persisted: list[bool] = [] + + async def persist(*_args: Any, **_kwargs: Any) -> dict[str, Any]: + persisted.append(True) + await coordinator.send("root", {"from": "reviewer", "content": "Check the conclusion"}) + return {"success": True, "scan_completed": True} + + monkeypatch.setattr(finish_module.asyncio, "to_thread", persist) + if not during_persistence: + await coordinator.send("root", {"from": "reviewer", "content": "Check the conclusion"}) + ctx = ToolContext( + context={"coordinator": coordinator, "agent_id": "root", "parent_id": None}, + tool_name="finish_scan", + tool_call_id="finish", + tool_arguments="{}", + ) + result = json.loads( + await finish_scan.on_invoke_tool( + ctx, + json.dumps( + dict.fromkeys( + ("executive_summary", "methodology", "technical_analysis", "recommendations"), + "Reviewed", + ) + ), + ) + ) + assert result["success"] is False + assert result["scan_completed"] is False + assert coordinator.statuses["root"] == "running" + assert bool(persisted) is during_persistence + + +async def test_checkpoint_preserves_an_existing_model_input_filter() -> None: + coordinator = await _graph() + calls: list[bool] = [] + + async def original(data: Any) -> Any: + calls.append(True) + return data.model_data + + config = RunConfig(call_model_input_filter=original) + wrapped = execution._with_mailbox_checkpoint(config, coordinator, "child") + data = CallModelData( + model_data=ModelInputData(input=[{"role": "user", "content": "task"}], instructions=None), + agent=Agent(name="Reporter"), + context={}, + ) + assert wrapped.call_model_input_filter is not None + result = await wrapped.call_model_input_filter(data) + assert result == data.model_data + assert calls == [True] + assert config.call_model_input_filter is original 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() diff --git a/tests/test_wait_for_user.py b/tests/test_wait_for_user.py index 29881f57..3165a207 100644 --- a/tests/test_wait_for_user.py +++ b/tests/test_wait_for_user.py @@ -6,6 +6,7 @@ import json from typing import Any import pytest +from agents.memory import SQLiteSession from agents.tool_context import ToolContext from strix.core.agents import AgentCoordinator @@ -66,9 +67,14 @@ async def test_rejected_in_an_autonomous_run() -> None: async def test_a_message_that_already_arrived_is_taken_instead_of_parking() -> None: context = await _context(interactive=True) coordinator = context["coordinator"] - await coordinator.send("root", {"from": "user", "content": "wait, one more thing"}) - - result = await _call(context) + session = SQLiteSession("root") + await coordinator.attach_runtime("root", session=session) + try: + await coordinator.send("root", {"from": "user", "content": "wait, one more thing"}) + result = await _call(context) + assert await session.get_items() == [{"role": "user", "content": "wait, one more thing"}] + finally: + session.close() assert result["wait_outcome"] == "message_arrived" assert result["pending_messages"] == 1