Skip to content

feat(agents): serialize the inline AgentTasks awaited from one turn #6862

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
u9g wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(agents): serialize the inline AgentTasks awaited from one turn #6862

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
80c37e5
refactor(agents): resolve the activity before arming the premature-fi…
u9g Aug 14, 2026
96d94d6
feat(agents): serialize the inline AgentTasks awaited from one turn
u9g Aug 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
240 changes: 109 additions & 131 deletions livekit-agents/livekit/agents/voice/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
from __future__ import annotations

import asyncio
import contextlib
import time
from collections.abc import AsyncGenerator, AsyncIterable, Coroutine, Generator
from dataclasses import dataclass
Expand Down Expand Up @@ -915,6 +914,15 @@ async def __await_impl(self) -> TaskResult_T:
f"{self.__class__.__name__} should only be awaited inside tool_functions or the on_enter/on_exit methods of an Agent" # noqa: E501
)

# imported at call time: agent_activity imports Agent at module scope, so the reverse can't
from .agent_activity import _AgentActivityContextVar, _SpeechHandleContextVar

speech_handle = _SpeechHandleContextVar.get(None)
old_activity = _AgentActivityContextVar.get()
old_agent = old_activity.agent
session = old_activity.session
self._old_agent = old_agent

def _handle_task_done(_: asyncio.Task[Any]) -> None:
if self.__fut.done():
return
Expand All @@ -931,31 +939,6 @@ def _handle_task_done(_: asyncio.Task[Any]) -> None:
)
)

current_task.add_done_callback(_handle_task_done)

from .agent_activity import _AgentActivityContextVar, _SpeechHandleContextVar

# TODO(theomonnom): add a global lock for inline tasks
# This may currently break in the case we use parallel tool calls.

