pipecat/tests/test_ui_job_lifecycle.py

#
# 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()