gracefully stop AgentTask and parent agents when session close

livekit-agents/livekit/agents/voice/agent.py

@@ -9,7 +9,7 @@
 from livekit import rtc
 
 from .. import inference, llm, stt, tokenize, tts, utils, vad
-from ..llm import ChatContext, RealtimeModel, find_function_tools
+from ..llm import ChatContext, RealtimeModel, ToolError, find_function_tools
 from ..llm.chat_context import _ReadOnlyChatContext
 from ..log import logger
 from ..types import NOT_GIVEN, FlushSentinel, NotGivenOr
@@ -685,10 +685,21 @@ def __init__(
 
         self.__started = False
         self.__fut = asyncio.Future[TaskResult_T]()
+        self.__inactive_ev = asyncio.Event()
+        self.__inactive_ev.set()  # set when the agent is not awaited or activity is closed
+
+        self._old_agent: Agent | None = None
 
     def done(self) -> bool:
         return self.__fut.done()
 
+    def cancel(self) -> None:
+        if self._activity:
+            self._activity.interrupt(force=True)
+        if self.__fut.done():
+            return
+        self.complete(ToolError(f"AgentTask {self.id} is cancelled"))
+
     def complete(self, result: TaskResult_T | Exception) -> None:
         if self.__fut.done():
             raise RuntimeError(f"{self.__class__.__name__} is already done")
@@ -755,6 +766,7 @@ def _handle_task_done(_: asyncio.Task[Any]) -> None:
         old_activity = _AgentActivityContextVar.get()
         old_agent = old_activity.agent
         session = old_activity.session
+        self._old_agent = old_agent
 
         old_allow_interruptions = True
         if speech_handle:
@@ -787,16 +799,30 @@ def _handle_task_done(_: asyncio.Task[Any]) -> None:
             )
 
         # TODO(theomonnom): could the RunResult watcher & the blocked_tasks share the same logic?
-        await session._update_activity(self, previous_activity="pause", blocked_tasks=blocked_tasks)
+        self.__inactive_ev.clear()
+        try:
+            await session._update_activity(
+                self, previous_activity="pause", blocked_tasks=blocked_tasks
+            )
 
-        # NOTE: _update_activity is calling the on_enter method, so the RunResult can capture all speeches
-        run_state = session._global_run_state
-        if speech_handle and run_state and not run_state.done():
-            # make sure to not deadlock on the current speech handle
-            run_state._unwatch_handle(speech_handle)
-            # it is OK to call _mark_done_if_needed here, the above _update_activity will call on_enter
-            # so handles added inside the on_enter will make sure we're not completing the run_state too early.
-            run_state._mark_done_if_needed(None)
+            if not self._activity and not self.done():
+                self.complete(
+                    ToolError(
+                        f"activity doesn't start for {self.id}, likely due to session closing"
+                    )
+                )
+
+            # NOTE: _update_activity is calling the on_enter method, so the RunResult can capture all speeches
+            run_state = session._global_run_state
+            if speech_handle and run_state and not run_state.done():
+                # make sure to not deadlock on the current speech handle
+                run_state._unwatch_handle(speech_handle)
+                # it is OK to call _mark_done_if_needed here, the above _update_activity will call on_enter
+                # so handles added inside the on_enter will make sure we're not completing the run_state too early.
+                run_state._mark_done_if_needed(None)
+        except Exception:
+            self.__inactive_ev.set()
+            raise
 
         try:
             return await asyncio.shield(self.__fut)
@@ -828,10 +854,14 @@ def _handle_task_done(_: asyncio.Task[Any]) -> None:
                 await session._update_activity(
                     old_agent, new_activity="resume", wait_on_enter=False
                 )
+            self.__inactive_ev.set()
 
     def __await__(self) -> Generator[None, None, TaskResult_T]:
         return self.__await_impl().__await__()
 
+    async def _wait_for_inactive(self) -> None:
+        await self.__inactive_ev.wait()
+
 
 @dataclass
 class _ActivityTaskInfo:

livekit-agents/livekit/agents/voice/agent_activity.py

@@ -636,10 +636,11 @@ async def _traceable_on_exit() -> None:
             await self._agent.on_exit()
 
         async with self._lock:
-            self._on_exit_task = task = self._create_speech_task(
-                _traceable_on_exit(), name="AgentTask_on_exit"
-            )
-            _set_activity_task_info(task, inline_task=True)
+            if self._on_exit_task is None:
+                self._on_exit_task = task = self._create_speech_task(
+                    _traceable_on_exit(), name="AgentTask_on_exit"
+                )
+                _set_activity_task_info(task, inline_task=True)
 
             self._cancel_preemptive_generation()
 
@@ -769,6 +770,9 @@ async def aclose(self) -> None:
             self._closed = True
             self._cancel_preemptive_generation()
 
+            # on_exit_task should be awaited in `drain`
+            self._on_exit_task = None
+
             await self._close_session()
             await asyncio.gather(*self._interrupt_background_speeches(force=False))
 
