b03247f360
Replaces every "task" identifier that referred to the BaseTask abstraction with "worker". Asyncio task plumbing (asyncio.Task, BaseTaskManager, TaskManager, create_task, cancel_task, etc.) stays untouched. Highlights: - Classes: BaseTask → BaseWorker, PipelineTask → PipelineWorker, LLMTask → LLMWorker, LLMContextTask → LLMContextWorker, TaskBus → WorkerBus, TaskRegistry → WorkerRegistry, TaskActivationArgs → WorkerActivationArgs, TaskReadyData → WorkerReadyData, TaskRegistryEntry → WorkerRegistryEntry, TaskObserver → WorkerObserver, all Bus*TaskMessage → Bus*WorkerMessage, BusAddTaskMessage.task field → worker, BusWorkerRegistryMessage.tasks field → workers. - Methods/decorators: activate_task → activate_worker, deactivate_task → deactivate_worker, add_task → add_worker, watch_task → watch_worker, @task_ready → @worker_ready, setup_pipeline_task hook → setup_pipeline_worker. - Params/fields: FrameProcessorSetup.pipeline_task and FunctionCallParams.pipeline_task → pipeline_worker. Parameter names like task_name → worker_name; spawn/run accept worker:. - Files: pipeline/base_task.py → base_worker.py, pipeline/task.py → worker.py (plus a re-export shim at pipeline/task.py), task_observer.py → worker_observer.py, task_ready_decorator.py → worker_ready_decorator.py, pipecat.tasks → pipecat.workers, llm_task.py → llm_worker.py, llm_context_task.py → llm_context_worker.py, examples/multi-task → examples/multi-worker. Back-compat: - PipelineTask kept as a deprecated subclass of PipelineWorker that warns on construction. - pipecat.pipeline.task re-exports PipelineWorker/PipelineTask/etc. so existing user imports keep working. - FrameProcessor.pipeline_task kept as a deprecated property that forwards to pipeline_worker. Local variables in examples that hold a worker (task = PipelineTask(...)) are renamed to worker = PipelineWorker(...). Asyncio-task locals (runner_task, etc.) are preserved.
740 lines
26 KiB
Python
740 lines
26 KiB
Python
#
|
|
# Copyright (c) 2026, Daily
|
|
#
|
|
# SPDX-License-Identifier: BSD 2-Clause License
|
|
#
|
|
|
|
import asyncio
|
|
import unittest
|
|
|
|
from pipecat.bus import (
|
|
AsyncQueueBus,
|
|
BusJobCancelMessage,
|
|
BusJobRequestMessage,
|
|
BusJobResponseMessage,
|
|
)
|
|
from pipecat.pipeline.base_worker import BaseWorker
|
|
from pipecat.pipeline.job_context import (
|
|
JobError,
|
|
JobEvent,
|
|
JobGroupError,
|
|
JobGroupEvent,
|
|
JobStatus,
|
|
)
|
|
from pipecat.registry import WorkerRegistry
|
|
from pipecat.registry.types import WorkerReadyData
|
|
from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams
|
|
|
|
|
|
class StubTask(BaseWorker):
|
|
pass
|
|
|
|
|
|
class JobWorkerTask(BaseWorker):
|
|
"""Worker that automatically responds to job requests via the bus."""
|
|
|
|
def __init__(self, name, *, response=None, status=JobStatus.COMPLETED):
|
|
super().__init__(name)
|
|
self._auto_response = response
|
|
self._auto_status = status
|
|
|
|
async def on_job_request(self, message):
|
|
await super().on_job_request(message)
|
|
await self.send_job_response(message.job_id, self._auto_response, status=self._auto_status)
|
|
|
|
|
|
class UrgentJobWorkerTask(BaseWorker):
|
|
"""Worker that responds urgently to job requests."""
|
|
|
|
def __init__(self, name, *, response=None, status=JobStatus.COMPLETED):
|
|
super().__init__(name)
|
|
self._auto_response = response
|
|
self._auto_status = status
|
|
|
|
async def on_job_request(self, message):
|
|
await super().on_job_request(message)
|
|
await self.send_job_response(
|
|
message.job_id, self._auto_response, status=self._auto_status, urgent=True
|
|
)
|
|
|
|
|
|
class UpdatingWorkerTask(BaseWorker):
|
|
"""Worker that sends updates before responding."""
|
|
|
|
def __init__(self, name, *, updates, response=None):
|
|
super().__init__(name)
|
|
self._updates = updates
|
|
self._auto_response = response
|
|
|
|
async def on_job_request(self, message):
|
|
await super().on_job_request(message)
|
|
for update in self._updates:
|
|
await self.send_job_update(message.job_id, update)
|
|
await self.send_job_response(message.job_id, self._auto_response)
|
|
|
|
|
|
class StreamingWorkerTask(BaseWorker):
|
|
"""Worker that streams data before responding."""
|
|
|
|
def __init__(self, name, *, chunks, response=None):
|
|
super().__init__(name)
|
|
self._chunks = chunks
|
|
self._auto_response = response
|
|
|
|
async def on_job_request(self, message):
|
|
await super().on_job_request(message)
|
|
await self.send_job_stream_start(message.job_id, {"content_type": "text"})
|
|
for chunk in self._chunks:
|
|
await self.send_job_stream_data(message.job_id, chunk)
|
|
await self.send_job_stream_end(message.job_id, self._auto_response)
|
|
|
|
|
|
class SlowWorkerTask(BaseWorker):
|
|
"""Worker that blocks during job execution until cancelled."""
|
|
|
|
def __init__(self, name):
|
|
super().__init__(name)
|
|
self.started = asyncio.Event()
|
|
self.was_cancelled = False
|
|
|
|
async def on_job_request(self, message):
|
|
await super().on_job_request(message)
|
|
self.started.set()
|
|
try:
|
|
# Block until cancelled
|
|
await asyncio.Event().wait()
|
|
except asyncio.CancelledError:
|
|
self.was_cancelled = True
|
|
|
|
|
|
async def create_test_env():
|
|
bus = AsyncQueueBus()
|
|
tm = TaskManager()
|
|
tm.setup(TaskManagerParams(loop=asyncio.get_running_loop()))
|
|
await bus.setup(tm)
|
|
await bus.start()
|
|
registry = WorkerRegistry(runner_name="test-runner")
|
|
return bus, tm, registry
|
|
|
|
|
|
async def setup_task(bus, registry, task):
|
|
"""Subscribe a task to the bus and register it as ready."""
|
|
task.attach(registry=registry, bus=bus)
|
|
await task.setup(bus.task_manager)
|
|
await bus.subscribe(task)
|
|
await registry.register(WorkerReadyData(worker_name=task.name, runner="test-runner"))
|
|
|
|
|
|
def capture_bus(bus):
|
|
sent = []
|
|
original_send = bus.send
|
|
|
|
async def capture_send(message):
|
|
sent.append(message)
|
|
await original_send(message)
|
|
|
|
bus.send = capture_send
|
|
return sent
|
|
|
|
|
|
class TestJobGroupContext(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.bus, self.tm, self.registry = await create_test_env()
|
|
|
|
async def asyncTearDown(self):
|
|
await self.bus.stop()
|
|
|
|
async def test_job_group_collects_responses(self):
|
|
"""job_group() context manager collects all responses."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
w1 = JobWorkerTask("w1", response={"a": 1})
|
|
w2 = JobWorkerTask("w2", response={"b": 2})
|
|
await setup_task(self.bus, self.registry, w1)
|
|
await setup_task(self.bus, self.registry, w2)
|
|
|
|
async with parent.job_group("w1", "w2", payload={"work": True}) as tg:
|
|
pass
|
|
|
|
self.assertEqual(tg.responses, {"w1": {"a": 1}, "w2": {"b": 2}})
|
|
|
|
async def test_job_group_sends_request(self):
|
|
"""job_group() sends BusJobRequestMessage to each task."""
|
|
sent = capture_bus(self.bus)
|
|
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
w1 = JobWorkerTask("w1", response={"ok": True})
|
|
await setup_task(self.bus, self.registry, w1)
|
|
|
|
async with parent.job_group("w1", payload={"data": 1}) as tg:
|
|
pass
|
|
|
|
request_msgs = [m for m in sent if isinstance(m, BusJobRequestMessage)]
|
|
self.assertEqual(len(request_msgs), 1)
|
|
self.assertEqual(request_msgs[0].payload, {"data": 1})
|
|
|
|
async def test_job_group_raises_on_cancel(self):
|
|
"""job_group() raises JobGroupError when group is cancelled."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
# StubTask doesn't auto-respond, so we can cancel manually
|
|
worker = StubTask("worker")
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
with self.assertRaises(JobGroupError) as ctx:
|
|
async with parent.job_group("worker") as tg:
|
|
asyncio.ensure_future(parent.cancel_job_group(tg.job_id, reason="manual cancel"))
|
|
|
|
# Let the event loop schedule handler tasks spawned by the bus
|
|
await asyncio.sleep(0)
|
|
self.assertIn("manual cancel", str(ctx.exception))
|
|
|
|
async def test_job_group_raises_on_timeout(self):
|
|
"""job_group() raises JobGroupError on timeout."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
# StubTask doesn't auto-respond, so timeout fires
|
|
worker = StubTask("worker")
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
with self.assertRaises(JobGroupError) as ctx:
|
|
async with parent.job_group("worker", timeout=0.05) as tg:
|
|
pass
|
|
|
|
self.assertIn("timeout", str(ctx.exception))
|
|
|
|
async def test_job_group_raises_on_ready_timeout(self):
|
|
"""job_group() raises JobGroupError when tasks aren't ready in time."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
# "ghost" is never registered, so the ready-wait times out
|
|
with self.assertRaises(JobGroupError) as ctx:
|
|
async with parent.job_group("ghost", timeout=0.05) as tg:
|
|
pass
|
|
|
|
self.assertIn("not ready", str(ctx.exception))
|
|
|
|
async def test_job_group_cancels_on_block_exception(self):
|
|
"""job_group() cancels remaining jobs when the block raises."""
|
|
sent = capture_bus(self.bus)
|
|
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = StubTask("worker")
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
with self.assertRaises(ValueError):
|
|
async with parent.job_group("worker") as tg:
|
|
raise ValueError("something went wrong")
|
|
|
|
cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)]
|
|
self.assertEqual(len(cancel_msgs), 1)
|
|
self.assertEqual(cancel_msgs[0].reason, "context exited with error")
|
|
|
|
async def test_job_group_raises_on_worker_error(self):
|
|
"""job_group() raises JobGroupError when a worker errors with cancel_on_error."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = JobWorkerTask("worker", response={"error": "failed"}, status=JobStatus.ERROR)
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
with self.assertRaises(JobGroupError):
|
|
async with parent.job_group("worker") as tg:
|
|
pass
|
|
|
|
# Error response is tracked in partial responses
|
|
self.assertEqual(tg.responses, {"worker": {"error": "failed"}})
|
|
|
|
async def test_job_group_on_job_error_fires(self):
|
|
"""on_job_error fires when a worker errors with cancel_on_error."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = JobWorkerTask("worker", response={"error": "boom"}, status=JobStatus.ERROR)
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
errors = []
|
|
|
|
@parent.event_handler("on_job_error")
|
|
async def on_error(task, message):
|
|
errors.append(message)
|
|
|
|
with self.assertRaises(JobGroupError):
|
|
async with parent.job_group("worker") as tg:
|
|
pass
|
|
|
|
self.assertEqual(len(errors), 1)
|
|
self.assertEqual(errors[0].source, "worker")
|
|
self.assertEqual(errors[0].response, {"error": "boom"})
|
|
self.assertEqual(errors[0].status, JobStatus.ERROR)
|
|
|
|
async def test_job_group_partial_responses_on_error(self):
|
|
"""Partial responses from successful workers are available after error."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
# w1 responds successfully, w2 responds with error.
|
|
# Order depends on bus dispatch, but both are registered.
|
|
w1 = JobWorkerTask("w1", response={"ok": True})
|
|
w2 = JobWorkerTask("w2", response={"error": "fail"}, status=JobStatus.ERROR)
|
|
await setup_task(self.bus, self.registry, w1)
|
|
await setup_task(self.bus, self.registry, w2)
|
|
|
|
with self.assertRaises(JobGroupError):
|
|
async with parent.job_group("w1", "w2") as tg:
|
|
pass
|
|
|
|
# w2's error response should be in partial responses
|
|
self.assertIn("w2", tg.responses)
|
|
self.assertEqual(tg.responses["w2"], {"error": "fail"})
|
|
|
|
async def test_job_group_job_id_available(self):
|
|
"""job_id is available inside the async with block."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = JobWorkerTask("worker", response={})
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
captured_job_id = None
|
|
|
|
async with parent.job_group("worker") as tg:
|
|
captured_job_id = tg.job_id
|
|
|
|
self.assertIsNotNone(captured_job_id)
|
|
self.assertEqual(captured_job_id, tg.job_id)
|
|
|
|
async def test_job_group_on_job_completed_still_fires(self):
|
|
"""on_job_completed callback still fires when using job_group()."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = JobWorkerTask("w1", response={"ok": True})
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
completed = []
|
|
|
|
@parent.event_handler("on_job_completed")
|
|
async def on_completed(task, result):
|
|
completed.append(result)
|
|
|
|
async with parent.job_group("w1") as tg:
|
|
pass
|
|
|
|
self.assertEqual(len(completed), 1)
|
|
self.assertEqual(completed[0].responses, {"w1": {"ok": True}})
|
|
|
|
async def test_job_group_iterates_updates(self):
|
|
"""async for yields update events from workers."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = UpdatingWorkerTask(
|
|
"worker",
|
|
updates=[{"progress": 25}, {"progress": 75}],
|
|
response={"result": "done"},
|
|
)
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
events = []
|
|
async with parent.job_group("worker", payload={"work": True}) as tg:
|
|
async for event in tg:
|
|
events.append(event)
|
|
|
|
self.assertEqual(len(events), 2)
|
|
self.assertEqual(events[0].type, JobGroupEvent.UPDATE)
|
|
self.assertEqual(events[0].worker_name, "worker")
|
|
self.assertEqual(events[0].data, {"progress": 25})
|
|
self.assertEqual(events[1].data, {"progress": 75})
|
|
self.assertEqual(tg.responses, {"worker": {"result": "done"}})
|
|
|
|
async def test_job_group_iterates_stream_events(self):
|
|
"""async for yields stream_start, stream_data, and stream_end events."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = StreamingWorkerTask(
|
|
"worker",
|
|
chunks=[{"text": "hello "}, {"text": "world"}],
|
|
response={"final": True},
|
|
)
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
events = []
|
|
async with parent.job_group("worker") as tg:
|
|
async for event in tg:
|
|
events.append(event)
|
|
|
|
types = [e.type for e in events]
|
|
self.assertEqual(
|
|
types,
|
|
[
|
|
JobGroupEvent.STREAM_START,
|
|
JobGroupEvent.STREAM_DATA,
|
|
JobGroupEvent.STREAM_DATA,
|
|
JobGroupEvent.STREAM_END,
|
|
],
|
|
)
|
|
self.assertEqual(events[0].data, {"content_type": "text"})
|
|
self.assertEqual(events[1].data, {"text": "hello "})
|
|
self.assertEqual(events[2].data, {"text": "world"})
|
|
self.assertEqual(events[3].data, {"final": True})
|
|
|
|
async def test_job_group_iterates_mixed_events(self):
|
|
"""async for yields events from multiple workers interleaved."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
w1 = UpdatingWorkerTask("w1", updates=[{"status": "working"}], response={"a": 1})
|
|
w2 = JobWorkerTask("w2", response={"b": 2})
|
|
await setup_task(self.bus, self.registry, w1)
|
|
await setup_task(self.bus, self.registry, w2)
|
|
|
|
events = []
|
|
async with parent.job_group("w1", "w2") as tg:
|
|
async for event in tg:
|
|
events.append(event)
|
|
|
|
self.assertEqual(len(events), 1)
|
|
self.assertEqual(events[0].type, JobGroupEvent.UPDATE)
|
|
self.assertEqual(events[0].worker_name, "w1")
|
|
self.assertEqual(tg.responses, {"w1": {"a": 1}, "w2": {"b": 2}})
|
|
|
|
async def test_job_group_no_iteration_still_works(self):
|
|
"""job_group() works without iterating (pass body)."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = UpdatingWorkerTask("worker", updates=[{"progress": 50}], response={"ok": True})
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
async with parent.job_group("worker") as tg:
|
|
pass
|
|
|
|
self.assertEqual(tg.responses, {"worker": {"ok": True}})
|
|
|
|
async def test_urgent_job_response_collected(self):
|
|
"""Urgent job responses are collected like normal responses."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = UrgentJobWorkerTask("worker", response={"urgent": True})
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
async with parent.job_group("worker") as tg:
|
|
pass
|
|
|
|
self.assertEqual(tg.responses, {"worker": {"urgent": True}})
|
|
|
|
async def test_urgent_job_response_triggers_on_job_error(self):
|
|
"""Urgent error response triggers on_job_error and cancels group."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = UrgentJobWorkerTask(
|
|
"worker", response={"error": "critical"}, status=JobStatus.ERROR
|
|
)
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
errors = []
|
|
|
|
@parent.event_handler("on_job_error")
|
|
async def on_error(task, message):
|
|
errors.append(message)
|
|
|
|
with self.assertRaises(JobGroupError):
|
|
async with parent.job_group("worker") as tg:
|
|
pass
|
|
|
|
self.assertEqual(len(errors), 1)
|
|
self.assertEqual(errors[0].response, {"error": "critical"})
|
|
|
|
async def test_urgent_response_has_system_priority(self):
|
|
"""Urgent job response is delivered before normal data messages."""
|
|
from pipecat.bus import BusDataMessage, BusJobResponseUrgentMessage
|
|
|
|
bus = self.bus
|
|
|
|
# Queue data messages first, then an urgent response
|
|
parent = StubTask("parent")
|
|
await setup_task(bus, self.registry, parent)
|
|
|
|
received = []
|
|
|
|
@parent.event_handler("on_bus_message")
|
|
async def on_msg(task, message):
|
|
received.append(message)
|
|
|
|
# Send data messages before starting dispatch
|
|
for i in range(3):
|
|
await bus.send(BusDataMessage(source=f"data_{i}"))
|
|
await bus.send(
|
|
BusJobResponseUrgentMessage(
|
|
source="worker", target="parent", job_id="t1", status=JobStatus.COMPLETED
|
|
)
|
|
)
|
|
|
|
# Start dispatch — urgent should arrive first
|
|
# bus is already started in setUp; restart by stopping and starting
|
|
await bus.stop()
|
|
await bus.start()
|
|
await asyncio.sleep(0.1)
|
|
await bus.stop()
|
|
|
|
# First message should be the urgent response
|
|
self.assertIsInstance(received[0], BusJobResponseUrgentMessage)
|
|
|
|
|
|
class TestJobContext(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.bus, self.tm, self.registry = await create_test_env()
|
|
|
|
async def asyncTearDown(self):
|
|
await self.bus.stop()
|
|
|
|
async def test_job_collects_response(self):
|
|
"""job() context manager collects the worker's response."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = JobWorkerTask("worker", response={"result": 42})
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
async with parent.job("worker", payload={"x": 1}) as t:
|
|
pass
|
|
|
|
self.assertEqual(t.response, {"result": 42})
|
|
|
|
async def test_job_sends_request(self):
|
|
"""job() sends a BusJobRequestMessage to the task."""
|
|
sent = capture_bus(self.bus)
|
|
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = JobWorkerTask("worker", response={"ok": True})
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
async with parent.job("worker") as t:
|
|
pass
|
|
|
|
request_msgs = [m for m in sent if isinstance(m, BusJobRequestMessage)]
|
|
self.assertEqual(len(request_msgs), 1)
|
|
self.assertEqual(request_msgs[0].target, "worker")
|
|
self.assertEqual(request_msgs[0].source, "parent")
|
|
|
|
async def test_job_iterates_events(self):
|
|
"""job() yields intermediate events via async for."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = UpdatingWorkerTask("worker", updates=[{"progress": 50}], response={"done": True})
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
events = []
|
|
async with parent.job("worker") as t:
|
|
async for event in t:
|
|
events.append(event)
|
|
|
|
self.assertEqual(len(events), 1)
|
|
self.assertEqual(events[0].type, JobEvent.UPDATE)
|
|
self.assertIsInstance(events[0], JobEvent)
|
|
self.assertEqual(t.response, {"done": True})
|
|
|
|
async def test_job_streams_data(self):
|
|
"""job() yields stream events from a streaming worker."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = StreamingWorkerTask(
|
|
"worker",
|
|
chunks=[{"text": "hello"}, {"text": "world"}],
|
|
response={"ok": True},
|
|
)
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
events = []
|
|
async with parent.job("worker") as t:
|
|
async for event in t:
|
|
events.append(event)
|
|
|
|
types = [e.type for e in events]
|
|
self.assertEqual(
|
|
types,
|
|
[
|
|
JobEvent.STREAM_START,
|
|
JobEvent.STREAM_DATA,
|
|
JobEvent.STREAM_DATA,
|
|
JobEvent.STREAM_END,
|
|
],
|
|
)
|
|
self.assertEqual(t.response, {"ok": True})
|
|
|
|
async def test_job_cancels_on_exception(self):
|
|
"""job() cancels the job if the block raises."""
|
|
sent = capture_bus(self.bus)
|
|
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = StubTask("worker")
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
with self.assertRaises(ValueError):
|
|
async with parent.job("worker") as t:
|
|
raise ValueError("something went wrong")
|
|
|
|
cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)]
|
|
self.assertEqual(len(cancel_msgs), 1)
|
|
self.assertEqual(cancel_msgs[0].reason, "context exited with error")
|
|
|
|
async def test_job_raises_on_worker_error(self):
|
|
"""job() raises JobError when the worker errors."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = JobWorkerTask("worker", response={"error": "boom"}, status=JobStatus.ERROR)
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
with self.assertRaises(JobError):
|
|
async with parent.job("worker") as t:
|
|
pass
|
|
|
|
async def test_job_exposes_job_id(self):
|
|
"""job() exposes the job_id inside the context."""
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = JobWorkerTask("worker", response={"ok": True})
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
async with parent.job("worker") as t:
|
|
self.assertIsInstance(t.job_id, str)
|
|
self.assertTrue(len(t.job_id) > 0)
|
|
|
|
async def test_job_group_cancels_on_cancelled_error(self):
|
|
"""job_group() cancels workers when CancelledError is raised (e.g. tool interruption)."""
|
|
sent = capture_bus(self.bus)
|
|
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = StubTask("worker")
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
async def run_job_group():
|
|
async with parent.job_group("worker") as tg:
|
|
# Simulate tool cancellation while waiting
|
|
raise asyncio.CancelledError()
|
|
|
|
task = asyncio.create_task(run_job_group())
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await task
|
|
|
|
cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)]
|
|
self.assertEqual(len(cancel_msgs), 1)
|
|
self.assertEqual(cancel_msgs[0].reason, "context exited with error")
|
|
# Job group should be cleaned up
|
|
self.assertEqual(len(parent.job_groups), 0)
|
|
|
|
async def test_job_cancels_on_cancelled_error(self):
|
|
"""job() cancels the worker when CancelledError is raised (e.g. tool interruption)."""
|
|
sent = capture_bus(self.bus)
|
|
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = StubTask("worker")
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
async def run_job():
|
|
async with parent.job("worker") as t:
|
|
raise asyncio.CancelledError()
|
|
|
|
task = asyncio.create_task(run_job())
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await task
|
|
|
|
cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)]
|
|
self.assertEqual(len(cancel_msgs), 1)
|
|
self.assertEqual(cancel_msgs[0].reason, "context exited with error")
|
|
self.assertEqual(len(parent.job_groups), 0)
|
|
|
|
async def test_fire_and_forget_jobs_cancelled_manually(self):
|
|
"""User cancels fire-and-forget jobs in a CancelledError handler."""
|
|
sent = capture_bus(self.bus)
|
|
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
w1 = StubTask("w1")
|
|
w2 = StubTask("w2")
|
|
await setup_task(self.bus, self.registry, w1)
|
|
await setup_task(self.bus, self.registry, w2)
|
|
|
|
# Simulate what a user would do in on_function_calls_cancelled:
|
|
# track job IDs and cancel only the ones started here
|
|
try:
|
|
job_ids = []
|
|
job_ids.append(await parent.request_job("w1", payload={"job": 1}))
|
|
job_ids.append(await parent.request_job("w2", payload={"job": 2}))
|
|
self.assertEqual(len(parent.job_groups), 2)
|
|
raise asyncio.CancelledError()
|
|
except asyncio.CancelledError:
|
|
for jid in job_ids:
|
|
await parent.cancel_job_group(jid, reason="tool cancelled")
|
|
|
|
cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)]
|
|
self.assertEqual(len(cancel_msgs), 2)
|
|
cancelled_targets = {m.target for m in cancel_msgs}
|
|
self.assertEqual(cancelled_targets, {"w1", "w2"})
|
|
for m in cancel_msgs:
|
|
self.assertEqual(m.reason, "tool cancelled")
|
|
self.assertEqual(len(parent.job_groups), 0)
|
|
|
|
async def test_cancel_interrupts_running_handler(self):
|
|
"""Cancelling a job interrupts a handler that is currently executing."""
|
|
sent = capture_bus(self.bus)
|
|
|
|
parent = StubTask("parent")
|
|
await setup_task(self.bus, self.registry, parent)
|
|
|
|
worker = SlowWorkerTask("worker")
|
|
await setup_task(self.bus, self.registry, worker)
|
|
|
|
job_id = await parent.request_job("worker", payload={"job": 1})
|
|
|
|
# Wait for the worker to start executing
|
|
await asyncio.wait_for(worker.started.wait(), timeout=2.0)
|
|
|
|
# Cancel while the handler is blocked
|
|
await parent.cancel_job_group(job_id, reason="no longer needed")
|
|
|
|
# Give the event loop time to process the cancellation
|
|
await asyncio.sleep(0.1)
|
|
|
|
# The handler should have been cancelled
|
|
self.assertTrue(worker.was_cancelled)
|
|
|
|
# A cancel message should have been sent
|
|
cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)]
|
|
self.assertEqual(len(cancel_msgs), 1)
|
|
self.assertEqual(cancel_msgs[0].reason, "no longer needed")
|
|
|
|
# A CANCELLED response should have been sent back
|
|
response_msgs = [m for m in sent if isinstance(m, BusJobResponseMessage)]
|
|
self.assertEqual(len(response_msgs), 1)
|
|
self.assertEqual(response_msgs[0].status, JobStatus.CANCELLED)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|