test(tutorials): poll task state instead of a fixed 1s sleep

examples/tutorials/10_async/00_base/010_multiturn/tests/test_agent.py

@@ -24,6 +24,7 @@
 import pytest_asyncio
 from test_utils.async_utils import (
     stream_agent_response,
+    wait_for_state_messages,
     send_event_and_poll_yielding,
 )
 
@@ -124,11 +125,7 @@ async def test_send_event_and_poll(self, client: AsyncAgentex, agent_id: str):
         assert user_message_found, "User message not found"
         assert agent_response_found, "Agent response not found"
 
-        await asyncio.sleep(1)  # wait for state to be updated
-        states = await client.states.list(agent_id=agent_id, task_id=task.id)
-        assert len(states) == 1
-        state = states[0].state
-        messages = state.get("messages", [])
+        messages = await wait_for_state_messages(client, agent_id, task.id, expected_count=3)
 
         assert isinstance(messages, list)
         assert len(messages) == 3
@@ -207,11 +204,7 @@ async def stream_messages() -> None:
         assert agent_response_found, "Agent response not found in stream"
 
         # Verify the state has been updated
-        await asyncio.sleep(1)  # wait for state to be updated
-        states = await client.states.list(agent_id=agent_id, task_id=task.id)
-        assert len(states) == 1
-        state = states[0].state
-        messages = state.get("messages", [])
+        messages = await wait_for_state_messages(client, agent_id, task.id, expected_count=3)
 
         assert isinstance(messages, list)
         assert len(messages) == 3

examples/tutorials/10_async/00_base/020_streaming/tests/test_agent.py

@@ -24,6 +24,7 @@
 import pytest_asyncio
 from test_utils.async_utils import (
     stream_agent_response,
+    wait_for_state_messages,
     send_event_and_poll_yielding,
 )
 
@@ -122,11 +123,7 @@ async def test_send_event_and_poll(self, client: AsyncAgentex, agent_id: str):
         assert agent_response_found, "Agent response not found"
 
         # assert the state has been updated
-        await asyncio.sleep(1)  # wait for state to be updated
-        states = await client.states.list(agent_id=agent_id, task_id=task.id)
-        assert len(states) == 1
-        state = states[0].state
-        messages = state.get("messages", [])
+        messages = await wait_for_state_messages(client, agent_id, task.id, expected_count=3)
 
         assert isinstance(messages, list)
         assert len(messages) == 3
@@ -205,11 +202,7 @@ async def stream_messages() -> None:
         assert delta_messages_found, "Delta messages not found in stream (streaming response expected)"
 
         # Verify the state has been updated
-        await asyncio.sleep(1)  # wait for state to be updated
-        states = await client.states.list(agent_id=agent_id, task_id=task.id)
-        assert len(states) == 1
-        state: dict[str, object] = states[0].state
-        messages = state.get("messages", [])
+        messages = await wait_for_state_messages(client, agent_id, task.id, expected_count=3)
 
         assert isinstance(messages, list)
         assert len(messages) == 3

examples/tutorials/test_utils/async_utils.py

@@ -156,6 +156,42 @@ async def poll_messages(
         await asyncio.sleep(sleep_interval)
 
 
+async def wait_for_state_messages(
+    client: AsyncAgentex,
+    agent_id: str,
+    task_id: str,
+    expected_count: int,
+    timeout: float = 30,
+    sleep_interval: float = 0.5,
+) -> list:
+    """
+    Poll the task state until its "messages" list reaches expected_count.
+
+    Agents typically send their reply before they persist the turn to task state,
+    so reading the state right after the reply arrives can see the previous turn.
+
+    Args:
+        client: AgentEx client instance
+        agent_id: The agent ID
+        task_id: The task ID
+        expected_count: Number of state messages to wait for
+        timeout: Maximum seconds to poll (default: 30)
+        sleep_interval: Seconds to sleep between polls (default: 0.5)
+
+    Returns:
+        The state's "messages" list from the last poll. Callers assert on it, so a
+        state that never reaches expected_count still fails with the real length.
+    """
+    deadline = time.monotonic() + timeout
+    while True:
+        states = await client.states.list(agent_id=agent_id, task_id=task_id)
+        assert len(states) == 1, f"Expected exactly one state, got {len(states)}"
+        messages = states[0].state.get("messages", [])
+        if len(messages) >= expected_count or time.monotonic() >= deadline:
+            return messages
+        await asyncio.sleep(sleep_interval)
+
+
 async def send_event_and_stream(
     client: AsyncAgentex,
     agent_id: str,

tests/lib/test_tutorial_async_utils.py

@@ -0,0 +1,79 @@
+"""Tests for wait_for_state_messages in the tutorials' test_utils.
+
+The tutorials' `test_utils` package shares its name with `tests/test_utils`, so the
+module is loaded from its file path instead of imported by name.
+"""
+
+from __future__ import annotations
+
+import time
+import importlib.util
+from types import ModuleType, SimpleNamespace
+from pathlib import Path
+
+import pytest
+
+_ASYNC_UTILS = Path(__file__).resolve().parents[2] / "examples" / "tutorials" / "test_utils" / "async_utils.py"
+
+
+def _load_async_utils() -> ModuleType:
+    spec = importlib.util.spec_from_file_location("tutorial_async_utils", _ASYNC_UTILS)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+wait_for_state_messages = _load_async_utils().wait_for_state_messages
+
+
+class _FakeStates:
+    """Returns one state whose message count grows from 1 to 3 once `write_after` seconds pass."""
+
+    def __init__(self, write_after: float | None):
+        self._start = time.monotonic()
+        self._write_after = write_after
+        self.calls = 0
+
+    async def list(self, agent_id: str, task_id: str) -> list[SimpleNamespace]:
+        del agent_id, task_id
+        self.calls += 1
+        written = self._write_after is not None and time.monotonic() - self._start >= self._write_after
+        return [SimpleNamespace(state={"messages": [{"role": "system"}] * (3 if written else 1)})]
+
+
+def _client(write_after: float | None) -> SimpleNamespace:
+    return SimpleNamespace(states=_FakeStates(write_after))
+
+
+@pytest.mark.asyncio
+async def test_waits_for_a_state_write_that_lands_after_the_reply() -> None:
+    client = _client(write_after=0.3)
+
+    messages = await wait_for_state_messages(client, "agent", "task", expected_count=3, timeout=5, sleep_interval=0.05)
+
+    assert len(messages) == 3
+    assert client.states.calls > 1
+
+
+@pytest.mark.asyncio
+async def test_returns_the_real_state_at_the_deadline_when_it_never_updates() -> None:
+    client = _client(write_after=None)
+    start = time.monotonic()
+
+    messages = await wait_for_state_messages(client, "agent", "task", expected_count=3, timeout=0.3, sleep_interval=0.05)
+
+    # It kept polling until the deadline instead of returning early. No upper bound:
+    # a busy worker may wake late, and returning at all already proves it stopped.
+    assert len(messages) == 1
+    assert time.monotonic() - start >= 0.3
+
+
+@pytest.mark.asyncio
+async def test_returns_on_the_first_poll_when_the_state_is_already_written() -> None:
+    client = _client(write_after=0)
+
+    messages = await wait_for_state_messages(client, "agent", "task", expected_count=3)
+
+    assert len(messages) == 3
+    assert client.states.calls == 1