2828)
2929
3030from agentex import AsyncAgentex
31- from agentex .types import TaskMessage , TextContent
3231from agentex .types .agent_rpc_params import ParamsCreateTaskRequest
3332from agentex .types .text_content_param import TextContentParam
3433
@@ -89,6 +88,10 @@ async def test_send_event_and_poll(self, client: AsyncAgentex, agent_id: str):
8988
9089 user_message = "Hello! Here is my test message"
9190 messages = []
91+
92+ # Flags to track what we've received
93+ user_message_found = False
94+ agent_response_found = False
9295 async for message in send_event_and_poll_yielding (
9396 client = client ,
9497 agent_id = agent_id ,
@@ -98,23 +101,25 @@ async def test_send_event_and_poll(self, client: AsyncAgentex, agent_id: str):
98101 sleep_interval = 1.0 ,
99102 yield_updates = False ,
100103 ):
101-
102104 messages .append (message )
103105
104- assert len (messages ) > 0
105- # the first message should be the agent re-iterating what the user sent
106- assert isinstance (messages , List )
107- assert len (messages ) == 2
108- first_message : TaskMessage = messages [0 ]
109- assert first_message .content == TextContent (
110- author = "user" ,
111- content = user_message ,
112- type = "text" ,
113- )
114-
115- second_message : TaskMessage = messages [1 ]
116- assert second_message .content is not None
117- assert second_message .content .author == "agent"
106+ # Validate messages as they come in
107+ if message .content and hasattr (message .content , "author" ):
108+ if message .content .author == "user" and message .content .content == user_message :
109+ user_message_found = True
110+ elif message .content .author == "agent" :
111+ # Agent response should come after user message
112+ assert user_message_found , "Agent response arrived before user message"
113+ agent_response_found = True
114+
115+ # Exit early if we've found all expected messages
116+ if user_message_found and agent_response_found :
117+ break
118+
119+ # Validate we received expected messages
120+ assert len (messages ) >= 2 , "Expected at least 2 messages (user + agent)"
121+ assert user_message_found , "User message not found"
122+ assert agent_response_found , "Agent response not found"
118123
119124 # assert the state has been updated
120125 await asyncio .sleep (1 ) # wait for state to be updated
@@ -158,41 +163,43 @@ async def test_send_event_and_stream(self, client: AsyncAgentex, agent_id: str):
158163 # Collect events from stream
159164 all_events = []
160165
166+ # Flags to track what we've received
167+ user_message_found = False
168+ full_agent_message_found = False
169+ delta_messages_found = False
161170 async def stream_messages () -> None :
171+ nonlocal user_message_found , full_agent_message_found , delta_messages_found
162172 async for event in stream_agent_response (
163173 client = client ,
164174 task_id = task .id ,
165175 timeout = 15 ,
166176 ):
167177 all_events .append (event )
168178
169- stream_task = asyncio .create_task (stream_messages ())
179+ # Check events as they arrive
180+ event_type = event .get ("type" )
181+ if event_type == "full" :
182+ content = event .get ("content" , {})
183+ if content .get ("content" ) == user_message and content .get ("author" ) == "user" :
184+ user_message_found = True
185+ elif content .get ("author" ) == "agent" :
186+ full_agent_message_found = True
187+ elif event_type == "delta" :
188+ delta_messages_found = True
189+ elif event_type == "done" :
190+ break
191+
192+ # Exit early if we've found all expected messages
193+ if user_message_found and full_agent_message_found and delta_messages_found :
194+ break
170195
196+ stream_task = asyncio .create_task (stream_messages ())
171197 event_content = TextContentParam (type = "text" , author = "user" , content = user_message )
172198 await client .agents .send_event (agent_id = agent_id , params = {"task_id" : task .id , "content" : event_content })
173-
174- # Wait for streaming to complete
175199 await stream_task
176200
177201 # Validate we received events
178202 assert len (all_events ) > 0 , "No events received in streaming response"
179-
180- # Check for user message, full agent response, and delta messages
181- user_message_found = False
182- full_agent_message_found = False
183- delta_messages_found = False
184-
185- for event in all_events :
186- event_type = event .get ("type" )
187- if event_type == "full" :
188- content = event .get ("content" , {})
189- if content .get ("content" ) == user_message and content .get ("author" ) == "user" :
190- user_message_found = True
191- elif content .get ("author" ) == "agent" :
192- full_agent_message_found = True
193- elif event_type == "delta" :
194- delta_messages_found = True
195-
196203 assert user_message_found , "User message not found in stream"
197204 assert full_agent_message_found , "Full agent message not found in stream"
198205 assert delta_messages_found , "Delta messages not found in stream (streaming response expected)"
0 commit comments