From 12dcc87030fec464001770796009f41fae88ea1f Mon Sep 17 00:00:00 2001 From: Filipi Fuchter Date: Tue, 23 Sep 2025 16:23:57 -0300 Subject: [PATCH] Fixed an issue in BaseOpenAILLMService that could cause the pipeline to freeze under certain race conditions. --- CHANGELOG.md | 3 + src/pipecat/services/openai/base_llm.py | 29 +++++++-- src/pipecat/utils/cancellable_stream.py | 87 +++++++++++++++++++++++++ 3 files changed, 114 insertions(+), 5 deletions(-) create mode 100644 src/pipecat/utils/cancellable_stream.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ad49ae5d7..76d729817 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -102,6 +102,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Fixed an issue in BaseOpenAILLMService that could cause the pipeline to freeze + under certain race conditions. + - Fixed a `BaseOutputTransport` issue that could produce large saved `AudioBufferProcessor` files when using an audio mixer. diff --git a/src/pipecat/services/openai/base_llm.py b/src/pipecat/services/openai/base_llm.py index 47127b77a..187050d6d 100644 --- a/src/pipecat/services/openai/base_llm.py +++ b/src/pipecat/services/openai/base_llm.py @@ -26,6 +26,7 @@ from pydantic import BaseModel, Field from pipecat.adapters.services.open_ai_adapter import OpenAILLMInvocationParams from pipecat.frames.frames import ( Frame, + InterruptionFrame, LLMContextFrame, LLMFullResponseEndFrame, LLMFullResponseStartFrame, @@ -41,6 +42,7 @@ from pipecat.processors.aggregators.openai_llm_context import ( ) from pipecat.processors.frame_processor import FrameDirection from pipecat.services.llm_service import FunctionCallFromLLM, LLMService +from pipecat.utils.cancellable_stream import CancellableStream from pipecat.utils.tracing.service_decorators import traced_llm @@ -138,6 +140,7 @@ class BaseOpenAILLMService(LLMService): default_headers=default_headers, **kwargs, ) + self.chunk_stream: Optional[CancellableStream[AsyncStream[AsyncStream]]] = None def create_client( self, @@ -329,13 +332,15 @@ class BaseOpenAILLMService(LLMService): await self.start_ttfb_metrics() # Generate chat completions using either OpenAILLMContext or universal LLMContext - chunk_stream = await ( - self._stream_chat_completions_specific_context(context) - if isinstance(context, OpenAILLMContext) - else self._stream_chat_completions_universal_context(context) + self.chunk_stream = CancellableStream( + await ( + self._stream_chat_completions_specific_context(context) + if isinstance(context, OpenAILLMContext) + else self._stream_chat_completions_universal_context(context) + ) ) - async for chunk in chunk_stream: + async for chunk in self.chunk_stream: if chunk.usage: tokens = LLMTokenUsage( prompt_tokens=chunk.usage.prompt_tokens, @@ -389,6 +394,8 @@ class BaseOpenAILLMService(LLMService): ): await self.push_frame(LLMTextFrame(chunk.choices[0].delta.audio["transcript"])) + self.chunk_stream = None + # if we got a function name and arguments, check to see if it's a function with # a registered handler. If so, run the registered callback, save the result to # the context, and re-prompt to get a chat answer. If we don't have a registered @@ -427,6 +434,18 @@ class BaseOpenAILLMService(LLMService): frame: The frame to process. direction: The direction of frame processing. """ + # OpenAI's client swallows asyncio.CancelledError internally, which prevents proper + # task cancellation propagation. To ensure proper cancellation behavior: + # 1. We check if there's an active chunk stream when receiving an interruption + # 2. We explicitly cancel the chunk stream first + # 3. This allows the task to be cancelled cleanly afterwards + # This approach ensures we don't get stuck in case there was a chunk processing loop + # when cancellation is requested. + if isinstance(frame, InterruptionFrame) and self.chunk_stream: + logger.debug(f"{self}: Cancelling chunk stream due to interruption") + await self.chunk_stream.cancel() + self.chunk_stream = None + await super().process_frame(frame, direction) context = None diff --git a/src/pipecat/utils/cancellable_stream.py b/src/pipecat/utils/cancellable_stream.py new file mode 100644 index 000000000..4b2be6645 --- /dev/null +++ b/src/pipecat/utils/cancellable_stream.py @@ -0,0 +1,87 @@ +"""Provides a wrapper class for making async streams cancellable. + +This module implements a CancellableStream class that wraps any async iterator, +adding the ability to safely cancel iteration at any point. This is particularly +useful for handling cleanup of async streams in scenarios where early termination +is required. + +The module provides functionality for: +- Wrapping any async iterator with cancellation capabilities +- Safe termination of async streams +- Proper cleanup and synchronization during cancellation +""" + +import asyncio +from typing import AsyncIterator, Generic, TypeVar + +T = TypeVar("T") + + +class CancellableStream(Generic[T]): + """A wrapper around an async stream that can be cancelled.""" + + def __init__(self, stream: AsyncIterator[T]) -> None: + """Initialize a cancellable stream wrapper. + + Creates a wrapper around an async iterator that can be cancelled mid-iteration. + The wrapper maintains state about cancellation requests and iteration status. + + Args: + stream: The async iterator to wrap. This is the source stream that will + be iterated over until either exhaustion or cancellation. + + Attributes: + _stream: The wrapped async iterator + _cancel_future: A future that completes when cancellation is acknowledged + _cancel_requested: Flag indicating if cancellation has been requested + _iter_started: Flag indicating if iteration has begun + """ + self._stream: AsyncIterator[T] = stream + self._cancel_future: asyncio.Future[None] | None = None + self._cancel_requested: bool = False + self._iter_started: bool = False + + async def cancel(self) -> None: + """Request stream cancellation and wait for acknowledgment. + + Sets up cancellation state and waits until the cancellation is acknowledged + by the next iteration attempt. If iteration hasn't started yet, the + cancellation is acknowledged immediately. + + The method will: + 1. Create a cancellation future if one doesn't exist + 2. Mark the stream as cancelled + 3. If iteration hasn't started, complete the future immediately + 4. Otherwise, wait for the next iteration to acknowledge cancellation + + Returns: + None: The method returns when cancellation is acknowledged. + """ + if self._cancel_future is None: + self._cancel_future = asyncio.get_event_loop().create_future() + self._cancel_requested = True + + # If iteration has not started, we complete the future immediately + if not self._iter_started: + self._cancel_future.set_result(None) + + await self._cancel_future + + def __aiter__(self) -> "CancellableStream[T]": + self._iter_started = True + return self + + async def __anext__(self) -> T: + if self._cancel_requested: + # Complete the future if cancellation was requested + if self._cancel_future and not self._cancel_future.done(): + self._cancel_future.set_result(None) + raise StopAsyncIteration + + try: + return await self._stream.__anext__() + except StopAsyncIteration: + # also complete cancel future if iteration naturally ends + if self._cancel_future and not self._cancel_future.done(): + self._cancel_future.set_result(None) + raise