Fix parallel tool handoff reply race

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

@@ -3664,12 +3664,15 @@ async def _next_segment() -> _SpeechSegment | None:
                 fnc_executed_ev.function_calls.append(sanitized_out.fnc_call)
                 fnc_executed_ev.function_call_outputs.append(sanitized_out.fnc_call_out)
 
-                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")
-                    ignore_task_switch = True
-                    # TODO(long): should we mark the function call as failed to notify the LLM?
-
-                new_agent_task = sanitized_out.agent_task
+                if sanitized_out.agent_task is not None:
+                    if new_agent_task is not None:
+                        logger.error(
+                            "expected to receive only one AgentTask from the tool executions"
+                        )
+                        ignore_task_switch = True
+                        # TODO(long): should we mark the function call as failed to notify the LLM?
+                    else:
+                        new_agent_task = sanitized_out.agent_task
 
             if new_agent_task and not ignore_task_switch:
                 fnc_executed_ev._handoff_required = True
@@ -3687,7 +3690,11 @@ async def _next_segment() -> _SpeechSegment | None:
                 self._agent._chat_ctx.insert(tool_messages)
                 self._session._tool_items_added(tool_messages)
 
-            if fnc_executed_ev.has_tool_reply and not speech_handle.interrupted:
+            if (
+                fnc_executed_ev.has_tool_reply
+                and not fnc_executed_ev.has_agent_handoff
+                and not speech_handle.interrupted
+            ):
                 # forwarding chat_ctx to the tool reply: drop the in-progress placeholders
                 # (the next turn re-injects from the live running set)
                 _strip_running_tool_calls(chat_ctx)
