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:
@@ -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", ""))
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user