@@ -76,8 +76,12 @@ async def engine_chat(self, message: str, **kwargs: Any) -> AsyncIterator[Stream
7676
7777
7878class _SlowFirstChunkService (_FakeService ):
79+ def __init__ (self , first_chunk_delay : float = 0.16 ) -> None :
80+ super ().__init__ ()
81+ self ._first_chunk_delay = first_chunk_delay
82+
7983 async def engine_chat (self , message : str , ** kwargs : Any ) -> AsyncIterator [StreamChunk ]:
80- await asyncio .sleep (0.16 )
84+ await asyncio .sleep (self . _first_chunk_delay )
8185 yield StreamChunk (request_id = "" , content = f"slow { message } " , event_type = "chunk" )
8286 yield StreamChunk (request_id = "" , content = "done" , event_type = "final" )
8387
@@ -617,11 +621,14 @@ async def test_daemon_client_stream_heartbeat_prevents_idle_timeout() -> None:
617621 with tempfile .TemporaryDirectory (prefix = "lfd-" , dir = _short_tempdir ()) as root :
618622 server , task , runtime_dir = await _start_server (
619623 Path (root ) / "runtime" ,
620- service = _SlowFirstChunkService (),
621- stream_heartbeat_s = 0.03 ,
624+ service = _SlowFirstChunkService (first_chunk_delay = 0.6 ),
625+ stream_heartbeat_s = 0.1 ,
622626 )
623627 socket_path = get_transport ().readiness_path (runtime_dir )
624- client = DaemonClient (socket_path , timeout_s = 0.1 )
628+ # The read timeout must stay smaller than the first-chunk delay (so an
629+ # idle stream would die) but large enough that slow CI dispatch latency
630+ # cannot starve the first heartbeat; 0.3s keeps both invariants.
631+ client = DaemonClient (socket_path , timeout_s = 0.3 )
625632
626633 try :
627634 events = [event async for event in client .engine_chat ("world" )]
0 commit comments