@@ -1040,14 +1044,16 @@ def _schedule_speech(self, speech: SpeechHandle, priority: int, force: bool = Fa
         # This allows for tool responses to be generated before the AgentActivity is finalized.
 
         if self._scheduling_paused and not force:
+            speech.interrupt(force=True)
             raise RuntimeError(
-                "cannot schedule new speech, the speech scheduling is draining/pausing"
+                "cannot schedule new speech, the speech scheduling is draining/pausing, the speech will be cancelled"
             )
 
         if self._scheduling_atask and self._scheduling_atask.done():
             logger.warning(
-                "attempting to schedule a new SpeechHandle, but the scheduling_task is not running."
+                "attempting to schedule a new SpeechHandle, but the scheduling_task is not running, the speech will be cancelled"
             )
+            speech.interrupt(force=True)
             return
 
         while True:
@@ -2681,7 +2687,7 @@ def _create_assistant_message(
 
                 if new_agent_task is not None and sanitized_out.agent_task is not None:
                     logger.error(
-                        "expected to receive only one AgentTask from the tool executions",
+                        "expected to receive only one Agent from the tool executions",
                     )
                     ignore_task_switch = True
 

livekit-agents/livekit/agents/voice/agent_session.py

@@ -36,7 +36,7 @@
 from ..utils.misc import is_given
 from . import io, room_io
 from ._utils import _set_participant_attributes
-from .agent import Agent
+from .agent import Agent, AgentTask
 from .agent_activity import AgentActivity
 from .audio_recognition import TurnDetectionMode
 from .client_events import ClientEventsHandler
@@ -782,19 +782,30 @@ async def _aclose_impl(
             self._closing = True
             self._cancel_user_away_timer()
 
-            if self._activity is not None:
+            activity = self._activity
+            while activity and isinstance(agent_task := activity.agent, AgentTask):
+                # notify AgentTask to complete and wait it to resume the parent agent
+                agent_task.cancel()
+                await agent_task._wait_for_inactive()
+
+                if old_agent := agent_task._old_agent:
+                    activity = old_agent._activity
+                else:
+                    break
+
+            if activity is not None:
                 if not drain:
                     try:
                         # force interrupt speeches when closing the session
-                        await self._activity.interrupt(force=True)
+                        await activity.interrupt(force=True)
                     except RuntimeError:
                         # uninterruptible speech
                         pass
-                await self._activity.drain()
+                await activity.drain()
 
                 # wait any uninterruptible speech to finish
-                if self._activity.current_speech:
-                    await self._activity.current_speech
+                if activity.current_speech:
+                    await activity.current_speech
 
                 # detach the inputs and outputs
                 self.input.audio = None
@@ -804,13 +815,13 @@ async def _aclose_impl(
 
                 if (
                     reason != CloseReason.ERROR
-                    and (audio_recognition := self._activity._audio_recognition) is not None
+                    and (audio_recognition := activity._audio_recognition) is not None
                 ):
                     # wait for the user transcript to be committed
                     audio_recognition.commit_user_turn(audio_detached=True, transcript_timeout=2.0)
 
-                await self._activity.aclose()
-                self._activity = None
+                await activity.aclose()
+            self._activity = None
 
             if self._agent_speaking_span:
                 self._agent_speaking_span.end()
@@ -1068,12 +1079,21 @@ async def _update_activity(
                 otel_context.attach(self._root_span_context)
 
             previous_activity_v = self._activity
-            if self._activity is not None:
+            if (activity := self._activity) is not None:
                 if previous_activity == "close":
-                    await self._activity.drain()
-                    await self._activity.aclose()
+                    await activity.drain()
+                    await activity.aclose()
                 elif previous_activity == "pause":
-                    await self._activity.pause(blocked_tasks=blocked_tasks or [])
+                    await activity.pause(blocked_tasks=blocked_tasks or [])
+
+            if self._closing and new_activity == "start":
+                # disallow starting a new activity when the session is closing
+                logger.warning(
+                    f"session is closing, skipping {new_activity} activity of {self._next_activity.agent.id}",
+                )
+                self._next_activity = None
+                self._activity = None
+                return
 
             self._activity = self._next_activity
             self._next_activity = None

livekit-agents/livekit/agents/voice/generation.py

@@ -20,6 +20,7 @@
     ChatContext,
     StopResponse,
     ToolContext,
+    ToolError,
     utils as llm_utils,
 )
 from ..log import logger
@@ -608,7 +609,16 @@ async def _traceable_fnc_tool(
                         val = await function_callable()
                         output = make_tool_output(fnc_call=fnc_call, output=val, exception=None)
                     except BaseException as e:
-                        if not isinstance(e, StopResponse):
+                        if isinstance(e, ToolError):
+                            logger.warning(
+                                "ToolError while executing tool: %s",
+                                e.message,
+                                extra={
+                                    "function": fnc_call.name,
+                                    "speech_id": speech_handle.id,
+                                },
+                            )
+                        elif not isinstance(e, StopResponse):
                             logger.exception(
                                 "exception occurred while executing tool",
                                 extra={"function": fnc_call.name, "speech_id": speech_handle.id},