diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index 29654ee..2d986fc 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.8.0" +__version__ = "1.8.1" diff --git a/assemblyai/tts/v1/async_session.py b/assemblyai/tts/v1/async_session.py index e241f57..094502a 100644 --- a/assemblyai/tts/v1/async_session.py +++ b/assemblyai/tts/v1/async_session.py @@ -101,7 +101,7 @@ def __init__( self._queue_ready: Any = None self._next_flush_id = 0 self._cancel_ids: List[int] = [] - self._cancelled_before = 0 + self._cancelled_before = -1 self._drop_warned = False self.begin: Optional[TTSBeginEvent] = None self.termination: Optional[TTSTerminationEvent] = None @@ -139,7 +139,7 @@ async def connect(self) -> TTSBeginEvent: self._queue_ready = asyncio.Event() self._next_flush_id = 0 self._cancel_ids = [] - self._cancelled_before = 0 + self._cancelled_before = -1 self._drop_warned = False self.begin = None self.termination = None @@ -306,7 +306,9 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: while True: if self._queue is not connection: return - if flush_id < self._cancelled_before: # dropped by a Cancel + # A Cancel drops every request before it and the text being + # written, whose id it takes. + if flush_id <= self._cancelled_before: return item = await self._get(timeout=0.5) if item is None: @@ -323,7 +325,7 @@ async def audio(self, flush_id: Optional[int] = None) -> AsyncIterator[bytes]: if ident == flush_id: return elif kind == "cancelled": - if ident > flush_id: + if ident >= flush_id: return elif kind == "error": raise self._surfaced() diff --git a/assemblyai/tts/v1/session.py b/assemblyai/tts/v1/session.py index 1eabaa6..0251ddb 100644 --- a/assemblyai/tts/v1/session.py +++ b/assemblyai/tts/v1/session.py @@ -161,8 +161,8 @@ def _reset(self) -> None: self._queue_ready = threading.Condition() self._next_flush_id = 0 self._cancel_ids: List[int] = [] - # Requests with a lower flush_id than this were dropped by a Cancel. - self._cancelled_before = 0 + # Requests up to this flush_id were dropped by a Cancel (-1: none). + self._cancelled_before = -1 self._drop_warned = False self.begin = None self.termination = None @@ -401,7 +401,9 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: while True: if self._queue is not connection: return - if flush_id < self._cancelled_before: # dropped by a Cancel + # A Cancel drops every request before it and the text being + # written, whose id it takes. + if flush_id <= self._cancelled_before: return item = self._get(timeout=0.5) if item is None: @@ -418,7 +420,7 @@ def audio(self, flush_id: Optional[int] = None) -> Iterator[bytes]: if ident == flush_id: return elif kind == "cancelled": - if ident > flush_id: # a Cancel sent after this request + if ident >= flush_id: # a Cancel sent after this request return elif kind == "error": raise self._surfaced() diff --git a/tests/unit/test_tts_session.py b/tests/unit/test_tts_session.py index b96608a..5d9985c 100644 --- a/tests/unit/test_tts_session.py +++ b/tests/unit/test_tts_session.py @@ -995,3 +995,25 @@ def on_flush_done(event): session.generate("late") assert exc_info.value.error_code == "input_rate_exceeded" assert blocked_for.get("s", 0) < 0.5 + + +def test_cancel_ends_a_reader_following_the_request_being_written(server): + import threading + + server.answer_flush = False + with aai.TTSSession() as session: + got: list = [] + reader = threading.Thread(target=lambda: got.extend(session.audio())) + reader.start() # no id: follows the text being written + session.generate("Interrupted before the flush.") + session.cancel() + reader.join(timeout=3) + assert not reader.is_alive() # ended at once instead of hanging + assert got == [] + server.answer_flush = True + assert b"".join(session.speak("Next.")) == b"Next." + + +def test_the_first_request_is_not_dropped_by_the_initial_watermark(server): + with aai.TTSSession() as session: + assert b"".join(session.speak("Zero.")) == b"Zero." diff --git a/tests/unit/test_tts_session_async.py b/tests/unit/test_tts_session_async.py index 4c4c5f3..40d9638 100644 --- a/tests/unit/test_tts_session_async.py +++ b/tests/unit/test_tts_session_async.py @@ -683,3 +683,19 @@ async def connect(uri, additional_headers): assert [c async for c in stale] == [] assert b"".join([c async for c in session.speak("NEW")]) == b"NEW" await session.terminate() + + +async def test_cancel_ends_a_reader_following_the_request_being_written(server): + server.answer_flush = False + async with aai.AsyncTTSSession() as session: + reading = asyncio.ensure_future(_collect(session.audio())) + await asyncio.sleep(0) # let the reader start before the text + await session.generate("Interrupted before the flush.") + await session.cancel() + assert await asyncio.wait_for(reading, 3) == [] + server.answer_flush = True + assert b"".join([c async for c in session.speak("Next.")]) == b"Next." + + +async def _collect(iterator): + return [chunk async for chunk in iterator]