improve gracefully_cancel logic

.changeset/afraid-dolls-enjoy.md

@@ -0,0 +1,5 @@
+---
+"livekit-agents": patch
+---
+
+improve gracefully_cancel logic

livekit-agents/livekit/agents/utils/aio/__init__.py

@@ -1,5 +1,5 @@
 import asyncio
-import contextlib
+import functools
 
 from . import debug, duplex_unix, itertools
 from .channel import Chan, ChanClosed, ChanReceiver, ChanSender
@@ -9,11 +9,28 @@
 
 
 async def gracefully_cancel(*futures: asyncio.Future):
-    for f in futures:
-        f.cancel()
+    loop = asyncio.get_running_loop()
+    waiters = []
 
-    with contextlib.suppress(asyncio.CancelledError):
-        await asyncio.gather(*futures)
+    for fut in futures:
+        waiter = loop.create_future()
+        cb = functools.partial(_release_waiter, waiter)
+        waiters.append((waiter, cb))
+        fut.add_done_callback(cb)
+        fut.cancel()
+
+    try:
+        for waiter in waiters:
+            await waiter
+    finally:
+        for i, fut in enumerate(futures):
+            _, cb = waiters[i]
+            fut.remove_done_callback(cb)
+
+
+def _release_waiter(waiter, *args):
+    if not waiter.done():
+        waiter.set_result(None)
 
 
 __all__ = [

livekit-plugins/livekit-plugins-google/livekit/plugins/google/stt.py

@@ -16,7 +16,6 @@
 
 import asyncio
 import dataclasses
-import os
 from dataclasses import dataclass
 from typing import AsyncIterable, List, Union
 
@@ -121,6 +120,7 @@ def _recognizer(self) -> str:
             project_id = self._ensure_client().transport._credentials.project_id  # type: ignore
         except AttributeError:
             from google.auth import default as ga_default
+
             _, project_id = ga_default()
         return f"projects/{project_id}/locations/global/recognizers/_"