From 6e56b4e329777ed7e88cc9750fce57489c1671bf Mon Sep 17 00:00:00 2001 From: Guy Basson Date: Mon, 12 Oct 2026 18:07:24 +0300 Subject: [PATCH] fix(agents): deliver headless inbox messages before model turns --- strix/core/agents.py | 107 ++++++-- strix/core/execution.py | 41 +++ strix/tools/agents_graph/tools.py | 76 ++--- strix/tools/finish/tool.py | 14 +- tests/test_agent_graph_coordination.py | 4 +- tests/test_execution_request_log.py | 6 +- tests/test_execution_transient_retry.py | 4 +- tests/test_headless_mailbox.py | 351 ++++++++++++++++++++++++ tests/test_wait_for_user.py | 12 +- 9 files changed, 540 insertions(+), 75 deletions(-) create mode 100644 tests/test_headless_mailbox.py 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..ae9b8fa5 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 @@ -770,6 +797,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 +858,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/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..0148c2a6 100644 --- a/strix/tools/finish/tool.py +++ b/strix/tools/finish/tool.py @@ -334,6 +334,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 +357,9 @@ 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) return json.dumps(result, ensure_ascii=False, default=str) 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_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