Merge 6a0441e072 into 717ffc8f4c

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,

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()