From 55d8c63dc0f85461ca514b717f84ab5a1d4494e1 Mon Sep 17 00:00:00 2001 From: Ahmed Allam Date: Wed, 5 Aug 2026 20:27:19 +0000 Subject: [PATCH] fix(agents): collapse repeated waits queued inside one model turn --- strix/agents/prompts/system_prompt.jinja | 1 + strix/core/hooks.py | 3 + strix/tools/agents_graph/tools.py | 25 +++++ tests/test_wait_dedupe.py | 126 +++++++++++++++++++++++ 4 files changed, 155 insertions(+) create mode 100644 tests/test_wait_dedupe.py diff --git a/strix/agents/prompts/system_prompt.jinja b/strix/agents/prompts/system_prompt.jinja index 0be45595..69fc633b 100644 --- a/strix/agents/prompts/system_prompt.jinja +++ b/strix/agents/prompts/system_prompt.jinja @@ -50,6 +50,7 @@ AUTONOMOUS BEHAVIOR: - NEVER send an empty or blank message. If you have no content to output or need to wait for subagent results, you MUST call the wait_for_agents tool (or another appropriate tool) instead of emitting an empty response. - There is no user attached to this run, so there is nobody to ask and nothing to yield to. If there is nothing left to execute: do NOT send filler/repetitive text — either call wait_for_agents (only if you are genuinely expecting another agent to message you) or finish your work (subagents: agent_finish; root: finish_scan) - While the agent loop is running, almost every output MUST be a tool call. Do NOT send plain text messages; act via tools. If waiting on another agent, use wait_for_agents; when done, use agent_finish (subagents) or finish_scan (root) +- wait_for_agents blocks and resumes you automatically, so it is never a poll you repeat: issue exactly ONE wait, then stop and react to what it returns. Never write out a wait/check loop (wait → view_agent_graph → wait → ...) ahead of time — those extra calls only strand you and are collapsed anyway - A text-only turn does nothing: it neither ends the run nor yields — it just wastes a turn and forces a retry. The lifecycle tools (``finish_scan`` for root, ``agent_finish`` for subagents) are the ONLY way to terminate, and the report flows through them. If you find yourself wanting to say "Done!" or "Scan complete" without a tool call, call the lifecycle tool instead. {% endif %} diff --git a/strix/core/hooks.py b/strix/core/hooks.py index bfeafd45..21400c0b 100644 --- a/strix/core/hooks.py +++ b/strix/core/hooks.py @@ -20,6 +20,8 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +LLM_TURN_KEY = "llm_turn" + _STAGE_LABELS: tuple[str, ...] = ("NOTICE", "URGENT", "CRITICAL") _TURN_WARN_BANDS: tuple[float, ...] = (0.70, 0.85, 0.95) _ROOT_BUDGET_WARN_BANDS: tuple[float, ...] = (0.70, 0.85, 0.95) @@ -144,6 +146,7 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]): system_prompt: str | None, # noqa: ARG002 input_items: list[TResponseInputItem], ) -> None: + context.context[LLM_TURN_KEY] = int(context.context.get(LLM_TURN_KEY, 0)) + 1 try: self._maybe_warn_turns(context, input_items) self._maybe_warn_budget(context, input_items) diff --git a/strix/tools/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index ebfb3846..d4acbc57 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -14,6 +14,7 @@ from agents import RunContextWrapper, function_tool from strix.core.agents import Status, coordinator_from_context from strix.core.execution import notify_parent_on_terminal +from strix.core.hooks import LLM_TURN_KEY from strix.skills import validate_requested_skills @@ -224,6 +225,7 @@ _WAIT_DEFAULT_TIMEOUT_S = 300 # ``timeout_seconds`` the model asks for. One second of headroom lets the # tool's own timeout fire first and return a clean result. _WAIT_HARD_CEILING_S = _WAIT_DEFAULT_TIMEOUT_S + 1 +_WAITED_TURN_KEY = "waited_llm_turn" @function_tool(timeout=_WAIT_HARD_CEILING_S) @@ -239,6 +241,11 @@ async def wait_for_agents( # noqa: PLR0911 completion reports. You resume the instant any message arrives, so size ``timeout_seconds`` to the work you're awaiting. + **Issue exactly one wait, then stop and react to what it returns.** + This call blocks and resumes on its own; it is not a poll you repeat. + Do not write out a wait/check loop ahead of time — a second wait in + the same turn returns immediately without waiting. + **This tool is only for waiting on other agents.** Two things it is NOT for: @@ -290,6 +297,24 @@ async def wait_for_agents( # noqa: PLR0911 default=str, ) + turn = inner.get(LLM_TURN_KEY) + if turn is not None and inner.get(_WAITED_TURN_KEY) == turn: + return json.dumps( + { + "success": True, + "wait_outcome": "already_waited", + "reason": reason, + "note": ( + "You already waited in this turn. A single wait_for_agents blocks and " + "resumes on its own, so queueing more waits only strands you — issue one " + "wait, then react to what it returns." + ), + }, + ensure_ascii=False, + default=str, + ) + inner[_WAITED_TURN_KEY] = turn + async with coordinator._lock: stopped = coordinator.statuses.get(me) == "stopped" if stopped: diff --git a/tests/test_wait_dedupe.py b/tests/test_wait_dedupe.py new file mode 100644 index 00000000..db97fb4d --- /dev/null +++ b/tests/test_wait_dedupe.py @@ -0,0 +1,126 @@ +"""Tests for collapsing repeated waits queued inside one model turn. + +An orchestrator that writes out its whole poll loop ahead of time queues +many ``wait_for_agents`` calls in a single response. Each one parks for its +full timeout, so the agent stops reacting for hours while its children run +unsupervised. Only the first wait of a turn parks; the rest return at once. +""" + +from __future__ import annotations + +import asyncio +import json +import time +from typing import TYPE_CHECKING, Any, cast + +import pytest +from agents import RunContextWrapper +from agents.tool_context import ToolContext + +from strix.core.agents import AgentCoordinator +from strix.core.hooks import LLM_TURN_KEY, ReportUsageHooks +from strix.tools.agents_graph.tools import wait_for_agents + + +if TYPE_CHECKING: + from collections.abc import Iterator + + +_WAIT_SECONDS = 2 + + +@pytest.fixture +def _fast_wait(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + # The real ceiling is 300s per wait; the shape of the bug is the same. + monkeypatch.setattr( + "strix.tools.agents_graph.tools._WAIT_DEFAULT_TIMEOUT_S", _WAIT_SECONDS, raising=True + ) + yield + + +async def _context() -> dict[str, Any]: + coordinator = AgentCoordinator() + await coordinator.register("root", "strix", parent_id=None) + return {"agent_id": "root", "coordinator": coordinator} + + +async def _wait(inner: dict[str, Any]) -> dict[str, Any]: + ctx = ToolContext( + context=inner, + tool_name="wait_for_agents", + tool_call_id="call-1", + tool_arguments="{}", + ) + raw: str = await wait_for_agents.on_invoke_tool( + ctx, json.dumps({"reason": "waiting for wave 1", "timeout_seconds": _WAIT_SECONDS}) + ) + return cast("dict[str, Any]", json.loads(raw)) + + +@pytest.mark.asyncio +async def test_waits_queued_in_one_turn_each_park_without_the_guard(_fast_wait: None) -> None: + # Repro: no turn marker in context (as before the fix) — every queued wait + # parks for its full timeout, so N waits cost N x timeout. + inner = await _context() + + started = time.monotonic() + outcomes = [(await _wait(inner))["wait_outcome"] for _ in range(3)] + elapsed = time.monotonic() - started + + assert outcomes == ["timeout", "timeout", "timeout"] + assert elapsed >= 3 * _WAIT_SECONDS + + +@pytest.mark.asyncio +async def test_repeated_waits_in_one_turn_are_collapsed(_fast_wait: None) -> None: + inner = await _context() + inner[LLM_TURN_KEY] = 1 + + started = time.monotonic() + outcomes = [(await _wait(inner))["wait_outcome"] for _ in range(3)] + elapsed = time.monotonic() - started + + assert outcomes == ["timeout", "already_waited", "already_waited"] + assert elapsed < 2 * _WAIT_SECONDS + + +@pytest.mark.asyncio +async def test_a_wait_in_the_next_turn_still_parks(_fast_wait: None) -> None: + inner = await _context() + inner[LLM_TURN_KEY] = 1 + assert (await _wait(inner))["wait_outcome"] == "timeout" + assert (await _wait(inner))["wait_outcome"] == "already_waited" + + inner[LLM_TURN_KEY] = 2 + + assert (await _wait(inner))["wait_outcome"] == "timeout" + + +@pytest.mark.asyncio +async def test_each_model_turn_bumps_the_turn_marker() -> None: + hooks = ReportUsageHooks(model="gw-model") + context: RunContextWrapper[dict[str, Any]] = RunContextWrapper(context={}) + agent = cast("Any", None) + + await hooks.on_llm_start(context, agent, None, []) + await hooks.on_llm_start(context, agent, None, []) + + assert context.context[LLM_TURN_KEY] == 2 + + +@pytest.mark.asyncio +async def test_a_collapsed_wait_still_reports_arriving_messages(_fast_wait: None) -> None: + inner = await _context() + inner[LLM_TURN_KEY] = 1 + coordinator = cast("AgentCoordinator", inner["coordinator"]) + + async def _send() -> None: + await asyncio.sleep(0.1) + await coordinator.send("root", {"type": "information", "content": "child done"}) + + task = asyncio.create_task(_send()) + first = await _wait(inner) + await task + + assert first["wait_outcome"] == "message_arrived" + assert (await _wait(inner))["wait_outcome"] == "already_waited"