add asyncio timeout() context manager

This is an implementation of Python's `asyncio.timeout()` which was introduced
in 3.11. We strongly encourage the use of the following approach instead of
`asyncio.wait_for()`.

    async with timeout(delay):
      # Your async code

In Python 3.12, `asyncio.wait_for()` is implemented in terms of
`asyncio.timeout()` which fixed a bunch of issues. However, this was never
backported (because of the lack of `async.timeout()`) and there are still many
remainig issues, specially in Python 3.10, in `async.wait_for()`.

See python/cpython#98518

Author: aconchillo

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added `pipecat.utils.timeout.timeout()` asynchronous context manager. This
+  should be preferred over `asyncio.wait_for()`.
+
 - Allow passing custom pipeline sink and source processors to a
   `Pipeline`. Pipeline source and sink processors are used to know and control
   what's coming in and out of a `Pipeline` processor.
@@ -46,6 +49,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Replaced `asyncio.wait_for()` for `async with timeout(delay):`. In Python
+  3.10, `asyncio.wait_for()` has issues regarding task cancellation
+  (i.e. cancellation is never propagated).
+  See https://bugs.python.org/issue42130
+
 - Fixed an `AudioBufferProcessor` issues that would cause audio overlap when
   setting a max buffer size.
 

scripts/evals/eval.py

@@ -43,6 +43,7 @@
 from pipecat.services.llm_service import FunctionCallParams
 from pipecat.services.openai.llm import OpenAILLMService
 from pipecat.transports.services.daily import DailyParams, DailyTransport
+from pipecat.utils.asyncio.timeout import timeout
 
 SCRIPT_DIR = Path(__file__).resolve().parent
 
