Keep a send_audio() chunk whose send fails out of the realtime SpeechPart.audio recorded while it was sending (#10058)

pydantic_ai_slim/pydantic_ai/realtime/_core.py

@@ -113,16 +113,14 @@ class InputWithdrawn:
 
 @dataclass(frozen=True, kw_only=True)
 class AudioSent:
-    """Input audio the caller streamed, retained for the spoken turn it belongs to."""
+    """Input audio the caller streamed went out, retained for the spoken turn it belongs to.
+
+    Reported once its send completed: a turn the provider ended while it was on its way was ended without it.
+    """
 
     data: bytes
 
 
-@dataclass(frozen=True, kw_only=True)
-class AudioUnsent:
-    """The input audio last reported as sent never went out: it is taken back, as much of it as is still there."""
-
-
 @dataclass(frozen=True, kw_only=True)
 class AudioCleared:
     """The caller discarded the buffered input audio."""
@@ -185,7 +183,6 @@ Command: TypeAlias = (
     InputSent
     | InputWithdrawn
     | AudioSent
-    | AudioUnsent
     | AudioCleared
     | ToolReturned
     | ToolCallRefused
@@ -344,9 +341,6 @@ class SessionCore:
         self._recorded_audio: deque[_Response | _UserTurn] = deque()
         """Recorded entities whose message still has audio, in the order they joined the conversation: what eviction
         takes audio from first, oldest first, without passing over anything that has none or isn't recorded yet."""
-        self._unsent_audio: tuple[_UserTurn | None, int] | None = None
-        """The latest input audio chunk, until the next: the turn it was cut for since (or `None` while it is in the
-        buffer), and its length, so `AudioUnsent` can take it back if its send fails."""
         self._placements = 0
         self._user_speaking = False
         self._speech_segmented = False
@@ -433,7 +427,7 @@ class SessionCore:
             self._input_sent(command)
         elif isinstance(command, InputWithdrawn):
             self._withdraw(command.input_ids)
-        elif isinstance(command, (AudioSent, AudioUnsent, AudioCleared)):
+        elif isinstance(command, (AudioSent, AudioCleared)):
             self._input_audio_command(command)
         elif isinstance(command, ToolReturned):
             self._returns[command.tool_call_id] = command.request
@@ -458,38 +452,17 @@ class SessionCore:
         else:
             assert_never(command)
 
-    def _input_audio_command(self, command: AudioSent | AudioUnsent | AudioCleared) -> None:
+    def _input_audio_command(self, command: AudioSent | AudioCleared) -> None:
         if isinstance(command, AudioSent):
             if self._retain_input:
                 self._input_audio.extend(command.data)
-                self._unsent_audio = (None, len(command.data))
                 self._bound_retained_audio()
-        elif isinstance(command, AudioUnsent):
-            self._take_back_audio()
         else:
-            self._unsent_audio = None
             self._input_audio.clear()
             # A turn the user was saying that hadn't joined the conversation never will now.
             for turn_id in [turn_id for turn_id, turn in self._turns.items() if not turn.ended]:
                 self._drop_turn(turn_id)
 
-    def _take_back_audio(self) -> None:
-        """Take back the latest input audio chunk, whose send failed: from the buffer, or the turn it was cut for.
-
-        By length, as the session takes it back from its own buffer: what of it the budget already trimmed, or
-        that went into history with its turn, is no longer there to take back.
-        """
-        if self._unsent_audio is None:
-            return
-        turn, length = self._unsent_audio
-        self._unsent_audio = None
-        if turn is None:
-            del self._input_audio[len(self._input_audio) - min(length, len(self._input_audio)) :]
-        elif turn.id in self._waiting_turn_audio and turn.audio:
-            length = min(length, len(turn.audio))
-            turn.audio = turn.audio[: len(turn.audio) - length]
-            self._waiting_turn_audio_bytes -= length
-
     def _input_sent(self, command: InputSent) -> None:
         if command.request is not None:
             self._inputs[command.input_id] = self._unplaced[command.input_id] = _Input(
@@ -726,8 +699,6 @@ class SessionCore:
         self._input_audio.clear()
         self._waiting_turn_audio[turn.id] = turn
         self._waiting_turn_audio_bytes += len(turn.audio)
-        if self._unsent_audio is not None and self._unsent_audio[0] is None:
-            self._unsent_audio = (turn, self._unsent_audio[1])
 
     def _release_turn_audio(self, turn: _UserTurn) -> None:
         if self._waiting_turn_audio.pop(turn.id, None) is not None and turn.audio:

pydantic_ai_slim/pydantic_ai/realtime/_session.py

@@ -86,7 +86,6 @@ from ..usage import RequestUsage, RunUsage, UsageLimits
 from ._core import (
     AudioCleared,
     AudioSent,
-    AudioUnsent,
     Closed,
     CoreInput,
     ExchangeAbandoned,
@@ -2075,27 +2074,18 @@ class RealtimeSession:
         audio_was_uncommitted = self._audio_uncommitted
         self._held_commit.has_audio = True
         self._open_user_turn_for_audio()
-        previous_length: int | None = None
-        if self._retain_input:
-            # Buffer the raw input so the finalized user turn can retain it. A per-item speech-stopped
-            # boundary later cuts this into that turn's own segment (see `_segment_input_audio`); only the
-            # exact split at the boundary is approximate (see `audio_retention`).
-            previous_length = len(self._input_audio)
+        retained: bytes | None = None
         try:
-            # The extend lives inside the try so a chunk that is not bytes-like (possible since
-            # `send_audio` accepts an async iterable of chunks) rolls the user-turn state back below
-            # instead of leaving `_user_turn_active` set with nothing buffered or sent. A failing
-            # `bytearray.extend` leaves the buffer unchanged, so the length guard in the handler
-            # correctly skips the truncation for that case.
+            # Copied inside the try so a chunk that is not bytes-like (possible since `send_audio` accepts
+            # an async iterable of chunks) fails before it goes out and rolls the user-turn state back
+            # below, instead of leaving `_user_turn_active` set with nothing buffered or sent.
             if self._retain_input:
-                self._input_audio.extend(data)
+                retained = memoryview(data).tobytes()
             self._audio_uncommitted = True
             await self._send_frame(BinaryAudio(data=data, media_type='audio/pcm'))
         except BaseException as e:
             self._user_turn_active = user_turn_was_active
             self._audio_uncommitted = audio_was_uncommitted
-            if previous_length is not None and len(self._input_audio) == previous_length + len(data):
-                del self._input_audio[previous_length:]
             if (
                 isinstance(e, RealtimeError)
                 and isinstance(e.__cause__, self._connection.transport_errors)
@@ -2108,7 +2098,14 @@ class RealtimeSession:
                 # the reconnect fails, the next chunk raises again.
                 return
             raise
-        if self._retain_input:
+        if retained is not None:
+            # Buffer the raw input so the finalized user turn can retain it, once it has gone out (with no
+            # await since, so still in send order), as the session core takes it: a boundary the pump handled
+            # while the chunk was on its way was drawn without it, and a chunk whose send failed belongs to
+            # no turn. A per-item speech-stopped boundary later cuts this into that turn's own segment (see
+            # `_segment_input_audio`); only the exact split at the boundary is approximate (see
+            # `audio_retention`).
+            self._input_audio.extend(retained)
             self._bound_retained_audio()
 
     def _open_user_turn_for_audio(self) -> None:
@@ -2408,7 +2405,6 @@ class RealtimeSession:
         # which is a caller that asked to send.
         async with self._send_lock:
             first_input = self._inputs_sent
-            audio_sent_to_core = False
             try:
                 for position, content in enumerate(contents):
                     # Numbered before the call, and whether or not it raises, matching how
@@ -2426,17 +2422,15 @@ class RealtimeSession:
                             tool_call_id=content.tool_call_id if isinstance(content, ToolResult) else None,
                         )
                     )
-                    if isinstance(content, BinaryAudio):
-                        # Under the lock, so the core takes audio in the order it goes out and a failed send
-                        # takes back its own chunk; before the send, as the pump can end the turn meanwhile.
-                        self._apply_core(AudioSent(data=bytes(content.data)))
-                        audio_sent_to_core = True
                     await self._connection.send(content)
+                    if isinstance(content, BinaryAudio):
+                        # Under the lock, so the core takes audio in the order it goes out; only once it has
+                        # gone out, as a turn the provider ended while it was on its way (which the pump can
+                        # record meanwhile) was ended without it, and a chunk whose send fails is in no turn.
+                        self._apply_core(AudioSent(data=bytes(content.data)))
             except BaseException as e:
                 # The caller takes back what it sent, so the core does too.
                 self._apply_core(InputWithdrawn(input_ids=tuple(range(first_input, self._inputs_sent))))
-                if audio_sent_to_core:
-                    self._apply_core(AudioUnsent())
                 if not isinstance(e, self._connection.transport_errors):
                     raise
                 # A send that fails because the link is gone is the same failure the receive side

tests/realtime/test_openai.py

@@ -5375,7 +5375,7 @@ async def test_a_conversation_id_resolved_late_reaches_the_core_history() -> Non
 
 
 @pytest.mark.anyio
-async def test_audio_that_fails_to_go_out_is_taken_back_from_the_core_too() -> None:
+async def test_audio_that_fails_to_go_out_never_reaches_the_core() -> None:
     class _FailingAudio(_QueuedWebSocket):
         async def send(self, data: str) -> None:
             raise OSError('gone')  # (the audio is all this session sends)
@@ -5397,8 +5397,8 @@ async def test_audio_that_fails_to_go_out_is_taken_back_from_the_core_too() -> N
 
 
 @pytest.mark.anyio
-async def test_a_failed_send_takes_back_its_own_audio_while_another_waits_to_go_out() -> None:
-    """Two chunks sent at once go out one after the other: the first failing takes back its own audio, not the second's."""
+async def test_a_failed_send_leaves_its_audio_out_while_another_waits_to_go_out() -> None:
+    """Two chunks sent at once go out one after the other: the first failing never adds its audio, and the second's goes in."""
 
     class _FirstAudioFails(_QueuedWebSocket):
         def __init__(self) -> None:
@@ -5437,6 +5437,74 @@ async def test_a_failed_send_takes_back_its_own_audio_while_another_waits_to_go_
         assert bytes(core._input_audio) == second  # pyright: ignore[reportPrivateUsage]
 
 
+@pytest.mark.anyio
+async def test_audio_whose_send_fails_after_its_turn_is_recorded_stays_out_of_that_turn() -> None:
+    """A turn recorded while a chunk is still going out doesn't keep that chunk when its send then fails.
+
+    The provider committed its buffer before it had the chunk, so the chunk is no part of that turn whether or not
+    its send succeeds: history and the streamed part both hold only the audio that went out ahead of the commit.
+    """
+
+    class _SecondAudioFails(_QueuedWebSocket):
+        def __init__(self) -> None:
+            super().__init__()
+            self.sending = asyncio.Event()
+            self.fail = asyncio.Event()
+            self.sends = 0
+
+        async def send(self, data: str) -> None:
+            self.sends += 1
+            if self.sends == 2:  # (the audio is all this session sends)
+                self.sending.set()
+                await self.fail.wait()
+                raise OSError('gone')
+            await super().send(data)
+
+    ws = _SecondAudioFails()
+    connection = OpenAIRealtimeConnection(ws)  # type: ignore[arg-type]
+    session = RealtimeSession(
+        connection,
+        model=FakeRealtimeModel(connection, system='openai'),
+        tool_manager=make_tool_manager(),
+        audio_retention='input_audio',
+    )
+    sent, unsent = b'\x01\x00' * 100, b'\x02\x00' * 50
+    streamed: list[SpeechPart] = []
+    async with session:
+        await session.send_audio(sent)
+        failing_send = asyncio.create_task(session.send_audio(unsent))
+        await ws.sending.wait()
+        for frame in [
+            {'type': 'input_audio_buffer.speech_started', 'item_id': 'item_u1', 'audio_start_ms': 0},
+            {'type': 'input_audio_buffer.speech_stopped', 'item_id': 'item_u1', 'audio_end_ms': 500},
+            {'type': 'input_audio_buffer.committed', 'item_id': 'item_u1', 'previous_item_id': None},
+            {
+                'type': 'conversation.item.input_audio_transcription.completed',
+                'item_id': 'item_u1',
+                'content_index': 0,
+                'transcript': 'Hi.',
+            },
+        ]:
+            ws.push(frame)
+        async for event in session:  # pragma: no branch
+            if isinstance(event, PartEndEvent) and isinstance(event.part, SpeechPart):
+                streamed.append(event.part)
+                break
+        # The turn is recorded while the chunk is still going out; then its send fails.
+        assert session.all_messages()
+        ws.fail.set()
+        with pytest.raises(RealtimeError):
+            await failing_send
+        history = session.all_messages()
+
+    [request] = history
+    [part] = request.parts
+    assert isinstance(part, SpeechPart) and part.transcript == 'Hi.' and part.audio is not None
+    assert part.audio.data[44:] == sent
+    [streamed_part] = streamed
+    assert streamed_part.audio is not None and streamed_part.audio.data[44:] == sent
+
+
 @pytest.mark.anyio
 async def test_retained_audio_is_kept_once_when_the_core_keeps_history() -> None:
     """The session core's history keeps retained audio; the session's own record (what tools see) lets its copy go.

tests/realtime/test_session.py

@@ -5863,6 +5863,12 @@ async def test_text_request_reserved_before_response_finishes_during_send() -> N
 
 
 async def test_send_audio_reserved_before_speech_boundary_during_send() -> None:
+    """The turn audio opens is reserved before the chunk goes out, so a boundary handled meanwhile still records it.
+
+    The chunk itself isn't in that turn: the provider drew the boundary before it had the chunk, so the chunk
+    is retained for the next turn instead.
+    """
+
     class _BoundarySend(FakeRealtimeConnection):
         async def send(self, content: RealtimeInput) -> None:
             self.sent.append(content)
@@ -5872,8 +5878,9 @@ async def test_send_audio_reserved_before_speech_boundary_during_send() -> None:
     session = RealtimeSession(conn, audio_retention='input_audio')
     await session.send_audio(b'\xaa\xbb')
     assert session.new_messages() == snapshot(
-        [ModelRequest(parts=[SpeechPart(speaker='user', audio=_wav_content(b'\xaa\xbb'))], timestamp=IsDatetime())]
+        [ModelRequest(parts=[SpeechPart(speaker='user')], timestamp=IsDatetime())]
     )
+    assert bytes(session._input_audio) == b'\xaa\xbb'  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_failed_sends_leave_no_phantom_history_or_audio() -> None:
@@ -5913,6 +5920,35 @@ async def test_failed_sends_leave_no_phantom_history_or_audio() -> None:
     assert unretained.new_messages() == []
 
 
+async def test_audio_whose_send_fails_after_its_turn_is_recorded_stays_out_of_that_turn() -> None:
+    """A turn recorded while a chunk is still going out doesn't keep that chunk when its send then fails."""
+
+    class _TurnEndsDuringFailingSend(FakeRealtimeConnection):
+        fail = False
+
+        async def send(self, content: RealtimeInput) -> None:
+            if self.fail:
+                session._translate_event(InputTranscript(text='Hi.', is_final=True))  # pyright: ignore[reportPrivateUsage]
+                raise RuntimeError('send failed')
+            self.sent.append(content)
+
+    conn = _TurnEndsDuringFailingSend([])
+    session = RealtimeSession(conn, audio_retention='input_audio')
+    await session.send_audio(b'\xaa\xbb')
+    conn.fail = True
+    with pytest.raises(RuntimeError, match='send failed'):
+        await session.send_audio(b'\xcc\xdd')
+    assert session.new_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[SpeechPart(speaker='user', transcript='Hi.', audio=_wav_content(b'\xaa\xbb'))],
+                timestamp=IsDatetime(),
+            )
+        ]
+    )
+    assert not session._input_audio  # pyright: ignore[reportPrivateUsage]
+
+
 async def test_transport_failure_while_sending_becomes_a_realtime_error() -> None:
     # A send that fails because the link is gone is the same failure the receive side reports, so it
     # surfaces as `RealtimeError` (a `ModelAPIError`) instead of leaking the transport's own exception.

