Handle response cancellation by draining before next inference

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.
This commit is contained in:
Paul Kompfner
2026-03-25 17:58:29 -04:00
parent 670ce30a1c
commit 26f85687d6
2 changed files with 218 additions and 10 deletions

View File

@@ -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", ""))

View File

@@ -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