From 26f85687d607e5a9f1d07b55d847adee8de324d7 Mon Sep 17 00:00:00 2001 From: Paul Kompfner Date: Wed, 25 Mar 2026 17:58:29 -0400 Subject: [PATCH] Handle response cancellation by draining before next inference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Instead of trying to filter stale events inline (unreliable — the API doesn't provide a way to correlate events to a specific response), drain remaining events from a cancelled response before starting the next one. On cancellation, send response.cancel and set a drain flag. At the start of the next _process_context, read and discard events until a terminal event arrives, ensuring a clean connection. Falls back to reconnecting if draining times out. --- src/pipecat/services/openai/responses/llm.py | 125 +++++++++++++++++-- tests/test_openai_responses_websocket.py | 103 +++++++++++++++ 2 files changed, 218 insertions(+), 10 deletions(-) diff --git a/src/pipecat/services/openai/responses/llm.py b/src/pipecat/services/openai/responses/llm.py index 184b44d87..e29a0a0ee 100644 --- a/src/pipecat/services/openai/responses/llm.py +++ b/src/pipecat/services/openai/responses/llm.py @@ -6,6 +6,7 @@ """OpenAI Responses API LLM service implementations (WebSocket and HTTP).""" +import asyncio import hashlib import json from contextlib import asynccontextmanager @@ -388,6 +389,11 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): self._previous_input_length: Optional[int] = None self._previous_response_output: Optional[list] = None + # Response cancellation state + self._current_response_id: Optional[str] = None + self._cancel_pending_response: bool = False + self._needs_drain: bool = False + # -- lifecycle ------------------------------------------------------------ async def start(self, frame: StartFrame): @@ -444,6 +450,7 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): await self._websocket.close() self._websocket = None self._clear_previous_response_state() + self._clear_cancellation_state() self._disconnecting = False except Exception as e: await self.push_error(error_msg=f"Error disconnecting from WebSocket: {e}", exception=e) @@ -503,9 +510,8 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): The (possibly modified) params dict. """ if self._previous_response_id is None: - logger.trace( - f"{self}: Sending full context ({len(full_input)} items) — no previous response" - ) + logger.debug(f"{self}: Sending full context ({len(full_input)} items)") + logger.trace(f"{self}: Reason: no previous response") return params if ( @@ -513,18 +519,18 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): or self._previous_input_hash is None or len(full_input) <= self._previous_input_length ): + logger.debug(f"{self}: Sending full context ({len(full_input)} items)") logger.trace( - f"{self}: Sending full context ({len(full_input)} items) — " - f"input not longer than previous ({self._previous_input_length})" + f"{self}: Reason: input not longer than previous ({self._previous_input_length})" ) return params prefix = full_input[: self._previous_input_length] prefix_hash = self._hash_input_items(prefix) if prefix_hash != self._previous_input_hash: + logger.debug(f"{self}: Sending full context ({len(full_input)} items)") logger.trace( - f"{self}: Sending full context ({len(full_input)} items) — " - f"input prefix hash mismatch " + f"{self}: Reason: input prefix hash mismatch " f"(previous input: {json.dumps(prefix, indent=2, default=str)}, " f"expected hash: {self._previous_input_hash}, " f"actual hash: {prefix_hash})" @@ -535,9 +541,9 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): response_output = self._previous_response_output or [] if not self._starts_with_response_output(items_after_prefix, response_output): + logger.debug(f"{self}: Sending full context ({len(full_input)} items)") logger.trace( - f"{self}: Sending full context ({len(full_input)} items) — " - f"response output mismatch after prefix " + f"{self}: Reason: response output mismatch after prefix " f"(previous response output: {json.dumps(response_output, indent=2, default=str)}, " f"items after prefix: {json.dumps(items_after_prefix, indent=2, default=str)})" ) @@ -548,7 +554,7 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): cached = self._previous_input_length + len(response_output) params["input"] = items_to_send params["previous_response_id"] = self._previous_response_id - logger.trace( + logger.debug( f"{self}: Sending incremental context via previous_response_id " f"({len(items_to_send)} new items, {cached} cached)" ) @@ -643,6 +649,68 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): self._previous_input_hash = None self._previous_response_output = None + # -- response cancellation ------------------------------------------------ + + def _clear_cancellation_state(self): + """Clear response cancellation tracking state.""" + self._current_response_id = None + self._cancel_pending_response = False + self._needs_drain = False + + async def _drain_cancelled_response(self): + """Drain events from a cancelled response before starting a new one. + + After a cancellation, the WebSocket may still have in-flight events + from the cancelled response. We must drain them before sending a + new ``response.create`` — we can't simply filter them inline because + the API doesn't provide a reliable way to correlate events to a + specific response (e.g. delta events carry neither a + ``response_id`` nor any intermediary identifier that could be + traced back to one). + + This method reads and discards events until a terminal event + (``response.completed``, ``response.failed``, or + ``response.incomplete``) arrives, ensuring the connection is clean. + Falls back to reconnecting if draining takes too long. + """ + logger.debug(f"{self}: Draining cancelled response events") + try: + while True: + raw = await asyncio.wait_for(self._websocket.recv(), timeout=5.0) + event = json.loads(raw) + event_type = event.get("type") + + # If we were cancelled before response.created, the first + # event here will be response.created for the cancelled + # request — send cancel now that we have the id. + if event_type == "response.created" and self._cancel_pending_response: + response_id = event.get("response", {}).get("id") + logger.debug( + f"{self}: Received response.created for pending-cancel " + f"response {response_id} — sending response.cancel" + ) + self._cancel_pending_response = False + if response_id: + try: + await self._ws_send( + {"type": "response.cancel", "response_id": response_id} + ) + except Exception: + pass + continue + + if event_type in ("response.completed", "response.failed", "response.incomplete"): + logger.debug( + f"{self}: Cancelled response terminated with {event_type} — " + f"connection is clean" + ) + self._needs_drain = False + return + except asyncio.TimeoutError: + logger.warning(f"{self}: Timed out draining cancelled response — reconnecting") + self._needs_drain = False + await self._reconnect() + # -- frame processing ----------------------------------------------------- async def process_frame(self, frame: Frame, direction: FrameDirection): @@ -665,6 +733,33 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): await self.push_frame(LLMFullResponseStartFrame()) await self.start_processing_metrics() await self._process_context(context) + except asyncio.CancelledError: + # The pipeline cancelled us (e.g. due to an interruption). + # Ask the server to stop generating and flag that we need + # to drain stale events before the next inference. We + # can't just send a new response.create and filter stale + # events inline — the API doesn't provide a reliable way + # to correlate events to a specific response. + if self._current_response_id: + logger.debug( + f"{self}: Cancelled during response {self._current_response_id} " + f"— sending response.cancel" + ) + try: + await self._ws_send( + {"type": "response.cancel", "response_id": self._current_response_id} + ) + except Exception: + pass + else: + logger.debug( + f"{self}: Cancelled before response.created " + f"— will cancel on next response.created" + ) + self._cancel_pending_response = True + self._current_response_id = None + self._needs_drain = True + raise except Exception as e: await self.push_error(error_msg=f"Error during completion: {e}", exception=e) finally: @@ -680,6 +775,11 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): Args: context: The LLM context containing conversation history. """ + # If a previous response was cancelled, drain its remaining events + # before starting a new one. + if self._needs_drain: + await self._drain_cancelled_response() + adapter: OpenAIResponsesLLMAdapter = self.get_llm_adapter() logger.debug( f"{self}: Generating response from universal context " @@ -767,6 +867,11 @@ class OpenAIResponsesLLMService(_BaseOpenAIResponsesLLMService): event = json.loads(raw) event_type = event.get("type") + if event_type == "response.created": + self._current_response_id = event.get("response", {}).get("id") + logger.debug(f"{self}: Response started: {self._current_response_id}") + continue + if event_type == "response.output_text.delta": await self.stop_ttfb_metrics() await self._push_llm_text(event.get("delta", "")) diff --git a/tests/test_openai_responses_websocket.py b/tests/test_openai_responses_websocket.py index ef69ec2bd..8623da0ba 100644 --- a/tests/test_openai_responses_websocket.py +++ b/tests/test_openai_responses_websocket.py @@ -6,12 +6,14 @@ """Tests for the WebSocket variant of OpenAIResponsesLLMService.""" +import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch import pytest from pipecat.processors.aggregators.llm_context import LLMContext +from pipecat.processors.frame_processor import FrameDirection from pipecat.services.openai.responses.llm import OpenAIResponsesLLMService @@ -578,6 +580,79 @@ class TestReceiveResponseEventsErrors: assert "Internal server error" in service.push_error.call_args.kwargs["error_msg"] +class TestDrainCancelledResponse: + @pytest.mark.asyncio + async def test_drain_discards_events_until_terminal(self): + """Draining should discard events until a terminal event arrives.""" + service = _make_service() + service._needs_drain = True + + ws = _ws_events( + {"type": "response.output_text.delta", "delta": "stale"}, + {"type": "response.output_text.delta", "delta": "also stale"}, + {"type": "response.completed", "response": {"id": "resp_old"}}, + ) + service._websocket = ws + + await service._drain_cancelled_response() + + assert not service._needs_drain + + @pytest.mark.asyncio + async def test_drain_handles_pending_cancel(self): + """If cancelled before response.created, drain should send cancel + once it sees the response.created, then continue draining.""" + service = _make_service() + service._needs_drain = True + service._cancel_pending_response = True + + mock_ws = AsyncMock() + mock_ws.recv = AsyncMock( + side_effect=[ + json.dumps({"type": "response.created", "response": {"id": "resp_late"}}), + json.dumps({"type": "response.output_text.delta", "delta": "stale"}), + json.dumps({"type": "response.failed", "response": {"id": "resp_late"}}), + ] + ) + mock_ws.send = AsyncMock() + service._websocket = mock_ws + + await service._drain_cancelled_response() + + assert not service._needs_drain + assert not service._cancel_pending_response + # Should have sent response.cancel + cancel_calls = [ + call for call in mock_ws.send.call_args_list if "response.cancel" in call.args[0] + ] + assert len(cancel_calls) == 1 + + @pytest.mark.asyncio + async def test_drain_timeout_triggers_reconnect(self): + """If draining takes too long, should fall back to reconnecting.""" + service = _make_service() + service._needs_drain = True + service.stop_all_metrics = AsyncMock() + service.push_error = AsyncMock() + + mock_ws = AsyncMock() + # recv() never returns a terminal event — times out + mock_ws.recv = AsyncMock(side_effect=asyncio.TimeoutError) + mock_ws.close = AsyncMock() + service._websocket = mock_ws + + with patch( + "pipecat.services.openai.responses.llm.websocket_connect", + new_callable=AsyncMock, + return_value=AsyncMock(), + ): + await service._drain_cancelled_response() + + assert not service._needs_drain + # Should have reconnected (old ws closed) + mock_ws.close.assert_called_once() + + # --------------------------------------------------------------------------- # Connection lifecycle # --------------------------------------------------------------------------- @@ -618,6 +693,34 @@ class TestConnectionLifecycle: assert service._previous_response_id is None mock_ws.close.assert_called_once() + @pytest.mark.asyncio + async def test_cancellation_preserves_connection_and_sets_drain(self): + """When process_frame is cancelled (e.g. interruption), the WebSocket + connection should be preserved and _needs_drain set.""" + service = _make_service() + service.stop_processing_metrics = AsyncMock() + service.push_frame = AsyncMock() + + mock_ws = AsyncMock() + mock_ws.recv = AsyncMock(side_effect=asyncio.CancelledError) + mock_ws.send = AsyncMock() + service._websocket = mock_ws + + context = MagicMock(spec=LLMContext) + context.tools = None + context.tool_choice = None + context.messages = [{"role": "user", "content": "hi"}] + + from pipecat.frames.frames import LLMContextFrame + + with pytest.raises(asyncio.CancelledError): + await service.process_frame(LLMContextFrame(context=context), FrameDirection.DOWNSTREAM) + + # Connection should be preserved, not closed + assert service._websocket is mock_ws + # Should be flagged for draining before next inference + assert service._needs_drain + @pytest.mark.asyncio async def test_ensure_connected_raises_on_failure(self): from pipecat.services.openai.responses.llm import _RetryableError