@@ -156,7 +156,7 @@ async def test_multi_turn_conversation(self, client: AsyncAgentex, agent_id: str
156156 agent_id = agent_id ,
157157 task_id = task .id ,
158158 user_message = user_message_1 ,
159- timeout = 20 ,
159+ timeout = 30 ,
160160 sleep_interval = 1.0 ,
161161 ):
162162 assert isinstance (message , TaskMessage )
@@ -215,13 +215,14 @@ async def test_send_event_and_stream_with_reasoning(self, client: AsyncAgentex,
215215 # Check for user message and agent response
216216 user_message_found = False
217217 agent_response_found = False
218+ reasoning_found = False
218219
219220 async def stream_messages () -> None : # noqa: ANN101
220- nonlocal user_message_found , agent_response_found
221+ nonlocal user_message_found , agent_response_found , reasoning_found
221222 async for event in stream_agent_response (
222223 client = client ,
223224 task_id = task .id ,
224- timeout = 60 ,
225+ timeout = 90 , # Increased timeout for CI environments
225226 ):
226227 msg_type = event .get ("type" )
227228 if msg_type == "full" :
@@ -241,22 +242,40 @@ async def stream_messages() -> None: # noqa: ANN101
241242 ):
242243 agent_response_found = True
243244 elif finished_message .content and finished_message .content .type == "reasoning" :
244- tool_response_found = True
245+ reasoning_found = True
246+
247+ # Exit early if we have what we need
248+ if user_message_found and agent_response_found :
249+ break
250+
245251 elif msg_type == "done" :
246252 task_message_update = StreamTaskMessageDone .model_validate (event )
247253 if task_message_update .parent_task_message and task_message_update .parent_task_message .id :
248254 finished_message = await client .messages .retrieve (task_message_update .parent_task_message .id )
249255 if finished_message .content and finished_message .content .type == "reasoning" :
256+ reasoning_found = True
257+ elif (
258+ finished_message .content
259+ and finished_message .content .type == "text"
260+ and finished_message .content .author == "agent"
261+ ):
250262 agent_response_found = True
251- continue
263+
264+ # Exit early if we have what we need
265+ if user_message_found and agent_response_found :
266+ break
252267
253268 stream_task = asyncio .create_task (stream_messages ())
254269
255270 event_content = TextContentParam (type = "text" , author = "user" , content = user_message )
256271 await client .agents .send_event (agent_id = agent_id , params = {"task_id" : task .id , "content" : event_content })
257272
258- # Wait for streaming to complete
259- await stream_task
273+ # Wait for streaming to complete with timeout
274+ try :
275+ await asyncio .wait_for (stream_task , timeout = 120 ) # Overall timeout for CI
276+ except asyncio .TimeoutError :
277+ stream_task .cancel ()
278+ pytest .fail ("Test timed out waiting for streaming response" )
260279
261280 assert user_message_found , "User message not found in stream"
262281 assert agent_response_found , "Agent response not found in stream"
0 commit comments