speech_handle = _SpeechHandleContextVar.get(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:
if speech_handle.interrupted:
raise RuntimeError(
f"{self.__class__.__name__} cannot be awaited inside a function tool that is already interrupted"
)

# lock the speech handle to prevent interruptions until the task is complete
# there should be no await before this line to avoid race conditions
old_allow_interruptions = speech_handle.allow_interruptions
speech_handle.allow_interruptions = False

blocked_tasks = [current_task]
if (
old_activity._on_enter_task
Expand All @@ -964,119 +947,114 @@ def _handle_task_done(_: asyncio.Task[Any]) -> None:
):
blocked_tasks.append(old_activity._on_enter_task)

# register before any await so a concurrent drain (e.g. session close)
# won't wait for tasks blocked on this handoff
old_activity._add_drain_blocked_tasks(blocked_tasks)

# watch the blocked tasks so an active run won't complete mid-handoff
# (the parent speech may predate the run, e.g. created in on_enter)
if (run_state := session._global_run_state) and not run_state.done():
for task in blocked_tasks:
run_state._watch_handle(task)

if (
task_info.function_call
and isinstance(old_activity.llm, RealtimeModel)
and not old_activity.llm.capabilities.manual_function_calls
async with old_activity._inline_task_slot(
speech_handle=speech_handle, blocked_tasks=blocked_tasks
):
logger.error(
f"Realtime model '{old_activity.llm.label}' does not support resuming function calls from chat context, "
"using AgentTask inside a function tool may have unexpected behavior."
)

# TODO(theomonnom): could the RunResult watcher & the blocked_tasks share the same logic?
self.__inactive_ev.clear()
suspended_handles: list[SpeechHandle | asyncio.Future[Any]] = []
pending_on_enter_task: asyncio.Task[None] | None = None
try:
# use wait_on_enter=False to avoid deadlock: on_enter may spawn nested
# AgentTasks that require user input, but session.run() can't return until
# all watched handles complete — creating a circular wait.
await session._update_activity(
self, previous_activity="pause", blocked_tasks=blocked_tasks, wait_on_enter=False
)
current_task.add_done_callback(_handle_task_done)

if (
task_info.function_call
and isinstance(old_activity.llm, RealtimeModel)
and not old_activity.llm.capabilities.manual_function_calls
):
logger.error(
f"Realtime model '{old_activity.llm.label}' does not support resuming function calls from chat context, "
"using AgentTask inside a function tool may have unexpected behavior."
)

if not self._activity and not self.done():
self.complete(
ToolError(
f"activity doesn't start for {self.id}, likely due to session closing"
)
# TODO(theomonnom): could the RunResult watcher & the blocked_tasks share the same logic?
self.__inactive_ev.clear()
suspended_handles: list[SpeechHandle | asyncio.Future[Any]] = []
pending_on_enter_task: asyncio.Task[None] | None = None
try:
# use wait_on_enter=False to avoid deadlock: on_enter may spawn nested
# AgentTasks that require user input, but session.run() can't return until
# all watched handles complete — creating a circular wait.
await session._update_activity(
self,
previous_activity="pause",
blocked_tasks=blocked_tasks,
wait_on_enter=False,
)

run_state = session._global_run_state
if not self._activity and not self.done():
self.complete(
ToolError(
f"activity doesn't start for {self.id}, likely due to session closing"
)
)

if self._activity and (on_enter_task := self._activity._on_enter_task):
run_state = session._global_run_state

if self._activity and (on_enter_task := self._activity._on_enter_task):
if run_state and not run_state.done():
# watch the on_enter task as a guard so RunResult won't complete
# before on_enter has registered its own speech handles
run_state._watch_handle(on_enter_task)
pending_on_enter_task = on_enter_task
else:
# no active run to guard — just wait for on_enter directly
await asyncio.shield(on_enter_task)

# now unwatch the parent speech handle and blocked tasks that belong to the
# old activity — they can't complete while this AgentTask is running, and
# keeping them watched would block RunResult from completing. A foreground
# hold waiting on this task is in the same position, so its guard suspends too.
if run_state and not run_state.done():
if speech_handle and run_state._unwatch_handle(speech_handle):
suspended_handles.append(speech_handle)
for blocked in [*blocked_tasks, *session._foreground_guards]:
if run_state._unwatch_handle(blocked):
suspended_handles.append(blocked)
if suspended_handles:
run_state._mark_done_if_needed(None)
# asyncio.CancelledError derives from BaseException, not Exception
except BaseException:
self.__inactive_ev.set()
raise

try:
return await asyncio.shield(self.__fut)

finally:
# run_state could have changed after self.__fut
run_state = session._global_run_state

# re-watch the suspended handles so the resumed parent activity
# is tracked by the current RunResult again
if run_state and not run_state.done():
# watch the on_enter task as a guard so RunResult won't complete
# before on_enter has registered its own speech handles
run_state._watch_handle(on_enter_task)
pending_on_enter_task = on_enter_task
for handle in suspended_handles:
run_state._watch_handle(handle)

if pending_on_enter_task:
try:
await asyncio.shield(pending_on_enter_task)
except BaseException:
logger.exception("error in on_enter task of agent %s", self.id)

if session._closing and self._activity is None:
# the activity never started (session closing), skip the handoff;
# the close path owns the previous activity
pass
elif session.current_agent != self:
logger.warning(
f"{self.__class__.__name__} completed, but the agent has changed in the meantime. "
"Ignoring handoff to the previous agent, likely due to `AgentSession.update_agent` being invoked."
)
await old_activity.aclose()
else:
# no active run to guard — just wait for on_enter directly
await asyncio.shield(on_enter_task)

# now unwatch the parent speech handle and blocked tasks that belong to the
# old activity — they can't complete while this AgentTask is running, and
# keeping them watched would block RunResult from completing. A foreground
# hold waiting on this task is in the same position, so its guard suspends too.
if run_state and not run_state.done():
if speech_handle and run_state._unwatch_handle(speech_handle):
suspended_handles.append(speech_handle)
for blocked in [*blocked_tasks, *session._foreground_guards]:
if run_state._unwatch_handle(blocked):
suspended_handles.append(blocked)
if suspended_handles:
run_state._mark_done_if_needed(None)
except Exception:
self.__inactive_ev.set()
raise

try:
return await asyncio.shield(self.__fut)

finally:
if speech_handle:
with contextlib.suppress(RuntimeError):
speech_handle.allow_interruptions = old_allow_interruptions

# run_state could have changed after self.__fut
run_state = session._global_run_state

# re-watch the suspended handles so the resumed parent activity
# is tracked by the current RunResult again
if run_state and not run_state.done():
for handle in suspended_handles:
run_state._watch_handle(handle)

if pending_on_enter_task:
try:
await asyncio.shield(pending_on_enter_task)
except BaseException:
logger.exception("error in on_enter task of agent %s", self.id)

if session._closing and self._activity is None:
# the activity never started (session closing), skip the handoff;
# the close path owns the previous activity
pass
elif session.current_agent != self:
logger.warning(
f"{self.__class__.__name__} completed, but the agent has changed in the meantime. "
"Ignoring handoff to the previous agent, likely due to `AgentSession.update_agent` being invoked."
)
await old_activity.aclose()
else:
merged_chat_ctx = old_agent.chat_ctx.merge(
self.chat_ctx,
exclude_function_call=not self._preserve_function_call_history,
exclude_instructions=True,
)
# set the chat_ctx directly, `session._update_activity` will sync it to the rt_session if needed
old_agent._chat_ctx.items[:] = merged_chat_ctx.items
merged_chat_ctx = old_agent.chat_ctx.merge(
self.chat_ctx,
exclude_function_call=not self._preserve_function_call_history,
exclude_instructions=True,
)
# set the chat_ctx directly, `session._update_activity` will sync it to the rt_session if needed
old_agent._chat_ctx.items[:] = merged_chat_ctx.items

await session._update_activity(
old_agent, new_activity="resume", wait_on_enter=False
)
self.__inactive_ev.set()
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__()
Expand Down
54 changes: 54 additions & 0 deletions livekit-agents/livekit/agents/voice/agent_activity.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import asyncio
import contextlib
import contextvars
import heapq
import json
Expand All @@ -20,6 +21,7 @@
from ..llm.realtime_fallback_adapter import _FallbackRealtimeSession
from ..llm.tool_context import (
StopResponse,
ToolError,
ToolFlag,
get_fnc_tool_names,
)
Expand Down Expand Up @@ -213,6 +215,8 @@ def __init__(self, agent: Agent, sess: AgentSession) -> None:
self._realtime_spans: utils.BoundedDict[str, trace.Span] | None = None
self._audio_recognition: AudioRecognition | None = None
self._lock = asyncio.Lock()
# one awaited inline AgentTask may pause this activity at a time
self._inline_task_lock = asyncio.Lock()
self._tool_choice: llm.ToolChoice | None = None

self._started = False
Expand Down Expand Up @@ -1216,6 +1220,56 @@ async def _pause_scheduling_task(
# we still wait for the entire execution (e.g function_tools)
await asyncio.shield(self._scheduling_atask)

@contextlib.asynccontextmanager
async def _inline_task_slot(
self,
*,
speech_handle: SpeechHandle | None,
blocked_tasks: list[asyncio.Task[Any]],
) -> AsyncGenerator[None, None]:
"""Grants the floor to one awaited inline AgentTask at a time.

Pausing an activity is a single slot: concurrent handoffs would leave every switch
but the last overwritten, and those tasks awaiting a result nothing can produce. The
inline tasks of a turn's parallel tool calls queue here and take the slot in turn.

The queue is per activity, which is the thing a handoff contends for. A nested task
pauses the activity of the task it is nested in, so it takes that activity's slot
and never waits on the one its parent is holding.

Each step below is ordered against the queue, and inverting any of them hangs the
session in its own way, so they belong together rather than at the call site.
"""
# before the queue: the tasks behind it share this speech, and a hold released
# between them lets one task's user turns interrupt it out from under the rest
if speech_handle is not None:
if speech_handle.interrupted:
raise RuntimeError("the speech that awaited the inline task is interrupted")

speech_handle._hold_interruptions()

try:
# before the queue: a queued task absent from the drain set makes session close
# wait on the slot it is still queued for
self._add_drain_blocked_tasks(blocked_tasks)
Comment thread
devin-ai-integration[bot] marked this conversation as resolved.

async with self._inline_task_lock:
if self._closed or self._session._closing:
# reported to the model as a tool failure, the way a tool awaiting an
# inline task through a session close has always been
raise ToolError("the activity that awaited the inline task is closing")

# past the queue: a run watching a task still waiting its turn waits for
# the user input the task ahead of it needs
if (run_state := self._session._global_run_state) and not run_state.done():
for task in blocked_tasks:
run_state._watch_handle(task)

yield
finally:
if speech_handle is not None:
speech_handle._release_interruptions()

def _add_drain_blocked_tasks(self, tasks: list[asyncio.Task[Any]]) -> None:
# tasks blocked on an agent handoff are excluded from the drain wait,
# otherwise drain would wait for them while the handoff waits for drain's lock
Expand Down
25 changes: 25 additions & 0 deletions livekit-agents/livekit/agents/voice/speech_handle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ def __init__(
) -> None:
self._id = speech_id
self._allow_interruptions = allow_interruptions
self._interruption_holds = 0
self._interruption_holds_restore = allow_interruptions
self._input_details = input_details

self._interrupt_fut = asyncio.Future[None]()
Expand Down Expand Up @@ -132,6 +134,29 @@ def allow_interruptions(self, value: bool) -> None:

self._allow_interruptions = value

def _hold_interruptions(self) -> None:
"""Disallow interruptions until every hold taken here is released.

Counted rather than set, because the holders of one speech are not serialised
against each other: the inline tasks awaited from a turn's parallel tool calls run
one at a time, and a hold released between them would let the user turns of one
task's sub-conversation interrupt the speech the rest are still anchored to. An
interrupted handle can no longer disallow interruptions, so those tasks could then
never run. The first holder owns the value the last one restores.
"""
if self._interruption_holds == 0:
self._interruption_holds_restore = self._allow_interruptions
self.allow_interruptions = False

self._interruption_holds += 1

def _release_interruptions(self) -> None:
self._interruption_holds -= 1
if self._interruption_holds == 0:
# a forced interrupt lands regardless of the hold, and leaves nothing to restore
with contextlib.suppress(RuntimeError):
self.allow_interruptions = self._interruption_holds_restore

@property
def chat_items(self) -> list[llm.ChatItem]:
return self._chat_items
Expand Down
Loading