@@ -4341,25 +4348,31 @@ def _create_assistant_message(
                 self._agent._chat_ctx._upsert_item(sanitized_out.fnc_call_out)
                 self._session._tool_items_added([sanitized_out.fnc_call_out])
 
-                if new_agent_task is not None and sanitized_out.agent_task is not None:
-                    logger.error(
-                        "expected to receive only one Agent from the tool executions",
-                    )
-                    ignore_task_switch = True
-
-                new_agent_task = sanitized_out.agent_task
+                if sanitized_out.agent_task is not None:
+                    if new_agent_task is not None:
+                        logger.error(
+                            "expected to receive only one Agent from the tool executions",
+                        )
+                        ignore_task_switch = True
+                    else:
+                        new_agent_task = sanitized_out.agent_task
 
             if new_agent_task and not ignore_task_switch:
                 fnc_executed_ev._handoff_required = True
 
             self._session.emit("function_tools_executed", fnc_executed_ev)
 
+            tool_reply_expected = (
+                fnc_executed_ev.has_tool_reply and not fnc_executed_ev.has_agent_handoff
+            )
             draining = self.scheduling_paused
             if fnc_executed_ev._handoff_required and new_agent_task and not ignore_task_switch:
                 self._session.update_agent(new_agent_task)
                 draining = True
 
-            if len(new_fnc_outputs) > 0:
+            # Sending results to an outgoing realtime session may trigger its automatic
+            # tool reply, so let the handoff close that session without syncing them.
+            if len(new_fnc_outputs) > 0 and not fnc_executed_ev.has_agent_handoff:
                 # wait all speeches played before updating the tool output and generating the response
                 # most realtime models don't support generating multiple responses at the same time
                 while self._current_speech or self._speech_q:
@@ -4377,7 +4390,7 @@ def _create_assistant_message(
                 auto_reply_fut: asyncio.Future[None] | None = None
                 if (
                     self._rt_session.capabilities.auto_tool_reply_generation
-                    and fnc_executed_ev.has_tool_reply
+                    and tool_reply_expected
                     and self._pending_auto_tool_reply_fut is None
                     and (run_state := self._session._global_run_state) is not None
                     and not run_state.done()
@@ -4415,7 +4428,6 @@ async def _wait_for_auto_tool_reply() -> None:
                             self._pending_auto_tool_reply_fut = None
                         auto_reply_fut.set_result(None)
 
-            tool_reply_expected = fnc_executed_ev.has_tool_reply
             if tool_reply_expected and not self._rt_session.capabilities.auto_tool_reply_generation:
                 self._rt_session.interrupt()
 

tests/test_realtime_parallel_handoff.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from livekit.agents import Agent, AgentSession, function_tool, utils
+from livekit.agents.llm import ChatContext, FunctionCall, GenerationCreatedEvent, MessageGeneration
+
+from .fake_realtime import FakeRealtimeModel, fake_capabilities
+
+pytestmark = [pytest.mark.unit, pytest.mark.virtual_time, pytest.mark.no_concurrent]
+
+
+class ReceivingAgent(Agent):
+    def __init__(self) -> None:
+        super().__init__(instructions="You are the receiving agent.")
+        self.entered = asyncio.Event()
+
+    async def on_enter(self) -> None:
+        self.entered.set()
+
+
+class ParallelHandoffAgent(Agent):
+    def __init__(self, receiver: ReceivingAgent) -> None:
+        super().__init__(instructions="You are the routing agent.")
+        self._receiver = receiver
+
+    @function_tool
+    async def transfer(self) -> Agent:
+        """Transfer the conversation to the receiving agent."""
+        return self._receiver
+
+    @function_tool
+    async def lookup(self) -> str:
+        """Look up the account before transfer."""
+        return "lookup complete"
+
+
+def _tool_generation(tool_names: tuple[str, str]) -> GenerationCreatedEvent:
+    message_ch = utils.aio.Chan[MessageGeneration]()
+    function_ch = utils.aio.Chan[FunctionCall]()
+    message_ch.close()
+    for name in tool_names:
+        function_ch.send_nowait(FunctionCall(call_id=f"{name}-1", name=name, arguments="{}"))
+    function_ch.close()
+
+    return GenerationCreatedEvent(
+        message_stream=message_ch,
+        function_stream=function_ch,
+        user_initiated=True,
+        response_id="parallel-handoff",
+    )
+
+
+@pytest.mark.parametrize(
+    "tool_names",
+    [("lookup", "transfer"), ("transfer", "lookup")],
+    ids=["handoff_last", "handoff_first"],
+)
+@pytest.mark.parametrize("auto_tool_reply_generation", [True, False], ids=["auto", "manual"])
+async def test_realtime_parallel_tool_reply_does_not_race_agent_handoff(
+    tool_names: tuple[str, str],
+    auto_tool_reply_generation: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model = FakeRealtimeModel(
+        capabilities=fake_capabilities(
+            auto_tool_reply_generation=auto_tool_reply_generation,
+            audio_output=False,
+        )
+    )
+    receiver = ReceivingAgent()
+    routing_agent = ParallelHandoffAgent(receiver)
+    tools_executed = asyncio.Event()
+
+    async with AgentSession(llm=model) as session:
+        session.on("function_tools_executed", lambda _: tools_executed.set())
+        await session.start(routing_agent)
+
+        outgoing_rt_session = model.active_session
+        chat_ctx_updates: list[ChatContext] = []
+        update_chat_ctx = outgoing_rt_session.update_chat_ctx
+
+        async def _record_chat_ctx_update(chat_ctx: ChatContext) -> None:
+            chat_ctx_updates.append(chat_ctx.copy())
+            await update_chat_ctx(chat_ctx)
+
+        monkeypatch.setattr(outgoing_rt_session, "update_chat_ctx", _record_chat_ctx_update)
+        speech_handle = session.generate_reply()
+
+        async def _wait_for_reply_request() -> None:
+            while not outgoing_rt_session._reply_futs:
+                await asyncio.sleep(0)
+
+        await asyncio.wait_for(_wait_for_reply_request(), timeout=5.0)
+        outgoing_rt_session._reply_futs[0].set_result(_tool_generation(tool_names))
+
+        await asyncio.wait_for(tools_executed.wait(), timeout=5.0)
+        await asyncio.wait_for(receiver.entered.wait(), timeout=5.0)
+        await asyncio.wait_for(speech_handle.wait_for_playout(), timeout=5.0)
+
+        assert session.current_agent is receiver
+        assert outgoing_rt_session.generate_reply_calls == 1
+        assert not any(
+            item.type == "function_call_output"
+            for chat_ctx in chat_ctx_updates
+            for item in chat_ctx.items
+        )
+
+        for label, items in (
+            ("routing agent chat_ctx", routing_agent.chat_ctx.items),
+            ("session history", session.history.items),
+        ):
+            calls = [item for item in items if item.type == "function_call"]
+            outputs = [item for item in items if item.type == "function_call_output"]
+            assert {call.call_id for call in calls} == {"transfer-1", "lookup-1"}, label
+            assert {output.call_id for output in outputs} == {"transfer-1", "lookup-1"}, label

tests/test_tool_output_per_call.py

@@ -35,6 +35,32 @@ async def do_the_thing(self) -> Any:
         return self._behavior()
 
 
+class ReceivingAgent(Agent):
+    def __init__(self) -> None:
+        super().__init__(instructions="You are the receiving agent.")
+        self.entered = asyncio.Event()
+
+    async def on_enter(self) -> None:
+        self.entered.set()
+        await self.session.generate_reply(instructions="receiver_greeting")
+
+
+class ParallelHandoffAgent(Agent):
+    def __init__(self, receiver: ReceivingAgent) -> None:
+        super().__init__(instructions="You are the routing agent.")
+        self._receiver = receiver
+
+    @function_tool
+    async def transfer(self) -> Agent:
+        """Transfer the conversation to the receiving agent."""
+        return self._receiver
+
+    @function_tool
+    async def lookup(self) -> str:
+        """Look up the account before transfer."""
+        return "lookup complete"
+
+
 async def _run(behavior: Callable[[], Any]) -> tuple[Agent, AgentSession, FunctionCallOutput]:
     """Run one turn whose single tool call is answered, and return that answer."""
     actions = FakeActions()
@@ -104,3 +130,49 @@ async def test_bare_handoff_answers_the_call_and_asks_for_no_reply() -> None:
     assert session.current_agent is not agent, "the handoff must be applied"
     assert output.output == ""
     assert not output.reply_required
+
+
+@pytest.mark.parametrize(
+    "tool_names",
+    [("lookup", "transfer"), ("transfer", "lookup")],
+    ids=["handoff_last", "handoff_first"],
+)
+async def test_parallel_tool_reply_does_not_race_agent_handoff(
+    tool_names: tuple[str, str],
+) -> None:
+    """A sibling tool result must not make the old agent reply after a handoff."""
+    actions = FakeActions()
+    actions.add_user_speech(0.5, 2.5, "Transfer and look up the account.")
+    actions.add_llm(
+        content="",
+        tool_calls=[
+            FunctionToolCall(name=name, arguments="{}", call_id=f"{name}-1") for name in tool_names
+        ],
+    )
+    actions.add_llm(content="old agent reply", input="lookup complete")
+    actions.add_tts(0.5, input="old agent reply")
+    actions.add_llm(content="receiver reply", input="receiver_greeting")
+    actions.add_tts(0.5, input="receiver reply")
+
+    session = create_session(actions)
+    receiver = ReceivingAgent()
+    routing_agent = ParallelHandoffAgent(receiver)
+    await asyncio.wait_for(run_session(session, routing_agent), timeout=SESSION_TIMEOUT)
+
+    assert receiver.entered.is_set(), "the receiving agent must run on_enter"
+    assistant_text = [
+        item.text_content
+        for item in session.history.items
+        if item.type == "message" and item.role == "assistant"
+    ]
+    assert "receiver reply" in assistant_text
+    assert "old agent reply" not in assistant_text
+
+    for label, items in (
+        ("routing agent chat_ctx", routing_agent.chat_ctx.items),
+        ("session history", session.history.items),
+    ):
+        calls = [item for item in items if item.type == "function_call"]
+        outputs = [item for item in items if item.type == "function_call_output"]
+        assert {call.call_id for call in calls} == {"transfer-1", "lookup-1"}, label
+        assert {output.call_id for output in outputs} == {"transfer-1", "lookup-1"}, label