Add UIWorker
UIWorker is an LLMContextWorker that observes and drives a client GUI over the RTVI UI channel: it stores accessibility snapshots, auto-injects <ui_state> at the start of each respond job, dispatches client events to @on_ui_event handlers, sends UI commands back to the client, and surfaces fan-out work as cancellable task cards via user_job_group(). The optional ReplyToolMixin exposes a bundled reply tool. The prompt_guide parameter auto-appends the UI wire-format guide to the LLM's system instruction (default UI_STATE_PROMPT_GUIDE; override with a string or disable with None), so the LLM can parse the injected <ui_state> / <ui_event> messages without the app concatenating the guide by hand.
This commit is contained in:
Mark Backman
366
tests/test_ui_job_lifecycle.py
Normal file
366
tests/test_ui_job_lifecycle.py
Normal file
@@ -0,0 +1,366 @@
|
||||
#
|
||||
# Copyright (c) 2026, Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
"""Tests for the UIWorker user-job-group lifecycle.
|
||||
|
||||
Covers:
|
||||
- ``UIWorker.on_bus_message`` forwarding of worker job updates/responses
|
||||
for registered user job groups as ``BusUITask*`` carriers.
|
||||
- The reserved ``__cancel_task`` client event routing to
|
||||
``cancel_job_group``.
|
||||
- ``UserJobGroupContext`` publishing ``group_started`` / ``group_completed``
|
||||
envelopes and (de)registering the group on the worker.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from pipecat.bus.messages import BusJobResponseMessage, BusJobUpdateMessage
|
||||
from pipecat.bus.ui_messages import (
|
||||
_UI_CANCEL_TASK_BUS_EVENT_NAME,
|
||||
BusUIEventMessage,
|
||||
BusUITaskCompletedMessage,
|
||||
BusUITaskGroupCompletedMessage,
|
||||
BusUITaskGroupStartedMessage,
|
||||
BusUITaskUpdateMessage,
|
||||
)
|
||||
from pipecat.frames.frames import LLMMessagesAppendFrame
|
||||
from pipecat.pipeline.job_context import JobGroup, JobStatus
|
||||
from pipecat.processors.frame_processor import FrameDirection
|
||||
from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams
|
||||
from pipecat.workers.ui import UIWorker
|
||||
|
||||
|
||||
async def _make_solo_worker(**kwargs) -> UIWorker:
|
||||
"""A UIWorker with a task manager and a ``queue_frame`` spy.
|
||||
|
||||
Suitable for testing forwarding logic by directly invoking
|
||||
``on_bus_message`` and asserting on captured ``send_bus_message``
|
||||
calls.
|
||||
"""
|
||||
worker = UIWorker("ui", llm=MagicMock(), active=False, **kwargs)
|
||||
tm = TaskManager()
|
||||
tm.setup(TaskManagerParams(loop=asyncio.get_running_loop()))
|
||||
worker._task_manager = tm
|
||||
|
||||
recorded: list = []
|
||||
|
||||
async def _record(frame, direction=FrameDirection.DOWNSTREAM):
|
||||
recorded.append(frame)
|
||||
|
||||
worker.queue_frame = _record # type: ignore[method-assign]
|
||||
worker._recorded = recorded # type: ignore[attr-defined]
|
||||
return worker
|
||||
|
||||
|
||||
class TestUIWorkerForwarding(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_unregistered_job_update_is_not_forwarded(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker.send_bus_message = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusJobUpdateMessage(
|
||||
source="worker", target=worker.name, job_id="t-unknown", update={"x": 1}
|
||||
)
|
||||
)
|
||||
|
||||
forwarded = [
|
||||
c.args[0]
|
||||
for c in worker.send_bus_message.await_args_list
|
||||
if isinstance(c.args[0], BusUITaskUpdateMessage)
|
||||
]
|
||||
self.assertEqual(forwarded, [])
|
||||
|
||||
async def test_registered_job_update_is_forwarded(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker._register_user_job_group(
|
||||
job_id="t1", worker_names=["worker"], label="hello", cancellable=True
|
||||
)
|
||||
worker.send_bus_message = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusJobUpdateMessage(
|
||||
source="worker",
|
||||
target=worker.name,
|
||||
job_id="t1",
|
||||
update={"kind": "tool_call", "tool": "WebSearch"},
|
||||
)
|
||||
)
|
||||
|
||||
forwarded = [
|
||||
c.args[0]
|
||||
for c in worker.send_bus_message.await_args_list
|
||||
if isinstance(c.args[0], BusUITaskUpdateMessage)
|
||||
]
|
||||
self.assertEqual(len(forwarded), 1)
|
||||
self.assertEqual(forwarded[0].task_id, "t1")
|
||||
self.assertEqual(forwarded[0].agent_name, "worker")
|
||||
self.assertEqual(forwarded[0].data, {"kind": "tool_call", "tool": "WebSearch"})
|
||||
|
||||
async def test_registered_job_response_is_forwarded(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker._register_user_job_group(
|
||||
job_id="t1", worker_names=["worker"], label=None, cancellable=True
|
||||
)
|
||||
worker.send_bus_message = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusJobResponseMessage(
|
||||
source="worker",
|
||||
target=worker.name,
|
||||
job_id="t1",
|
||||
status=JobStatus.COMPLETED,
|
||||
response={"answer": 42},
|
||||
)
|
||||
)
|
||||
|
||||
forwarded = [
|
||||
c.args[0]
|
||||
for c in worker.send_bus_message.await_args_list
|
||||
if isinstance(c.args[0], BusUITaskCompletedMessage)
|
||||
]
|
||||
self.assertEqual(len(forwarded), 1)
|
||||
self.assertEqual(forwarded[0].task_id, "t1")
|
||||
self.assertEqual(forwarded[0].agent_name, "worker")
|
||||
self.assertEqual(forwarded[0].status, "completed")
|
||||
self.assertEqual(forwarded[0].response, {"answer": 42})
|
||||
|
||||
async def test_response_status_serializes_for_cancelled_and_error(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker._register_user_job_group(
|
||||
job_id="t1", worker_names=["w"], label=None, cancellable=True
|
||||
)
|
||||
worker.send_bus_message = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusJobResponseMessage(
|
||||
source="w", target=worker.name, job_id="t1", status=JobStatus.CANCELLED
|
||||
)
|
||||
)
|
||||
await worker.on_bus_message(
|
||||
BusJobResponseMessage(
|
||||
source="w", target=worker.name, job_id="t1", status=JobStatus.ERROR
|
||||
)
|
||||
)
|
||||
|
||||
statuses = [
|
||||
c.args[0].status
|
||||
for c in worker.send_bus_message.await_args_list
|
||||
if isinstance(c.args[0], BusUITaskCompletedMessage)
|
||||
]
|
||||
self.assertEqual(statuses, ["cancelled", "error"])
|
||||
|
||||
|
||||
class TestCancelJobEvent(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_cancel_event_routes_to_cancel_job_group(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker._register_user_job_group(
|
||||
job_id="t1", worker_names=["w"], label=None, cancellable=True
|
||||
)
|
||||
worker.cancel_job_group = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusUIEventMessage(
|
||||
source="bridge",
|
||||
target=worker.name,
|
||||
event_name=_UI_CANCEL_TASK_BUS_EVENT_NAME,
|
||||
payload={"task_id": "t1", "reason": "user clicked cancel"},
|
||||
)
|
||||
)
|
||||
|
||||
worker.cancel_job_group.assert_awaited_once_with("t1", reason="user clicked cancel")
|
||||
|
||||
async def test_cancel_event_default_reason_when_omitted(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker._register_user_job_group(
|
||||
job_id="t1", worker_names=["w"], label=None, cancellable=True
|
||||
)
|
||||
worker.cancel_job_group = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusUIEventMessage(
|
||||
source="bridge",
|
||||
target=worker.name,
|
||||
event_name=_UI_CANCEL_TASK_BUS_EVENT_NAME,
|
||||
payload={"task_id": "t1"},
|
||||
)
|
||||
)
|
||||
|
||||
worker.cancel_job_group.assert_awaited_once()
|
||||
self.assertEqual(worker.cancel_job_group.await_args.kwargs["reason"], "cancelled by user")
|
||||
|
||||
async def test_non_cancellable_group_is_ignored(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker._register_user_job_group(
|
||||
job_id="t1", worker_names=["w"], label=None, cancellable=False
|
||||
)
|
||||
worker.cancel_job_group = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusUIEventMessage(
|
||||
source="bridge",
|
||||
target=worker.name,
|
||||
event_name=_UI_CANCEL_TASK_BUS_EVENT_NAME,
|
||||
payload={"task_id": "t1"},
|
||||
)
|
||||
)
|
||||
|
||||
worker.cancel_job_group.assert_not_awaited()
|
||||
|
||||
async def test_unknown_job_id_is_ignored(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker.cancel_job_group = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusUIEventMessage(
|
||||
source="bridge",
|
||||
target=worker.name,
|
||||
event_name=_UI_CANCEL_TASK_BUS_EVENT_NAME,
|
||||
payload={"task_id": "nope"},
|
||||
)
|
||||
)
|
||||
|
||||
worker.cancel_job_group.assert_not_awaited()
|
||||
|
||||
async def test_missing_or_bad_payload_is_ignored(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker.cancel_job_group = AsyncMock() # type: ignore[method-assign]
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusUIEventMessage(
|
||||
source="bridge",
|
||||
target=worker.name,
|
||||
event_name=_UI_CANCEL_TASK_BUS_EVENT_NAME,
|
||||
payload=None,
|
||||
)
|
||||
)
|
||||
await worker.on_bus_message(
|
||||
BusUIEventMessage(
|
||||
source="bridge",
|
||||
target=worker.name,
|
||||
event_name=_UI_CANCEL_TASK_BUS_EVENT_NAME,
|
||||
payload={"task_id": 42},
|
||||
)
|
||||
)
|
||||
|
||||
worker.cancel_job_group.assert_not_awaited()
|
||||
|
||||
|
||||
class TestForwardingDoesNotInjectLLMContext(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_job_update_forwarding_does_not_queue_append_frames(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker._register_user_job_group(
|
||||
job_id="t1", worker_names=["w"], label=None, cancellable=True
|
||||
)
|
||||
|
||||
await worker.on_bus_message(
|
||||
BusJobUpdateMessage(source="w", target=worker.name, job_id="t1", update={"x": 1})
|
||||
)
|
||||
|
||||
appends = [f for f in worker._recorded if isinstance(f, LLMMessagesAppendFrame)]
|
||||
self.assertEqual(appends, [])
|
||||
|
||||
|
||||
def _stub_job_group(worker, job_id="t1", worker_names=("w1",)):
|
||||
"""Make ``create_job_group_and_request_job`` return a self-completing group.
|
||||
|
||||
The group completes on the next loop tick so both the context-manager
|
||||
and fire-and-forget paths terminate without a running bus or workers.
|
||||
"""
|
||||
|
||||
async def _fake_create(names, *, name=None, payload=None, timeout=None, cancel_on_error=True):
|
||||
group = JobGroup(job_id=job_id, worker_names=set(names))
|
||||
|
||||
async def _finish():
|
||||
# Yield so JobGroupContext.__aenter__ can set event_queue first.
|
||||
await asyncio.sleep(0)
|
||||
group.complete()
|
||||
|
||||
asyncio.create_task(_finish())
|
||||
return group
|
||||
|
||||
worker.create_job_group_and_request_job = _fake_create # type: ignore[method-assign]
|
||||
|
||||
|
||||
class TestUserJobGroupContext(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_context_publishes_started_and_completed(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker.send_bus_message = AsyncMock() # type: ignore[method-assign]
|
||||
_stub_job_group(worker)
|
||||
|
||||
async with worker.user_job_group("w1", label="My research") as tg:
|
||||
self.assertEqual(tg.job_id, "t1")
|
||||
self.assertIn("t1", worker._user_job_groups)
|
||||
|
||||
self.assertNotIn("t1", worker._user_job_groups)
|
||||
|
||||
kinds = [type(c.args[0]).__name__ for c in worker.send_bus_message.await_args_list]
|
||||
self.assertEqual(
|
||||
kinds,
|
||||
["BusUITaskGroupStartedMessage", "BusUITaskGroupCompletedMessage"],
|
||||
)
|
||||
|
||||
started = worker.send_bus_message.await_args_list[0].args[0]
|
||||
self.assertIsInstance(started, BusUITaskGroupStartedMessage)
|
||||
self.assertEqual(started.task_id, "t1")
|
||||
self.assertEqual(started.agents, ["w1"])
|
||||
self.assertEqual(started.label, "My research")
|
||||
self.assertTrue(started.cancellable)
|
||||
|
||||
completed = worker.send_bus_message.await_args_list[1].args[0]
|
||||
self.assertIsInstance(completed, BusUITaskGroupCompletedMessage)
|
||||
self.assertEqual(completed.task_id, "t1")
|
||||
|
||||
async def test_non_cancellable_group_sets_flag_in_started_message(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker.send_bus_message = AsyncMock() # type: ignore[method-assign]
|
||||
_stub_job_group(worker)
|
||||
|
||||
async with worker.user_job_group("w1", cancellable=False):
|
||||
pass
|
||||
|
||||
started = worker.send_bus_message.await_args_list[0].args[0]
|
||||
self.assertFalse(started.cancellable)
|
||||
|
||||
async def test_unregisters_on_exit(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker.send_bus_message = AsyncMock() # type: ignore[method-assign]
|
||||
_stub_job_group(worker)
|
||||
|
||||
async with worker.user_job_group("w1") as tg:
|
||||
pass
|
||||
|
||||
self.assertNotIn(tg.job_id, worker._user_job_groups)
|
||||
|
||||
async def test_start_user_job_group_returns_id_and_publishes(self):
|
||||
worker = await _make_solo_worker()
|
||||
worker.send_bus_message = AsyncMock() # type: ignore[method-assign]
|
||||
_stub_job_group(worker)
|
||||
|
||||
job_id = await worker.start_user_job_group("w1", label="Background work")
|
||||
self.assertEqual(job_id, "t1")
|
||||
|
||||
started = worker.send_bus_message.await_args_list[0].args[0]
|
||||
self.assertIsInstance(started, BusUITaskGroupStartedMessage)
|
||||
self.assertEqual(started.label, "Background work")
|
||||
|
||||
# The background runner drains the group and publishes completion.
|
||||
for _ in range(50):
|
||||
await asyncio.sleep(0)
|
||||
if any(
|
||||
isinstance(c.args[0], BusUITaskGroupCompletedMessage)
|
||||
for c in worker.send_bus_message.await_args_list
|
||||
):
|
||||
break
|
||||
else:
|
||||
self.fail("group_completed envelope was not published")
|
||||
|
||||
self.assertNotIn("t1", worker._user_job_groups)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user