improve EndCallTool (#4563)

longcw · davidzhao · web-flow · commit 7fe642d2a3bf · 2026-01-23T14:59:47.000+08:00
Co-authored-by: David Zhao &lt;dz@livekit.io&gt;
diff --git a/examples/voice_agents/session_close_callback.py b/examples/voice_agents/session_close_callback.py
@@ -2,18 +2,9 @@
 
 from dotenv import load_dotenv
 
-from livekit.agents import (
-    Agent,
-    AgentServer,
-    AgentSession,
-    CloseEvent,
-    JobContext,
-    cli,
-    room_io,
-    utils,
-)
+from livekit.agents import Agent, AgentServer, AgentSession, CloseEvent, JobContext, cli
 from livekit.agents.beta.tools import EndCallTool
-from livekit.plugins import silero
+from livekit.plugins import google, silero  # noqa: F401
 
 logger = logging.getLogger("my-worker")
 logger.setLevel(logging.INFO)
@@ -25,26 +16,25 @@
 # or when the worker is shutting down. When closing the session, agent will be interrupted
 # and the last agent message will be added to the chat context.
 
-server = AgentServer()
-
 
 class MyAgent(Agent):
     def __init__(self):
         super().__init__(
             instructions="You are a helpful assistant.",
-            tools=[EndCallTool()],
+            stt="assemblyai/universal-streaming",
+            llm="openai/gpt-4.1-mini",
+            tts="cartesia/sonic-3",
+            # llm=google.realtime.RealtimeModel(),
+            tools=[
+                EndCallTool(
+                    end_instructions="thanks the user for calling and tell them goodbye",
+                    delete_room=True,  # this will disconnect all remote participants, including SIP callers
+                )
+            ],
         )
 
-    @utils.log_exceptions(logger=logger)
-    async def on_exit(self) -> None:
-        logger.info("exiting the agent")
-        if self.session.current_speech:
-            await self.session.current_speech
-
-        logger.info("generating goodbye message")
-        await self.session.generate_reply(
-            instructions="say goodbye to the user", tool_choice="none"
-        )
+    async def on_enter(self) -> None:
+        self.session.generate_reply(instructions="say hello to the user")
 
 
 server = AgentServer()
@@ -53,22 +43,10 @@ async def on_exit(self) -> None:
 @server.rtc_session()
 async def entrypoint(ctx: JobContext):
     session = AgentSession(
-        stt="assemblyai/universal-streaming",
-        llm="openai/gpt-4.1-mini",
-        tts="rime/arcana",
         vad=silero.VAD.load(),
     )
 
-    # session will be closed automatically when the linked participant disconnects
-    # with reason CLIENT_INITIATED, ROOM_DELETED, or USER_REJECTED
-    # or you can disable it by setting the RoomInputOptions.close_on_disconnect to False
-    await session.start(
-        agent=MyAgent(),
-        room=ctx.room,
-        room_options=room_io.RoomOptions(
-            delete_room_on_close=True,
-        ),
-    )
+    await session.start(agent=MyAgent(), room=ctx.room)
 
     @session.on("close")
     def on_close(ev: CloseEvent):
diff --git a/livekit-agents/livekit/agents/beta/tools/end_call.py b/livekit-agents/livekit/agents/beta/tools/end_call.py
@@ -1,28 +1,132 @@
-from ...llm import Tool, Toolset, function_tool
+import asyncio
+from collections.abc import Awaitable
+from typing import Any, Callable
+
+from ...job import get_job_context
+from ...llm import RealtimeModel, Tool, Toolset, function_tool
 from ...log import logger
-from ...voice.events import RunContext
+from ...voice.events import CloseEvent, RunContext, SpeechCreatedEvent
+from ...voice.speech_handle import SpeechHandle
+
+END_CALL_DESCRIPTION = """
+Ends the current call and disconnects immediately.
+
+Call when:
+- The user clearly indicates they are done (e.g., “that’s all, bye”).
+- The agent determines the conversation is complete and should end.
+
+Do not call when:
+- The user asks to pause, hold, or transfer.
+- Intent is unclear.
+
+This is the final action the agent can take.
+Once called, no further interaction is possible with the user.
+Don't generate any other text or response when the tool is called.
+"""
 
 
 class EndCallTool(Toolset):
-    @function_tool(name="end_call")
-    async def _end_call(self, ctx: RunContext) -> None:
+    def __init__(
+        self,
+        *,
+        extra_description: str = "",
+        delete_room: bool = True,
+        end_instructions: str | None = "say goodbye to the user",
+        on_tool_called: Callable[[Toolset.ToolCalledEvent], Awaitable[None]] | None = None,
+        on_tool_completed: Callable[[Toolset.ToolCompletedEvent], Awaitable[None]] | None = None,
+    ):
         """
-        Ends the current call and disconnects immediately.
+        This tool allows the agent to end the call and disconnect from the room.
 
-        Call when:
-        - The user clearly indicates they are done (e.g., “that’s all, bye”).
-        - The agent determines the conversation is complete and should end.
+        Args:
+            extra_description: Additional description to add to the end call tool.
+            delete_room: Whether to delete the room when the user ends the call. deleting the room disconnects all remote users, including SIP callers.
+            end_instructions: Tool output to the LLM for generating the tool response.
+            on_tool_called: Callback to call when the tool is called.
+            on_tool_completed: Callback to call when the tool is completed.
+        """
+        super().__init__()
+        self._delete_room = delete_room
+        self._extra_description = extra_description
 
-        Do not call when:
-        - The user asks to pause, hold, or transfer.
-        - Intent is unclear.
+        self._end_instructions = end_instructions
+        self._on_tool_called = on_tool_called
+        self._on_tool_completed = on_tool_completed
 
-        This is the final action the agent can take.
-        Once called, no further interaction is possible with the user.
-        """
+        self._end_call_tool = function_tool(
+            self._end_call,
+            name="end_call",
+            description=f"{END_CALL_DESCRIPTION}\n{extra_description}",
+        )
+        self._shutdown_session_task: asyncio.Task[None] | None = None
+
+    async def _end_call(self, ctx: RunContext) -> Any | None:
         logger.debug("end_call tool called")
-        ctx.session.shutdown()
+        llm_v = ctx.session.current_agent._get_activity_or_raise().llm
+
+        def _on_speech_done(_: SpeechHandle) -> None:
+            if (
+                not isinstance(llm_v, RealtimeModel)
+                or not llm_v.capabilities.auto_tool_reply_generation
+            ):
+                # tool reply will reuse the same speech handle, so we can shutdown the session
+                # directly after this speech handle is done
+                ctx.session.shutdown()
+            else:
+                self._shutdown_session_task = asyncio.create_task(
+                    self._delayed_session_shutdown(ctx)
+                )
+
+        ctx.speech_handle.add_done_callback(_on_speech_done)
+        ctx.session.once("close", self._on_session_close)
+
+        if self._on_tool_called:
+            await self._on_tool_called(Toolset.ToolCalledEvent(ctx=ctx, arguments={}))
+
+        completed_ev = Toolset.ToolCompletedEvent(ctx=ctx, output=self._end_instructions)
+        if self._on_tool_completed:
+            await self._on_tool_completed(completed_ev)
+
+        return completed_ev.output
+
+    async def _delayed_session_shutdown(self, ctx: RunContext) -> None:
+        """Shutdown the session after the tool reply is played out"""
+        speech_created_fut = asyncio.Future[SpeechHandle]()
+
+        @ctx.session.once("speech_created")
+        def _on_speech_created(ev: SpeechCreatedEvent) -> None:
+            if not speech_created_fut.done():
+                speech_created_fut.set_result(ev.speech_handle)
+
+        try:
+            speech_handle = await asyncio.wait_for(speech_created_fut, timeout=5.0)
+            await speech_handle
+        except asyncio.TimeoutError:
+            logger.warning("tool reply timed out, shutting down session")
+        finally:
+            ctx.session.off("speech_created", _on_speech_created)
+            ctx.session.shutdown()
+
+    def _on_session_close(self, ev: CloseEvent) -> None:
+        """Close the job process when AgentSession is closed"""
+        if self._shutdown_session_task:
+            # cleanup
+            self._shutdown_session_task.cancel()
+            self._shutdown_session_task = None
+
+        job_ctx = get_job_context()
+
+        if self._delete_room:
+
+            async def _on_shutdown() -> None:
+                logger.info("deleting the room because the user ended the call")
+                await job_ctx.delete_room()
+
+            job_ctx.add_shutdown_callback(_on_shutdown)
+
+        # shutdown the job process
+        job_ctx.shutdown(reason=ev.reason.value)
 
     @property
     def tools(self) -> list[Tool]:
-        return [self._end_call]
+        return [self._end_call_tool]
diff --git a/livekit-agents/livekit/agents/llm/tool_context.py b/livekit-agents/livekit/agents/llm/tool_context.py
@@ -21,12 +21,15 @@
 from collections.abc import Awaitable, Sequence
 from dataclasses import dataclass
 from enum import Flag, auto
-from typing import Any, Callable, Generic, Literal, TypeVar, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, TypeVar, Union, overload
 
 from typing_extensions import NotRequired, ParamSpec, Required, Self, TypedDict, TypeGuard
 
 from . import _provider_format
 
+if TYPE_CHECKING:
+    from ..voice.events import RunContext
+
 
 class Tool(ABC):  # noqa: B024
     pass
@@ -39,6 +42,16 @@ class ProviderTool(Tool):
 
 
 class Toolset(ABC):
+    @dataclass
+    class ToolCalledEvent:
+        ctx: RunContext
+        arguments: dict[str, Any]
+
+    @dataclass
+    class ToolCompletedEvent:
+        ctx: RunContext
+        output: Any | Exception | None
+
     @property
     @abstractmethod
     def tools(self) -> list[Tool]:
