diff --git a/strix/tools/agents_graph/tools.py b/strix/tools/agents_graph/tools.py index ad16fe03..83676810 100644 --- a/strix/tools/agents_graph/tools.py +++ b/strix/tools/agents_graph/tools.py @@ -228,13 +228,23 @@ def _session_items_payload(items: list[Any]) -> list[dict[str, Any]]: _WAIT_DEFAULT_TIMEOUT_S = 300 -# Enforced by the SDK around the whole tool call, so it caps an oversized -# ``timeout_seconds`` the model asks for. One second of headroom lets the -# tool's own timeout fire first and return a clean result. +_WAIT_MIN_TIMEOUT_S = 1 +# Enforced by the SDK around the whole tool call. Requested timeouts are +# clamped to the default, so this 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" +def _effective_wait_timeout(requested: Any) -> int: + """Clamp a requested wait timeout into the range the tool can honour.""" + try: + seconds = int(requested) + except (TypeError, ValueError): + return _WAIT_DEFAULT_TIMEOUT_S + return max(_WAIT_MIN_TIMEOUT_S, min(seconds, _WAIT_DEFAULT_TIMEOUT_S)) + + @function_tool(timeout=_WAIT_HARD_CEILING_S) async def wait_for_agents( # noqa: PLR0911 ctx: RunContextWrapper, @@ -280,8 +290,8 @@ async def wait_for_agents( # noqa: PLR0911 reason: One-line note shown in graph snapshots while you're waiting (helps a human or sibling agent debug who's stuck on what). - timeout_seconds: Max seconds to wait (default 300, and values above - that are cut short by a hard ceiling). This is only + timeout_seconds: Max seconds to wait (default 300, which is also the + maximum — larger requests are clamped to it). This is only a cap — the tool returns the INSTANT a message arrives, so a larger value never makes you wait longer when the reply does come. Right-size it to what you're waiting on: a short wait @@ -293,6 +303,7 @@ async def wait_for_agents( # noqa: PLR0911 elapses. On timeout the tool returns and you decide whether to keep working or wait again. """ + effective_timeout = _effective_wait_timeout(timeout_seconds) inner = _ctx(ctx) coordinator = coordinator_from_context(inner) me = inner.get("agent_id") @@ -366,16 +377,19 @@ async def wait_for_agents( # noqa: PLR0911 await coordinator.park_waiting(me, wait_kind="agents") try: - await asyncio.wait_for(coordinator.wait_for_message(me), timeout_seconds) + await asyncio.wait_for(coordinator.wait_for_message(me), effective_timeout) except TimeoutError: await coordinator.mark_running(me) + note = "No messages within timeout — continue work or call agent_finish." + if effective_timeout != timeout_seconds: + note += f" Requested timeout was clamped to the {effective_timeout}s maximum." return json.dumps( { "success": True, "wait_outcome": "timeout", - "timeout_seconds": timeout_seconds, + "timeout_seconds": effective_timeout, "reason": reason, - "note": "No messages within timeout — continue work or call agent_finish.", + "note": note, }, ensure_ascii=False, default=str, diff --git a/tests/test_wait_dedupe.py b/tests/test_wait_dedupe.py index db97fb4d..8764d7a5 100644 --- a/tests/test_wait_dedupe.py +++ b/tests/test_wait_dedupe.py @@ -108,6 +108,83 @@ async def test_each_model_turn_bumps_the_turn_marker() -> None: assert context.context[LLM_TURN_KEY] == 2 +async def _wait_with(inner: dict[str, Any], timeout_seconds: 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": timeout_seconds}) + ) + return cast("dict[str, Any]", json.loads(raw)) + + +@pytest.fixture +def _captured_timeouts(monkeypatch: pytest.MonkeyPatch) -> Iterator[list[float]]: + """Record the timeout handed to the inner wait, without ever waiting.""" + captured: list[float] = [] + + async def fake_wait_for(awaitable: Any, timeout: float) -> None: + captured.append(timeout) + awaitable.close() + raise TimeoutError + + monkeypatch.setattr(asyncio, "wait_for", fake_wait_for) + yield captured + + +@pytest.mark.asyncio +async def test_oversized_timeout_still_returns_the_clean_timeout_payload(_fast_wait: None) -> None: + # A model asking to wait "indefinitely" must not blow past the SDK's + # tool-call ceiling and lose the payload telling it what to do next. + result = await _wait_with(await _context(), 300_000) + + assert result["success"] is True + assert result["wait_outcome"] == "timeout" + assert result["timeout_seconds"] == _WAIT_SECONDS + assert "clamped" in result["note"] + + +@pytest.mark.asyncio +async def test_inner_wait_gets_the_clamped_timeout(_captured_timeouts: list[float]) -> None: + assert (await _wait_with(await _context(), 300_000))["timeout_seconds"] == 300 + assert _captured_timeouts == [300] + + +@pytest.mark.asyncio +async def test_non_positive_timeout_is_floored(_captured_timeouts: list[float]) -> None: + assert (await _wait_with(await _context(), 0))["timeout_seconds"] == 1 + assert (await _wait_with(await _context(), -30))["timeout_seconds"] == 1 + assert _captured_timeouts == [1, 1] + + +@pytest.mark.asyncio +async def test_in_range_timeout_passes_through(_captured_timeouts: list[float]) -> None: + result = await _wait_with(await _context(), 60) + + assert result["timeout_seconds"] == 60 + assert "clamped" not in result["note"] + assert _captured_timeouts == [60] + + +@pytest.mark.asyncio +async def test_early_returns_are_unaffected_by_an_oversized_timeout(_fast_wait: None) -> None: + inner = await _context() + inner[LLM_TURN_KEY] = 1 + coordinator = cast("AgentCoordinator", inner["coordinator"]) + await coordinator.send("root", {"type": "information", "content": "child done"}) + + assert (await _wait_with(inner, 300_000))["wait_outcome"] == "message_arrived" + assert (await _wait_with(inner, 300_000))["wait_outcome"] == "already_waited" + + inner[LLM_TURN_KEY] = 2 + await coordinator.set_status("root", "stopped") + + assert (await _wait_with(inner, 300_000))["wait_outcome"] == "stopped" + + @pytest.mark.asyncio async def test_a_collapsed_wait_still_reports_arriving_messages(_fast_wait: None) -> None: inner = await _context()