Fix memory leak in pending job assignments during timeout/cancellation/shutdown

livekit-agents/livekit/agents/worker.py

@@ -903,6 +903,11 @@ async def aclose(self) -> None:
 
             self._closed = True
 
+            for _, fut in list(self._pending_assignments.items()):
+                if not fut.done():
+                    fut.cancel()
+            self._pending_assignments.clear()
+
             if self._conn_task is not None:
                 await utils.aio.cancel_and_wait(self._conn_task)
 
@@ -1177,11 +1182,22 @@ async def _on_accept(args: JobAcceptArguments) -> None:
             try:
                 await asyncio.wait_for(wait_assignment, ASSIGNMENT_TIMEOUT)
             except asyncio.TimeoutError:
+                self._pending_assignments.pop(job_req.id, None)
                 logger.warning(
                     f"assignment for job {job_req.id} timed out",
                     extra={"job_request": job_req, "agent_name": self._agent_name},
                 )
                 raise AssignmentTimeoutError() from None
+            except asyncio.CancelledError:
+                self._pending_assignments.pop(job_req.id, None)
+                logger.debug(
+                    f"assignment for job {job_req.id} was cancelled",
+                    extra={"job_request": job_req, "agent_name": self._agent_name},
+                )
+                raise
+            except Exception:
+                self._pending_assignments.pop(job_req.id, None)
+                raise
 
             job_assign = wait_assignment.result()
             running_info = RunningJobInfo(