diff --git a/agentex/src/domain/use_cases/streams_use_case.py b/agentex/src/domain/use_cases/streams_use_case.py index 48eec2ce..529d8830 100644 --- a/agentex/src/domain/use_cases/streams_use_case.py +++ b/agentex/src/domain/use_cases/streams_use_case.py @@ -101,14 +101,15 @@ async def stream_task_events( task_id = task.id stream_topic = get_task_event_stream_topic(task_id=task_id) + # Snapshot the read cursor BEFORE yielding "connected". "connected" is + # the client's cue to send its message, which makes the agent start + # XADD-ing deltas. Snapshotting after the yield lets a congested relay + # fall behind far enough that those deltas land before the snapshot and + # are never read. Snapshotting first resolves to "0-0" (stream is empty + # until the client sends), so we read from the beginning. + last_id = await self.stream_repository.get_stream_tail_id(stream_topic) # Send initial connection data yield f"data: {TaskStreamConnectedEventEntity(type='connected', taskId=task_id).model_dump_json()}\n\n" - # Snapshot the tail once on entry rather than passing "$" to every - # XREAD. "$" re-resolves to the current tail on each call, so any - # entry XADD'd in the gap between BLOCKing reads lands behind the - # new "$" and is unreachable — silently dropping deltas from - # fast-emitting agents. - last_id = await self.stream_repository.get_stream_tail_id(stream_topic) last_message_time = asyncio.get_running_loop().time() ping_interval = float( self.environment_variables.SSE_KEEPALIVE_PING_INTERVAL diff --git a/agentex/tests/integration/test_task_stream.py b/agentex/tests/integration/test_task_stream.py index f2a0d762..7da23edd 100644 --- a/agentex/tests/integration/test_task_stream.py +++ b/agentex/tests/integration/test_task_stream.py @@ -563,6 +563,54 @@ async def reader(): pass repo.read_messages = original_read_messages + async def test_event_xadded_after_connected_is_delivered( + self, test_agent_and_task, streams_use_case + ): + """ + Regression for the SSE entry-gap drop bug: a delta XADDed right after + "connected" must still be delivered. + + We step the generator manually so the XADD lands while it's suspended + at the "connected" yield. If the cursor is snapshotted after that yield + the delta is lost; snapshotting before it (the fix) delivers it. + """ + from src.utils.stream_topics import get_task_event_stream_topic + + _agent, task = test_agent_and_task + stream_topic = get_task_event_stream_topic(task_id=task.id) + repo = streams_use_case.stream_repository + + gen = streams_use_case.stream_task_events(task_id=task.id) + try: + # First yielded event is "connected"; generator is now suspended here. + first = await asyncio.wait_for(gen.__anext__(), timeout=5) + assert "connected" in first, f"expected connected event first, got: {first}" + + # Emit a delta while the generator is suspended at the yield. + sentinel = "after-connected-sentinel" + await repo.send_data( + stream_topic, {"type": "error", "message": sentinel} + ) + + # The delta must be delivered; a silent stream means it was dropped. + received = False + for _ in range(10): + try: + evt = await asyncio.wait_for(gen.__anext__(), timeout=4) + except (TimeoutError, asyncio.TimeoutError): + break # stream went silent — sentinel was dropped + if evt.startswith("data: ") and sentinel in evt: + received = True + break + + assert received, ( + "A delta XADDed after 'connected' was not delivered. The read " + "cursor was snapshotted after signaling connected, so the entry " + "landed behind the tail and was lost (the SSE entry-gap bug)." + ) + finally: + await gen.aclose() + async def test_stream_sends_keepalive_pings_during_idle_periods( self, test_agent_and_task, streams_use_case ):