PipelineTask: automatically cancel tasks if pipeline is idle

tests/test_pipeline.py

@@ -5,6 +5,7 @@
 #
 
 import asyncio
+import time
 import unittest
 
 from pipecat.frames.frames import EndFrame, HeartbeatFrame, StartFrame, TextFrame
@@ -100,7 +101,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
 
         identity = IdentityFilter()
         pipeline = Pipeline([identity])
-        task = PipelineTask(pipeline)
+        task = PipelineTask(pipeline, cancel_on_idle_timeout=False)
         task.set_event_loop(asyncio.get_event_loop())
         task.set_reached_upstream_filter((TextFrame,))
         task.set_reached_downstream_filter((TextFrame,))
@@ -123,7 +124,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
         await task.queue_frame(TextFrame(text="Hello Downstream!"))
 
         try:
-            await asyncio.wait_for(task.run(), timeout=1.0)
+            await asyncio.wait_for(asyncio.shield(task.run()), timeout=1.0)
         except asyncio.TimeoutError:
             pass
 
@@ -149,6 +150,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
                 heartbeats_period_secs=0.2,
             ),
             observers=[heartbeats_observer],
+            cancel_on_idle_timeout=False,
         )
         task.set_event_loop(asyncio.get_event_loop())
 
@@ -156,7 +158,90 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
 
         await task.queue_frame(TextFrame(text="Hello!"))
         try:
-            await asyncio.wait_for(task.run(), timeout=1.0)
+            await asyncio.wait_for(asyncio.shield(task.run()), timeout=1.0)
         except asyncio.TimeoutError:
             pass
         assert heartbeats_counter == expected_heartbeats
+
+    async def test_idle_task(self):
+        identity = IdentityFilter()
+        pipeline = Pipeline([identity])
+        task = PipelineTask(pipeline, idle_timeout_secs=0.2)
+        task.set_event_loop(asyncio.get_event_loop())
+        await task.run()
+        assert True
+
+    async def test_no_idle_task(self):
+        identity = IdentityFilter()
+        pipeline = Pipeline([identity])
+        task = PipelineTask(pipeline, idle_timeout_secs=0.2, cancel_on_idle_timeout=False)
+        task.set_event_loop(asyncio.get_event_loop())
+        try:
+            await asyncio.wait_for(asyncio.shield(task.run()), timeout=0.3)
+        except asyncio.TimeoutError:
+            assert True
+        else:
+            assert False
+
+    async def test_idle_task_heartbeats(self):
+        identity = IdentityFilter()
+        pipeline = Pipeline([identity])
+        task = PipelineTask(
+            pipeline,
+            params=PipelineParams(
+                enable_heartbeats=True,
+                heartbeats_period_secs=0.1,
+            ),
+            idle_timeout_secs=0.3,
+        )
+        task.set_event_loop(asyncio.get_event_loop())
+        await task.run()
+        assert True
+
+    async def test_idle_task_event_handler(self):
+        identity = IdentityFilter()
+        pipeline = Pipeline([identity])
+        task = PipelineTask(pipeline, idle_timeout_secs=0.2, cancel_on_idle_timeout=False)
+        task.set_event_loop(asyncio.get_event_loop())
+
+        idle_timeout = False
+
+        @task.event_handler("on_idle_timeout")
+        async def on_idle_timeout(task: PipelineTask):
+            nonlocal idle_timeout
+            idle_timeout = True
+            await task.cancel()
+
+        await task.run()
+        assert True
+
+    async def test_idle_task_frames(self):
+        idle_timeout_secs = 0.2
+        sleep_time_secs = idle_timeout_secs / 2
+
+        identity = IdentityFilter()
+        pipeline = Pipeline([identity])
+        task = PipelineTask(
+            pipeline,
+            idle_timeout_secs=idle_timeout_secs,
+            idle_timeout_frames=(TextFrame,),
+        )
+        task.set_event_loop(asyncio.get_event_loop())
+
+        async def delayed_frames():
+            await asyncio.sleep(sleep_time_secs)
+            await task.queue_frame(TextFrame("Hello Pipecat!"))
+            await asyncio.sleep(sleep_time_secs)
+            await task.queue_frame(TextFrame("Hello Pipecat!"))
+            await asyncio.sleep(sleep_time_secs)
+            await task.queue_frame(TextFrame("Hello Pipecat!"))
+
+        start_time = time.time()
+
+        tasks = {asyncio.create_task(task.run()), asyncio.create_task(delayed_frames())}
+
+        await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+
+        diff_time = time.time() - start_time
+
+        self.assertGreater(diff_time, sleep_time_secs * 3)