wehub-resource-sync
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from livekit.agents import Agent, AgentSession, AgentTask
|
|
from livekit.agents.llm import ToolError
|
|
|
|
from .fake_llm import FakeLLM
|
|
|
|
pytestmark = [pytest.mark.unit, pytest.mark.no_concurrent]
|
|
|
|
|
|
class _SimpleTask(AgentTask):
|
|
def __init__(self) -> None:
|
|
super().__init__(instructions="simple task")
|
|
|
|
|
|
class _ParentAgent(Agent):
|
|
def __init__(self) -> None:
|
|
super().__init__(instructions="parent agent")
|
|
self.entered = asyncio.Event()
|
|
self.exit_started = asyncio.Event()
|
|
self.task_error: BaseException | None = None
|
|
|
|
async def on_enter(self) -> None:
|
|
self.entered.set()
|
|
# hold on_enter until the session close has started draining this activity
|
|
# (on_exit runs inside drain() while it holds the activity lock), reproducing
|
|
# a participant disconnect racing `await AgentTask()`
|
|
await self.exit_started.wait()
|
|
try:
|
|
await _SimpleTask()
|
|
except ToolError as e:
|
|
self.task_error = e
|
|
|
|
async def on_exit(self) -> None:
|
|
self.exit_started.set()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aclose_while_on_enter_awaits_agent_task() -> None:
|
|
"""Closing the session while on_enter awaits an AgentTask must not deadlock:
|
|
drain() waits for the on_enter task, which waits for the activity handoff,
|
|
which waits for the activity lock held by drain()."""
|
|
session = AgentSession(llm=FakeLLM())
|
|
agent = _ParentAgent()
|
|
await session.start(agent)
|
|
await asyncio.wait_for(agent.entered.wait(), timeout=5.0)
|
|
|
|
await asyncio.wait_for(session.aclose(), timeout=10.0)
|
|
|
|
assert isinstance(agent.task_error, ToolError)
|