mirror of
https://github.com/usestrix/strix.git
synced 2026-10-15 04:18:10 +00:00
Merge 618e676928 into 87f2d783d6
This commit is contained in:
commit
3bc0557318
12 changed files with 763 additions and 81 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
351
tests/test_headless_mailbox.py
Normal file
351
tests/test_headless_mailbox.py
Normal file
|
|
@ -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
|
||||
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()
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue