fix(workers): cancel jobs when interrupted during context exit

changelog/5721.fixed.md

@@ -0,0 +1 @@
+- Fixed `job()` and `job_group()` leaving child jobs running when the caller is cancelled while waiting for responses on context exit.

src/pipecat/pipeline/job_context.py

@@ -261,8 +261,8 @@ class JobGroupContext:
 
     On normal completion, results are available via ``responses``.
     On worker error (with ``cancel_on_error=True``) or timeout, raises
-    ``JobGroupError``. If the ``async with`` block raises, remaining
-    jobs are cancelled.
+    ``JobGroupError``. If the ``async with`` block raises, or the caller
+    is cancelled while waiting for responses on exit, remaining jobs are cancelled.
 
     Example::
 
@@ -343,7 +343,16 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
             return False
 
         assert self._group is not None
-        await self._group.wait()
+        try:
+            await self._group.wait()
+        except asyncio.CancelledError:
+            if self._group.job_id in self._worker.job_groups:
+                await asyncio.shield(
+                    self._worker.cancel_job_group(
+                        self._group.job_id, reason="context exited with error"
+                    )
+                )
+            raise
         return False
 
 
@@ -356,7 +365,8 @@ class JobContext:
 
     On normal completion, the result is available via ``response``.
     On worker error or timeout, raises ``JobError``. If the
-    ``async with`` block raises, the job is cancelled.
+    ``async with`` block raises, or the caller is cancelled while waiting
+    for the response on exit, the job is cancelled.
 
     Example::
 
@@ -436,6 +446,14 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
         assert self._group is not None
         try:
             await self._group.wait()
+        except asyncio.CancelledError:
+            if self._group.job_id in self._worker.job_groups:
+                await asyncio.shield(
+                    self._worker.cancel_job_group(
+                        self._group.job_id, reason="context exited with error"
+                    )
+                )
+            raise
         except JobGroupError as e:
             raise JobError(str(e)) from e
         return False

tests/test_job_group.py

@@ -100,6 +100,7 @@ def __init__(self, name):
         super().__init__(name)
         self.started = asyncio.Event()
         self.was_cancelled = False
+        self.cancelled = asyncio.Event()
 
     async def on_job_request(self, message):
         await super().on_job_request(message)
@@ -109,6 +110,7 @@ async def on_job_request(self, message):
             await asyncio.Event().wait()
         except asyncio.CancelledError:
             self.was_cancelled = True
+            self.cancelled.set()
 
 
 async def create_test_env():
@@ -745,6 +747,53 @@ async def test_fire_and_forget_jobs_cancelled_manually(self):
             self.assertEqual(m.reason, "tool cancelled")
         self.assertEqual(len(parent.job_groups), 0)
 
+    async def test_job_cancels_while_waiting_on_context_exit(self):
+        await self._assert_cancellation_during_context_exit(single_worker=True)
+
+    async def test_job_group_cancels_while_waiting_on_context_exit(self):
+        await self._assert_cancellation_during_context_exit(single_worker=False)
+
+    async def _assert_cancellation_during_context_exit(self, *, single_worker: bool):
+        sent = capture_bus(self.bus)
+        parent = StubTask("parent")
+        await setup_task(self.bus, self.registry, parent)
+        names = ("worker",) if single_worker else ("worker-1", "worker-2")
+        workers = [SlowWorkerTask(name) for name in names]
+        for worker in workers:
+            await setup_task(self.bus, self.registry, worker)
+
+        context = parent.job(names[0]) if single_worker else parent.job_group(*names)
+        exiting = asyncio.Event()
+
+        async def run_job():
+            async with context:
+                # No await after set(): the caller reaches __aexit__ before the test resumes.
+                exiting.set()
+
+        task = self.tm.create_task(run_job(), "job-context-exit")
+        try:
+            await asyncio.wait_for(exiting.wait(), timeout=2.0)
+            await asyncio.wait_for(
+                asyncio.gather(*(worker.started.wait() for worker in workers)), timeout=2.0
+            )
+            group = parent.job_groups[context.job_id]
+            task.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await task
+
+            cancel_msgs = [message for message in sent if isinstance(message, BusJobCancelMessage)]
+            self.assertEqual(len(cancel_msgs), len(workers))
+            self.assertEqual({message.target for message in cancel_msgs}, set(names))
+            self.assertTrue(all(message.job_id == context.job_id for message in cancel_msgs))
+            self.assertNotIn(context.job_id, parent.job_groups)
+            self.assertTrue(group.is_done)
+            await asyncio.wait_for(
+                asyncio.gather(*(worker.cancelled.wait() for worker in workers)), timeout=2.0
+            )
+        finally:
+            await self.tm.cancel_task(task)
+            await parent.cancel_job_group(context.job_id, reason="test cleanup")
+
     async def test_cancel_interrupts_running_handler(self):
         """Cancelling a job interrupts a handler that is currently executing."""
         sent = capture_bus(self.bus)