reverting those changes

Author: RoxyFarhad

examples/tutorials/10_async/00_base/000_hello_acp/tests/test_agent.py

@@ -15,22 +15,24 @@
 - AGENT_NAME: Name of the agent to test (default: ab000-hello-acp)
 """
 
+import asyncio
 import os
 import uuid
 
 import pytest
 import pytest_asyncio
-from test_utils.async_utils import (
-    poll_messages,
-    stream_agent_response,
-    send_event_and_poll_yielding,
-)
 
 from agentex import AsyncAgentex
 from agentex.types import TaskMessage
 from agentex.types.agent_rpc_params import ParamsCreateTaskRequest
 from agentex.types.text_content_param import TextContentParam
 
+from test_utils.async_utils import (
+    poll_messages,
+    send_event_and_poll_yielding,
+    stream_agent_response,
+)
+
 # Configuration from environment variables
 AGENTEX_API_BASE_URL = os.environ.get("AGENTEX_API_BASE_URL", "http://localhost:5003")
 AGENT_NAME = os.environ.get("AGENT_NAME", "ab000-hello-acp")
@@ -143,36 +145,43 @@ async def test_send_event_and_stream(self, client: AsyncAgentex, agent_id: str):
         user_echo_found = False
         agent_response_found = False
 
+        async def stream_messages() -> None:
+            nonlocal user_echo_found, agent_response_found
+            async for event in stream_agent_response(
+                client=client,
+                task_id=task.id,
+                timeout=stream_timeout,
+            ):
+                all_events.append(event)
+                # Check events as they arrive
+                event_type = event.get("type")
+                if event_type == "full":
+                    content = event.get("content", {})
+                    if content.get("content") is None:
+                        continue  # Skip empty content
+                    if content.get("type") == "text" and content.get("author") == "agent":
+                        # Check for agent response to user message
+                        if "Hello! I've received your message" in content.get("content", ""):
+                            # Agent response should come after user echo
+                            assert user_echo_found, "Agent response arrived before user message echo (incorrect order)"
+                            agent_response_found = True
+                    elif content.get("type") == "text" and content.get("author") == "user":
+                        # Check for user message echo
+                        if content.get("content") == user_message:
+                            user_echo_found = True
+                elif event_type == "done":
+                    break
+
+                # Exit early if we've found all expected messages
+                if user_echo_found and agent_response_found:
+                    break
+
+        stream_task = asyncio.create_task(stream_messages())
+
         # Send the event
         event_content = TextContentParam(type="text", author="user", content=user_message)
         await client.agents.send_event(agent_id=agent_id, params={"task_id": task.id, "content": event_content})
-
-        async for event in stream_agent_response(
-            client=client,
-            task_id=task.id,
-            timeout=stream_timeout,
-        ):
-            all_events.append(event)
-            # Check events as they arrive
-            event_type = event.get("type")
-            if event_type == "full":
-                content = event.get("content", {})
-                if content.get("content") is None:
-                    continue  # Skip empty content
-                if content.get("type") == "text" and content.get("author") == "agent":
-                    # Check for agent response to user message
-                    if "Hello! I've received your message" in content.get("content", ""):
-                        # Agent response should come after user echo
-                        assert user_echo_found, "Agent response arrived before user message echo (incorrect order)"
-                        agent_response_found = True
-                elif content.get("type") == "text" and content.get("author") == "user":
-                    # Check for user message echo
-                    if content.get("content") == user_message:
-                        user_echo_found = True
-
-            # Exit early if we've found all expected messages
-            if user_echo_found and agent_response_found:
-                break
+        await stream_task
 
         # Verify all expected messages were received (fail if stream ended without finding them)
         assert user_echo_found, "User message echo not found in stream"

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

@@ -167,31 +167,38 @@ async def test_send_event_and_stream(self, client: AsyncAgentex, agent_id: str):
         # Flags to track what we've received
         user_message_found = False
         agent_response_found = False
+        async def stream_messages() -> None:
+            nonlocal user_message_found, agent_response_found
+            async for event in stream_agent_response(
+                client=client,
+                task_id=task.id,
+                timeout=15,
+            ):
+                all_events.append(event)
+
+                # Check events as they arrive
+                event_type = event.get("type")
+                if event_type == "full":
+                    content = event.get("content", {})
+                    if content.get("content") == user_message and content.get("author") == "user":
+                        # User message should come before agent response
+                        assert not agent_response_found, "User message arrived after agent response (incorrect order)"
+                        user_message_found = True
+                    elif content.get("author") == "agent":
+                        # Agent response should come after user message
+                        assert user_message_found, "Agent response arrived before user message (incorrect order)"
+                        agent_response_found = True
+                elif event_type == "done":
+                    break
+
+                # Exit early if we've found both messages
+                if user_message_found and agent_response_found:
+                    break
+
+        stream_task = asyncio.create_task(stream_messages())
         event_content = TextContentParam(type="text", author="user", content=user_message)
         await client.agents.send_event(agent_id=agent_id, params={"task_id": task.id, "content": event_content})
-        async for event in stream_agent_response(
-            client=client,
-            task_id=task.id,
-            timeout=15,
-        ):
-            all_events.append(event)
-
-            # Check events as they arrive
-            event_type = event.get("type")
-            if event_type == "full":
-                content = event.get("content", {})
-                if content.get("content") == user_message and content.get("author") == "user":
-                    # User message should come before agent response
-                    assert not agent_response_found, "User message arrived after agent response (incorrect order)"
-                    user_message_found = True
-                elif content.get("author") == "agent":
-                    # Agent response should come after user message
-                    assert user_message_found, "Agent response arrived before user message (incorrect order)"
-                    agent_response_found = True
-
-            # Exit early if we've found both messages
-            if user_message_found and agent_response_found:
-                break
+        await stream_task
 
         # Validate we received events
         assert len(all_events) > 0, "No events received in streaming response"

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

@@ -167,29 +167,36 @@ async def test_send_event_and_stream(self, client: AsyncAgentex, agent_id: str):
         user_message_found = False
         full_agent_message_found = False
         delta_messages_found = False
+        async def stream_messages() -> None:
+            nonlocal user_message_found, full_agent_message_found, delta_messages_found
+            async for event in stream_agent_response(
+                client=client,
+                task_id=task.id,
+                timeout=15,
+            ):
+                all_events.append(event)
+
+                # Check events as they arrive
+                event_type = event.get("type")
+                if event_type == "full":
+                    content = event.get("content", {})
+                    if content.get("content") == user_message and content.get("author") == "user":
+                        user_message_found = True
+                    elif content.get("author") == "agent":
+                        full_agent_message_found = True
+                elif event_type == "delta":
+                    delta_messages_found = True
+                elif event_type == "done":
+                    break
+
+                # Exit early if we've found all expected messages
+                if user_message_found and full_agent_message_found and delta_messages_found:
+                    break
+
+        stream_task = asyncio.create_task(stream_messages())
         event_content = TextContentParam(type="text", author="user", content=user_message)
         await client.agents.send_event(agent_id=agent_id, params={"task_id": task.id, "content": event_content})
-        async for event in stream_agent_response(
-            client=client,
-            task_id=task.id,
-            timeout=15,
-        ):
-            all_events.append(event)
-
-            # Check events as they arrive
-            event_type = event.get("type")
-            if event_type == "full":
-                content = event.get("content", {})
-                if content.get("content") == user_message and content.get("author") == "user":
-                    user_message_found = True
-                elif content.get("author") == "agent":
-                    full_agent_message_found = True
-            elif event_type == "delta":
-                delta_messages_found = True
-
-            # Exit early if we've found all expected messages
-            if user_message_found and full_agent_message_found and delta_messages_found:
-                break
+        await stream_task
 
         # Validate we received events
         assert len(all_events) > 0, "No events received in streaming response"

examples/tutorials/10_async/00_base/040_other_sdks/tests/test_agent.py

@@ -280,24 +280,32 @@ async def test_send_event_and_stream_simple(self, client: AsyncAgentex, agent_id
         # Collect events from stream
         # Check for user message and delta messages
         user_message_found = False
+        async def stream_messages() -> None:
+            nonlocal user_message_found
+            async for event in stream_agent_response(
+                client=client,
+                task_id=task.id,
+                timeout=20,
+            ):
+                msg_type = event.get("type")
+                # For full messages, content is at the top level
+                # For delta messages, we need to check parent_task_message
+                if msg_type == "full":
+                    if (
+                        event.get("content", {}).get("type") == "text"
+                        and event.get("content", {}).get("author") == "user"
+                    ):
+                        user_message_found = True
+                elif msg_type == "done":
+                    break
+
+                if user_message_found:
+                    break
+
+        stream_task = asyncio.create_task(stream_messages())
         event_content = TextContentParam(type="text", author="user", content=user_message)
         await client.agents.send_event(agent_id=agent_id, params={"task_id": task.id, "content": event_content})
-        async for event in stream_agent_response(
-            client=client,
-            task_id=task.id,
-            timeout=20,
-        ):
-            msg_type = event.get("type")
-            # For full messages, content is at the top level
-            # For delta messages, we need to check parent_task_message
-            if msg_type == "full":
-                if (
-                    event.get("content", {}).get("type") == "text"
-                    and event.get("content", {}).get("author") == "user"
-                ):
-                    user_message_found = True
-            elif msg_type == "done":
-                break
+        await stream_task
         assert user_message_found, "User message found in stream"
         ## keep polling the states for 10 seconds for the input_list and turn_number to be updated
         for i in range(10):
@@ -333,47 +341,51 @@ async def test_send_event_and_stream_with_tools(self, client: AsyncAgentex, agen
         tool_requests_seen = []
         tool_responses_seen = []
         text_deltas_seen = []
+        async def stream_messages() -> None:
+            async for event in stream_agent_response(
+                client=client,
+                task_id=task.id,
+                timeout=45,
+            ):
+                msg_type = event.get("type")
+
+                # For full messages, content is at the top level
+                # For delta messages, we need to check parent_task_message
+                if msg_type == "delta":
+                    parent_msg = event.get("parent_task_message", {})
+                    content = parent_msg.get("content", {})
+                    delta = event.get("delta", {})
+                    content_type = content.get("type")
+
+                    if content_type == "text":
+                        text_deltas_seen.append(delta.get("text_delta", ""))
+                elif msg_type == "full":
+                    # For full messages
+                    content = event.get("content", {})
+                    content_type = content.get("type")
+
+                    if content_type == "tool_request":
+                        tool_requests_seen.append(
+                            {
+                                "name": content.get("name"),
+                                "tool_call_id": content.get("tool_call_id"),
+                                "streaming_type": msg_type,
+                            }
+                        )
+                    elif content_type == "tool_response":
+                        tool_responses_seen.append(
+                            {
+                                "tool_call_id": content.get("tool_call_id"),
+                                "streaming_type": msg_type,
+                            }
+                        )
+                elif msg_type == "done":
+                    break
+
+        stream_task = asyncio.create_task(stream_messages())
         event_content = TextContentParam(type="text", author="user", content=user_message)
         await client.agents.send_event(agent_id=agent_id, params={"task_id": task.id, "content": event_content})
-        async for event in stream_agent_response(
-            client=client,
-            task_id=task.id,
-            timeout=45,
-        ):
-            msg_type = event.get("type")
-
-            # For full messages, content is at the top level
-            # For delta messages, we need to check parent_task_message
-            if msg_type == "delta":
-                parent_msg = event.get("parent_task_message", {})
-                content = parent_msg.get("content", {})
-                delta = event.get("delta", {})
-                content_type = content.get("type")
-
-                if content_type == "text":
-                    text_deltas_seen.append(delta.get("text_delta", ""))
-            elif msg_type == "full":
-                # For full messages
-                content = event.get("content", {})
-                content_type = content.get("type")
-
-                if content_type == "tool_request":
-                    tool_requests_seen.append(
-                        {
-                            "name": content.get("name"),
-                            "tool_call_id": content.get("tool_call_id"),
-                            "streaming_type": msg_type,
-                        }
-                    )
-                elif content_type == "tool_response":
-                    tool_responses_seen.append(
-                        {
-                            "tool_call_id": content.get("tool_call_id"),
-                            "streaming_type": msg_type,
-                        }
-                    )
-            elif msg_type == "done":
-                break
+        await stream_task
 
         # Verify we saw tool usage (if the agent decided to use tools)
         # Note: The agent may or may not use tools depending on its reasoning