tests/realtime/test_session_core.py

@@ -33,7 +33,6 @@ from pydantic_ai.messages import (
 from pydantic_ai.realtime._core import (
     AudioCleared,
     AudioSent,
-    AudioUnsent,
     Closed,
     CoreInput,
     ExchangeAbandoned,
@@ -894,23 +893,6 @@ def test_a_conversation_id_resolved_late_reaches_what_the_core_recorded() -> Non
     assert {message.conversation_id for message in session_core.all_messages()} == {'c1'}
 
 
-def test_input_audio_that_never_went_out_is_taken_back() -> None:
-    session_core = feed(
-        core(input_transcription_enabled=False, retain_input_audio=True),
-        AudioSent(data=b'\x01\x00'),
-        AudioSent(data=b'\x02\x00'),
-        AudioUnsent(),
-        # Only the latest chunk is taken back, once.
-        AudioUnsent(),
-        UserTurnStarted(turn_id='u1'),
-        UserTurnEnded(turn_id='u1'),
-    )
-    [request] = session_core.all_messages()
-    [part] = request.parts
-    assert isinstance(part, SpeechPart) and part.audio is not None
-    assert part.audio.data[44:] == b'\x01\x00'
-
-
 def test_a_reconnect_replays_a_held_turns_audio_without_counting_it() -> None:
     """The replay builds the held turn's request with its audio, which isn't retained history, so isn't counted."""
     session_core = feed(
@@ -951,25 +933,6 @@ def test_retained_audio_eviction_passes_over_what_is_not_recorded() -> None:
     assert list(session_core._recorded_audio) == [session_core._turns['u5']]  # pyright: ignore[reportPrivateUsage]
 
 
-def test_input_audio_cut_for_a_turn_before_its_send_failed_is_taken_back_from_the_turn() -> None:
-    session_core = feed(
-        core(retain_input_audio=True),
-        AudioSent(data=b'\x01\x00'),
-        AudioSent(data=b'\x02\x00'),
-        RealtimeInputSpeechStartEvent(item_id='u1'),
-        UserTurnStarted(turn_id='u1'),
-        # The speech end cuts the turn's audio while the last chunk is still on its way, and its send then fails.
-        RealtimeInputSpeechEndEvent(item_id='u1'),
-        AudioUnsent(),
-        UserTurnEnded(turn_id='u1'),
-        InputTranscript('Hi.', item_id='u1', is_final=True),
-    )
-    [request] = session_core.all_messages()
-    [part] = request.parts
-    assert isinstance(part, SpeechPart) and part.audio is not None
-    assert part.audio.data[44:] == b'\x01\x00'
-
-
 def test_retained_audio_trims_the_oldest_audio_coming_in_first() -> None:
     """A response whose incoming audio was emptied counts as newest when more comes in for it."""
     session_core = feed(
@@ -1007,38 +970,27 @@ def test_retained_audio_evicted_is_let_go_by_the_core_too() -> None:
     assert len(session_core._recorded_audio) == 2  # pyright: ignore[reportPrivateUsage]
 
 
-def test_a_failed_send_of_audio_trimmed_by_the_budget_leaves_nothing_of_it() -> None:
-    """A chunk longer than the whole budget is trimmed before its send fails: what is left of it is taken back."""
-    session_core = feed(
-        core(input_transcription_enabled=False, retain_input_audio=True, retain_audio_max_seconds=0.1),
-        AudioSent(data=_tenth_of_a_second(1) * 2),
-        AudioUnsent(),
-        AudioSent(data=b'\x05\x00'),
-        UserTurnStarted(turn_id='u1'),
-        UserTurnEnded(turn_id='u1'),
-    )
-    [request] = session_core.all_messages()
-    [part] = request.parts
-    assert isinstance(part, SpeechPart) and part.audio is not None
-    assert part.audio.data[44:] == b'\x05\x00'
+def test_input_audio_that_went_out_after_its_turn_was_recorded_is_left_to_the_next() -> None:
+    """Audio is reported once its send completed, so a chunk still on its way when a turn is recorded comes after it.
 
-
-def test_audio_taken_back_from_a_turn_no_longer_waiting_leaves_it_alone() -> None:
-    """Once its turn is recorded, the core doesn't change the message to take the audio back.
-
-    A known gap, not the intent (the current session core has it too): a turn recorded while the send of its last
-    chunk is still under way keeps that chunk's audio if the send then fails.
+    The provider ended that turn without the chunk, which is the next turn's start, or, if its send fails, no
+    turn's: the recorded turn keeps only the audio that went out ahead of it either way.
     """
     session_core = feed(
         core(input_transcription_enabled=False, retain_input_audio=True),
         AudioSent(data=b'\x01\x00'),
         UserTurnStarted(turn_id='u1'),
         UserTurnEnded(turn_id='u1'),
-        AudioUnsent(),
+        AudioSent(data=b'\x02\x00'),
+        UserTurnStarted(turn_id='u2'),
+        UserTurnEnded(turn_id='u2'),
     )
-    [request] = session_core.all_messages()
-    [part] = request.parts
-    assert isinstance(part, SpeechPart) and part.audio is not None
+    assert [
+        part.audio.data[44:]
+        for message in session_core.all_messages()
+        for part in message.parts
+        if isinstance(part, SpeechPart) and part.audio is not None
+    ] == [b'\x01\x00', b'\x02\x00']
 
 
 def test_retained_audio_eviction_does_not_walk_what_waits_or_has_no_audio(monkeypatch: pytest.MonkeyPatch) -> None: