code-assistant: work around CancelledError swallow in ClaudeSDKClient

claude_agent_sdk's _AsyncioTaskHandle.wait() uses
`with suppress(asyncio.CancelledError)` to silence the inner read
task's expected cancellation, but it also swallows the outer task's
cancellation if it lands on the same await — causing cancel_task to
time out.

Bypass `async with ClaudeSDKClient` and drive connect/disconnect
ourselves so disconnect() runs in a finally where the outer
CancelledError has already been raised and suspended by Python's
exception machinery, out of reach of the SDK's suppress.
This commit is contained in:
Aleix Conchillo Flaqué
2026-05-14 22:44:06 -07:00
parent a5bb9f65de
commit de1bd7cb7e

View File

@@ -10,11 +10,7 @@ import asyncio
from loguru import logger from loguru import logger
from pipecat.bus import ( from pipecat.bus import BusJobRequestMessage
BusCancelTaskMessage,
BusEndTaskMessage,
BusJobRequestMessage,
)
from pipecat.pipeline.base_task import BaseTask from pipecat.pipeline.base_task import BaseTask
from pipecat.pipeline.job_context import JobStatus from pipecat.pipeline.job_context import JobStatus
@@ -66,7 +62,7 @@ class CodeWorker(BaseTask):
async def start(self) -> None: async def start(self) -> None:
"""Launch the Claude SDK worker loop alongside the standard task start.""" """Launch the Claude SDK worker loop alongside the standard task start."""
await super().start() await super().start()
self._worker_task = self.create_task(self._worker_loop(), f"{self.name}::worker") self._worker_task = self.create_task(self._worker_loop(), "worker")
async def stop(self) -> None: async def stop(self) -> None:
"""Cancel the worker loop before tearing down the task.""" """Cancel the worker loop before tearing down the task."""
@@ -81,19 +77,15 @@ class CodeWorker(BaseTask):
logger.info(f"Worker '{self.name}': queued '{message.payload['question']}'") logger.info(f"Worker '{self.name}': queued '{message.payload['question']}'")
self._queue.put_nowait(message) self._queue.put_nowait(message)
async def _handle_task_end(self, message: BusEndTaskMessage) -> None:
"""Signal the run loop to finish on a graceful end."""
await super()._handle_task_end(message)
self._finished_event.set()
async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None:
"""Signal the run loop to finish on cancellation."""
await super()._handle_task_cancel(message)
self._finished_event.set()
async def _worker_loop(self): async def _worker_loop(self):
client = ClaudeSDKClient(options=self._claude_options)
try:
await client.connect()
except Exception as e:
logger.error(f"Worker '{self.name}': failed to start Claude SDK: {e}")
return
try: try:
async with ClaudeSDKClient(options=self._claude_options) as client:
while True: while True:
message = await self._queue.get() message = await self._queue.get()
question = message.payload["question"] question = message.payload["question"]
@@ -116,5 +108,11 @@ class CodeWorker(BaseTask):
await self.send_job_response( await self.send_job_response(
message.job_id, {"error": str(e)}, status=JobStatus.ERROR message.job_id, {"error": str(e)}, status=JobStatus.ERROR
) )
except Exception as e: finally:
logger.error(f"Worker '{self.name}': failed to start Claude SDK: {e}") # Bypass `async with ClaudeSDKClient` and call disconnect()
# ourselves: __aexit__ → Query.close() → _read_task.wait() uses
# `with suppress(asyncio.CancelledError)`, which would swallow the
# outer task's cancellation. By the time this finally runs, our
# CancelledError has already been raised once, so _must_cancel is
# cleared and disconnect()'s awaits proceed normally.
await client.disconnect()