From 0387c086c373da9d39dfeb2df5144d7a058850e0 Mon Sep 17 00:00:00 2001 From: devteamaegis Date: Tue, 26 May 2026 01:58:20 -0400 Subject: [PATCH] fix(livekit): run tool.func in thread pool so status-update task can execute MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The LiveKit tool handler creates a _status_update task to notify the user while a tool is executing, then immediately calls tool.func() synchronously. Because tool.func() is a blocking HTTP call (using the sync AgentMail client), it never yields to the event loop — so the _status_update task is never scheduled. By the time tool.func() returns, the task is cancelled before it runs even once. The status update therefore silently never fires. Fix: off-load tool.func() to a thread via asyncio.to_thread(). This yields the event loop back to other tasks while the HTTP call runs in a worker thread, giving _status_update the opportunity to execute as intended. Adds python/tests/test_livekit_status_update.py with two test cases: - status update fires during a deliberately slow synchronous tool - tool result JSON is returned correctly after the fix --- python/src/agentmail_toolkit/livekit.py | 6 +- python/tests/test_livekit_status_update.py | 175 +++++++++++++++++++++ 2 files changed, 180 insertions(+), 1 deletion(-) create mode 100644 python/tests/test_livekit_status_update.py diff --git a/python/src/agentmail_toolkit/livekit.py b/python/src/agentmail_toolkit/livekit.py index ea09fb8..a3d8f55 100644 --- a/python/src/agentmail_toolkit/livekit.py +++ b/python/src/agentmail_toolkit/livekit.py @@ -22,7 +22,11 @@ async def _status_update(): status_update_task = asyncio.create_task(_status_update()) - result = tool.func(self.client, raw_arguments).model_dump_json() + # Run the synchronous HTTP call in a thread pool so the event + # loop stays free and _status_update can actually execute. + result = await asyncio.to_thread( + lambda: tool.func(self.client, raw_arguments).model_dump_json() + ) status_update_task.cancel() diff --git a/python/tests/test_livekit_status_update.py b/python/tests/test_livekit_status_update.py new file mode 100644 index 0000000..b68b9c7 --- /dev/null +++ b/python/tests/test_livekit_status_update.py @@ -0,0 +1,175 @@ +"""Tests that the LiveKit tool's status-update task actually runs. + +Before the fix, `tool.func` was called synchronously inside the async `f` +handler. Because a synchronous (blocking) call never yields to the event +loop, the `_status_update` task created just before it could never be +scheduled — and was cancelled before it had a chance to run. + +After the fix, `tool.func` is off-loaded to a thread via `asyncio.to_thread`, +which yields control back to the event loop and lets `_status_update` run. +""" + +import asyncio +import threading +import unittest +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + + +# --------------------------------------------------------------------------- +# Minimal stubs for livekit.agents and agentmail so the import works in +# environments that don't have those packages installed. +# --------------------------------------------------------------------------- + +import sys +import types + +# Stub heavy optional dependencies before any agentmail_toolkit import +for _mod in [ + "livekit", + "livekit.agents", + "pymupdf", + "filetype", + "docx", + "langchain", + "langchain.tools", + "agents", +]: + if _mod not in sys.modules: + sys.modules[_mod] = types.ModuleType(_mod) + +livekit_agents = sys.modules["livekit.agents"] +livekit_agents.FunctionTool = object # type: ignore[attr-defined] +livekit_agents.RunContext = object # type: ignore[attr-defined] +livekit_agents.ToolError = Exception # type: ignore[attr-defined] + + +def _noop_function_tool(**kwargs): + """Stub for livekit.agents.function_tool — just return a sentinel.""" + func = kwargs.get("f") + + class _FakeFunctionTool: + def __init__(self): + self._f = func + + async def __call__(self, *args, **kw): + return await self._f(*args, **kw) + + return _FakeFunctionTool() + + +livekit_agents.function_tool = _noop_function_tool # type: ignore[attr-defined] + +if "agentmail" not in sys.modules: + sys.modules["agentmail"] = types.ModuleType("agentmail") + sys.modules["agentmail"].AgentMail = MagicMock # type: ignore[attr-defined] + +# Patch the toolkit base so we don't try to build all tools at init time. +from agentmail_toolkit.tools import Tool # noqa: E402 +from agentmail_toolkit.livekit import AgentMailToolkit # noqa: E402 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_tool(func) -> Tool: + """Build a minimal Tool wrapping the given callable.""" + from pydantic import BaseModel + + class _Params(BaseModel): + pass + + return Tool( + name="test_tool", + description="test description", + params_schema=_Params, + func=func, + ) + + +def _make_context(reply_mock: AsyncMock) -> MagicMock: + ctx = MagicMock() + ctx.session.generate_reply = reply_mock + return ctx + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestLiveKitStatusUpdateRuns(unittest.IsolatedAsyncioTestCase): + """Verify that _status_update is given a chance to run during tool execution.""" + + async def test_status_update_fires_during_blocking_tool(self) -> None: + """_status_update must run before tool.func returns. + + If tool.func blocks the event loop (old code), _status_update is + cancelled before it can be scheduled and `generate_reply` is never + called. With asyncio.to_thread, the task runs concurrently. + """ + reply_called = asyncio.Event() + + async def mock_generate_reply(**kwargs): + reply_called.set() + + reply_mock = AsyncMock(side_effect=mock_generate_reply) + ctx = _make_context(reply_mock) + + # Synchronous tool that deliberately yields by sleeping in a thread — + # we wait until status_update has a chance to fire, then proceed. + gate = threading.Event() + + def slow_func(client, kwargs): + # Block just long enough for the event loop to schedule tasks. + gate.wait(timeout=2) + result = MagicMock() + result.model_dump_json.return_value = '{"ok": true}' + return result + + tool = _make_tool(slow_func) + + toolkit = object.__new__(AgentMailToolkit) + toolkit.client = MagicMock() + built = toolkit._build_tool(tool) + + # Schedule the tool call. We allow _status_update to run by releasing + # the gate only after the event loop gets control back. + async def run(): + task = asyncio.create_task(built._f({}, ctx)) + # Let the event loop tick once so _status_update can start. + await asyncio.sleep(0) + gate.set() # Release slow_func + return await task + + await run() + self.assertTrue( + reply_called.is_set(), + "generate_reply was never called — status update did not run", + ) + + async def test_tool_result_returned_correctly(self) -> None: + """The tool result must be the JSON-serialised return value of func.""" + reply_mock = AsyncMock() + ctx = _make_context(reply_mock) + + expected = '{"message": "hello"}' + + def func(client, kwargs): + result = MagicMock() + result.model_dump_json.return_value = expected + return result + + tool = _make_tool(func) + toolkit = object.__new__(AgentMailToolkit) + toolkit.client = MagicMock() + built = toolkit._build_tool(tool) + + result = await built._f({}, ctx) + self.assertEqual(result, expected) + + +if __name__ == "__main__": + unittest.main()