Skip to content

async tools work outside AsyncToolset via activity-scoped manager#5711

Open
longcw wants to merge 6 commits into
mainfrom
longc/async-task-manager
Open

async tools work outside AsyncToolset via activity-scoped manager#5711
longcw wants to merge 6 commits into
mainfrom
longc/async-task-manager

Conversation

@longcw
Copy link
Copy Markdown
Contributor

@longcw longcw commented May 12, 2026

Summary

Extract _AsyncToolManager from AsyncToolset so async-tool lifecycle (AsyncRunContext creation, background task spawn, duplicate handling, progress-reply coalescing) is usable independently of the toolset wrapper. Each AgentActivity now owns one manager (on_duplicate_call="reject"), and at tool dispatch any tool whose signature includes AsyncRunContext is routed through it. No upfront grouping or AsyncToolset wrapping needed — a bare @function_tool async def foo(ctx: AsyncRunContext, ...) in agent.tools Just Works.

Backward compatibility

  • AsyncToolset(tools=[...]) unchanged — _wrap_tool body now delegates to its own _AsyncToolManager but external behavior (duplicate handling, progress narration, cancel) is preserved.

Example

from livekit.agents import Agent, AsyncRunContext
from livekit.agents.llm import function_tool

class MyAgent(Agent):
    @function_tool
    async def search_news(self, ctx: AsyncRunContext, topic: str) -> str:
        await ctx.update(f"searching the news for {topic}...")
        await asyncio.sleep(10)
        return f"Top story on {topic}: nothing dramatic happened today."

No AsyncToolset wrapping. The agent acknowledges immediately via the update, then narrates the final result.

See examples/voice_agents/basic_agent.py for the full example.

Extract `_AsyncToolManager` from `AsyncToolset` so the lifecycle pieces —
AsyncRunContext creation, background task spawn, duplicate handling, progress
reply coalescing — can be used independently of the toolset wrapper. Each
`AgentActivity` now owns one such manager (`on_duplicate_call="reject"`), and
at tool dispatch any tool whose signature includes `AsyncRunContext` is routed
through that manager — no upfront grouping or `AsyncToolset` wrapping needed.

`AsyncToolset` continues to work unchanged: it delegates to its own internal
manager and still always exposes `get_running_tasks` / `cancel_task`. The
activity surfaces them dynamically — only while at least one activity-managed
task is live — so agents that never invoke an async tool see no extra tools
in their LLM context. `ToolContext` dedupes by name, so both can coexist.

Tools placed inside an explicit `AsyncToolset` retain its richer semantics
(toolset-scoped lifetime that survives agent handoff, configurable
`on_duplicate_call`). Tools placed directly in `agent.tools` get the
activity-scoped default — tasks die on handoff.

`AsyncRunContext` is now re-exported from `livekit.agents` next to
`RunContext`, so tools can declare it without reaching into the internal
module path.
@chenghao-mou chenghao-mou requested a review from a team May 12, 2026 09:57
devin-ai-integration[bot]

This comment was marked as resolved.

longcw added 3 commits May 13, 2026 09:22
…ync tools

Two PR-review fixes:

1. _close_session() now aclose()s the activity's _AsyncToolManager so the
   pause() path also cancels in-flight bare async tool tasks. Previously only
   aclose() did this, so a late ctx.update() / return from a paused activity
   could touch the next agent's chat context via _enqueue_reply.

2. mock_tools() now takes precedence over the activity manager. The check
   moved ahead of the use_async_manager branch in _execute_tools_task. For
   async tools with a mock, we pass raw JSON kwargs straight through —
   _run_mock trims to the mock's actual signature, so the AsyncRunContext
   param of the original tool isn't required on the mock.
Two layers of cleanup on top of the activity-scoped manager:

1. `prepare_function_arguments` now raises `ToolError` directly for argument
   validation failures (bad JSON, ValidationError, missing required params,
   context-type mismatch). Removes the matching try/except boilerplate from
   `execute_function_call`, `tool_proxy.py::_handle_call`, and the now-deleted
   helper in `generation.py`. One place owns the validation→ToolError
   contract; the message LLMs see is consistent.

2. `_execute_tools_task` now uses a single `_execute(ctx)` closure for sync
   and async tools. The closure preps args against the tool's signature and
   either runs the mock or the real tool. Sync path partial-binds with a
   plain `RunContext`; async path hands the closure to
   `async_tool_manager.spawn`, which calls it with an `AsyncRunContext`.
   `_run_mock` moves to module level in `run_result.py` so it's shared
   between dispatch sites.

3. `_AsyncToolManager.spawn` now takes a `function_callable` (the body) and
   `run_ctx` only — no `tool` or `raw_arguments`. Callers build the closure
   that captures whatever they need. The manager is a pure async-lifecycle
   service.

4. AsyncToolset wrappers handle their own mock lookup so mocks of wrapped
   async tools run through the toolset's manager (full async lifecycle).
   `_is_async_toolset_wrapper(tool)` is the internal contract that tells
   dispatch in `_execute_tools_task` to defer mock handling.

5. Activity-scoped manager teardown moved into `_close_session` so the
   pause() path (not just aclose) cancels in-flight bare async tool tasks.
   Otherwise a late `ctx.update()` from a paused activity could touch the
   next agent's chat context.

Tests collapsed into `tests/test_async_tool.py`:
- 13 unit tests for the manager and dispatch helpers, no real session.
- 2 e2e tests driving a real `AgentSession` with a scripted `FakeLLM` and
  `mock_tools` — including the post-update-raise path that proves errors
  surface to the LLM via `_enqueue_reply` → `REPLY_INSTRUCTIONS` →
  `generate_reply`.
devin-ai-integration[bot]

This comment was marked as resolved.

longcw added 2 commits May 13, 2026 18:12
prepare_function_arguments now wraps validation failures as ToolError
("Error parsing arguments for ..."). Update the three tests that still
expected ValueError or the old "invalid parameters" JSON message.
ToolError from arg parsing is already logged inside prepare_function_arguments;
user-raised ToolError is an intentional message to the LLM. Neither needs a
traceback at logger.exception level.
@longcw longcw requested a review from theomonnom May 13, 2026 10:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants