fix(room_io): keep a re-finalized user transcript in its own segment

livekit-agents/livekit/agents/voice/room_io/_output.py

@@ -220,9 +220,35 @@ def __init__(
         self._flush_task: asyncio.Task[None] | None = None
         self._closed = False
 
+        # see _ParticipantStreamTranscriptionOutput._capture_item
+        self._current_item_id: str | None = None
+        self._bound_item_id: str | None = None
+        self._bound_segment_id: str | None = None
+        self._next_segment_id: str | None = None
+
         self._reset_state()
         self.set_participant(participant)
 
+    def _capture_item(self, item_id: str | None) -> None:
+        """Tie the text that follows to a provider item, reusing its segment on a revision."""
+        # see _ParticipantStreamTranscriptionOutput._capture_item
+        self._current_item_id = item_id
+        reuse = item_id is not None and item_id == self._bound_item_id and self._bound_segment_id
+        self._next_segment_id = self._bound_segment_id if reuse else None
+
+    def _remember_item_segment(self) -> None:
+        """Record the segment the text being captured is published under.
+
+        Done per capture rather than in ``_reset_state`` or only when a segment opens:
+        the legacy output resets after flushing (which would point the item at a segment
+        never published), and in the realtime flow the segment is opened by the empty
+        update sent when the user stops speaking, before the transcript that carries the
+        item id arrives.
+        """
+        if self._current_item_id is not None:
+            self._bound_item_id = self._current_item_id
+            self._bound_segment_id = self._current_id
+
     def set_participant(
         self,
         participant: rtc.Participant | str | None,
@@ -253,7 +279,8 @@ def set_participant(
         self._reset_state()
 
     def _reset_state(self) -> None:
-        self._current_id = utils.shortuuid("SG_")
+        self._current_id = self._next_segment_id or utils.shortuuid("SG_")
+        self._next_segment_id = None
         self._capturing = False
         self._pushed_text = ""
 
@@ -269,6 +296,8 @@ async def capture_text(self, text: str) -> None:
             self._reset_state()
             self._capturing = True
 
+        self._remember_item_segment()
+
         if self._is_delta_stream:
             self._pushed_text += text
         else:
@@ -386,9 +415,47 @@ def __init__(
         self._flush_atask: asyncio.Task[None] | None = None
         self._closed = False
 
+        # a provider item can be finalized more than once; the segment it was
+        # published under is reused so the revision lands in place
+        self._current_item_id: str | None = None
+        self._bound_item_id: str | None = None
+        self._bound_segment_id: str | None = None
+        self._next_segment_id: str | None = None
+
         self._reset_state()
         self.set_participant(participant)
 
+    def _capture_item(self, item_id: str | None) -> None:
+        """Tie the text that follows to a provider item.
+
+        A realtime provider may finalize the same input-audio transcription more than
+        once, each time with a longer transcript. Without this the second final opens a
+        new segment and the revisions stack up client-side instead of replacing the one
+        already shown.
+
+        Both fields are re-derived on every call. Nothing guarantees a capture consumes
+        the reservation (capture_text returns early until a participant is known), and
+        text that carries no item - the empty update sent when the user stops speaking,
+        or any stt path - must not inherit the previous item's association, or a late
+        revision of that item would land in a newer utterance's segment.
+        """
+        self._current_item_id = item_id
+        reuse = item_id is not None and item_id == self._bound_item_id and self._bound_segment_id
+        self._next_segment_id = self._bound_segment_id if reuse else None
+
+    def _remember_item_segment(self) -> None:
+        """Record the segment the text being captured is published under.
+
+        Done per capture rather than in ``_reset_state`` or only when a segment opens:
+        the legacy output resets after flushing (which would point the item at a segment
+        never published), and in the realtime flow the segment is opened by the empty
+        update sent when the user stops speaking, before the transcript that carries the
+        item id arrives.
+        """
+        if self._current_item_id is not None:
+            self._bound_item_id = self._current_item_id
+            self._bound_segment_id = self._current_id
+
     def set_participant(
         self,
         participant: rtc.Participant | str | None,
@@ -409,7 +476,8 @@ def set_participant(
         self._reset_state()
 
     def _reset_state(self) -> None:
-        self._current_id = utils.shortuuid("SG_")
+        self._current_id = self._next_segment_id or utils.shortuuid("SG_")
+        self._next_segment_id = None
         self._capturing = False
         self._latest_text = ""
         # per-segment markup stripping: delta streams strip incrementally (buffering a tag
@@ -470,6 +538,8 @@ async def capture_text(self, text: str) -> None:
             self._reset_state()
             self._capturing = True
 
+        self._remember_item_segment()
+
         # the raw text (expressive markup intact) arrives here; publish only the visible
         # text. Skip a chunk that strips to nothing (a partial tag still buffering, or a
         # markup-only token) so the transcript cadence isn't disturbed.
@@ -619,6 +689,10 @@ def set_participant(self, participant: rtc.Participant | str | None) -> None:
         for source in self.__outputs:
             source.set_participant(participant)
 
+    def _capture_item(self, item_id: str | None) -> None:
+        for source in self.__outputs:
+            source._capture_item(item_id)
+
     async def capture_text(self, text: str) -> None:
         await asyncio.gather(*[sink.capture_text(text) for sink in self.__outputs])
 

livekit-agents/livekit/agents/voice/room_io/room_io.py

@@ -366,6 +366,9 @@ async def _forward_user_transcript(
             if self._user_tr_output is None:
                 continue
 
+            # the chat context upserts revisions of an item onto one message; the
+            # rendered transcript has to follow, or the two views diverge
+            self._user_tr_output._capture_item(ev.item_id)
             await self._user_tr_output.capture_text(ev.transcript)
             if ev.is_final:
                 self._user_tr_output.flush()

tests/test_room_io.py

@@ -428,3 +428,156 @@ async def test_audio_output_playback_started_fires_once_across_pause_resume() ->
     assert len(started) == 1
 
     await utils.aio.cancel_and_wait(forward_task)
+
+
+@pytest.mark.asyncio
+async def test_refinalized_item_keeps_its_transcript_segment() -> None:
+    # a realtime provider can finalize the same input-audio item repeatedly, each time
+    # with a longer transcript. The chat context upserts those onto one message, so the
+    # rendered transcript has to reuse the segment too - otherwise the revisions stack
+    # up client-side as N segments, each containing its predecessors (#6710)
+    room = _FakeRoom()
+    output = _ParticipantStreamTranscriptionOutput(room=room, participant="user")
+
+    async def _finalize(item_id: str, transcript: str) -> str:
+        output._capture_item(item_id)
+        await output.capture_text(transcript)
+        segment_id = output._current_id
+        output.flush()
+        if output._flush_atask is not None:
+            await output._flush_atask
+        return segment_id
+
+    first = await _finalize("item_a", "hello")
+    revised = await _finalize("item_a", "hello there, how are you")
+
+    assert revised == first, "a revision of the same item must land in its own segment"
+
+    other = await _finalize("item_b", "different item")
+    assert other != first, "a new item must get a new segment"
+
+
+@pytest.mark.asyncio
+async def test_transcript_segments_are_unique_without_item_ids() -> None:
+    # STT paths don't carry an item id; each final has to open its own segment
+    room = _FakeRoom()
+    output = _ParticipantStreamTranscriptionOutput(room=room, participant="user")
+
+    segments = []
+    for transcript in ("first utterance", "second utterance"):
+        output._capture_item(None)
+        await output.capture_text(transcript)
+        segments.append(output._current_id)
+        output.flush()
+        if output._flush_atask is not None:
+            await output._flush_atask
+
+    assert segments[0] != segments[1]
+
+
+@pytest.mark.asyncio
+async def test_refinalized_item_keeps_its_segment_on_both_channels() -> None:
+    # the legacy output resets after flushing too, so the item's segment has to be
+    # recorded when a capture starts rather than on every reset - otherwise the legacy
+    # channel keeps stacking revisions while the stream channel replaces them
+    room = _FakeRoom()
+    output = _ParticipantTranscriptionOutput(room=room, is_delta_stream=False, participant="user")
+    legacy_output, stream_output = output._ParticipantTranscriptionOutput__outputs
+    legacy_output._track_id = "TR_fake"
+
+    async def _finalize(item_id: str, transcript: str) -> tuple[str, str]:
+        output._capture_item(item_id)
+        await output.capture_text(transcript)
+        ids = (legacy_output._current_id, stream_output._current_id)
+        output.flush()
+        if legacy_output._flush_task is not None:
+            await legacy_output._flush_task
+        if stream_output._flush_atask is not None:
+            await stream_output._flush_atask
+        return ids
+
+    first_legacy, first_stream = await _finalize("item_a", "hello")
+    revised_legacy, revised_stream = await _finalize("item_a", "hello there")
+
+    assert revised_legacy == first_legacy, "the legacy channel stacked the revision"
+    assert revised_stream == first_stream, "the stream channel stacked the revision"
+
+    other_legacy, other_stream = await _finalize("item_b", "next")
+    assert other_legacy != first_legacy
+    assert other_stream != first_stream
+
+
+@pytest.mark.asyncio
+async def test_unconsumed_reservation_does_not_leak_to_the_next_item() -> None:
+    # capture_text returns early until a participant is known, so a reservation made by
+    # _capture_item can go unused. It must not survive for whatever text comes next, or
+    # a new utterance publishes into the previous one's segment and replaces it on screen
+    room = _FakeRoom()
+    output = _ParticipantStreamTranscriptionOutput(room=room, participant="user")
+
+    output._capture_item("item_a")
+    await output.capture_text("hello")
+    first = output._current_id
+    output.flush()
+    if output._flush_atask is not None:
+        await output._flush_atask
+
+    # a revision is announced but never captured (no participant yet, empty chunk, ...)
+    output._capture_item("item_a")
+    # then a different item arrives
+    output._capture_item("item_b")
+    await output.capture_text("a new utterance")
+    second = output._current_id
+
+    assert second != first, "the new utterance overwrote the previous one's segment"
+
+
+@pytest.mark.asyncio
+async def test_realtime_flow_binds_the_item_to_the_open_segment() -> None:
+    # the realtime session opens the segment with an empty, item-less update when the
+    # user stops speaking; the transcript carrying the item id lands in that already
+    # open capture, so the item has to be bound then rather than only when a segment starts
+    room = _FakeRoom()
+    output = _ParticipantStreamTranscriptionOutput(room=room, participant="user")
+
+    output._capture_item(None)  # InputSpeechStopped: transcript="", is_final=False
+    await output.capture_text("")
+    output._capture_item("item_a")
+    await output.capture_text("hello")
+    first = output._current_id
+    output.flush()
+    if output._flush_atask is not None:
+        await output._flush_atask
+
+    output._capture_item("item_a")  # the provider revises the same item
+    await output.capture_text("hello there, how are you")
+    assert output._current_id == first, "the first revision still opened a second line"
+
+
+@pytest.mark.asyncio
+async def test_late_revision_does_not_hijack_a_newer_utterance() -> None:
+    # a correction arriving after the next utterance has started must not be written
+    # into that newer utterance's line
+    room = _FakeRoom()
+    output = _ParticipantStreamTranscriptionOutput(room=room, participant="user")
+
+    output._capture_item("item_a")
+    await output.capture_text("first utterance")
+    output.flush()
+    if output._flush_atask is not None:
+        await output._flush_atask
+
+    # the next utterance opens with the empty item-less update, then its own item
+    output._capture_item(None)
+    await output.capture_text("")
+    output._capture_item("item_b")
+    await output.capture_text("second utterance")
+    newer = output._current_id
+    output.flush()
+    if output._flush_atask is not None:
+        await output._flush_atask
+
+    output._capture_item("item_a")  # late correction of the older item
+    await output.capture_text("first utterance, corrected")
+
+    assert output._current_id != newer, "the late correction overwrote a newer utterance"