@@ -122,7 +123,8 @@ async def run_eval(self, example_file: str, prompt: EvalPrompt, eval: Optional[s
             logger.error(f"ERROR: Unable to run {example_file}: {e}")
 
         try:
-            result = await asyncio.wait_for(self._queue.get(), timeout=1.0)
+            async with timeout(1.0):
+                result = await self._queue.get()
         except asyncio.TimeoutError:
             result = False
 

src/pipecat/pipeline/task.py

@@ -51,6 +51,7 @@
     TaskManager,
     TaskManagerParams,
 )
+from pipecat.utils.asyncio.timeout import timeout
 from pipecat.utils.asyncio.watchdog_queue import WatchdogQueue
 from pipecat.utils.tracing.setup import is_tracing_available
 from pipecat.utils.tracing.turn_trace_observer import TurnTraceObserver
@@ -654,7 +655,8 @@ async def _heartbeat_monitor_handler(self):
         wait_time = HEARTBEAT_MONITOR_SECONDS
         while True:
             try:
-                frame = await asyncio.wait_for(self._heartbeat_queue.get(), timeout=wait_time)
+                async with timeout(wait_time):
+                    frame = await self._heartbeat_queue.get()
                 process_time = (self._clock.get_time() - frame.timestamp) / 1_000_000_000
                 logger.trace(f"{self}: heartbeat frame processed in {process_time} seconds")
                 self._heartbeat_queue.task_done()
@@ -677,9 +679,8 @@ async def _idle_monitor_handler(self):
 
         while running:
             try:
-                frame = await asyncio.wait_for(
-                    self._idle_queue.get(), timeout=self._idle_timeout_secs
-                )
+                async with timeout(self._idle_timeout_secs):
+                    frame = await self._idle_queue.get()
 
                 if not isinstance(frame, InputAudioRawFrame):
                     frame_buffer.append(frame)

src/pipecat/processors/aggregators/dtmf_aggregator.py

@@ -25,6 +25,7 @@
     TranscriptionFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
+from pipecat.utils.asyncio.timeout import timeout
 from pipecat.utils.time import time_now_iso8601
 
 
@@ -136,7 +137,8 @@ async def _aggregation_task_handler(self):
         """Background task that handles timeout-based flushing."""
         while True:
             try:
-                await asyncio.wait_for(self._digit_event.wait(), timeout=self._idle_timeout)
+                async with timeout(self._idle_timeout):
+                    await self._digit_event.wait()
                 self._digit_event.clear()
             except asyncio.TimeoutError:
                 self.reset_watchdog()

src/pipecat/processors/aggregators/llm_response.py

@@ -60,6 +60,7 @@
     OpenAILLMContextFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
+from pipecat.utils.asyncio.timeout import timeout
 from pipecat.utils.time import time_now_iso8601
 
 
@@ -660,17 +661,18 @@ async def _aggregation_task_handler(self):
                 #   consistent user experience between real VAD detection and
                 #   emulated VAD scenarios.
                 if not self._emulating_vad:
-                    timeout = self._params.aggregation_timeout
+                    delay = self._params.aggregation_timeout
                 elif self._turn_params:
-                    timeout = self._params.turn_emulated_vad_timeout
+                    delay = self._params.turn_emulated_vad_timeout
                 else:
                     # Use VAD stop_secs when no turn analyzer is present, fallback if no VAD params
-                    timeout = (
+                    delay = (
                         self._vad_params.stop_secs
                         if self._vad_params
                         else self._params.turn_emulated_vad_timeout
                     )
-                await asyncio.wait_for(self._aggregation_event.wait(), timeout)
+                async with timeout(delay):
+                    await self._aggregation_event.wait()
                 await self._maybe_emulate_user_speaking()
             except asyncio.TimeoutError:
                 if not self._user_speaking:

src/pipecat/processors/idle_frame_processor.py

@@ -11,6 +11,7 @@
 
 from pipecat.frames.frames import Frame, StartFrame
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
+from pipecat.utils.asyncio.timeout import timeout
 from pipecat.utils.asyncio.watchdog_event import WatchdogEvent
 
 
@@ -85,7 +86,8 @@ async def _idle_task_handler(self):
         """Handle idle timeout monitoring and callback execution."""
         while True:
             try:
-                await asyncio.wait_for(self._idle_event.wait(), timeout=self._timeout)
+                async with timeout(self._timeout):
+                    await self._idle_event.wait()
             except asyncio.TimeoutError:
                 await self._callback(self)
             finally:

src/pipecat/processors/user_idle_processor.py

@@ -22,6 +22,7 @@
     UserStoppedSpeakingFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
+from pipecat.utils.asyncio.timeout import timeout
 from pipecat.utils.asyncio.watchdog_event import WatchdogEvent
 
 
@@ -191,7 +192,8 @@ async def _idle_task_handler(self) -> None:
         """
         while True:
             try:
-                await asyncio.wait_for(self._idle_event.wait(), timeout=self._timeout)
+                async with timeout(self._timeout):
+                    await self._idle_event.wait()
             except asyncio.TimeoutError:
                 if not self._interrupted:
                     self._retry_count += 1

src/pipecat/services/anthropic/llm.py

@@ -53,6 +53,7 @@
 )
 from pipecat.processors.frame_processor import FrameDirection
 from pipecat.services.llm_service import FunctionCallFromLLM, LLMService
+from pipecat.utils.asyncio.timeout import timeout
 from pipecat.utils.asyncio.watchdog_async_iterator import WatchdogAsyncIterator
 from pipecat.utils.tracing.service_decorators import traced_llm
 
@@ -185,9 +186,8 @@ async def _create_message_stream(self, api_call, params):
         """
         if self._retry_on_timeout:
             try:
-                response = await asyncio.wait_for(
-                    api_call(**params), timeout=self._retry_timeout_secs
-                )
+                async with timeout(self._retry_timeout_secs):
+                    response = await api_call(**params)
                 return response
             except (APITimeoutError, asyncio.TimeoutError):
                 # Retry, this time without a timeout so we get a response

src/pipecat/services/assemblyai/stt.py

@@ -31,6 +31,7 @@
 from pipecat.processors.frame_processor import FrameDirection
 from pipecat.services.stt_service import STTService
 from pipecat.transcriptions.language import Language
+from pipecat.utils.asyncio.timeout import timeout
 from pipecat.utils.time import time_now_iso8601
 from pipecat.utils.tracing.service_decorators import traced_stt
 
@@ -219,10 +220,8 @@ async def _disconnect(self):
                 await self._websocket.send(json.dumps({"type": "Terminate"}))
 
                 try:
-                    await asyncio.wait_for(
-                        self._termination_event.wait(),
-                        timeout=5.0,
-                    )
+                    async with timeout(5.0):
+                        await self._termination_event.wait()
                 except asyncio.TimeoutError:
                     logger.warning("Timed out waiting for termination message from server")
 
@@ -247,7 +246,8 @@ async def _receive_task_handler(self):
         try:
             while self._connected:
                 try:
-                    message = await asyncio.wait_for(self._websocket.recv(), timeout=1.0)
+                    async with timeout(1.0):
+                        message = await self._websocket.recv()
                     data = json.loads(message)
                     await self._handle_message(data)
                 except asyncio.TimeoutError:

src/pipecat/services/aws/llm.py

@@ -52,6 +52,7 @@
 )
 from pipecat.processors.frame_processor import FrameDirection
 from pipecat.services.llm_service import LLMService
+from pipecat.utils.asyncio.timeout import timeout
 from pipecat.utils.tracing.service_decorators import traced_llm
 
 try:
@@ -801,9 +802,8 @@ async def _create_converse_stream(self, client, request_params):
         """
         if self._retry_on_timeout:
             try:
-                response = await asyncio.wait_for(
-                    client.converse_stream(**request_params), timeout=self._retry_timeout_secs
-                )
+                async with timeout(self._retry_timeout_secs):
+                    response = await client.converse_stream(**request_params)
                 return response
             except (ReadTimeoutError, asyncio.TimeoutError) as e:
                 # Retry, this time without a timeout so we get a response