fix(tts): cancel() ends a reader following the request being written

assemblyai/__version__.py

@@ -1 +1 @@
-__version__ = "1.8.0"
+__version__ = "1.8.1"

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()

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()

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."

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]